nangs 0.0.4

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.

Getting Started

We have GPU and CPU images

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, 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': 100, 'activations': 'relu'}
pde.buildSolution(mlp)

# set optimization parameters
pde.compile(lr=0.01, epochs=20, batch_size=32)

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

Epoch 1/20 Losses 0.40368 PDE 0.01568 boco 0.04150 initial_condition 0.34651 Val 0.03113

Epoch 2/20 Losses 0.26363 PDE 0.03761 boco 0.05959 initial_condition 0.16643 Val 0.02898

Epoch 3/20 Losses 0.14324 PDE 0.03913 boco 0.04351 initial_condition 0.06060 Val 0.04284

Epoch 4/20 Losses 0.08736 PDE 0.02872 boco 0.03349 initial_condition 0.02515 Val 0.02061

Epoch 5/20 Losses 0.05867 PDE 0.01603 boco 0.02591 initial_condition 0.01674 Val 0.01160

Epoch 6/20 Losses 0.04793 PDE 0.00891 boco 0.02590 initial_condition 0.01311 Val 0.00647

Epoch 7/20 Losses 0.03881 PDE 0.00493 boco 0.01901 initial_condition 0.01487 Val 0.00607

Epoch 8/20 Losses 0.01393 PDE 0.00450 boco 0.00616 initial_condition 0.00327 Val 0.00378

Epoch 9/20 Losses 0.00570 PDE 0.00183 boco 0.00265 initial_condition 0.00121 Val 0.00114

Epoch 10/20 Losses 0.01189 PDE 0.00136 boco 0.00679 initial_condition 0.00375 Val 0.00126

Epoch 11/20 Losses 0.00379 PDE 0.00088 boco 0.00197 initial_condition 0.00094 Val 0.00067

Epoch 12/20 Losses 0.00154 PDE 0.00067 boco 0.00057 initial_condition 0.00030 Val 0.00047

Epoch 13/20 Losses 0.00076 PDE 0.00048 boco 0.00013 initial_condition 0.00015 Val 0.00039

Epoch 14/20 Losses 0.00071 PDE 0.00033 boco 0.00022 initial_condition 0.00016 Val 0.00028

Epoch 15/20 Losses 0.00072 PDE 0.00030 boco 0.00024 initial_condition 0.00018 Val 0.00029

Epoch 16/20 Losses 0.01209 PDE 0.00050 boco 0.00769 initial_condition 0.00391 Val 0.00049

Epoch 17/20 Losses 0.00379 PDE 0.00043 boco 0.00219 initial_condition 0.00118 Val 0.00026

Epoch 18/20 Losses 0.00163 PDE 0.00029 boco 0.00090 initial_condition 0.00044 Val 0.00022

Epoch 19/20 Losses 0.00054 PDE 0.00020 boco 0.00018 initial_condition 0.00016 Val 0.00027

Epoch 20/20 Losses 0.00056 PDE 0.00018 boco 0.00027 initial_condition 0.00010 Val 0.00013

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.legend()
for boco in pde.bocos:
    ax2.plot(hist['bocos'][boco.name], label=boco.name)
ax2.legend()
ax2.grid(True)
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))
    _p = pde.evaluate({'x': x, 't': np.array([_t])}, device)
    _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, 1])
    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.