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

jinns/nn/_spinn.py

@@ -21,12 +21,12 @@ class SPINN(AbstractPINN):
         used for non-stationnary equations.
     r : int
         An integer. The dimension of the embedding.
-    eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
+    eq_type : Literal["ODE", "PDEStatio", "PDENonStatio"]
         A string with three possibilities.
         "ODE": the PINN is called with one input `t`.
-        "statio_PDE": the PINN is called with one input `x`, `x`
+        "PDEStatio": the PINN is called with one input `x`, `x`
         can be high dimensional.
-        "nonstatio_PDE": the PINN is called with two inputs `t` and `x`, `x`
+        "PDENonStatio": the PINN is called with two inputs `t` and `x`, `x`
         can be high dimensional.
         **Note**: the input dimension as given in eqx_list has to match the sum
         of the dimension of `t` + the dimension of `x`.
@@ -49,7 +49,7 @@ class SPINN(AbstractPINN):
 
     """
 
-    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"] = eqx.field(
+    eq_type: Literal["ODE", "PDEStatio", "PDENonStatio"] = eqx.field(
         static=True, kw_only=True
     )
     d: int = eqx.field(static=True, kw_only=True)

jinns/nn/_spinn_mlp.py

@@ -78,7 +78,7 @@ class SPINN_MLP(SPINN):
         d: int,
         r: int,
         eqx_list: tuple[tuple[Callable, int, int] | tuple[Callable], ...],
-        eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+        eq_type: Literal["ODE", "PDEStatio", "PDENonStatio"],
         m: int = 1,
         filter_spec: PyTree[Union[bool, Callable[[Any], bool]]] = None,
     ) -> tuple[Self, SPINN]:
@@ -114,12 +114,12 @@ class SPINN_MLP(SPINN):
                 (jax.nn.tanh,),
                 (eqx.nn.Linear, 20, r * m)
             )`.
-        eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
+        eq_type : Literal["ODE", "PDEStatio", "PDENonStatio"]
             A string with three possibilities.
             "ODE": the PINN is called with one input `t`.
-            "statio_PDE": the PINN is called with one input `x`, `x`
+            "PDEStatio": the PINN is called with one input `x`, `x`
             can be high dimensional.
-            "nonstatio_PDE": the PINN is called with two inputs `t` and `x`, `x`
+            "PDENonStatio": the PINN is called with two inputs `t` and `x`, `x`
             can be high dimensional.
             **Note**: the input dimension as given in eqx_list has to match the sum
             of the dimension of `t` + the dimension of `x`.
@@ -150,11 +150,11 @@ class SPINN_MLP(SPINN):
         Raises
         ------
         RuntimeError
-            If the parameter value for eq_type is not in `["ODE", "statio_PDE",
-            "nonstatio_PDE"]` and for various failing checks
+            If the parameter value for eq_type is not in `["ODE", "PDEStatio",
+            "PDENonStatio"]` and for various failing checks
         """
 
-        if eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
+        if eq_type not in ["ODE", "PDEStatio", "PDENonStatio"]:
             raise RuntimeError("Wrong parameter value for eq_type")
 
         def element_is_layer(element: tuple) -> TypeGuard[tuple[Callable, int, int]]:

jinns/parameters/_derivative_keys.py

@@ -47,9 +47,9 @@ class DerivativeKeysODE(eqx.Module):
     [`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,
+    To do so, user can either use strings or a `Params[bool]` object
+    with PyTree structure matching the parameters of the problem (`Params[Array]`) at
+    hand, and leaves being 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.
 
@@ -156,12 +156,12 @@ class DerivativeKeysODE(eqx.Module):
         """
         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
+        parameters, the equation parameters or both. The `Params[Array]` 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.
+            You can mix strings and `Params[bool]` if you need granularity.
 
         Parameters
         ----------
@@ -498,7 +498,14 @@ def _set_derivatives(
         `Params(nn_params=True | False, eq_params={"alpha":True | False,
         "beta":True | False})`.
         """
-
+        assert jax.tree.structure(params_.eq_params) == jax.tree.structure(
+            derivative_mask.eq_params
+        ), (
+            "The derivative "
+            "mask for eq_params does not have the same tree structure as "
+            "Params.eq_params. This is often due to a wrong Params[bool] "
+            "passed when initializing the derivative key object."
+        )
         return Params(
             nn_params=jax.lax.cond(
                 derivative_mask.nn_params,

jinns/parameters/_params.py

@@ -2,6 +2,7 @@
 Formalize the data structure for the parameters
 """
 
+from __future__ import annotations
 from dataclasses import fields
 from typing import Generic, TypeVar
 import equinox as eqx
@@ -60,6 +61,15 @@ class Params(eqx.Module, Generic[T]):
         else:
             self.eq_params = eq_params
 
+    def partition(self, mask: Params[bool] | None):
+        """
+        following the boolean mask, partition into two Params
+        """
+        if mask is not None:
+            return eqx.partition(self, mask)
+        else:
+            return self, None
+
 
 def update_eq_params(
     params: Params[Array],

jinns/solver/_rar.py

@@ -10,6 +10,7 @@ from jax import vmap
 import jax.numpy as jnp
 import equinox as eqx
 
+from jinns.loss._DynamicLossAbstract import ODE, PDEStatio, PDENonStatio
 from jinns.data._DataGeneratorODE import DataGeneratorODE
 from jinns.data._CubicMeshPDEStatio import CubicMeshPDEStatio
 from jinns.data._CubicMeshPDENonStatio import CubicMeshPDENonStatio
@@ -176,16 +177,25 @@ def _rar_step_init(
             )
 
             data = eqx.tree_at(lambda m: m.key, data, new_key)
-
-        v_dyn_loss = vmap(
-            lambda inputs: loss.dynamic_loss.evaluate(inputs, loss.u, params),
-        )
-        dyn_on_s = v_dyn_loss(new_samples)
-
-        if dyn_on_s.ndim > 1:
-            mse_on_s = (jnp.linalg.norm(dyn_on_s, axis=-1) ** 2).flatten()
         else:
-            mse_on_s = dyn_on_s**2
+            raise ValueError("Wrong DataGenerator type")
+
+        v_dyn_loss = jax.tree.map(
+            lambda d: vmap(
+                lambda inputs: d.evaluate(inputs, loss.u, params),
+            ),
+            loss.dynamic_loss,
+            is_leaf=lambda x: isinstance(x, (ODE, PDEStatio, PDENonStatio)),
+        )
+        dyn_on_s = jax.tree.map(lambda d: d(new_samples), v_dyn_loss)
+
+        mse_on_s = jax.tree.reduce(
+            jnp.add,
+            jax.tree.map(
+                lambda v: (jnp.linalg.norm(v, axis=-1) ** 2).flatten(), dyn_on_s
+            ),
+            0,
+        )
 
         ## Select the m points with higher dynamic loss
         higher_residual_idx = jax.lax.dynamic_slice(