jinns 0.6.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

jinns/loss/_DynamicLoss.py

@@ -48,6 +48,7 @@ class FisherKPP(PDENonStatio):
         """
         super().__init__(Tmax, eq_params_heterogeneity)
 
+    @PDENonStatio.evaluate_heterogeneous_parameters
     def evaluate(self, t, x, u, params):
         r"""
         Evaluate the dynamic loss at :math:`(t,x)`.
@@ -67,16 +68,12 @@ class FisherKPP(PDENonStatio):
             differential equation parameters and the neural network parameter
         """
         if isinstance(u, PINN):
-            params["eq_params"] = self._eval_heterogeneous_parameters(
-                params["eq_params"], t, x, self.eq_params_heterogeneity
-            )
-
             # Note that the last dim of u is nec. 1
             u_ = lambda t, x: u(t, x, params)[0]
 
             du_dt = grad(u_, 0)(t, x)
 
-            lap = _laplacian_rev(u, params, x, t)[..., None]
+            lap = _laplacian_rev(t, x, u, params)[..., None]
 
             return du_dt + self.Tmax * (
                 -params["eq_params"]["D"] * lap
@@ -87,17 +84,12 @@ class FisherKPP(PDENonStatio):
                 )
             )
         if isinstance(u, SPINN):
-            x_grid = _get_grid(x)
-            params["eq_params"] = self._eval_heterogeneous_parameters(
-                params["eq_params"], t, x_grid, self.eq_params_heterogeneity
-            )
-
             u_tx, du_dt = jax.jvp(
                 lambda t: u(t, x, params),
                 (t,),
                 (jnp.ones_like(t),),
             )
-            lap = _laplacian_fwd(u, params, x, t)[..., None]
+            lap = _laplacian_fwd(t, x, u, params)[..., None]
             return du_dt + self.Tmax * (
                 -params["eq_params"]["D"] * lap
                 - u_tx
@@ -160,12 +152,8 @@ class BurgerEquation(PDENonStatio):
             differential equation parameters and the neural network parameter
         """
         if isinstance(u, PINN):
-            params["eq_params"] = self._eval_heterogeneous_parameters(
-                params["eq_params"], t, x, self.eq_params_heterogeneity
-            )
-
             # Note that the last dim of u is nec. 1
-            u_ = lambda t, x: u(t, x, params)[0]
+            u_ = lambda t, x: jnp.squeeze(u(t, x, params)[u.slice_solution])
             du_dt = grad(u_, 0)
             du_dx = grad(u_, 1)
             d2u_dx2 = grad(
@@ -179,10 +167,6 @@ class BurgerEquation(PDENonStatio):
             )
 
         if isinstance(u, SPINN):
-            x_grid = _get_grid(x)
-            params["eq_params"] = self._eval_heterogeneous_parameters(
-                params["eq_params"], t, x_grid, self.eq_params_heterogeneity
-            )
             # d=2 JVP calls are expected since we have time and x
             # then with a batch of size B, we then have Bd JVP calls
             u_tx, du_dt = jax.jvp(
@@ -352,10 +336,6 @@ class FPENonStatioLoss2D(PDENonStatio):
             differential equation parameters and the neural network parameter
         """
         if isinstance(u, PINN):
-            params["eq_params"] = self._eval_heterogeneous_parameters(
-                params["eq_params"], t, x, self.eq_params_heterogeneity
-            )
-
             # Note that the last dim of u is nec. 1
             u_ = lambda t, x: u(t, x, params)[0]
 
@@ -411,10 +391,6 @@ class FPENonStatioLoss2D(PDENonStatio):
 
         if isinstance(u, SPINN):
             x_grid = _get_grid(x)
-            params["eq_params"] = self._eval_heterogeneous_parameters(
-                params["eq_params"], t, x_grid, self.eq_params_heterogeneity
-            )
-
             _, du_dt = jax.jvp(
                 lambda t: u(t, x, params),
                 (t,),
@@ -634,12 +610,12 @@ class MassConservation2DStatio(PDEStatio):
         if isinstance(u_dict[self.nn_key], PINN):
             u = u_dict[self.nn_key]
 
-            return _div_rev(u, params, x)[..., None]
+            return _div_rev(None, x, u, params)[..., None]
 
         if isinstance(u_dict[self.nn_key], SPINN):
             u = u_dict[self.nn_key]
 
-            return _div_fwd(u, params, x)[..., None]
+            return _div_fwd(None, x, u, params)[..., None]
         raise ValueError("u is not among the recognized types (PINN or SPINN)")
 
 
@@ -724,12 +700,12 @@ class NavierStokes2DStatio(PDEStatio):
         if isinstance(u_dict[self.u_key], PINN):
             u = u_dict[self.u_key]
 
-            u_dot_nabla_x_u = _u_dot_nabla_times_u_rev(u, u_params, x)
+            u_dot_nabla_x_u = _u_dot_nabla_times_u_rev(None, x, u, u_params)
 
             p = lambda x: u_dict[self.p_key](x, p_params)
             jac_p = jacrev(p, 0)(x)  # compute the gradient
 
-            vec_laplacian_u = _vectorial_laplacian(u, u_params, x, u_vec_ndim=2)
+            vec_laplacian_u = _vectorial_laplacian(None, x, u, u_params, u_vec_ndim=2)
 
             # dynamic loss on x axis
             result_x = (
@@ -751,7 +727,7 @@ class NavierStokes2DStatio(PDEStatio):
         if isinstance(u_dict[self.u_key], SPINN):
             u = u_dict[self.u_key]
 
-            u_dot_nabla_x_u = _u_dot_nabla_times_u_fwd(u, u_params, x)
+            u_dot_nabla_x_u = _u_dot_nabla_times_u_fwd(None, x, u, u_params)
 
             p = lambda x: u_dict[self.p_key](x, p_params)
 
@@ -761,7 +737,7 @@ class NavierStokes2DStatio(PDEStatio):
             _, dp_dy = jax.jvp(p, (x,), (tangent_vec_1,))
 
             vec_laplacian_u = jnp.moveaxis(
-                _vectorial_laplacian(u, u_params, x, u_vec_ndim=2),
+                _vectorial_laplacian(None, x, u, u_params, u_vec_ndim=2),
                 source=0,
                 destination=-1,
             )

jinns/loss/_DynamicLossAbstract.py

@@ -23,8 +23,8 @@ class DynamicLoss:
             Default None. A dict with the keys being the same as in eq_params
             and the value being either None (no heterogeneity) or a function
             which encodes for the spatio-temporal heterogeneity of the parameter.
-            Such a function must be jittable and take three arguments `t`,
-            `x` and `params["eq_params"]` even if one is not used. Therefore,
+            Such a function must be jittable and take four arguments `t`, `x`,
+            `u` and `params` even if one is not used. Therefore,
             one can introduce spatio-temporal covariates upon which a particular
             parameter can depend, e.g. in a GLM fashion. The effect of these
             covariables can themselves be estimated by being in `eq_params` too.
@@ -35,20 +35,24 @@ class DynamicLoss:
         self.Tmax = Tmax
         self.eq_params_heterogeneity = eq_params_heterogeneity
 
-    def _eval_heterogeneous_parameters(
-        self, eq_params, t, x, eq_params_heterogeneity=None
-    ):
+    @staticmethod
+    def _eval_heterogeneous_parameters(t, x, u, params, eq_params_heterogeneity=None):
         eq_params_ = {}
         if eq_params_heterogeneity is None:
-            return eq_params
-        for k, p in eq_params.items():
+            return params["eq_params"]
+        for k, p in params["eq_params"].items():
             try:
                 if eq_params_heterogeneity[k] is None:
                     eq_params_[k] = p
                 else:
-                    eq_params_[k] = eq_params_heterogeneity[k](
-                        t, x, eq_params  # heterogeneity encoded through a function
-                    )
+                    if t is None:
+                        eq_params_[k] = eq_params_heterogeneity[k](
+                            x, u, params  # heterogeneity encoded through a function
+                        )
+                    else:
+                        eq_params_[k] = eq_params_heterogeneity[k](
+                            t, x, u, params  # heterogeneity encoded through a function
+                        )
             except KeyError:
                 # we authorize missing eq_params_heterogeneity key
                 # is its heterogeneity is None anyway
@@ -79,6 +83,31 @@ class ODE(DynamicLoss):
         """
         super().__init__(Tmax, eq_params_heterogeneity)
 
+    def eval_heterogeneous_parameters(self, t, u, params, eq_params_heterogeneity=None):
+        return super()._eval_heterogeneous_parameters(
+            t, None, u, params, eq_params_heterogeneity
+        )
+
+    @staticmethod
+    def evaluate_heterogeneous_parameters(evaluate):
+        """
+        Decorator which aims to decorate the evaluate methods of Dynamic losses
+        in order. It calls _eval_heterogeneous_parameters which applies the
+        user defined rules to obtain spatially / temporally heterogeneous
+        parameters
+        """
+
+        def wrapper(*args):
+            self, t, u, params = args
+            params["eq_params"] = self.eval_heterogeneous_parameters(
+                t, u, params, self.eq_params_heterogeneity
+            )
+            new_args = args[:-1] + (params,)
+            res = evaluate(*new_args)
+            return res
+
+        return wrapper
+
 
 class PDEStatio(DynamicLoss):
     r"""
@@ -99,6 +128,31 @@ class PDEStatio(DynamicLoss):
         """
         super().__init__(eq_params_heterogeneity=eq_params_heterogeneity)
 
+    def eval_heterogeneous_parameters(self, x, u, params, eq_params_heterogeneity=None):
+        return super()._eval_heterogeneous_parameters(
+            None, x, u, params, eq_params_heterogeneity
+        )
+
+    @staticmethod
+    def evaluate_heterogeneous_parameters(evaluate):
+        """
+        Decorator which aims to decorate the evaluate methods of Dynamic losses
+        in order. It calls _eval_heterogeneous_parameters which applies the
+        user defined rules to obtain spatially / temporally heterogeneous
+        parameters
+        """
+
+        def wrapper(*args):
+            self, x, u, params = args
+            params["eq_params"] = self.eval_heterogeneous_parameters(
+                x, u, params, self.eq_params_heterogeneity
+            )
+            new_args = args[:-1] + (params,)
+            res = evaluate(*new_args)
+            return res
+
+        return wrapper
+
 
 class PDENonStatio(DynamicLoss):
     r"""
@@ -122,3 +176,30 @@ class PDENonStatio(DynamicLoss):
             heterogeneity for no parameters.
         """
         super().__init__(Tmax, eq_params_heterogeneity)
+
+    def eval_heterogeneous_parameters(
+        self, t, x, u, params, eq_params_heterogeneity=None
+    ):
+        return super()._eval_heterogeneous_parameters(
+            t, x, u, params, eq_params_heterogeneity
+        )
+
+    @staticmethod
+    def evaluate_heterogeneous_parameters(evaluate):
+        """
+        Decorator which aims to decorate the evaluate methods of Dynamic losses
+        in order. It calls _eval_heterogeneous_parameters which applies the
+        user defined rules to obtain spatially / temporally heterogeneous
+        parameters
+        """
+
+        def wrapper(*args):
+            self, t, x, u, params = args
+            params["eq_params"] = self.eval_heterogeneous_parameters(
+                t, x, u, params, self.eq_params_heterogeneity
+            )
+            new_args = args[:-1] + (params,)
+            res = evaluate(*new_args)
+            return res
+
+        return wrapper

jinns/loss/_LossPDE.py

@@ -254,6 +254,7 @@ class LossPDEStatio(LossPDEAbstract):
         derivative_keys=None,
         omega_boundary_fun=None,
         omega_boundary_condition=None,
+        omega_boundary_dim=None,
         norm_key=None,
         norm_borders=None,
         norm_samples=None,
@@ -310,6 +311,12 @@ class LossPDEStatio(LossPDEAbstract):
             enforce a particular boundary condition on this facet.
             The facet called "xmin", resp. "xmax" etc., in 2D,
             refers to the set of 2D points with fixed "xmin", resp. "xmax", etc.
+        omega_boundary_dim
+            Either None, or a jnp.s_ or a dict of jnp.s_ with keys following
+            the logic of omega_boundary_fun. It indicates which dimension(s) of
+            the PINN will be forced to match the boundary condition
+            Note that it must be a slice and not an integer (a preprocessing of the
+            user provided argument takes care of it)
         norm_key
             Jax random key to draw samples in for the Monte Carlo computation
             of the normalization constant. Default is None
@@ -345,7 +352,7 @@ class LossPDEStatio(LossPDEAbstract):
                 raise ValueError(
                     f"obs_batch must be a list of size 2. You gave {len(obs_batch)}"
                 )
-            if any(isinstance(b, jnp.ndarray) for b in obs_batch):
+            if not all(isinstance(b, jnp.ndarray) for b in obs_batch):
                 raise ValueError("Every element of obs_batch should be a jnp.array.")
             n_obs = obs_batch[0].shape[0]
             if any(b.shape[0] != n_obs for b in obs_batch):
@@ -411,16 +418,49 @@ class LossPDEStatio(LossPDEAbstract):
 
         self.omega_boundary_fun = omega_boundary_fun
         self.omega_boundary_condition = omega_boundary_condition
+
+        self.omega_boundary_dim = omega_boundary_dim
+        if isinstance(self.omega_boundary_fun, dict):
+            if self.omega_boundary_dim is None:
+                self.omega_boundary_dim = {
+                    k: jnp.s_[::] for k in self.omega_boundary_fun.keys()
+                }
+            if list(self.omega_boundary_dim.keys()) != list(
+                self.omega_boundary_fun.keys()
+            ):
+                raise ValueError(
+                    "If omega_boundary_fun is a dict,"
+                    " omega_boundary_dim should be a dict with the same keys"
+                )
+            for k, v in self.omega_boundary_dim.items():
+                if isinstance(v, int):
+                    # rewrite it as a slice to ensure that axis does not disappear when
+                    # indexing
+                    self.omega_boundary_dim[k] = jnp.s_[v : v + 1]
+
+        else:
+            if self.omega_boundary_dim is None:
+                self.omega_boundary_dim = jnp.s_[::]
+            if isinstance(self.omega_boundary_dim, int):
+                # rewrite it as a slice to ensure that axis does not disappear when
+                # indexing
+                self.omega_boundary_dim = jnp.s_[
+                    self.omega_boundary_dim : self.omega_boundary_dim + 1
+                ]
+            if not isinstance(self.omega_boundary_dim, slice):
+                raise ValueError("self.omega_boundary_dim must be a jnp.s_" " object")
+
         self.dynamic_loss = dynamic_loss
         self.obs_batch = obs_batch
 
-        if sobolev_m is not None:
+        self.sobolev_m = sobolev_m
+        if self.sobolev_m is not None:
             self.sobolev_reg = _sobolev(
-                self.u, sobolev_m
+                self.u, self.sobolev_m
             )  # we return a function, that way
             # the order of sobolev_m is static and the conditional in the recursive
             # function is properly set
-            self.sobolev_m = sobolev_m
+            self.sobolev_m = self.sobolev_m
         else:
             self.sobolev_reg = None
 
@@ -512,7 +552,7 @@ class LossPDEStatio(LossPDEAbstract):
         if self.normalization_loss is not None:
             if isinstance(self.u, PINN):
                 v_u = vmap(
-                    lambda x: self.u(x, params_),
+                    lambda x: self.u(x, params_)[self.u.slice_solution],
                     (0),
                     0,
                 )
@@ -560,6 +600,7 @@ class LossPDEStatio(LossPDEAbstract):
                                 self.u,
                                 params_,
                                 idx,
+                                self.omega_boundary_dim[facet],
                             )
                         )
             else:
@@ -574,6 +615,7 @@ class LossPDEStatio(LossPDEAbstract):
                             self.u,
                             params_,
                             facet,
+                            self.omega_boundary_dim,
                         )
                     )
         else:
@@ -585,11 +627,11 @@ class LossPDEStatio(LossPDEAbstract):
             # TODO implement for SPINN
             if isinstance(self.u, PINN):
                 v_u = vmap(
-                    lambda x: self.u(x, params_),
+                    lambda x: self.u(x, params_)[self.u.slice_solution],
                     0,
                     0,
                 )
-                val = v_u(self.obs_batch[0][:, None])
+                val = v_u(self.obs_batch[0])
                 mse_observation_loss = jnp.mean(
                     self.loss_weights["observations"]
                     * jnp.sum(
@@ -655,6 +697,7 @@ class LossPDEStatio(LossPDEAbstract):
             "derivative_keys": self.derivative_keys,
             "omega_boundary_fun": self.omega_boundary_fun,
             "omega_boundary_condition": self.omega_boundary_condition,
+            "omega_boundary_dim": self.omega_boundary_dim,
             "norm_borders": self.norm_borders,
             "sobolev_m": self.sobolev_m,
         }
@@ -670,6 +713,7 @@ class LossPDEStatio(LossPDEAbstract):
             aux_data["derivative_keys"],
             aux_data["omega_boundary_fun"],
             aux_data["omega_boundary_condition"],
+            aux_data["omega_boundary_dim"],
             norm_key,
             aux_data["norm_borders"],
             norm_samples,
@@ -706,6 +750,7 @@ class LossPDENonStatio(LossPDEStatio):
         derivative_keys=None,
         omega_boundary_fun=None,
         omega_boundary_condition=None,
+        omega_boundary_dim=None,
         initial_condition_fun=None,
         norm_key=None,
         norm_borders=None,
@@ -760,6 +805,12 @@ class LossPDENonStatio(LossPDEStatio):
             enforce a particular boundary condition on this facet.
             The facet called "xmin", resp. "xmax" etc., in 2D,
             refers to the set of 2D points with fixed "xmin", resp. "xmax", etc.
+        omega_boundary_dim
+            Either None, or a jnp.s_ or a dict of jnp.s_ with keys following
+            the logic of omega_boundary_fun. It indicates which dimension(s) of
+            the PINN will be forced to match the boundary condition
+            Note that it must be a slice and not an integer (a preprocessing of the
+            user provided argument takes care of it)
         initial_condition_fun
             A function representing the temporal initial condition. If None
             (default) then no initial condition is applied
@@ -815,6 +866,7 @@ class LossPDENonStatio(LossPDEStatio):
             derivative_keys,
             omega_boundary_fun,
             omega_boundary_condition,
+            omega_boundary_dim,
             norm_key,
             norm_borders,
             norm_samples,
@@ -978,6 +1030,7 @@ class LossPDENonStatio(LossPDEStatio):
                                 self.u,
                                 params_,
                                 idx,
+                                self.omega_boundary_dim[facet],
                             )
                         )
             else:
@@ -993,6 +1046,7 @@ class LossPDENonStatio(LossPDEStatio):
                             self.u,
                             params_,
                             facet,
+                            self.omega_boundary_dim,
                         )
                     )
         else:
@@ -1036,7 +1090,7 @@ class LossPDENonStatio(LossPDEStatio):
             # TODO implement for SPINN
             if isinstance(self.u, PINN):
                 v_u = vmap(
-                    lambda t, x: self.u(t, x, params_),
+                    lambda t, x: self.u(t, x, params_)[self.u.slice_solution],
                     (0, 0),
                     0,
                 )
@@ -1108,6 +1162,7 @@ class LossPDENonStatio(LossPDEStatio):
             "derivative_keys": self.derivative_keys,
             "omega_boundary_fun": self.omega_boundary_fun,
             "omega_boundary_condition": self.omega_boundary_condition,
+            "omega_boundary_dim": self.omega_boundary_dim,
             "initial_condition_fun": self.initial_condition_fun,
             "norm_borders": self.norm_borders,
             "sobolev_m": self.sobolev_m,
@@ -1124,6 +1179,7 @@ class LossPDENonStatio(LossPDEStatio):
             aux_data["derivative_keys"],
             aux_data["omega_boundary_fun"],
             aux_data["omega_boundary_condition"],
+            aux_data["omega_boundary_dim"],
             aux_data["initial_condition_fun"],
             norm_key,
             aux_data["norm_borders"],
@@ -1163,6 +1219,7 @@ class SystemLossPDE:
         derivative_keys_dict=None,
         omega_boundary_fun_dict=None,
         omega_boundary_condition_dict=None,
+        omega_boundary_dim_dict=None,
         initial_condition_fun_dict=None,
         norm_key_dict=None,
         norm_borders_dict=None,
@@ -1199,15 +1256,17 @@ class SystemLossPDE:
             are not specified then the default behaviour for `derivative_keys`
             of LossODE is used
         omega_boundary_fun_dict
-            A dict of functions to be matched in the border condition, or a
-            dict of dict of functions (see doc for `omega_boundary_fun` in
-            LossPDEStatio or LossPDENonStatio).
-            Must share the keys of `u_dict`
+            A dict of dict of functions (see doc for `omega_boundary_fun` in
+            LossPDEStatio or LossPDENonStatio). Default is None.
+            Must share the keys of `u_dict`.
         omega_boundary_condition_dict
-            A dict of either None (no condition), or a string defining the boundary
-            condition e.g. Dirichlet or Von Neumann, or a dict of dict of
-            strings (see doc for `omega_boundary_fun` in
-            LossPDEStatio or LossPDENonStatio).
+            A dict of dict of strings (see doc for
+            `omega_boundary_condition_dict` in
+            LossPDEStatio or LossPDENonStatio). Default is None.
+            Must share the keys of `u_dict`
+        omega_boundary_dim_dict
+            A dict of dict of slices (see doc for `omega_boundary_dim` in
+            LossPDEStatio or LossPDENonStatio). Default is None.
             Must share the keys of `u_dict`
         initial_condition_fun_dict
             A dict of functions representing the temporal initial condition. If None
@@ -1265,6 +1324,10 @@ class SystemLossPDE:
             self.omega_boundary_condition_dict = {k: None for k in u_dict.keys()}
         else:
             self.omega_boundary_condition_dict = omega_boundary_condition_dict
+        if omega_boundary_dim_dict is None:
+            self.omega_boundary_dim_dict = {k: None for k in u_dict.keys()}
+        else:
+            self.omega_boundary_dim_dict = omega_boundary_dim_dict
         if initial_condition_fun_dict is None:
             self.initial_condition_fun_dict = {k: None for k in u_dict.keys()}
         else:
@@ -1309,6 +1372,7 @@ class SystemLossPDE:
             or u_dict.keys() != self.obs_batch_dict.keys()
             or u_dict.keys() != self.omega_boundary_fun_dict.keys()
             or u_dict.keys() != self.omega_boundary_condition_dict.keys()
+            or u_dict.keys() != self.omega_boundary_dim_dict.keys()
             or u_dict.keys() != self.initial_condition_fun_dict.keys()
             or u_dict.keys() != self.norm_key_dict.keys()
             or u_dict.keys() != self.norm_borders_dict.keys()
@@ -1345,6 +1409,7 @@ class SystemLossPDE:
                     derivative_keys=self.derivative_keys_dict[i],
                     omega_boundary_fun=self.omega_boundary_fun_dict[i],
                     omega_boundary_condition=self.omega_boundary_condition_dict[i],
+                    omega_boundary_dim=self.omega_boundary_dim_dict[i],
                     norm_key=self.norm_key_dict[i],
                     norm_borders=self.norm_borders_dict[i],
                     norm_samples=self.norm_samples_dict[i],
@@ -1366,6 +1431,7 @@ class SystemLossPDE:
                     derivative_keys=self.derivative_keys_dict[i],
                     omega_boundary_fun=self.omega_boundary_fun_dict[i],
                     omega_boundary_condition=self.omega_boundary_condition_dict[i],
+                    omega_boundary_dim=self.omega_boundary_dim_dict[i],
                     initial_condition_fun=self.initial_condition_fun_dict[i],
                     norm_key=self.norm_key_dict[i],
                     norm_borders=self.norm_borders_dict[i],

jinns/loss/_boundary_conditions.py

@@ -11,7 +11,7 @@ from jinns.utils._spinn import SPINN
 
 
 def _compute_boundary_loss_statio(
-    boundary_condition_type, f, border_batch, u, params, facet
+    boundary_condition_type, f, border_batch, u, params, facet, dim_to_apply
 ):
     r"""A generic function that will compute the mini-batch MSE of a
     boundary condition in the stationary case, given by:
@@ -44,6 +44,9 @@ def _compute_boundary_loss_statio(
     facet:
         An integer which represents the id of the facet which is currently
         considered (in the order provided wy the DataGenerator which is fixed)
+    dim_to_apply
+        A jnp.s_ object which indicates which dimension(s) of u will be forced
+        to match the boundary condition
 
     Returns
     -------
@@ -51,17 +54,24 @@ def _compute_boundary_loss_statio(
         the MSE computed on `border_batch`
     """
     if boundary_condition_type.lower() in "dirichlet":
-        mse = boundary_dirichlet_statio(f, border_batch, u, params)
+        mse = boundary_dirichlet_statio(f, border_batch, u, params, dim_to_apply)
     elif any(
         boundary_condition_type.lower() in s
         for s in ["von neumann", "vn", "vonneumann"]
     ):
-        mse = boundary_neumann_statio(f, border_batch, u, params, facet)
+        mse = boundary_neumann_statio(f, border_batch, u, params, facet, dim_to_apply)
     return mse
 
 
 def _compute_boundary_loss_nonstatio(
-    boundary_condition_type, f, times_batch, border_batch, u, params, facet
+    boundary_condition_type,
+    f,
+    times_batch,
+    border_batch,
+    u,
+    params,
+    facet,
+    dim_to_apply,
 ):
     r"""A generic function that will compute the mini-batch MSE of a
     boundary condition in the non-stationary case, given by:
@@ -95,6 +105,8 @@ def _compute_boundary_loss_nonstatio(
     facet:
         An integer which represents the id of the facet which is currently
         considered (in the order provided wy the DataGenerator which is fixed)
+    dim_to_apply
+        A jnp.s_ object. The dimension of u on which to apply the boundary condition
 
     Returns
     -------
@@ -102,16 +114,20 @@ def _compute_boundary_loss_nonstatio(
         the MSE computed on `border_batch`
     """
     if boundary_condition_type.lower() in "dirichlet":
-        mse = boundary_dirichlet_nonstatio(f, times_batch, border_batch, u, params)
+        mse = boundary_dirichlet_nonstatio(
+            f, times_batch, border_batch, u, params, dim_to_apply
+        )
     elif any(
         boundary_condition_type.lower() in s
         for s in ["von neumann", "vn", "vonneumann"]
     ):
-        mse = boundary_neumann_nonstatio(f, times_batch, border_batch, u, params, facet)
+        mse = boundary_neumann_nonstatio(
+            f, times_batch, border_batch, u, params, facet, dim_to_apply
+        )
     return mse
 
 
-def boundary_dirichlet_statio(f, border_batch, u, params):
+def boundary_dirichlet_statio(f, border_batch, u, params, dim_to_apply):
     r"""
     This omega boundary condition enforces a solution that is equal to f on
     border batch.
@@ -129,17 +145,19 @@ def boundary_dirichlet_statio(f, border_batch, u, params):
         Typically, it is a dictionary of
         dictionaries: `eq_params` and `nn_params``, respectively the
         differential equation parameters and the neural network parameter
+    dim_to_apply
+        A jnp.s_ object. The dimension of u on which to apply the boundary condition
     """
     if isinstance(u, PINN):
         v_u_boundary = vmap(
-            lambda dx: u(dx, params) - f(dx),
+            lambda dx: u(dx, params)[dim_to_apply] - f(dx),
             (0),
             0,
         )
 
         mse_u_boundary = jnp.sum((v_u_boundary(border_batch)) ** 2, axis=-1)
     elif isinstance(u, SPINN):
-        values = u(border_batch, params)
+        values = u(border_batch, params)[..., dim_to_apply]
         x_grid = _get_grid(border_batch)
         boundaries = _check_user_func_return(f(x_grid), values.shape)
         res = values - boundaries
@@ -150,7 +168,7 @@ def boundary_dirichlet_statio(f, border_batch, u, params):
     return mse_u_boundary
 
 
-def boundary_neumann_statio(f, border_batch, u, params, facet):
+def boundary_neumann_statio(f, border_batch, u, params, facet, dim_to_apply):
     r"""
     This omega boundary condition enforces a solution where :math:`\nabla u\cdot
     n` is equal to `f` on omega borders. :math:`n` is the unitary
@@ -173,6 +191,8 @@ def boundary_neumann_statio(f, border_batch, u, params, facet):
     facet:
         An integer which represents the id of the facet which is currently
         considered (in the order provided wy the DataGenerator which is fixed)
+    dim_to_apply
+        A jnp.s_ object. The dimension of u on which to apply the boundary condition
     """
     # We resort to the shape of the border_batch to determine the dimension as
     # described in the border_batch function
@@ -186,7 +206,7 @@ def boundary_neumann_statio(f, border_batch, u, params, facet):
         n = jnp.array([[-1, 1, 0, 0], [0, 0, -1, 1]])
 
     if isinstance(u, PINN):
-        u_ = lambda x, params: u(x, params)[0]
+        u_ = lambda x, params: jnp.squeeze(u(x, params)[dim_to_apply])
         v_neumann = vmap(
             lambda dx: jnp.dot(
                 grad(u_, 0)(dx, params),
@@ -206,7 +226,7 @@ def boundary_neumann_statio(f, border_batch, u, params, facet):
                 lambda x: u(
                     x,
                     params,
-                ),
+                )[..., dim_to_apply],
                 (border_batch,),
                 (jnp.ones_like(border_batch),),
             )
@@ -246,7 +266,9 @@ def boundary_neumann_statio(f, border_batch, u, params, facet):
     return mse_u_boundary
 
 
-def boundary_dirichlet_nonstatio(f, times_batch, omega_border_batch, u, params):
+def boundary_dirichlet_nonstatio(
+    f, times_batch, omega_border_batch, u, params, dim_to_apply
+):
     """
     This omega boundary condition enforces a solution that is equal to f
     at times_batch x omega borders
@@ -266,6 +288,8 @@ def boundary_dirichlet_nonstatio(f, times_batch, omega_border_batch, u, params):
         Typically, it is a dictionary of
         dictionaries: `eq_params` and `nn_params``, respectively the
         differential equation parameters and the neural network parameter
+    dim_to_apply
+        A jnp.s_ object. The dimension of u on which to apply the boundary condition
     """
     if isinstance(u, PINN):
         tile_omega_border_batch = jnp.tile(
@@ -280,7 +304,7 @@ def boundary_dirichlet_nonstatio(f, times_batch, omega_border_batch, u, params):
                 t,
                 dx,
                 params,
-            )
+            )[dim_to_apply]
             - f(t, dx),
             (0, 0),
             0,
@@ -304,7 +328,7 @@ def boundary_dirichlet_nonstatio(f, times_batch, omega_border_batch, u, params):
         # otherwise we require batches to have same shape and we do not need
         # this operation
 
-        values = u(times_batch, tile_omega_border_batch, params)
+        values = u(times_batch, tile_omega_border_batch, params)[..., dim_to_apply]
         tx_grid = _get_grid(jnp.concatenate([times_batch, omega_border_batch], axis=-1))
         boundaries = _check_user_func_return(
             f(tx_grid[..., 0:1], tx_grid[..., 1:]), values.shape
@@ -314,7 +338,9 @@ def boundary_dirichlet_nonstatio(f, times_batch, omega_border_batch, u, params):
     return mse_u_boundary
 
 
-def boundary_neumann_nonstatio(f, times_batch, omega_border_batch, u, params, facet):
+def boundary_neumann_nonstatio(
+    f, times_batch, omega_border_batch, u, params, facet, dim_to_apply
+):
     r"""
     This omega boundary condition enforces a solution where :math:`\nabla u\cdot
     n` is equal to `f` at time_batch x omega borders. :math:`n` is the unitary
@@ -338,6 +364,8 @@ def boundary_neumann_nonstatio(f, times_batch, omega_border_batch, u, params, fa
     facet:
         An integer which represents the id of the facet which is currently
         considered (in the order provided wy the DataGenerator which is fixed)
+    dim_to_apply
+        A jnp.s_ object. The dimension of u on which to apply the boundary condition
     """
     # We resort to the shape of the border_batch to determine the dimension as
     # described in the border_batch function
@@ -358,7 +386,7 @@ def boundary_neumann_nonstatio(f, times_batch, omega_border_batch, u, params, fa
         def rep_times(k):
             return jnp.repeat(times_batch, k, axis=0)
 
-        u_ = lambda t, x, params: u(t, x, params)[0]
+        u_ = lambda t, x, params: jnp.squeeze(u(t, x, params)[dim_to_apply])
         v_neumann = vmap(
             lambda t, dx: jnp.dot(
                 grad(u_, 1)(t, dx, params),
@@ -389,7 +417,7 @@ def boundary_neumann_nonstatio(f, times_batch, omega_border_batch, u, params, fa
         # high dim output at once
         if omega_border_batch.shape[0] == 1:  # i.e. case 1D
             _, du_dx = jax.jvp(
-                lambda x: u(times_batch, x, params),
+                lambda x: u(times_batch, x, params)[..., dim_to_apply],
                 (omega_border_batch,),
                 (jnp.ones_like(omega_border_batch),),
             )
@@ -402,12 +430,12 @@ def boundary_neumann_nonstatio(f, times_batch, omega_border_batch, u, params, fa
                 jnp.array([0.0, 1.0])[None], omega_border_batch.shape[0], axis=0
             )
             _, du_dx1 = jax.jvp(
-                lambda x: u(times_batch, x, params),
+                lambda x: u(times_batch, x, params)[..., dim_to_apply],
                 (omega_border_batch,),
                 (tangent_vec_0,),
             )
             _, du_dx2 = jax.jvp(
-                lambda x: u(times_batch, x, params),
+                lambda x: u(times_batch, x, params)[..., dim_to_apply],
                 (omega_border_batch,),
                 (tangent_vec_1,),
             )

jinns/loss/_operators.py

@@ -9,7 +9,7 @@ from jinns.utils._pinn import PINN
 from jinns.utils._spinn import SPINN
 
 
-def _div_rev(u, params, x, t=None):
+def _div_rev(t, x, u, params):
     r"""
     Compute the divergence of a vector field :math:`\mathbf{u}`, i.e.,
     :math:`\nabla \cdot \mathbf{u}(\mathbf{x})` with :math:`\mathbf{u}` a vector
@@ -21,14 +21,14 @@ def _div_rev(u, params, x, t=None):
         if t is None:
             du_dxi = grad(lambda x, params: u(x, params)[i], 0)(x, params)[i]
         else:
-            du_dxi = grad(lambda t, x, params: u(x, params)[i], 1)(x, params)[i]
+            du_dxi = grad(lambda t, x, params: u(t, x, params)[i], 1)(t, x, params)[i]
         return _, du_dxi
 
     _, accu = jax.lax.scan(scan_fun, {}, jnp.arange(x.shape[0]))
     return jnp.sum(accu)
 
 
-def _div_fwd(u, params, x, t=None):
+def _div_fwd(t, x, u, params):
     r"""
     Compute the divergence of a **batched** vector field :math:`\mathbf{u}`, i.e.,
     :math:`\nabla \cdot \mathbf{u}(\mathbf{x})` with :math:`\mathbf{u}` a vector
@@ -55,7 +55,7 @@ def _div_fwd(u, params, x, t=None):
     return jnp.sum(accu, axis=0)
 
 
-def _laplacian_rev(u, params, x, t=None):
+def _laplacian_rev(t, x, u, params):
     r"""
     Compute the Laplacian of a scalar field :math:`u` (from :math:`\mathbb{R}^d`
     to :math:`\mathbb{R}`) for :math:`\mathbf{x}` of arbitrary dimension, i.e.,
@@ -98,7 +98,7 @@ def _laplacian_rev(u, params, x, t=None):
     # return jnp.sum(trace_hessian)
 
 
-def _laplacian_fwd(u, params, x, t=None):
+def _laplacian_fwd(t, x, u, params):
     r"""
     Compute the Laplacian of a **batched** scalar field :math:`u`
     (from :math:`\mathbb{R}^{b\times d}` to :math:`\mathbb{R}^{b\times b}`)
@@ -134,7 +134,7 @@ def _laplacian_fwd(u, params, x, t=None):
     return jnp.sum(trace_hessian, axis=0)
 
 
-def _vectorial_laplacian(u, params, x, t=None, u_vec_ndim=None):
+def _vectorial_laplacian(t, x, u, params, u_vec_ndim=None):
     r"""
     Compute the vectorial Laplacian of a vector field :math:`\mathbf{u}` (from
     :math:`\mathbb{R}^d`
@@ -163,7 +163,7 @@ def _vectorial_laplacian(u, params, x, t=None, u_vec_ndim=None):
                 uj = lambda x, params: jnp.expand_dims(u(x, params)[j], axis=-1)
             else:
                 uj = lambda t, x, params: jnp.expand_dims(u(t, x, params)[j], axis=-1)
-            lap_on_j = _laplacian_rev(uj, params, x, t)
+            lap_on_j = _laplacian_rev(t, x, uj, params)
         elif isinstance(u, SPINN):
             if t is None:
                 uj = lambda x, params: jnp.expand_dims(u(x, params)[..., j], axis=-1)
@@ -171,7 +171,7 @@ def _vectorial_laplacian(u, params, x, t=None, u_vec_ndim=None):
                 uj = lambda t, x, params: jnp.expand_dims(
                     u(t, x, params)[..., j], axis=-1
                 )
-            lap_on_j = _laplacian_fwd(uj, params, x, t)
+            lap_on_j = _laplacian_fwd(t, x, uj, params)
 
         return _, lap_on_j
 
@@ -179,7 +179,7 @@ def _vectorial_laplacian(u, params, x, t=None, u_vec_ndim=None):
     return vec_lap
 
 
-def _u_dot_nabla_times_u_rev(u, params, x, t=None):
+def _u_dot_nabla_times_u_rev(t, x, u, params):
     r"""
     Implement :math:`((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})` for
     :math:`\mathbf{x}` of arbitrary
@@ -224,7 +224,7 @@ def _u_dot_nabla_times_u_rev(u, params, x, t=None):
     raise NotImplementedError("x.ndim must be 2")
 
 
-def _u_dot_nabla_times_u_fwd(u, params, x, t=None):
+def _u_dot_nabla_times_u_fwd(t, x, u, params):
     r"""
     Implement :math:`((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})` for
     :math:`\mathbf{x}` of arbitrary dimension **with a batch dimension**.

jinns/utils/_pinn.py

@@ -64,16 +64,44 @@ class PINN:
     The function create_PINN has the role to population the `__call__` function
     """
 
-    def __init__(self, key, eqx_list, output_slice=None):
+    def __init__(
+        self,
+        key,
+        eqx_list,
+        slice_solution,
+        eq_type,
+        input_transform,
+        output_transform,
+        output_slice=None,
+    ):
         _pinn = _MLP(key, eqx_list)
         self.params, self.static = eqx.partition(_pinn, eqx.is_inexact_array)
+        self.slice_solution = slice_solution
+        self.eq_type = eq_type
+        self.input_transform = input_transform
+        self.output_transform = output_transform
         self.output_slice = output_slice
 
     def init_params(self):
         return self.params
 
-    def __call__(self, *args, **kwargs):
-        return self.apply_fn(self, *args, **kwargs)
+    def __call__(self, *args):
+        if self.eq_type == "ODE":
+            (t, params) = args
+            t = t[None]  #  Add dimension which is lacking for the ODE batches
+            return self._eval_nn(
+                t, params, self.input_transform, self.output_transform
+            ).squeeze()
+        if self.eq_type == "statio_PDE":
+            (x, params) = args
+            return self._eval_nn(x, params, self.input_transform, self.output_transform)
+        if self.eq_type == "nonstatio_PDE":
+            (t, x, params) = args
+            t_x = jnp.concatenate([t, x], axis=-1)
+            return self._eval_nn(
+                t_x, params, self.input_transform, self.output_transform
+            )
+        raise ValueError("Wrong value for self.eq_type")
 
     def _eval_nn(self, inputs, params, input_transform, output_transform):
         """
@@ -82,7 +110,7 @@ class PINN:
         """
         try:
             model = eqx.combine(params["nn_params"], self.static)
-        except:  # give more flexibility
+        except (KeyError, TypeError) as e:  # give more flexibility
             model = eqx.combine(params, self.static)
         res = output_transform(inputs, model(input_transform(inputs, params)).squeeze())
 
@@ -103,6 +131,7 @@ def create_PINN(
     input_transform=None,
     output_transform=None,
     shared_pinn_outputs=None,
+    slice_solution=None,
 ):
     """
     Utility function to create a standard PINN neural network with the equinox
@@ -152,6 +181,13 @@ def create_PINN(
         network. In this case we return a list of PINNs, one for each output in
         shared_pinn_outputs. This is useful to create PINNs that share the
         same network and same parameters. Default is None, we only return one PINN.
+    slice_solution
+        A jnp.s_ object which indicates which axis of the PINN output is
+        dedicated to the actual equation solution. Default None
+        means that slice_solution = the whole PINN output. This argument is useful
+        when the PINN is also used to output equation parameters for example
+        Note that it must be a slice and not an integer (a preprocessing of the
+        user provided argument takes care of it)
 
 
     Returns
@@ -182,6 +218,19 @@ def create_PINN(
     if eq_type != "ODE" and dim_x == 0:
         raise RuntimeError("Wrong parameter combination eq_type and dim_x")
 
+    try:
+        nb_outputs_declared = eqx_list[-1][2]  # normally we look for 3rd ele of
+        # last layer
+    except IndexError:
+        nb_outputs_declared = eqx_list[-2][2]
+
+    if slice_solution is None:
+        slice_solution = jnp.s_[0:nb_outputs_declared]
+    if isinstance(slice_solution, int):
+        # rewrite it as a slice to ensure that axis does not disappear when
+        # indexing
+        slice_solution = jnp.s_[slice_solution : slice_solution + 1]
+
     if input_transform is None:
 
         def input_transform(_in, _params):
@@ -192,34 +241,19 @@ def create_PINN(
         def output_transform(_in_pinn, _out_pinn):
             return _out_pinn
 
-    if eq_type == "ODE":
-
-        def apply_fn(self, t, params):
-            t = t[
-                None
-            ]  # Note that we added a dimension to t which is lacking for the ODE batches
-            return self._eval_nn(t, params, input_transform, output_transform).squeeze()
-
-    elif eq_type == "statio_PDE":
-        # Here we add an argument `x` which can be high dimensional
-        def apply_fn(self, x, params):
-            return self._eval_nn(x, params, input_transform, output_transform)
-
-    elif eq_type == "nonstatio_PDE":
-        # Here we add an argument `x` which can be high dimensional
-        def apply_fn(self, t, x, params):
-            t_x = jnp.concatenate([t, x], axis=-1)
-            return self._eval_nn(t_x, params, input_transform, output_transform)
-
-    else:
-        raise RuntimeError("Wrong parameter value for eq_type")
-
     if shared_pinn_outputs is not None:
         pinns = []
         static = None
         for output_slice in shared_pinn_outputs:
-            pinn = PINN(key, eqx_list, output_slice)
-            pinn.apply_fn = apply_fn
+            pinn = PINN(
+                key,
+                eqx_list,
+                slice_solution,
+                eq_type,
+                input_transform,
+                output_transform,
+                output_slice,
+            )
             # all the pinns are in fact the same so we share the same static
             if static is None:
                 static = pinn.static
@@ -227,6 +261,7 @@ def create_PINN(
                 pinn.static = static
             pinns.append(pinn)
         return pinns
-    pinn = PINN(key, eqx_list)
-    pinn.apply_fn = apply_fn
+    pinn = PINN(
+        key, eqx_list, slice_solution, eq_type, input_transform, output_transform
+    )
     return pinn

jinns/utils/_spinn.py

@@ -82,16 +82,35 @@ class SPINN:
     The function create_SPINN has the role to population the `__call__` function
     """
 
-    def __init__(self, key, d, r, eqx_list, m=1):
+    def __init__(self, key, d, r, eqx_list, eq_type, m=1):
         self.d, self.r, self.m = d, r, m
         _spinn = _SPINN(key, d, r, eqx_list, m)
         self.params, self.static = eqx.partition(_spinn, eqx.is_inexact_array)
+        self.eq_type = eq_type
 
     def init_params(self):
         return self.params
 
-    def __call__(self, *args, **kwargs):
-        return self.apply_fn(self, *args, **kwargs)
+    def __call__(self, *args):
+        if self.eq_type == "statio_PDE":
+            (x, params) = args
+            try:
+                spinn = eqx.combine(params["nn_params"], self.static)
+            except (KeyError, TypeError) as e:
+                spinn = eqx.combine(params, self.static)
+            v_model = jax.vmap(spinn, (0))
+            res = v_model(t=None, x=x)
+            return self._eval_nn(res)
+        if self.eq_type == "nonstatio_PDE":
+            (t, x, params) = args
+            try:
+                spinn = eqx.combine(params["nn_params"], self.static)
+            except (KeyError, TypeError) as e:
+                spinn = eqx.combine(params, self.static)
+            v_model = jax.vmap(spinn, ((0, 0)))
+            res = v_model(t, x)
+            return self._eval_nn(res)
+        raise RuntimeError("Wrong parameter value for eq_type")
 
     def _eval_nn(self, res):
         """
@@ -209,32 +228,6 @@ def create_SPINN(key, d, r, eqx_list, eq_type, m=1):
             "Too many dimensions, not enough letters available in jnp.einsum"
         )
 
-    if eq_type == "statio_PDE":
-
-        def apply_fn(self, x, params):
-            try:
-                spinn = eqx.combine(params["nn_params"], self.static)
-            except:  # give more flexibility
-                spinn = eqx.combine(params, self.static)
-            v_model = jax.vmap(spinn, (0))
-            res = v_model(t=None, x=x)
-            return self._eval_nn(res)
-
-    elif eq_type == "nonstatio_PDE":
-
-        def apply_fn(self, t, x, params):
-            try:
-                spinn = eqx.combine(params["nn_params"], self.static)
-            except:  # give more flexibility
-                spinn = eqx.combine(params, self.static)
-            v_model = jax.vmap(spinn, ((0, 0)))
-            res = v_model(t, x)
-            return self._eval_nn(res)
-
-    else:
-        raise RuntimeError("Wrong parameter value for eq_type")
-
-    spinn = SPINN(key, d, r, eqx_list, m)
-    spinn.apply_fn = apply_fn
+    spinn = SPINN(key, d, r, eqx_list, eq_type, m)
 
     return spinn

{jinns-0.6.0.dist-info → jinns-0.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.6.0
+Version: 0.6.1
 Summary: Physics Informed Neural Network with JAX
 Author-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
 Maintainer-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-0.6.0.dist-info → jinns-0.6.1.dist-info}/RECORD

@@ -2,24 +2,24 @@ jinns/__init__.py,sha256=Nw5pdlmDhJwco3bXX3YttkeCF8czX_6m0poh8vu0lDQ,113
 jinns/data/_DataGenerators.py,sha256=nIuKtkX4V4ckfT4-g0bjlY7BLkgcok5JbI9OzJn73mA,44461
 jinns/data/__init__.py,sha256=S13J59Fxuph4uNJ542fP_Mj8U72ilhb5t_UQ-c1k3nY,232
 jinns/data/_display.py,sha256=NfINLJAGmQSPz30cWVaeQFpabzCXprp4RNH6Iycx-VU,7722
-jinns/loss/_DynamicLoss.py,sha256=_2ezagpJhnOKxH8--r-Dld4YN_9GAGtGbBO2ehr8uyM,28250
-jinns/loss/_DynamicLossAbstract.py,sha256=Sldyu3C8VKLUKhj7Hb4dU84dq0RFvwa3DuKFRD1mK7A,4819
+jinns/loss/_DynamicLoss.py,sha256=JOn6rpVcM825QHFWBQBOGCrR9tvStDaZiaB6pxZ0fhE,27336
+jinns/loss/_DynamicLossAbstract.py,sha256=NnvS8WKJZZZiGMfidxpGlkmWtINn_UDj3BNXe9l9O84,7787
 jinns/loss/_LossODE.py,sha256=nYWzUVO6vpZMYXZkeUpDRMZSowyGfGWUdwzBNciQkyo,20639
-jinns/loss/_LossPDE.py,sha256=RrZRhMEIHbQn9Fc1SjkJoGiaOHHFaQJdFBbx7V4Ug1E,68822
+jinns/loss/_LossPDE.py,sha256=gPi75x4MTjNSxY0Vz7swnsG1serPOtIPeKYsV5uYkr4,72221
 jinns/loss/__init__.py,sha256=pFNYUxns-NPXBFdqrEVSiXkQLfCtKw-t2trlhvLzpYE,355
-jinns/loss/_boundary_conditions.py,sha256=nKxVSgq9GfmxDQlSilzFFq3uV0t2GjxvUirCtBZvcE0,15235
-jinns/loss/_operators.py,sha256=ye4JP8SVJ3Shr7nWCiIxa5KmgPMZcCcHduRgq8cm-J4,11098
+jinns/loss/_boundary_conditions.py,sha256=K4cLYe5ecW2gJznv51nM3zOvE4UO-ZvVmVp3VIiuddM,16296
+jinns/loss/_operators.py,sha256=zDGJqYqeYH7xd-4dtGX9PS-pf0uSOpUUXGo5SVjIJ4o,11069
 jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jinns/solver/_rar.py,sha256=ZoklcHq5MFLUsIiWMUH_426MRxq7BUZkD86jcQYau0I,13918
 jinns/solver/_seq2seq.py,sha256=ihHnb6UpvShgEXOynUhsrPyt0wqNXIjeEL8HxvgRRHE,5985
 jinns/solver/_solve.py,sha256=-Ji16NJ9CMBcFDKLFpBibvQwow9ysAB4pVGwRRqoXRE,10138
 jinns/utils/__init__.py,sha256=bksGuq0mNoqciKIhEA8wOHkUppYQYJqOiHK2SEn5jac,227
 jinns/utils/_optim.py,sha256=Q8a_dMqb6rjdv1qGibgvl7cv8DENfXsthDJ56RziWJc,4441
-jinns/utils/_pinn.py,sha256=oxAuqTt3biGx66sOZAg_NrioaxqL9BAWI0Jj6lQfl7o,8142
-jinns/utils/_spinn.py,sha256=RB3erlaD2iQ1_SalsakIslBF-3Iizatb9NLG97SZ5es,7943
+jinns/utils/_pinn.py,sha256=wEbCU4aDT8nd1sz72ZrgR8e-xlR-P0iczHk_hnnZQBo,9427
+jinns/utils/_spinn.py,sha256=RZveR5R2HKMyQKuWQE_H6WKAfxABCkB6YGKyDZK051w,7889
 jinns/utils/_utils.py,sha256=b0-R6jtyMescaDfod0DuzeLMpm1C7gTwpU1MZEW2468,6060
-jinns-0.6.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.6.0.dist-info/METADATA,sha256=gezTMTKxIyTGKWMKkhylm2mNxNiZY4w_NysNbyiizn4,2388
-jinns-0.6.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-jinns-0.6.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.6.0.dist-info/RECORD,,
+jinns-0.6.1.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.6.1.dist-info/METADATA,sha256=yUHlhcVpsm8b1vHJSy5_TZjEUaaufI3niTvWZNXehks,2388
+jinns-0.6.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+jinns-0.6.1.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.6.1.dist-info/RECORD,,