r"""Solves the electrostatics equation in 2D on a bi-material domain using PINNs. .. math:: -\sigma (\partial_{xx} u_1 - \partial_{yy} u_1) = 0 in \Omega The domain is defined as a bi-material region (a rectangle within a square) with different material properties (:math:`\sigma = 1` in the outer domain and :math:`\sigma = 5` in the inner domain); the sixteen boundary and interface conditions are all enforced weakly. The neural network used is a simple MLP (Multilayer Perceptron); the optimization is performed using either Natural Gradient Descent. """ import timeit import jax import jax.numpy as jnp import matplotlib.pyplot as plt from scimba_jax.domains.meshless_domains.domains_2d import Square2D from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import ( ApproximationSpace, ) from scimba_jax.nonlinear_approximation.integration.monte_carlo import ( DomainSampler, TensorizedSampler, ) from scimba_jax.nonlinear_approximation.model_class.funcparam_field import ( ParamFieldFunction, ) from scimba_jax.nonlinear_approximation.model_class.funcparam_vectorial import ( ParamVecFunction, ) from scimba_jax.nonlinear_approximation.networks.mlp import MLP from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector from scimba_jax.physical_models.abstract_physical_model import AbstractPhysicalModel from scimba_jax.physical_models.abstract_residuals import ( BoundaryResidual, InteriorResidual, ) from scimba_jax.plots.plots_nd import ( plot_abstract_approx_space, ) N_EPOCHS_ENG = 400 N_COLLOC = 2000 N_BC_COLLOC = 6000 ######################################################## # Define the physical problem # ######################################################## class BiMaterialLaplacian(InteriorResidual): def __init__(self, fact, component, domain, f_rhs): super().__init__(domain=domain, size=1, model_type="x", f_rhs=f_rhs) self.fact = fact self.component = component def construct_residual(self, *vars): rho = vars[0] assert isinstance(rho, ParamVecFunction) # split outputs rho_c = rho.component(self.component) # compute laplacian lap_c = rho_c.laplacian("x") # scale by factor return -self.fact * lap_c class NeumannResidual(BoundaryResidual): def __init__( self, index, domain, size=1, model_type="x", f_rhs=None, time_domain=None ): super().__init__( domain=domain, size=size, model_type=model_type, f_rhs=f_rhs, time_domain=time_domain, ) self.index = index def construct_residual(self, *vars): rho = vars[0] assert isinstance(rho, ParamVecFunction) rho_index = rho.component(self.index) # rho_with_n = rho_index.with_normals() grad_x_rho = rho_index.gradient("x") n_model = ParamFieldFunction( rho.dims, lambda *args: args[rho.argnums["n"]], rho.f_type, ) grad_n_rho = grad_x_rho.dot(n_model) return grad_n_rho class DirichletResidual(BoundaryResidual): def __init__( self, index, domain, size=1, model_type="x", f_rhs=None, time_domain=None ): super().__init__( domain=domain, size=size, model_type=model_type, f_rhs=f_rhs, time_domain=time_domain, ) self.index = index def construct_residual(self, *vars): rho = vars[0] assert isinstance(rho, ParamVecFunction) rho_index = rho.component(self.index) return rho_index class InterfaceResidual(BoundaryResidual): def __init__(self, fact, index, domain, f_rhs): super().__init__(domain=domain, size=2, model_type="x", f_rhs=f_rhs) self.fact = fact self.index = index def construct_residual(self, *vars): rho = vars[0] assert isinstance(rho, ParamVecFunction) # split outputs rho_i = rho.component(0) rho_o = rho.component(1) # compute partial derivative with respect to symbol rho_i_index = rho_i.partial_derivative("x", component_index=self.index) rho_o_index = rho_o.partial_derivative("x", component_index=self.index) # scale by factor equality = rho_i - rho_o discontinuity = rho_i_index - self.fact * rho_o_index return ParamVecFunction.cat([equality, discontinuity]) class BiMaterialElectrostatic(AbstractPhysicalModel): """BiMaterialElectrostatic model.""" def __init__(self, main_domain, f_rhs, f_bc_rhs): super().__init__(main_domain=main_domain) self.e_r = 3.3 # inner residuals self.physical_residuals = { "outer": BiMaterialLaplacian( fact=1.0, component=0, domain=main_domain, f_rhs=f_rhs["outer"] ), "inner": BiMaterialLaplacian( fact=5.0, component=1, domain=main_domain, f_rhs=f_rhs["inner"] ), } # boundary residuals self.physical_residuals["bc SN"] = NeumannResidual( index=0, domain=self.boundaries["bc SN"], size=1, f_rhs=f_bc_rhs["bc SN"], ) self.physical_residuals["bc E"] = DirichletResidual( index=0, domain=self.boundaries["bc E"], size=1, f_rhs=f_bc_rhs["bc E"], ) self.physical_residuals["bc W"] = DirichletResidual( index=0, domain=self.boundaries["bc W"], size=1, f_rhs=f_bc_rhs["bc W"], ) self.physical_residuals["bc inner_SN"] = InterfaceResidual( fact=self.e_r, index=1, domain=self.boundaries["bc inner_SN"], f_rhs=f_bc_rhs["bc inner_SN"], ) self.physical_residuals["bc inner_WE"] = InterfaceResidual( fact=self.e_r, index=0, domain=self.boundaries["bc inner_WE"], f_rhs=f_bc_rhs["bc inner_WE"], ) f_rhs = { "inner": lambda *args: 0.0, "outer": lambda *args: 0.0, } f_bc_rhs = { "bc SN": lambda *args: 0.0, "bc E": lambda *args: -1.0, "bc W": lambda *args: 1.0, "bc inner_SN": lambda *args: jnp.zeros((2,)), "bc inner_WE": lambda *args: jnp.zeros((2,)), } ######################################################## # Define the geometric domain # ######################################################## domain_x = [(-1.0, 1.0), (-1.0, 1.0)] domain_ix = [(-0.5, 0.5), (-0.25, 0.25)] big_square = Square2D(domain_x, is_main_domain=True, label_str="outer") small_square = Square2D(domain_ix, is_main_domain=False, label_str="inner") big_square.add_subdomain(small_square) big_square.set_boundaries_dict( { "bc SN": ["bc south", "bc north"], "bc E": ["bc east"], "bc W": ["bc west"], "bc inner_SN": ["bc inner south", "bc inner north"], "bc inner_WE": ["bc inner west", "bc inner east"], } ) sampler = TensorizedSampler([DomainSampler(big_square)], bc=True) ######################################################## # Instantiate an approximation space # ######################################################## key = jax.random.PRNGKey(0) nn = MLP(in_size=2, out_size=2, hidden_sizes=[16, 16, 16], key=key) space = ApproximationSpace({"x": 2}, [(nn, "vec", 2)], model_type="x") ######################################################## # Instantiate the problem and the PINN # ######################################################## model = BiMaterialElectrostatic(big_square, f_rhs, f_bc_rhs) print("\n\n") print("@@@@@@@@@@@@@@@ train with ENG @@@@@@@@@@@@@@@@@@@@@@") pinn = Projector( model, space, sampler, matrix_regularization=5e-4, ) key, sample_dict = sampler.sample(key, N_COLLOC, N_BC_COLLOC) loss0 = pinn.evaluate_loss(space, sample_dict) print("initial loss: ", loss0) ######################################################## # Train the PINN and plot the approximation # ######################################################## # start = timeit.default_timer() key, pinn = pinn.project(key, space, N_EPOCHS_ENG, N_COLLOC, N_BC_COLLOC) end = timeit.default_timer() print("best loss: ", pinn.best_loss) print("time for %d epochs: " % N_EPOCHS_ENG, end - start) plot_abstract_approx_space( pinn.space, big_square, loss=pinn.losses, residual=pinn.model, derivatives=["ux"], components=[{"outer": 0, "inner": 1}], draw_contours=True, n_drawn_contours=20, title="Approximate Bimaterial Electrostatic with a PINN trained with ENG", loss_groups=["bc"], ) plt.show()