jinns 1.4.0__py3-none-any.whl → 1.5.1__py3-none-any.whl

jinns/loss/_LossODE.py

@@ -7,7 +7,7 @@ from __future__ import (
 )  # https://docs.python.org/3/library/typing.html#constant
 
 from dataclasses import InitVar
-from typing import TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING, TypedDict, Callable
 from types import EllipsisType
 import abc
 import warnings
@@ -19,6 +19,7 @@ from jaxtyping import Float, Array
 from jinns.loss._loss_utils import (
     dynamic_loss_apply,
     observations_loss_apply,
+    initial_condition_check,
 )
 from jinns.parameters._params import (
     _get_vmap_in_axes_params,
@@ -27,10 +28,11 @@ from jinns.parameters._params import (
 from jinns.parameters._derivative_keys import _set_derivatives, DerivativeKeysODE
 from jinns.loss._loss_weights import LossWeightsODE
 from jinns.loss._abstract_loss import AbstractLoss
+from jinns.loss._loss_components import ODEComponents
+from jinns.parameters._params import Params
 
 if TYPE_CHECKING:
     # imports only used in type hints
-    from jinns.parameters._params import Params
     from jinns.data._Batchs import ODEBatch
     from jinns.nn._abstract_pinn import AbstractPINN
     from jinns.loss import ODE
@@ -42,22 +44,32 @@ if TYPE_CHECKING:
 
 
 class _LossODEAbstract(AbstractLoss):
-    """
+    r"""
     Parameters
     ----------
 
     loss_weights : LossWeightsODE, default=None
         The loss weights for the differents term : dynamic loss,
-        initial condition and eventually observations if any. All fields are
-        set to 1.0 by default.
+        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
+        Default is None meaning no update for loss weights. Otherwise a string
     derivative_keys : DerivativeKeysODE, 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 | Float[Array, " 1"], Float[Array, " dim"]], default=None
-        tuple of length 2 with initial condition $(t_0, u_0)$.
+    initial_condition : tuple[
+            Float[Array, "n_cond "],
+            Float[Array, "n_cond dim"]
+        ] |
+        tuple[int | float | Float[Array, " "],
+              int | float | Float[Array, " dim"]
+        ] | 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 | None, default=None
         Slice object specifying the begininning/ending
         slice of u output(s) that is observed. This is useful for
@@ -74,7 +86,9 @@ class _LossODEAbstract(AbstractLoss):
     derivative_keys: DerivativeKeysODE | None = eqx.field(kw_only=True, default=None)
     loss_weights: LossWeightsODE | None = eqx.field(kw_only=True, default=None)
     initial_condition: (
-        tuple[float | Float[Array, " 1"], Float[Array, " dim"]] | None
+        tuple[Float[Array, " n_cond 1"], Float[Array, " n_cond dim"]]
+        | tuple[int | float | Float[Array, " "], int | float | Float[Array, " dim"]]
+        | None
     ) = eqx.field(kw_only=True, default=None)
     obs_slice: EllipsisType | slice | None = eqx.field(
         kw_only=True, default=None, static=True
@@ -108,18 +122,60 @@ class _LossODEAbstract(AbstractLoss):
                     "Initial condition should be a tuple of len 2 with (t0, u0), "
                     f"{self.initial_condition} was passed."
                 )
-            # some checks/reshaping for t0
+            # some checks/reshaping for t0 and u0
             t0, u0 = self.initial_condition
             if isinstance(t0, Array):
-                if not t0.shape:  # e.g. user input: jnp.array(0.)
-                    t0 = jnp.array([t0])
-                elif t0.shape != (1,):
+                # at the end we want to end up with t0 of shape (:, 1) to account for
+                # possibly several data points
+                if t0.ndim <= 1:
+                    # in this case we assume t0 belongs one (initial)
+                    # condition
+                    t0 = initial_condition_check(t0, dim_size=1)[
+                        None, :
+                    ]  # make a (1, 1) here
+                if t0.ndim > 2:
+                    raise ValueError(
+                        "It t0 is an Array, it represents n_cond"
+                        " imposed conditions and must be of shape (n_cond, 1)"
+                    )
+            else:
+                # in this case t0 clearly represents one (initial) condition
+                t0 = initial_condition_check(t0, dim_size=1)[
+                    None, :
+                ]  # make a (1, 1) here
+            if isinstance(u0, Array):
+                # at the end we want to end up with u0 of shape (:, dim) to account for
+                # possibly several data points
+                if not u0.shape:
+                    # in this case we assume u0 belongs to one (initial)
+                    # condition
+                    u0 = initial_condition_check(u0, dim_size=1)[
+                        None, :
+                    ]  # make a (1, 1) here
+                elif u0.ndim == 1:
+                    # in this case we assume u0 belongs to one (initial)
+                    # condition
+                    u0 = initial_condition_check(u0, dim_size=u0.shape[0])[
+                        None, :
+                    ]  # make a (1, dim) here
+                if u0.ndim > 2:
                     raise ValueError(
-                        f"Wrong t0 input (self.initial_condition[0]) It should be"
-                        f"a float or an array of shape (1,). Got shape: {t0.shape}"
+                        "It u0 is an Array, it represents n_cond "
+                        "imposed conditions and must be of shape (n_cond, dim)"
                     )
-            if isinstance(t0, float):  # e.g. user input: 0
-                t0 = jnp.array([t0])
+            else:
+                # at the end we want to end up with u0 of shape (:, dim) to account for
+                # possibly several data points
+                u0 = initial_condition_check(u0, dim_size=None)[
+                    None, :
+                ]  # make a (1, 1) here
+
+            if t0.shape[0] != u0.shape[0] or t0.ndim != u0.ndim:
+                raise ValueError(
+                    "t0 and u0 must represent a same number of initial"
+                    " conditial conditions"
+                )
+
             self.initial_condition = (t0, u0)
 
         if self.obs_slice is None:
@@ -154,8 +210,11 @@ class LossODE(_LossODEAbstract):
     ----------
     loss_weights : LossWeightsODE, default=None
         The loss weights for the differents term : dynamic loss,
-        initial condition and eventually observations if any. All fields are
-        set to 1.0 by default.
+        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
+        Default is None meaning no update for loss weights. Otherwise a string
     derivative_keys : DerivativeKeysODE, default=None
         Specify which field of `params` should be differentiated for each
         composant of the total loss. Particularily useful for inverse problems.
@@ -204,21 +263,26 @@ class LossODE(_LossODEAbstract):
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
 
-    def evaluate(
+    def evaluate_by_terms(
         self, params: Params[Array], batch: ODEBatch
-    ) -> tuple[Float[Array, " "], LossDictODE]:
+    ) -> tuple[
+        ODEComponents[Float[Array, " "] | None], ODEComponents[Float[Array, " "] | None]
+    ]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
+        We retrieve two PyTrees with loss values and gradients for each term
 
         Parameters
         ---------
         params
             Parameters at which the loss is evaluated
         batch
-            Composed of a batch of time points
-            at which to evaluate the differential operator. An optional additional batch of parameters (eg. for metamodeling) and an optional additional batch of observed inputs/outputs/parameters can
-            be supplied.
+            Composed of a batch of points in the
+            domain, a batch of points in the domain
+            border and an optional additional batch of parameters (eg. for
+            metamodeling) and an optional additional batch of observed
+            inputs/outputs/parameters
         """
         temporal_batch = batch.temporal_batch
 
@@ -233,71 +297,120 @@ class LossODE(_LossODEAbstract):
 
         ## dynamic part
         if self.dynamic_loss is not None:
-            mse_dyn_loss = dynamic_loss_apply(
-                self.dynamic_loss.evaluate,
+            dyn_loss_fun = lambda p: dynamic_loss_apply(
+                self.dynamic_loss.evaluate,  # type: ignore
                 self.u,
                 temporal_batch,
-                _set_derivatives(params, self.derivative_keys.dyn_loss),  # type: ignore
+                _set_derivatives(p, self.derivative_keys.dyn_loss),  # type: ignore
                 self.vmap_in_axes + vmap_in_axes_params,
-                self.loss_weights.dyn_loss,  # type: ignore
             )
         else:
-            mse_dyn_loss = jnp.array(0.0)
+            dyn_loss_fun = None
 
         # initial condition
         if self.initial_condition is not None:
-            vmap_in_axes = (None,) + vmap_in_axes_params
-            if not jax.tree_util.tree_leaves(vmap_in_axes):
-                # test if only None in vmap_in_axes to avoid the value error:
-                # `vmap must have at least one non-None value in in_axes`
-                v_u = self.u
-            else:
-                v_u = vmap(self.u, (None,) + vmap_in_axes_params)
             t0, u0 = self.initial_condition
             u0 = jnp.array(u0)
-            mse_initial_condition = jnp.mean(
-                self.loss_weights.initial_condition  # type: ignore
-                * jnp.sum(
-                    (
-                        v_u(
-                            t0,
-                            _set_derivatives(
-                                params,
-                                self.derivative_keys.initial_condition,  # type: ignore
-                            ),
-                        )
-                        - u0
+
+            # first construct the plain init loss no vmaping
+            initial_condition_fun__ = lambda t, u, p: jnp.sum(
+                (
+                    self.u(
+                        t,
+                        _set_derivatives(
+                            p,
+                            self.derivative_keys.initial_condition,  # type: ignore
+                        ),
                     )
-                    ** 2,
-                    axis=-1,
+                    - u
                 )
+                ** 2,
+                axis=0,
+            )
+            # now vmap over the number of conditions (first dim of t0 and u0)
+            # and take the mean
+            initial_condition_fun_ = lambda p: jnp.mean(
+                vmap(initial_condition_fun__, (0, 0, None))(t0, u0, p)
             )
+            # now vmap over the the possible batch of parameters and take the
+            # average. Note that we then finally have a cartesian product
+            # between the batch of parameters (if any) and the number of
+            # conditions (if any)
+            if not jax.tree_util.tree_leaves(vmap_in_axes_params):
+                # 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_
+            else:
+                initial_condition_fun = lambda p: jnp.mean(
+                    vmap(initial_condition_fun_, vmap_in_axes_params)(p)
+                )
         else:
-            mse_initial_condition = jnp.array(0.0)
+            initial_condition_fun = None
 
         if batch.obs_batch_dict is not None:
             # update params with the batches of observed params
-            params = _update_eq_params_dict(params, batch.obs_batch_dict["eq_params"])
+            params_obs = _update_eq_params_dict(
+                params, batch.obs_batch_dict["eq_params"]
+            )
 
             # MSE loss wrt to an observed batch
-            mse_observation_loss = observations_loss_apply(
+            obs_loss_fun = lambda po: observations_loss_apply(
                 self.u,
                 batch.obs_batch_dict["pinn_in"],
-                _set_derivatives(params, self.derivative_keys.observations),  # type: ignore
+                _set_derivatives(po, self.derivative_keys.observations),  # type: ignore
                 self.vmap_in_axes + vmap_in_axes_params,
                 batch.obs_batch_dict["val"],
-                self.loss_weights.observations,  # type: ignore
                 self.obs_slice,
             )
         else:
-            mse_observation_loss = jnp.array(0.0)
-
-        # total loss
-        total_loss = mse_dyn_loss + mse_initial_condition + mse_observation_loss
-        return total_loss, (
-            {
-                "dyn_loss": mse_dyn_loss,
-                "initial_condition": mse_initial_condition,
-                "observations": mse_observation_loss,
-            }
+            params_obs = None
+            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] = (
+            ODEComponents(dyn_loss_fun, initial_condition_fun, obs_loss_fun)
+        )
+        all_params: ODEComponents[Params[Array] | None] = ODEComponents(
+            params, params, params_obs
         )
+        mses_grads = jax.tree.map(
+            lambda fun, params: self.get_gradients(fun, params),
+            all_funs,
+            all_params,
+            is_leaf=lambda x: x is None,
+        )
+
+        mses = jax.tree.map(
+            lambda leaf: leaf[0], mses_grads, is_leaf=lambda x: isinstance(x, tuple)
+        )
+        grads = jax.tree.map(
+            lambda leaf: leaf[1], mses_grads, is_leaf=lambda x: isinstance(x, tuple)
+        )
+
+        return mses, grads
+
+    def evaluate(
+        self, params: Params[Array], batch: ODEBatch
+    ) -> tuple[Float[Array, " "], ODEComponents[Float[Array, " "] | None]]:
+        """
+        Evaluate the loss function at a batch of points for given parameters.
+
+        We retrieve the total value itself and a PyTree with loss values for each term
+
+        Parameters
+        ---------
+        params
+            Parameters at which the loss is evaluated
+        batch
+            Composed of a batch of points in the
+            domain, a batch of points in the domain
+            border and an optional additional batch of parameters (eg. for
+            metamodeling) and an optional additional batch of observed
+            inputs/outputs/parameters
+        """
+        loss_terms, _ = self.evaluate_by_terms(params, batch)
+
+        loss_val = self.ponderate_and_sum_loss(loss_terms)
+
+        return loss_val, loss_terms