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

jinns/loss/_abstract_loss.py

@@ -1,15 +1,27 @@
 from __future__ import annotations
 
 import abc
-from typing import Self, Literal, Callable, TypeVar, Generic, Any
-from jaxtyping import PRNGKeyArray, Array, PyTree, Float
+import warnings
+from typing import Self, Literal, Callable, TypeVar, Generic, Any, get_args
+from dataclasses import InitVar
+from jaxtyping import Array, PyTree, Float, PRNGKeyArray
 import equinox as eqx
 import jax
 import jax.numpy as jnp
 import optax
 from jinns.parameters._params import Params
-from jinns.loss._loss_weight_updates import soft_adapt, lr_annealing, ReLoBRaLo
-from jinns.utils._types import AnyLossComponents, AnyBatch, AnyLossWeights
+from jinns.loss._loss_weight_updates import (
+    soft_adapt,
+    lr_annealing,
+    ReLoBRaLo,
+    prior_loss,
+)
+from jinns.utils._types import (
+    AnyLossComponents,
+    AnyBatch,
+    AnyLossWeights,
+    AnyDerivativeKeys,
+)
 
 L = TypeVar(
     "L", bound=AnyLossWeights
@@ -25,31 +37,85 @@ C = TypeVar(
     "C", bound=AnyLossComponents[Array | None]
 )  # The above comment also works with Unions (https://docs.python.org/3/library/typing.html#typing.TypeVar)
 
+DK = TypeVar("DK", bound=AnyDerivativeKeys)
+
 # In the cases above, without the bound, we could not have covariance on
 # the type because it would break LSP. Note that covariance on the return type
 # is authorized in LSP hence we do not need the same TypeVar instruction for
 # the return types of evaluate_by_terms for example!
 
 
-class AbstractLoss(eqx.Module, Generic[L, B, C]):
+AvailableUpdateWeightMethods = Literal[
+    "softadapt", "soft_adapt", "prior_loss", "lr_annealing", "ReLoBRaLo"
+]
+
+
+class AbstractLoss(eqx.Module, Generic[L, B, C, DK]):
     """
     About the call:
     https://github.com/patrick-kidger/equinox/issues/1002 + https://docs.kidger.site/equinox/pattern/
     """
 
+    derivative_keys: eqx.AbstractVar[DK]
     loss_weights: eqx.AbstractVar[L]
-    update_weight_method: Literal["soft_adapt", "lr_annealing", "ReLoBRaLo"] | None = (
-        eqx.field(kw_only=True, default=None, static=True)
+    loss_weight_scales: L = eqx.field(init=False)
+    update_weight_method: AvailableUpdateWeightMethods | None = eqx.field(
+        kw_only=True, default=None, static=True
+    )
+    vmap_in_axes: tuple[int] = eqx.field(static=True)
+    keep_initial_loss_weight_scales: InitVar[bool] = eqx.field(
+        default=True, kw_only=True
     )
 
+    def __init__(
+        self,
+        *,
+        loss_weights,
+        derivative_keys,
+        vmap_in_axes,
+        update_weight_method=None,
+        keep_initial_loss_weight_scales: bool = False,
+    ):
+        if update_weight_method is not None and update_weight_method not in get_args(
+            AvailableUpdateWeightMethods
+        ):
+            raise ValueError(f"{update_weight_method=} is not a valid method")
+        self.update_weight_method = update_weight_method
+        self.loss_weights = loss_weights
+        self.derivative_keys = derivative_keys
+        self.vmap_in_axes = vmap_in_axes
+        if keep_initial_loss_weight_scales:
+            self.loss_weight_scales = self.loss_weights
+            if self.update_weight_method is not None:
+                warnings.warn(
+                    "Loss weights out from update_weight_method will still be"
+                    " multiplied by the initial input loss_weights"
+                )
+        else:
+            self.loss_weight_scales = optax.tree_utils.tree_ones_like(self.loss_weights)
+            # self.loss_weight_scales will contain None where self.loss_weights
+            # has None
+
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         return self.evaluate(*args, **kwargs)
 
     @abc.abstractmethod
-    def evaluate_by_terms(self, params: Params[Array], batch: B) -> tuple[C, C]:
+    def evaluate_by_terms(
+        self,
+        opt_params: Params[Array],
+        batch: B,
+        *,
+        non_opt_params: Params[Array] | None = None,
+    ) -> tuple[C, C]:
         pass
 
-    def evaluate(self, params: Params[Array], batch: B) -> tuple[Float[Array, " "], C]:
+    def evaluate(
+        self,
+        opt_params: Params[Array],
+        batch: B,
+        *,
+        non_opt_params: Params[Array] | None = None,
+    ) -> tuple[Float[Array, " "], C]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
@@ -57,16 +123,20 @@ class AbstractLoss(eqx.Module, Generic[L, B, C]):
 
         Parameters
         ---------
-        params
-            Parameters at which the loss is evaluated
+        opt_params
+            Parameters, which are optimized, at which the loss is evaluated
         batch
             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
+        non_opt_params
+            Parameters, which are non optimized, at which the loss is evaluated
         """
-        loss_terms, _ = self.evaluate_by_terms(params, batch)
+        loss_terms, _ = self.evaluate_by_terms(
+            opt_params, batch, non_opt_params=non_opt_params
+        )
 
         loss_val = self.ponderate_and_sum_loss(loss_terms)
 
@@ -102,10 +172,14 @@ class AbstractLoss(eqx.Module, Generic[L, B, C]):
         raise ValueError(
             "The numbers of declared loss weights and "
             "declared loss terms do not concord "
-            f" got {len(weights)} and {len(terms_list)}"
+            f" got {len(weights)} and {len(terms_list)}. "
+            "If you passed tuple of dyn_loss, make sure to pass "
+            "tuple of loss weights at LossWeights.dyn_loss."
+            "If you passed tuple of obs datasets, make sure to pass "
+            "tuple of loss weights at LossWeights.observations."
         )
 
-    def ponderate_and_sum_gradient(self, terms: C) -> C:
+    def ponderate_and_sum_gradient(self, terms: C) -> Params[Array | None]:
         """
         Get total gradients from individual loss gradients and weights
         for each parameter
@@ -146,6 +220,8 @@ class AbstractLoss(eqx.Module, Generic[L, B, C]):
             new_weights = soft_adapt(
                 self.loss_weights, iteration_nb, loss_terms, stored_loss_terms
             )
+        elif self.update_weight_method == "prior_loss":
+            new_weights = prior_loss(self.loss_weights, iteration_nb, stored_loss_terms)
         elif self.update_weight_method == "lr_annealing":
             new_weights = lr_annealing(self.loss_weights, grad_terms)
         elif self.update_weight_method == "ReLoBRaLo":
@@ -158,6 +234,13 @@ class AbstractLoss(eqx.Module, Generic[L, B, C]):
         # Below we update the non None entry in the PyTree self.loss_weights
         # we directly get the non None entries because None is not treated as a
         # leaf
+
+        new_weights = jax.lax.cond(
+            iteration_nb == 0,
+            lambda nw: nw,
+            lambda nw: jnp.array(jax.tree.leaves(self.loss_weight_scales)) * nw,
+            new_weights,
+        )
         return eqx.tree_at(
             lambda pt: jax.tree.leaves(pt.loss_weights), self, new_weights
         )

jinns/loss/_boundary_conditions.py

@@ -227,7 +227,7 @@ def boundary_neumann(
     if isinstance(u, PINN):
         u_ = lambda inputs, params: jnp.squeeze(u(inputs, params)[dim_to_apply])
 
-        if u.eq_type == "statio_PDE":
+        if u.eq_type == "PDEStatio":
             v_neumann = vmap(
                 lambda inputs, params: _subtract_with_check(
                     f(inputs),
@@ -240,7 +240,7 @@ def boundary_neumann(
                 vmap_in_axes,
                 0,
             )
-        elif u.eq_type == "nonstatio_PDE":
+        elif u.eq_type == "PDENonStatio":
             v_neumann = vmap(
                 lambda inputs, params: _subtract_with_check(
                     f(inputs),
@@ -274,14 +274,14 @@ def boundary_neumann(
         if (batch_array.shape[0] == 1 and isinstance(batch, PDEStatioBatch)) or (
             batch_array.shape[-1] == 2 and isinstance(batch, PDENonStatioBatch)
         ):
-            if u.eq_type == "statio_PDE":
+            if u.eq_type == "PDEStatio":
                 _, du_dx = jax.jvp(
                     lambda inputs: u(inputs, params)[..., dim_to_apply],
                     (batch_array,),
                     (jnp.ones_like(batch_array),),
                 )
                 values = du_dx * n[facet]
-            if u.eq_type == "nonstatio_PDE":
+            if u.eq_type == "PDENonStatio":
                 _, du_dx = jax.jvp(
                     lambda inputs: u(inputs, params)[..., dim_to_apply],
                     (batch_array,),
@@ -291,7 +291,7 @@ def boundary_neumann(
         elif (batch_array.shape[-1] == 2 and isinstance(batch, PDEStatioBatch)) or (
             batch_array.shape[-1] == 3 and isinstance(batch, PDENonStatioBatch)
         ):
-            if u.eq_type == "statio_PDE":
+            if u.eq_type == "PDEStatio":
                 tangent_vec_0 = jnp.repeat(
                     jnp.array([1.0, 0.0])[None], batch_array.shape[0], axis=0
                 )
@@ -309,7 +309,7 @@ def boundary_neumann(
                     (tangent_vec_1,),
                 )
                 values = du_dx1 * n[0, facet] + du_dx2 * n[1, facet]  # dot product
-            if u.eq_type == "nonstatio_PDE":
+            if u.eq_type == "PDENonStatio":
                 tangent_vec_0 = jnp.repeat(
                     jnp.array([0.0, 1.0, 0.0])[None], batch_array.shape[0], axis=0
                 )

jinns/loss/_loss_utils.py

@@ -26,12 +26,12 @@ from jinns.data._Batchs import PDEStatioBatch, PDENonStatioBatch
 from jinns.parameters._params import Params
 
 if TYPE_CHECKING:
-    from jinns.utils._types import BoundaryConditionFun
+    from jinns.utils._types import BoundaryConditionFun, AnyBatch
     from jinns.nn._abstract_pinn import AbstractPINN
 
 
 def dynamic_loss_apply(
-    dyn_loss: Callable,
+    dyn_loss: Callable[[AnyBatch, AbstractPINN, Params[Array]], Array],
     u: AbstractPINN,
     batch: (
         Float[Array, " batch_size 1"]

jinns/loss/_loss_weight_updates.py

@@ -13,6 +13,36 @@ if TYPE_CHECKING:
     from jinns.utils._types import AnyLossComponents, AnyLossWeights
 
 
+def prior_loss(
+    loss_weights: AnyLossWeights,
+    iteration_nb: int,
+    stored_loss_terms: AnyLossComponents,
+) -> Array:
+    """
+    Simple adaptative weights according to the prior loss idea:
+    the ponderation in front of a loss term is given by the inverse of the
+    value of that loss term at the previous iteration
+    """
+
+    def do_nothing(loss_weights, _):
+        return jnp.array(
+            jax.tree.leaves(loss_weights, is_leaf=eqx.is_inexact_array), dtype=float
+        )
+
+    def _prior_loss(_, stored_loss_terms):
+        new_weights = jax.tree.map(
+            lambda slt: 1 / (slt[iteration_nb - 1] + 1e-6), stored_loss_terms
+        )
+        return jnp.array(jax.tree.leaves(new_weights), dtype=float)
+
+    return jax.lax.cond(
+        iteration_nb == 0,
+        lambda op: do_nothing(*op),
+        lambda op: _prior_loss(*op),
+        (loss_weights, stored_loss_terms),
+    )
+
+
 def soft_adapt(
     loss_weights: AnyLossWeights,
     iteration_nb: int,

jinns/loss/_loss_weights.py

@@ -18,6 +18,10 @@ from jinns.loss._loss_components import (
 def lw_converter(x: Array | None) -> Array | None:
     if x is None:
         return x
+    elif isinstance(x, tuple):
+        # user might input tuple of scalar loss weights to account for cases
+        # when dyn loss is also a tuple of (possibly 1D) dyn_loss
+        return tuple(jnp.asarray(x_) for x_ in x)
     else:
         return jnp.asarray(x)
 

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