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

jinns/loss/_DynamicLossAbstract.py

@@ -21,53 +21,22 @@ else:
     from equinox import AbstractClassVar
 
 
-def _decorator_heteregeneous_params(evaluate, eq_type):
+def _decorator_heteregeneous_params(evaluate):
 
-    def wrapper_ode(*args):
-        self, t, u, params = args
+    def wrapper(*args):
+        self, inputs, u, params = args
         _params = eqx.tree_at(
             lambda p: p.eq_params,
             params,
             self._eval_heterogeneous_parameters(
-                t, None, u, params, self.eq_params_heterogeneity
+                inputs, u, params, self.eq_params_heterogeneity
             ),
         )
         new_args = args[:-1] + (_params,)
         res = evaluate(*new_args)
         return res
 
-    def wrapper_pde_statio(*args):
-        self, x, u, params = args
-        _params = eqx.tree_at(
-            lambda p: p.eq_params,
-            params,
-            self._eval_heterogeneous_parameters(
-                None, x, u, params, self.eq_params_heterogeneity
-            ),
-        )
-        new_args = args[:-1] + (_params,)
-        res = evaluate(*new_args)
-        return res
-
-    def wrapper_pde_non_statio(*args):
-        self, t, x, u, params = args
-        _params = eqx.tree_at(
-            lambda p: p.eq_params,
-            params,
-            self._eval_heterogeneous_parameters(
-                t, x, u, params, self.eq_params_heterogeneity
-            ),
-        )
-        new_args = args[:-1] + (_params,)
-        res = evaluate(*new_args)
-        return res
-
-    if eq_type == "ODE":
-        return wrapper_ode
-    elif eq_type == "Statio PDE":
-        return wrapper_pde_statio
-    elif eq_type == "Non-statio PDE":
-        return wrapper_pde_non_statio
+    return wrapper
 
 
 class DynamicLoss(eqx.Module):
@@ -110,8 +79,7 @@ class DynamicLoss(eqx.Module):
 
     def _eval_heterogeneous_parameters(
         self,
-        t: Float[Array, "1"],
-        x: Float[Array, "dim"],
+        inputs: Float[Array, "1"] | Float[Array, "dim"] | Float[Array, "1+dim"],
         u: eqx.Module,
         params: Params | ParamsDict,
         eq_params_heterogeneity: Dict[str, Callable | None] = None,
@@ -124,14 +92,7 @@ class DynamicLoss(eqx.Module):
                 if eq_params_heterogeneity[k] is None:
                     eq_params_[k] = p
                 else:
-                    # heterogeneity encoded through a function whose
-                    # signature will vary according to _eq_type
-                    if self._eq_type == "ODE":
-                        eq_params_[k] = eq_params_heterogeneity[k](t, u, params)
-                    elif self._eq_type == "Statio PDE":
-                        eq_params_[k] = eq_params_heterogeneity[k](x, u, params)
-                    elif self._eq_type == "Non-statio PDE":
-                        eq_params_[k] = eq_params_heterogeneity[k](t, x, u, params)
+                    eq_params_[k] = eq_params_heterogeneity[k](inputs, u, params)
             except KeyError:
                 # we authorize missing eq_params_heterogeneity key
                 # if its heterogeneity is None anyway
@@ -140,22 +101,17 @@ class DynamicLoss(eqx.Module):
 
     def _evaluate(
         self,
-        t: Float[Array, "1"],
-        x: Float[Array, "dim"],
+        inputs: Float[Array, "1"] | Float[Array, "dim"] | Float[Array, "1+dim"],
         u: eqx.Module,
         params: Params | ParamsDict,
     ) -> float:
-        # Here we handle the various possible signature
-        if self._eq_type == "ODE":
-            ans = self.equation(t, u, params)
-        elif self._eq_type == "Statio PDE":
-            ans = self.equation(x, u, params)
-        elif self._eq_type == "Non-statio PDE":
-            ans = self.equation(t, x, u, params)
-        else:
-            raise NotImplementedError("the equation type is not handled.")
-
-        return ans
+        evaluation = self.equation(inputs, u, params)
+        if len(evaluation.shape) == 0:
+            raise ValueError(
+                "The output of dynamic loss must be vectorial, "
+                "i.e. of shape (d,) with d >= 1"
+            )
+        return evaluation
 
     @abc.abstractmethod
     def equation(self, *args, **kwargs):
@@ -191,7 +147,7 @@ class ODE(DynamicLoss):
 
     _eq_type: ClassVar[str] = "ODE"
 
-    @partial(_decorator_heteregeneous_params, eq_type="ODE")
+    @partial(_decorator_heteregeneous_params)
     def evaluate(
         self,
         t: Float[Array, "1"],
@@ -199,7 +155,7 @@ class ODE(DynamicLoss):
         params: Params | ParamsDict,
     ) -> float:
         """Here we call DynamicLoss._evaluate with x=None"""
-        return self._evaluate(t, None, u, params)
+        return self._evaluate(t, u, params)
 
     @abc.abstractmethod
     def equation(
@@ -260,12 +216,12 @@ class PDEStatio(DynamicLoss):
 
     _eq_type: ClassVar[str] = "Statio PDE"
 
-    @partial(_decorator_heteregeneous_params, eq_type="Statio PDE")
+    @partial(_decorator_heteregeneous_params)
     def evaluate(
         self, x: Float[Array, "dimension"], u: eqx.Module, params: Params | ParamsDict
     ) -> float:
         """Here we call the DynamicLoss._evaluate with t=None"""
-        return self._evaluate(None, x, u, params)
+        return self._evaluate(x, u, params)
 
     @abc.abstractmethod
     def equation(
@@ -325,23 +281,21 @@ class PDENonStatio(DynamicLoss):
 
     _eq_type: ClassVar[str] = "Non-statio PDE"
 
-    @partial(_decorator_heteregeneous_params, eq_type="Non-statio PDE")
+    @partial(_decorator_heteregeneous_params)
     def evaluate(
         self,
-        t: Float[Array, "1"],
-        x: Float[Array, "dim"],
+        t_x: Float[Array, "1 + dim"],
         u: eqx.Module,
         params: Params | ParamsDict,
     ) -> float:
         """Here we call the DynamicLoss._evaluate with full arguments"""
-        ans = self._evaluate(t, x, u, params)
+        ans = self._evaluate(t_x, u, params)
         return ans
 
     @abc.abstractmethod
     def equation(
         self,
-        t: Float[Array, "1"],
-        x: Float[Array, "dim"],
+        t_x: Float[Array, "1 + dim"],
         u: eqx.Module,
         params: Params | ParamsDict,
     ) -> float:
@@ -353,10 +307,8 @@ class PDENonStatio(DynamicLoss):
 
         Parameters
         ----------
-        t : Float[Array, "1"]
-            A 1-dimensional jnp.array representing the time point.
-        x : Float[Array, "d"]
-            A `d` dimensional jnp.array representing a point in the spatial domain $\Omega$.
+        t_x : Float[Array, "1 + dim"]
+            A jnp array containing the concatenation of a time point and a point in $\Omega$
         u : eqx.Module
             The neural network.
         params : Params | ParamsDict

jinns/loss/_LossODE.py

@@ -28,7 +28,7 @@ from jinns.parameters._params import (
 from jinns.parameters._derivative_keys import _set_derivatives, DerivativeKeysODE
 from jinns.loss._loss_weights import LossWeightsODE, LossWeightsODEDict
 from jinns.loss._DynamicLossAbstract import ODE
-from jinns.utils._pinn import PINN
+from jinns.nn._pinn import PINN
 
 if TYPE_CHECKING:
     from jinns.utils._types import *
@@ -209,7 +209,7 @@ class LossODE(_LossODEAbstract):
             mse_dyn_loss = dynamic_loss_apply(
                 self.dynamic_loss.evaluate,
                 self.u,
-                (temporal_batch,),
+                temporal_batch,
                 _set_derivatives(params, self.derivative_keys.dyn_loss),
                 self.vmap_in_axes + vmap_in_axes_params,
                 self.loss_weights.dyn_loss,
@@ -227,7 +227,7 @@ class LossODE(_LossODEAbstract):
             else:
                 v_u = vmap(self.u, (None,) + vmap_in_axes_params)
             t0, u0 = self.initial_condition  # pylint: disable=unpacking-non-sequence
-            t0 = jnp.array(t0)
+            t0 = jnp.array([t0])
             u0 = jnp.array(u0)
             mse_initial_condition = jnp.mean(
                 self.loss_weights.initial_condition
@@ -522,7 +522,7 @@ class SystemLossODE(eqx.Module):
             return dynamic_loss_apply(
                 dyn_loss.evaluate,
                 self.u_dict,
-                (temporal_batch,),
+                temporal_batch,
                 _set_derivatives(params_dict, self.derivative_keys_dyn_loss.dyn_loss),
                 vmap_in_axes_t + vmap_in_axes_params,
                 loss_weight,

jinns/loss/_LossPDE.py

@@ -22,9 +22,7 @@ from jinns.loss._loss_utils import (
     initial_condition_apply,
     constraints_system_loss_apply,
 )
-from jinns.data._DataGenerators import (
-    append_obs_batch,
-)
+from jinns.data._DataGenerators import append_obs_batch
 from jinns.parameters._params import (
     _get_vmap_in_axes_params,
     _update_eq_params_dict,
@@ -40,8 +38,8 @@ from jinns.loss._loss_weights import (
     LossWeightsPDEDict,
 )
 from jinns.loss._DynamicLossAbstract import PDEStatio, PDENonStatio
-from jinns.utils._pinn import PINN
-from jinns.utils._spinn import SPINN
+from jinns.nn._pinn import PINN
+from jinns.nn._spinn import SPINN
 from jinns.data._Batchs import PDEStatioBatch, PDENonStatioBatch
 
 
@@ -94,12 +92,16 @@ class _LossPDEAbstract(eqx.Module):
         Note that it must be a slice and not an integer
         (but a preprocessing of the user provided argument takes care of it)
     norm_samples : Float[Array, "nb_norm_samples dimension"], default=None
-        Fixed sample point in the space over which to compute the
+        Monte-Carlo sample points for computing the
         normalization constant. Default is None.
-    norm_int_length : float, default=None
-        A float. Must be provided if `norm_samples` is provided. The domain area
-        (or interval length in 1D) upon which we perform the numerical
-        integration. Default None
+    norm_weights : Float[Array, "nb_norm_samples"] | float | int, default=None
+        The importance sampling weights for Monte-Carlo integration of the
+        normalization constant. Must be provided if `norm_samples` is provided.
+        `norm_weights` should have the same leading dimension as
+        `norm_samples`.
+        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
         slice object specifying the begininning/ending of the PINN output
         that is observed (this is then useful for multidim PINN). Default is None.
@@ -129,7 +131,9 @@ class _LossPDEAbstract(eqx.Module):
     norm_samples: Float[Array, "nb_norm_samples dimension"] | None = eqx.field(
         kw_only=True, default=None
     )
-    norm_int_length: float | None = eqx.field(kw_only=True, default=None)
+    norm_weights: Float[Array, "nb_norm_samples"] | float | int | None = eqx.field(
+        kw_only=True, default=None
+    )
     obs_slice: slice | None = eqx.field(kw_only=True, default=None, static=True)
 
     params: InitVar[Params] = eqx.field(kw_only=True, default=None)
@@ -253,8 +257,25 @@ class _LossPDEAbstract(eqx.Module):
             if not isinstance(self.omega_boundary_dim, slice):
                 raise ValueError("self.omega_boundary_dim must be a jnp.s_ object")
 
-        if self.norm_samples is not None and self.norm_int_length is None:
-            raise ValueError("self.norm_samples and norm_int_length must be provided")
+        if self.norm_samples is not None:
+            if self.norm_weights is None:
+                raise ValueError(
+                    "`norm_weights` must be provided when `norm_samples` is used!"
+                )
+            try:
+                assert self.norm_weights.shape[0] == self.norm_samples.shape[0]
+            except (AssertionError, AttributeError):
+                if isinstance(self.norm_weights, (int, float)):
+                    self.norm_weights = jnp.array(
+                        [self.norm_weights], dtype=jax.dtypes.canonicalize_dtype(float)
+                    )
+                else:
+                    raise ValueError(
+                        "`norm_weights` should have the same leading dimension"
+                        " as `norm_samples`,"
+                        f" got shape {self.norm_weights.shape} and"
+                        f" shape {self.norm_samples.shape}."
+                    )
 
     @abc.abstractmethod
     def evaluate(
@@ -326,12 +347,16 @@ class LossPDEStatio(_LossPDEAbstract):
         Note that it must be a slice and not an integer
         (but a preprocessing of the user provided argument takes care of it)
     norm_samples : Float[Array, "nb_norm_samples dimension"], default=None
-        Fixed sample point in the space over which to compute the
+        Monte-Carlo sample points for computing the
         normalization constant. Default is None.
-    norm_int_length : float, default=None
-        A float. Must be provided if `norm_samples` is provided. The domain area
-        (or interval length in 1D) upon which we perform the numerical
-        integration. Default None
+    norm_weights : Float[Array, "nb_norm_samples"] | float | int, default=None
+        The importance sampling weights for Monte-Carlo integration of the
+        normalization constant. Must be provided if `norm_samples` is provided.
+        `norm_weights` should have the same leading dimension as
+        `norm_samples`.
+        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
         slice object specifying the begininning/ending of the PINN output
         that is observed (this is then useful for multidim PINN). Default is None.
@@ -370,8 +395,8 @@ class LossPDEStatio(_LossPDEAbstract):
 
     def _get_dynamic_loss_batch(
         self, batch: PDEStatioBatch
-    ) -> tuple[Float[Array, "batch_size dimension"]]:
-        return (batch.inside_batch,)
+    ) -> Float[Array, "batch_size dimension"]:
+        return batch.domain_batch
 
     def _get_normalization_loss_batch(
         self, _
@@ -432,8 +457,8 @@ class LossPDEStatio(_LossPDEAbstract):
                 self.u,
                 self._get_normalization_loss_batch(batch),
                 _set_derivatives(params, self.derivative_keys.norm_loss),
-                self.vmap_in_axes + vmap_in_axes_params,
-                self.norm_int_length,
+                vmap_in_axes_params,
+                self.norm_weights,
                 self.loss_weights.norm_loss,
             )
         else:
@@ -551,12 +576,16 @@ class LossPDENonStatio(LossPDEStatio):
         Note that it must be a slice and not an integer
         (but a preprocessing of the user provided argument takes care of it)
     norm_samples : Float[Array, "nb_norm_samples dimension"], default=None
-        Fixed sample point in the space over which to compute the
+        Monte-Carlo sample points for computing the
         normalization constant. Default is None.
-    norm_int_length : float, default=None
-        A float. Must be provided if `norm_samples` is provided. The domain area
-        (or interval length in 1D) upon which we perform the numerical
-        integration. Default None
+    norm_weights : Float[Array, "nb_norm_samples"] | float | int, default=None
+        The importance sampling weights for Monte-Carlo integration of the
+        normalization constant. Must be provided if `norm_samples` is provided.
+        `norm_weights` should have the same leading dimension as
+        `norm_samples`.
+        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
         slice object specifying the begininning/ending of the PINN output
         that is observed (this is then useful for multidim PINN). Default is None.
@@ -575,6 +604,9 @@ class LossPDENonStatio(LossPDEStatio):
         kw_only=True, default=None, static=True
     )
 
+    _max_norm_samples_omega: Int = eqx.field(init=False, static=True)
+    _max_norm_time_slices: Int = eqx.field(init=False, static=True)
+
     def __post_init__(self, params=None):
         """
         Note that neither __init__ or __post_init__ are called when udating a
@@ -585,7 +617,7 @@ class LossPDENonStatio(LossPDEStatio):
         )  # because __init__ or __post_init__ of Base
         # class is not automatically called
 
-        self.vmap_in_axes = (0, 0)  # for t and x
+        self.vmap_in_axes = (0,)  # for t_x
 
         if self.initial_condition_fun is None:
             warnings.warn(
@@ -593,28 +625,28 @@ class LossPDENonStatio(LossPDEStatio):
                 "case (e.g by. hardcoding it into the PINN output)."
             )
 
+        # witht the variables below we avoid memory overflow since a cartesian
+        # product is taken
+        self._max_norm_time_slices = 100
+        self._max_norm_samples_omega = 1000
+
     def _get_dynamic_loss_batch(
         self, batch: PDENonStatioBatch
-    ) -> tuple[Float[Array, "batch_size 1"], Float[Array, "batch_size dimension"]]:
-        times_batch = batch.times_x_inside_batch[:, 0:1]
-        omega_batch = batch.times_x_inside_batch[:, 1:]
-        return (times_batch, omega_batch)
+    ) -> Float[Array, "batch_size 1+dimension"]:
+        return batch.domain_batch
 
     def _get_normalization_loss_batch(
         self, batch: PDENonStatioBatch
-    ) -> tuple[Float[Array, "batch_size 1"], Float[Array, "nb_norm_samples dimension"]]:
+    ) -> Float[Array, "nb_norm_time_slices nb_norm_samples dimension"]:
         return (
-            batch.times_x_inside_batch[:, 0:1],
-            self.norm_samples,
+            batch.domain_batch[: self._max_norm_time_slices, 0:1],
+            self.norm_samples[: self._max_norm_samples_omega],
         )
 
     def _get_observations_loss_batch(
         self, batch: PDENonStatioBatch
     ) -> tuple[Float[Array, "batch_size 1"], Float[Array, "batch_size dimension"]]:
-        return (
-            batch.obs_batch_dict["pinn_in"][:, 0:1],
-            batch.obs_batch_dict["pinn_in"][:, 1:],
-        )
+        return (batch.obs_batch_dict["pinn_in"],)
 
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
@@ -637,7 +669,7 @@ class LossPDENonStatio(LossPDEStatio):
             of parameters (eg. for metamodeling) and an optional additional batch of observed
             inputs/outputs/parameters
         """
-        omega_batch = batch.times_x_inside_batch[:, 1:]
+        omega_batch = batch.initial_batch
 
         # Retrieve the optional eq_params_batch
         # and update eq_params with the latter
@@ -660,7 +692,6 @@ class LossPDENonStatio(LossPDEStatio):
                 _set_derivatives(params, self.derivative_keys.initial_condition),
                 (0,) + vmap_in_axes_params,
                 self.initial_condition_fun,
-                omega_batch.shape[0],
                 self.loss_weights.initial_condition,
             )
         else:
@@ -730,15 +761,21 @@ class SystemLossPDE(eqx.Module):
         (default) then no temporal boundary condition is applied
         Must share the keys of `u_dict`
     norm_samples_dict : Dict[str, Float[Array, "nb_norm_samples dimension"] | None, default=None
-        A dict of fixed sample point in the space over which to compute the
-        normalization constant. Default is None
-        Must share the keys of `u_dict`
-    norm_int_length_dict : Dict[str, float | None] | None, default=None
-        A dict of Float. The domain area
-        (or interval length in 1D) upon which we perform the numerical
-        integration for each element of u_dict.
-        Default is None
+        A dict of Monte-Carlo sample points for computing the
+        normalization constant. Default is None.
         Must share the keys of `u_dict`
+    norm_weights_dict : Dict[str, Array[Float, "nb_norm_samples"] | float | int | None] | None, default=None
+        A dict of jnp.array with the same keys as `u_dict`. The importance
+        sampling weights for Monte-Carlo integration of the
+        normalization constant for each element of u_dict. Must be provided if
+        `norm_samples_dict` is provided.
+        `norm_weights_dict[key]` should have the same leading dimension as
+        `norm_samples_dict[key]` for each `key`.
+        Alternatively, the user can pass a float or an integer.
+        For each key, an array of similar shape to `norm_samples_dict[key]`
+        or shape `(1,)` is expected. 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. Default is None
     obs_slice_dict : Dict[str, slice | None] | None, default=None
         dict of obs_slice, with keys from `u_dict` to designate the
         output(s) channels that are forced to observed values, for each
@@ -774,9 +811,9 @@ class SystemLossPDE(eqx.Module):
     norm_samples_dict: Dict[str, Float[Array, "nb_norm_samples dimension"]] | None = (
         eqx.field(kw_only=True, default=None)
     )
-    norm_int_length_dict: Dict[str, float | None] | None = eqx.field(
-        kw_only=True, default=None
-    )
+    norm_weights_dict: (
+        Dict[str, Float[Array, "nb_norm_samples dimension"] | float | int | None] | None
+    ) = eqx.field(kw_only=True, default=None)
     obs_slice_dict: Dict[str, slice | None] | None = eqx.field(
         kw_only=True, default=None, static=True
     )
@@ -819,8 +856,8 @@ class SystemLossPDE(eqx.Module):
             self.initial_condition_fun_dict = self.u_dict_with_none
         if self.norm_samples_dict is None:
             self.norm_samples_dict = self.u_dict_with_none
-        if self.norm_int_length_dict is None:
-            self.norm_int_length_dict = self.u_dict_with_none
+        if self.norm_weights_dict is None:
+            self.norm_weights_dict = self.u_dict_with_none
         if self.obs_slice_dict is None:
             self.obs_slice_dict = {k: jnp.s_[...] for k in self.u_dict.keys()}
             if self.u_dict.keys() != self.obs_slice_dict.keys():
@@ -861,7 +898,7 @@ class SystemLossPDE(eqx.Module):
             or self.u_dict.keys() != self.omega_boundary_dim_dict.keys()
             or self.u_dict.keys() != self.initial_condition_fun_dict.keys()
             or self.u_dict.keys() != self.norm_samples_dict.keys()
-            or self.u_dict.keys() != self.norm_int_length_dict.keys()
+            or self.u_dict.keys() != self.norm_weights_dict.keys()
         ):
             raise ValueError("All the dicts concerning the PINNs should have same keys")
 
@@ -890,7 +927,7 @@ class SystemLossPDE(eqx.Module):
                     omega_boundary_condition=self.omega_boundary_condition_dict[i],
                     omega_boundary_dim=self.omega_boundary_dim_dict[i],
                     norm_samples=self.norm_samples_dict[i],
-                    norm_int_length=self.norm_int_length_dict[i],
+                    norm_weights=self.norm_weights_dict[i],
                     obs_slice=self.obs_slice_dict[i],
                 )
             elif self.u_dict[i].eq_type == "nonstatio_PDE":
@@ -911,7 +948,7 @@ class SystemLossPDE(eqx.Module):
                     omega_boundary_dim=self.omega_boundary_dim_dict[i],
                     initial_condition_fun=self.initial_condition_fun_dict[i],
                     norm_samples=self.norm_samples_dict[i],
-                    norm_int_length=self.norm_int_length_dict[i],
+                    norm_weights=self.norm_weights_dict[i],
                     obs_slice=self.obs_slice_dict[i],
                 )
             else:
@@ -1016,19 +1053,7 @@ class SystemLossPDE(eqx.Module):
         if self.u_dict.keys() != params_dict.nn_params.keys():
             raise ValueError("u_dict and params_dict[nn_params] should have same keys ")
 
-        if isinstance(batch, PDEStatioBatch):
-            omega_batch, _ = batch.inside_batch, batch.border_batch
-            vmap_in_axes_x_or_x_t = (0,)
-
-            batches = (omega_batch,)
-        elif isinstance(batch, PDENonStatioBatch):
-            times_batch = batch.times_x_inside_batch[:, 0:1]
-            omega_batch = batch.times_x_inside_batch[:, 1:]
-
-            batches = (omega_batch, times_batch)
-            vmap_in_axes_x_or_x_t = (0, 0)
-        else:
-            raise ValueError("Wrong type of batch")
+        vmap_in_axes = (0,)
 
         # Retrieve the optional eq_params_batch
         # and update eq_params with the latter
@@ -1036,7 +1061,6 @@ class SystemLossPDE(eqx.Module):
         if batch.param_batch_dict is not None:
             eq_params_batch_dict = batch.param_batch_dict
 
-            # TODO
             # feed the eq_params with the batch
             for k in eq_params_batch_dict.keys():
                 params_dict.eq_params[k] = eq_params_batch_dict[k]
@@ -1050,11 +1074,15 @@ class SystemLossPDE(eqx.Module):
             return dynamic_loss_apply(
                 dyn_loss.evaluate,
                 self.u_dict,
-                batches,
+                (
+                    batch.domain_batch
+                    if isinstance(batch, PDEStatioBatch)
+                    else batch.domain_batch
+                ),
                 _set_derivatives(params_dict, self.derivative_keys_dyn_loss.dyn_loss),
-                vmap_in_axes_x_or_x_t + vmap_in_axes_params,
+                vmap_in_axes + vmap_in_axes_params,
                 loss_weight,
-                u_type=type(list(self.u_dict.values())[0]),
+                u_type=list(self.u_dict.values())[0].__class__.__base__,
             )
 
         dyn_loss_mse_dict = jax.tree_util.tree_map(

jinns/loss/__init__.py

@@ -3,7 +3,7 @@ from ._LossODE import LossODE, SystemLossODE
 from ._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
 from ._DynamicLoss import (
     GeneralizedLotkaVolterra,
-    BurgerEquation,
+    BurgersEquation,
     FPENonStatioLoss2D,
     OU_FPENonStatioLoss2D,
     FisherKPP,
@@ -19,9 +19,10 @@ from ._loss_weights import (
 )
 
 from ._operators import (
-    _div_fwd,
-    _div_rev,
-    _laplacian_fwd,
-    _laplacian_rev,
-    _vectorial_laplacian,
+    divergence_fwd,
+    divergence_rev,
+    laplacian_fwd,
+    laplacian_rev,
+    vectorial_laplacian_fwd,
+    vectorial_laplacian_rev,
 )