jinns 1.0.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

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,
@@ -103,6 +101,9 @@ class _LossPDEAbstract(eqx.Module):
     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.
+    params : InitVar[Params], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
     """
 
     # NOTE static=True only for leaf attributes that are not valid JAX types
@@ -129,18 +130,26 @@ class _LossPDEAbstract(eqx.Module):
     norm_int_length: float | None = eqx.field(kw_only=True, default=None)
     obs_slice: slice | None = eqx.field(kw_only=True, default=None, static=True)
 
-    def __post_init__(self):
+    params: InitVar[Params] = eqx.field(kw_only=True, default=None)
+
+    def __post_init__(self, params=None):
         """
         Note that neither __init__ or __post_init__ are called when udating a
         Module with eqx.tree_at
         """
         if self.derivative_keys is None:
             # be default we only take gradient wrt nn_params
-            self.derivative_keys = (
-                DerivativeKeysPDENonStatio()
-                if isinstance(self, LossPDENonStatio)
-                else DerivativeKeysPDEStatio()
-            )
+            try:
+                self.derivative_keys = (
+                    DerivativeKeysPDENonStatio(params=params)
+                    if isinstance(self, LossPDENonStatio)
+                    else DerivativeKeysPDEStatio(params=params)
+                )
+            except ValueError as exc:
+                raise ValueError(
+                    "Problem at self.derivative_keys initialization "
+                    f"received {self.derivative_keys=} and {params=}"
+                ) from exc
 
         if self.loss_weights is None:
             self.loss_weights = (
@@ -324,6 +333,9 @@ class LossPDEStatio(_LossPDEAbstract):
     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.
+    params : InitVar[Params], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
 
 
     Raises
@@ -342,20 +354,22 @@ class LossPDEStatio(_LossPDEAbstract):
 
     vmap_in_axes: tuple[Int] = eqx.field(init=False, static=True)
 
-    def __post_init__(self):
+    def __post_init__(self, params=None):
         """
         Note that neither __init__ or __post_init__ are called when udating a
         Module with eqx.tree_at!
         """
-        super().__post_init__()  # because __init__ or __post_init__ of Base
+        super().__post_init__(
+            params=params
+        )  # because __init__ or __post_init__ of Base
         # class is not automatically called
 
         self.vmap_in_axes = (0,)  # for x only here
 
     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, _
@@ -416,7 +430,7 @@ 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,
+                vmap_in_axes_params,
                 self.norm_int_length,
                 self.loss_weights.norm_loss,
             )
@@ -547,6 +561,9 @@ class LossPDENonStatio(LossPDEStatio):
     initial_condition_fun : Callable, default=None
         A function representing the temporal initial condition. If None
         (default) then no initial condition is applied
+    params : InitVar[Params], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
 
     """
 
@@ -556,15 +573,20 @@ class LossPDENonStatio(LossPDEStatio):
         kw_only=True, default=None, static=True
     )
 
-    def __post_init__(self):
+    _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
         Module with eqx.tree_at!
         """
-        super().__post_init__()  # because __init__ or __post_init__ of Base
+        super().__post_init__(
+            params=params
+        )  # 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(
@@ -572,28 +594,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)
@@ -616,8 +638,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
@@ -640,7 +661,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:
@@ -725,6 +745,9 @@ class SystemLossPDE(eqx.Module):
         PINNs. Default is None. But if a value is given, all the entries of
         `u_dict` must be represented here with default value `jnp.s_[...]`
         if no particular slice is to be given
+    params : InitVar[ParamsDict], default=None
+        The main Params object of the problem needed to instanciate the
+        DerivativeKeysODE if the latter is not specified.
 
     """
 
@@ -763,22 +786,20 @@ class SystemLossPDE(eqx.Module):
     loss_weights: InitVar[LossWeightsPDEDict | None] = eqx.field(
         kw_only=True, default=None
     )
+    params_dict: InitVar[ParamsDict] = eqx.field(kw_only=True, default=None)
 
     # following have init=False and are set in the __post_init__
     u_constraints_dict: Dict[str, LossPDEStatio | LossPDENonStatio] = eqx.field(
         init=False
     )
-    derivative_keys_u_dict: Dict[
-        str, DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio
-    ] = eqx.field(init=False)
-    derivative_keys_dyn_loss_dict: Dict[
-        str, DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio
-    ] = eqx.field(init=False)
+    derivative_keys_dyn_loss: DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio = (
+        eqx.field(init=False)
+    )
     u_dict_with_none: Dict[str, None] = eqx.field(init=False)
     # internally the loss weights are handled with a dictionary
     _loss_weights: Dict[str, dict] = eqx.field(init=False)
 
-    def __post_init__(self, loss_weights):
+    def __post_init__(self, loss_weights=None, params_dict=None):
         # a dictionary that will be useful at different places
         self.u_dict_with_none = {k: None for k in self.u_dict.keys()}
         # First, for all the optional dict,
@@ -818,24 +839,19 @@ class SystemLossPDE(eqx.Module):
             # iterating on dynamic_loss_dict. So each time we will require dome
             # derivative_keys_dict
 
-        # but then if the user did not provide anything, we must at least have
-        # a default value for the dynamic_loss_dict keys entries in
-        # self.derivative_keys_dict since the computation of dynamic losses is
-        # made without create a loss object that would provide the
-        # default values
-        for k in self.dynamic_loss_dict.keys():
+        # derivative keys for the u_constraints. Note that we create missing
+        # DerivativeKeysODE around a Params object and not ParamsDict
+        # this works because u_dict.keys == params_dict.nn_params.keys()
+        for k in self.u_dict.keys():
             if self.derivative_keys_dict[k] is None:
-                try:
-                    if self.u_dict[k].eq_type == "statio_PDE":
-                        self.derivative_keys_dict[k] = DerivativeKeysPDEStatio()
-                    else:
-                        self.derivative_keys_dict[k] = DerivativeKeysPDENonStatio()
-                except KeyError:  # We are in a key that is not in u_dict but in
-                    # dynamic_loss_dict
-                    if isinstance(self.dynamic_loss_dict[k], PDEStatio):
-                        self.derivative_keys_dict[k] = DerivativeKeysPDEStatio()
-                    else:
-                        self.derivative_keys_dict[k] = DerivativeKeysPDENonStatio()
+                if self.u_dict[k].eq_type == "statio_PDE":
+                    self.derivative_keys_dict[k] = DerivativeKeysPDEStatio(
+                        params=params_dict.extract_params(k)
+                    )
+                else:
+                    self.derivative_keys_dict[k] = DerivativeKeysPDENonStatio(
+                        params=params_dict.extract_params(k)
+                    )
 
         # Second we make sure that all the dicts (except dynamic_loss_dict) have the same keys
         if (
@@ -904,16 +920,11 @@ class SystemLossPDE(eqx.Module):
                     f"got {self.u_dict[i].eq_type[i]}"
                 )
 
-        # for convenience in the tree_map of evaluate,
-        # we separate the two derivative keys dict
-        self.derivative_keys_dyn_loss_dict = {
-            k: self.derivative_keys_dict[k]
-            for k in self.dynamic_loss_dict.keys() & self.derivative_keys_dict.keys()
-        }
-        self.derivative_keys_u_dict = {
-            k: self.derivative_keys_dict[k]
-            for k in self.u_dict.keys() & self.derivative_keys_dict.keys()
-        }
+        # derivative keys for the dynamic loss. Note that we create a
+        # DerivativeKeysODE around a ParamsDict object because a whole
+        # params_dict is feed to DynamicLoss.evaluate functions (extract_params
+        # happen inside it)
+        self.derivative_keys_dyn_loss = DerivativeKeysPDENonStatio(params=params_dict)
 
         # also make sure we only have PINNs or SPINNs
         if not (
@@ -1005,19 +1016,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
@@ -1025,7 +1024,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]
@@ -1034,14 +1032,18 @@ class SystemLossPDE(eqx.Module):
             batch.param_batch_dict, params_dict
         )
 
-        def dyn_loss_for_one_key(dyn_loss, derivative_key, loss_weight):
+        def dyn_loss_for_one_key(dyn_loss, loss_weight):
             """The function used in tree_map"""
             return dynamic_loss_apply(
                 dyn_loss.evaluate,
                 self.u_dict,
-                batches,
-                _set_derivatives(params_dict, derivative_key.dyn_loss),
-                vmap_in_axes_x_or_x_t + vmap_in_axes_params,
+                (
+                    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 + vmap_in_axes_params,
                 loss_weight,
                 u_type=type(list(self.u_dict.values())[0]),
             )
@@ -1049,7 +1051,6 @@ class SystemLossPDE(eqx.Module):
         dyn_loss_mse_dict = jax.tree_util.tree_map(
             dyn_loss_for_one_key,
             self.dynamic_loss_dict,
-            self.derivative_keys_dyn_loss_dict,
             self._loss_weights["dyn_loss"],
             is_leaf=lambda x: isinstance(
                 x, (PDEStatio, PDENonStatio)

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,
 )