jinns 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

jinns/loss/_DynamicLossAbstract.py

@@ -21,53 +21,22 @@ else:
     from equinox import AbstractClassVar
 
 
-def _decorator_heteregeneous_params(evaluate, eq_type):
+def _decorator_heteregeneous_params(evaluate):
 
-    def wrapper_ode(*args):
-        self, t, u, params = args
+    def wrapper(*args):
+        self, inputs, u, params = args
         _params = eqx.tree_at(
             lambda p: p.eq_params,
             params,
             self._eval_heterogeneous_parameters(
-                t, None, u, params, self.eq_params_heterogeneity
+                inputs, u, params, self.eq_params_heterogeneity
             ),
         )
         new_args = args[:-1] + (_params,)
         res = evaluate(*new_args)
         return res
 
-    def wrapper_pde_statio(*args):
-        self, x, u, params = args
-        _params = eqx.tree_at(
-            lambda p: p.eq_params,
-            params,
-            self._eval_heterogeneous_parameters(
-                None, x, u, params, self.eq_params_heterogeneity
-            ),
-        )
-        new_args = args[:-1] + (_params,)
-        res = evaluate(*new_args)
-        return res
-
-    def wrapper_pde_non_statio(*args):
-        self, t, x, u, params = args
-        _params = eqx.tree_at(
-            lambda p: p.eq_params,
-            params,
-            self._eval_heterogeneous_parameters(
-                t, x, u, params, self.eq_params_heterogeneity
-            ),
-        )
-        new_args = args[:-1] + (_params,)
-        res = evaluate(*new_args)
-        return res
-
-    if eq_type == "ODE":
-        return wrapper_ode
-    elif eq_type == "Statio PDE":
-        return wrapper_pde_statio
-    elif eq_type == "Non-statio PDE":
-        return wrapper_pde_non_statio
+    return wrapper
 
 
 class DynamicLoss(eqx.Module):
@@ -110,8 +79,7 @@ class DynamicLoss(eqx.Module):
 
     def _eval_heterogeneous_parameters(
         self,
-        t: Float[Array, "1"],
-        x: Float[Array, "dim"],
+        inputs: Float[Array, "1"] | Float[Array, "dim"] | Float[Array, "1+dim"],
         u: eqx.Module,
         params: Params | ParamsDict,
         eq_params_heterogeneity: Dict[str, Callable | None] = None,
@@ -124,14 +92,7 @@ class DynamicLoss(eqx.Module):
                 if eq_params_heterogeneity[k] is None:
                     eq_params_[k] = p
                 else:
-                    # heterogeneity encoded through a function whose
-                    # signature will vary according to _eq_type
-                    if self._eq_type == "ODE":
-                        eq_params_[k] = eq_params_heterogeneity[k](t, u, params)
-                    elif self._eq_type == "Statio PDE":
-                        eq_params_[k] = eq_params_heterogeneity[k](x, u, params)
-                    elif self._eq_type == "Non-statio PDE":
-                        eq_params_[k] = eq_params_heterogeneity[k](t, x, u, params)
+                    eq_params_[k] = eq_params_heterogeneity[k](inputs, u, params)
             except KeyError:
                 # we authorize missing eq_params_heterogeneity key
                 # if its heterogeneity is None anyway
@@ -140,22 +101,17 @@ class DynamicLoss(eqx.Module):
 
     def _evaluate(
         self,
-        t: Float[Array, "1"],
-        x: Float[Array, "dim"],
+        inputs: Float[Array, "1"] | Float[Array, "dim"] | Float[Array, "1+dim"],
         u: eqx.Module,
         params: Params | ParamsDict,
     ) -> float:
-        # Here we handle the various possible signature
-        if self._eq_type == "ODE":
-            ans = self.equation(t, u, params)
-        elif self._eq_type == "Statio PDE":
-            ans = self.equation(x, u, params)
-        elif self._eq_type == "Non-statio PDE":
-            ans = self.equation(t, x, u, params)
-        else:
-            raise NotImplementedError("the equation type is not handled.")
-
-        return ans
+        evaluation = self.equation(inputs, u, params)
+        if len(evaluation.shape) == 0:
+            raise ValueError(
+                "The output of dynamic loss must be vectorial, "
+                "i.e. of shape (d,) with d >= 1"
+            )
+        return evaluation
 
     @abc.abstractmethod
     def equation(self, *args, **kwargs):
@@ -191,7 +147,7 @@ class ODE(DynamicLoss):
 
     _eq_type: ClassVar[str] = "ODE"
 
-    @partial(_decorator_heteregeneous_params, eq_type="ODE")
+    @partial(_decorator_heteregeneous_params)
     def evaluate(
         self,
         t: Float[Array, "1"],
@@ -199,7 +155,7 @@ class ODE(DynamicLoss):
         params: Params | ParamsDict,
     ) -> float:
         """Here we call DynamicLoss._evaluate with x=None"""
-        return self._evaluate(t, None, u, params)
+        return self._evaluate(t, u, params)
 
     @abc.abstractmethod
     def equation(
@@ -260,12 +216,12 @@ class PDEStatio(DynamicLoss):
 
     _eq_type: ClassVar[str] = "Statio PDE"
 
-    @partial(_decorator_heteregeneous_params, eq_type="Statio PDE")
+    @partial(_decorator_heteregeneous_params)
     def evaluate(
         self, x: Float[Array, "dimension"], u: eqx.Module, params: Params | ParamsDict
     ) -> float:
         """Here we call the DynamicLoss._evaluate with t=None"""
-        return self._evaluate(None, x, u, params)
+        return self._evaluate(x, u, params)
 
     @abc.abstractmethod
     def equation(
@@ -325,23 +281,21 @@ class PDENonStatio(DynamicLoss):
 
     _eq_type: ClassVar[str] = "Non-statio PDE"
 
-    @partial(_decorator_heteregeneous_params, eq_type="Non-statio PDE")
+    @partial(_decorator_heteregeneous_params)
     def evaluate(
         self,
-        t: Float[Array, "1"],
-        x: Float[Array, "dim"],
+        t_x: Float[Array, "1 + dim"],
         u: eqx.Module,
         params: Params | ParamsDict,
     ) -> float:
         """Here we call the DynamicLoss._evaluate with full arguments"""
-        ans = self._evaluate(t, x, u, params)
+        ans = self._evaluate(t_x, u, params)
         return ans
 
     @abc.abstractmethod
     def equation(
         self,
-        t: Float[Array, "1"],
-        x: Float[Array, "dim"],
+        t_x: Float[Array, "1 + dim"],
         u: eqx.Module,
         params: Params | ParamsDict,
     ) -> float:
@@ -353,10 +307,8 @@ class PDENonStatio(DynamicLoss):
 
         Parameters
         ----------
-        t : Float[Array, "1"]
-            A 1-dimensional jnp.array representing the time point.
-        x : Float[Array, "d"]
-            A `d` dimensional jnp.array representing a point in the spatial domain $\Omega$.
+        t_x : Float[Array, "1 + dim"]
+            A jnp array containing the concatenation of a time point and a point in $\Omega$
         u : eqx.Module
             The neural network.
         params : Params | ParamsDict

jinns/loss/_LossODE.py

@@ -209,7 +209,7 @@ class LossODE(_LossODEAbstract):
             mse_dyn_loss = dynamic_loss_apply(
                 self.dynamic_loss.evaluate,
                 self.u,
-                (temporal_batch,),
+                temporal_batch,
                 _set_derivatives(params, self.derivative_keys.dyn_loss),
                 self.vmap_in_axes + vmap_in_axes_params,
                 self.loss_weights.dyn_loss,
@@ -227,7 +227,7 @@ class LossODE(_LossODEAbstract):
             else:
                 v_u = vmap(self.u, (None,) + vmap_in_axes_params)
             t0, u0 = self.initial_condition  # pylint: disable=unpacking-non-sequence
-            t0 = jnp.array(t0)
+            t0 = jnp.array([t0])
             u0 = jnp.array(u0)
             mse_initial_condition = jnp.mean(
                 self.loss_weights.initial_condition
@@ -522,7 +522,7 @@ class SystemLossODE(eqx.Module):
             return dynamic_loss_apply(
                 dyn_loss.evaluate,
                 self.u_dict,
-                (temporal_batch,),
+                temporal_batch,
                 _set_derivatives(params_dict, self.derivative_keys_dyn_loss.dyn_loss),
                 vmap_in_axes_t + vmap_in_axes_params,
                 loss_weight,

jinns/loss/_LossPDE.py

@@ -22,9 +22,7 @@ from jinns.loss._loss_utils import (
     initial_condition_apply,
     constraints_system_loss_apply,
 )
-from jinns.data._DataGenerators import (
-    append_obs_batch,
-)
+from jinns.data._DataGenerators import append_obs_batch
 from jinns.parameters._params import (
     _get_vmap_in_axes_params,
     _update_eq_params_dict,
@@ -370,8 +368,8 @@ class LossPDEStatio(_LossPDEAbstract):
 
     def _get_dynamic_loss_batch(
         self, batch: PDEStatioBatch
-    ) -> tuple[Float[Array, "batch_size dimension"]]:
-        return (batch.inside_batch,)
+    ) -> Float[Array, "batch_size dimension"]:
+        return batch.domain_batch
 
     def _get_normalization_loss_batch(
         self, _
@@ -432,7 +430,7 @@ class LossPDEStatio(_LossPDEAbstract):
                 self.u,
                 self._get_normalization_loss_batch(batch),
                 _set_derivatives(params, self.derivative_keys.norm_loss),
-                self.vmap_in_axes + vmap_in_axes_params,
+                vmap_in_axes_params,
                 self.norm_int_length,
                 self.loss_weights.norm_loss,
             )
@@ -575,6 +573,9 @@ class LossPDENonStatio(LossPDEStatio):
         kw_only=True, default=None, static=True
     )
 
+    _max_norm_samples_omega: Int = eqx.field(init=False, static=True)
+    _max_norm_time_slices: Int = eqx.field(init=False, static=True)
+
     def __post_init__(self, params=None):
         """
         Note that neither __init__ or __post_init__ are called when udating a
@@ -585,7 +586,7 @@ class LossPDENonStatio(LossPDEStatio):
         )  # because __init__ or __post_init__ of Base
         # class is not automatically called
 
-        self.vmap_in_axes = (0, 0)  # for t and x
+        self.vmap_in_axes = (0,)  # for t_x
 
         if self.initial_condition_fun is None:
             warnings.warn(
@@ -593,28 +594,28 @@ class LossPDENonStatio(LossPDEStatio):
                 "case (e.g by. hardcoding it into the PINN output)."
             )
 
+        # witht the variables below we avoid memory overflow since a cartesian
+        # product is taken
+        self._max_norm_time_slices = 100
+        self._max_norm_samples_omega = 1000
+
     def _get_dynamic_loss_batch(
         self, batch: PDENonStatioBatch
-    ) -> tuple[Float[Array, "batch_size 1"], Float[Array, "batch_size dimension"]]:
-        times_batch = batch.times_x_inside_batch[:, 0:1]
-        omega_batch = batch.times_x_inside_batch[:, 1:]
-        return (times_batch, omega_batch)
+    ) -> Float[Array, "batch_size 1+dimension"]:
+        return batch.domain_batch
 
     def _get_normalization_loss_batch(
         self, batch: PDENonStatioBatch
-    ) -> tuple[Float[Array, "batch_size 1"], Float[Array, "nb_norm_samples dimension"]]:
+    ) -> Float[Array, "nb_norm_time_slices nb_norm_samples dimension"]:
         return (
-            batch.times_x_inside_batch[:, 0:1],
-            self.norm_samples,
+            batch.domain_batch[: self._max_norm_time_slices, 0:1],
+            self.norm_samples[: self._max_norm_samples_omega],
         )
 
     def _get_observations_loss_batch(
         self, batch: PDENonStatioBatch
     ) -> tuple[Float[Array, "batch_size 1"], Float[Array, "batch_size dimension"]]:
-        return (
-            batch.obs_batch_dict["pinn_in"][:, 0:1],
-            batch.obs_batch_dict["pinn_in"][:, 1:],
-        )
+        return (batch.obs_batch_dict["pinn_in"],)
 
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
@@ -637,7 +638,7 @@ class LossPDENonStatio(LossPDEStatio):
             of parameters (eg. for metamodeling) and an optional additional batch of observed
             inputs/outputs/parameters
         """
-        omega_batch = batch.times_x_inside_batch[:, 1:]
+        omega_batch = batch.initial_batch
 
         # Retrieve the optional eq_params_batch
         # and update eq_params with the latter
@@ -660,7 +661,6 @@ class LossPDENonStatio(LossPDEStatio):
                 _set_derivatives(params, self.derivative_keys.initial_condition),
                 (0,) + vmap_in_axes_params,
                 self.initial_condition_fun,
-                omega_batch.shape[0],
                 self.loss_weights.initial_condition,
             )
         else:
@@ -1016,19 +1016,7 @@ class SystemLossPDE(eqx.Module):
         if self.u_dict.keys() != params_dict.nn_params.keys():
             raise ValueError("u_dict and params_dict[nn_params] should have same keys ")
 
-        if isinstance(batch, PDEStatioBatch):
-            omega_batch, _ = batch.inside_batch, batch.border_batch
-            vmap_in_axes_x_or_x_t = (0,)
-
-            batches = (omega_batch,)
-        elif isinstance(batch, PDENonStatioBatch):
-            times_batch = batch.times_x_inside_batch[:, 0:1]
-            omega_batch = batch.times_x_inside_batch[:, 1:]
-
-            batches = (omega_batch, times_batch)
-            vmap_in_axes_x_or_x_t = (0, 0)
-        else:
-            raise ValueError("Wrong type of batch")
+        vmap_in_axes = (0,)
 
         # Retrieve the optional eq_params_batch
         # and update eq_params with the latter
@@ -1036,7 +1024,6 @@ class SystemLossPDE(eqx.Module):
         if batch.param_batch_dict is not None:
             eq_params_batch_dict = batch.param_batch_dict
 
-            # TODO
             # feed the eq_params with the batch
             for k in eq_params_batch_dict.keys():
                 params_dict.eq_params[k] = eq_params_batch_dict[k]
@@ -1050,9 +1037,13 @@ class SystemLossPDE(eqx.Module):
             return dynamic_loss_apply(
                 dyn_loss.evaluate,
                 self.u_dict,
-                batches,
+                (
+                    batch.domain_batch
+                    if isinstance(batch, PDEStatioBatch)
+                    else batch.domain_batch
+                ),
                 _set_derivatives(params_dict, self.derivative_keys_dyn_loss.dyn_loss),
-                vmap_in_axes_x_or_x_t + vmap_in_axes_params,
+                vmap_in_axes + vmap_in_axes_params,
                 loss_weight,
                 u_type=type(list(self.u_dict.values())[0]),
             )

jinns/loss/__init__.py

@@ -3,7 +3,7 @@ from ._LossODE import LossODE, SystemLossODE
 from ._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
 from ._DynamicLoss import (
     GeneralizedLotkaVolterra,
-    BurgerEquation,
+    BurgersEquation,
     FPENonStatioLoss2D,
     OU_FPENonStatioLoss2D,
     FisherKPP,
@@ -19,9 +19,10 @@ from ._loss_weights import (
 )
 
 from ._operators import (
-    _div_fwd,
-    _div_rev,
-    _laplacian_fwd,
-    _laplacian_rev,
-    _vectorial_laplacian,
+    divergence_fwd,
+    divergence_rev,
+    laplacian_fwd,
+    laplacian_rev,
+    vectorial_laplacian_fwd,
+    vectorial_laplacian_rev,
 )