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

jinns/loss/_operators.py

@@ -18,8 +18,8 @@ from jinns.nn._abstract_pinn import AbstractPINN
 
 def _get_eq_type(
     u: AbstractPINN | Callable[[Array, Params[Array]], Array],
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None,
-) -> Literal["nonstatio_PDE", "statio_PDE"]:
+    eq_type: Literal["PDENonStatio", "PDEStatio"] | None,
+) -> Literal["PDENonStatio", "PDEStatio"]:
     """
     But we filter out ODE from eq_type because we only have operators that does
     not work with ODEs so far
@@ -36,7 +36,7 @@ def divergence_rev(
     inputs: Float[Array, " dim"] | Float[Array, " 1+dim"],
     u: AbstractPINN | Callable[[Array, Params[Array]], Array],
     params: Params[Array],
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+    eq_type: Literal["PDENonStatio", "PDEStatio"] | None = None,
 ) -> Float[Array, " "]:
     r"""
     Compute the divergence of a vector field $\mathbf{u}$, i.e.,
@@ -64,7 +64,7 @@ def divergence_rev(
     eq_type = _get_eq_type(u, eq_type)
 
     def scan_fun(_, i):
-        if eq_type == "nonstatio_PDE":
+        if eq_type == "PDENonStatio":
             du_dxi = grad(lambda inputs, params: u(inputs, params)[1 + i])(
                 inputs, params
             )[1 + i]
@@ -74,9 +74,9 @@ def divergence_rev(
             ]
         return _, du_dxi
 
-    if eq_type == "nonstatio_PDE":
+    if eq_type == "PDENonStatio":
         _, accu = jax.lax.scan(scan_fun, {}, jnp.arange(inputs.shape[0] - 1))
-    elif eq_type == "statio_PDE":
+    elif eq_type == "PDEStatio":
         _, accu = jax.lax.scan(scan_fun, {}, jnp.arange(inputs.shape[0]))
     else:
         raise ValueError("Unexpected u.eq_type!")
@@ -87,7 +87,7 @@ def divergence_fwd(
     inputs: Float[Array, " batch_size dim"] | Float[Array, " batch_size 1+dim"],
     u: AbstractPINN | Callable[[Array, Params[Array]], Array],
     params: Params[Array],
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+    eq_type: Literal["PDENonStatio", "PDEStatio"] | None = None,
 ) -> Float[Array, " batch_size * (1+dim) 1"] | Float[Array, " batch_size * (dim) 1"]:
     r"""
     Compute the divergence of a **batched** vector field $\mathbf{u}$, i.e.,
@@ -120,7 +120,7 @@ def divergence_fwd(
     eq_type = _get_eq_type(u, eq_type)
 
     def scan_fun(_, i):
-        if eq_type == "nonstatio_PDE":
+        if eq_type == "PDENonStatio":
             tangent_vec = jnp.repeat(
                 jax.nn.one_hot(i + 1, inputs.shape[-1])[None],
                 inputs.shape[0],
@@ -140,9 +140,9 @@ def divergence_fwd(
             )
         return _, du_dxi
 
-    if eq_type == "nonstatio_PDE":
+    if eq_type == "PDENonStatio":
         _, accu = jax.lax.scan(scan_fun, {}, jnp.arange(inputs.shape[1] - 1))
-    elif eq_type == "statio_PDE":
+    elif eq_type == "PDEStatio":
         _, accu = jax.lax.scan(scan_fun, {}, jnp.arange(inputs.shape[1]))
     else:
         raise ValueError("Unexpected u.eq_type!")
@@ -154,7 +154,7 @@ def laplacian_rev(
     u: AbstractPINN | Callable[[Array, Params[Array]], Array],
     params: Params[Array],
     method: Literal["trace_hessian_x", "trace_hessian_t_x", "loop"] = "trace_hessian_x",
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+    eq_type: Literal["PDENonStatio", "PDEStatio"] | None = None,
 ) -> Float[Array, " "]:
     r"""
     Compute the Laplacian of a scalar field $u$ from $\mathbb{R}^d$
@@ -196,22 +196,22 @@ def laplacian_rev(
         # computation and then discarding elements but for higher order derivatives
         # it might not be worth it. See other options below for computating the
         # Laplacian
-        if eq_type == "nonstatio_PDE":
+        if eq_type == "PDENonStatio":
             u_ = lambda x: jnp.squeeze(
                 u(jnp.concatenate([inputs[:1], x], axis=0), params)
             )
             return jnp.sum(jnp.diag(jax.hessian(u_)(inputs[1:])))
-        if eq_type == "statio_PDE":
+        if eq_type == "PDEStatio":
             u_ = lambda inputs: jnp.squeeze(u(inputs, params))
             return jnp.sum(jnp.diag(jax.hessian(u_)(inputs)))
         raise ValueError("Unexpected eq_type!")
     if method == "trace_hessian_t_x":
         # NOTE that it is unclear whether it is better to vectorially compute the
         # Hessian (despite a useless time dimension) as below
-        if eq_type == "nonstatio_PDE":
+        if eq_type == "PDENonStatio":
             u_ = lambda inputs: jnp.squeeze(u(inputs, params))
             return jnp.sum(jnp.diag(jax.hessian(u_)(inputs))[1:])
-        if eq_type == "statio_PDE":
+        if eq_type == "PDEStatio":
             u_ = lambda inputs: jnp.squeeze(u(inputs, params))
             return jnp.sum(jnp.diag(jax.hessian(u_)(inputs)))
         raise ValueError("Unexpected eq_type!")
@@ -225,7 +225,7 @@ def laplacian_rev(
         u_ = lambda inputs: u(inputs, params).squeeze()
 
         def scan_fun(_, i):
-            if eq_type == "nonstatio_PDE":
+            if eq_type == "PDENonStatio":
                 d2u_dxi2 = grad(
                     lambda inputs: grad(u_)(inputs)[1 + i],
                 )(inputs)[1 + i]
@@ -236,11 +236,11 @@ def laplacian_rev(
                 )(inputs)[i]
             return _, d2u_dxi2
 
-        if eq_type == "nonstatio_PDE":
+        if eq_type == "PDENonStatio":
             _, trace_hessian = jax.lax.scan(
                 scan_fun, {}, jnp.arange(inputs.shape[0] - 1)
             )
-        elif eq_type == "statio_PDE":
+        elif eq_type == "PDEStatio":
             _, trace_hessian = jax.lax.scan(scan_fun, {}, jnp.arange(inputs.shape[0]))
         else:
             raise ValueError("Unexpected eq_type!")
@@ -253,7 +253,7 @@ def laplacian_fwd(
     u: AbstractPINN | Callable[[Array, Params[Array]], Array],
     params: Params[Array],
     method: Literal["trace_hessian_t_x", "trace_hessian_x", "loop"] = "loop",
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+    eq_type: Literal["PDENonStatio", "PDEStatio"] | None = None,
 ) -> Float[Array, " batch_size * (1+dim) 1"] | Float[Array, " batch_size * (dim) 1"]:
     r"""
     Compute the Laplacian of a **batched** scalar field $u$
@@ -302,7 +302,7 @@ def laplacian_fwd(
     if method == "loop":
 
         def scan_fun(_, i):
-            if eq_type == "nonstatio_PDE":
+            if eq_type == "PDENonStatio":
                 tangent_vec = jnp.repeat(
                     jax.nn.one_hot(i + 1, inputs.shape[-1])[None],
                     inputs.shape[0],
@@ -323,17 +323,17 @@ def laplacian_fwd(
             __, d2u_dxi2 = jax.jvp(du_dxi_fun, (inputs,), (tangent_vec,))
             return _, d2u_dxi2
 
-        if eq_type == "nonstatio_PDE":
+        if eq_type == "PDENonStatio":
             _, trace_hessian = jax.lax.scan(
                 scan_fun, {}, jnp.arange(inputs.shape[-1] - 1)
             )
-        elif eq_type == "statio_PDE":
+        elif eq_type == "PDEStatio":
             _, trace_hessian = jax.lax.scan(scan_fun, {}, jnp.arange(inputs.shape[-1]))
         else:
             raise ValueError("Unexpected eq_type!")
         return jnp.sum(trace_hessian, axis=0)
     if method == "trace_hessian_t_x":
-        if eq_type == "nonstatio_PDE":
+        if eq_type == "PDENonStatio":
             # compute the Hessian including the batch dimension, get rid of the
             # (..,1,..) axis that is here because of the scalar output
             # if inputs.shape==(10,3) (1 for time, 2 for x_dim)
@@ -351,7 +351,7 @@ def laplacian_fwd(
             res_dims = "".join([f"{chr(97 + d)}" for d in range(inputs.shape[-1])])
             lap = jnp.einsum(res_dims + "ii->" + res_dims, r)
             return lap[..., None]
-        if eq_type == "statio_PDE":
+        if eq_type == "PDEStatio":
             # compute the Hessian including the batch dimension, get rid of the
             # (..,1,..) axis that is here because of the scalar output
             # if inputs.shape==(10,2), r.shape=(10,10,1,10,2,10,2)
@@ -369,7 +369,7 @@ def laplacian_fwd(
             return lap[..., None]
         raise ValueError("Unexpected eq_type!")
     if method == "trace_hessian_x":
-        if eq_type == "statio_PDE":
+        if eq_type == "PDEStatio":
             # compute the Hessian including the batch dimension, get rid of the
             # (..,1,..) axis that is here because of the scalar output
             # if inputs.shape==(10,2), r.shape=(10,10,1,10,2,10,2)
@@ -394,7 +394,7 @@ def vectorial_laplacian_rev(
     u: AbstractPINN | Callable[[Array, Params[Array]], Array],
     params: Params[Array],
     dim_out: int | None = None,
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+    eq_type: Literal["PDENonStatio", "PDEStatio"] | None = None,
 ) -> Float[Array, " dim_out"]:
     r"""
     Compute the vectorial Laplacian of a vector field $\mathbf{u}$ from
@@ -448,7 +448,7 @@ def vectorial_laplacian_fwd(
     u: AbstractPINN | Callable[[Array, Params[Array]], Array],
     params: Params[Array],
     dim_out: int | None = None,
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+    eq_type: Literal["PDENonStatio", "PDEStatio"] | None = None,
 ) -> Float[Array, " batch_size * (1+dim) n"] | Float[Array, " batch_size * (dim) n"]:
     r"""
     Compute the vectorial Laplacian of a vector field $\mathbf{u}$ when

jinns/nn/_abstract_pinn.py

@@ -13,7 +13,7 @@ class AbstractPINN(eqx.Module):
     https://github.com/patrick-kidger/equinox/issues/1002 + https://docs.kidger.site/equinox/pattern/
     """
 
-    eq_type: eqx.AbstractVar[Literal["ODE", "statio_PDE", "nonstatio_PDE"]]
+    eq_type: eqx.AbstractVar[Literal["ODE", "PDEStatio", "PDENonStatio"]]
 
     @abc.abstractmethod
     def __call__(self, inputs: Any, params: Params[Array], *args, **kwargs) -> Any:

jinns/nn/_hyperpinn.py

@@ -67,9 +67,9 @@ class HyperPINN(PINN):
     eq_type : str
         A string with three possibilities.
         "ODE": the HyperPINN is called with one input `t`.
-        "statio_PDE": the HyperPINN is called with one input `x`, `x`
+        "PDEStatio": the HyperPINN is called with one input `x`, `x`
         can be high dimensional.
-        "nonstatio_PDE": the HyperPINN is called with two inputs `t` and `x`, `x`
+        "PDENonStatio": the HyperPINN 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` or the output dimension
@@ -192,7 +192,7 @@ class HyperPINN(PINN):
         hyper = eqx.combine(params.nn_params, self.static_hyper)
 
         eq_params_batch = jnp.concatenate(
-            [getattr(params.eq_params, k).flatten() for k in self.hyperparams],
+            [getattr(params.eq_params, k).flatten() for k in self.hyperparams],  # pylint: disable=E1133
             axis=0,
         )
 
@@ -214,7 +214,7 @@ class HyperPINN(PINN):
     def create(
         cls,
         *,
-        eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+        eq_type: Literal["ODE", "PDEStatio", "PDENonStatio"],
         hyperparams: list[str],
         hypernet_input_size: int,
         key: PRNGKeyArray | None = None,
@@ -257,9 +257,9 @@ class HyperPINN(PINN):
         eq_type
             A string with three possibilities.
             "ODE": the HyperPINN is called with one input `t`.
-            "statio_PDE": the HyperPINN is called with one input `x`, `x`
+            "PDEStatio": the HyperPINN is called with one input `x`, `x`
             can be high dimensional.
-            "nonstatio_PDE": the HyperPINN is called with two inputs `t` and `x`, `x`
+            "PDENonStatio": the HyperPINN 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` or the output dimension

jinns/nn/_mlp.py

@@ -95,7 +95,7 @@ class PINN_MLP(PINN):
     def create(
         cls,
         *,
-        eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+        eq_type: Literal["ODE", "PDEStatio", "PDENonStatio"],
         key: PRNGKeyArray | None = None,
         eqx_network: eqx.nn.MLP | MLP | None = None,
         eqx_list: tuple[tuple[Callable, int, int] | tuple[Callable], ...] | None = None,
@@ -130,9 +130,9 @@ class PINN_MLP(PINN):
         eq_type
             A string with three possibilities.
             "ODE": the MLP is called with one input `t`.
-            "statio_PDE": the MLP is called with one input `x`, `x`
+            "PDEStatio": the MLP is called with one input `x`, `x`
             can be high dimensional.
-            "nonstatio_PDE": the MLP is called with two inputs `t` and `x`, `x`
+            "PDENonStatio": the MLP 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` or the output dimension

jinns/nn/_pinn.py

@@ -50,12 +50,12 @@ class PINN(AbstractPINN):
         when the PINN is also used to output equation parameters for example
         Note that it must be a slice and not an integer (a preprocessing of the
         user provided argument takes care of it).
-    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` or the output dimension
@@ -83,11 +83,11 @@ class PINN(AbstractPINN):
     Raises
     ------
     RuntimeError
-        If the parameter value for eq_type is not in `["ODE", "statio_PDE",
-        "nonstatio_PDE"]`
+        If the parameter value for eq_type is not in `["ODE", "PDEStatio",
+        "PDENonStatio"]`
     """
 
-    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"] = eqx.field(
+    eq_type: Literal["ODE", "PDEStatio", "PDENonStatio"] = eqx.field(
         static=True, kw_only=True
     )
     slice_solution: slice = eqx.field(static=True, kw_only=True, default=None)
@@ -108,7 +108,7 @@ class PINN(AbstractPINN):
     static: PINN = eqx.field(init=False, static=True)
 
     def __post_init__(self, eqx_network):
-        if self.eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
+        if self.eq_type not in ["ODE", "PDEStatio", "PDENonStatio"]:
             raise RuntimeError("Wrong parameter value for eq_type")
         # saving the static part of the model and initial parameters
 

jinns/nn/_ppinn.py

@@ -31,12 +31,12 @@ class PPINN_MLP(PINN):
         when the PINN is also used to output equation parameters for example
         Note that it must be a slice and not an integer (a preprocessing of the
         user provided argument takes care of it).
-    eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
+    eq_type : Literal["ODE", "PDEStatio", "PDENonStatio"]
         A string with three possibilities.
         "ODE": the PPINN is called with one input `t`.
-        "statio_PDE": the PPINN is called with one input `x`, `x`
+        "PDEStatio": the PPINN is called with one input `x`, `x`
         can be high dimensional.
-        "nonstatio_PDE": the PPINN is called with two inputs `t` and `x`, `x`
+        "PDENonStatio": the PPINN 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` or the output dimension
@@ -125,7 +125,7 @@ class PPINN_MLP(PINN):
         cls,
         *,
         key: PRNGKeyArray | None = None,
-        eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+        eq_type: Literal["ODE", "PDEStatio", "PDENonStatio"],
         eqx_network_list: list[eqx.nn.MLP | MLP] | None = None,
         eqx_list_list: (
             list[tuple[tuple[Callable, int, int] | tuple[Callable], ...]] | None
@@ -158,9 +158,9 @@ class PPINN_MLP(PINN):
         eq_type
             A string with three possibilities.
             "ODE": the PPINN MLP is called with one input `t`.
-            "statio_PDE": the PPINN MLP is called with one input `x`, `x`
+            "PDEStatio": the PPINN MLP is called with one input `x`, `x`
             can be high dimensional.
-            "nonstatio_PDE": the PPINN MLP is called with two inputs `t` and `x`, `x`
+            "PDENonStatio": the PPINN MLP 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` or the output dimension

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/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(

jinns/solver/_solve.py

@@ -272,6 +272,7 @@ def solve(
         params=init_params,
         last_non_nan_params=init_params,
         opt_state=opt_state,
+        # params_mask=params_mask,
     )
     optimization_extra = OptimizationExtraContainer(
         curr_seq=curr_seq,
@@ -430,7 +431,9 @@ def solve(
         return (
             i,
             loss,
-            OptimizationContainer(params, last_non_nan_params, opt_state),
+            OptimizationContainer(
+                params, last_non_nan_params, opt_state
+            ),  # , params_mask),
             OptimizationExtraContainer(
                 curr_seq,
                 best_iter_id,

jinns/solver/_utils.py

@@ -606,18 +606,24 @@ def _check_batch_size(other_data, main_data, attr_name):
                         " vectorization"
                     )
     if isinstance(main_data, DataGeneratorParameter):
-        if main_data.param_batch_size is not None:
-            if getattr(other_data, attr_name) != main_data.param_batch_size:
-                raise ValueError(
-                    f"{other_data.__class__}.{attr_name} must be equal"
-                    f" to {main_data.__class__}.param_batch_size for correct"
-                    " vectorization"
-                )
-        else:
-            if main_data.n is not None:
-                if getattr(other_data, attr_name) != main_data.n:
+        batch_size = getattr(other_data, attr_name)  # this can be a tuple with
+        # DataGeneratorObservations
+        if not isinstance(batch_size, tuple):
+            batch_size = (batch_size,)
+
+        for bs in batch_size:
+            if main_data.param_batch_size is not None:
+                if bs != main_data.param_batch_size:
                     raise ValueError(
                         f"{other_data.__class__}.{attr_name} must be equal"
-                        f" to {main_data.__class__}.n for correct"
+                        f" to {main_data.__class__}.param_batch_size for correct"
                         " vectorization"
                     )
+            else:
+                if main_data.n is not None:
+                    if bs != main_data.n:
+                        raise ValueError(
+                            f"{other_data.__class__}.{attr_name} must be equal"
+                            f" to {main_data.__class__}.n for correct"
+                            " vectorization"
+                        )

{jinns-1.7.0.dist-info → jinns-1.7.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jinns
-Version: 1.7.0
+Version: 1.7.1
 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>
@@ -20,14 +20,24 @@ Requires-Dist: optax>=0.2.6
 Requires-Dist: equinox>=0.13.2
 Requires-Dist: matplotlib
 Requires-Dist: jaxtyping
+Requires-Dist: pyright>=1.1.407
+Requires-Dist: pytest>=9.0.2
 Provides-Extra: notebook
 Requires-Dist: jupyter; extra == "notebook"
 Requires-Dist: seaborn; extra == "notebook"
 Requires-Dist: pandas; extra == "notebook"
-Requires-Dist: pytest; extra == "notebook"
-Requires-Dist: pre-commit; extra == "notebook"
-Requires-Dist: pyright; extra == "notebook"
 Requires-Dist: diffrax; extra == "notebook"
+Provides-Extra: gpu-cuda13
+Requires-Dist: jax[cuda13]>=0.8.1; extra == "gpu-cuda13"
+Provides-Extra: gpu-cuda12
+Requires-Dist: jax[cuda12]>=0.8.1; extra == "gpu-cuda12"
+Provides-Extra: cpu
+Requires-Dist: jax>=0.8.1; extra == "cpu"
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pylint; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
 Dynamic: license-file
 
 jinns

{jinns-1.7.0.dist-info → jinns-1.7.1.dist-info}/RECORD

@@ -1,36 +1,36 @@
 jinns/__init__.py,sha256=f8ZCZH95C7U8NfJ3M-00kf1IBbdgbwRNBhlluiOCGJE,530
 jinns/data/_AbstractDataGenerator.py,sha256=le5plNhOE7hV72SC6p2xhWxljnGeStPc2kitzeMR_ts,535
-jinns/data/_Batchs.py,sha256=UZfJWIKAFLQJxOtWN1ybJdoSQCO6PWk3SAgV9YmNVnI,2809
+jinns/data/_Batchs.py,sha256=mtHDlsen1vBBmFgGzynPmFfmBFC7zKcnnVEg_JnNlig,2857
 jinns/data/_CubicMeshPDENonStatio.py,sha256=dtev6j7HuVZ_IBGnKVKOlJTyZVjnW_8G1VRGpcL40VU,22922
 jinns/data/_CubicMeshPDEStatio.py,sha256=g4H9wI0FkA9sdnlSxp00QArq9IjhPrmsGgSleLdNpI4,23314
-jinns/data/_DataGeneratorODE.py,sha256=XOGeX9m5J9CkkGx2alJ5G7FZYGn5cBUaFpO-3-MKvdc,7724
-jinns/data/_DataGeneratorObservations.py,sha256=h2bS-oOUi1rYvhDHy8xT-N9voqtb_IOelL2asnXCybw,9765
+jinns/data/_DataGeneratorODE.py,sha256=7m3mhMk-o02Jq_ZieWktFG5w6mVwnWPs1XYiYqoJ0DA,7731
+jinns/data/_DataGeneratorObservations.py,sha256=Lw1Fno4UDJ0LRkhFJ0tDtkkW-yWJLeswCWja504pUDo,27781
 jinns/data/_DataGeneratorParameter.py,sha256=m7WCWBtmNptF8qQSWGPTQEFVzkJb8kJpidKAFBmoBiQ,10205
 jinns/data/__init__.py,sha256=DEzEmoD5TmjHCaG6pS2jmFDzuhZn1ZDpFdVa2v0jCe0,703
 jinns/data/_utils.py,sha256=xKHf7NksJ_AmrtEcpJsh7WSEvI3Yk98_cM5kmXSfmx0,5596
 jinns/experimental/__init__.py,sha256=DT9e57zbjfzPeRnXemGUqnGd--MhV77FspChT0z4YrE,410
 jinns/experimental/_diffrax_solver.py,sha256=upMr3kTTNrxEiSUO_oLvCXcjS9lPxSjvbB81h3qlhaU,6813
 jinns/loss/_DynamicLoss.py,sha256=FXVraH97tIWA5ByO-V-r54Nhwd_EvHOfr0e2qjrfdf0,22667
-jinns/loss/_DynamicLossAbstract.py,sha256=bTiDZHqkHDdvfj6hcjPcESYzwzqOd2LtR_h7GPAqQpY,13602
-jinns/loss/_LossODE.py,sha256=c2SBstJgy1yTTGcLZrnrM3XFYS67AZ8SwbcPj0cdfXE,15894
-jinns/loss/_LossPDE.py,sha256=uhOGJ9440WaCaLasMbqp8pxf9obXYbvFWbqJMYuxJUU,39719
+jinns/loss/_DynamicLossAbstract.py,sha256=mm687ly-r7ONZPJaQN6kQa9CwM05msiZt9cCkQ4Ue54,13699
+jinns/loss/_LossODE.py,sha256=mbsK6n61J3Fx4zl1VvJCLigowV-UPxIJoxn-Zf_IwH0,18039
+jinns/loss/_LossPDE.py,sha256=rjy5mKEkFOrTqpRf7W51JEVnIysnD7RhQ33aAvZMcfg,43056
 jinns/loss/__init__.py,sha256=z5xYgBipNFf66__5BqQc6R_8r4F6A3TXL60YjsM8Osk,1287
-jinns/loss/_abstract_loss.py,sha256=wBpB2mw_JizimiNHYIs-431jZ87mxvzamTFkX_rnQZ4,6872
-jinns/loss/_boundary_conditions.py,sha256=9HGw1cGLfmEilP4V4B2T0zl0YP1kNtrtXVLQNiBmWgc,12464
+jinns/loss/_abstract_loss.py,sha256=MG-nC_ApYNhnTGLvW1_vx8VXXJLJ2q3Ird9sbUZN5fk,9000
+jinns/loss/_boundary_conditions.py,sha256=WqPAIcG5DbY-K_sYYvi_cr4TPoi0ajIGvnmz9KSw3IE,12458
 jinns/loss/_loss_components.py,sha256=Y5GGNuAS_tHl-AzIMBxf3yKj0N2QCdVdTY4jC7NDuic,442
-jinns/loss/_loss_utils.py,sha256=eJ4JcBm396LHx7Tti88ZQrLcKqVL1oSfFGT23VNkytQ,11949
-jinns/loss/_loss_weight_updates.py,sha256=Q5RJkulMjP5G3twKEGmc46QeQJSCb7yyS7OudzRqQkY,6775
-jinns/loss/_loss_weights.py,sha256=x0URuCaFXULlVDwlL9ZEy45L33tj1Fa0VJUQlSPhGkA,2416
-jinns/loss/_operators.py,sha256=Ds5yRH7hu-jaGRp7PYbt821BgYuEvgWHufWhYgdMjw0,22909
+jinns/loss/_loss_utils.py,sha256=thnV4rA77tH-exwe5Y2dbPOTDAmw44sQJmYyRGe8cRY,12007
+jinns/loss/_loss_weight_updates.py,sha256=WaQZl4TE1eXybLGjMXaXRQhUtA_7Cnyc1r5vTyr4_98,7683
+jinns/loss/_loss_weights.py,sha256=mPfIXdCoTAcnxsgNK_KjMTBUbk_Q3m4BmFF229ghEpQ,2640
+jinns/loss/_operators.py,sha256=RsBJzvvADB1uhJQSbvg1e90KPw8csQR5GKIidu5Xsp8,22874
 jinns/nn/__init__.py,sha256=gwE48oqB_FsSIE-hUvCLz0jPaqX350LBxzH6ueFWYk4,456
-jinns/nn/_abstract_pinn.py,sha256=E17TFjxgs63BHl3dUKwajqkBuggVbyNVtovX0ePPzr4,602
-jinns/nn/_hyperpinn.py,sha256=x_hBWqokBbURMOdIG_0EBNGJjxAzjlUZ4__fyb-uasA,20245
-jinns/nn/_mlp.py,sha256=ajVWRTtGVnmRjuaylEMlsK-3H5LMqKDvhUsMNnH_rIM,8935
-jinns/nn/_pinn.py,sha256=nRBIdjTir1ViI9hBO4_4DY6uS-1P8Xv-TDLU-O_qzTk,8332
-jinns/nn/_ppinn.py,sha256=qBMenzj0TpapGxyl7nSWb6mER-Ny42JrcOd7Wld19hU,10144
+jinns/nn/_abstract_pinn.py,sha256=jp1Wi2LDBvSkJempliQq4m7OaIbYr8d16elnJX8movk,600
+jinns/nn/_hyperpinn.py,sha256=A4nNnID7XVsU7kDafIZo7OchRx9EmoSHseoH4EmY0m8,20264
+jinns/nn/_mlp.py,sha256=VqObd7gLGIO25N4XVfHK_XwvhR2dgeTR-kOcRLTH0Os,8931
+jinns/nn/_pinn.py,sha256=f5OyiZCgyq_wubI10rNYSlY00letMsXYDSHcISuvvwM,8322
+jinns/nn/_ppinn.py,sha256=wvzeT8CyYeiQ2q6BTvNTz3L7g7FPemOQcPIGVDy34Rs,10136
 jinns/nn/_save_load.py,sha256=4pGDPZPKzDmY5UzRtA_zxpDK4hFL2LuDxexWW2PnHtw,7656
-jinns/nn/_spinn.py,sha256=a0ZsnzyXi-oNd6i8pAwF1tdjJQU51pbrGg608O6ykQw,4352
-jinns/nn/_spinn_mlp.py,sha256=7w-WsdcoLHVWwkK2DtXRKRtSrcJkEXJ2_n27VmXI2j4,7651
+jinns/nn/_spinn.py,sha256=MGaPQRA-qVltshESB7IP_kft0kINjijPhEV5kBfKDvI,4346
+jinns/nn/_spinn_mlp.py,sha256=NWSNqVOUPz-Zk2JNYZKbmcYWtXHDrxN2aIYiJH4jK1k,7641
 jinns/nn/_utils.py,sha256=SoaYvNM6V7M3GcXK3diqO7B9_zoBXqSux2fsTOcouSA,1018
 jinns/parameters/__init__.py,sha256=T_BhKT77X5nei2pYj-QIOKCyD6j2UCgSgAY3LnNdkCo,333
 jinns/parameters/_derivative_keys.py,sha256=laMZoT7ZZa3ZSTuvA-TKM6-VX32vKZFfA5ThDWhj4v4,21084
@@ -38,10 +38,10 @@ jinns/parameters/_params.py,sha256=A6F5HAbxaKPonsH5IrT3ERNSy60OGLUD0eV9hYNj28U,4
 jinns/plot/__init__.py,sha256=KPHX0Um4FbciZO1yD8kjZbkaT8tT964Y6SE2xCQ4eDU,135
 jinns/plot/_plot.py,sha256=-A5auNeElaz2_8UzVQJQE4143ZFg0zgMjStU7kwttEY,11565
 jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jinns/solver/_rar.py,sha256=vSVTnCGCusI1vTZCvIkP2_G8we44G_42yZHx2sOK9DE,10291
-jinns/solver/_solve.py,sha256=QX1vBec9pk9UZdjA4oqcqWNnvwb8p10Qzm24szwoZJs,20046
+jinns/solver/_rar.py,sha256=IFaH0-b-uPp7fs1rX14LY2oy3cXH2NN6QqjC4rdm6n8,10659
+jinns/solver/_solve.py,sha256=3JDlX19o3UCpIfrTyFLgWh3ihAwGbl1oB0NsdbBke30,20130
 jinns/solver/_solve_alternate.py,sha256=9R1idmQMhkLz_QsxwjsfWcS_f_Cvg71NS43MfNbpXSo,31664
-jinns/solver/_utils.py,sha256=dNr-XipxSy3Vxo_DqudMKUHNgc6rzvZ_efTy--GnesQ,22766
+jinns/solver/_utils.py,sha256=UMq7jhrShVeVUbcJqvjNdLvczDOP1x6XolF0jWjNs0A,23002
 jinns/utils/_DictToModuleMeta.py,sha256=kdCAZ6U7L_ZNsGnvTfk3r-85RaV19Qq6y6zdhvI_hGM,3043
 jinns/utils/_ItemizableModule.py,sha256=KdjsisFddGOUTA7-1jXlrkFxnX5sIVdtyNDxOaEcEYw,634
 jinns/utils/__init__.py,sha256=k2kWDahXsLhSlsbeB_W73WLaVSvPG_S6mrNG_5fW_Mo,121
@@ -50,9 +50,9 @@ jinns/utils/_types.py,sha256=PftmBr-hV33U-aB_TzTYNLD42ogRWmZKA02CPdxQWBg,2482
 jinns/utils/_utils.py,sha256=M7NXX9ok-BkH5o_xo74PB1_Cc8XiDipSl51rq82dTH4,2821
 jinns/validation/__init__.py,sha256=FTyUO-v1b8Tv-FDSQsntrH7zl9E0ENexqKMT_dFRkYo,124
 jinns/validation/_validation.py,sha256=8p6sMKiBAvA6JNm65hjkMj0997LJ0BkyCREEh0AnPVE,4803
-jinns-1.7.0.dist-info/licenses/AUTHORS,sha256=7NwCj9nU-HNG1asvy4qhQ2w7oZHrn-Lk5_wK_Ve7a3M,80
-jinns-1.7.0.dist-info/licenses/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-1.7.0.dist-info/METADATA,sha256=nv5KCuvdeLV2I3kWjckhMCsKe7KwRJ3NoPgiwmK1WBk,5651
-jinns-1.7.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-jinns-1.7.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-1.7.0.dist-info/RECORD,,
+jinns-1.7.1.dist-info/licenses/AUTHORS,sha256=7NwCj9nU-HNG1asvy4qhQ2w7oZHrn-Lk5_wK_Ve7a3M,80
+jinns-1.7.1.dist-info/licenses/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-1.7.1.dist-info/METADATA,sha256=BuWK-l7JNmKFJ9silTxgpeJQqNgX3JRLZfo0WgmGb4Y,5982
+jinns-1.7.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+jinns-1.7.1.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-1.7.1.dist-info/RECORD,,