nangs 0.0.8

Solving Partial Differential Equations with Neural Networks

Welcome to nangs

Solving Partial Differential Equations with Neural Networks.

Nangs is a Python library built on top of Pytorch to solve Partial Differential Equations.

Our objective is to develop a new tool for simulating nature, using Neural Networks as solution approximation to Partial Differential Equations, increasing accuracy and optimziation speed while reducing computational cost.

Read our paper to know more.

Read the docs.

Installing

nangs is on PyPI so you can just run:

pip install nangs

You will also need to insall Pytorch.

Getting Started

Let's assume we want to solve the following PDE:

Different numerical techniques that solve this problem exist, and all of them are based on finding an approximate function that satisfies the PDE. Traditional numerical methods discretize the domain into small elements where a form of the solutions is assumed (for example, a constant) and then the final solution is composed as a piece-wise, discontinuous function.

Nangs uses the property of neural networks (NNs) as universal function approximators to find a continuous and derivable solution to the PDE, that requires significant less computing resources compared with traditional techniques and with the advantage of including the free-parameters as part of the solution.

The independen variables (i.e, x and t) are used as input values for the NN, and the solution (i.e. p) is the output. In order to find the solution, at each step the NN outputs are derived w.r.t the inputs. Then, a loss function that matches the PDE is built and the weights are updated accordingly. If the loss function goes to zero, we can assume that our NN is indeed the solution to our PDE.

import math
import numpy as np 
import matplotlib.pyplot as plt 
import torch

# import nangs
from nangs.pde import PDE
from nangs.bocos import PeriodicBoco, DirichletBoco
from nangs.solutions import MLP

# define custom PDE
class MyPDE(PDE):
    def __init__(self, inputs, outputs, params=None):
        super().__init__(inputs, outputs, params)
    def computePDELoss(self, grads, inputs, outputs, params): 
        # here is where the magic happens
        dpdt, dpdx = grads['p']['t'], grads['p']['x']
        u = params['u']
        return [dpdt + u*dpdx]

# instanciate pde
pde = MyPDE(inputs=['x', 't'], outputs=['p'], params=['u'])

# define input values for training
x = np.linspace(0,1,40)
t = np.linspace(0,1,30)
u = np.array([1.0])
pde.setValues({'x': x, 't': t, 'u': u})

# define input values for testing
x_v = np.linspace(0,1,25)
t_v = np.linspace(0,1,15)
pde.setValues({'x': x_v, 't': t_v}, train=False)

# periodic b.c for the space dimension
x1, x2 = np.array([0]), np.array([1])
boco = PeriodicBoco('boco', {'x': x1, 't': t}, {'x': x2, 't': t})
pde.addBoco(boco)

# initial condition (dirichlet for temporal dimension)
p0 = np.sin(2.*math.pi*x)
boco = DirichletBoco('initial_condition', {'x': x, 't': np.array([0])}, {'p': p0})
pde.addBoco(boco)

# define solution topology
mlp = MLP(pde.n_inputs, pde.n_outputs, 3, 256)
optimizer = torch.optim.Adam(mlp.parameters(), lr=3e-4)
pde.compile(mlp, optimizer)

# find the solution
hist = pde.solve()

Epoch 1/30 Losses 0.40837 PDE [ 0.00777 ] boco 0.02863 initial_condition 0.37197 Val [ 0.01869 ]

Epoch 2/30 Losses 0.27972 PDE [ 0.02495 ] boco 0.05397 initial_condition 0.20080 Val [ 0.02539 ]

Epoch 3/30 Losses 0.15535 PDE [ 0.02909 ] boco 0.03968 initial_condition 0.08659 Val [ 0.02680 ]

Epoch 4/30 Losses 0.08142 PDE [ 0.02092 ] boco 0.03128 initial_condition 0.02922 Val [ 0.01671 ]

Epoch 5/30 Losses 0.04809 PDE [ 0.01526 ] boco 0.02045 initial_condition 0.01238 Val [ 0.01072 ]

Epoch 6/30 Losses 0.03523 PDE [ 0.01239 ] boco 0.01436 initial_condition 0.00848 Val [ 0.00939 ]

Epoch 7/30 Losses 0.02426 PDE [ 0.00980 ] boco 0.00917 initial_condition 0.00529 Val [ 0.00886 ]

Epoch 8/30 Losses 0.01611 PDE [ 0.00713 ] boco 0.00579 initial_condition 0.00319 Val [ 0.00503 ]

Epoch 9/30 Losses 0.00978 PDE [ 0.00515 ] boco 0.00300 initial_condition 0.00164 Val [ 0.00502 ]

Epoch 10/30 Losses 0.00626 PDE [ 0.00340 ] boco 0.00183 initial_condition 0.00103 Val [ 0.00265 ]

Epoch 11/30 Losses 0.00404 PDE [ 0.00265 ] boco 0.00091 initial_condition 0.00049 Val [ 0.00240 ]

Epoch 12/30 Losses 0.00291 PDE [ 0.00217 ] boco 0.00046 initial_condition 0.00028 Val [ 0.00274 ]

Epoch 13/30 Losses 0.00270 PDE [ 0.00216 ] boco 0.00033 initial_condition 0.00020 Val [ 0.00160 ]

Epoch 14/30 Losses 0.00212 PDE [ 0.00165 ] boco 0.00028 initial_condition 0.00019 Val [ 0.00181 ]

Epoch 15/30 Losses 0.00193 PDE [ 0.00151 ] boco 0.00026 initial_condition 0.00015 Val [ 0.00157 ]

Epoch 16/30 Losses 0.00188 PDE [ 0.00150 ] boco 0.00022 initial_condition 0.00016 Val [ 0.00158 ]

Epoch 17/30 Losses 0.00171 PDE [ 0.00122 ] boco 0.00033 initial_condition 0.00016 Val [ 0.00114 ]

Epoch 18/30 Losses 0.00121 PDE [ 0.00094 ] boco 0.00019 initial_condition 0.00008 Val [ 0.00102 ]

Epoch 19/30 Losses 0.00110 PDE [ 0.00085 ] boco 0.00017 initial_condition 0.00009 Val [ 0.00096 ]

Epoch 20/30 Losses 0.00116 PDE [ 0.00080 ] boco 0.00026 initial_condition 0.00010 Val [ 0.00078 ]

Epoch 21/30 Losses 0.00092 PDE [ 0.00073 ] boco 0.00012 initial_condition 0.00007 Val [ 0.00073 ]

Epoch 22/30 Losses 0.00085 PDE [ 0.00068 ] boco 0.00010 initial_condition 0.00006 Val [ 0.00074 ]

Epoch 23/30 Losses 0.00101 PDE [ 0.00065 ] boco 0.00027 initial_condition 0.00009 Val [ 0.00071 ]

Epoch 24/30 Losses 0.00090 PDE [ 0.00063 ] boco 0.00019 initial_condition 0.00008 Val [ 0.00060 ]

Epoch 25/30 Losses 0.00113 PDE [ 0.00080 ] boco 0.00021 initial_condition 0.00012 Val [ 0.00059 ]

Epoch 26/30 Losses 0.00077 PDE [ 0.00053 ] boco 0.00017 initial_condition 0.00007 Val [ 0.00067 ]

Epoch 27/30 Losses 0.00107 PDE [ 0.00052 ] boco 0.00043 initial_condition 0.00013 Val [ 0.00047 ]

Epoch 28/30 Losses 0.00057 PDE [ 0.00044 ] boco 0.00008 initial_condition 0.00005 Val [ 0.00052 ]

Epoch 29/30 Losses 0.00050 PDE [ 0.00042 ] boco 0.00005 initial_condition 0.00003 Val [ 0.00052 ]

Epoch 30/30 Losses 0.00066 PDE [ 0.00042 ] boco 0.00018 initial_condition 0.00006 Val [ 0.00038 ]

/opt/conda/lib/python3.7/site-packages/numpy/core/fromnumeric.py:3257: RuntimeWarning: Mean of empty slice.
  out=out, **kwargs)
/opt/conda/lib/python3.7/site-packages/numpy/core/_methods.py:161: RuntimeWarning: invalid value encountered in double_scalars
  ret = ret.dtype.type(ret / rcount)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15,5))
ax1.plot(hist['train_loss'], label="train_loss")
ax1.plot(hist['val_loss'], label="val_loss")
ax1.grid(True)
ax1.set_yscale("log")
ax1.legend()
for boco in pde.bocos:
    ax2.plot(hist['bocos'][boco.name], label=boco.name)
ax2.legend()
ax2.grid(True)
ax2.set_yscale("log")
plt.show()

# evaluate the solution
x = np.linspace(0,1,50)
t = np.linspace(0,1,100)
p, p0, l2 = [], [], []
for _t in t:
    _p0 = np.sin(2.*math.pi*(x-u*_t)) 
    pde.evaluate({'x': x, 't': np.array([_t])})
    _p = pde.outputs['p']
    _l2 = np.mean((_p - _p0)**2)
    p.append(_p)
    p0.append(_p0)
    l2.append(_l2)

from matplotlib import animation, rc
rc('animation', html='html5')

def plot(x, p, p0, t, l2):
    ax.clear()
    tit = ax.set_title(f"t = {t:.2f}, l2 = {l2:.5f}", fontsize=14)
    ax.plot(x, p0, "-k", label="Exact")
    ax.plot(x, p, "g^", label="NN")
    ax.set_xlabel("x", fontsize=14)
    ax.set_ylabel("p", fontsize=14, rotation=np.pi/2)
    ax.legend(loc="upper left")
    ax.grid(True)
    ax.set_xlim([0, 1])
    ax.set_ylim([-1.2, 1.2])
    return [tit]

def get_anim(fig, ax, x, p, p0, t, l2):
    def anim(i):
        return plot(x, p[i], p0[i], t[i], l2[i])
    return anim

fig = plt.figure(figsize=(10,5))
ax = fig.add_subplot(111, autoscale_on=False)
animate = get_anim(fig, ax, x, p, p0, t, l2)
anim = animation.FuncAnimation(fig, animate, frames=len(t), interval=100, blit=True)

anim

Your browser does not support the video tag.

Examples

Check the examples to learn more about using nangs to solve PDEs with NNs.

Copyright

Copyright 2020 onwards, SensioAI. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this project's files except in compliance with the License. A copy of the License is provided in the LICENSE file in this repository.