Time discrete: heat equationΒΆ
Solves a 2D heat equation using discrete pinns implicit.
\[\partial_t u - \partial_{xx} u - \partial_{yy} u = 0\]
The equation is solved using a discrete pinns with implicit scheme and an Natural Gradient optimizer.
[1]:
import math
import matplotlib.pyplot as plt
import torch
from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_2d import Square2D
from scimba_torch.integration.monte_carlo import DomainSampler, TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import UniformParametricSampler
from scimba_torch.neural_nets.coordinates_based_nets.features import PeriodicMLP
from scimba_torch.numerical_solvers.temporal_pde.discrete_pinn import (
DiscretePINNImplicit,
)
from scimba_torch.numerical_solvers.temporal_pde.time_discrete import (
TimeDiscreteImplicitNaturalGradientProjector,
TimeDiscreteNaturalGradientProjector,
)
from scimba_torch.physical_models.temporal_pde.heat_equation import (
HeatEquation2DStrongForm,
HeatEquation2DStrongFormImplicit,
)
from scimba_torch.plots.plot_time_discrete_scheme import plot_time_discrete_scheme
from scimba_torch.utils.scimba_tensors import LabelTensor
def f_ini(x: LabelTensor, mu):
x1, x2 = x.get_components()
f = torch.sin(torch.pi * x1) * torch.sin(torch.pi * x2)
return f
def f_rhs(w, t, x: LabelTensor, mu: LabelTensor):
x1, x2 = x.get_components()
return 0 * x1
def f_bc(x: LabelTensor, mu: LabelTensor):
x1, _ = x.get_components()
return x1 * 0.0
def f_exact(t: LabelTensor, x: LabelTensor, mu: LabelTensor):
x1, x2 = x.get_components()
D = torch.tensor(0.02)
f = (
torch.exp(-2 * D * torch.pi**2 * t.x)
* torch.sin(torch.pi * x1)
* torch.sin(torch.pi * x2)
)
return f
domain_x = Square2D([(-1, 1), (-1, 1)], is_main_domain=True)
domain_mu = [(0.02, 0.02 + 1e-5)]
sampler = TensorizedSampler(
[DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
)
space = NNxSpace(
1,
1,
PeriodicMLP,
domain_x,
sampler,
layer_sizes=[20] * 2,
parameters_bounds=domain_mu,
)
gamma_sdirk2 = 1.0 - 1.0 / math.sqrt(2.0)
dt = 1e-2
alpha = gamma_sdirk2
type_scheme = "implicit"
pde = HeatEquation2DStrongForm(space, init=f_ini, f=f_rhs, g=f_bc)
pde_imp = HeatEquation2DStrongFormImplicit(
space, init=f_ini, f=f_rhs, g=f_bc, dt=dt, alpha=alpha
)
if type_scheme == "explicit":
projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)
else:
projector = TimeDiscreteImplicitNaturalGradientProjector(
pde_imp, rhs=f_ini, dt=dt, alpha=alpha
)
scheme = DiscretePINNImplicit(
pde, pde_imp, projector, type_scheme=type_scheme, scheme="sdirk2"
)
retrain = True
if retrain:
scheme.initialization(epochs=100, verbose=True, n_collocation=3000)
scheme.projector.save("ini_heat2D_imp")
else:
scheme.projector.load("ini_heat2D_imp")
T = 1e-1
scheme.projector.space.load_from_best_approx()
scheme.solve(dt=dt, final_time=T, n_collocation=3000, epochs=10, verbose=True)
plot_time_discrete_scheme(
scheme,
solution=f_exact,
error=f_exact,
draw_contours=True,
)
plt.show()
activating optimizer ScimbaSGD
epoch: 0, best loss: 2.436e-01
epoch: 0, loss: 2.436e-01
epoch: 1, best loss: 2.365e-01
epoch: 2, best loss: 2.282e-01
epoch: 3, best loss: 1.173e-01
epoch: 4, best loss: 9.008e-02
epoch: 5, best loss: 3.224e-02
epoch: 6, best loss: 1.850e-03
epoch: 7, best loss: 6.583e-05
epoch: 8, best loss: 1.678e-05
epoch: 9, best loss: 1.731e-06
epoch: 10, best loss: 3.266e-07
epoch: 11, best loss: 1.794e-07
epoch: 12, best loss: 3.613e-08
epoch: 13, best loss: 2.083e-08
epoch: 14, best loss: 1.366e-08
epoch: 15, best loss: 1.068e-08
epoch: 16, best loss: 9.335e-09
epoch: 17, best loss: 7.511e-09
epoch: 18, best loss: 5.732e-09
epoch: 19, best loss: 4.789e-09
epoch: 20, best loss: 3.724e-09
epoch: 21, best loss: 3.399e-09
epoch: 22, best loss: 3.156e-09
epoch: 23, best loss: 2.832e-09
epoch: 24, best loss: 2.359e-09
epoch: 25, best loss: 2.179e-09
epoch: 26, best loss: 1.943e-09
epoch: 28, best loss: 1.774e-09
epoch: 29, best loss: 1.540e-09
epoch: 30, best loss: 1.442e-09
epoch: 32, best loss: 1.289e-09
epoch: 33, best loss: 1.226e-09
epoch: 35, best loss: 1.200e-09
epoch: 36, best loss: 1.080e-09
epoch: 38, best loss: 1.053e-09
epoch: 39, best loss: 9.348e-10
epoch: 40, best loss: 9.248e-10
epoch: 41, best loss: 8.838e-10
epoch: 42, best loss: 8.530e-10
epoch: 43, best loss: 8.366e-10
epoch: 44, best loss: 8.057e-10
epoch: 45, best loss: 7.858e-10
epoch: 47, best loss: 7.227e-10
epoch: 48, best loss: 7.052e-10
epoch: 49, best loss: 6.694e-10
epoch: 51, best loss: 6.417e-10
epoch: 52, best loss: 6.341e-10
epoch: 53, best loss: 6.306e-10
epoch: 54, best loss: 5.692e-10
epoch: 55, best loss: 5.311e-10
epoch: 58, best loss: 4.999e-10
epoch: 59, best loss: 4.919e-10
epoch: 60, best loss: 4.791e-10
epoch: 61, best loss: 4.501e-10
epoch: 64, best loss: 4.017e-10
epoch: 66, best loss: 3.747e-10
epoch: 68, best loss: 3.723e-10
epoch: 70, best loss: 3.665e-10
epoch: 72, best loss: 3.381e-10
epoch: 74, best loss: 3.347e-10
epoch: 75, best loss: 3.220e-10
epoch: 76, best loss: 3.067e-10
epoch: 77, best loss: 3.057e-10
epoch: 80, best loss: 2.937e-10
epoch: 82, best loss: 2.854e-10
epoch: 83, best loss: 2.835e-10
epoch: 84, best loss: 2.643e-10
epoch: 88, best loss: 2.609e-10
epoch: 89, best loss: 2.586e-10
epoch: 90, best loss: 2.565e-10
epoch: 91, best loss: 2.351e-10
epoch: 95, best loss: 2.202e-10
epoch: 98, best loss: 2.198e-10
epoch: 99, best loss: 2.165e-10
Training done!
Final loss value: 2.137e-10
Best loss value: 2.165e-10
epoch: 0, best loss: 3.271e-07
epoch: 0, loss: 3.271e-07
epoch: 1, best loss: 3.451e-12
epoch: 2, best loss: 1.575e-13
epoch: 3, best loss: 9.744e-14
epoch: 4, best loss: 8.967e-14
epoch: 5, best loss: 8.638e-14
epoch: 6, best loss: 7.692e-14
epoch: 7, best loss: 7.445e-14
Training done!
Final loss value: 7.359e-14
Best loss value: 7.445e-14
epoch: 0, best loss: 1.952e-06
epoch: 0, loss: 1.952e-06
epoch: 1, best loss: 2.053e-11
epoch: 2, best loss: 2.911e-12
epoch: 3, best loss: 1.042e-12
epoch: 4, best loss: 9.205e-13
epoch: 5, best loss: 8.939e-13
epoch: 7, best loss: 8.746e-13
epoch: 9, best loss: 8.549e-13
Training done!
Final loss value: 8.928e-13
Best loss value: 8.549e-13
epoch: 0, best loss: 3.317e-07
epoch: 0, loss: 3.317e-07
epoch: 1, best loss: 2.982e-12
epoch: 2, best loss: 1.266e-13
epoch: 3, best loss: 9.406e-14
epoch: 4, best loss: 8.786e-14
epoch: 5, best loss: 7.905e-14
epoch: 6, best loss: 7.667e-14
epoch: 7, best loss: 7.539e-14
epoch: 9, best loss: 7.158e-14
Training done!
Final loss value: 7.071e-14
Best loss value: 7.158e-14
epoch: 0, best loss: 2.036e-06
epoch: 0, loss: 2.036e-06
epoch: 1, best loss: 2.272e-11
epoch: 2, best loss: 1.409e-12
epoch: 3, best loss: 9.227e-13
epoch: 4, best loss: 8.922e-13
epoch: 5, best loss: 8.305e-13
epoch: 6, best loss: 7.981e-13
epoch: 9, best loss: 7.924e-13
Training done!
Final loss value: 8.007e-13
Best loss value: 7.924e-13
epoch: 0, best loss: 3.389e-07
epoch: 0, loss: 3.389e-07
epoch: 1, best loss: 3.093e-12
epoch: 2, best loss: 1.225e-13
epoch: 3, best loss: 9.059e-14
epoch: 4, best loss: 7.952e-14
epoch: 5, best loss: 7.432e-14
epoch: 7, best loss: 6.906e-14
Training done!
Final loss value: 7.300e-14
Best loss value: 6.906e-14
epoch: 0, best loss: 1.908e-06
epoch: 0, loss: 1.908e-06
epoch: 1, best loss: 2.064e-11
epoch: 2, best loss: 1.895e-12
epoch: 3, best loss: 9.299e-13
epoch: 4, best loss: 8.735e-13
epoch: 5, best loss: 8.295e-13
epoch: 6, best loss: 8.195e-13
epoch: 7, best loss: 7.771e-13
epoch: 8, best loss: 7.612e-13
Training done!
Final loss value: 7.596e-13
Best loss value: 7.612e-13
epoch: 0, best loss: 3.127e-07
epoch: 0, loss: 3.127e-07
epoch: 1, best loss: 3.125e-12
epoch: 2, best loss: 1.696e-13
epoch: 3, best loss: 9.237e-14
epoch: 4, best loss: 8.351e-14
epoch: 5, best loss: 7.727e-14
epoch: 6, best loss: 7.356e-14
epoch: 7, best loss: 7.166e-14
epoch: 8, best loss: 6.964e-14
epoch: 9, best loss: 6.524e-14
Training done!
Final loss value: 6.865e-14
Best loss value: 6.524e-14
epoch: 0, best loss: 1.891e-06
epoch: 0, loss: 1.891e-06
epoch: 1, best loss: 4.466e-11
epoch: 2, best loss: 1.750e-12
epoch: 3, best loss: 9.102e-13
epoch: 4, best loss: 8.600e-13
epoch: 5, best loss: 8.111e-13
epoch: 6, best loss: 7.784e-13
epoch: 8, best loss: 7.685e-13
epoch: 9, best loss: 7.647e-13
Training done!
Final loss value: 7.144e-13
Best loss value: 7.647e-13
epoch: 0, best loss: 3.261e-07
epoch: 0, loss: 3.261e-07
epoch: 1, best loss: 3.249e-12
epoch: 2, best loss: 1.146e-13
epoch: 3, best loss: 8.147e-14
epoch: 4, best loss: 7.647e-14
epoch: 5, best loss: 6.815e-14
epoch: 6, best loss: 6.752e-14
epoch: 7, best loss: 6.708e-14
epoch: 8, best loss: 6.352e-14
Training done!
Final loss value: 6.332e-14
Best loss value: 6.352e-14
epoch: 0, best loss: 1.868e-06
epoch: 0, loss: 1.868e-06
epoch: 1, best loss: 1.954e-11
epoch: 2, best loss: 5.073e-12
epoch: 3, best loss: 9.518e-13
epoch: 4, best loss: 7.804e-13
epoch: 6, best loss: 7.690e-13
epoch: 8, best loss: 7.503e-13
Training done!
Final loss value: 7.580e-13
Best loss value: 7.503e-13
epoch: 0, best loss: 3.234e-07
epoch: 0, loss: 3.234e-07
epoch: 1, best loss: 2.917e-12
epoch: 2, best loss: 1.248e-13
epoch: 3, best loss: 9.132e-14
epoch: 4, best loss: 8.092e-14
epoch: 5, best loss: 7.262e-14
epoch: 6, best loss: 6.963e-14
epoch: 7, best loss: 6.481e-14
epoch: 8, best loss: 6.357e-14
Training done!
Final loss value: 6.409e-14
Best loss value: 6.357e-14
epoch: 0, best loss: 1.828e-06
epoch: 0, loss: 1.828e-06
epoch: 1, best loss: 2.613e-11
epoch: 2, best loss: 2.757e-12
epoch: 3, best loss: 8.578e-13
epoch: 4, best loss: 7.842e-13
epoch: 5, best loss: 7.781e-13
epoch: 7, best loss: 7.511e-13
epoch: 8, best loss: 7.333e-13
epoch: 9, best loss: 7.262e-13
Training done!
Final loss value: 7.263e-13
Best loss value: 7.262e-13
epoch: 0, best loss: 3.336e-07
epoch: 0, loss: 3.336e-07
epoch: 1, best loss: 3.032e-12
epoch: 2, best loss: 1.296e-13
epoch: 3, best loss: 7.887e-14
epoch: 4, best loss: 6.793e-14
epoch: 5, best loss: 6.417e-14
epoch: 6, best loss: 6.138e-14
epoch: 8, best loss: 6.038e-14
Training done!
Final loss value: 6.127e-14
Best loss value: 6.038e-14
epoch: 0, best loss: 1.845e-06
epoch: 0, loss: 1.845e-06
epoch: 1, best loss: 1.933e-11
epoch: 2, best loss: 9.847e-13
epoch: 3, best loss: 7.952e-13
epoch: 4, best loss: 7.362e-13
epoch: 6, best loss: 7.312e-13
epoch: 7, best loss: 7.182e-13
epoch: 8, best loss: 6.828e-13
epoch: 9, best loss: 6.824e-13
Training done!
Final loss value: 7.024e-13
Best loss value: 6.824e-13
epoch: 0, best loss: 3.169e-07
epoch: 0, loss: 3.169e-07
epoch: 1, best loss: 3.670e-12
epoch: 2, best loss: 1.403e-13
epoch: 3, best loss: 8.145e-14
epoch: 4, best loss: 7.028e-14
epoch: 5, best loss: 7.005e-14
epoch: 6, best loss: 6.711e-14
epoch: 7, best loss: 6.205e-14
epoch: 8, best loss: 5.953e-14
Training done!
Final loss value: 6.127e-14
Best loss value: 5.953e-14
epoch: 0, best loss: 1.907e-06
epoch: 0, loss: 1.907e-06
epoch: 1, best loss: 2.538e-11
epoch: 2, best loss: 3.484e-12
epoch: 3, best loss: 8.376e-13
epoch: 4, best loss: 7.261e-13
epoch: 5, best loss: 7.110e-13
epoch: 6, best loss: 6.865e-13
epoch: 9, best loss: 6.486e-13
Training done!
Final loss value: 6.896e-13
Best loss value: 6.486e-13
epoch: 0, best loss: 3.182e-07
epoch: 0, loss: 3.182e-07
epoch: 1, best loss: 2.747e-12
epoch: 2, best loss: 8.412e-14
epoch: 3, best loss: 6.455e-14
epoch: 4, best loss: 6.322e-14
epoch: 5, best loss: 6.104e-14
epoch: 6, best loss: 5.976e-14
epoch: 7, best loss: 5.657e-14
epoch: 9, best loss: 5.599e-14
Training done!
Final loss value: 5.299e-14
Best loss value: 5.599e-14
epoch: 0, best loss: 1.897e-06
epoch: 0, loss: 1.897e-06
epoch: 1, best loss: 2.077e-11
epoch: 2, best loss: 9.470e-13
epoch: 3, best loss: 7.576e-13
epoch: 4, best loss: 7.211e-13
epoch: 5, best loss: 6.797e-13
epoch: 8, best loss: 6.792e-13
epoch: 9, best loss: 6.529e-13
Training done!
Final loss value: 6.556e-13
Best loss value: 6.529e-13
epoch: 0, best loss: 3.168e-07
epoch: 0, loss: 3.168e-07
epoch: 1, best loss: 2.744e-12
epoch: 2, best loss: 8.666e-14
epoch: 3, best loss: 7.094e-14
epoch: 4, best loss: 6.572e-14
epoch: 5, best loss: 6.282e-14
epoch: 7, best loss: 5.806e-14
epoch: 9, best loss: 5.651e-14
Training done!
Final loss value: 5.650e-14
Best loss value: 5.651e-14
epoch: 0, best loss: 1.752e-06
epoch: 0, loss: 1.752e-06
epoch: 1, best loss: 1.685e-11
epoch: 2, best loss: 1.261e-12
epoch: 3, best loss: 7.239e-13
epoch: 4, best loss: 6.852e-13
epoch: 5, best loss: 6.744e-13
epoch: 6, best loss: 6.722e-13
epoch: 7, best loss: 6.545e-13
epoch: 8, best loss: 6.325e-13
epoch: 9, best loss: 6.259e-13
Training done!
Final loss value: 6.466e-13
Best loss value: 6.259e-13
[ ]: