246 lines
8.6 KiB
Plaintext
Vendored
246 lines
8.6 KiB
Plaintext
Vendored
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"# Tutorial 3: resolution of wave equation with custom Network"
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"### The problem solution "
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"In this tutorial we present how to solve the wave equation using the `SpatialProblem` and `TimeDependentProblem` class, and the `Network` class for building custom **torch** networks.\n",
|
|
"\n",
|
|
"The problem is written in the following form:\n",
|
|
"\n",
|
|
"\\begin{equation}\n",
|
|
"\\begin{cases}\n",
|
|
"\\Delta u(x,y,t) = \\frac{\\partial^2}{\\partial t^2} u(x,y,t) \\quad \\text{in } D, \\\\\\\\\n",
|
|
"u(x, y, t=0) = \\sin(\\pi x)\\sin(\\pi y), \\\\\\\\\n",
|
|
"u(x, y, t) = 0 \\quad \\text{on } \\Gamma_1 \\cup \\Gamma_2 \\cup \\Gamma_3 \\cup \\Gamma_4,\n",
|
|
"\\end{cases}\n",
|
|
"\\end{equation}\n",
|
|
"\n",
|
|
"where $D$ is a square domain $[0,1]^2$, and $\\Gamma_i$, with $i=1,...,4$, are the boundaries of the square, and the velocity in the standard wave equation is fixed to one."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"First of all, some useful imports."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"import torch\n",
|
|
"\n",
|
|
"from pina.problem import SpatialProblem, TimeDependentProblem\n",
|
|
"from pina.operators import nabla, grad\n",
|
|
"from pina.model import Network\n",
|
|
"from pina import Condition, Span, PINN, Plotter"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"Now, the wave problem is written in PINA code as a class, inheriting from `SpatialProblem` and `TimeDependentProblem` since we deal with spatial, and time dependent variables. The equations are written as `conditions` that should be satisfied in the corresponding domains. `truth_solution` is the exact solution which will be compared with the predicted one."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"class Wave(TimeDependentProblem, SpatialProblem):\n",
|
|
" output_variables = ['u']\n",
|
|
" spatial_domain = Span({'x': [0, 1], 'y': [0, 1]})\n",
|
|
" temporal_domain = Span({'t': [0, 1]})\n",
|
|
"\n",
|
|
" def wave_equation(input_, output_):\n",
|
|
" u_t = grad(output_, input_, components=['u'], d=['t'])\n",
|
|
" u_tt = grad(u_t, input_, components=['dudt'], d=['t'])\n",
|
|
" nabla_u = nabla(output_, input_, components=['u'], d=['x', 'y'])\n",
|
|
" return nabla_u - u_tt\n",
|
|
"\n",
|
|
" def nil_dirichlet(input_, output_):\n",
|
|
" value = 0.0\n",
|
|
" return output_.extract(['u']) - value\n",
|
|
"\n",
|
|
" def initial_condition(input_, output_):\n",
|
|
" u_expected = (torch.sin(torch.pi*input_.extract(['x'])) *\n",
|
|
" torch.sin(torch.pi*input_.extract(['y'])))\n",
|
|
" return output_.extract(['u']) - u_expected\n",
|
|
"\n",
|
|
" conditions = {\n",
|
|
" 'gamma1': Condition(location=Span({'x': [0, 1], 'y': 1, 't': [0, 1]}), function=nil_dirichlet),\n",
|
|
" 'gamma2': Condition(location=Span({'x': [0, 1], 'y': 0, 't': [0, 1]}), function=nil_dirichlet),\n",
|
|
" 'gamma3': Condition(location=Span({'x': 1, 'y': [0, 1], 't': [0, 1]}), function=nil_dirichlet),\n",
|
|
" 'gamma4': Condition(location=Span({'x': 0, 'y': [0, 1], 't': [0, 1]}), function=nil_dirichlet),\n",
|
|
" 't0': Condition(location=Span({'x': [0, 1], 'y': [0, 1], 't': 0}), function=initial_condition),\n",
|
|
" 'D': Condition(location=Span({'x': [0, 1], 'y': [0, 1], 't': [0, 1]}), function=wave_equation),\n",
|
|
" }\n",
|
|
"\n",
|
|
" def wave_sol(self, pts):\n",
|
|
" return (torch.sin(torch.pi*pts.extract(['x'])) *\n",
|
|
" torch.sin(torch.pi*pts.extract(['y'])) *\n",
|
|
" torch.cos(torch.sqrt(torch.tensor(2.))*torch.pi*pts.extract(['t'])))\n",
|
|
"\n",
|
|
" truth_solution = wave_sol\n",
|
|
"\n",
|
|
"problem = Wave()"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"After the problem, a **torch** model is needed to solve the PINN. With the `Network` class the users can convert any **torch** model in a **PINA** model which uses label tensors with a single line of code. We will write a simple residual network using linear layers. Here we implement a simple residual network composed by linear torch layers.\n",
|
|
"\n",
|
|
"This neural network takes as input the coordinates (in this case $x$, $y$ and $t$) and provides the unkwown field of the Wave problem. The residual of the equations are evaluated at several sampling points (which the user can manipulate using the method `span_pts`) and the loss minimized by the neural network is the sum of the residuals."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"class TorchNet(torch.nn.Module):\n",
|
|
" \n",
|
|
" def __init__(self):\n",
|
|
" super().__init__()\n",
|
|
" \n",
|
|
" self.residual = torch.nn.Sequential(torch.nn.Linear(3, 24),\n",
|
|
" torch.nn.Tanh(),\n",
|
|
" torch.nn.Linear(24, 3))\n",
|
|
" \n",
|
|
" self.mlp = torch.nn.Sequential(torch.nn.Linear(3, 64),\n",
|
|
" torch.nn.Tanh(),\n",
|
|
" torch.nn.Linear(64, 1))\n",
|
|
" def forward(self, x):\n",
|
|
" residual_x = self.residual(x)\n",
|
|
" return self.mlp(x + residual_x)\n",
|
|
"\n",
|
|
"# model definition\n",
|
|
"model = Network(model = TorchNet(),\n",
|
|
" input_variables=problem.input_variables,\n",
|
|
" output_variables=problem.output_variables,\n",
|
|
" extra_features=None)"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"In this tutorial, the neural network is trained for 2000 epochs with a learning rate of 0.001. These parameters can be modified as desired.\n",
|
|
"We highlight that the generation of the sampling points and the train is here encapsulated within the function `generate_samples_and_train`, but only for saving some lines of code in the next cells; that function is not mandatory in the **PINA** framework. The training takes approximately one minute."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"def generate_samples_and_train(model, problem):\n",
|
|
" # generate pinn object\n",
|
|
" pinn = PINN(problem, model, lr=0.001)\n",
|
|
"\n",
|
|
" pinn.span_pts(1000, 'random', locations=['D','t0', 'gamma1', 'gamma2', 'gamma3', 'gamma4'])\n",
|
|
" pinn.train(1500, 150)\n",
|
|
" return pinn\n",
|
|
"\n",
|
|
"\n",
|
|
"pinn = generate_samples_and_train(model, problem)"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"After the training is completed one can now plot some results using the `Plotter` class of **PINA**."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"plotter = Plotter()\n",
|
|
"\n",
|
|
"# plotting at fixed time t = 0.6\n",
|
|
"plotter.plot(pinn, fixed_variables={'t': 0.6})\n"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"We can also plot the pinn loss during the training to see the decrease."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"import matplotlib.pyplot as plt\n",
|
|
"\n",
|
|
"plt.figure(figsize=(16, 6))\n",
|
|
"plotter.plot_loss(pinn, label='Loss')\n",
|
|
"\n",
|
|
"plt.grid()\n",
|
|
"plt.legend()\n",
|
|
"plt.show()"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"You can now trying improving the training by changing network, optimizer and its parameters, changin the sampling points,or adding extra features!"
|
|
],
|
|
"metadata": {}
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"name": "python3",
|
|
"display_name": "Python 3.9.16 64-bit ('dl': conda)"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.9.16"
|
|
},
|
|
"interpreter": {
|
|
"hash": "56be7540488f3dc66429ddf54a0fa9de50124d45fcfccfaf04c4c3886d735a3a"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
} |