nangs 0.0.5

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.

Alternatively, you can use one of our Docker images. You will need Docker 19.03 and, if you have an NVIDIA GPU, the NVIDIA Drivers (you do not need CUDA) and nvidia-docker. We have GPU and CPU images.

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
cuda = False
device = "cuda" if torch.cuda.is_available() and cuda else "cpu"

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

# 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,30)
t = np.linspace(0,1,20)
u = np.array([1.0])
pde.setValues({'x': x, 't': t, 'u': u})

# define input values for testing
x = np.linspace(0,1,25)
t = np.linspace(0,1,15)
pde.setValues({'x': x, 't': t}, 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 = {'layers': 3, 'neurons': 256, 'activations': 'relu'}
pde.buildSolution(mlp)

# set optimization parameters
pde.compile(lr=1e-3, epochs=30, batch_size=32)

# find the solution
hist = pde.solve(device, 'best_solution.pth')

Epoch 1/30 Losses 0.38618 PDE [ 0.01395 ] boco 0.03983 initial_condition 0.33240 Val [ 0.02610 ]

Epoch 2/30 Losses 0.27415 PDE [ 0.03017 ] boco 0.05479 initial_condition 0.18920 Val [ 0.02214 ]

Epoch 3/30 Losses 0.16190 PDE [ 0.03392 ] boco 0.04238 initial_condition 0.08560 Val [ 0.03257 ]

Epoch 4/30 Losses 0.09087 PDE [ 0.03058 ] boco 0.03144 initial_condition 0.02886 Val [ 0.02977 ]

Epoch 5/30 Losses 0.06263 PDE [ 0.02181 ] boco 0.02406 initial_condition 0.01676 Val [ 0.01644 ]

Epoch 6/30 Losses 0.06183 PDE [ 0.01522 ] boco 0.02975 initial_condition 0.01686 Val [ 0.01179 ]

Epoch 7/30 Losses 0.04133 PDE [ 0.01304 ] boco 0.01755 initial_condition 0.01074 Val [ 0.00905 ]

Epoch 8/30 Losses 0.03224 PDE [ 0.01251 ] boco 0.01203 initial_condition 0.00769 Val [ 0.00972 ]

Epoch 9/30 Losses 0.02400 PDE [ 0.00979 ] boco 0.00877 initial_condition 0.00544 Val [ 0.00977 ]

Epoch 10/30 Losses 0.01900 PDE [ 0.00741 ] boco 0.00734 initial_condition 0.00426 Val [ 0.00658 ]

Epoch 11/30 Losses 0.01900 PDE [ 0.00640 ] boco 0.00861 initial_condition 0.00399 Val [ 0.00672 ]

Epoch 12/30 Losses 0.01135 PDE [ 0.00696 ] boco 0.00266 initial_condition 0.00173 Val [ 0.00671 ]

Epoch 13/30 Losses 0.00613 PDE [ 0.00394 ] boco 0.00127 initial_condition 0.00092 Val [ 0.00250 ]

Epoch 14/30 Losses 0.01266 PDE [ 0.00258 ] boco 0.00829 initial_condition 0.00179 Val [ 0.00277 ]

Epoch 15/30 Losses 0.00486 PDE [ 0.00252 ] boco 0.00171 initial_condition 0.00063 Val [ 0.00278 ]

Epoch 16/30 Losses 0.00310 PDE [ 0.00243 ] boco 0.00033 initial_condition 0.00035 Val [ 0.00537 ]

Epoch 17/30 Losses 0.00416 PDE [ 0.00333 ] boco 0.00036 initial_condition 0.00046 Val [ 0.00219 ]

Epoch 18/30 Losses 0.00224 PDE [ 0.00165 ] boco 0.00030 initial_condition 0.00029 Val [ 0.00166 ]

Epoch 19/30 Losses 0.00181 PDE [ 0.00149 ] boco 0.00014 initial_condition 0.00018 Val [ 0.00108 ]

Epoch 20/30 Losses 0.00149 PDE [ 0.00097 ] boco 0.00033 initial_condition 0.00019 Val [ 0.00114 ]

Epoch 21/30 Losses 0.00144 PDE [ 0.00094 ] boco 0.00034 initial_condition 0.00016 Val [ 0.00097 ]

Epoch 22/30 Losses 0.00172 PDE [ 0.00105 ] boco 0.00048 initial_condition 0.00018 Val [ 0.00092 ]

Epoch 23/30 Losses 0.00343 PDE [ 0.00079 ] boco 0.00220 initial_condition 0.00045 Val [ 0.00116 ]

Epoch 24/30 Losses 0.00270 PDE [ 0.00089 ] boco 0.00143 initial_condition 0.00039 Val [ 0.00116 ]

Epoch 25/30 Losses 0.00114 PDE [ 0.00074 ] boco 0.00026 initial_condition 0.00013 Val [ 0.00080 ]

Epoch 26/30 Losses 0.00098 PDE [ 0.00078 ] boco 0.00012 initial_condition 0.00008 Val [ 0.00190 ]

Epoch 27/30 Losses 0.00186 PDE [ 0.00103 ] boco 0.00058 initial_condition 0.00025 Val [ 0.00102 ]

Epoch 28/30 Losses 0.00268 PDE [ 0.00111 ] boco 0.00115 initial_condition 0.00041 Val [ 0.00090 ]

Epoch 29/30 Losses 0.00186 PDE [ 0.00092 ] boco 0.00066 initial_condition 0.00028 Val [ 0.00130 ]

Epoch 30/30 Losses 0.00135 PDE [ 0.00066 ] boco 0.00048 initial_condition 0.00021 Val [ 0.00074 ]

/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
pde.load_state_dict('best_solution.pth')
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])}, device)
    _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.