jinns 1.5.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

jinns/data/_CubicMeshPDEStatio.py

@@ -8,8 +8,11 @@ from __future__ import (
 import warnings
 import equinox as eqx
 import jax
+import numpy as np
 import jax.numpy as jnp
-from jaxtyping import Key, Array, Float
+from scipy.stats import qmc
+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
 from jinns.data._AbstractDataGenerator import AbstractDataGenerator
@@ -22,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
@@ -50,11 +53,13 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
         A tuple of maximum values of the domain along each dimension. For a sampling
         in `n` dimension, this represents $(x_{1, max}, x_{2,max}, ...,
         x_{n,max})$
-    method : str, default="uniform"
-        Either `grid` or `uniform`, default is `uniform`.
+    method : Literal["grid", "uniform", "sobol", "halton"], default="uniform"
+        Either "grid", "uniform", "sobol" or "halton", default is `uniform`.
         The method that generates the `nt` time points. `grid` means
         regularly spaced points over the domain. `uniform` means uniformly
-        sampled points over the domain
+        sampled points over the domain.
+        **Note** that Sobol and Halton approaches use scipy modules and will not
+        be JIT compatible.
     rar_parameters : dict[str, int], default=None
         Defaults to None: do not use Residual Adaptative Resampling.
         Otherwise a dictionary with keys
@@ -75,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: str = 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)
@@ -111,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)
 
@@ -120,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)
@@ -132,6 +158,22 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
                 )
             self.n = perfect_sq
 
+        if self.method in ["sobol", "halton"]:
+            log2_n = jnp.log2(self.n)
+            lower_pow = 2 ** jnp.floor(log2_n)
+            higher_pow = 2 ** jnp.ceil(log2_n)
+            closest_two_power = (
+                lower_pow
+                if (self.n - lower_pow) < (higher_pow - self.n)
+                else higher_pow
+            )
+            if self.n != closest_two_power:
+                warnings.warn(
+                    f"QuasiMonteCarlo sampling with {self.method} requires sample size to be a power fo 2."
+                    f"Modfiying self.n from {self.n} to {closest_two_power}.",
+                )
+                self.n = int(closest_two_power)
+
         if self.omega_batch_size is None:
             self.curr_omega_idx = 0
         else:
@@ -174,29 +216,53 @@ 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
+                )
+
+            return jnp.concatenate(
+                [
+                    jax.random.uniform(
+                        keys[i],
+                        (sample_size, 1),
+                        minval=self.min_pts[i],
+                        maxval=self.max_pts[i],
+                    )
+                    for i in range(self.dim)
+                ],
+                axis=-1,
+            )
+        else:
+            return self._qmc_in_omega_domain(keys[0], sample_size)
+
+    def _qmc_in_omega_domain(
+        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:
-            xmin, xmax = self.min_pts[0], self.max_pts[0]
-            return jax.random.uniform(
-                keys, shape=(sample_size, 1), minval=xmin, maxval=xmax
+            qmc_seq = qmc_generator(
+                d=self.dim,
+                scramble=True,
+                rng=np.random.default_rng(np.uint32(subkey)),
             )
-        # keys = jax.random.split(key, self.dim)
-        return jnp.concatenate(
-            [
-                jax.random.uniform(
-                    keys[i],
-                    (sample_size, 1),
-                    minval=self.min_pts[i],
-                    maxval=self.max_pts[i],
-                )
-                for i in range(self.dim)
-            ],
-            axis=-1,
+            u = qmc_seq.random(n=sample_size)
+            return jnp.array(
+                qmc.scale(u, l_bounds=self.min_pts[0], u_bounds=self.max_pts[0])
+            )
+        sampler = qmc.Sobol(
+            d=self.dim, scramble=True, rng=np.random.default_rng(np.uint32(subkey))
         )
+        samples = sampler.random(n=sample_size)
+        samples = qmc.scale(samples, l_bounds=self.min_pts, u_bounds=self.max_pts)
+        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:
@@ -206,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.
@@ -260,10 +327,67 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
             + f"implemented yet. You are asking for generation in dimension d={self.dim}."
         )
 
+    def qmc_in_omega_border_domain(
+        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
+        if sample_size is None:
+            return None
+        if self.dim == 1:
+            xmin = self.min_pts[0]
+            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.
+            facet_n = sample_size // (2 * self.dim)
+
+            def generate_qmc_sample(key, min_val, max_val):
+                qmc_seq = qmc_generator(
+                    d=1,
+                    scramble=True,
+                    rng=np.random.default_rng(np.uint32(key)),
+                )
+                u = qmc_seq.random(n=facet_n)
+                return jnp.array(qmc.scale(u, l_bounds=min_val, u_bounds=max_val))
+
+            xmin = jnp.hstack(
+                [
+                    self.min_pts[0] * jnp.ones((facet_n, 1)),
+                    generate_qmc_sample(keys[0], self.min_pts[1], self.max_pts[1]),
+                ]
+            )
+            xmax = jnp.hstack(
+                [
+                    self.max_pts[0] * jnp.ones((facet_n, 1)),
+                    generate_qmc_sample(keys[1], self.min_pts[1], self.max_pts[1]),
+                ]
+            )
+            ymin = jnp.hstack(
+                [
+                    generate_qmc_sample(keys[2], self.min_pts[0], self.max_pts[0]),
+                    self.min_pts[1] * jnp.ones((facet_n, 1)),
+                ]
+            )
+            ymax = jnp.hstack(
+                [
+                    generate_qmc_sample(keys[3], self.min_pts[0], self.max_pts[0]),
+                    self.max_pts[1] * jnp.ones((facet_n, 1)),
+                ]
+            )
+            return jnp.stack([xmin, xmax, ymin, ymax], axis=-1)
+        raise NotImplementedError(
+            "Generation of the border of a cube in dimension > 2 is not "
+            + f"implemented yet. You are asking for generation in dimension d={self.dim}."
+        )
+
     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"""
@@ -290,20 +414,21 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
                 )
                 xyz_ = [a.reshape((data_size, 1)) for a in xyz_]
                 omega = jnp.concatenate(xyz_, axis=-1)
-        elif self.method == "uniform":
-            if self.dim == 1:
-                key, subkeys = jax.random.split(key, 2)
+        elif self.method in ["uniform", "sobol", "halton"]:
+            if self.dim == 1 or self.method in ["sobol", "halton"]:
+                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"""
@@ -317,15 +442,24 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
             key, *subkeys = jax.random.split(key, 5)
         else:
             subkeys = None
-        omega_border = self.sample_in_omega_border_domain(
-            subkeys, sample_size=data_size
-        )
 
+        if self.method in ["grid", "uniform"]:
+            omega_border = self.sample_in_omega_border_domain(
+                subkeys, sample_size=data_size
+            )
+        elif self.method in ["sobol", "halton"]:
+            omega_border = self.qmc_in_omega_border_domain(
+                subkeys, sample_size=data_size
+            )
+        else:
+            raise ValueError("Method " + self.method + " is not implemented.")
         return key, omega_border
 
     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,
@@ -367,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(
@@ -379,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,
@@ -443,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,
         )

jinns/data/_DataGeneratorODE.py

@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Key, Array, Float
+from jaxtyping import PRNGKeyArray, Array, Float
 from jinns.data._Batchs import ODEBatch
 from jinns.data._utils import _check_and_set_rar_parameters, _reset_or_increment
 from jinns.data._AbstractDataGenerator import AbstractDataGenerator
@@ -24,7 +24,7 @@ class DataGeneratorODE(AbstractDataGenerator):
 
     Parameters
     ----------
-    key : Key
+    key : PRNGKeyArray
         Jax random key to sample new time points and to shuffle batches
     nt : int
         The number of total time points that will be divided in
@@ -42,10 +42,10 @@ class DataGeneratorODE(AbstractDataGenerator):
         The method that generates the `nt` time points. `grid` means
         regularly spaced points over the domain. `uniform` means uniformly
         sampled points over the domain
-    rar_parameters : RarParameterDict, default=None
+    rar_parameters : None | RarParameterDict, default=None
        A TypedDict to specify the Residual Adaptative Resampling procedure. See
        the docstring from RarParameterDict
-    n_start : int, default=None
+    n_start : None | int, default=None
         Defaults to None. The effective size of nt used at start time.
         This value must be
         provided when rar_parameters is not None. Otherwise we set internally
@@ -54,25 +54,43 @@ class DataGeneratorODE(AbstractDataGenerator):
         then corresponds to the initial number of points we train the PINN.
     """
 
-    key: Key = eqx.field(kw_only=True)
-    nt: int = eqx.field(kw_only=True, static=True)
-    tmin: Float = eqx.field(kw_only=True)
-    tmax: Float = eqx.field(kw_only=True)
-    temporal_batch_size: int | None = eqx.field(static=True, default=None, kw_only=True)
-    method: str = eqx.field(
-        static=True, kw_only=True, default_factory=lambda: "uniform"
-    )
-    rar_parameters: dict[str, int] = eqx.field(default=None, kw_only=True)
-    n_start: int = eqx.field(static=True, default=None, kw_only=True)
-
-    # all the init=False fields are set in __post_init__
+    key: PRNGKeyArray
+    nt: int = eqx.field(static=True)
+    tmin: float
+    tmax: float
+    temporal_batch_size: int | None = eqx.field(static=True)
+    method: str = eqx.field(static=True)
+    rar_parameters: None | dict[str, int]
+    n_start: None | int
+
+    # --- Below fields are not passed as arguments to __init__
     p: Float[Array, " nt 1"] | 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)
     curr_time_idx: int = eqx.field(init=False)
     times: Float[Array, " nt 1"] = eqx.field(init=False)
 
-    def __post_init__(self):
+    def __init__(
+        self,
+        *,
+        key: PRNGKeyArray,
+        nt: int,
+        tmin: float,
+        tmax: float,
+        temporal_batch_size: int | None,
+        method: str = "uniform",
+        rar_parameters: None | dict[str, int] = None,
+        n_start: None | int = None,
+    ):
+        self.key = key
+        self.nt = nt
+        self.tmin = tmin
+        self.tmax = tmax
+        self.temporal_batch_size = temporal_batch_size
+        self.method = method
+        self.n_start = n_start
+        self.rar_parameters = rar_parameters
+
         (
             self.n_start,
             self.p,
@@ -97,7 +115,7 @@ class DataGeneratorODE(AbstractDataGenerator):
         # above way for the key.
 
     def sample_in_time_domain(
-        self, key: Key, sample_size: int | None = None
+        self, key: PRNGKeyArray, sample_size: int | None = None
     ) -> Float[Array, " nt 1"]:
         return jax.random.uniform(
             key,
@@ -106,7 +124,9 @@ class DataGeneratorODE(AbstractDataGenerator):
             maxval=self.tmax,
         )
 
-    def generate_time_data(self, key: Key) -> tuple[Key, Float[Array, " nt"]]:
+    def generate_time_data(
+        self, key: PRNGKeyArray
+    ) -> tuple[PRNGKeyArray, Float[Array, " nt"]]:
         """
         Construct a complete set of `self.nt` time points according to the
         specified `self.method`
@@ -125,7 +145,11 @@ class DataGeneratorODE(AbstractDataGenerator):
     def _get_time_operands(
         self,
     ) -> tuple[
-        Key, Float[Array, " nt 1"], int, int | None, Float[Array, " nt 1"] | None
+        PRNGKeyArray,
+        Float[Array, " nt 1"],
+        int,
+        int | None,
+        Float[Array, " nt 1"] | None,
     ]:
         return (
             self.key,
@@ -150,7 +174,7 @@ class DataGeneratorODE(AbstractDataGenerator):
         bend = bstart + self.temporal_batch_size
 
         # Compute the effective number of used collocation points
-        if self.rar_parameters is not None:
+        if self.rar_parameters is not None and self.n_start is not None:
             nt_eff = (
                 self.n_start
                 + self.rar_iter_nb  # type: ignore
@@ -167,7 +191,9 @@ class DataGeneratorODE(AbstractDataGenerator):
             # handled above
         )
         new = eqx.tree_at(
-            lambda m: (m.key, m.times, m.curr_time_idx), self, new_attributes
+            lambda m: (m.key, m.times, m.curr_time_idx),  # type: ignore
+            self,
+            new_attributes,
         )
 
         # commands below are equivalent to