r"""Solves the advection of a 2D bump function on a disk with a discrete PINN. ..math:: \partial_t u + a \cdot \nabla u = 0 where :math:`u: \mathbb{R}^2 \times (0, T) \times \mathbb{R}^2 \to \mathbb{R}` is the unknown function, depending on space, time, and two parameters (the position and the variance of the initial bump). The equation is solved using a discrete PINN, with either a classical Adam optimizer, or a natural gradient preconditioning, or Anagram preconditioning. """ from typing import Callable import matplotlib.pyplot as plt import torch from scimba_torch.approximation_space.abstract_space import AbstractApproxSpace from scimba_torch.approximation_space.nn_space import NNxSpace from scimba_torch.domain.meshless_domain.domain_2d import Disk2D 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.mlp import GenericMLP from scimba_torch.numerical_solvers.temporal_pde.discrete_pinn import DiscretePINN from scimba_torch.numerical_solvers.temporal_pde.time_discrete import ( TimeDiscreteAnagramProjector, TimeDiscreteCollocationProjector, TimeDiscreteNaturalGradientProjector, ) from scimba_torch.physical_models.temporal_pde.abstract_temporal_pde import TemporalPDE from scimba_torch.plots.plot_time_discrete_scheme import plot_time_discrete_scheme from scimba_torch.utils.scimba_tensors import LabelTensor, MultiLabelTensor class RotatingTransport(TemporalPDE): def __init__(self, space: AbstractApproxSpace, a: Callable, f_ic: Callable): self.space = space self.a = a self.f_ic = f_ic def grad(self, w: torch.Tensor, z: torch.Tensor) -> torch.Tensor: # returns the gradient of w with respect to z return self.space.grad(w, z) def rhs( # the right hand side of the equation; w is the evaluation of the approximation at t, x, y, mu self, w: MultiLabelTensor, t: LabelTensor, xy: LabelTensor, mu: LabelTensor ) -> torch.Tensor: x, _ = xy.get_components() return 0 * x def bc_rhs( # the left hand side of the boundary conditions self, w: MultiLabelTensor, t: LabelTensor, xy: LabelTensor, n: LabelTensor, mu: LabelTensor, ) -> torch.Tensor: x, _ = xy.get_components() return torch.ones_like(x) def initial_condition( # the left hand side of the initial conditions self, xy: LabelTensor, mu: LabelTensor ) -> torch.Tensor: return self.f_ic(xy, mu) def time_operator( # the time dependant part of the left hand side of the equation self, w: MultiLabelTensor, t: LabelTensor, xy: LabelTensor, mu: LabelTensor ) -> torch.Tensor: return self.grad(w, t) def space_operator( # the space dependant part of the left hand side of the equation self, w: MultiLabelTensor, t: LabelTensor, xy: LabelTensor, mu: LabelTensor ) -> torch.Tensor: grad_u = torch.cat(list(self.grad(w.get_components(), xy)), dim=1) return torch.einsum("bi,bi-> b", self.a(t, xy, mu), grad_u)[..., None] def operator( # already defined in TemporalPDE; given here for completeness self, w: MultiLabelTensor, t: LabelTensor, xy: LabelTensor, mu: LabelTensor ) -> torch.Tensor: return self.space_operator(w, t, xy, mu) + self.time_operator(w, t, xy, mu) # Dirichlet condition: u(t, x, mu) = 0 def bc_operator( # the left hand side of the equation self, w: MultiLabelTensor, t: LabelTensor, x: LabelTensor, n: LabelTensor, mu: LabelTensor, ) -> torch.Tensor: return w.get_components() def exact_solution(t: LabelTensor, xy: LabelTensor, mu: LabelTensor) -> torch.Tensor: x, y = xy.get_components() t_ = t.get_components() c, v = mu.get_components() x_t = torch.cos(2 * torch.pi * t_) y_t = torch.sin(2 * torch.pi * t_) return 1 + torch.exp(-0.5 / v**2 * ((x - c * x_t) ** 2 + (y - c * y_t) ** 2)) def initial_condition(xy: LabelTensor, mu: LabelTensor) -> torch.Tensor: t = LabelTensor(torch.zeros((xy.shape[0], 1))) return exact_solution(t, xy, mu) def a(t: LabelTensor, xy: LabelTensor, mu: LabelTensor) -> torch.Tensor: x, y = xy.get_components() return torch.cat((-2 * torch.pi * y, 2 * torch.pi * x), dim=1) torch.random.manual_seed(0) def solve_with_discrete_pinns( dt: float, T: float, time_integration_scheme: str = "euler_exp", bool_weak_bc: bool = False, precond: str = "ENG", ): domain_x = Disk2D((0.0, 0.0), 1.0, is_main_domain=True) domain_mu = [[0.3, 0.3], [0.1, 0.1]] sampler = TensorizedSampler( [ DomainSampler(domain_x), UniformParametricSampler(domain_mu), ] ) space = NNxSpace(1, 2, GenericMLP, domain_x, sampler, layer_sizes=[8, 16, 8]) pde = RotatingTransport(space, a, initial_condition) if precond == "None": projector_init = TimeDiscreteCollocationProjector(pde) projector = TimeDiscreteCollocationProjector(pde, bool_weak_bc=bool_weak_bc) elif precond == "ENG": projector_init = TimeDiscreteNaturalGradientProjector(pde) projector = TimeDiscreteNaturalGradientProjector(pde, bool_weak_bc=bool_weak_bc) else: projector_init = TimeDiscreteAnagramProjector(pde, svd_threshold=1e-2) projector = TimeDiscreteAnagramProjector( pde, bool_weak_bc=bool_weak_bc, svd_threshold=1e-2 ) scheme = DiscretePINN( pde, projector, projector_init=projector_init, scheme=time_integration_scheme, bool_weak_bc=bool_weak_bc, ) scheme.initialization(epochs=100, n_collocation=15_000, verbose=True) scheme.solve( dt=dt, final_time=T, epochs=100, n_collocation=15_000, n_bc_collocation=15_000, verbose=True, ) plot_time_discrete_scheme( scheme, solution=exact_solution, error=exact_solution, ) plt.show() if __name__ == "__main__": ### No boundary conditions, RK4 scheme, ENG precond # solve_with_discrete_pinns(0.02, 0.04, "rk4", False, "ENG") ### Boundary conditions, Euler Exp scheme, ENG preconditioning # solve_with_discrete_pinns(0.02, 0.04, "euler_exp", True, "ENG") ### Boundary conditions, Euler Exp scheme, Anagram preconditioning # solve_with_discrete_pinns(0.02, 0.04, "euler_exp", True, "Anagram") ## No boundary conditions, RK2 scheme, ENG preconditioning # solve_with_discrete_pinns(0.02, 0.04, "rk2", False, "ENG") ## Boundary conditions, RK2 scheme, ENG preconditioning solve_with_discrete_pinns(0.02, 0.04, "rk2", True, "ENG")