jinns 1.4.0__py3-none-any.whl → 1.5.1__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/_CubicMeshPDENonStatio.py

@@ -7,8 +7,10 @@ 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 scipy.stats import qmc
 from jaxtyping import Key, Array, Float
 from jinns.data._Batchs import PDENonStatioBatch
 from jinns.data._utils import (
@@ -65,11 +67,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
@@ -150,9 +154,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:
@@ -182,21 +188,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:
@@ -209,14 +222,30 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             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 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_two_power = (
+                    lower_pow
+                    if (self.ni - lower_pow) < (higher_pow - self.ni)
+                    else higher_pow
                 )
-            self.ni = perfect_sq
+                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.ni} to {closest_two_power}.",
+                    )
+                self.ni = int(closest_two_power)
             self.key, self.initial = self.generate_omega_data(
                 self.key, data_size=self.ni
             )
@@ -245,16 +274,115 @@ 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,
+        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"]]:
+        """
+        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: Key, sample_size: int | None = None
+    ) -> tuple[Key, 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(

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 scipy.stats import qmc
 from jaxtyping import Key, 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
@@ -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
@@ -94,7 +99,7 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
     # 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(
+    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)
@@ -132,6 +137,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:
@@ -176,24 +197,48 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
     def sample_in_omega_domain(
         self, keys: Key, 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, sample_size)
+
+    def _qmc_in_omega_domain(
+        self, subkey: Key, 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
@@ -260,6 +305,62 @@ class CubicMeshPDEStatio(AbstractDataGenerator):
             + f"implemented yet. You are asking for generation in dimension d={self.dim}."
         )
 
+    def qmc_in_omega_border_domain(
+        self, keys: Key, 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:
+            # 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
     ) -> tuple[
@@ -290,8 +391,8 @@ 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:
+        elif self.method in ["uniform", "sobol", "halton"]:
+            if self.dim == 1 or self.method in ["sobol", "halton"]:
                 key, subkeys = jax.random.split(key, 2)
             else:
                 key, *subkeys = jax.random.split(key, self.dim + 1)
@@ -317,10 +418,17 @@ 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(

jinns/loss/_DynamicLossAbstract.py

@@ -6,11 +6,13 @@ from __future__ import (
     annotations,
 )  # https://docs.python.org/3/library/typing.html#constant
 
+import warnings
 import abc
 from functools import partial
 from typing import Callable, TYPE_CHECKING, ClassVar, Generic, TypeVar
 import equinox as eqx
 from jaxtyping import Float, Array
+import jax.numpy as jnp
 
 
 # See : https://docs.kidger.site/equinox/api/module/advanced_fields/#equinox.AbstractClassVar--known-issues
@@ -70,14 +72,31 @@ class DynamicLoss(eqx.Module, Generic[InputDim]):
         A value can be missing, in this case there is no heterogeneity (=None).
         Default None, meaning there is no heterogeneity in the equation
         parameters.
+    vectorial_dyn_loss_ponderation : Float[Array, " dim"], default=None
+        Add a different ponderation weight to each of the dimension to the
+        dynamic loss. This array must have the same dimension as the output of
+        the dynamic loss equation or an error is raised. Default is None which
+        means that a ponderation of 1 is applied on each dimension.
+        `vectorial_dyn_loss_ponderation`
+        is different from loss weights, which are attributes of Loss
+        classes and which implement scalar (and possibly dynamic)
+        ponderations for each term of the total loss.
+        `vectorial_dyn_loss_ponderation` can be used with loss weights.
     """
 
     _eq_type = AbstractClassVar[str]  # class variable denoting the type of
     # differential equation
     Tmax: Float = eqx.field(kw_only=True, default=1)
-    eq_params_heterogeneity: dict[str, Callable | None] = eqx.field(
+    eq_params_heterogeneity: dict[str, Callable | None] | None = eqx.field(
         kw_only=True, default=None, static=True
     )
+    vectorial_dyn_loss_ponderation: Float[Array, " dim"] | None = eqx.field(
+        kw_only=True, default=None
+    )
+
+    def __post_init__(self):
+        if self.vectorial_dyn_loss_ponderation is None:
+            self.vectorial_dyn_loss_ponderation = jnp.array(1.0)
 
     def _eval_heterogeneous_parameters(
         self,
@@ -110,12 +129,21 @@ class DynamicLoss(eqx.Module, Generic[InputDim]):
         u: AbstractPINN,
         params: Params[Array],
     ) -> float:
-        evaluation = self.equation(inputs, u, params)
+        evaluation = self.vectorial_dyn_loss_ponderation * self.equation(
+            inputs, u, params
+        )
         if len(evaluation.shape) == 0:
             raise ValueError(
                 "The output of dynamic loss must be vectorial, "
                 "i.e. of shape (d,) with d >= 1"
             )
+        if len(evaluation.shape) > 1:
+            warnings.warn(
+                "Return value from DynamicLoss' equation has more "
+                "than one dimension. This is in general a mistake (probably from "
+                "an unfortunate broadcast in jnp.array computations) resulting in "
+                "bad reduction operations in losses."
+            )
         return evaluation
 
     @abc.abstractmethod