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

jinns/loss/_DynamicLossAbstract.py

@@ -21,61 +21,32 @@ 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):
     r"""
     Abstract base class for dynamic losses. Implements the physical term:
+
     $$
         \mathcal{N}[u](t, x) = 0
     $$
+
     for **one** point $t$, $x$ or $(t, x)$, depending on the context.
 
     Parameters
@@ -108,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,
@@ -122,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
@@ -138,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):
@@ -189,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"],
@@ -197,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(
@@ -258,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(
@@ -323,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:
@@ -351,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

@@ -56,6 +56,9 @@ class _LossODEAbstract(eqx.Module):
         slice of u output(s) that is observed. This is useful for
         multidimensional PINN, with partially observed outputs.
         Default is None (whole output is observed).
+    params : InitVar[Params], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
     """
 
     # NOTE static=True only for leaf attributes that are not valid JAX types
@@ -66,13 +69,21 @@ class _LossODEAbstract(eqx.Module):
     initial_condition: tuple | None = eqx.field(kw_only=True, default=None)
     obs_slice: slice | None = eqx.field(kw_only=True, default=None, static=True)
 
-    def __post_init__(self):
+    params: InitVar[Params] = eqx.field(default=None, kw_only=True)
+
+    def __post_init__(self, params=None):
         if self.loss_weights is None:
             self.loss_weights = LossWeightsODE()
 
         if self.derivative_keys is None:
-            # be default we only take gradient wrt nn_params
-            self.derivative_keys = DerivativeKeysODE()
+            try:
+                # be default we only take gradient wrt nn_params
+                self.derivative_keys = DerivativeKeysODE(params=params)
+            except ValueError as exc:
+                raise ValueError(
+                    "Problem at self.derivative_keys initialization "
+                    f"received {self.derivative_keys=} and {params=}"
+                ) from exc
         if self.initial_condition is None:
             warnings.warn(
                 "Initial condition wasn't provided. Be sure to cover for that"
@@ -131,6 +142,9 @@ class LossODE(_LossODEAbstract):
         slice of u output(s) that is observed. This is useful for
         multidimensional PINN, with partially observed outputs.
         Default is None (whole output is observed).
+    params : InitVar[Params], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
     u : eqx.Module
         the PINN
     dynamic_loss : DynamicLoss
@@ -152,8 +166,10 @@ class LossODE(_LossODEAbstract):
 
     vmap_in_axes: tuple[Int] = eqx.field(init=False, static=True)
 
-    def __post_init__(self):
-        super().__post_init__()  # because __init__ or __post_init__ of Base
+    def __post_init__(self, params=None):
+        super().__post_init__(
+            params=params
+        )  # because __init__ or __post_init__ of Base
         # class is not automatically called
 
         self.vmap_in_axes = (0,)
@@ -193,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,
@@ -211,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
@@ -300,6 +316,9 @@ class SystemLossODE(eqx.Module):
         PINNs. Default is None. But if a value is given, all the entries of
         `u_dict` must be represented here with default value `jnp.s_[...]`
         if no particular slice is to be given.
+    params_dict : InitVar[ParamsDict], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
 
     Raises
     ------
@@ -332,14 +351,16 @@ class SystemLossODE(eqx.Module):
     loss_weights: InitVar[LossWeightsODEDict | None] = eqx.field(
         kw_only=True, default=None
     )
+    params_dict: InitVar[ParamsDict] = eqx.field(kw_only=True, default=None)
+
     u_constraints_dict: Dict[str, LossODE] = eqx.field(init=False)
-    derivative_keys_dyn_loss_dict: Dict[str, DerivativeKeysODE] = eqx.field(init=False)
+    derivative_keys_dyn_loss: DerivativeKeysODE = eqx.field(init=False)
 
     u_dict_with_none: Dict[str, None] = eqx.field(init=False)
     # internally the loss weights are handled with a dictionary
     _loss_weights: Dict[str, dict] = eqx.field(init=False)
 
-    def __post_init__(self, loss_weights):
+    def __post_init__(self, loss_weights=None, params_dict=None):
         # a dictionary that will be useful at different places
         self.u_dict_with_none = {k: None for k in self.u_dict.keys()}
         if self.initial_condition_dict is None:
@@ -369,14 +390,14 @@ class SystemLossODE(eqx.Module):
             # iterating on dynamic_loss_dict. So each time we will require dome
             # derivative_keys_dict
 
-        # but then if the user did not provide anything, we must at least have
-        # a default value for the dynamic_loss_dict keys entries in
-        # self.derivative_keys_dict since the computation of dynamic losses is
-        # made without create a lossODE object that would provide the
-        # default values
-        for k in self.dynamic_loss_dict.keys():
+        # derivative keys for the u_constraints. Note that we create missing
+        # DerivativeKeysODE around a Params object and not ParamsDict
+        # this works because u_dict.keys == params_dict.nn_params.keys()
+        for k in self.u_dict.keys():
             if self.derivative_keys_dict[k] is None:
-                self.derivative_keys_dict[k] = DerivativeKeysODE()
+                self.derivative_keys_dict[k] = DerivativeKeysODE(
+                    params=params_dict.extract_params(k)
+                )
 
         self._loss_weights = self.set_loss_weights(loss_weights)
 
@@ -397,12 +418,11 @@ class SystemLossODE(eqx.Module):
                 obs_slice=self.obs_slice_dict[i],
             )
 
-        # for convenience in the tree_map of evaluate
-        self.derivative_keys_dyn_loss_dict = {
-            k: self.derivative_keys_dict[k]
-            for k in self.dynamic_loss_dict.keys()  # & self.derivative_keys_dict.keys()
-            # comment because intersection is neceserily fulfilled right?
-        }
+        # derivative keys for the dynamic loss. Note that we create a
+        # DerivativeKeysODE around a ParamsDict object because a whole
+        # params_dict is feed to DynamicLoss.evaluate functions (extract_params
+        # happen inside it)
+        self.derivative_keys_dyn_loss = DerivativeKeysODE(params=params_dict)
 
     def set_loss_weights(self, loss_weights_init):
         """
@@ -497,13 +517,13 @@ class SystemLossODE(eqx.Module):
             batch.param_batch_dict, params_dict
         )
 
-        def dyn_loss_for_one_key(dyn_loss, derivative_key, loss_weight):
+        def dyn_loss_for_one_key(dyn_loss, loss_weight):
             """This function is used in tree_map"""
             return dynamic_loss_apply(
                 dyn_loss.evaluate,
                 self.u_dict,
-                (temporal_batch,),
-                _set_derivatives(params_dict, derivative_key.dyn_loss),
+                temporal_batch,
+                _set_derivatives(params_dict, self.derivative_keys_dyn_loss.dyn_loss),
                 vmap_in_axes_t + vmap_in_axes_params,
                 loss_weight,
                 u_type=PINN,
@@ -512,7 +532,6 @@ class SystemLossODE(eqx.Module):
         dyn_loss_mse_dict = jax.tree_util.tree_map(
             dyn_loss_for_one_key,
             self.dynamic_loss_dict,
-            self.derivative_keys_dyn_loss_dict,
             self._loss_weights["dyn_loss"],
             is_leaf=lambda x: isinstance(x, ODE),  # before when dynamic losses
             # where plain (unregister pytree) node classes, we could not traverse