<jupyter_start><jupyter_text>One-dimensional advection equationWe want to solve the following PDE:\begin{equation}    \frac{\partial \phi}{\partial t} + u \frac{\partial \phi}{\partial x} = 0\end{equation}The independen variables (i.e, $x$ and $t$) are used as input values for the NN, and the solution (i.e. $\phi$) 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. We will try to find a general solution for different values of $u$, so it will be set also as an input.<jupyter_code># autoreload nangs
%reload_ext autoreload
%autoreload 2
%matplotlib inline
# imports
import math
import numpy as np 
import matplotlib.pyplot as plt 

import torch
cuda = True
device = "cuda" if torch.cuda.is_available() and cuda else "cpu"<jupyter_output><empty_output><jupyter_text>Define PDEFirst we define our PDE and set the values for training.<jupyter_code># imports
from nangs.pde import PDE
from nangs.bocos import PeriodicBoco, DirichletBoco
from nangs.solutions import MLP

# define our PDE
class MyPDE(PDE):
    def __init__(self, inputs=None, outputs=None):
        super().__init__(inputs, outputs)
    # the loss of our NN is the PDE !
    def computePDELoss(self, grads, inputs, outputs, params): 
        dpdt, dpdx = grads['p']['t'], grads['p']['x']
        u = inputs['u']
        return [dpdt + u*dpdx]

# instanciate pde with keys for inputs/outputs
pde = MyPDE(inputs=['x', 't', 'u'], outputs=['p'])

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

# periodic b.c for the space dimension
boco = PeriodicBoco('boco', {'x': x[:1], 't': t[:-1], 'u': u}, {'x': x[-1:], 't': t[:-1], 'u': u})
pde.addBoco(boco)

# initial condition (dirichlet for temporal dimension)
p00, p0 = np.sin(2.*math.pi*x), np.array([])
for i in u:
    p0 = np.concatenate((p0,p00)) # one for each value of 'u', keeping the order (x, t, u)
boco = DirichletBoco('initial_condition', {'x': x, 't': t[:1], 'u': u}, {'p': p0})
pde.addBoco(boco)

# define input values for validation
x_v = np.linspace(0,1,25)
t_v = np.linspace(0,1,15)
u_v = np.linspace(0,1,5)
pde.setValues({'x': x_v, 't': t_v[:-1], 'u': u_v}, train=False)<jupyter_output><empty_output><jupyter_text>Now we define a topology for our solution and set the training parameters. Then we can find a solution for our PDE.<jupyter_code># define solution topology
mlp = MLP(pde.n_inputs, pde.n_outputs, 5, 2048)
optimizer = torch.optim.Adam(mlp.parameters(), lr=3e-4)
pde.compile(mlp, optimizer)
# find the solution
hist = pde.solve() 
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()
ax1.set_yscale("log")
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()<jupyter_output><empty_output><jupyter_text>Finally, we can evaluate our solution. Since we trained it for also the parameter *u* we can obtain a solution for any value of *u* without having to perform additional computations.<jupyter_code># evaluate the solution
x = np.linspace(0,1,50)
t = np.linspace(0,1,50)
u = np.linspace(0,1,5)
p, p0, l2 = [], [], []
for _t in t:
    _p, _p0, _l2 = [], [], []
    for _u in u:
        _p0.append(np.sin(2.*math.pi*(x-_u*_t)))
        pde.evaluate({'x': x, 't': np.array([_t]), 'u': np.array([_u])}, device)
        _p.append(pde.outputs['p'])
        _l2.append(np.mean((pde.outputs['p']-np.sin(2.*math.pi*(x-_u*_t)))**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, u):
    ax.clear()
    #tit = ax.set_title(f"t = {t:.2f}, l2 = {l2:.5f}", fontsize=14)
    tit = ax.set_title(f"t = {t:.2f}", fontsize=14)
    for i, _u in enumerate(u):
        ax.plot(x, p0[i], label=f"Exact (u = {_u})")
        ax.plot(x, p[i], ".k", label=f"NN (u = {_u}, l2 = {l2[i]:.5f})")
    ax.set_xlabel("x", fontsize=14)
    ax.set_ylabel("p", fontsize=14, rotation=np.pi/2)
    ax.legend(loc="upper right")
    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, u):
    def anim(i):
        return plot(x, p[i], p0[i], t[i], l2[i], u)
    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, u)
anim = animation.FuncAnimation(fig, animate, frames=len(t), interval=100, blit=True)
anim<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.17212945784374356, 'fraction_numerical': 0.026798598227169654, 'mean_word_length': 4.273913043478261, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '189736', 'license': 'Apache-2.0', 'n_tokens_mistral': 1974, 'n_tokens_neox': 1914, 'n_words': 536, 'path': 'examples/01_adv1d.ipynb', 'repo_name': 'smatkovi/nangs', 'size': '4851'}