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

jinns/__init__.py

@@ -1,10 +1,3 @@
-# import jinns.data
-# import jinns.loss
-# import jinns.solver
-# import jinns.utils
-# import jinns.experimental
-# import jinns.parameters
-# import jinns.plot
 from jinns import data as data
 from jinns import loss as loss
 from jinns import solver as solver
@@ -16,3 +9,10 @@ from jinns import nn as nn
 from jinns.solver._solve import solve
 
 __all__ = ["nn", "solve"]
+
+import warnings
+
+warnings.filterwarnings(
+    action="ignore",
+    message=r"Using `field\(init=False\)`",
+)

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

@@ -7,9 +7,11 @@ from __future__ import (
 )  # https://docs.python.org/3/library/typing.html#constant
 import warnings
 import equinox as eqx
+import numpy as np
 import jax
 import jax.numpy as jnp
-from jaxtyping import Key, Array, Float
+from scipy.stats import qmc
+from jaxtyping import PRNGKeyArray, Array, Float
 from jinns.data._Batchs import PDENonStatioBatch
 from jinns.data._utils import (
     make_cartesian_product,
@@ -27,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
@@ -48,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], ...]
@@ -65,11 +66,13 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         The minimum value of the time domain to consider
     tmax : float
         The maximum value of the time domain to consider
-    method : str, default="uniform"
+    method : Literal["uniform", "grid", "sobol", "halton"], default="uniform"
         Either `grid` or `uniform`, 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
@@ -90,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)
@@ -106,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
@@ -140,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,
@@ -150,9 +173,11 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         elif self.method == "uniform":
             self.key, domain_times = self.generate_time_data(self.key, self.n)
             self.domain = jnp.concatenate([domain_times, self.omega], axis=1)
+        elif self.method in ["sobol", "halton"]:
+            self.key, self.domain = self.qmc_in_time_omega_domain(self.key, self.n)
         else:
             raise ValueError(
-                f'Bad value for method. Got {self.method}, expected "grid" or "uniform"'
+                f'Bad value for method. Got {self.method}, expected "grid" or "uniform" or "sobol" or "halton"'
             )
 
         if self.domain_batch_size is None:
@@ -172,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
@@ -182,21 +207,28 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
                     "number of points per facets (nb//2*self.dim)"
                     " cannot be lower than border batch size"
                 )
-            self.key, boundary_times = self.generate_time_data(
-                self.key, self.nb // (2 * self.dim)
-            )
-            boundary_times = boundary_times.reshape(-1, 1, 1)
-            boundary_times = jnp.repeat(
-                boundary_times, self.omega_border.shape[-1], axis=2
-            )
-            if self.dim == 1:
-                self.border = make_cartesian_product(
-                    boundary_times, self.omega_border[None, None]
+            if self.method in ["grid", "uniform"]:
+                self.key, boundary_times = self.generate_time_data(
+                    self.key, self.nb // (2 * self.dim)
+                )
+                boundary_times = boundary_times.reshape(-1, 1, 1)
+                boundary_times = jnp.repeat(
+                    boundary_times, self.omega_border.shape[-1], axis=2
                 )
+                if self.dim == 1:
+                    self.border = make_cartesian_product(
+                        boundary_times, self.omega_border[None, None]
+                    )
+                else:
+                    self.border = jnp.concatenate(
+                        [boundary_times, self.omega_border], axis=1
+                    )
             else:
-                self.border = jnp.concatenate(
-                    [boundary_times, self.omega_border], axis=1
+                self.key, self.border = self.qmc_in_time_omega_border_domain(
+                    self.key,
+                    self.nb,  # type: ignore (see inside the fun)
                 )
+
             if self.border_batch_size is None:
                 self.curr_border_idx = 0
             else:
@@ -208,15 +240,31 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             self.border_batch_size = None
             self.curr_border_idx = 0
 
-        if self.ni is not None:
-            perfect_sq = int(jnp.round(jnp.sqrt(self.ni)) ** 2)
-            if self.ni != perfect_sq:
-                warnings.warn(
-                    "Grid sampling is requested in dimension 2 with a non"
-                    f" perfect square dataset size (self.ni = {self.ni})."
-                    f" Modifying self.ni to self.ni = {perfect_sq}."
+        if ni is not None:
+            if self.method == "grid":
+                perfect_sq = int(jnp.round(jnp.sqrt(self.ni)) ** 2)
+                if self.ni != perfect_sq:
+                    warnings.warn(
+                        "Grid sampling is requested in dimension 2 with a non"
+                        f" perfect square dataset size (self.ni = {self.ni})."
+                        f" Modifying self.ni to self.ni = {perfect_sq}."
+                    )
+                self.ni = perfect_sq
+            if self.method in ["sobol", "halton"]:
+                log2_n = jnp.log2(self.ni)
+                lower_pow = 2 ** jnp.floor(log2_n)
+                higher_pow = 2 ** jnp.ceil(log2_n)
+                closest_power_of_two = (
+                    lower_pow
+                    if (self.ni - lower_pow) < (higher_pow - self.ni)
+                    else higher_pow
                 )
-            self.ni = perfect_sq
+                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_power_of_two}.",
+                    )
+                self.ni = int(closest_power_of_two)
             self.key, self.initial = self.generate_omega_data(
                 self.key, data_size=self.ni
             )
@@ -235,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`
@@ -245,21 +293,122 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         if self.method == "grid":
             partial_times = (self.tmax - self.tmin) / nt
             return key, jnp.arange(self.tmin, self.tmax, partial_times)[:, None]
-        if self.method == "uniform":
+        elif self.method in ["uniform", "sobol", "halton"]:
             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"]:
-        return jax.random.uniform(
-            key,
-            (nt, 1),
-            minval=self.tmin,
-            maxval=self.tmax,
+    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: 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
+        """
+        key, subkey = jax.random.split(key, 2)
+        qmc_generator = qmc.Sobol if self.method == "sobol" else qmc.Halton
+        sampler = qmc_generator(
+            d=self.dim + 1, scramble=True, rng=np.random.default_rng(np.uint32(subkey))
+        )
+        samples = sampler.random(n=sample_size)
+        samples[:, 1:] = qmc.scale(
+            samples[:, 1:], l_bounds=self.min_pts, u_bounds=self.max_pts
+        )  # We scale omega domain to be in (min_pts, max_pts)
+        return key, jnp.array(samples)
+
+    def qmc_in_time_omega_border_domain(
+        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.
+
+        We need to do some type ignore in this function because we have lost
+        the type narrowing from post_init,  type checkers only narrow at function level and because we cannot narrow a class attribute.
+        """
+        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:
+            key, subkey = jax.random.split(key, 2)
+            qmc_seq = qmc_generator(
+                d=1, scramble=True, rng=np.random.default_rng(np.uint32(subkey))
+            )
+            boundary_times = jnp.array(
+                qmc_seq.random(self.nb // (2 * self.dim))  # type: ignore
+            )
+            boundary_times = boundary_times.reshape(-1, 1, 1)
+            boundary_times = jnp.repeat(
+                boundary_times,
+                self.omega_border.shape[-1],  # type: ignore
+                axis=2,
+            )
+            return key, make_cartesian_product(
+                boundary_times,
+                self.omega_border[None, None],  # type: ignore
+            )
+        if self.dim == 2:
+            # 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.
+            key, *subkeys = jax.random.split(key, 5)
+            facet_n = sample_size // (2 * self.dim)
+
+            def generate_qmc_sample(key, min_val, max_val):
+                qmc_seq = qmc_generator(
+                    d=2,
+                    scramble=True,
+                    rng=np.random.default_rng(np.uint32(key)),
+                )
+                u = qmc_seq.random(n=facet_n)
+                u[:, 1:2] = qmc.scale(u[:, 1:2], l_bounds=min_val, u_bounds=max_val)
+                return jnp.array(u)
+
+            xmin_sample = generate_qmc_sample(
+                subkeys[0], self.min_pts[1], self.max_pts[1]
+            )  # [t,x,y]
+            xmin = jnp.hstack(
+                [
+                    xmin_sample[:, 0:1],
+                    self.min_pts[0] * jnp.ones((facet_n, 1)),
+                    xmin_sample[:, 1:2],
+                ]
+            )
+            xmax_sample = generate_qmc_sample(
+                subkeys[1], self.min_pts[1], self.max_pts[1]
+            )
+            xmax = jnp.hstack(
+                [
+                    xmax_sample[:, 0:1],
+                    self.max_pts[0] * jnp.ones((facet_n, 1)),
+                    xmax_sample[:, 1:2],
+                ]
+            )
+            ymin = jnp.hstack(
+                [
+                    generate_qmc_sample(subkeys[2], self.min_pts[0], self.max_pts[0]),
+                    self.min_pts[1] * jnp.ones((facet_n, 1)),
+                ]
+            )
+            ymax = jnp.hstack(
+                [
+                    generate_qmc_sample(subkeys[3], self.min_pts[0], self.max_pts[0]),
+                    self.max_pts[1] * jnp.ones((facet_n, 1)),
+                ]
+            )
+            return key, 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 _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,
@@ -296,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,
         )
@@ -309,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,
@@ -355,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,
         )
@@ -372,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,
@@ -401,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,
         )