r"""Solves a rigid body ODE using a PINN. .. math:: \frac{dx}{dt} = y z + f_x(t), \\ \frac{dy}{dt} = - x z + f_y(t), \\ \frac{dz}{dt} = x y + f_z(t), \\ where :math:`(x, y, z): (0, T) \times (\mu_{\min}, \mu_{\max}) \to \mathbb{R}^3` is the unknown function, with :math:`(0, T) \subset \mathbb{R}` the time domain and :math:`(\mu_{\min}, \mu_{\max}) \subset \mathbb{R}` the parameter domain. The initial condition is :math:`(x(0, \mu), y(0, \mu), z(0, \mu)) = (0, 1, 0)` for all :math:`\mu \in (\mu_{\min}, \mu_{\max})`. The exact solution is .. math:: x(t, \mu) = \sin(t), \\ y(t, \mu) = \cos(t), \\ z(t, \mu) = \sin(t), and the right-hand side functions :math:`f_x`, :math:`f_y`, and :math:`f_z` are chosen such that the exact solution satisfies the ODE. Two training strategies are compared: energy natural gradient with weak IC and energy natural gradient with strong IC. """ # %% import matplotlib.pyplot as plt import torch from scimba_torch.approximation_space.nn_space import NNtSpace from scimba_torch.integration.monte_carlo import TensorizedSampler from scimba_torch.integration.monte_carlo_parameters import UniformParametricSampler from scimba_torch.integration.monte_carlo_time import UniformTimeSampler from scimba_torch.neural_nets.coordinates_based_nets.mlp import ( GenericMLP, ) from scimba_torch.numerical_solvers.ode.pinns import ( NaturalGradientODEPinns, ) from scimba_torch.physical_models.ode.rigid_body import RigidBody from scimba_torch.utils.scimba_tensors import LabelTensor # %% N_EQUATIONS = 3 N_PARAMETERS = 3 VAR_NAMES = ["x", "y", "z"] def exact_x(t, i_x): return i_x * torch.sin(t) def exact_dx_dt(t, i_x): return i_x * torch.cos(t) def exact_y(t, i_y): return i_y * torch.tanh(t / 2) def exact_dy_dt(t, i_y): return i_y * (1 - torch.tanh(t / 2) ** 2) / 2 def exact_z(t, i_z): return i_z * torch.exp(-t) def exact_dz_dt(t, i_z): return -i_z * torch.exp(-t) def exact_sol(t, mu): i_x, i_y, i_z = mu[:, 0], mu[:, 1], mu[:, 2] return torch.stack([exact_x(t, i_x), exact_y(t, i_y), exact_z(t, i_z)], dim=-1) def f_ini(mu): i_x, i_y, i_z = mu.get_components() zero = torch.zeros_like(i_x) return torch.cat( [exact_x(zero, i_x), exact_y(zero, i_y), exact_z(zero, i_z)], dim=-1 ) def initial(mu): i_x, i_y, i_z = mu[0], mu[1], mu[2] zero = torch.zeros_like(i_x) return torch.stack( [exact_x(zero, i_x), exact_y(zero, i_y), exact_z(zero, i_z)], dim=-1 ) def exact(t, mu): t_ = t.get_components() i_x, i_y, i_z = mu.get_components() return torch.cat([exact_x(t_, i_x), exact_y(t_, i_y), exact_z(t_, i_z)], dim=-1) def rhs_func(w, t, mu): t_ = t.get_components() i_x, i_y, i_z = mu.get_components() f = RigidBody.compute_eqs( exact_dx_dt(t_, i_x), exact_dy_dt(t_, i_y), exact_dz_dt(t_, i_z), exact_x(t_, i_x), exact_y(t_, i_y), exact_z(t_, i_z), i_x, i_y, i_z, ) return torch.cat(f, dim=-1) t_min, t_max = 0.0, 2 * torch.pi domain_t = (t_min, t_max) i_x_min, i_x_max = 2.25, 2.75 i_y_min, i_y_max = 1.25, 1.75 i_z_min, i_z_max = 0.25, 0.75 domain_mu = [(i_x_min, i_x_max), (i_y_min, i_y_max), (i_z_min, i_z_max)] sampler = TensorizedSampler( [UniformTimeSampler(domain_t), UniformParametricSampler(domain_mu)] ) def post_processing( inputs: torch.Tensor, t: LabelTensor, mu: LabelTensor ) -> torch.Tensor: t_ = t.get_components() i_x, i_y, i_z = mu.get_components() # set factors according to the slopes of the exact solution at 0 factors = torch.cat((i_x * t_, i_y * t_ / 2, i_z * (1 - t_)), dim=-1) return f_ini(mu) + (factors - f_ini(mu)) * inputs def functional_post_processing( func, t: torch.Tensor, mu: torch.Tensor, theta: torch.Tensor ) -> torch.Tensor: factors = torch.stack((mu[0] * t[0], mu[1] * t[0] / 2, mu[2] * (1 - t[0])), dim=-1) return initial(mu) + (factors - initial(mu)) * func(t, mu, theta) def create_ode(): space = NNtSpace( N_EQUATIONS, N_PARAMETERS, GenericMLP, sampler, layer_sizes=[16] * 3 ) return RigidBody(space, init=f_ini, rhs_func=rhs_func) def create_ode_pp(): space = NNtSpace( N_EQUATIONS, N_PARAMETERS, GenericMLP, sampler, layer_sizes=[16] * 3, post_processing=post_processing, ) return RigidBody(space, init=f_ini, rhs_func=rhs_func) def plot_pinn( pinn, title: str, i_x_val: float = (domain_mu[0][0] + domain_mu[0][1]) / 2, i_y_val: float = (domain_mu[1][0] + domain_mu[1][1]) / 2, i_z_val: float = (domain_mu[2][0] + domain_mu[2][1]) / 2, ): """Plot the results of the PINN.""" with torch.no_grad(): fig, ax = plt.subplots(N_EQUATIONS, 3, figsize=(15, 5 + 2 * N_EQUATIONS)) t = LabelTensor(torch.linspace(t_min, t_max, 1000)[:, None]) mu = LabelTensor( torch.tensor([[i_x_val, i_y_val, i_z_val]]).expand(t.shape[0], -1) ) u = pinn.evaluate(t, mu).w u_exact = exact(t, mu) error = u - u_exact pinn.losses.plot(ax[0, 0]) for i in range(N_EQUATIONS): ax[i, 1].plot(t.x, u[:, i], label=f"approximation, {VAR_NAMES[i]}(t)") ax[i, 1].plot(t.x, u_exact[:, i], label=f"exact, {VAR_NAMES[i]}(t)") ax[i, 1].set_title(f"{title}") ax[i, 1].legend() ax[i, 2].plot(t.x, error[:, i], label=f"error, {VAR_NAMES[i]}(t)") ax[i, 2].set_title(f"Error, L2 = {torch.sqrt(torch.mean(error**2)):.2e}") ax[i, 2].legend() plt.show() # %% pinn_eng_weak = NaturalGradientODEPinns( create_ode(), ic_type="weak", ic_weight=200, one_loss_by_equation=True, matrix_regularization=1e-4, ) resume_solve = True if resume_solve or not pinn_eng_weak.load(__file__, "simple_ode_eng"): pinn_eng_weak.solve(epochs=200, n_collocation=4_000, n_ic_collocation=4_000) pinn_eng_weak.save(__file__, "simple_ode_eng") plot_pinn(pinn_eng_weak, "ENG with weak IC") # %% pinn_eng_strong = NaturalGradientODEPinns( create_ode_pp(), matrix_regularization=1e-3, functional_post_processing=functional_post_processing, one_loss_by_equation=True, ) resume_solve = True if resume_solve or not pinn_eng_strong.load(__file__, "simple_ode_eng_strong"): pinn_eng_strong.solve(epochs=200, n_collocation=4_000) pinn_eng_strong.save(__file__, "simple_ode_eng_strong") plot_pinn(pinn_eng_strong, "ENG with strong IC") # %%