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

jinns/loss/_LossPDE.py

@@ -1,29 +1,52 @@
+# pylint: disable=unsubscriptable-object, no-member
 """
 Main module to implement a PDE loss in jinns
 """
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
 
+import abc
+from dataclasses import InitVar, fields
+from typing import TYPE_CHECKING, Dict, Callable
 import warnings
 import jax
 import jax.numpy as jnp
-from jax.tree_util import register_pytree_node_class
-from jinns.loss._Losses import (
+import equinox as eqx
+from jaxtyping import Float, Array, Key, Int
+from jinns.loss._loss_utils import (
     dynamic_loss_apply,
     boundary_condition_apply,
     normalization_loss_apply,
     observations_loss_apply,
-    sobolev_reg_apply,
     initial_condition_apply,
     constraints_system_loss_apply,
 )
-from jinns.data._DataGenerators import PDEStatioBatch, PDENonStatioBatch
-from jinns.utils._utils import (
+from jinns.data._DataGenerators import (
+    append_obs_batch,
+)
+from jinns.parameters._params import (
     _get_vmap_in_axes_params,
-    _set_derivatives,
     _update_eq_params_dict,
 )
+from jinns.parameters._derivative_keys import (
+    _set_derivatives,
+    DerivativeKeysPDEStatio,
+    DerivativeKeysPDENonStatio,
+)
+from jinns.loss._loss_weights import (
+    LossWeightsPDEStatio,
+    LossWeightsPDENonStatio,
+    LossWeightsPDEDict,
+)
+from jinns.loss._DynamicLossAbstract import PDEStatio, PDENonStatio
 from jinns.utils._pinn import PINN
 from jinns.utils._spinn import SPINN
-from jinns.loss._operators import _sobolev
+from jinns.data._Batchs import PDEStatioBatch, PDENonStatioBatch
+
+
+if TYPE_CHECKING:
+    from jinns.utils._types import *
 
 _IMPLEMENTED_BOUNDARY_CONDITIONS = [
     "dirichlet",
@@ -31,375 +54,167 @@ _IMPLEMENTED_BOUNDARY_CONDITIONS = [
     "vonneumann",
 ]
 
-_LOSS_WEIGHT_KEYS_PDESTATIO = [
-    "sobolev",
-    "observations",
-    "norm_loss",
-    "boundary_loss",
-    "dyn_loss",
-]
-
-_LOSS_WEIGHT_KEYS_PDENONSTATIO = _LOSS_WEIGHT_KEYS_PDESTATIO + ["initial_condition"]
 
-
-@register_pytree_node_class
-class LossPDEAbstract:
+class _LossPDEAbstract(eqx.Module):
     """
-    Super class for the actual Pinn loss classes. This class should not be
-    used. It serves for common attributes between LossPDEStatio and
-    LossPDENonStatio
-
-
-    **Note:** LossPDEAbstract is jittable. Hence it implements the tree_flatten() and
-    tree_unflatten methods.
+    Parameters
+    ----------
+
+    loss_weights : LossWeightsPDEStatio | LossWeightsPDENonStatio, default=None
+        The loss weights for the differents term : dynamic loss,
+        initial condition (if LossWeightsPDENonStatio), boundary conditions if
+        any, normalization loss if any and observations if any.
+        All fields are set to 1.0 by default.
+    derivative_keys : DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio, 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.
+    omega_boundary_fun : Callable | Dict[str, Callable], default=None
+         The function to be matched in the border condition (can be None) or a
+         dictionary of such functions as values and keys as described
+         in `omega_boundary_condition`.
+    omega_boundary_condition : str | Dict[str, str], default=None
+        Either None (no condition, by default), or a string defining
+        the boundary condition (Dirichlet or Von Neumann),
+        or a dictionary with such strings as values. In this case,
+        the keys are the facets and must be in the following order:
+        1D -> [“xmin”, “xmax”], 2D -> [“xmin”, “xmax”, “ymin”, “ymax”].
+        Note that high order boundaries are currently not implemented.
+        A value in the dict can be None, this means we do not enforce
+        a particular boundary condition on this facet.
+        The facet called “xmin”, resp. “xmax” etc., in 2D,
+        refers to the set of 2D points with fixed “xmin”, resp. “xmax”, etc.
+    omega_boundary_dim : slice | Dict[str, slice], default=None
+        Either None, or a slice object or a dictionary of slice objects as
+        values and keys as described in `omega_boundary_condition`.
+        `omega_boundary_dim` indicates which dimension(s) of the PINN
+        will be forced to match the boundary condition.
+        Note that it must be a slice and not an integer
+        (but a preprocessing of the user provided argument takes care of it)
+    norm_samples : Float[Array, "nb_norm_samples dimension"], default=None
+        Fixed sample point in the space over which to compute the
+        normalization constant. Default is None.
+    norm_int_length : float, default=None
+        A float. Must be provided if `norm_samples` is provided. The domain area
+        (or interval length in 1D) upon which we perform the numerical
+        integration. Default None
+    obs_slice : slice, default=None
+        slice object specifying the begininning/ending of the PINN output
+        that is observed (this is then useful for multidim PINN). Default is None.
+    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,
-        derivative_keys=None,
-        norm_key=None,
-        norm_borders=None,
-        norm_samples=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: DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio | None = (
+        eqx.field(kw_only=True, default=None)
+    )
+    loss_weights: LossWeightsPDEStatio | LossWeightsPDENonStatio | None = eqx.field(
+        kw_only=True, default=None
+    )
+    omega_boundary_fun: Callable | Dict[str, Callable] | None = eqx.field(
+        kw_only=True, default=None, static=True
+    )
+    omega_boundary_condition: str | Dict[str, str] | None = eqx.field(
+        kw_only=True, default=None, static=True
+    )
+    omega_boundary_dim: slice | Dict[str, slice] | None = eqx.field(
+        kw_only=True, default=None, static=True
+    )
+    norm_samples: Float[Array, "nb_norm_samples dimension"] | None = eqx.field(
+        kw_only=True, default=None
+    )
+    norm_int_length: float | 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(kw_only=True, default=None)
+
+    def __post_init__(self, params=None):
         """
-        Parameters
-        ----------
-        u
-            the PINN object
-        loss_weights
-            a dictionary with values used to ponderate each term in the loss
-            function. Valid keys are `dyn_loss`, `norm_loss`, `boundary_loss`
-            and `observations`
-            Note that we can have jnp.arrays with the same dimension of
-            `u` which then ponderates each output of `u`
-        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
-        norm_key
-            Jax random key to draw samples in for the Monte Carlo computation
-            of the normalization constant. Default is None
-        norm_borders
-            tuple of (min, max) of the boundaray values of the space over which
-            to integrate in the computation of the normalization constant.
-            A list of tuple for higher dimensional problems. Default None.
-        norm_samples
-            Fixed sample point in the space over which to compute the
-            normalization constant. Default is None
-
-        Raises
-        ------
-        RuntimeError
-            When provided an invalid combination of `norm_key`, `norm_borders`
-            and `norm_samples`. See note below.
-
-        **Note:** If `norm_key` and `norm_borders` and `norm_samples` are `None`
-        then no normalization loss in enforced.
-        If `norm_borders` and `norm_samples` are given while
-        `norm_samples` is `None` then samples are drawn at each loss evaluation.
-        Otherwise, if `norm_samples` is given, those samples are used.
+        Note that neither __init__ or __post_init__ are called when udating a
+        Module with eqx.tree_at
         """
-
-        self.u = u
-        if derivative_keys is None:
+        if self.derivative_keys is None:
             # be default we only take gradient wrt nn_params
-            derivative_keys = {
-                k: ["nn_params"]
-                for k in [
-                    "dyn_loss",
-                    "boundary_loss",
-                    "norm_loss",
-                    "initial_condition",
-                    "observations",
-                    "sobolev",
-                ]
-            }
-        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",
-                    "boundary_loss",
-                    "norm_loss",
-                    "initial_condition",
-                    "observations",
-                    "sobolev",
-                ]
-            }
-
-        self.derivative_keys = derivative_keys
-        self.loss_weights = loss_weights
-        self.norm_borders = norm_borders
-        self.norm_key = norm_key
-        self.norm_samples = norm_samples
-
-        if norm_key is None and norm_borders is None and norm_samples is None:
-            # if there is None of the 3 above, that means we don't consider
-            # normalization loss
-            self.normalization_loss = None
-        elif (
-            norm_key is not None and norm_borders is not None and norm_samples is None
-        ):  # this ordering so that by default priority is to given mc_samples
-            self.norm_sample_method = "generate"
-            if not isinstance(self.norm_borders[0], tuple):
-                self.norm_borders = (self.norm_borders,)
-            self.norm_xmin, self.norm_xmax = [], []
-            for i, _ in enumerate(self.norm_borders):
-                self.norm_xmin.append(self.norm_borders[i][0])
-                self.norm_xmax.append(self.norm_borders[i][1])
-            self.int_length = jnp.prod(
-                jnp.array(
-                    [
-                        self.norm_xmax[i] - self.norm_xmin[i]
-                        for i in range(len(self.norm_borders))
-                    ]
-                )
-            )
-            self.normalization_loss = True
-        elif norm_samples is None:
-            raise RuntimeError(
-                "norm_borders should always provided then either"
-                " norm_samples (fixed norm_samples) or norm_key (random norm_samples)"
-                " is required."
-            )
-        else:
-            # ok, we are sure we have norm_samples given by the user
-            self.norm_sample_method = "user"
-            if not isinstance(self.norm_borders[0], tuple):
-                self.norm_borders = (self.norm_borders,)
-            self.norm_xmin, self.norm_xmax = [], []
-            for i, _ in enumerate(self.norm_borders):
-                self.norm_xmin.append(self.norm_borders[i][0])
-                self.norm_xmax.append(self.norm_borders[i][1])
-            self.int_length = jnp.prod(
-                jnp.array(
-                    [
-                        self.norm_xmax[i] - self.norm_xmin[i]
-                        for i in range(len(self.norm_borders))
-                    ]
+            try:
+                self.derivative_keys = (
+                    DerivativeKeysPDENonStatio(params=params)
+                    if isinstance(self, LossPDENonStatio)
+                    else DerivativeKeysPDEStatio(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.loss_weights is None:
+            self.loss_weights = (
+                LossWeightsPDENonStatio()
+                if isinstance(self, LossPDENonStatio)
+                else LossWeightsPDEStatio()
             )
-            self.normalization_loss = True
 
-    def get_norm_samples(self):
-        """
-        Returns a batch of points in the domain for integration when the
-        normalization constraint is enforced. The batch of points is either
-        fixed (provided by the user) or regenerated at each iteration.
-        """
-        if self.norm_sample_method == "user":
-            return self.norm_samples
-        if self.norm_sample_method == "generate":
-            ## NOTE TODO CHECK the performances of this for loop
-            norm_samples = []
-            for d in range(len(self.norm_borders)):
-                self.norm_key, subkey = jax.random.split(self.norm_key)
-                norm_samples.append(
-                    jax.random.uniform(
-                        subkey,
-                        shape=(1000, 1),
-                        minval=self.norm_xmin[d],
-                        maxval=self.norm_xmax[d],
-                    )
-                )
-            self.norm_samples = jnp.concatenate(norm_samples, axis=-1)
-            return self.norm_samples
-        raise RuntimeError("Problem with the value of self.norm_sample_method")
-
-    def tree_flatten(self):
-        children = (self.norm_key, self.norm_samples, self.loss_weights)
-        aux_data = {
-            "norm_borders": self.norm_borders,
-            "derivative_keys": self.derivative_keys,
-            "u": self.u,
-        }
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(self, aux_data, children):
-        (norm_key, norm_samples, loss_weights) = children
-        pls = self(
-            aux_data["u"],
-            loss_weights,
-            aux_data["derivative_keys"],
-            norm_key,
-            aux_data["norm_borders"],
-            norm_samples,
-        )
-        return pls
-
-
-@register_pytree_node_class
-class LossPDEStatio(LossPDEAbstract):
-    r"""Loss object for a stationary partial differential equation
-
-    .. math::
-        \mathcal{N}[u](x) = 0, \forall x  \in \Omega
-
-    where :math:`\mathcal{N}[\cdot]` is a differential operator and the
-    boundary condition is :math:`u(x)=u_b(x)` The additional condition of
-    integrating to 1 can be included, i.e. :math:`\int u(x)\mathrm{d}x=1`.
-
-
-    **Note:** LossPDEStatio is jittable. Hence it implements the tree_flatten() and
-    tree_unflatten methods.
-    """
-
-    def __init__(
-        self,
-        u,
-        loss_weights,
-        dynamic_loss,
-        derivative_keys=None,
-        omega_boundary_fun=None,
-        omega_boundary_condition=None,
-        omega_boundary_dim=None,
-        norm_key=None,
-        norm_borders=None,
-        norm_samples=None,
-        sobolev_m=None,
-        obs_slice=None,
-    ):
-        r"""
-        Parameters
-        ----------
-        u
-            the PINN object
-        loss_weights
-            a dictionary with values used to ponderate each term in the loss
-            function. Valid keys are `dyn_loss`, `norm_loss`, `boundary_loss`,
-            `observations` and `sobolev`.
-            Note that we can have jnp.arrays with the same dimension of
-            `u` which then ponderates each output of `u`
-        dynamic_loss
-            the stationary PDE 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
-        omega_boundary_fun
-            The function to be matched in the border condition (can be None)
-            or a dictionary of such function. In this case, the keys are the
-            facets and the values are the functions. The keys must be in the
-            following order: 1D -> ["xmin", "xmax"], 2D -> ["xmin", "xmax",
-            "ymin", "ymax"]. Note that high order boundaries are currently not
-            implemented. A value in the dict can be None, this means we do not
-            enforce a particular boundary condition on this facet.
-            The facet called "xmin", resp. "xmax" etc., in 2D,
-            refers to the set of 2D points with fixed "xmin", resp. "xmax", etc.
-        omega_boundary_condition
-            Either None (no condition), or a string defining the boundary
-            condition e.g. Dirichlet or Von Neumann, or a dictionary of such
-            strings. In this case, the keys are the
-            facets and the values are the strings. The keys must be in the
-            following order: 1D -> ["xmin", "xmax"], 2D -> ["xmin", "xmax",
-            "ymin", "ymax"]. Note that high order boundaries are currently not
-            implemented. A value in the dict can be None, this means we do not
-            enforce a particular boundary condition on this facet.
-            The facet called "xmin", resp. "xmax" etc., in 2D,
-            refers to the set of 2D points with fixed "xmin", resp. "xmax", etc.
-        omega_boundary_dim
-            Either None, or a jnp.s\_ or a dict of jnp.s\_ with keys following
-            the logic of omega_boundary_fun. It indicates which dimension(s) of
-            the PINN will be forced to match the boundary condition
-            Note that it must be a slice and not an integer (a preprocessing of the
-            user provided argument takes care of it)
-        norm_key
-            Jax random key to draw samples in for the Monte Carlo computation
-            of the normalization constant. Default is None
-        norm_borders
-            tuple of (min, max) of the boundaray values of the space over which
-            to integrate in the computation of the normalization constant.
-            A list of tuple for higher dimensional problems. Default None.
-        norm_samples
-            Fixed sample point in the space over which to compute the
-            normalization constant. Default is None
-        sobolev_m
-            An integer. Default is None.
-            It corresponds to the Sobolev regularization order as proposed in
-            *Convergence and error analysis of PINNs*,
-            Doumeche et al., 2023, https://arxiv.org/pdf/2305.01240.pdf
-        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 conditions on omega_boundary_condition and omega_boundary_fun
-            are not respected
-        """
+        if self.obs_slice is None:
+            self.obs_slice = jnp.s_[...]
 
-        super().__init__(
-            u, loss_weights, derivative_keys, norm_key, norm_borders, norm_samples
-        )
+        if (
+            isinstance(self.omega_boundary_fun, dict)
+            and not isinstance(self.omega_boundary_condition, dict)
+        ) or (
+            not isinstance(self.omega_boundary_fun, dict)
+            and isinstance(self.omega_boundary_condition, dict)
+        ):
+            raise ValueError(
+                "if one of self.omega_boundary_fun or "
+                "self.omega_boundary_condition is dict, the other should be too."
+            )
 
-        if omega_boundary_condition is None or omega_boundary_fun is None:
+        if self.omega_boundary_condition is None or self.omega_boundary_fun is None:
             warnings.warn(
                 "Missing boundary function or no boundary condition."
                 "Boundary function is thus ignored."
             )
         else:
-            if isinstance(omega_boundary_condition, dict):
-                for _, v in omega_boundary_condition.items():
+            if isinstance(self.omega_boundary_condition, dict):
+                for _, v in self.omega_boundary_condition.items():
                     if v is not None and not any(
                         v.lower() in s for s in _IMPLEMENTED_BOUNDARY_CONDITIONS
                     ):
                         raise NotImplementedError(
-                            f"The boundary condition {omega_boundary_condition} is not"
+                            f"The boundary condition {self.omega_boundary_condition} is not"
                             f"implemented yet. Try one of :"
                             f"{_IMPLEMENTED_BOUNDARY_CONDITIONS}."
                         )
             else:
                 if not any(
-                    omega_boundary_condition.lower() in s
+                    self.omega_boundary_condition.lower() in s
                     for s in _IMPLEMENTED_BOUNDARY_CONDITIONS
                 ):
                     raise NotImplementedError(
-                        f"The boundary condition {omega_boundary_condition} is not"
+                        f"The boundary condition {self.omega_boundary_condition} is not"
                         f"implemented yet. Try one of :"
                         f"{_IMPLEMENTED_BOUNDARY_CONDITIONS}."
                     )
-                if isinstance(omega_boundary_fun, dict) and isinstance(
-                    omega_boundary_condition, dict
+                if isinstance(self.omega_boundary_fun, dict) and isinstance(
+                    self.omega_boundary_condition, dict
                 ):
                     if (
                         not (
-                            list(omega_boundary_fun.keys()) == ["xmin", "xmax"]
-                            and list(omega_boundary_condition.keys())
+                            list(self.omega_boundary_fun.keys()) == ["xmin", "xmax"]
+                            and list(self.omega_boundary_condition.keys())
                             == ["xmin", "xmax"]
                         )
                     ) or (
                         not (
-                            list(omega_boundary_fun.keys())
+                            list(self.omega_boundary_fun.keys())
                             == ["xmin", "xmax", "ymin", "ymax"]
-                            and list(omega_boundary_condition.keys())
+                            and list(self.omega_boundary_condition.keys())
                             == ["xmin", "xmax", "ymin", "ymax"]
                         )
                     ):
@@ -408,10 +223,6 @@ class LossPDEStatio(LossPDEAbstract):
                             "boundary condition dictionaries is incorrect"
                         )
 
-        self.omega_boundary_fun = omega_boundary_fun
-        self.omega_boundary_condition = omega_boundary_condition
-
-        self.omega_boundary_dim = omega_boundary_dim
         if isinstance(self.omega_boundary_fun, dict):
             if self.omega_boundary_dim is None:
                 self.omega_boundary_dim = {
@@ -440,44 +251,144 @@ class LossPDEStatio(LossPDEAbstract):
                     self.omega_boundary_dim : self.omega_boundary_dim + 1
                 ]
             if not isinstance(self.omega_boundary_dim, slice):
-                raise ValueError("self.omega_boundary_dim must be a jnp.s_" " object")
+                raise ValueError("self.omega_boundary_dim must be a jnp.s_ object")
 
-        self.dynamic_loss = dynamic_loss
+        if self.norm_samples is not None and self.norm_int_length is None:
+            raise ValueError("self.norm_samples and norm_int_length must be provided")
 
-        self.sobolev_m = sobolev_m
-        if self.sobolev_m is not None:
-            self.sobolev_reg = _sobolev(
-                self.u, self.sobolev_m
-            )  # we return a function, that way
-            # the order of sobolev_m is static and the conditional in the recursive
-            # function is properly set
-        else:
-            self.sobolev_reg = None
+    @abc.abstractmethod
+    def evaluate(
+        self: eqx.Module,
+        params: Params,
+        batch: PDEStatioBatch | PDENonStatioBatch,
+    ) -> tuple[Float, dict]:
+        raise NotImplementedError
 
-        for k in _LOSS_WEIGHT_KEYS_PDESTATIO:
-            if k not in self.loss_weights.keys():
-                self.loss_weights[k] = 0
 
-        if (
-            isinstance(self.omega_boundary_fun, dict)
-            and not isinstance(self.omega_boundary_condition, dict)
-        ) or (
-            not isinstance(self.omega_boundary_fun, dict)
-            and isinstance(self.omega_boundary_condition, dict)
-        ):
-            raise ValueError(
-                "if one of self.omega_boundary_fun or "
-                "self.omega_boundary_condition is dict, the other should be too."
-            )
+class LossPDEStatio(_LossPDEAbstract):
+    r"""Loss object for a stationary partial differential equation
 
-        self.obs_slice = obs_slice
-        if self.obs_slice is None:
-            self.obs_slice = jnp.s_[...]
+    $$
+        \mathcal{N}[u](x) = 0, \forall x  \in \Omega
+    $$
+
+    where $\mathcal{N}[\cdot]$ is a differential operator and the
+    boundary condition is $u(x)=u_b(x)$ The additional condition of
+    integrating to 1 can be included, i.e. $\int u(x)\mathrm{d}x=1$.
+
+    Parameters
+    ----------
+    u : eqx.Module
+        the PINN
+    dynamic_loss : DynamicLoss
+        the stationary PDE dynamic part of the loss, basically the differential
+        operator $\mathcal{N}[u](x)$. Should implement a method
+        `dynamic_loss.evaluate(x, u, params)`.
+        Can be None in order to access only some part of the evaluate call
+        results.
+    key : Key
+        A JAX PRNG Key for the loss class treated as an attribute. Default is
+        None. This field is provided for future developments and additional
+        losses that might need some randomness. Note that special care must be
+        taken when splitting the key because in-place updates are forbidden in
+        eqx.Modules.
+    loss_weights : LossWeightsPDEStatio, default=None
+        The loss weights for the differents term : dynamic loss,
+        boundary conditions if any, normalization loss if any and
+        observations if any.
+        All fields are set to 1.0 by default.
+    derivative_keys : DerivativeKeysPDEStatio, 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.
+    omega_boundary_fun : Callable | Dict[str, Callable], default=None
+         The function to be matched in the border condition (can be None) or a
+         dictionary of such functions as values and keys as described
+         in `omega_boundary_condition`.
+    omega_boundary_condition : str | Dict[str, str], default=None
+        Either None (no condition, by default), or a string defining
+        the boundary condition (Dirichlet or Von Neumann),
+        or a dictionary with such strings as values. In this case,
+        the keys are the facets and must be in the following order:
+        1D -> [“xmin”, “xmax”], 2D -> [“xmin”, “xmax”, “ymin”, “ymax”].
+        Note that high order boundaries are currently not implemented.
+        A value in the dict can be None, this means we do not enforce
+        a particular boundary condition on this facet.
+        The facet called “xmin”, resp. “xmax” etc., in 2D,
+        refers to the set of 2D points with fixed “xmin”, resp. “xmax”, etc.
+    omega_boundary_dim : slice | Dict[str, slice], default=None
+        Either None, or a slice object or a dictionary of slice objects as
+        values and keys as described in `omega_boundary_condition`.
+        `omega_boundary_dim` indicates which dimension(s) of the PINN
+        will be forced to match the boundary condition.
+        Note that it must be a slice and not an integer
+        (but a preprocessing of the user provided argument takes care of it)
+    norm_samples : Float[Array, "nb_norm_samples dimension"], default=None
+        Fixed sample point in the space over which to compute the
+        normalization constant. Default is None.
+    norm_int_length : float, default=None
+        A float. Must be provided if `norm_samples` is provided. The domain area
+        (or interval length in 1D) upon which we perform the numerical
+        integration. Default None
+    obs_slice : slice, default=None
+        slice object specifying the begininning/ending of the PINN output
+        that is observed (this is then useful for multidim PINN). Default is None.
+    params : InitVar[Params], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
+
+
+    Raises
+    ------
+    ValueError
+        If conditions on omega_boundary_condition and omega_boundary_fun
+        are not respected
+    """
+
+    # 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
+    key: Key | None = eqx.field(kw_only=True, default=None)
+
+    vmap_in_axes: tuple[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
+        Module with eqx.tree_at!
+        """
+        super().__post_init__(
+            params=params
+        )  # because __init__ or __post_init__ of Base
+        # class is not automatically called
+
+        self.vmap_in_axes = (0,)  # for x only here
+
+    def _get_dynamic_loss_batch(
+        self, batch: PDEStatioBatch
+    ) -> tuple[Float[Array, "batch_size dimension"]]:
+        return (batch.inside_batch,)
+
+    def _get_normalization_loss_batch(
+        self, _
+    ) -> Float[Array, "nb_norm_samples dimension"]:
+        return (self.norm_samples,)
+
+    def _get_observations_loss_batch(
+        self, batch: PDEStatioBatch
+    ) -> Float[Array, "batch_size obs_dim"]:
+        return (batch.obs_batch_dict["pinn_in"],)
 
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
 
-    def evaluate(self, params, batch):
+    def evaluate(
+        self, params: Params, batch: PDEStatioBatch
+    ) -> tuple[Float[Array, "1"], dict[str, float]]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
@@ -485,22 +396,14 @@ class LossPDEStatio(LossPDEAbstract):
         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 PDEStatioBatch object.
-            Such a named tuple is composed of a batch of points in the
+            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
         """
-        omega_batch, _ = batch.inside_batch, batch.border_batch
-
-        vmap_in_axes_x = (0,)
-
         # Retrieve the optional eq_params_batch
         # and update eq_params with the latter
         # and update vmap_in_axes
@@ -511,44 +414,41 @@ class LossPDEStatio(LossPDEAbstract):
         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,
-                (omega_batch,),
-                params_,
-                vmap_in_axes_x + vmap_in_axes_params,
-                self.loss_weights["dyn_loss"],
+                self._get_dynamic_loss_batch(batch),
+                _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)
 
         # normalization part
-        params_ = _set_derivatives(params, "norm_loss", self.derivative_keys)
-        if self.normalization_loss is not None:
+        if self.norm_samples is not None:
             mse_norm_loss = normalization_loss_apply(
                 self.u,
-                (self.get_norm_samples(),),
-                params_,
-                vmap_in_axes_x + vmap_in_axes_params,
-                self.int_length,
-                self.loss_weights["norm_loss"],
+                self._get_normalization_loss_batch(batch),
+                _set_derivatives(params, self.derivative_keys.norm_loss),
+                self.vmap_in_axes + vmap_in_axes_params,
+                self.norm_int_length,
+                self.loss_weights.norm_loss,
             )
         else:
             mse_norm_loss = jnp.array(0.0)
 
         # boundary part
-        params_ = _set_derivatives(params, "boundary_loss", self.derivative_keys)
         if self.omega_boundary_condition is not None:
             mse_boundary_loss = boundary_condition_apply(
                 self.u,
                 batch,
-                params_,
+                _set_derivatives(params, self.derivative_keys.boundary_loss),
                 self.omega_boundary_fun,
                 self.omega_boundary_condition,
                 self.omega_boundary_dim,
-                self.loss_weights["boundary_loss"],
+                self.loss_weights.boundary_loss,
             )
         else:
             mse_boundary_loss = jnp.array(0.0)
@@ -558,40 +458,21 @@ class LossPDEStatio(LossPDEAbstract):
             # update params with the batches of observed params
             params = _update_eq_params_dict(params, batch.obs_batch_dict["eq_params"])
 
-            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_x + vmap_in_axes_params,
+                self._get_observations_loss_batch(batch),
+                _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:
             mse_observation_loss = jnp.array(0.0)
 
-        # Sobolev regularization
-        params_ = _set_derivatives(params, "sobolev", self.derivative_keys)
-        if self.sobolev_reg is not None:
-            mse_sobolev_loss = sobolev_reg_apply(
-                self.u,
-                (omega_batch,),
-                params_,
-                vmap_in_axes_x + vmap_in_axes_params,
-                self.sobolev_reg,
-                self.loss_weights["sobolev"],
-            )
-        else:
-            mse_sobolev_loss = jnp.array(0.0)
-
         # total loss
         total_loss = (
-            mse_dyn_loss
-            + mse_norm_loss
-            + mse_boundary_loss
-            + mse_observation_loss
-            + mse_sobolev_loss
+            mse_dyn_loss + mse_norm_loss + mse_boundary_loss + mse_observation_loss
         )
         return total_loss, (
             {
@@ -599,205 +480,148 @@ class LossPDEStatio(LossPDEAbstract):
                 "norm_loss": mse_norm_loss,
                 "boundary_loss": mse_boundary_loss,
                 "observations": mse_observation_loss,
-                "sobolev": mse_sobolev_loss,
                 "initial_condition": jnp.array(0.0),  # for compatibility in the
                 # tree_map of SystemLoss
             }
         )
 
-    def tree_flatten(self):
-        children = (self.norm_key, self.norm_samples, self.loss_weights)
-        aux_data = {
-            "u": self.u,
-            "dynamic_loss": self.dynamic_loss,
-            "derivative_keys": self.derivative_keys,
-            "omega_boundary_fun": self.omega_boundary_fun,
-            "omega_boundary_condition": self.omega_boundary_condition,
-            "omega_boundary_dim": self.omega_boundary_dim,
-            "norm_borders": self.norm_borders,
-            "sobolev_m": self.sobolev_m,
-            "obs_slice": self.obs_slice,
-        }
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        (norm_key, norm_samples, loss_weights) = children
-        pls = cls(
-            aux_data["u"],
-            loss_weights,
-            aux_data["dynamic_loss"],
-            aux_data["derivative_keys"],
-            aux_data["omega_boundary_fun"],
-            aux_data["omega_boundary_condition"],
-            aux_data["omega_boundary_dim"],
-            norm_key,
-            aux_data["norm_borders"],
-            norm_samples,
-            aux_data["sobolev_m"],
-            aux_data["obs_slice"],
-        )
-        return pls
 
-
-@register_pytree_node_class
 class LossPDENonStatio(LossPDEStatio):
     r"""Loss object for a stationary partial differential equation
 
-    .. math::
+    $$
         \mathcal{N}[u](t, x) = 0, \forall t \in I, \forall x \in \Omega
+    $$
 
-    where :math:`\mathcal{N}[\cdot]` is a differential operator.
-    The boundary condition is :math:`u(t, x)=u_b(t, x),\forall
-    x\in\delta\Omega, \forall t`.
-    The initial condition is :math:`u(0, x)=u_0(x), \forall x\in\Omega`
+    where $\mathcal{N}[\cdot]$ is a differential operator.
+    The boundary condition is $u(t, x)=u_b(t, x),\forall
+    x\in\delta\Omega, \forall t$.
+    The initial condition is $u(0, x)=u_0(x), \forall x\in\Omega$
     The additional condition of
-    integrating to 1 can be included, i.e., :math:`\int u(t, x)\mathrm{d}x=1`.
-
+    integrating to 1 can be included, i.e., $\int u(t, x)\mathrm{d}x=1$.
+
+    Parameters
+    ----------
+    u : eqx.Module
+        the PINN
+    dynamic_loss : DynamicLoss
+        the non stationary PDE dynamic part of the loss, basically the differential
+        operator $\mathcal{N}[u](t, x)$. Should implement a method
+        `dynamic_loss.evaluate(t, x, u, params)`.
+        Can be None in order to access only some part of the evaluate call
+        results.
+    key : Key
+        A JAX PRNG Key for the loss class treated as an attribute. Default is
+        None. This field is provided for future developments and additional
+        losses that might need some randomness. Note that special care must be
+        taken when splitting the key because in-place updates are forbidden in
+        eqx.Modules.
+        reason
+    loss_weights : LossWeightsPDENonStatio, default=None
+        The loss weights for the differents term : dynamic loss,
+        boundary conditions if any, initial condition, normalization loss if any and
+        observations if any.
+        All fields are set to 1.0 by default.
+    derivative_keys : DerivativeKeysPDENonStatio, 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.
+    omega_boundary_fun : Callable | Dict[str, Callable], default=None
+         The function to be matched in the border condition (can be None) or a
+         dictionary of such functions as values and keys as described
+         in `omega_boundary_condition`.
+    omega_boundary_condition : str | Dict[str, str], default=None
+        Either None (no condition, by default), or a string defining
+        the boundary condition (Dirichlet or Von Neumann),
+        or a dictionary with such strings as values. In this case,
+        the keys are the facets and must be in the following order:
+        1D -> [“xmin”, “xmax”], 2D -> [“xmin”, “xmax”, “ymin”, “ymax”].
+        Note that high order boundaries are currently not implemented.
+        A value in the dict can be None, this means we do not enforce
+        a particular boundary condition on this facet.
+        The facet called “xmin”, resp. “xmax” etc., in 2D,
+        refers to the set of 2D points with fixed “xmin”, resp. “xmax”, etc.
+    omega_boundary_dim : slice | Dict[str, slice], default=None
+        Either None, or a slice object or a dictionary of slice objects as
+        values and keys as described in `omega_boundary_condition`.
+        `omega_boundary_dim` indicates which dimension(s) of the PINN
+        will be forced to match the boundary condition.
+        Note that it must be a slice and not an integer
+        (but a preprocessing of the user provided argument takes care of it)
+    norm_samples : Float[Array, "nb_norm_samples dimension"], default=None
+        Fixed sample point in the space over which to compute the
+        normalization constant. Default is None.
+    norm_int_length : float, default=None
+        A float. Must be provided if `norm_samples` is provided. The domain area
+        (or interval length in 1D) upon which we perform the numerical
+        integration. Default None
+    obs_slice : slice, default=None
+        slice object specifying the begininning/ending of the PINN output
+        that is observed (this is then useful for multidim PINN). Default is None.
+    initial_condition_fun : Callable, default=None
+        A function representing the temporal initial condition. If None
+        (default) then no initial condition is applied
+    params : InitVar[Params], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
 
-    **Note:** LossPDENonStatio is jittable. Hence it implements the tree_flatten() and
-    tree_unflatten methods.
     """
 
-    def __init__(
-        self,
-        u,
-        loss_weights,
-        dynamic_loss,
-        derivative_keys=None,
-        omega_boundary_fun=None,
-        omega_boundary_condition=None,
-        omega_boundary_dim=None,
-        initial_condition_fun=None,
-        norm_key=None,
-        norm_borders=None,
-        norm_samples=None,
-        sobolev_m=None,
-        obs_slice=None,
-    ):
-        r"""
-        Parameters
-        ----------
-        u
-            the PINN object
-        loss_weights
-            dictionary of values for loss term ponderation
-            Note that we can have jnp.arrays with the same dimension of
-            `u` which then ponderates each output of `u`
-        dynamic_loss
-            A Dynamic loss object whose evaluate method corresponds to the
-            dynamic term in the loss
-            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
-        omega_boundary_fun
-            The function to be matched in the border condition (can be None)
-            or a dictionary of such function. In this case, the keys are the
-            facets and the values are the functions. The keys must be in the
-            following order: 1D -> ["xmin", "xmax"], 2D -> ["xmin", "xmax",
-            "ymin", "ymax"]. Note that high order boundaries are currently not
-            implemented. A value in the dict can be None, this means we do not
-            enforce a particular boundary condition on this facet.
-            The facet called "xmin", resp. "xmax" etc., in 2D,
-            refers to the set of 2D points with fixed "xmin", resp. "xmax", etc.
-        omega_boundary_condition
-            Either None (no condition), or a string defining the boundary
-            condition e.g. Dirichlet or Von Neumann, or a dictionary of such
-            strings. In this case, the keys are the
-            facets and the values are the strings. The keys must be in the
-            following order: 1D -> ["xmin", "xmax"], 2D -> ["xmin", "xmax",
-            "ymin", "ymax"]. Note that high order boundaries are currently not
-            implemented. A value in the dict can be None, this means we do not
-            enforce a particular boundary condition on this facet.
-            The facet called "xmin", resp. "xmax" etc., in 2D,
-            refers to the set of 2D points with fixed "xmin", resp. "xmax", etc.
-        omega_boundary_dim
-            Either None, or a jnp.s\_ or a dict of jnp.s\_ with keys following
-            the logic of omega_boundary_fun. It indicates which dimension(s) of
-            the PINN will be forced to match the boundary condition
-            Note that it must be a slice and not an integer (a preprocessing of the
-            user provided argument takes care of it)
-        initial_condition_fun
-            A function representing the temporal initial condition. If None
-            (default) then no initial condition is applied
-        norm_key
-            Jax random key to draw samples in for the Monte Carlo computation
-            of the normalization constant. Default is None
-        norm_borders
-            tuple of (min, max) of the boundaray values of the space over which
-            to integrate in the computation of the normalization constant.
-            A list of tuple for higher dimensional problems. Default None.
-        norm_samples
-            Fixed sample point in the space over which to compute the
-            normalization constant. Default is None
-        sobolev_m
-            An integer. Default is None.
-            It corresponds to the Sobolev regularization order as proposed in
-            *Convergence and error analysis of PINNs*,
-            Doumeche et al., 2023, https://arxiv.org/pdf/2305.01240.pdf
-        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.
-
+    # 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
+    initial_condition_fun: Callable | None = eqx.field(
+        kw_only=True, default=None, static=True
+    )
 
+    def __post_init__(self, params=None):
+        """
+        Note that neither __init__ or __post_init__ are called when udating a
+        Module with eqx.tree_at!
         """
+        super().__post_init__(
+            params=params
+        )  # because __init__ or __post_init__ of Base
+        # class is not automatically called
 
-        super().__init__(
-            u,
-            loss_weights,
-            dynamic_loss,
-            derivative_keys,
-            omega_boundary_fun,
-            omega_boundary_condition,
-            omega_boundary_dim,
-            norm_key,
-            norm_borders,
-            norm_samples,
-            sobolev_m=sobolev_m,
-            obs_slice=obs_slice,
-        )
-        if initial_condition_fun is None:
+        self.vmap_in_axes = (0, 0)  # for t and x
+
+        if self.initial_condition_fun 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)."
             )
-        self.initial_condition_fun = initial_condition_fun
-
-        self.sobolev_m = sobolev_m
-        if self.sobolev_m is not None:
-            # This overwrite the wrongly initialized self.sobolev_reg with
-            # statio=True in the LossPDEStatio init
-            self.sobolev_reg = _sobolev(self.u, self.sobolev_m, statio=False)
-            # we return a function, that way
-            # the order of sobolev_m is static and the conditional in the recursive
-            # function is properly set
-        else:
-            self.sobolev_reg = None
 
-        for k in _LOSS_WEIGHT_KEYS_PDENONSTATIO:
-            if k not in self.loss_weights.keys():
-                self.loss_weights[k] = 0
+    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)
+
+    def _get_normalization_loss_batch(
+        self, batch: PDENonStatioBatch
+    ) -> tuple[Float[Array, "batch_size 1"], Float[Array, "nb_norm_samples dimension"]]:
+        return (
+            batch.times_x_inside_batch[:, 0:1],
+            self.norm_samples,
+        )
+
+    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:],
+        )
 
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
 
     def evaluate(
-        self,
-        params,
-        batch,
-    ):
+        self, params: Params, batch: PDENonStatioBatch
+    ) -> tuple[Float[Array, "1"], dict[str, float]]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
@@ -805,191 +629,54 @@ class LossPDENonStatio(LossPDEStatio):
         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 PDENonStatioBatch object.
-            Such a named tuple is composed of a batch of points in
+            Composed of a batch of points in
             the domain, a batch of points in the domain
             border, a batch of time points and an optional additional batch
             of parameters (eg. for metamodeling) and an optional additional batch of observed
             inputs/outputs/parameters
         """
-
-        times_batch = batch.times_x_inside_batch[:, 0:1]
         omega_batch = batch.times_x_inside_batch[:, 1:]
-        n = omega_batch.shape[0]
-
-        vmap_in_axes_x_t = (0, 0)
 
         # Retrieve the optional eq_params_batch
         # and update eq_params with the latter
         # and update vmap_in_axes
         if batch.param_batch_dict is not None:
-            eq_params_batch_dict = batch.param_batch_dict
-
-            # feed the eq_params with the batch
-            for k in eq_params_batch_dict.keys():
-                params["eq_params"][k] = eq_params_batch_dict[k]
+            # update eq_params with the batches of generated params
+            params = _update_eq_params_dict(params, batch.param_batch_dict)
 
         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,
-                (times_batch, omega_batch),
-                params_,
-                vmap_in_axes_x_t + vmap_in_axes_params,
-                self.loss_weights["dyn_loss"],
-            )
-        else:
-            mse_dyn_loss = jnp.array(0.0)
-
-        # normalization part
-        params_ = _set_derivatives(params, "norm_loss", self.derivative_keys)
-        if self.normalization_loss is not None:
-            mse_norm_loss = normalization_loss_apply(
-                self.u,
-                (times_batch, self.get_norm_samples()),
-                params_,
-                vmap_in_axes_x_t + vmap_in_axes_params,
-                self.int_length,
-                self.loss_weights["norm_loss"],
-            )
-        else:
-            mse_norm_loss = jnp.array(0.0)
-
-        # boundary part
-        params_ = _set_derivatives(params, "boundary_loss", self.derivative_keys)
-        if self.omega_boundary_fun is not None:
-            mse_boundary_loss = boundary_condition_apply(
-                self.u,
-                batch,
-                params_,
-                self.omega_boundary_fun,
-                self.omega_boundary_condition,
-                self.omega_boundary_dim,
-                self.loss_weights["boundary_loss"],
-            )
-        else:
-            mse_boundary_loss = jnp.array(0.0)
+        # For mse_dyn_loss, mse_norm_loss, mse_boundary_loss,
+        # mse_observation_loss we use the evaluate from parent class
+        partial_mse, partial_mse_terms = super().evaluate(params, batch)
 
         # initial condition
-        params_ = _set_derivatives(params, "initial_condition", self.derivative_keys)
         if self.initial_condition_fun is not None:
             mse_initial_condition = initial_condition_apply(
                 self.u,
                 omega_batch,
-                params_,
+                _set_derivatives(params, self.derivative_keys.initial_condition),
                 (0,) + vmap_in_axes_params,
                 self.initial_condition_fun,
-                n,
-                self.loss_weights["initial_condition"],
+                omega_batch.shape[0],
+                self.loss_weights.initial_condition,
             )
         else:
             mse_initial_condition = jnp.array(0.0)
 
-        # Observation mse
-        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_ = _set_derivatives(params, "observations", self.derivative_keys)
-            mse_observation_loss = observations_loss_apply(
-                self.u,
-                (
-                    batch.obs_batch_dict["pinn_in"][:, 0:1],
-                    batch.obs_batch_dict["pinn_in"][:, 1:],
-                ),
-                params_,
-                vmap_in_axes_x_t + vmap_in_axes_params,
-                batch.obs_batch_dict["val"],
-                self.loss_weights["observations"],
-                self.obs_slice,
-            )
-        else:
-            mse_observation_loss = jnp.array(0.0)
-
-        # Sobolev regularization
-        params_ = _set_derivatives(params, "sobolev", self.derivative_keys)
-        if self.sobolev_reg is not None:
-            mse_sobolev_loss = sobolev_reg_apply(
-                self.u,
-                (omega_batch, times_batch),
-                params_,
-                vmap_in_axes_x_t + vmap_in_axes_params,
-                self.sobolev_reg,
-                self.loss_weights["sobolev"],
-            )
-        else:
-            mse_sobolev_loss = jnp.array(0.0)
-
         # total loss
-        total_loss = (
-            mse_dyn_loss
-            + mse_norm_loss
-            + mse_boundary_loss
-            + mse_initial_condition
-            + mse_observation_loss
-            + mse_sobolev_loss
-        )
-
-        return total_loss, (
-            {
-                "dyn_loss": mse_dyn_loss,
-                "norm_loss": mse_norm_loss,
-                "boundary_loss": mse_boundary_loss,
-                "initial_condition": mse_initial_condition,
-                "observations": mse_observation_loss,
-                "sobolev": mse_sobolev_loss,
-            }
-        )
+        total_loss = partial_mse + mse_initial_condition
 
-    def tree_flatten(self):
-        children = (self.norm_key, self.norm_samples, self.loss_weights)
-        aux_data = {
-            "u": self.u,
-            "dynamic_loss": self.dynamic_loss,
-            "derivative_keys": self.derivative_keys,
-            "omega_boundary_fun": self.omega_boundary_fun,
-            "omega_boundary_condition": self.omega_boundary_condition,
-            "omega_boundary_dim": self.omega_boundary_dim,
-            "initial_condition_fun": self.initial_condition_fun,
-            "norm_borders": self.norm_borders,
-            "sobolev_m": self.sobolev_m,
-            "obs_slice": self.obs_slice,
+        return total_loss, {
+            **partial_mse_terms,
+            "initial_condition": mse_initial_condition,
         }
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        (norm_key, norm_samples, loss_weights) = children
-        pls = cls(
-            aux_data["u"],
-            loss_weights,
-            aux_data["dynamic_loss"],
-            aux_data["derivative_keys"],
-            aux_data["omega_boundary_fun"],
-            aux_data["omega_boundary_condition"],
-            aux_data["omega_boundary_dim"],
-            aux_data["initial_condition_fun"],
-            norm_key,
-            aux_data["norm_borders"],
-            norm_samples,
-            aux_data["sobolev_m"],
-            aux_data["obs_slice"],
-        )
-        return pls
 
 
-@register_pytree_node_class
-class SystemLossPDE:
-    """
+class SystemLossPDE(eqx.Module):
+    r"""
     Class to implement a system of PDEs.
     The goal is to give maximum freedom to the user. The class is created with
     a dict of dynamic loss, and dictionaries of all the objects that are used
@@ -1003,190 +690,182 @@ class SystemLossPDE:
     Indeed, these dictionaries (except `dynamic_loss_dict`) are tied to one
     solution.
 
-    **Note:** SystemLossPDE is jittable. Hence it implements the tree_flatten() and
-    tree_unflatten methods.
+    Parameters
+    ----------
+    u_dict : Dict[str, eqx.Module]
+        dict of PINNs
+    loss_weights : LossWeightsPDEDict
+        A dictionary of LossWeightsODE
+    derivative_keys_dict : Dict[str, DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio], default=None
+        A dictionnary of DerivativeKeysPDEStatio or DerivativeKeysPDENonStatio
+        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.
+    dynamic_loss_dict : Dict[str, PDEStatio | PDENonStatio]
+        A dict of dynamic part of the loss, basically the differential
+        operator $\mathcal{N}[u](t, x)$ or $\mathcal{N}[u](x)$.
+    key_dict : Dict[str, Key], default=None
+        A dictionary of JAX PRNG keys. The dictionary keys of key_dict must
+        match that of u_dict. See LossPDEStatio or LossPDENonStatio for
+        more details.
+    omega_boundary_fun_dict : Dict[str, Callable | Dict[str, Callable] | None], default=None
+        A dict of of function or of dict of functions or of None
+        (see doc for `omega_boundary_fun` in
+        LossPDEStatio or LossPDENonStatio). Default is None.
+        Must share the keys of `u_dict`.
+    omega_boundary_condition_dict : Dict[str, str | Dict[str, str] | None], default=None
+        A dict of strings or of dict of strings or of None
+        (see doc for `omega_boundary_condition_dict` in
+        LossPDEStatio or LossPDENonStatio). Default is None.
+        Must share the keys of `u_dict`
+    omega_boundary_dim_dict : Dict[str, slice | Dict[str, slice] | None], default=None
+        A dict of slices or of dict of slices or of None
+        (see doc for `omega_boundary_dim` in
+        LossPDEStatio or LossPDENonStatio). Default is None.
+        Must share the keys of `u_dict`
+    initial_condition_fun_dict : Dict[str, Callable | None], default=None
+        A dict of functions representing the temporal initial condition (None
+        value is possible). If None
+        (default) then no temporal boundary condition is applied
+        Must share the keys of `u_dict`
+    norm_samples_dict : Dict[str, Float[Array, "nb_norm_samples dimension"] | None, default=None
+        A dict of fixed sample point in the space over which to compute the
+        normalization constant. Default is None
+        Must share the keys of `u_dict`
+    norm_int_length_dict : Dict[str, float | None] | None, default=None
+        A dict of Float. The domain area
+        (or interval length in 1D) upon which we perform the numerical
+        integration for each element of u_dict.
+        Default is None
+        Must share the keys of `u_dict`
+    obs_slice_dict : Dict[str, slice | None] | None, default=None
+        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
+    params : InitVar[ParamsDict], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
+
     """
 
-    def __init__(
-        self,
-        u_dict,
-        loss_weights,
-        dynamic_loss_dict,
-        nn_type_dict,
-        derivative_keys_dict=None,
-        omega_boundary_fun_dict=None,
-        omega_boundary_condition_dict=None,
-        omega_boundary_dim_dict=None,
-        initial_condition_fun_dict=None,
-        norm_key_dict=None,
-        norm_borders_dict=None,
-        norm_samples_dict=None,
-        sobolev_m_dict=None,
-        obs_slice_dict=None,
-    ):
-        r"""
-        Parameters
-        ----------
-        u_dict
-            A dict of PINNs
-        loss_weights
-            A dictionary of dictionaries with values used to
-            ponderate each term in the loss
-            function. 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`
-        dynamic_loss_dict
-            A dict of dynamic part of the loss, basically the differential
-            operator :math:`\mathcal{N}[u](t)`.
-        nn_type_dict
-            A dict whose keys are that of u_dict whose value is either
-            `nn_statio` or `nn_nonstatio` which signifies either the PINN has a
-            time component in input or not.
-        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
-        omega_boundary_fun_dict
-            A dict of dict of functions (see doc for `omega_boundary_fun` in
-            LossPDEStatio or LossPDENonStatio). Default is None.
-            Must share the keys of `u_dict`.
-        omega_boundary_condition_dict
-            A dict of dict of strings (see doc for
-            `omega_boundary_condition_dict` in
-            LossPDEStatio or LossPDENonStatio). Default is None.
-            Must share the keys of `u_dict`
-        omega_boundary_dim_dict
-            A dict of dict of slices (see doc for `omega_boundary_dim` in
-            LossPDEStatio or LossPDENonStatio). Default is None.
-            Must share the keys of `u_dict`
-        initial_condition_fun_dict
-            A dict of functions representing the temporal initial condition. If None
-            (default) then no temporal boundary condition is applied
-            Must share the keys of `u_dict`
-        norm_key_dict
-            A dict of Jax random keys to draw samples in for the Monte Carlo computation
-            of the normalization constant. Default is None
-            Must share the keys of `u_dict`
-        norm_borders_dict
-            A dict of tuples of (min, max) of the boundaray values of the space over which
-            to integrate in the computation of the normalization constant.
-            A list of tuple for higher dimensional problems. Default None.
-            Must share the keys of `u_dict`
-        norm_samples_dict
-            A dict of fixed sample point in the space over which to compute the
-            normalization constant. Default is None
-            Must share the keys of `u_dict`
-        sobolev_m
-            Default is None. A dictionary of integers, one per key which must
-            match `u_dict`.
-            It corresponds to the Sobolev regularization order as proposed in
-            *Convergence and error analysis of PINNs*,
-            Doumeche et al., 2023, https://arxiv.org/pdf/2305.01240.pdf
-        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, PDEStatio | PDENonStatio]
+    key_dict: Dict[str, Key] | None = eqx.field(kw_only=True, default=None)
+    derivative_keys_dict: Dict[
+        str, DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio | None
+    ] = eqx.field(kw_only=True, default=None)
+    omega_boundary_fun_dict: Dict[str, Callable | Dict[str, Callable] | None] | None = (
+        eqx.field(kw_only=True, default=None, static=True)
+    )
+    omega_boundary_condition_dict: Dict[str, str | Dict[str, str] | None] | None = (
+        eqx.field(kw_only=True, default=None, static=True)
+    )
+    omega_boundary_dim_dict: Dict[str, slice | Dict[str, slice] | None] | None = (
+        eqx.field(kw_only=True, default=None, static=True)
+    )
+    initial_condition_fun_dict: Dict[str, Callable | None] | None = eqx.field(
+        kw_only=True, default=None, static=True
+    )
+    norm_samples_dict: Dict[str, Float[Array, "nb_norm_samples dimension"]] | None = (
+        eqx.field(kw_only=True, default=None)
+    )
+    norm_int_length_dict: Dict[str, float | None] | 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
+    )
+
+    # For the user loss_weights are passed as a LossWeightsPDEDict (with internal
+    # dictionary having keys in u_dict and / or dynamic_loss_dict)
+    loss_weights: InitVar[LossWeightsPDEDict | None] = eqx.field(
+        kw_only=True, default=None
+    )
+    params_dict: InitVar[ParamsDict] = eqx.field(kw_only=True, default=None)
+
+    # following have init=False and are set in the __post_init__
+    u_constraints_dict: Dict[str, LossPDEStatio | LossPDENonStatio] = eqx.field(
+        init=False
+    )
+    derivative_keys_dyn_loss: DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio = (
+        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()}
+        self.u_dict_with_none = {k: None for k in self.u_dict.keys()}
         # First, for all the optional dict,
         # if the user did not provide at all this optional argument,
         # we make sure there is a null ponderating loss_weight and we
         # create a dummy dict with the required keys and all the values to
         # None
-        if omega_boundary_fun_dict is None:
+        if self.key_dict is None:
+            self.key_dict = self.u_dict_with_none
+        if self.omega_boundary_fun_dict is None:
             self.omega_boundary_fun_dict = self.u_dict_with_none
-        else:
-            self.omega_boundary_fun_dict = omega_boundary_fun_dict
-        if omega_boundary_condition_dict is None:
+        if self.omega_boundary_condition_dict is None:
             self.omega_boundary_condition_dict = self.u_dict_with_none
-        else:
-            self.omega_boundary_condition_dict = omega_boundary_condition_dict
-        if omega_boundary_dim_dict is None:
+        if self.omega_boundary_dim_dict is None:
             self.omega_boundary_dim_dict = self.u_dict_with_none
-        else:
-            self.omega_boundary_dim_dict = omega_boundary_dim_dict
-        if initial_condition_fun_dict is None:
+        if self.initial_condition_fun_dict is None:
             self.initial_condition_fun_dict = self.u_dict_with_none
-        else:
-            self.initial_condition_fun_dict = initial_condition_fun_dict
-        if norm_key_dict is None:
-            self.norm_key_dict = self.u_dict_with_none
-        else:
-            self.norm_key_dict = norm_key_dict
-        if norm_borders_dict is None:
-            self.norm_borders_dict = self.u_dict_with_none
-        else:
-            self.norm_borders_dict = norm_borders_dict
-        if norm_samples_dict is None:
+        if self.norm_samples_dict is None:
             self.norm_samples_dict = self.u_dict_with_none
-        else:
-            self.norm_samples_dict = norm_samples_dict
-        if sobolev_m_dict is None:
-            self.sobolev_m_dict = self.u_dict_with_none
-        else:
-            self.sobolev_m_dict = sobolev_m_dict
-        if obs_slice_dict is None:
-            self.obs_slice_dict = {k: jnp.s_[...] for k in u_dict.keys()}
-        else:
-            self.obs_slice_dict = obs_slice_dict
-            if u_dict.keys() != obs_slice_dict.keys():
+        if self.norm_int_length_dict is None:
+            self.norm_int_length_dict = self.u_dict_with_none
+        if self.obs_slice_dict is None:
+            self.obs_slice_dict = {k: jnp.s_[...] for k in self.u_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"]}
+                if self.u_dict[k].eq_type == "statio_PDE":
+                    self.derivative_keys_dict[k] = DerivativeKeysPDEStatio(
+                        params=params_dict.extract_params(k)
+                    )
+                else:
+                    self.derivative_keys_dict[k] = DerivativeKeysPDENonStatio(
+                        params=params_dict.extract_params(k)
+                    )
 
         # Second we make sure that all the dicts (except dynamic_loss_dict) have the same keys
         if (
-            u_dict.keys() != nn_type_dict.keys()
-            or u_dict.keys() != self.omega_boundary_fun_dict.keys()
-            or u_dict.keys() != self.omega_boundary_condition_dict.keys()
-            or u_dict.keys() != self.omega_boundary_dim_dict.keys()
-            or u_dict.keys() != self.initial_condition_fun_dict.keys()
-            or u_dict.keys() != self.norm_key_dict.keys()
-            or u_dict.keys() != self.norm_borders_dict.keys()
-            or u_dict.keys() != self.norm_samples_dict.keys()
-            or u_dict.keys() != self.sobolev_m_dict.keys()
+            self.u_dict.keys() != self.key_dict.keys()
+            or self.u_dict.keys() != self.omega_boundary_fun_dict.keys()
+            or self.u_dict.keys() != self.omega_boundary_condition_dict.keys()
+            or self.u_dict.keys() != self.omega_boundary_dim_dict.keys()
+            or self.u_dict.keys() != self.initial_condition_fun_dict.keys()
+            or self.u_dict.keys() != self.norm_samples_dict.keys()
+            or self.u_dict.keys() != self.norm_int_length_dict.keys()
         ):
             raise ValueError("All the dicts concerning the PINNs should have same keys")
 
-        self.dynamic_loss_dict = dynamic_loss_dict
-        self.u_dict = u_dict
-        # TODO nn_type should become a class attribute now that we have PINN
-        # class and SPINNs class
-        self.nn_type_dict = nn_type_dict
-
-        self.loss_weights = loss_weights  # This calls the setter
+        self._loss_weights = self.set_loss_weights(loss_weights)
 
         # Third, in order not to benefit from LossPDEStatio and
         # LossPDENonStatio and in order to factorize code, we create internally
@@ -1194,95 +873,93 @@ class SystemLossPDE:
         # We will not use the dynamic loss term
         self.u_constraints_dict = {}
         for i in self.u_dict.keys():
-            if self.nn_type_dict[i] == "nn_statio":
+            if self.u_dict[i].eq_type == "statio_PDE":
                 self.u_constraints_dict[i] = LossPDEStatio(
-                    u=u_dict[i],
-                    loss_weights={
-                        "dyn_loss": 0.0,
-                        "norm_loss": 1.0,
-                        "boundary_loss": 1.0,
-                        "observations": 1.0,
-                        "sobolev": 1.0,
-                    },
+                    u=self.u_dict[i],
+                    loss_weights=LossWeightsPDENonStatio(
+                        dyn_loss=0.0,
+                        norm_loss=1.0,
+                        boundary_loss=1.0,
+                        observations=1.0,
+                        initial_condition=1.0,
+                    ),
                     dynamic_loss=None,
+                    key=self.key_dict[i],
                     derivative_keys=self.derivative_keys_dict[i],
                     omega_boundary_fun=self.omega_boundary_fun_dict[i],
                     omega_boundary_condition=self.omega_boundary_condition_dict[i],
                     omega_boundary_dim=self.omega_boundary_dim_dict[i],
-                    norm_key=self.norm_key_dict[i],
-                    norm_borders=self.norm_borders_dict[i],
                     norm_samples=self.norm_samples_dict[i],
-                    sobolev_m=self.sobolev_m_dict[i],
+                    norm_int_length=self.norm_int_length_dict[i],
                     obs_slice=self.obs_slice_dict[i],
                 )
-            elif self.nn_type_dict[i] == "nn_nonstatio":
+            elif self.u_dict[i].eq_type == "nonstatio_PDE":
                 self.u_constraints_dict[i] = LossPDENonStatio(
-                    u=u_dict[i],
-                    loss_weights={
-                        "dyn_loss": 0.0,
-                        "norm_loss": 1.0,
-                        "boundary_loss": 1.0,
-                        "observations": 1.0,
-                        "initial_condition": 1.0,
-                        "sobolev": 1.0,
-                    },
+                    u=self.u_dict[i],
+                    loss_weights=LossWeightsPDENonStatio(
+                        dyn_loss=0.0,
+                        norm_loss=1.0,
+                        boundary_loss=1.0,
+                        observations=1.0,
+                        initial_condition=1.0,
+                    ),
                     dynamic_loss=None,
+                    key=self.key_dict[i],
                     derivative_keys=self.derivative_keys_dict[i],
                     omega_boundary_fun=self.omega_boundary_fun_dict[i],
                     omega_boundary_condition=self.omega_boundary_condition_dict[i],
                     omega_boundary_dim=self.omega_boundary_dim_dict[i],
                     initial_condition_fun=self.initial_condition_fun_dict[i],
-                    norm_key=self.norm_key_dict[i],
-                    norm_borders=self.norm_borders_dict[i],
                     norm_samples=self.norm_samples_dict[i],
-                    sobolev_m=self.sobolev_m_dict[i],
+                    norm_int_length=self.norm_int_length_dict[i],
+                    obs_slice=self.obs_slice_dict[i],
                 )
             else:
                 raise ValueError(
-                    f"Wrong value for nn_type_dict[i], got {nn_type_dict[i]}"
+                    "Wrong value for self.u_dict[i].eq_type[i], "
+                    f"got {self.u_dict[i].eq_type[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()
-        }
+        # 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 = DerivativeKeysPDENonStatio(params=params_dict)
 
         # also make sure we only have PINNs or SPINNs
         if not (
-            all(isinstance(value, PINN) for value in u_dict.values())
-            or all(isinstance(value, SPINN) for value in u_dict.values())
+            all(isinstance(value, PINN) for value in self.u_dict.values())
+            or all(isinstance(value, SPINN) for value in self.u_dict.values())
         ):
             raise ValueError(
                 "We only accept dictionary of PINNs or dictionary of SPINNs"
             )
 
-    @property
-    def loss_weights(self):
-        return self._loss_weights
-
-    @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: LossWeightsPDEDict
+    ) -> dict[str, dict]:
+        """
+        This rather complex function enables the user to specify a simple
+        loss_weights=LossWeightsPDEDict(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.keys():
                     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"
@@ -1290,51 +967,36 @@ class SystemLossPDE:
                         )
                 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"
                         )
+            if 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()}
-        # Some special checks below
-        if all(v is None for k, v in self.sobolev_m_dict.items()):
-            self._loss_weights["sobolev"] = {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()}
-        if all(v is None for k, v in self.omega_boundary_fun_dict.items()) or all(
-            v is None for k, v in self.omega_boundary_condition_dict.items()
-        ):
-            self._loss_weights["boundary_loss"] = {k: 0 for k in self.u_dict.keys()}
-        if (
-            all(v is None for k, v in self.norm_key_dict.items())
-            or all(v is None for k, v in self.norm_borders_dict.items())
-            or all(v is None for k, v in self.norm_samples_dict.items())
-        ):
-            self._loss_weights["norm_loss"] = {k: 0 for k in self.u_dict.keys()}
-        if all(v is None for k, v in self.initial_condition_fun_dict.items()):
-            self._loss_weights["initial_condition"] = {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,
-    ):
+        params_dict: ParamsDict,
+        batch: PDEStatioBatch | PDENonStatioBatch,
+    ) -> tuple[Float[Array, "1"], dict[str, float]]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
@@ -1342,12 +1004,8 @@ class SystemLossPDE:
         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
+            Parameters at which the losses of the system are evaluated
         batch
-            A PDEStatioBatch or PDENonStatioBatch object.
             Such named tuples are composed of  batch of points in the
             domain, a batch of points in the domain
             border, (a batch of time points a for PDENonStatioBatch) and an
@@ -1355,7 +1013,7 @@ class SystemLossPDE:
             and an optional additional batch of observed
             inputs/outputs/parameters
         """
-        if self.u_dict.keys() != params_dict["nn_params"].keys():
+        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):
@@ -1378,22 +1036,22 @@ class SystemLossPDE:
         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]
+                params_dict.eq_params[k] = eq_params_batch_dict[k]
 
         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):
             """The function used in tree_map"""
-            params_dict_ = _set_derivatives(params_dict, "dyn_loss", derivative_key)
             return dynamic_loss_apply(
                 dyn_loss.evaluate,
                 self.u_dict,
                 batches,
-                params_dict_,
+                _set_derivatives(params_dict, self.derivative_keys_dyn_loss.dyn_loss),
                 vmap_in_axes_x_or_x_t + vmap_in_axes_params,
                 loss_weight,
                 u_type=type(list(self.u_dict.values())[0]),
@@ -1402,8 +1060,14 @@ class SystemLossPDE:
         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, (PDEStatio, PDENonStatio)
+            ),  # 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)
@@ -1418,11 +1082,10 @@ class SystemLossPDE:
             "boundary_loss": "*",
             "observations": "*",
             "initial_condition": "*",
-            "sobolev": "*",
         }
         # 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,
@@ -1435,41 +1098,3 @@ class SystemLossPDE:
         total_loss += mse_dyn_loss
         res_dict["dyn_loss"] += mse_dyn_loss
         return total_loss, res_dict
-
-    def tree_flatten(self):
-        children = (
-            self.norm_key_dict,
-            self.norm_samples_dict,
-            self.initial_condition_fun_dict,
-            self._loss_weights,
-        )
-        aux_data = {
-            "u_dict": self.u_dict,
-            "dynamic_loss_dict": self.dynamic_loss_dict,
-            "norm_borders_dict": self.norm_borders_dict,
-            "omega_boundary_fun_dict": self.omega_boundary_fun_dict,
-            "omega_boundary_condition_dict": self.omega_boundary_condition_dict,
-            "nn_type_dict": self.nn_type_dict,
-            "sobolev_m_dict": self.sobolev_m_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):
-        (
-            norm_key_dict,
-            norm_samples_dict,
-            initial_condition_fun_dict,
-            loss_weights,
-        ) = children
-        loss_ode = cls(
-            loss_weights=loss_weights,
-            norm_key_dict=norm_key_dict,
-            norm_samples_dict=norm_samples_dict,
-            initial_condition_fun_dict=initial_condition_fun_dict,
-            **aux_data,
-        )
-
-        return loss_ode