nangs 0.0.7

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(epochs=50)

Epoch 1/50 Losses 0.40488 PDE [ 0.00865 ] boco 0.02703 initial_condition 0.36920 Val [ 0.02271 ]

Epoch 2/50 Losses 0.27958 PDE [ 0.02658 ] boco 0.05469 initial_condition 0.19832 Val [ 0.02176 ]

Epoch 3/50 Losses 0.16315 PDE [ 0.03167 ] boco 0.03852 initial_condition 0.09297 Val [ 0.02226 ]

Epoch 4/50 Losses 0.08097 PDE [ 0.02417 ] boco 0.02951 initial_condition 0.02729 Val [ 0.02281 ]

Epoch 5/50 Losses 0.04978 PDE [ 0.01821 ] boco 0.01974 initial_condition 0.01183 Val [ 0.01231 ]

Epoch 6/50 Losses 0.03304 PDE [ 0.01291 ] boco 0.01263 initial_condition 0.00749 Val [ 0.00973 ]

Epoch 7/50 Losses 0.02165 PDE [ 0.00996 ] boco 0.00743 initial_condition 0.00425 Val [ 0.01086 ]

Epoch 8/50 Losses 0.01205 PDE [ 0.00668 ] boco 0.00337 initial_condition 0.00201 Val [ 0.00553 ]

Epoch 9/50 Losses 0.00777 PDE [ 0.00500 ] boco 0.00162 initial_condition 0.00115 Val [ 0.00421 ]

Epoch 10/50 Losses 0.00519 PDE [ 0.00386 ] boco 0.00079 initial_condition 0.00054 Val [ 0.00359 ]

Epoch 11/50 Losses 0.00389 PDE [ 0.00290 ] boco 0.00062 initial_condition 0.00037 Val [ 0.00259 ]

Epoch 12/50 Losses 0.00303 PDE [ 0.00237 ] boco 0.00042 initial_condition 0.00025 Val [ 0.00190 ]

Epoch 13/50 Losses 0.00257 PDE [ 0.00200 ] boco 0.00035 initial_condition 0.00021 Val [ 0.00201 ]

Epoch 14/50 Losses 0.00209 PDE [ 0.00176 ] boco 0.00020 initial_condition 0.00013 Val [ 0.00157 ]

Epoch 15/50 Losses 0.00195 PDE [ 0.00156 ] boco 0.00024 initial_condition 0.00014 Val [ 0.00164 ]

Epoch 16/50 Losses 0.00165 PDE [ 0.00134 ] boco 0.00019 initial_condition 0.00012 Val [ 0.00129 ]

Epoch 17/50 Losses 0.00186 PDE [ 0.00132 ] boco 0.00039 initial_condition 0.00016 Val [ 0.00129 ]

Epoch 18/50 Losses 0.00143 PDE [ 0.00116 ] boco 0.00018 initial_condition 0.00009 Val [ 0.00221 ]

Epoch 19/50 Losses 0.00201 PDE [ 0.00140 ] boco 0.00042 initial_condition 0.00019 Val [ 0.00088 ]

Epoch 20/50 Losses 0.00135 PDE [ 0.00096 ] boco 0.00026 initial_condition 0.00013 Val [ 0.00088 ]

Epoch 21/50 Losses 0.00133 PDE [ 0.00094 ] boco 0.00027 initial_condition 0.00011 Val [ 0.00071 ]

Epoch 22/50 Losses 0.00095 PDE [ 0.00073 ] boco 0.00015 initial_condition 0.00007 Val [ 0.00064 ]

Epoch 23/50 Losses 0.00079 PDE [ 0.00063 ] boco 0.00011 initial_condition 0.00006 Val [ 0.00081 ]

Epoch 24/50 Losses 0.00090 PDE [ 0.00064 ] boco 0.00018 initial_condition 0.00008 Val [ 0.00081 ]

Epoch 25/50 Losses 0.00126 PDE [ 0.00066 ] boco 0.00046 initial_condition 0.00015 Val [ 0.00056 ]

Epoch 26/50 Losses 0.00109 PDE [ 0.00058 ] boco 0.00038 initial_condition 0.00013 Val [ 0.00055 ]

Epoch 27/50 Losses 0.00065 PDE [ 0.00057 ] boco 0.00005 initial_condition 0.00004 Val [ 0.00062 ]

Epoch 28/50 Losses 0.00080 PDE [ 0.00058 ] boco 0.00015 initial_condition 0.00007 Val [ 0.00058 ]

Epoch 29/50 Losses 0.00098 PDE [ 0.00050 ] boco 0.00036 initial_condition 0.00012 Val [ 0.00043 ]

Epoch 30/50 Losses 0.00083 PDE [ 0.00054 ] boco 0.00021 initial_condition 0.00009 Val [ 0.00048 ]

Epoch 31/50 Losses 0.00060 PDE [ 0.00042 ] boco 0.00012 initial_condition 0.00006 Val [ 0.00046 ]

Epoch 32/50 Losses 0.00099 PDE [ 0.00047 ] boco 0.00039 initial_condition 0.00014 Val [ 0.00061 ]

Epoch 33/50 Losses 0.00086 PDE [ 0.00045 ] boco 0.00031 initial_condition 0.00011 Val [ 0.00037 ]

Epoch 34/50 Losses 0.00044 PDE [ 0.00038 ] boco 0.00003 initial_condition 0.00003 Val [ 0.00034 ]

Epoch 35/50 Losses 0.00045 PDE [ 0.00035 ] boco 0.00006 initial_condition 0.00004 Val [ 0.00045 ]

Epoch 36/50 Losses 0.00092 PDE [ 0.00040 ] boco 0.00040 initial_condition 0.00012 Val [ 0.00047 ]

Epoch 37/50 Losses 0.00066 PDE [ 0.00036 ] boco 0.00022 initial_condition 0.00008 Val [ 0.00029 ]

Epoch 38/50 Losses 0.00069 PDE [ 0.00031 ] boco 0.00028 initial_condition 0.00009 Val [ 0.00033 ]

Epoch 39/50 Losses 0.00074 PDE [ 0.00032 ] boco 0.00033 initial_condition 0.00010 Val [ 0.00048 ]

Epoch 40/50 Losses 0.00054 PDE [ 0.00031 ] boco 0.00017 initial_condition 0.00007 Val [ 0.00032 ]

Epoch 41/50 Losses 0.00055 PDE [ 0.00032 ] boco 0.00017 initial_condition 0.00006 Val [ 0.00041 ]

Epoch 42/50 Losses 0.00059 PDE [ 0.00040 ] boco 0.00013 initial_condition 0.00006 Val [ 0.00041 ]

Epoch 43/50 Losses 0.00084 PDE [ 0.00033 ] boco 0.00040 initial_condition 0.00011 Val [ 0.00039 ]

Epoch 44/50 Losses 0.00061 PDE [ 0.00037 ] boco 0.00017 initial_condition 0.00007 Val [ 0.00042 ]

Epoch 45/50 Losses 0.00050 PDE [ 0.00030 ] boco 0.00015 initial_condition 0.00006 Val [ 0.00037 ]

Epoch 46/50 Losses 0.00040 PDE [ 0.00026 ] boco 0.00010 initial_condition 0.00004 Val [ 0.00040 ]

Epoch 47/50 Losses 0.00077 PDE [ 0.00028 ] boco 0.00036 initial_condition 0.00013 Val [ 0.00029 ]

Epoch 48/50 Losses 0.00045 PDE [ 0.00028 ] boco 0.00012 initial_condition 0.00005 Val [ 0.00028 ]

Epoch 49/50 Losses 0.00067 PDE [ 0.00027 ] boco 0.00031 initial_condition 0.00010 Val [ 0.00030 ]

Epoch 50/50 Losses 0.00062 PDE [ 0.00030 ] boco 0.00025 initial_condition 0.00008 Val [ 0.00030 ]

/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.