nangs 0.0.6

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
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.38070 PDE [ 0.01580 ] boco 0.03874 initial_condition 0.32616 Val [ 0.02477 ]

Epoch 2/30 Losses 0.24440 PDE [ 0.02929 ] boco 0.05514 initial_condition 0.15997 Val [ 0.03093 ]

Epoch 3/30 Losses 0.14186 PDE [ 0.03262 ] boco 0.04536 initial_condition 0.06388 Val [ 0.02594 ]

Epoch 4/30 Losses 0.08540 PDE [ 0.02532 ] boco 0.03465 initial_condition 0.02543 Val [ 0.02103 ]

Epoch 5/30 Losses 0.07090 PDE [ 0.01604 ] boco 0.03478 initial_condition 0.02009 Val [ 0.01865 ]

Epoch 6/30 Losses 0.04612 PDE [ 0.01343 ] boco 0.02101 initial_condition 0.01168 Val [ 0.00770 ]

Epoch 7/30 Losses 0.02798 PDE [ 0.00895 ] boco 0.01199 initial_condition 0.00704 Val [ 0.00675 ]

Epoch 8/30 Losses 0.02767 PDE [ 0.00834 ] boco 0.01428 initial_condition 0.00504 Val [ 0.00693 ]

Epoch 9/30 Losses 0.01288 PDE [ 0.00616 ] boco 0.00448 initial_condition 0.00224 Val [ 0.00389 ]

Epoch 10/30 Losses 0.00616 PDE [ 0.00346 ] boco 0.00139 initial_condition 0.00131 Val [ 0.00252 ]

Epoch 11/30 Losses 0.00526 PDE [ 0.00291 ] boco 0.00140 initial_condition 0.00095 Val [ 0.00318 ]

Epoch 12/30 Losses 0.00509 PDE [ 0.00246 ] boco 0.00196 initial_condition 0.00067 Val [ 0.00168 ]

Epoch 13/30 Losses 0.00370 PDE [ 0.00202 ] boco 0.00119 initial_condition 0.00049 Val [ 0.00206 ]

Epoch 14/30 Losses 0.00258 PDE [ 0.00148 ] boco 0.00084 initial_condition 0.00026 Val [ 0.00162 ]

Epoch 15/30 Losses 0.00269 PDE [ 0.00159 ] boco 0.00080 initial_condition 0.00029 Val [ 0.00129 ]

Epoch 16/30 Losses 0.00511 PDE [ 0.00107 ] boco 0.00345 initial_condition 0.00060 Val [ 0.00132 ]

Epoch 17/30 Losses 0.00233 PDE [ 0.00109 ] boco 0.00092 initial_condition 0.00032 Val [ 0.00088 ]

Epoch 18/30 Losses 0.00188 PDE [ 0.00088 ] boco 0.00077 initial_condition 0.00022 Val [ 0.00103 ]

Epoch 19/30 Losses 0.00168 PDE [ 0.00074 ] boco 0.00073 initial_condition 0.00022 Val [ 0.00084 ]

Epoch 20/30 Losses 0.00303 PDE [ 0.00091 ] boco 0.00179 initial_condition 0.00032 Val [ 0.00111 ]

Epoch 21/30 Losses 0.00211 PDE [ 0.00091 ] boco 0.00085 initial_condition 0.00035 Val [ 0.00107 ]

Epoch 22/30 Losses 0.00492 PDE [ 0.00073 ] boco 0.00363 initial_condition 0.00056 Val [ 0.00091 ]

Epoch 23/30 Losses 0.00270 PDE [ 0.00108 ] boco 0.00126 initial_condition 0.00036 Val [ 0.00093 ]

Epoch 24/30 Losses 0.00130 PDE [ 0.00081 ] boco 0.00032 initial_condition 0.00017 Val [ 0.00078 ]

Epoch 25/30 Losses 0.00097 PDE [ 0.00070 ] boco 0.00014 initial_condition 0.00013 Val [ 0.00115 ]

Epoch 26/30 Losses 0.00415 PDE [ 0.00079 ] boco 0.00229 initial_condition 0.00107 Val [ 0.00079 ]

Epoch 27/30 Losses 0.00288 PDE [ 0.00074 ] boco 0.00138 initial_condition 0.00076 Val [ 0.00071 ]

Epoch 28/30 Losses 0.00175 PDE [ 0.00071 ] boco 0.00079 initial_condition 0.00025 Val [ 0.00096 ]

Epoch 29/30 Losses 0.00165 PDE [ 0.00064 ] boco 0.00083 initial_condition 0.00018 Val [ 0.00064 ]

Epoch 30/30 Losses 0.00193 PDE [ 0.00060 ] boco 0.00112 initial_condition 0.00021 Val [ 0.00049 ]

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