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

jinns/data/_DataGenerators.py

@@ -54,7 +54,7 @@ def make_cartesian_product(
 
 
 def _reset_batch_idx_and_permute(
-    operands: tuple[Key, Float[Array, "n dimension"], Int, None, Float[Array, "n"]]
+    operands: tuple[Key, Float[Array, "n dimension"], Int, None, Float[Array, "n"]],
 ) -> tuple[Key, Float[Array, "n dimension"], Int]:
     key, domain, curr_idx, _, p = operands
     # resetting counter
@@ -78,7 +78,7 @@ def _reset_batch_idx_and_permute(
 
 
 def _increment_batch_idx(
-    operands: tuple[Key, Float[Array, "n dimension"], Int, None, Float[Array, "n"]]
+    operands: tuple[Key, Float[Array, "n dimension"], Int, None, Float[Array, "n"]],
 ) -> tuple[Key, Float[Array, "n dimension"], Int]:
     key, domain, curr_idx, batch_size, _ = operands
     # simply increases counter and get the batch

jinns/loss/_DynamicLoss.py

@@ -13,8 +13,8 @@ from jax import grad
 import jax.numpy as jnp
 import equinox as eqx
 
-from jinns.utils._pinn import PINN
-from jinns.utils._spinn import SPINN
+from jinns.nn._pinn import PINN
+from jinns.nn._spinn_mlp import SPINN
 
 from jinns.utils._utils import get_grid
 from jinns.loss._DynamicLossAbstract import ODE, PDEStatio, PDENonStatio

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 *

jinns/loss/_LossPDE.py

@@ -38,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
 
 
@@ -92,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.
@@ -127,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)
@@ -251,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(
@@ -324,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.
@@ -431,7 +458,7 @@ class LossPDEStatio(_LossPDEAbstract):
                 self._get_normalization_loss_batch(batch),
                 _set_derivatives(params, self.derivative_keys.norm_loss),
                 vmap_in_axes_params,
-                self.norm_int_length,
+                self.norm_weights,
                 self.loss_weights.norm_loss,
             )
         else:
@@ -549,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.
@@ -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:
@@ -1045,7 +1082,7 @@ class SystemLossPDE(eqx.Module):
                 _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]),
+                u_type=list(self.u_dict.values())[0].__class__.__base__,
             )
 
         dyn_loss_mse_dict = jax.tree_util.tree_map(

jinns/loss/_boundary_conditions.py

@@ -13,8 +13,8 @@ from jax import vmap, grad
 import equinox as eqx
 from jinns.utils._utils import get_grid, _subtract_with_check
 from jinns.data._Batchs import *
-from jinns.utils._pinn import PINN
-from jinns.utils._spinn import SPINN
+from jinns.nn._pinn import PINN
+from jinns.nn._spinn import SPINN
 
 if TYPE_CHECKING:
     from jinns.utils._types import *

jinns/loss/_loss_utils.py

@@ -19,9 +19,9 @@ from jinns.loss._boundary_conditions import (
 from jinns.utils._utils import _subtract_with_check, get_grid
 from jinns.data._DataGenerators import append_obs_batch, make_cartesian_product
 from jinns.parameters._params import _get_vmap_in_axes_params
-from jinns.utils._pinn import PINN
-from jinns.utils._spinn import SPINN
-from jinns.utils._hyperpinn import HYPERPINN
+from jinns.nn._pinn import PINN
+from jinns.nn._spinn import SPINN
+from jinns.nn._hyperpinn import HyperPINN
 from jinns.data._Batchs import *
 from jinns.parameters._params import Params, ParamsDict
 
@@ -40,13 +40,13 @@ def dynamic_loss_apply(
     params: Params | ParamsDict,
     vmap_axes: tuple[int | None, ...],
     loss_weight: float | Float[Array, "dyn_loss_dimension"],
-    u_type: PINN | HYPERPINN | None = None,
+    u_type: PINN | HyperPINN | None = None,
 ) -> float:
     """
     Sometimes when u is a lambda function a or dict we do not have access to
     its type here, hence the last argument
     """
-    if u_type == PINN or u_type == HYPERPINN or isinstance(u, (PINN, HYPERPINN)):
+    if u_type == PINN or u_type == HyperPINN or isinstance(u, (PINN, HyperPINN)):
         v_dyn_loss = vmap(
             lambda batch, params: dyn_loss(
                 batch, u, params  # we must place the params at the end
@@ -75,23 +75,25 @@ def normalization_loss_apply(
     ),
     params: Params | ParamsDict,
     vmap_axes_params: tuple[int | None, ...],
-    int_length: int,
+    norm_weights: Float[Array, "nb_norm_samples"],
     loss_weight: float,
 ) -> float:
     """
     Note the squeezing on each result. We expect unidimensional *PINN since
     they represent probability distributions
     """
-    if isinstance(u, (PINN, HYPERPINN)):
+    if isinstance(u, (PINN, HyperPINN)):
         if len(batches) == 1:
             v_u = vmap(
-                lambda b: u(b)[u.slice_solution],
+                lambda *b: u(*b)[u.slice_solution],
                 (0,) + vmap_axes_params,
                 0,
             )
             res = v_u(*batches, params)
+            assert res.shape[-1] == 1, "norm loss expects unidimensional *PINN"
+            # Monte-Carlo integration using importance sampling
             mse_norm_loss = loss_weight * (
-                jnp.abs(jnp.mean(res.squeeze()) * int_length - 1) ** 2
+                jnp.abs(jnp.mean(res.squeeze() * norm_weights) - 1) ** 2
             )
         else:
             # NOTE this cartesian product is costly
@@ -107,20 +109,23 @@ def normalization_loss_apply(
                 in_axes=(0,) + vmap_axes_params,
             )
             res = v_u(batches, params)
-            # Over all the times t, we perform a integration
+            assert res.shape[-1] == 1, "norm loss expects unidimensional *PINN"
+            # For all times t, we perform an integration. Then we average the
+            # losses over times.
             mse_norm_loss = loss_weight * jnp.mean(
-                jnp.abs(jnp.mean(res.squeeze(), axis=-1) * int_length - 1) ** 2
+                jnp.abs(jnp.mean(res.squeeze() * norm_weights, axis=-1) - 1) ** 2
             )
     elif isinstance(u, SPINN):
         if len(batches) == 1:
             res = u(*batches, params)
+            assert res.shape[-1] == 1, "norm loss expects unidimensional *SPINN"
             mse_norm_loss = (
                 loss_weight
                 * jnp.abs(
                     jnp.mean(
                         res.squeeze(),
                     )
-                    * int_length
+                    * norm_weights
                     - 1
                 )
                 ** 2
@@ -134,6 +139,7 @@ def normalization_loss_apply(
                 ),
                 params,
             )
+            assert res.shape[-1] == 1, "norm loss expects unidimensional *SPINN"
             # the outer mean() below is for the times stamps
             mse_norm_loss = loss_weight * jnp.mean(
                 jnp.abs(
@@ -141,7 +147,7 @@ def normalization_loss_apply(
                         res.squeeze(),
                         axis=(d + 1 for d in range(res.ndim - 2)),
                     )
-                    * int_length
+                    * norm_weights
                     - 1
                 )
                 ** 2
@@ -230,7 +236,7 @@ def observations_loss_apply(
     obs_slice: slice,
 ) -> float:
     # TODO implement for SPINN
-    if isinstance(u, (PINN, HYPERPINN)):
+    if isinstance(u, (PINN, HyperPINN)):
         v_u = vmap(
             lambda *args: u(*args)[u.slice_solution],
             vmap_axes,
@@ -264,7 +270,7 @@ def initial_condition_apply(
 ) -> float:
     n = omega_batch.shape[0]
     t0_omega_batch = jnp.concatenate([jnp.zeros((n, 1)), omega_batch], axis=1)
-    if isinstance(u, (PINN, HYPERPINN)):
+    if isinstance(u, (PINN, HyperPINN)):
         v_u_t0 = vmap(
             lambda t0_x, params: _subtract_with_check(
                 initial_condition_fun(t0_x[1:]),

jinns/loss/_operators.py

@@ -9,8 +9,6 @@ import jax.numpy as jnp
 from jax import grad
 import equinox as eqx
 from jaxtyping import Float, Array
-from jinns.utils._pinn import PINN
-from jinns.utils._spinn import SPINN
 from jinns.parameters._params import Params
 
 

jinns/nn/__init__.py

@@ -0,0 +1,7 @@
+from ._save_load import save_pinn, load_pinn
+from ._pinn import PINN
+from ._spinn import SPINN
+from ._mlp import PINN_MLP, MLP
+from ._spinn_mlp import SPINN_MLP, SMLP
+from ._hyperpinn import HyperPINN
+from ._ppinn import PPINN_MLP