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

jinns/loss/_LossPDE.py

@@ -122,7 +122,7 @@ class _LossPDEAbstract(
         made broadcastable to `norm_samples`.
         These corresponds to the weights $w_k = \frac{1}{q(x_k)}$ where
         $q(\cdot)$ is the proposal p.d.f. and $x_k$ are the Monte-Carlo samples.
-    obs_slice : EllipsisType | slice, default=None
+    obs_slice : tuple[EllipsisType | slice, ...] | EllipsisType | slice | None, 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.
     key : Key | None
@@ -136,7 +136,7 @@ class _LossPDEAbstract(
     # NOTE static=True only for leaf attributes that are not valid JAX types
     # (ie. jax.Array cannot be static) and that we do not expect to change
     u: eqx.AbstractVar[AbstractPINN]
-    dynamic_loss: eqx.AbstractVar[Y]
+    dynamic_loss: tuple[eqx.AbstractVar[Y] | None, ...]
     omega_boundary_fun: (
         BoundaryConditionFun | dict[str, BoundaryConditionFun] | None
     ) = eqx.field(static=True)
@@ -146,7 +146,7 @@ class _LossPDEAbstract(
     omega_boundary_dim: slice | dict[str, slice] = eqx.field(static=True)
     norm_samples: Float[Array, " nb_norm_samples dimension"] | None
     norm_weights: Float[Array, " nb_norm_samples"] | None
-    obs_slice: EllipsisType | slice = eqx.field(static=True)
+    obs_slice: tuple[EllipsisType | slice, ...] = eqx.field(static=True)
     key: PRNGKeyArray | None
 
     def __init__(
@@ -159,7 +159,10 @@ class _LossPDEAbstract(
         omega_boundary_dim: int | slice | dict[str, int | slice] | None = None,
         norm_samples: Float[Array, " nb_norm_samples dimension"] | None = None,
         norm_weights: Float[Array, " nb_norm_samples"] | float | int | None = None,
-        obs_slice: EllipsisType | slice | None = None,
+        obs_slice: tuple[EllipsisType | slice, ...]
+        | EllipsisType
+        | slice
+        | None = None,
         key: PRNGKeyArray | None = None,
         derivative_keys: DKPDE,
         **kwargs: Any,  # for arguments for super()
@@ -181,7 +184,9 @@ class _LossPDEAbstract(
             )
 
         if obs_slice is None:
-            self.obs_slice = jnp.s_[...]
+            self.obs_slice = (jnp.s_[...],)
+        elif not isinstance(obs_slice, tuple):
+            self.obs_slice = (obs_slice,)
         else:
             self.obs_slice = obs_slice
 
@@ -321,16 +326,23 @@ class _LossPDEAbstract(
 
     def _get_dyn_loss_fun(
         self, batch: B, vmap_in_axes_params: tuple[Params[int | None] | None]
-    ) -> Callable[[Params[Array]], Array] | None:
-        if self.dynamic_loss is not None:
-            dyn_loss_eval = self.dynamic_loss.evaluate
-            dyn_loss_fun: Callable[[Params[Array]], Array] | None = (
-                lambda p: dynamic_loss_apply(
-                    dyn_loss_eval,
-                    self.u,
-                    self._get_dynamic_loss_batch(batch),
-                    _set_derivatives(p, self.derivative_keys.dyn_loss),
-                    self.vmap_in_axes + vmap_in_axes_params,
+    ) -> tuple[Callable[[Params[Array]], Array], ...] | None:
+        # Note, for the record, multiple dynamic losses
+        # have been introduced in MR 92
+        if self.dynamic_loss != (None,):
+            dyn_loss_fun: tuple[Callable[[Params[Array]], Array], ...] | None = (
+                jax.tree.map(
+                    lambda d: lambda p: dynamic_loss_apply(
+                        d.evaluate,
+                        self.u,
+                        self._get_dynamic_loss_batch(batch),
+                        _set_derivatives(p, self.derivative_keys.dyn_loss),
+                        self.vmap_in_axes + vmap_in_axes_params,
+                    ),
+                    self.dynamic_loss,
+                    is_leaf=lambda x: isinstance(
+                        x, (PDEStatio, PDENonStatio)
+                    ),  # do not traverse further than first level
                 )
             )
         else:
@@ -340,9 +352,13 @@ class _LossPDEAbstract(
 
     def _get_norm_loss_fun(
         self, batch: B, vmap_in_axes_params: tuple[Params[int | None] | None]
-    ) -> Callable[[Params[Array]], Array] | None:
+    ) -> tuple[Callable[[Params[Array]], Array], ...] | None:
+        # Note that since MR 92
+        # norm_loss_fun is formed as a tuple for
+        # consistency with dynamic and observation
+        # losses and more modularity for later
         if self.norm_samples is not None:
-            norm_loss_fun: Callable[[Params[Array]], Array] | None = (
+            norm_loss_fun: tuple[Callable[[Params[Array]], Array], ...] | None = (
                 lambda p: normalization_loss_apply(
                     self.u,
                     cast(
@@ -351,7 +367,7 @@ class _LossPDEAbstract(
                     _set_derivatives(p, self.derivative_keys.norm_loss),
                     vmap_in_axes_params,
                     self.norm_weights,  # type: ignore -> can't get the __post_init__ narrowing here
-                )
+                ),
             )
         else:
             norm_loss_fun = None
@@ -359,20 +375,24 @@ class _LossPDEAbstract(
 
     def _get_boundary_loss_fun(
         self, batch: B
-    ) -> Callable[[Params[Array]], Array] | None:
+    ) -> tuple[Callable[[Params[Array]], Array], ...] | None:
+        # Note that since MR 92
+        # boundary_loss_fun is formed as a tuple for
+        # consistency with dynamic and observation
+        # losses and more modularity for later
         if (
             self.omega_boundary_condition is not None
             and self.omega_boundary_fun is not None
         ):
-            boundary_loss_fun: Callable[[Params[Array]], Array] | None = (
+            boundary_loss_fun: tuple[Callable[[Params[Array]], Array], ...] | None = (
                 lambda p: boundary_condition_apply(
                     self.u,
                     batch,
                     _set_derivatives(p, self.derivative_keys.boundary_loss),
                     self.omega_boundary_fun,  # type: ignore (we are in lambda)
                     self.omega_boundary_condition,  # type: ignore
-                    self.omega_boundary_dim,  # type: ignore
-                )
+                    self.omega_boundary_dim,
+                ),
             )
         else:
             boundary_loss_fun = None
@@ -384,24 +404,37 @@ class _LossPDEAbstract(
         batch: B,
         vmap_in_axes_params: tuple[Params[int | None] | None],
         params: Params[Array],
-    ) -> tuple[Params[Array] | None, Callable[[Params[Array]], Array] | None]:
+    ) -> tuple[
+        Params[Array] | None, tuple[Callable[[Params[Array]], Array], ...] | None
+    ]:
+        # Note, for the record, multiple DGObs
+        # (leading to batch.obs_batch_dict being tuple | None)
+        # have been introduced in MR 92
         if batch.obs_batch_dict is not None:
-            # update params with the batches of observed params
-            params_obs = update_eq_params(params, batch.obs_batch_dict["eq_params"])
-
-            pinn_in, val = (
-                batch.obs_batch_dict["pinn_in"],
-                batch.obs_batch_dict["val"],
+            if len(batch.obs_batch_dict) != len(self.obs_slice):
+                raise ValueError(
+                    "There must be the same number of "
+                    "observation datasets as the number of "
+                    "obs_slice"
+                )
+            params_obs = jax.tree.map(
+                lambda d: update_eq_params(params, d["eq_params"]),
+                batch.obs_batch_dict,
+                is_leaf=lambda x: isinstance(x, dict),
             )
-
-            obs_loss_fun: Callable[[Params[Array]], Array] | None = (
-                lambda po: observations_loss_apply(
-                    self.u,
-                    pinn_in,
-                    _set_derivatives(po, self.derivative_keys.observations),
-                    self.vmap_in_axes + vmap_in_axes_params,
-                    val,
+            obs_loss_fun: tuple[Callable[[Params[Array]], Array], ...] | None = (
+                jax.tree.map(
+                    lambda d, slice_: lambda po: observations_loss_apply(
+                        self.u,
+                        d["pinn_in"],
+                        _set_derivatives(po, self.derivative_keys.observations),
+                        self.vmap_in_axes + vmap_in_axes_params,
+                        d["val"],
+                        slice_,
+                    ),
+                    batch.obs_batch_dict,
                     self.obs_slice,
+                    is_leaf=lambda x: isinstance(x, dict),
                 )
             )
         else:
@@ -434,7 +467,7 @@ class LossPDEStatio(
     ----------
     u : AbstractPINN
         the PINN
-    dynamic_loss : PDEStatio | None
+    dynamic_loss : tuple[PDEStatio, ...] | PDEStatio | None
         the stationary PDE dynamic part of the loss, basically the differential
         operator $\mathcal{N}[u](x)$. Should implement a method
         `dynamic_loss.evaluate(x, u, params)`.
@@ -497,7 +530,7 @@ class LossPDEStatio(
         Alternatively, the user can pass a float or an integer.
         These corresponds to the weights $w_k = \frac{1}{q(x_k)}$ where
         $q(\cdot)$ is the proposal p.d.f. and $x_k$ are the Monte-Carlo samples.
-    obs_slice : slice, default=None
+    obs_slice : tuple[EllipsisType | slice, ...] | EllipsisType | slice | None, 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.
 
@@ -512,7 +545,7 @@ class LossPDEStatio(
     # (ie. jax.Array cannot be static) and that we do not expect to change
 
     u: AbstractPINN
-    dynamic_loss: PDEStatio | None
+    dynamic_loss: tuple[PDEStatio | None, ...]
     loss_weights: LossWeightsPDEStatio
     derivative_keys: DerivativeKeysPDEStatio
 
@@ -522,7 +555,7 @@ class LossPDEStatio(
         self,
         *,
         u: AbstractPINN,
-        dynamic_loss: PDEStatio | None,
+        dynamic_loss: tuple[PDEStatio, ...] | PDEStatio | None,
         loss_weights: LossWeightsPDEStatio | None = None,
         derivative_keys: DerivativeKeysPDEStatio | None = None,
         params: Params[Array] | None = None,
@@ -543,15 +576,16 @@ class LossPDEStatio(
                     "Problem at derivative_keys initialization "
                     f"received {derivative_keys=} and {params=}"
                 ) from exc
-        else:
-            derivative_keys = derivative_keys
 
         super().__init__(
             derivative_keys=derivative_keys,
             vmap_in_axes=(0,),
             **kwargs,
         )
-        self.dynamic_loss = dynamic_loss
+        if not isinstance(dynamic_loss, tuple):
+            self.dynamic_loss = (dynamic_loss,)
+        else:
+            self.dynamic_loss = dynamic_loss
 
     def _get_dynamic_loss_batch(
         self, batch: PDEStatioBatch
@@ -619,29 +653,44 @@ class LossPDEStatio(
         )
 
         # get the unweighted mses for each loss term as well as the gradients
-        all_funs: PDEStatioComponents[Callable[[Params[Array]], Array] | None] = (
+        all_funs: PDEStatioComponents[
+            tuple[Callable[[Params[Array]], Array], ...] | None
+        ] = PDEStatioComponents(
+            dyn_loss_fun,
+            norm_loss_fun,
+            boundary_loss_fun,
+            obs_loss_fun,
+        )
+        all_params: PDEStatioComponents[tuple[Params[Array], ...] | None] = (
             PDEStatioComponents(
-                dyn_loss_fun, norm_loss_fun, boundary_loss_fun, obs_loss_fun
+                jax.tree.map(lambda l: params, dyn_loss_fun),
+                jax.tree.map(lambda l: params, norm_loss_fun),
+                jax.tree.map(lambda l: params, boundary_loss_fun),
+                params_obs,
             )
         )
-        all_params: PDEStatioComponents[Params[Array] | None] = PDEStatioComponents(
-            params, params, params, params_obs
-        )
         mses_grads = jax.tree.map(
             self.get_gradients,
             all_funs,
             all_params,
             is_leaf=lambda x: x is None,
         )
+        # NOTE the is_leaf below is more complex since it must pass possible the tuple
+        # of dyn_loss and then stops (but also account it should not stop when
+        # the tuple of dyn_loss is of length 2)
         mses = jax.tree.map(
-            lambda leaf: leaf[0],  # type: ignore
+            lambda leaf: leaf[0],
             mses_grads,
-            is_leaf=lambda x: isinstance(x, tuple),
+            is_leaf=lambda x: isinstance(x, tuple)
+            and len(x) == 2
+            and isinstance(x[1], Params),
         )
         grads = jax.tree.map(
-            lambda leaf: leaf[1],  # type: ignore
+            lambda leaf: leaf[1],
             mses_grads,
-            is_leaf=lambda x: isinstance(x, tuple),
+            is_leaf=lambda x: isinstance(x, tuple)
+            and len(x) == 2
+            and isinstance(x[1], Params),
         )
 
         return mses, grads
@@ -673,7 +722,7 @@ class LossPDENonStatio(
     ----------
     u : AbstractPINN
         the PINN
-    dynamic_loss : PDENonStatio
+    dynamic_loss : tuple[PDENonStatio, ...] | PDENonStatio | None
         the non stationary PDE dynamic part of the loss, basically the differential
         operator $\mathcal{N}[u](t, x)$. Should implement a method
         `dynamic_loss.evaluate(t, x, u, params)`.
@@ -750,14 +799,13 @@ class LossPDENonStatio(
         Alternatively, the user can pass a float or an integer.
         These corresponds to the weights $w_k = \frac{1}{q(x_k)}$ where
         $q(\cdot)$ is the proposal p.d.f. and $x_k$ are the Monte-Carlo samples.
-    obs_slice : slice, default=None
+    obs_slice : tuple[EllipsisType | slice, ...] | EllipsisType | slice | None, 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.
-
     """
 
     u: AbstractPINN
-    dynamic_loss: PDENonStatio | None
+    dynamic_loss: tuple[PDENonStatio | None, ...]
     loss_weights: LossWeightsPDENonStatio
     derivative_keys: DerivativeKeysPDENonStatio
     params: InitVar[Params[Array] | None]
@@ -774,7 +822,7 @@ class LossPDENonStatio(
         self,
         *,
         u: AbstractPINN,
-        dynamic_loss: PDENonStatio | None,
+        dynamic_loss: tuple[PDENonStatio, ...] | PDENonStatio | None,
         loss_weights: LossWeightsPDENonStatio | None = None,
         derivative_keys: DerivativeKeysPDENonStatio | None = None,
         initial_condition_fun: Callable[[Float[Array, " dimension"]], Array]
@@ -800,8 +848,6 @@ class LossPDENonStatio(
                     "Problem at derivative_keys initialization "
                     f"received {derivative_keys=} and {params=}"
                 ) from exc
-        else:
-            derivative_keys = derivative_keys
 
         super().__init__(
             derivative_keys=derivative_keys,
@@ -809,7 +855,10 @@ class LossPDENonStatio(
             **kwargs,
         )
 
-        self.dynamic_loss = dynamic_loss
+        if not isinstance(dynamic_loss, tuple):
+            self.dynamic_loss = (dynamic_loss,)
+        else:
+            self.dynamic_loss = dynamic_loss
 
         if initial_condition_fun is None:
             warnings.warn(
@@ -904,7 +953,13 @@ class LossPDENonStatio(
 
         # initial condition
         if self.initial_condition_fun is not None:
-            mse_initial_condition_fun: Callable[[Params[Array]], Array] | None = (
+            # Note that since MR 92
+            # initial_condition_fun is formed as a tuple for
+            # consistency with dynamic and observation
+            # losses and more modularity for later
+            initial_condition_fun: (
+                tuple[Callable[[Params[Array]], Array], ...] | None
+            ) = (
                 lambda p: initial_condition_apply(
                     self.u,
                     omega_initial_batch,
@@ -912,39 +967,52 @@ class LossPDENonStatio(
                     (0,) + vmap_in_axes_params,
                     self.initial_condition_fun,  # type: ignore
                     self.t0,
-                )
+                ),
             )
         else:
-            mse_initial_condition_fun = None
+            initial_condition_fun = None
 
         # get the unweighted mses for each loss term as well as the gradients
-        all_funs: PDENonStatioComponents[Callable[[Params[Array]], Array] | None] = (
+        all_funs: PDENonStatioComponents[
+            tuple[Callable[[Params[Array]], Array], ...] | None
+        ] = PDENonStatioComponents(
+            dyn_loss_fun,
+            norm_loss_fun,
+            boundary_loss_fun,
+            obs_loss_fun,
+            initial_condition_fun,
+        )
+        all_params: PDENonStatioComponents[tuple[Params[Array], ...] | None] = (
             PDENonStatioComponents(
-                dyn_loss_fun,
-                norm_loss_fun,
-                boundary_loss_fun,
-                obs_loss_fun,
-                mse_initial_condition_fun,
+                jax.tree.map(lambda l: params, dyn_loss_fun),
+                jax.tree.map(lambda l: params, norm_loss_fun),
+                jax.tree.map(lambda l: params, boundary_loss_fun),
+                params_obs,
+                jax.tree.map(lambda l: params, initial_condition_fun),
             )
         )
-        all_params: PDENonStatioComponents[Params[Array] | None] = (
-            PDENonStatioComponents(params, params, params, params_obs, params)
-        )
         mses_grads = jax.tree.map(
             self.get_gradients,
             all_funs,
             all_params,
             is_leaf=lambda x: x is None,
         )
+        # NOTE the is_leaf below is more complex since it must pass possible the tuple
+        # of dyn_loss and then stops (but also account it should not stop when
+        # the tuple of dyn_loss is of length 2)
         mses = jax.tree.map(
-            lambda leaf: leaf[0],  # type: ignore
+            lambda leaf: leaf[0],
             mses_grads,
-            is_leaf=lambda x: isinstance(x, tuple),
+            is_leaf=lambda x: isinstance(x, tuple)
+            and len(x) == 2
+            and isinstance(x[1], Params),
         )
         grads = jax.tree.map(
-            lambda leaf: leaf[1],  # type: ignore
+            lambda leaf: leaf[1],
             mses_grads,
-            is_leaf=lambda x: isinstance(x, tuple),
+            is_leaf=lambda x: isinstance(x, tuple)
+            and len(x) == 2
+            and isinstance(x[1], Params),
         )
 
         return mses, grads

jinns/loss/_abstract_loss.py

@@ -1,14 +1,21 @@
 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.loss._loss_weight_updates import (
+    soft_adapt,
+    lr_annealing,
+    ReLoBRaLo,
+    prior_loss,
+)
 from jinns.utils._types import (
     AnyLossComponents,
     AnyBatch,
@@ -38,6 +45,11 @@ DK = TypeVar("DK", bound=AnyDerivativeKeys)
 # the return types of evaluate_by_terms for example!
 
 
+AvailableUpdateWeightMethods = Literal[
+    "softadapt", "soft_adapt", "prior_loss", "lr_annealing", "ReLoBRaLo"
+]
+
+
 class AbstractLoss(eqx.Module, Generic[L, B, C, DK]):
     """
     About the call:
@@ -46,10 +58,43 @@ class AbstractLoss(eqx.Module, Generic[L, B, C, DK]):
 
     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)
@@ -127,7 +172,11 @@ class AbstractLoss(eqx.Module, Generic[L, B, C, DK]):
         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) -> Params[Array | None]:
@@ -171,6 +220,8 @@ class AbstractLoss(eqx.Module, Generic[L, B, C, DK]):
             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":
@@ -183,6 +234,13 @@ class AbstractLoss(eqx.Module, Generic[L, B, C, DK]):
         # 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)