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

jinns/data/_AbstractDataGenerator.py

@@ -15,5 +15,5 @@ class AbstractDataGenerator(eqx.Module):
     """
 
     @abc.abstractmethod
-    def get_batch(self) -> tuple[type[Self], AnyBatch]:  # type: ignore
+    def get_batch(self) -> tuple[Self, AnyBatch]:
         pass

jinns/data/_Batchs.py

@@ -18,25 +18,59 @@ class ObsBatchDict(TypedDict):
 
     pinn_in: Float[Array, "  obs_batch_size input_dim"]
     val: Float[Array, "  obs_batch_size output_dim"]
-    eq_params: dict[str, Float[Array, "  obs_batch_size 1"]]
+    eq_params: (
+        eqx.Module | None
+    )  # None cause sometime user don't provide observed params
 
 
 class ODEBatch(eqx.Module):
     temporal_batch: Float[Array, "  batch_size"]
-    param_batch_dict: dict[str, Array] = eqx.field(default=None)
-    obs_batch_dict: ObsBatchDict = eqx.field(default=None)
-
-
-class PDENonStatioBatch(eqx.Module):
-    domain_batch: Float[Array, "  batch_size 1+dimension"]
-    border_batch: Float[Array, "  batch_size dimension n_facets"] | None
-    initial_batch: Float[Array, "  batch_size dimension"] | None
-    param_batch_dict: dict[str, Array] = eqx.field(default=None)
-    obs_batch_dict: ObsBatchDict = eqx.field(default=None)
+    param_batch_dict: eqx.Module | None = eqx.field(default=None)
+    obs_batch_dict: ObsBatchDict | None = eqx.field(default=None)
 
 
 class PDEStatioBatch(eqx.Module):
     domain_batch: Float[Array, "  batch_size dimension"]
     border_batch: Float[Array, "  batch_size dimension n_facets"] | None
-    param_batch_dict: dict[str, Array] = eqx.field(default=None)
-    obs_batch_dict: ObsBatchDict = eqx.field(default=None)
+    param_batch_dict: eqx.Module | None
+    obs_batch_dict: ObsBatchDict | None
+
+    # rewrite __init__ to be able to use inheritance for the NonStatio case
+    # below. That way PDENonStatioBatch is a subtype of PDEStatioBatch which
+    # 1) makes more sense and 2) CubicMeshPDENonStatio.get_batch passes pyright.
+    def __init__(
+        self,
+        *,
+        domain_batch: Float[Array, "  batch_size dimension"],
+        border_batch: Float[Array, "  batch_size dimension n_facets"] | None,
+        param_batch_dict: eqx.Module | None = None,
+        obs_batch_dict: ObsBatchDict | None = None,
+    ):
+        # TODO: document this ?
+        self.domain_batch = domain_batch
+        self.border_batch = border_batch
+        self.param_batch_dict = param_batch_dict
+        self.obs_batch_dict = obs_batch_dict
+
+
+class PDENonStatioBatch(PDEStatioBatch):
+    # TODO: document this ?
+    domain_batch: Float[Array, "  batch_size 1+dimension"]  # Override type
+    initial_batch: (
+        Float[Array, "  batch_size dimension"] | None
+    )  # why can it be None ? Examples?
+
+    def __init__(
+        self,
+        *,
+        domain_batch: Float[Array, "  batch_size 1+dimension"],
+        border_batch: Float[Array, "  batch_size dimension n_facets"] | None,
+        initial_batch: Float[Array, "  batch_size dimension"] | None,
+        param_batch_dict: eqx.Module | None = None,
+        obs_batch_dict: ObsBatchDict | None = None,
+    ):
+        self.domain_batch = domain_batch
+        self.border_batch = border_batch
+        self.initial_batch = initial_batch
+        self.param_batch_dict = param_batch_dict
+        self.obs_batch_dict = obs_batch_dict

jinns/data/_CubicMeshPDENonStatio.py

@@ -11,7 +11,7 @@ import numpy as np
 import jax
 import jax.numpy as jnp
 from scipy.stats import qmc
-from jaxtyping import Key, Array, Float
+from jaxtyping import PRNGKeyArray, Array, Float
 from jinns.data._Batchs import PDENonStatioBatch
 from jinns.data._utils import (
     make_cartesian_product,
@@ -29,7 +29,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
 
     Parameters
     ----------
-    key : Key
+    key : PRNGKeyArray
         Jax random key to sample new time points and to shuffle batches
     n : int
         The number of total $I\times \Omega$ points that will be divided in
@@ -50,9 +50,8 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         among the `nb` points. If None, `domain_batch_size` no
         mini-batches are used.
     initial_batch_size : int | None, default=None
-        The size of the batch of randomly selected points among
-        the `ni` points. If None no
-        mini-batches are used.
+        The number of randomly selected points among the `ni` initial spatial
+        points used for initial condition. If None, no mini-batches are used.
     dim : int
         An integer. Dimension of $\Omega$ domain.
     min_pts : tuple[tuple[Float, Float], ...]
@@ -94,13 +93,14 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         then corresponds to the initial number of omega points we train the PINN.
     """
 
-    tmin: Float = eqx.field(kw_only=True)
-    tmax: Float = eqx.field(kw_only=True)
-    ni: int = eqx.field(kw_only=True, static=True)
-    domain_batch_size: int | None = eqx.field(kw_only=True, static=True, default=None)
-    initial_batch_size: int | None = eqx.field(kw_only=True, static=True, default=None)
-    border_batch_size: int | None = eqx.field(kw_only=True, static=True, default=None)
+    tmin: float
+    tmax: float
+    ni: int = eqx.field(static=True)
+    domain_batch_size: int | None = eqx.field(static=True)
+    initial_batch_size: int | None = eqx.field(static=True)
+    border_batch_size: int | None = eqx.field(static=True)
 
+    # --- Below fields are not passed as arguments to __init__
     curr_domain_idx: int = eqx.field(init=False)
     curr_initial_idx: int = eqx.field(init=False)
     curr_border_idx: int = eqx.field(init=False)
@@ -110,13 +110,32 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
     )
     initial: Float[Array, " ni dim"] | None = eqx.field(init=False)
 
-    def __post_init__(self):
+    def __init__(
+        self,
+        tmin: float,
+        tmax: float,
+        ni: int,
+        domain_batch_size: int | None = None,
+        initial_batch_size: int | None = None,
+        border_batch_size: int | None = None,
+        **kwargs,  # kwargs for CubicMeshPDEStatio.__init__
+    ):
         """
         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
-        # class is not automatically called
+        # sanity check
+        if ni is None:
+            raise ValueError("`ni` cannot be None.")
+
+        super().__init__(**kwargs)
+        self.tmin = tmin
+        self.tmax = tmax
+        self.ni = ni
+
+        self.domain_batch_size = domain_batch_size
+        self.initial_batch_size = initial_batch_size
+        self.border_batch_size = border_batch_size
 
         if self.method == "grid":
             # NOTE we must redo the sampling with the square root number of samples
@@ -144,7 +163,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             )
             self.domain = make_cartesian_product(half_domain_times, half_domain_omega)
 
-            # NOTE
+            # NOTE below re-do CubicMeshPDE.__init__() ? Maybe useless?
             (
                 self.n_start,
                 self.p,
@@ -178,7 +197,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
                     " a multiple of 2xd (the # of faces of a d-dimensional cube)"
                 )
             # the check below concern omega_border_batch_size for dim > 1 in
-            # super.__post_init__. Here it concerns all dim values since our
+            # super.__init__. Here it concerns all dim values since our
             # border_batch is the concatenation or cartesian product with times
             if (
                 self.border_batch_size is not None
@@ -221,7 +240,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             self.border_batch_size = None
             self.curr_border_idx = 0
 
-        if self.ni is not None:
+        if ni is not None:
             if self.method == "grid":
                 perfect_sq = int(jnp.round(jnp.sqrt(self.ni)) ** 2)
                 if self.ni != perfect_sq:
@@ -235,17 +254,17 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
                 log2_n = jnp.log2(self.ni)
                 lower_pow = 2 ** jnp.floor(log2_n)
                 higher_pow = 2 ** jnp.ceil(log2_n)
-                closest_two_power = (
+                closest_power_of_two = (
                     lower_pow
                     if (self.ni - lower_pow) < (higher_pow - self.ni)
                     else higher_pow
                 )
-                if self.n != closest_two_power:
+                if self.n != closest_power_of_two:
                     warnings.warn(
                         f"QuasiMonteCarlo sampling with {self.method} requires sample size to be a power fo 2."
-                        f"Modfiying self.n from {self.ni} to {closest_two_power}.",
+                        f"Modfiying self.n from {self.ni} to {closest_power_of_two}.",
                     )
-                self.ni = int(closest_two_power)
+                self.ni = int(closest_power_of_two)
             self.key, self.initial = self.generate_omega_data(
                 self.key, data_size=self.ni
             )
@@ -264,8 +283,8 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         self.omega_border = None
 
     def generate_time_data(
-        self, key: Key, nt: int
-    ) -> tuple[Key, Float[Array, " nt 1"]]:
+        self, key: PRNGKeyArray, nt: int
+    ) -> tuple[PRNGKeyArray, Float[Array, " nt 1"]]:
         """
         Construct a complete set of `nt` time points according to the
         specified `self.method`
@@ -278,12 +297,14 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             return key, self.sample_in_time_domain(subkey, nt)
         raise ValueError("Method " + self.method + " is not implemented.")
 
-    def sample_in_time_domain(self, key: Key, nt: int) -> Float[Array, " nt 1"]:
+    def sample_in_time_domain(
+        self, key: PRNGKeyArray, nt: int
+    ) -> Float[Array, " nt 1"]:
         return jax.random.uniform(key, (nt, 1), minval=self.tmin, maxval=self.tmax)
 
     def qmc_in_time_omega_domain(
-        self, key: Key, sample_size: int
-    ) -> tuple[Key, Float[Array, "n 1+dim"]]:
+        self, key: PRNGKeyArray, sample_size: int
+    ) -> tuple[PRNGKeyArray, Float[Array, "n 1+dim"]]:
         """
         Because in Quasi-Monte Carlo sampling we cannot concatenate two vectors generated independently
         We generate time and omega samples jointly
@@ -300,8 +321,8 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         return key, jnp.array(samples)
 
     def qmc_in_time_omega_border_domain(
-        self, key: Key, sample_size: int | None = None
-    ) -> tuple[Key, Float[Array, "n 1+dim"]] | None:
+        self, key: PRNGKeyArray, sample_size: int | None = None
+    ) -> tuple[PRNGKeyArray, Float[Array, "n 1+dim"]] | None:
         """
         For each facet of the border we generate Quasi-MonteCarlo sequences jointy with time.
 
@@ -387,7 +408,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
 
     def _get_domain_operands(
         self,
-    ) -> tuple[Key, Float[Array, " n 1+dim"], int, int | None, Array | None]:
+    ) -> tuple[PRNGKeyArray, Float[Array, " n 1+dim"], int, int | None, Array | None]:
         return (
             self.key,
             self.domain,
@@ -424,7 +445,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             # handled above
         )
         new = eqx.tree_at(
-            lambda m: (m.key, m.domain, m.curr_domain_idx),
+            lambda m: (m.key, m.domain, m.curr_domain_idx),  # type: ignore
             self,
             new_attributes,
         )
@@ -437,7 +458,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
     def _get_border_operands(
         self,
     ) -> tuple[
-        Key,
+        PRNGKeyArray,
         Float[Array, " nb 1+1 2"] | Float[Array, " (nb//4) 2+1 4"] | None,
         int,
         int | None,
@@ -483,7 +504,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             # handled above
         )
         new = eqx.tree_at(
-            lambda m: (m.key, m.border, m.curr_border_idx),
+            lambda m: (m.key, m.border, m.curr_border_idx),  # type: ignore
             self,
             new_attributes,
         )
@@ -500,7 +521,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
 
     def _get_initial_operands(
         self,
-    ) -> tuple[Key, Float[Array, " ni dim"] | None, int, int | None, None]:
+    ) -> tuple[PRNGKeyArray, Float[Array, " ni dim"] | None, int, int | None, None]:
         return (
             self.key,
             self.initial,
@@ -529,7 +550,7 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             # handled above
         )
         new = eqx.tree_at(
-            lambda m: (m.key, m.initial, m.curr_initial_idx),
+            lambda m: (m.key, m.initial, m.curr_initial_idx),  # type: ignore
             self,
             new_attributes,
         )

jinns/data/_CubicMeshPDEStatio.py

@@ -11,7 +11,7 @@ import jax
 import numpy as np
 import jax.numpy as jnp
 from scipy.stats import qmc
-from jaxtyping import Key, Array, Float
+from jaxtyping import PRNGKeyArray, Array, Float
 from typing import Literal
 from jinns.data._Batchs import PDEStatioBatch
 from jinns.data._utils import _check_and_set_rar_parameters, _reset_or_increment
@@ -25,7 +25,7 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
 
     Parameters
     ----------
-    key : Key
+    key : PRNGKeyArray
         Jax random key to sample new time points and to shuffle batches
     n : int
         The number of total $\Omega$ points that will be divided in
@@ -80,32 +80,28 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
         then corresponds to the initial number of points we train the PINN on.
     """
 
-    # kw_only in base class is motivated here: https://stackoverflow.com/a/69822584
-    key: Key = eqx.field(kw_only=True)
-    n: int = eqx.field(kw_only=True, static=True)
+    key: PRNGKeyArray
+    n: int = eqx.field(static=True)
     nb: int | None = eqx.field(kw_only=True, static=True, default=None)
     omega_batch_size: int | None = eqx.field(
-        kw_only=True,
         static=True,
-        default=None,  # can be None as
+        # can be None as
         # CubicMeshPDENonStatio inherits but also if omega_batch_size=n
     )  # static cause used as a
     # shape in jax.lax.dynamic_slice
     omega_border_batch_size: int | None = eqx.field(
-        kw_only=True, static=True, default=None
+        static=True,
     )  # static cause used as a
     # shape in jax.lax.dynamic_slice
-    dim: int = eqx.field(kw_only=True, static=True)  # static cause used as a
+    dim: int = eqx.field(static=True)  # static cause used as a
     # shape in jax.lax.dynamic_slice
-    min_pts: tuple[float, ...] = eqx.field(kw_only=True)
-    max_pts: tuple[float, ...] = eqx.field(kw_only=True)
-    method: Literal["grid", "uniform", "sobol", "halton"] = eqx.field(
-        kw_only=True, static=True, default_factory=lambda: "uniform"
-    )
-    rar_parameters: dict[str, int] = eqx.field(kw_only=True, default=None)
-    n_start: int = eqx.field(kw_only=True, default=None, static=True)
+    min_pts: tuple[float, ...]
+    max_pts: tuple[float, ...]
+    method: Literal["grid", "uniform", "sobol", "halton"] = eqx.field(static=True)
+    rar_parameters: None | dict[str, int]
+    n_start: int = eqx.field(static=True)
 
-    # all the init=False fields are set in __post_init__
+    # --- Below fields are not passed as arguments to __init__
     p: Float[Array, " n"] | None = eqx.field(init=False)
     rar_iter_from_last_sampling: int | None = eqx.field(init=False)
     rar_iter_nb: int | None = eqx.field(init=False)
@@ -116,7 +112,32 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
         eqx.field(init=False)
     )
 
-    def __post_init__(self):
+    def __init__(
+        self,
+        *,
+        key: PRNGKeyArray,
+        n: int,
+        nb: int | None = None,
+        omega_batch_size: int | None = None,
+        omega_border_batch_size: int | None = None,
+        dim: int,
+        min_pts: tuple[float, ...],
+        max_pts: tuple[float, ...],
+        method: Literal["grid", "uniform", "sobol", "halton"] = "uniform",
+        rar_parameters: dict[str, int] | None = None,
+        n_start: int | None = None,
+    ):
+        self.key = key
+        self.n = n
+        self.nb = nb
+        self.omega_batch_size = omega_batch_size
+        self.omega_border_batch_size = omega_border_batch_size
+        self.dim = dim
+        self.min_pts = min_pts
+        self.max_pts = max_pts
+        self.method = method
+        self.rar_parameters = rar_parameters
+
         assert self.dim == len(self.min_pts) and isinstance(self.min_pts, tuple)
         assert self.dim == len(self.max_pts) and isinstance(self.max_pts, tuple)
 
@@ -125,7 +146,7 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
             self.p,
             self.rar_iter_from_last_sampling,
             self.rar_iter_nb,
-        ) = _check_and_set_rar_parameters(self.rar_parameters, self.n, self.n_start)
+        ) = _check_and_set_rar_parameters(self.rar_parameters, self.n, n_start)
 
         if self.method == "grid" and self.dim == 2:
             perfect_sq = int(jnp.round(jnp.sqrt(self.n)) ** 2)
@@ -195,13 +216,13 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
         self.key, self.omega_border = self.generate_omega_border_data(self.key)
 
     def sample_in_omega_domain(
-        self, keys: Key, sample_size: int
+        self, keys: list[PRNGKeyArray], sample_size: int
     ) -> Float[Array, " n dim"]:
         if self.method == "uniform":
             if self.dim == 1:
                 xmin, xmax = self.min_pts[0], self.max_pts[0]
                 return jax.random.uniform(
-                    keys, shape=(sample_size, 1), minval=xmin, maxval=xmax
+                    *keys, shape=(sample_size, 1), minval=xmin, maxval=xmax
                 )
 
             return jnp.concatenate(
@@ -217,10 +238,10 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
                 axis=-1,
             )
         else:
-            return self._qmc_in_omega_domain(keys, sample_size)
+            return self._qmc_in_omega_domain(keys[0], sample_size)
 
     def _qmc_in_omega_domain(
-        self, subkey: Key, sample_size: int
+        self, subkey: PRNGKeyArray, sample_size: int
     ) -> Float[Array, "n dim"]:
         qmc_generator = qmc.Sobol if self.method == "sobol" else qmc.Halton
         if self.dim == 1:
@@ -241,7 +262,7 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
         return jnp.array(samples)
 
     def sample_in_omega_border_domain(
-        self, keys: Key, sample_size: int | None = None
+        self, keys: list[PRNGKeyArray] | None, sample_size: int | None = None
     ) -> Float[Array, " 1 2"] | Float[Array, " (nb//4) 2 4"] | None:
         sample_size = self.nb if sample_size is None else sample_size
         if sample_size is None:
@@ -251,6 +272,7 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
             xmax = self.max_pts[0]
             return jnp.array([xmin, xmax]).astype(float)
         if self.dim == 2:
+            assert keys is not None
             # currently hard-coded the 4 edges for d==2
             # TODO : find a general & efficient way to sample from the border
             # (facets) of the hypercube in general dim.
@@ -306,7 +328,7 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
         )
 
     def qmc_in_omega_border_domain(
-        self, keys: Key, sample_size: int | None = None
+        self, keys: list[PRNGKeyArray] | None, sample_size: int | None = None
     ) -> Float[Array, " 1 2"] | Float[Array, " (nb//4) 2 4"] | None:
         qmc_generator = qmc.Sobol if self.method == "sobol" else qmc.Halton
         sample_size = self.nb if sample_size is None else sample_size
@@ -317,6 +339,7 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
             xmax = self.max_pts[0]
             return jnp.array([xmin, xmax]).astype(float)
         if self.dim == 2:
+            assert keys is not None
             # currently hard-coded the 4 edges for d==2
             # TODO : find a general & efficient way to sample from the border
             # (facets) of the hypercube in general dim.
@@ -362,9 +385,9 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
         )
 
     def generate_omega_data(
-        self, key: Key, data_size: int | None = None
+        self, key: PRNGKeyArray, data_size: int | None = None
     ) -> tuple[
-        Key,
+        PRNGKeyArray,
         Float[Array, " n dim"],
     ]:
         r"""
@@ -393,18 +416,19 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
                 omega = jnp.concatenate(xyz_, axis=-1)
         elif self.method in ["uniform", "sobol", "halton"]:
             if self.dim == 1 or self.method in ["sobol", "halton"]:
-                key, subkeys = jax.random.split(key, 2)
+                key, subkey = jax.random.split(key, 2)
+                omega = self.sample_in_omega_domain([subkey], sample_size=data_size)
             else:
                 key, *subkeys = jax.random.split(key, self.dim + 1)
-            omega = self.sample_in_omega_domain(subkeys, sample_size=data_size)
+                omega = self.sample_in_omega_domain(subkeys, sample_size=data_size)
         else:
             raise ValueError("Method " + self.method + " is not implemented.")
         return key, omega
 
     def generate_omega_border_data(
-        self, key: Key, data_size: int | None = None
+        self, key: PRNGKeyArray, data_size: int | None = None
     ) -> tuple[
-        Key,
+        PRNGKeyArray,
         Float[Array, " 1 2"] | Float[Array, " (nb//4) 2 4"] | None,
     ]:
         r"""
@@ -433,7 +457,9 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
 
     def _get_omega_operands(
         self,
-    ) -> tuple[Key, Float[Array, " n dim"], int, int | None, Float[Array, " n"] | None]:
+    ) -> tuple[
+        PRNGKeyArray, Float[Array, " n dim"], int, int | None, Float[Array, " n"] | None
+    ]:
         return (
             self.key,
             self.omega,
@@ -475,7 +501,9 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
             # handled above
         )
         new = eqx.tree_at(
-            lambda m: (m.key, m.omega, m.curr_omega_idx), self, new_attributes
+            lambda m: (m.key, m.omega, m.curr_omega_idx),  # type: ignore
+            self,
+            new_attributes,
         )
 
         return new, jax.lax.dynamic_slice(
@@ -487,7 +515,7 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
     def _get_omega_border_operands(
         self,
     ) -> tuple[
-        Key,
+        PRNGKeyArray,
         Float[Array, " 1 2"] | Float[Array, " (nb//4) 2 4"] | None,
         int,
         int | None,
@@ -551,7 +579,7 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
             # handled above
         )
         new = eqx.tree_at(
-            lambda m: (m.key, m.omega_border, m.curr_omega_border_idx),
+            lambda m: (m.key, m.omega_border, m.curr_omega_border_idx),  # type: ignore
             self,
             new_attributes,
         )