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

jinns/loss/_LossODE.py

@@ -1,143 +1,185 @@
+# pylint: disable=unsubscriptable-object, no-member
 """
 Main module to implement a ODE loss in jinns
 """
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
 
+from dataclasses import InitVar, fields
+from typing import TYPE_CHECKING, Dict
+import abc
 import warnings
 import jax
 import jax.numpy as jnp
 from jax import vmap
-from jax.tree_util import register_pytree_node_class
-from jinns.utils._utils import (
-    _get_vmap_in_axes_params,
-    _set_derivatives,
-    _update_eq_params_dict,
-)
-from jinns.loss._Losses import (
+import equinox as eqx
+from jaxtyping import Float, Array, Int
+from jinns.data._DataGenerators import append_obs_batch
+from jinns.loss._loss_utils import (
     dynamic_loss_apply,
     constraints_system_loss_apply,
     observations_loss_apply,
 )
+from jinns.parameters._params import (
+    _get_vmap_in_axes_params,
+    _update_eq_params_dict,
+)
+from jinns.parameters._derivative_keys import _set_derivatives, DerivativeKeysODE
+from jinns.loss._loss_weights import LossWeightsODE, LossWeightsODEDict
+from jinns.loss._DynamicLossAbstract import ODE
 from jinns.utils._pinn import PINN
 
-_LOSS_WEIGHT_KEYS_ODE = ["observations", "dyn_loss", "initial_condition"]
-
-
-@register_pytree_node_class
-class LossODE:
-    r"""Loss object for an ordinary differential equation
+if TYPE_CHECKING:
+    from jinns.utils._types import *
 
-    .. math::
-        \mathcal{N}[u](t) = 0, \forall t \in I
 
-    where :math:`\mathcal{N}[\cdot]` is a differential operator and the
-    initial condition is :math:`u(t_0)=u_0`.
-
-
-    **Note:** LossODE is jittable. Hence it implements the tree_flatten() and
-    tree_unflatten methods.
+class _LossODEAbstract(eqx.Module):
+    """
+    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.
+    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, default=None
+        tuple of length 2 with initial condition $(t_0, u_0)$.
+    obs_slice : Slice, 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], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
     """
 
-    def __init__(
-        self,
-        u,
-        loss_weights,
-        dynamic_loss,
-        derivative_keys=None,
-        initial_condition=None,
-        obs_slice=None,
-    ):
-        r"""
-        Parameters
-        ----------
-        u :
-            the PINN
-        loss_weights :
-            a dictionary with values used to ponderate each term in the loss
-            function. Valid keys are `dyn_loss`, `initial_condition` and `observations`
-            Note that we can have jnp.arrays with the same dimension of
-            `u` which then ponderates each output of `u`
-        dynamic_loss :
-            the ODE dynamic part of the loss, basically the differential
-            operator :math:`\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 results.
-        derivative_keys
-            A dict of lists of strings. In the dict, the key must correspond to
-            the loss term keywords. Then each of the values must correspond to keys in the parameter
-            dictionary (*at top level only of the parameter dictionary*).
-            It enables selecting the set of parameters
-            with respect to which the gradients of the dynamic
-            loss are computed. If nothing is provided, we set ["nn_params"] for all loss term
-            keywords, this is what is typically
-            done in solving forward problems, when we only estimate the
-            equation solution with a PINN. If some loss terms keywords are
-            missing we set their value to ["nn_params"] by default for the same
-            reason
-        initial_condition :
-            tuple of length 2 with initial condition :math:`(t0, u0)`.
-            Can be None in order to
-            access only some part of the evaluate call results.
-        obs_slice:
-            slice object specifying the begininning/ending
-            slice of u output(s) that is observed (this is then useful for
-            multidim PINN). Default is None.
-
-        Raises
-        ------
-        ValueError
-            if initial condition is not a tuple.
-        """
-        self.dynamic_loss = dynamic_loss
-        self.u = u
-        if derivative_keys is None:
-            # be default we only take gradient wrt nn_params
-            derivative_keys = {
-                k: ["nn_params"]
-                for k in [
-                    "dyn_loss",
-                    "initial_condition",
-                    "observations",
-                ]
-            }
-        if isinstance(derivative_keys, list):
-            # if the user only provided a list, this defines the gradient taken
-            # for all the loss entries
-            derivative_keys = {
-                k: derivative_keys
-                for k in [
-                    "dyn_loss",
-                    "initial_condition",
-                    "observations",
-                ]
-            }
-
-        self.derivative_keys = derivative_keys
-
-        if initial_condition is None:
+    # 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
+    # kw_only in base class is motivated here: https://stackoverflow.com/a/69822584
+    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 | None = eqx.field(kw_only=True, default=None)
+    obs_slice: slice | None = eqx.field(kw_only=True, default=None, static=True)
+
+    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:
+            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"
                 "case (e.g by. hardcoding it into the PINN output)."
             )
         else:
-            if not isinstance(initial_condition, tuple) or len(initial_condition) != 2:
+            if (
+                not isinstance(self.initial_condition, tuple)
+                or len(self.initial_condition) != 2
+            ):
                 raise ValueError(
-                    f"Initial condition should be a tuple of len 2 with (t0, u0), {initial_condition} was passed."
+                    "Initial condition should be a tuple of len 2 with (t0, u0), "
+                    f"{self.initial_condition} was passed."
                 )
-        self.initial_condition = initial_condition
-        self.loss_weights = loss_weights
-        self.obs_slice = obs_slice
+
         if self.obs_slice is None:
             self.obs_slice = jnp.s_[...]
 
-        for k in _LOSS_WEIGHT_KEYS_ODE:
-            if k not in self.loss_weights.keys():
-                self.loss_weights[k] = 0
+        if self.loss_weights is None:
+            self.loss_weights = LossWeightsODE()
+
+    @abc.abstractmethod
+    def evaluate(
+        self: eqx.Module, params: Params, batch: ODEBatch
+    ) -> tuple[Float, dict]:
+        raise NotImplementedError
+
+
+class LossODE(_LossODEAbstract):
+    r"""Loss object for an ordinary differential equation
+
+    $$
+    \mathcal{N}[u](t) = 0, \forall t \in I
+    $$
+
+    where $\mathcal{N}[\cdot]$ is a differential operator and the
+    initial condition is $u(t_0)=u_0$.
+
+
+    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.
+    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, default=None
+        tuple of length 2 with initial condition $(t_0, u_0)$.
+    obs_slice Slice, 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], 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
+        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.
+
+    Raises
+    ------
+    ValueError
+        if initial condition is not a tuple.
+    """
+
+    # 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: eqx.Module
+    dynamic_loss: DynamicLoss | None
+
+    vmap_in_axes: tuple[Int] = eqx.field(init=False, static=True)
+
+    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,)
 
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
 
-    def evaluate(self, params, batch):
+    def evaluate(
+        self, params: Params, batch: ODEBatch
+    ) -> tuple[Float[Array, "1"], dict[str, float]]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
@@ -145,21 +187,14 @@ class LossODE:
         Parameters
         ---------
         params
-            The dictionary of parameters of the model.
-            Typically, it is a dictionary of
-            dictionaries: `eq_params` and `nn_params``, respectively the
-            differential equation parameters and the neural network parameter
+            Parameters at which the loss is evaluated
         batch
-            A ODEBatch object.
-            Such a named tuple is composed of a batch of time points
-            at which to evaluate an optional additional batch of parameters (eg. for
-            metamodeling) and an optional additional batch of observed
-            inputs/outputs/parameters
+            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.
         """
         temporal_batch = batch.temporal_batch
 
-        vmap_in_axes_t = (0,)
-
         # Retrieve the optional eq_params_batch
         # and update eq_params with the latter
         # and update vmap_in_axes
@@ -170,21 +205,19 @@ class LossODE:
         vmap_in_axes_params = _get_vmap_in_axes_params(batch.param_batch_dict, params)
 
         ## dynamic part
-        params_ = _set_derivatives(params, "dyn_loss", self.derivative_keys)
         if self.dynamic_loss is not None:
             mse_dyn_loss = dynamic_loss_apply(
                 self.dynamic_loss.evaluate,
                 self.u,
                 (temporal_batch,),
-                params_,
-                vmap_in_axes_t + vmap_in_axes_params,
-                self.loss_weights["dyn_loss"],
+                _set_derivatives(params, self.derivative_keys.dyn_loss),
+                self.vmap_in_axes + vmap_in_axes_params,
+                self.loss_weights.dyn_loss,
             )
         else:
             mse_dyn_loss = jnp.array(0.0)
 
         # initial condition
-        params_ = _set_derivatives(params, "initial_condition", self.derivative_keys)
         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):
@@ -193,12 +226,24 @@ class LossODE:
                 v_u = self.u
             else:
                 v_u = vmap(self.u, (None,) + vmap_in_axes_params)
-            t0, u0 = self.initial_condition
+            t0, u0 = self.initial_condition  # pylint: disable=unpacking-non-sequence
             t0 = jnp.array(t0)
             u0 = jnp.array(u0)
             mse_initial_condition = jnp.mean(
-                self.loss_weights["initial_condition"]
-                * jnp.sum((v_u(t0, params_) - u0) ** 2, axis=-1)
+                self.loss_weights.initial_condition
+                * jnp.sum(
+                    (
+                        v_u(
+                            t0,
+                            _set_derivatives(
+                                params, self.derivative_keys.initial_condition
+                            ),
+                        )
+                        - u0
+                    )
+                    ** 2,
+                    axis=-1,
+                )
             )
         else:
             mse_initial_condition = jnp.array(0.0)
@@ -208,14 +253,13 @@ class LossODE:
             params = _update_eq_params_dict(params, batch.obs_batch_dict["eq_params"])
 
             # MSE loss wrt to an observed batch
-            params_ = _set_derivatives(params, "observations", self.derivative_keys)
             mse_observation_loss = observations_loss_apply(
                 self.u,
                 (batch.obs_batch_dict["pinn_in"],),
-                params_,
-                vmap_in_axes_t + vmap_in_axes_params,
+                _set_derivatives(params, self.derivative_keys.observations),
+                self.vmap_in_axes + vmap_in_axes_params,
                 batch.obs_batch_dict["val"],
-                self.loss_weights["observations"],
+                self.loss_weights.observations,
                 self.obs_slice,
             )
         else:
@@ -231,193 +275,178 @@ class LossODE:
             }
         )
 
-    def tree_flatten(self):
-        children = (self.initial_condition, self.loss_weights)
-        aux_data = {
-            "u": self.u,
-            "dynamic_loss": self.dynamic_loss,
-            "obs_slice": self.obs_slice,
-            "derivative_keys": self.derivative_keys,
-        }
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        (initial_condition, loss_weights) = children
-        loss_ode = cls(
-            loss_weights=loss_weights,
-            initial_condition=initial_condition,
-            **aux_data,
-        )
-        return loss_ode
 
-
-@register_pytree_node_class
-class SystemLossODE:
-    """
+class SystemLossODE(eqx.Module):
+    r"""
     Class to implement a system of ODEs.
     The goal is to give maximum freedom to the user. The class is created with
-    a dict of dynamic loss and a dict of initial conditions. When then iterate
-    over the dynamic losses that compose the system. All the PINNs with all the
-    parameter dictionaries are passed as arguments to each dynamic loss
-    evaluate functions; it is inside the dynamic loss that specification are
-    performed.
+    a dict of dynamic loss and a dict of initial conditions. Then, it iterates
+    over the dynamic losses that compose the system. All PINNs are passed as
+    arguments to each dynamic loss evaluate functions, along with all the
+    parameter dictionaries. All specification is left to the responsability
+    of the user, inside the dynamic loss.
 
     **Note:** All the dictionaries (except `dynamic_loss_dict`) must have the same keys.
     Indeed, these dictionaries (except `dynamic_loss_dict`) are tied to one
     solution.
 
-    **Note:** SystemLossODE is jittable. Hence it implements the tree_flatten() and
-    tree_unflatten methods.
+    Parameters
+    ----------
+    u_dict : Dict[str, eqx.Module]
+        dict of PINNs
+    loss_weights : LossWeightsODEDict
+        A dictionary of LossWeightsODE
+    derivative_keys_dict : Dict[str, DerivativeKeysODE], default=None
+        A dictionnary of DerivativeKeysODE specifying what field of `params`
+        should be used during gradient computations for each of the terms of
+        the total loss, for each of the loss in the system. Default is
+        `"nn_params`" everywhere.
+    initial_condition_dict : Dict[str, tuple], default=None
+        dict of tuple of length 2 with initial condition $(t_0, u_0)$
+        Must share the keys of `u_dict`. Default is None. No initial
+        condition is permitted when the initial condition is hardcoded in
+        the PINN architecture for example
+    dynamic_loss_dict : Dict[str, ODE]
+        dict of dynamic part of the loss, basically the differential
+        operator $\mathcal{N}[u](t)$. Should implement a method
+        `dynamic_loss.evaluate(t, u, params)`
+    obs_slice_dict : Dict[str, Slice]
+        dict of obs_slice, with keys from `u_dict` to designate the
+        output(s) channels that are observed, for each
+        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
+    ------
+    ValueError
+        if initial condition is not a dict of tuple.
+    ValueError
+        if the dictionaries that should share the keys of u_dict do not.
     """
 
-    def __init__(
-        self,
-        u_dict,
-        loss_weights,
-        dynamic_loss_dict,
-        derivative_keys_dict=None,
-        initial_condition_dict=None,
-        obs_slice_dict=None,
-    ):
-        r"""
-        Parameters
-        ----------
-        u_dict
-            dict of PINNs
-        loss_weights
-            A dictionary of dictionaries with values used to
-            ponderate each term in the loss
-            function. Valid keys in the first dictionary are `dyn_loss`,
-            `initial_condition` and `observations`. The keys of the nested
-            dictionaries must share the keys of `u_dict`. Note that the values
-            at the leaf level can have jnp.arrays with the same dimension of
-            `u` which then ponderates each output of `u`
-        derivative_keys_dict
-            A dict of derivative keys as defined in LossODE. The key of this
-            dict must be that of `dynamic_loss_dict` at least and specify how
-            to compute gradient for the `dyn_loss` loss term at least (see the
-            check at the beginning of the present `__init__` function.
-            Other keys of this dict might be that of `u_dict` to specify how to
-            compute gradients for all the different constraints. If those keys
-            are not specified then the default behaviour for `derivative_keys`
-            of LossODE is used
-        initial_condition_dict
-            dict of tuple of length 2 with initial condition :math:`(t_0, u_0)`
-            Must share the keys of `u_dict`. Default is None. No initial
-            condition is permitted when the initial condition is hardcoded in
-            the PINN architecture for example
-        dynamic_loss_dict
-            dict of dynamic part of the loss, basically the differential
-            operator :math:`\mathcal{N}[u](t)`. Should implement a method
-            `dynamic_loss.evaluate(t, u, params)`
-        obs_slice_dict
-            dict of obs_slice, with keys from `u_dict` to designate the
-            output(s) channels that are forced to observed values, for each
-            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
-
-        Raises
-        ------
-        ValueError
-            if initial condition is not a dict of tuple
-        ValueError
-            if the dictionaries that should share the keys of u_dict do not
-        """
-
+    # 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_dict: Dict[str, eqx.Module]
+    dynamic_loss_dict: Dict[str, ODE]
+    derivative_keys_dict: Dict[str, DerivativeKeysODE | None] | None = eqx.field(
+        kw_only=True, default=None
+    )
+    initial_condition_dict: Dict[str, tuple] | None = eqx.field(
+        kw_only=True, default=None
+    )
+
+    obs_slice_dict: Dict[str, slice | None] | None = eqx.field(
+        kw_only=True, default=None, static=True
+    )  # We are at an "leaf" attribute here (slice, not valid JAX type). Since
+    # we do not expect it to change with put a static=True here. But note that
+    # this is the only static for all the SystemLossODE attribute, since all
+    # other are composed of more complex structures ("non-leaf")
+
+    # For the user loss_weights are passed as a LossWeightsODEDict (with internal
+    # dictionary having keys in u_dict and / or dynamic_loss_dict)
+    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: 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=None, params_dict=None):
         # a dictionary that will be useful at different places
-        self.u_dict_with_none = {k: None for k in u_dict.keys()}
-        if initial_condition_dict is None:
+        self.u_dict_with_none = {k: None for k in self.u_dict.keys()}
+        if self.initial_condition_dict is None:
             self.initial_condition_dict = self.u_dict_with_none
         else:
-            self.initial_condition_dict = initial_condition_dict
-            if u_dict.keys() != initial_condition_dict.keys():
+            if self.u_dict.keys() != self.initial_condition_dict.keys():
                 raise ValueError(
                     "initial_condition_dict should have same keys as u_dict"
                 )
-        if obs_slice_dict is None:
-            self.obs_slice_dict = {k: jnp.s_[...] for k in u_dict.keys()}
+        if self.obs_slice_dict is None:
+            self.obs_slice_dict = {k: jnp.s_[...] for k in self.u_dict.keys()}
         else:
-            self.obs_slice_dict = obs_slice_dict
-            if u_dict.keys() != obs_slice_dict.keys():
+            if self.u_dict.keys() != self.obs_slice_dict.keys():
                 raise ValueError("obs_slice_dict should have same keys as u_dict")
 
-        if derivative_keys_dict is None:
+        if self.derivative_keys_dict is None:
             self.derivative_keys_dict = {
                 k: None
-                for k in set(list(dynamic_loss_dict.keys()) + list(u_dict.keys()))
+                for k in set(
+                    list(self.dynamic_loss_dict.keys()) + list(self.u_dict.keys())
+                )
             }
             # set() because we can have duplicate entries and in this case we
             # say it corresponds to the same derivative_keys_dict entry
-        else:
-            self.derivative_keys_dict = 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 dynamic_loss_dict.keys():
+            # we need both because the constraints (all but dyn_loss) will be
+            # done by iterating on u_dict while the dyn_loss will be by
+            # iterating on dynamic_loss_dict. So each time we will require dome
+            # derivative_keys_dict
+
+        # 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] = {"dyn_loss": ["nn_params"]}
-
-        self.dynamic_loss_dict = dynamic_loss_dict
-        self.u_dict = u_dict
+                self.derivative_keys_dict[k] = DerivativeKeysODE(
+                    params=params_dict.extract_params(k)
+                )
 
-        self.loss_weights = loss_weights  # We call the setter
-        # note that self.initial_condition_dict must be
-        # initialized beforehand
+        self._loss_weights = self.set_loss_weights(loss_weights)
 
         # The constaints on the solutions will be implemented by reusing a
         # LossODE class without dynamic loss term
         self.u_constraints_dict = {}
         for i in self.u_dict.keys():
             self.u_constraints_dict[i] = LossODE(
-                u=u_dict[i],
-                loss_weights={
-                    "dyn_loss": 0.0,
-                    "initial_condition": 1.0,
-                    "observations": 1.0,
-                },
+                u=self.u_dict[i],
+                loss_weights=LossWeightsODE(
+                    dyn_loss=0.0,
+                    initial_condition=1.0,
+                    observations=1.0,
+                ),
                 dynamic_loss=None,
                 derivative_keys=self.derivative_keys_dict[i],
                 initial_condition=self.initial_condition_dict[i],
                 obs_slice=self.obs_slice_dict[i],
             )
 
-        # for convenience in the tree_map of evaluate,
-        # we separate the two derivative keys dict
-        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()
-        }
-        self.derivative_keys_u_dict = {
-            k: self.derivative_keys_dict[k]
-            for k in self.u_dict.keys() & self.derivative_keys_dict.keys()
-        }
-
-    @property
-    def loss_weights(self):
-        return self._loss_weights
+        # 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)
 
-    @loss_weights.setter
-    def loss_weights(self, value):
-        self._loss_weights = {}
-        for k, v in value.items():
+    def set_loss_weights(self, loss_weights_init):
+        """
+        This rather complex function enables the user to specify a simple
+        loss_weights=LossWeightsODEDict(dyn_loss=1., initial_condition=Tmax)
+        for ponderating values being applied to all the equations of the
+        system... So all the transformations are handled here
+        """
+        _loss_weights = {}
+        for k in fields(loss_weights_init):
+            v = getattr(loss_weights_init, k.name)
             if isinstance(v, dict):
-                for kk, vv in v.items():
+                for vv in v.values():
                     if not isinstance(vv, (int, float)) and not (
-                        isinstance(vv, jnp.ndarray)
+                        isinstance(vv, Array)
                         and ((vv.shape == (1,) or len(vv.shape) == 0))
                     ):
                         # TODO improve that
                         raise ValueError(
                             f"loss values cannot be vectorial here, got {vv}"
                         )
-                if k == "dyn_loss":
+                if k.name == "dyn_loss":
                     if v.keys() == self.dynamic_loss_dict.keys():
-                        self._loss_weights[k] = v
+                        _loss_weights[k.name] = v
                     else:
                         raise ValueError(
                             "Keys in nested dictionary of loss_weights"
@@ -425,48 +454,41 @@ class SystemLossODE:
                         )
                 else:
                     if v.keys() == self.u_dict.keys():
-                        self._loss_weights[k] = v
+                        _loss_weights[k.name] = v
                     else:
                         raise ValueError(
                             "Keys in nested dictionary of loss_weights"
                             " do not match u_dict keys"
                         )
+            elif v is None:
+                _loss_weights[k.name] = {kk: 0 for kk in self.u_dict.keys()}
             else:
                 if not isinstance(v, (int, float)) and not (
-                    isinstance(v, jnp.ndarray)
-                    and ((v.shape == (1,) or len(v.shape) == 0))
+                    isinstance(v, Array) and ((v.shape == (1,) or len(v.shape) == 0))
                 ):
                     # TODO improve that
                     raise ValueError(f"loss values cannot be vectorial here, got {v}")
-                if k == "dyn_loss":
-                    self._loss_weights[k] = {
+                if k.name == "dyn_loss":
+                    _loss_weights[k.name] = {
                         kk: v for kk in self.dynamic_loss_dict.keys()
                     }
                 else:
-                    self._loss_weights[k] = {kk: v for kk in self.u_dict.keys()}
-        if all(v is None for k, v in self.initial_condition_dict.items()):
-            self._loss_weights["initial_condition"] = {k: 0 for k in self.u_dict.keys()}
-        if "observations" not in value.keys():
-            self._loss_weights["observations"] = {k: 0 for k in self.u_dict.keys()}
+                    _loss_weights[k.name] = {kk: v for kk in self.u_dict.keys()}
+
+        return _loss_weights
 
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
 
-    def evaluate(self, params_dict, batch):
+    def evaluate(self, params_dict: ParamsDict, batch: ODEBatch) -> Float[Array, "1"]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
 
         Parameters
         ---------
-        params_dict
-            A dictionary of dictionaries of parameters of the model.
-            Typically, it is a dictionary of dictionaries of
-            dictionaries: `eq_params` and `nn_params``, respectively the
-            differential equation parameters and the neural network parameter.
-            Note that params_dict["nn_params"] need not be a dictionary anymore
-            but can directly be the parameters. It is useful when working with
-            neural networks sharing the same parameters
+        params
+            A ParamsDict object
         batch
             A ODEBatch object.
             Such a named tuple is composed of a batch of time points
@@ -475,10 +497,10 @@ class SystemLossODE:
             inputs/outputs/parameters
         """
         if (
-            isinstance(params_dict["nn_params"], dict)
-            and self.u_dict.keys() != params_dict["nn_params"].keys()
+            isinstance(params_dict.nn_params, dict)
+            and self.u_dict.keys() != params_dict.nn_params.keys()
         ):
-            raise ValueError("u_dict and params_dict[nn_params] should have same keys ")
+            raise ValueError("u_dict and params_dict.nn_params should have same keys ")
 
         temporal_batch = batch.temporal_batch
 
@@ -489,20 +511,19 @@ class SystemLossODE:
         # and update vmap_in_axes
         if batch.param_batch_dict is not None:
             # update params with the batches of generated params
-            params_dict = _update_eq_params_dict(params_dict, batch.param_batch_dict)
+            params = _update_eq_params_dict(params, batch.param_batch_dict)
 
         vmap_in_axes_params = _get_vmap_in_axes_params(
             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"""
-            params_dict_ = _set_derivatives(params_dict, "dyn_loss", derivative_key)
             return dynamic_loss_apply(
                 dyn_loss.evaluate,
                 self.u_dict,
                 (temporal_batch,),
-                params_dict_,
+                _set_derivatives(params_dict, self.derivative_keys_dyn_loss.dyn_loss),
                 vmap_in_axes_t + vmap_in_axes_params,
                 loss_weight,
                 u_type=PINN,
@@ -511,8 +532,12 @@ class SystemLossODE:
         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
+            # this level. Now that dynamic losses are eqx.Module they can be
+            # traversed by tree map recursion. Hence we need to specify to that
+            # we want to stop at this level
         )
         mse_dyn_loss = jax.tree_util.tree_reduce(
             lambda x, y: x + y, jax.tree_util.tree_leaves(dyn_loss_mse_dict)
@@ -527,7 +552,8 @@ class SystemLossODE:
 
         # we need to do the following for the tree_mapping to work
         if batch.obs_batch_dict is None:
-            batch = batch._replace(obs_batch_dict=self.u_dict_with_none)
+            batch = append_obs_batch(batch, self.u_dict_with_none)
+
         total_loss, res_dict = constraints_system_loss_apply(
             self.u_constraints_dict,
             batch,
@@ -540,27 +566,3 @@ class SystemLossODE:
         total_loss += mse_dyn_loss
         res_dict["dyn_loss"] += mse_dyn_loss
         return total_loss, res_dict
-
-    def tree_flatten(self):
-        children = (
-            self.initial_condition_dict,
-            self._loss_weights,
-        )
-        aux_data = {
-            "u_dict": self.u_dict,
-            "dynamic_loss_dict": self.dynamic_loss_dict,
-            "derivative_keys_dict": self.derivative_keys_dict,
-            "obs_slice_dict": self.obs_slice_dict,
-        }
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        (initial_condition_dict, loss_weights) = children
-        loss_ode = cls(
-            loss_weights=loss_weights,
-            initial_condition_dict=initial_condition_dict,
-            **aux_data,
-        )
-
-        return loss_ode