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

jinns/loss/_DynamicLossAbstract.py

@@ -73,9 +73,11 @@ def _decorator_heteregeneous_params(evaluate, eq_type):
 class DynamicLoss(eqx.Module):
     r"""
     Abstract base class for dynamic losses. Implements the physical term:
+
     $$
         \mathcal{N}[u](t, x) = 0
     $$
+
     for **one** point $t$, $x$ or $(t, x)$, depending on the context.
 
     Parameters

jinns/loss/_LossODE.py

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

jinns/loss/_LossPDE.py

@@ -103,6 +103,9 @@ class _LossPDEAbstract(eqx.Module):
     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.
     """
 
     # NOTE static=True only for leaf attributes that are not valid JAX types
@@ -129,18 +132,26 @@ class _LossPDEAbstract(eqx.Module):
     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)
 
-    def __post_init__(self):
+    params: InitVar[Params] = eqx.field(kw_only=True, default=None)
+
+    def __post_init__(self, params=None):
         """
         Note that neither __init__ or __post_init__ are called when udating a
         Module with eqx.tree_at
         """
         if self.derivative_keys is None:
             # be default we only take gradient wrt nn_params
-            self.derivative_keys = (
-                DerivativeKeysPDENonStatio()
-                if isinstance(self, LossPDENonStatio)
-                else DerivativeKeysPDEStatio()
-            )
+            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 = (
@@ -324,6 +335,9 @@ class LossPDEStatio(_LossPDEAbstract):
     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
@@ -342,12 +356,14 @@ class LossPDEStatio(_LossPDEAbstract):
 
     vmap_in_axes: tuple[Int] = eqx.field(init=False, static=True)
 
-    def __post_init__(self):
+    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__()  # because __init__ or __post_init__ of Base
+        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
@@ -547,6 +563,9 @@ class LossPDENonStatio(LossPDEStatio):
     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.
 
     """
 
@@ -556,12 +575,14 @@ class LossPDENonStatio(LossPDEStatio):
         kw_only=True, default=None, static=True
     )
 
-    def __post_init__(self):
+    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__()  # because __init__ or __post_init__ of Base
+        super().__post_init__(
+            params=params
+        )  # because __init__ or __post_init__ of Base
         # class is not automatically called
 
         self.vmap_in_axes = (0, 0)  # for t and x
@@ -616,7 +637,6 @@ class LossPDENonStatio(LossPDEStatio):
             of parameters (eg. for metamodeling) and an optional additional batch of observed
             inputs/outputs/parameters
         """
-
         omega_batch = batch.times_x_inside_batch[:, 1:]
 
         # Retrieve the optional eq_params_batch
@@ -725,6 +745,9 @@ class SystemLossPDE(eqx.Module):
         PINNs. Default is None. But if a value is given, all the entries of
         `u_dict` must be represented here with default value `jnp.s_[...]`
         if no particular slice is to be given
+    params : InitVar[ParamsDict], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
 
     """
 
@@ -763,22 +786,20 @@ class SystemLossPDE(eqx.Module):
     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_u_dict: Dict[
-        str, DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio
-    ] = eqx.field(init=False)
-    derivative_keys_dyn_loss_dict: Dict[
-        str, DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio
-    ] = 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):
+    def __post_init__(self, loss_weights=None, params_dict=None):
         # a dictionary that will be useful at different places
         self.u_dict_with_none = {k: None for k in self.u_dict.keys()}
         # First, for all the optional dict,
@@ -818,24 +839,19 @@ class SystemLossPDE(eqx.Module):
             # iterating on dynamic_loss_dict. So each time we will require dome
             # derivative_keys_dict
 
-        # but then if the user did not provide anything, we must at least have
-        # a default value for the dynamic_loss_dict keys entries in
-        # self.derivative_keys_dict since the computation of dynamic losses is
-        # made without create a loss object that would provide the
-        # default values
-        for k in self.dynamic_loss_dict.keys():
+        # derivative keys for the u_constraints. Note that we create missing
+        # DerivativeKeysODE around a Params object and not ParamsDict
+        # this works because u_dict.keys == params_dict.nn_params.keys()
+        for k in self.u_dict.keys():
             if self.derivative_keys_dict[k] is None:
-                try:
-                    if self.u_dict[k].eq_type == "statio_PDE":
-                        self.derivative_keys_dict[k] = DerivativeKeysPDEStatio()
-                    else:
-                        self.derivative_keys_dict[k] = DerivativeKeysPDENonStatio()
-                except KeyError:  # We are in a key that is not in u_dict but in
-                    # dynamic_loss_dict
-                    if isinstance(self.dynamic_loss_dict[k], PDEStatio):
-                        self.derivative_keys_dict[k] = DerivativeKeysPDEStatio()
-                    else:
-                        self.derivative_keys_dict[k] = DerivativeKeysPDENonStatio()
+                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 (
@@ -904,16 +920,11 @@ class SystemLossPDE(eqx.Module):
                     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 (
@@ -1034,13 +1045,13 @@ class SystemLossPDE(eqx.Module):
             batch.param_batch_dict, params_dict
         )
 
-        def dyn_loss_for_one_key(dyn_loss, derivative_key, loss_weight):
+        def dyn_loss_for_one_key(dyn_loss, loss_weight):
             """The function used in tree_map"""
             return dynamic_loss_apply(
                 dyn_loss.evaluate,
                 self.u_dict,
                 batches,
-                _set_derivatives(params_dict, derivative_key.dyn_loss),
+                _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]),
@@ -1049,7 +1060,6 @@ class SystemLossPDE(eqx.Module):
         dyn_loss_mse_dict = jax.tree_util.tree_map(
             dyn_loss_for_one_key,
             self.dynamic_loss_dict,
-            self.derivative_keys_dyn_loss_dict,
             self._loss_weights["dyn_loss"],
             is_leaf=lambda x: isinstance(
                 x, (PDEStatio, PDENonStatio)

jinns/loss/_loss_utils.py

@@ -297,12 +297,12 @@ def constraints_system_loss_apply(
     if isinstance(params_dict.nn_params, dict):
 
         def apply_u_constraint(
-            u_constraint, nn_params, loss_weights_for_u, obs_batch_u
+            u_constraint, nn_params, eq_params, loss_weights_for_u, obs_batch_u
         ):
             res_dict_for_u = u_constraint.evaluate(
                 Params(
                     nn_params=nn_params,
-                    eq_params=params_dict.eq_params,
+                    eq_params=eq_params,
                 ),
                 append_obs_batch(batch, obs_batch_u),
             )[1]
@@ -319,6 +319,11 @@ def constraints_system_loss_apply(
             apply_u_constraint,
             u_constraints_dict,
             params_dict.nn_params,
+            (
+                params_dict.eq_params
+                if params_dict.eq_params.keys() == params_dict.nn_params.keys()
+                else {k: params_dict.eq_params for k in params_dict.nn_params.keys()}
+            ),  # this manipulation is needed since we authorize eq_params not to have the same structure as nn_params in ParamsDict
             loss_weights_T,
             batch.obs_batch_dict,
             is_leaf=lambda x: (

jinns/parameters/_derivative_keys.py

@@ -2,50 +2,468 @@
 Formalize the data structure for the derivative keys
 """
 
-from dataclasses import fields
+from functools import partial
+from dataclasses import fields, InitVar
 from typing import Literal
 import jax
 import equinox as eqx
 
-from jinns.parameters._params import Params
+from jinns.parameters._params import Params, ParamsDict
+
+
+def _get_masked_parameters(
+    derivative_mask_str: str, params: Params | ParamsDict
+) -> Params | ParamsDict:
+    """
+    Creates the Params object with True values where we want to differentiate
+    """
+    if isinstance(params, Params):
+        # start with a params object with True everywhere. We will update to False
+        # for parameters wrt which we do want not to differentiate the loss
+        diff_params = jax.tree.map(
+            lambda x: True,
+            params,
+            is_leaf=lambda x: isinstance(x, eqx.Module)
+            and not isinstance(x, Params),  # do not travers nn_params, more
+            # granularity could be imagined here, in the future
+        )
+        if derivative_mask_str == "both":
+            return diff_params
+        if derivative_mask_str == "eq_params":
+            return eqx.tree_at(lambda p: p.nn_params, diff_params, False)
+        if derivative_mask_str == "nn_params":
+            return eqx.tree_at(
+                lambda p: p.eq_params,
+                diff_params,
+                jax.tree.map(lambda x: False, params.eq_params),
+            )
+        raise ValueError(
+            "Bad value for DerivativeKeys. Got "
+            f'{derivative_mask_str}, expected "both", "nn_params" or '
+            ' "eq_params"'
+        )
+    elif isinstance(params, ParamsDict):
+        # do not travers nn_params, more
+        # granularity could be imagined here, in the future
+        diff_params = ParamsDict(
+            nn_params=True, eq_params=jax.tree.map(lambda x: True, params.eq_params)
+        )
+        if derivative_mask_str == "both":
+            return diff_params
+        if derivative_mask_str == "eq_params":
+            return eqx.tree_at(lambda p: p.nn_params, diff_params, False)
+        if derivative_mask_str == "nn_params":
+            return eqx.tree_at(
+                lambda p: p.eq_params,
+                diff_params,
+                jax.tree.map(lambda x: False, params.eq_params),
+            )
+        raise ValueError(
+            "Bad value for DerivativeKeys. Got "
+            f'{derivative_mask_str}, expected "both", "nn_params" or '
+            ' "eq_params"'
+        )
+
+    else:
+        raise ValueError(
+            f"Bad value for params. Got {type(params)}, expected Params "
+            " or ParamsDict"
+        )
 
 
 class DerivativeKeysODE(eqx.Module):
-    # we use static = True because all fields are string, hence should be
-    # invisible by JAX transforms (JIT, etc.)
-    dyn_loss: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
-        kw_only=True, default="nn_params", static=True
-    )
-    observations: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
-        kw_only=True, default="nn_params", static=True
+    """
+    A class that specifies with repect to which parameter(s) each term of the
+    loss is differentiated. For example, you can specify that the
+    [`DynamicLoss`][jinns.loss.DynamicLoss] should be differentiated both with
+    respect to the neural network parameters *and* the equation parameters, or only some of them.
+
+    To do so, user can either use strings or a `Params` object
+    with PyTree structure matching the parameters of the problem at
+    hand, and booleans indicating if gradient is to be taken or not. Internally,
+    a `jax.lax.stop_gradient()` is appropriately set to each `True` node when
+    computing each loss term.
+
+    !!! note
+
+         1. For unspecified loss term, the default is to differentiate with
+        respect to `"nn_params"` only.
+         2. No granularity inside `Params.nn_params` is currently supported.
+         3. Note that the main Params or ParamsDict object of the problem is mandatory if initialization via `from_str()`.
+
+    A typical specification is of the form:
+    ```python
+    Params(
+        nn_params=True | False,
+        eq_params={
+            "alpha":True | False,
+            "beta":True | False,
+            ...
+        }
     )
-    initial_condition: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
-        kw_only=True, default="nn_params", static=True
+    ```
+
+    Parameters
+    ----------
+    dyn_loss : Params | ParamsDict | None, default=None
+        Tell wrt which node of `Params` we will differentiate the
+        dynamic loss. To do so, the fields of `Params` contain True (if
+        differentiation) or False (if no differentiation).
+    observations : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        observation loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+    initial_condition : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        initial condition loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+    params : InitVar[Params | ParamsDict], default=None
+        The main Params object of the problem. It is required
+        if some terms are unspecified (None). This is because, jinns cannot
+        infer the content of `Params.eq_params`.
+    """
+
+    dyn_loss: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
+    observations: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
+    initial_condition: Params | ParamsDict | None = eqx.field(
+        kw_only=True, default=None
     )
 
+    params: InitVar[Params | ParamsDict] = eqx.field(kw_only=True, default=None)
+
+    def __post_init__(self, params=None):
+        if self.dyn_loss is None:
+            try:
+                self.dyn_loss = _get_masked_parameters("nn_params", params)
+            except AttributeError:
+                raise ValueError(
+                    "self.dyn_loss is None, hence params should be " "passed"
+                )
+        if self.observations is None:
+            try:
+                self.observations = _get_masked_parameters("nn_params", params)
+            except AttributeError:
+                raise ValueError(
+                    "self.observations is None, hence params should be " "passed"
+                )
+        if self.initial_condition is None:
+            try:
+                self.initial_condition = _get_masked_parameters("nn_params", params)
+            except AttributeError:
+                raise ValueError(
+                    "self.initial_condition is None, hence params should be " "passed"
+                )
+
+    @classmethod
+    def from_str(
+        cls,
+        params: Params | ParamsDict,
+        dyn_loss: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+        observations: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+        initial_condition: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+    ):
+        """
+        Construct the DerivativeKeysODE from strings. For each term of the
+        loss, specify whether to differentiate wrt the neural network
+        parameters, the equation parameters or both. The `Params` object, which
+        contains the actual array of parameters must be passed to
+        construct the fields with the appropriate PyTree structure.
+
+        !!! note
+            You can mix strings and `Params` if you need granularity.
+
+        Parameters
+        ----------
+        params
+            The actual Params or ParamsDict object of the problem.
+        dyn_loss
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the dynamic loss. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        observations
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the observations. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        initial_condition
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the initial condition. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        """
+        return DerivativeKeysODE(
+            dyn_loss=(
+                _get_masked_parameters(dyn_loss, params)
+                if isinstance(dyn_loss, str)
+                else dyn_loss
+            ),
+            observations=(
+                _get_masked_parameters(observations, params)
+                if isinstance(observations, str)
+                else observations
+            ),
+            initial_condition=(
+                _get_masked_parameters(initial_condition, params)
+                if isinstance(initial_condition, str)
+                else initial_condition
+            ),
+        )
+
 
 class DerivativeKeysPDEStatio(eqx.Module):
+    """
+    See [jinns.parameters.DerivativeKeysODE][].
 
-    dyn_loss: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
-        kw_only=True, default="nn_params", static=True
-    )
-    observations: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
-        kw_only=True, default="nn_params", static=True
-    )
-    boundary_loss: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
-        kw_only=True, default="nn_params", static=True
-    )
-    norm_loss: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
-        kw_only=True, default="nn_params", static=True
-    )
+    Parameters
+    ----------
+     dyn_loss : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        dynamic loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+     observations : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        observation loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+     boundary_loss : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        boundary loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+     norm_loss : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        normalization loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+     params : InitVar[Params | ParamsDict], default=None
+        The main Params object of the problem. It is required
+        if some terms are unspecified (None). This is because, jinns cannot infer the
+        content of `Params.eq_params`.
+    """
+
+    dyn_loss: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
+    observations: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
+    boundary_loss: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
+    norm_loss: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
+
+    params: InitVar[Params | ParamsDict] = eqx.field(kw_only=True, default=None)
+
+    def __post_init__(self, params=None):
+        if self.dyn_loss is None:
+            try:
+                self.dyn_loss = _get_masked_parameters("nn_params", params)
+            except AttributeError:
+                raise ValueError("self.dyn_loss is None, hence params should be passed")
+        if self.observations is None:
+            try:
+                self.observations = _get_masked_parameters("nn_params", params)
+            except AttributeError:
+                raise ValueError(
+                    "self.observations is None, hence params should be passed"
+                )
+        if self.boundary_loss is None:
+            try:
+                self.boundary_loss = _get_masked_parameters("nn_params", params)
+            except AttributeError:
+                raise ValueError(
+                    "self.boundary_loss is None, hence params should be passed"
+                )
+        if self.norm_loss is None:
+            try:
+                self.norm_loss = _get_masked_parameters("nn_params", params)
+            except AttributeError:
+                raise ValueError(
+                    "self.norm_loss is None, hence params should be passed"
+                )
+
+    @classmethod
+    def from_str(
+        cls,
+        params: Params | ParamsDict,
+        dyn_loss: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+        observations: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+        boundary_loss: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+        norm_loss: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+    ):
+        """
+        See [jinns.parameters.DerivativeKeysODE.from_str][].
+
+        Parameters
+        ----------
+        params
+            The actual Param or ParamsDict object of the problem.
+        dyn_loss
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the dynamic loss. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        observations
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the observations. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        boundary_loss
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the boundary loss. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        norm_loss
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the normalization loss. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        """
+        return DerivativeKeysPDEStatio(
+            dyn_loss=(
+                _get_masked_parameters(dyn_loss, params)
+                if isinstance(dyn_loss, str)
+                else dyn_loss
+            ),
+            observations=(
+                _get_masked_parameters(observations, params)
+                if isinstance(observations, str)
+                else observations
+            ),
+            boundary_loss=(
+                _get_masked_parameters(boundary_loss, params)
+                if isinstance(boundary_loss, str)
+                else boundary_loss
+            ),
+            norm_loss=(
+                _get_masked_parameters(norm_loss, params)
+                if isinstance(norm_loss, str)
+                else norm_loss
+            ),
+        )
 
 
 class DerivativeKeysPDENonStatio(DerivativeKeysPDEStatio):
+    """
+    See [jinns.parameters.DerivativeKeysODE][].
 
-    initial_condition: Literal["nn_params", "eq_params", "both"] = eqx.field(
-        kw_only=True, default="nn_params", static=True
+    Parameters
+    ----------
+    dyn_loss : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        dynamic loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+    observations : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        observation loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+    boundary_loss : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        boundary loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+    norm_loss : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        normalization loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+    initial_condition : Params | ParamsDict | None, default=None
+        Tell wrt which parameters among Params we will differentiate the
+        initial_condition loss. To do so, the fields of Params contain True (if
+        differentiation) or False (if no differentiation).
+    params : InitVar[Params | ParamsDict], default=None
+        The main Params object of the problem. It is required
+        if some terms are unspecified (None). This is because, jinns cannot infer the
+        content of `Params.eq_params`.
+    """
+
+    initial_condition: Params | ParamsDict | None = eqx.field(
+        kw_only=True, default=None
     )
 
+    def __post_init__(self, params=None):
+        super().__post_init__(params=params)
+        if self.initial_condition is None:
+            try:
+                self.initial_condition = _get_masked_parameters("nn_params", params)
+            except AttributeError:
+                raise ValueError(
+                    "self.initial_condition is None, hence params should be passed"
+                )
+
+    @classmethod
+    def from_str(
+        cls,
+        params: Params | ParamsDict,
+        dyn_loss: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+        observations: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+        boundary_loss: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+        norm_loss: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+        initial_condition: (
+            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+        ) = "nn_params",
+    ):
+        """
+        See [jinns.parameters.DerivativeKeysODE.from_str][].
+
+        Parameters
+        ----------
+        params
+            The actual Params | ParamsDict object of the problem.
+        dyn_loss
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the dynamic loss. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        observations
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the observations. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        boundary_loss
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the boundary loss. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        norm_loss
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the normalization loss. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        initial_condition
+            Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
+            `"both"` we will differentiate the initial_condition loss. Default is
+            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+        """
+        return DerivativeKeysPDENonStatio(
+            dyn_loss=(
+                _get_masked_parameters(dyn_loss, params)
+                if isinstance(dyn_loss, str)
+                else dyn_loss
+            ),
+            observations=(
+                _get_masked_parameters(observations, params)
+                if isinstance(observations, str)
+                else observations
+            ),
+            boundary_loss=(
+                _get_masked_parameters(boundary_loss, params)
+                if isinstance(boundary_loss, str)
+                else boundary_loss
+            ),
+            norm_loss=(
+                _get_masked_parameters(norm_loss, params)
+                if isinstance(norm_loss, str)
+                else norm_loss
+            ),
+            initial_condition=(
+                _get_masked_parameters(initial_condition, params)
+                if isinstance(initial_condition, str)
+                else initial_condition
+            ),
+        )
+
 
 def _set_derivatives(params, derivative_keys):
     """
@@ -53,42 +471,51 @@ def _set_derivatives(params, derivative_keys):
     has a copy of the params with appropriate derivatives set
     """
 
-    def _set_derivatives_(loss_term_derivative):
-        if loss_term_derivative == "both":
-            return params
-        # the next line put a stop_gradient around the fields that do not
-        # appear in loss_term_derivative. Currently there are only two possible
-        # values nn_params and eq_params but there might be more in the future
-        return eqx.tree_at(
-            lambda p: tuple(
-                getattr(p, f.name)
-                for f in fields(Params)
-                if f.name != loss_term_derivative
+    def _set_derivatives_ParamsDict(params_, derivative_mask):
+        """
+        The next lines put a stop_gradient around the fields that do not
+        differentiate the loss term
+        **Note:** **No granularity inside `ParamsDict.nn_params` is currently
+        supported.**
+        This means a typical Params specification is of the form:
+        `ParamsDict(nn_params=True | False, eq_params={"0":{"alpha":True | False,
+        "beta":True | False}}, "1":{"alpha":True | False, "beta":True | False}})`.
+        """
+        # a ParamsDict object is reconstructed by hand since we do not want to
+        # traverse nn_params, for now...
+        return ParamsDict(
+            nn_params=jax.lax.cond(
+                derivative_mask.nn_params,
+                lambda p: p,
+                jax.lax.stop_gradient,
+                params_.nn_params,
+            ),
+            eq_params=jax.tree.map(
+                lambda p, d: jax.lax.cond(d, lambda p: p, jax.lax.stop_gradient, p),
+                params_.eq_params,
+                derivative_mask.eq_params,
             ),
-            params,
-            replace_fn=jax.lax.stop_gradient,
         )
 
-    def _set_derivatives_dict(loss_term_derivative):
-        if loss_term_derivative == "both":
-            return params
-        # the next line put a stop_gradient around the fields that do not
-        # appear in loss_term_derivative. Currently there are only two possible
-        # values nn_params and eq_params but there might be more in the future
-        return {
-            k: eqx.tree_at(
-                lambda p: tuple(
-                    getattr(p, f.name)
-                    for f in fields(Params)
-                    if f.name != loss_term_derivative
-                ),
-                params_,
-                replace_fn=jax.lax.stop_gradient,
-            )
-            for k, params_ in params
-        }
+    def _set_derivatives_(params_, derivative_mask):
+        """
+        The next lines put a stop_gradient around the fields that do not
+        differentiate the loss term
+        **Note:** **No granularity inside `Params.nn_params` is currently
+        supported.**
+        This means a typical Params specification is of the form:
+        `Params(nn_params=True | False, eq_params={"alpha":True | False,
+        "beta":True | False})`.
+        """
+        return jax.tree.map(
+            lambda p, d: jax.lax.cond(d, lambda p: p, jax.lax.stop_gradient, p),
+            params_,
+            derivative_mask,
+            is_leaf=lambda x: isinstance(x, eqx.Module)
+            and not isinstance(x, Params),  # do not travers nn_params, more
+            # granularity could be imagined here, in the future
+        )
 
-    if not isinstance(params, dict):
-        return _set_derivatives_(derivative_keys)
-    else:
-        return _set_derivatives_dict(derivative_keys)
+    if isinstance(params, ParamsDict):
+        return _set_derivatives_ParamsDict(params, derivative_keys)
+    return _set_derivatives_(params, derivative_keys)

jinns-1.1.0.dist-info/AUTHORS

@@ -0,0 +1,2 @@
+Hugo Gangloff <hugo.gangloff@inrae.fr>
+Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-1.0.0.dist-info → jinns-1.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 1.0.0
+Version: 1.1.0
 Summary: Physics Informed Neural Network with JAX
 Author-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
 Maintainer-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
+License-File: AUTHORS
 Requires-Dist: numpy
 Requires-Dist: jax
 Requires-Dist: jaxopt

{jinns-1.0.0.dist-info → jinns-1.1.0.dist-info}/RECORD

@@ -5,16 +5,16 @@ jinns/data/__init__.py,sha256=TRCH0Z4-SQZ50MbSf46CUYWBkWVDmXCyez9T-EGiv_8,338
 jinns/experimental/__init__.py,sha256=3jCIy2R2i_0Erwxg-HwISdH79Nt1XCXhS9yY1F5awiY,208
 jinns/experimental/_diffrax_solver.py,sha256=upMr3kTTNrxEiSUO_oLvCXcjS9lPxSjvbB81h3qlhaU,6813
 jinns/loss/_DynamicLoss.py,sha256=WGbAuWnNfsbzUlWEiW_ARd4kI3jmHwdqPjxLC-wCA6s,25753
-jinns/loss/_DynamicLossAbstract.py,sha256=UGFAZWCxbk7yTLdngNasaVawz51xZhu7LfCg4MTqj74,13765
-jinns/loss/_LossODE.py,sha256=4Tcl_Fpn7htoJvggreL3vv5xfQw2VR1tshLGVX7BUx4,22480
-jinns/loss/_LossPDE.py,sha256=zmnLblA6tK-KUKHFbwoo_N01CU1KwCdZt5lUCo-_ecM,48528
+jinns/loss/_DynamicLossAbstract.py,sha256=Xyt28Oej_zlhcV3f6cw2vnAKyRJhXBiA63CsdL3PihU,13767
+jinns/loss/_LossODE.py,sha256=ThWPse6Gn5crM3_tzwZCBx-usoD0xWu6y1n0GVl2dpI,23422
+jinns/loss/_LossPDE.py,sha256=R9kNQiaFbFx2eCMdjB7ie7UJ9pJW7PmvHijioNgu-bs,49117
 jinns/loss/__init__.py,sha256=Fm4QAHaVmp0CA7HSwb7KUctwdXnNZ9v5KmTqpeoYPaE,669
 jinns/loss/_boundary_conditions.py,sha256=O0D8eWsFfvNNeO20PQ0rUKBI_MDqaBvqChfXaztZoL4,16679
-jinns/loss/_loss_utils.py,sha256=DI9KifH1aObrFau86BgAjtBNS8ecSRYB-_H6fduzGNg,13339
+jinns/loss/_loss_utils.py,sha256=44J-VF6dxT_o5BcNWFOiLpY40c35YnAxxZkoNtdtcZc,13689
 jinns/loss/_loss_weights.py,sha256=F0Fgji2XpVk3pr9oIryGuXcG1FGQo4Dv6WFgze2BtA0,2201
 jinns/loss/_operators.py,sha256=o-Ljp_9_HXB9Mhm-ANh6ouNw4_PsqLJAha7dFDGl_nQ,10781
 jinns/parameters/__init__.py,sha256=1gxNLoAXUjhUzBWuh86YjU5pYy8SOboCs8TrKcU1wZc,158
-jinns/parameters/_derivative_keys.py,sha256=MbbyywuemjvgX-QVG2SSqlEreB7bMDCLnXmRZuGFkjE,3294
+jinns/parameters/_derivative_keys.py,sha256=UyEcgfNF1vwPcGWD2ShAZkZiq4thzRDm_OUJzOfjjiY,21909
 jinns/parameters/_params.py,sha256=wK9ZSqoL9KnjOWqc_ZhJ09ffbsgeUEcttc1Rhme0lLk,3550
 jinns/plot/__init__.py,sha256=Q279h5veYWNLQyttsC8_tDOToqUHh8WaRON90CiWXqk,81
 jinns/plot/_plot.py,sha256=ZGIJdGwEd3NlHRTq_2sOfEH_CtOkvPwdgCMct-nQlJE,11691
@@ -31,8 +31,9 @@ jinns/utils/_types.py,sha256=P_dS0odrHbyalYJ0FjS6q0tkXAGr-4GArsiyJYrB1ho,1878
 jinns/utils/_utils.py,sha256=Ow8xB516E7yHDZatokVJHHFNPDu6fXr9-NmraUXjjyw,1819
 jinns/validation/__init__.py,sha256=Jv58mzgC3F7cRfXA6caicL1t_U0UAhbwLrmMNVg6E7s,66
 jinns/validation/_validation.py,sha256=bvqL2poTFJfn9lspWqMqXvQGcQIodKwKrC786QtEZ7A,4700
-jinns-1.0.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-1.0.0.dist-info/METADATA,sha256=XD9Az0b4tUbBNS1xReORk0j1HMkonZGs9_B0FfHFq34,2514
-jinns-1.0.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jinns-1.0.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-1.0.0.dist-info/RECORD,,
+jinns-1.1.0.dist-info/AUTHORS,sha256=7NwCj9nU-HNG1asvy4qhQ2w7oZHrn-Lk5_wK_Ve7a3M,80
+jinns-1.1.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-1.1.0.dist-info/METADATA,sha256=3Qk885oguf6S_WPHd9KCFVIWP21nJqX9zFWoS9ZI-T0,2536
+jinns-1.1.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jinns-1.1.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-1.1.0.dist-info/RECORD,,