jinns 1.7.0__py3-none-any.whl → 1.7.1__py3-none-any.whl

jinns/loss/_DynamicLossAbstract.py

@@ -16,6 +16,7 @@ from jaxtyping import Float, Array, PyTree
 import jax
 import jax.numpy as jnp
 from jinns.parameters._params import EqParams
+from jinns.nn import SPINN
 
 
 # See : https://docs.kidger.site/equinox/api/module/advanced_fields/#equinox.AbstractClassVar--known-issues
@@ -38,6 +39,7 @@ def _decorator_heteregeneous_params(evaluate):
             self._eval_heterogeneous_parameters(
                 inputs, u, params, self.eq_params_heterogeneity
             ),
+            is_leaf=lambda x: x is None,
         )
         new_args = args[:-1] + (_params,)
         res = evaluate(*new_args)
@@ -152,7 +154,7 @@ class DynamicLoss(eqx.Module, Generic[InputDim]):
                 "The output of dynamic loss must be vectorial, "
                 "i.e. of shape (d,) with d >= 1"
             )
-        if len(evaluation.shape) > 1:
+        if len(evaluation.shape) > 1 and not isinstance(u, SPINN):
             warnings.warn(
                 "Return value from DynamicLoss' equation has more "
                 "than one dimension. This is in general a mistake (probably from "

jinns/loss/_LossODE.py

@@ -28,13 +28,13 @@ from jinns.parameters._derivative_keys import _set_derivatives, DerivativeKeysOD
 from jinns.loss._loss_weights import LossWeightsODE
 from jinns.loss._abstract_loss import AbstractLoss
 from jinns.loss._loss_components import ODEComponents
+from jinns.loss import ODE
 from jinns.parameters._params import Params
 from jinns.data._Batchs import ODEBatch
 
 if TYPE_CHECKING:
     # imports only used in type hints
     from jinns.nn._abstract_pinn import AbstractPINN
-    from jinns.loss import ODE
 
     InitialConditionUser = (
         tuple[Float[Array, " n_cond "], Float[Array, " n_cond dim"]]
@@ -65,39 +65,42 @@ class LossODE(
     ----------
     u : eqx.Module
         the PINN
-    dynamic_loss : ODE
+    dynamic_loss : tuple[ODE, ...] | ODE | None
         the ODE dynamic part of the loss, basically the differential
         operator $\mathcal{N}[u](t)$. Should implement a method
         `dynamic_loss.evaluate(t, u, params)`.
         Can be None in order to access only some part of the evaluate call.
-    loss_weights : LossWeightsODE, default=None
+    loss_weights : LossWeightsODE | None, default=None
         The loss weights for the differents term : dynamic loss,
         initial condition and eventually observations if any.
         Can be updated according to a specific algorithm. See
         `update_weight_method`
-    update_weight_method : Literal['soft_adapt', 'lr_annealing', 'ReLoBRaLo'], default=None
+    update_weight_method : Literal['soft_adapt', 'lr_annealing', 'ReLoBRaLo'] | None, default=None
         Default is None meaning no update for loss weights. Otherwise a string
-    derivative_keys : DerivativeKeysODE, default=None
+    keep_initial_loss_weight_scales : bool, default=True
+        Only used if an update weight method is specified. It decides whether
+        the updated loss weights are multiplied by the initial `loss_weights`
+        passed by the user at initialization. This is useful to force some
+        scale difference between the adaptative loss weights even after the
+        update method is applied.
+    derivative_keys : DerivativeKeysODE | None, default=None
         Specify which field of `params` should be differentiated for each
         composant of the total loss. Particularily useful for inverse problems.
         Fields can be "nn_params", "eq_params" or "both". Those that should not
         be updated will have a `jax.lax.stop_gradient` called on them. Default
         is `"nn_params"` for each composant of the loss.
-    initial_condition : tuple[
-            Float[Array, "n_cond "],
-            Float[Array, "n_cond dim"]
-        ] |
-        tuple[int | float | Float[Array, " "],
-              int | float | Float[Array, " dim"]
-        ], default=None
+    initial_condition : InitialConditionUser | None, default=None
         Most of the time, a tuple of length 2 with initial condition $(t_0, u_0)$.
         From jinns v1.5.1 we accept tuples of jnp arrays with shape (n_cond, 1) for t0 and (n_cond, dim) for u0. This is useful to include observed conditions at different time points, such as *e.g* final conditions. It was designed to implement $\mathcal{L}^{aux}$ from _Systems biology informed deep learning for inferring parameters and hidden dynamics_, Alireza Yazdani et al., 2020
-    obs_slice : EllipsisType | slice, default=None
+    obs_slice : tuple[EllipsisType | slice, ...] | EllipsisType | slice | None, default=None
         Slice object specifying the begininning/ending
         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[Array]], default=None
+        **Note**: If several observation datasets are passed this arguments need to be set as a
+        tuple of jnp.slice objects with the same length as the number of
+        observation datasets
+    params : InitVar[Params[Array]] | None, default=None
         The main Params object of the problem needed to instanciate the
         DerivativeKeysODE if the latter is not specified.
     Raises
@@ -109,22 +112,26 @@ class LossODE(
     # NOTE static=True only for leaf attributes that are not valid JAX types
     # (ie. jax.Array cannot be static) and that we do not expect to change
     u: AbstractPINN
-    dynamic_loss: ODE | None
+    dynamic_loss: tuple[ODE | None, ...]
+    vmap_in_axes: tuple[int] = eqx.field(static=True)
     derivative_keys: DerivativeKeysODE
     loss_weights: LossWeightsODE
     initial_condition: InitialCondition | None
-    obs_slice: EllipsisType | slice = eqx.field(static=True)
+    obs_slice: tuple[EllipsisType | slice, ...] = eqx.field(static=True)
     params: InitVar[Params[Array] | None]
 
     def __init__(
         self,
         *,
         u: AbstractPINN,
-        dynamic_loss: ODE | None,
+        dynamic_loss: tuple[ODE, ...] | ODE | None,
         loss_weights: LossWeightsODE | None = None,
         derivative_keys: DerivativeKeysODE | None = None,
         initial_condition: InitialConditionUser | None = None,
-        obs_slice: EllipsisType | slice | None = None,
+        obs_slice: tuple[EllipsisType | slice, ...]
+        | EllipsisType
+        | slice
+        | None = None,
         params: Params[Array] | None = None,
         **kwargs: Any,  # this is for arguments for super()
     ):
@@ -141,8 +148,6 @@ class LossODE(
                     f"received {derivative_keys=} and {params=}"
                 )
             derivative_keys = DerivativeKeysODE(params=params)
-        else:
-            derivative_keys = derivative_keys
 
         super().__init__(
             loss_weights=self.loss_weights,
@@ -151,7 +156,10 @@ class LossODE(
             **kwargs,
         )
         self.u = u
-        self.dynamic_loss = dynamic_loss
+        if not isinstance(dynamic_loss, tuple):
+            self.dynamic_loss = (dynamic_loss,)
+        else:
+            self.dynamic_loss = dynamic_loss
         if self.update_weight_method is not None and jnp.any(
             jnp.array(jax.tree.leaves(self.loss_weights)) == 0
         ):
@@ -231,10 +239,15 @@ class LossODE(
             self.initial_condition = (t0, u0)
 
         if obs_slice is None:
-            self.obs_slice = jnp.s_[...]
+            self.obs_slice = (jnp.s_[...],)
+        elif not isinstance(obs_slice, tuple):
+            self.obs_slice = (obs_slice,)
         else:
             self.obs_slice = obs_slice
 
+        if self.loss_weights is None:
+            self.loss_weights = LossWeightsODE()
+
     def evaluate_by_terms(
         self,
         opt_params: Params[Array],
@@ -280,23 +293,29 @@ class LossODE(
             cast(eqx.Module, batch.param_batch_dict), params
         )
 
-        ## dynamic part
-        if self.dynamic_loss is not None:
-            dyn_loss_eval = self.dynamic_loss.evaluate
-            dyn_loss_fun: Callable[[Params[Array]], Array] | None = (
-                lambda p: dynamic_loss_apply(
-                    dyn_loss_eval,
-                    self.u,
-                    temporal_batch,
-                    _set_derivatives(p, self.derivative_keys.dyn_loss),
-                    self.vmap_in_axes + vmap_in_axes_params,
+        if self.dynamic_loss != (None,):
+            # Note, for the record, multiple dynamic losses
+            # have been introduced in MR 92
+            dyn_loss_fun: tuple[Callable[[Params[Array]], Array], ...] | None = (
+                jax.tree.map(
+                    lambda d: lambda p: dynamic_loss_apply(
+                        d.evaluate,
+                        self.u,
+                        temporal_batch,
+                        _set_derivatives(p, self.derivative_keys.dyn_loss),
+                        self.vmap_in_axes + vmap_in_axes_params,
+                    ),
+                    self.dynamic_loss,
+                    is_leaf=lambda x: isinstance(x, ODE),  # do not traverse
+                    # further than first level
                 )
             )
         else:
             dyn_loss_fun = None
 
         if self.initial_condition is not None:
-            # initial condition
+            # Note, for the record, multiple initial conditions for LossODEs
+            # have been introduced in MR 77
             t0, u0 = self.initial_condition
 
             # first construct the plain init loss no vmaping
@@ -331,34 +350,50 @@ class LossODE(
                 # if there is no parameter batch to vmap over we cannot call
                 # vmap because calling vmap must be done with at least one non
                 # None in_axes or out_axes
-                initial_condition_fun = initial_condition_fun_
+                initial_condition_fun = (initial_condition_fun_,)
             else:
-                initial_condition_fun: Callable[[Params[Array]], Array] | None = (
+                initial_condition_fun: (
+                    tuple[Callable[[Params[Array]], Array], ...] | None
+                ) = (
                     lambda p: jnp.mean(
                         vmap(initial_condition_fun_, vmap_in_axes_params)(p)
-                    )
+                    ),
                 )
+            # Note that since MR 92
+            # initial_condition_fun is formed as a tuple for
+            # consistency with dynamic and observation
+            # losses and more modularity for later
         else:
             initial_condition_fun = None
 
         if batch.obs_batch_dict is not None:
-            # update params with the batches of observed params
-            params_obs = update_eq_params(params, batch.obs_batch_dict["eq_params"])
-
-            pinn_in, val = (
-                batch.obs_batch_dict["pinn_in"],
-                batch.obs_batch_dict["val"],
-            )  # the reason for this intruction is https://github.com/microsoft/pyright/discussions/8340
-
-            # MSE loss wrt to an observed batch
-            obs_loss_fun: Callable[[Params[Array]], Array] | None = (
-                lambda po: observations_loss_apply(
-                    self.u,
-                    pinn_in,
-                    _set_derivatives(po, self.derivative_keys.observations),
-                    self.vmap_in_axes + vmap_in_axes_params,
-                    val,
+            # Note, for the record, multiple DGObs
+            # (leading to batch.obs_batch_dict being tuple | None)
+            # have been introduced in MR 92
+            if len(batch.obs_batch_dict) != len(self.obs_slice):
+                raise ValueError(
+                    "There must be the same number of "
+                    "observation datasets as the number of "
+                    "obs_slice"
+                )
+            params_obs = jax.tree.map(
+                lambda d: update_eq_params(params, d["eq_params"]),
+                batch.obs_batch_dict,
+                is_leaf=lambda x: isinstance(x, dict),
+            )
+            obs_loss_fun: tuple[Callable[[Params[Array]], Array], ...] | None = (
+                jax.tree.map(
+                    lambda d, slice_: lambda po: observations_loss_apply(
+                        self.u,
+                        d["pinn_in"],
+                        _set_derivatives(po, self.derivative_keys.observations),
+                        self.vmap_in_axes + vmap_in_axes_params,
+                        d["val"],
+                        slice_,
+                    ),
+                    batch.obs_batch_dict,
                     self.obs_slice,
+                    is_leaf=lambda x: isinstance(x, dict),
                 )
             )
         else:
@@ -366,33 +401,36 @@ class LossODE(
             obs_loss_fun = None
 
         # get the unweighted mses for each loss term as well as the gradients
-        all_funs: ODEComponents[Callable[[Params[Array]], Array] | None] = (
+        all_funs: ODEComponents[tuple[Callable[[Params[Array]], Array], ...] | None] = (
             ODEComponents(dyn_loss_fun, initial_condition_fun, obs_loss_fun)
         )
-        all_params: ODEComponents[Params[Array] | None] = ODEComponents(
-            params, params, params_obs
+        all_params: ODEComponents[tuple[Params[Array], ...] | None] = ODEComponents(
+            jax.tree.map(lambda l: params, dyn_loss_fun),
+            jax.tree.map(lambda l: params, initial_condition_fun),
+            params_obs,
         )
 
-        # Note that the lambda functions below are with type: ignore just
-        # because the lambda are not type annotated, but there is no proper way
-        # to do this and we should assign the lambda to a type hinted variable
-        # before hand: this is not practical, let us not get mad at this
         mses_grads = jax.tree.map(
             self.get_gradients,
             all_funs,
             all_params,
             is_leaf=lambda x: x is None,
         )
-
+        # NOTE the is_leaf below is more complex since it must pass possible the tuple
+        # of dyn_loss and then stops (but also account it should not stop when
+        # the tuple of dyn_loss is of length 2)
         mses = jax.tree.map(
-            lambda leaf: leaf[0],  # type: ignore
+            lambda leaf: leaf[0],
             mses_grads,
-            is_leaf=lambda x: isinstance(x, tuple),
+            is_leaf=lambda x: isinstance(x, tuple)
+            and len(x) == 2
+            and isinstance(x[1], Params),
         )
         grads = jax.tree.map(
-            lambda leaf: leaf[1],  # type: ignore
+            lambda leaf: leaf[1],
             mses_grads,
-            is_leaf=lambda x: isinstance(x, tuple),
+            is_leaf=lambda x: isinstance(x, tuple)
+            and len(x) == 2
+            and isinstance(x[1], Params),
         )
-
         return mses, grads