jinns 1.5.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/_LossODE.py

@@ -19,6 +19,7 @@ from jaxtyping import Float, Array
 from jinns.loss._loss_utils import (
     dynamic_loss_apply,
     observations_loss_apply,
+    initial_condition_check,
 )
 from jinns.parameters._params import (
     _get_vmap_in_axes_params,
@@ -43,7 +44,7 @@ if TYPE_CHECKING:
 
 
 class _LossODEAbstract(AbstractLoss):
-    """
+    r"""
     Parameters
     ----------
 
@@ -60,8 +61,15 @@ class _LossODEAbstract(AbstractLoss):
         Fields can be "nn_params", "eq_params" or "both". Those that should not
         be updated will have a `jax.lax.stop_gradient` called on them. Default
         is `"nn_params"` for each composant of the loss.
-    initial_condition : tuple[float | Float[Array, " 1"], Float[Array, " dim"]], default=None
-        tuple of length 2 with initial condition $(t_0, u_0)$.
+    initial_condition : tuple[
+            Float[Array, "n_cond "],
+            Float[Array, "n_cond dim"]
+        ] |
+        tuple[int | float | Float[Array, " "],
+              int | float | Float[Array, " dim"]
+        ] | None, default=None
+        Most of the time, a tuple of length 2 with initial condition $(t_0, u_0)$.
+        From jinns v1.5.1 we accept tuples of jnp arrays with shape (n_cond, 1) for t0 and (n_cond, dim) for u0. This is useful to include observed conditions at different time points, such as *e.g* final conditions. It was designed to implement $\mathcal{L}^{aux}$ from _Systems biology informed deep learning for inferring parameters and hidden dynamics_, Alireza Yazdani et al., 2020
     obs_slice : EllipsisType | slice | None, default=None
         Slice object specifying the begininning/ending
         slice of u output(s) that is observed. This is useful for
@@ -78,7 +86,9 @@ class _LossODEAbstract(AbstractLoss):
     derivative_keys: DerivativeKeysODE | None = eqx.field(kw_only=True, default=None)
     loss_weights: LossWeightsODE | None = eqx.field(kw_only=True, default=None)
     initial_condition: (
-        tuple[float | Float[Array, " 1"], Float[Array, " dim"]] | None
+        tuple[Float[Array, " n_cond 1"], Float[Array, " n_cond dim"]]
+        | tuple[int | float | Float[Array, " "], int | float | Float[Array, " dim"]]
+        | None
     ) = eqx.field(kw_only=True, default=None)
     obs_slice: EllipsisType | slice | None = eqx.field(
         kw_only=True, default=None, static=True
@@ -112,20 +122,60 @@ class _LossODEAbstract(AbstractLoss):
                     "Initial condition should be a tuple of len 2 with (t0, u0), "
                     f"{self.initial_condition} was passed."
                 )
-            # some checks/reshaping for t0
+            # some checks/reshaping for t0 and u0
             t0, u0 = self.initial_condition
             if isinstance(t0, Array):
-                if not t0.shape:  # e.g. user input: jnp.array(0.)
-                    t0 = jnp.array([t0])
-                elif t0.shape != (1,):
+                # at the end we want to end up with t0 of shape (:, 1) to account for
+                # possibly several data points
+                if t0.ndim <= 1:
+                    # in this case we assume t0 belongs one (initial)
+                    # condition
+                    t0 = initial_condition_check(t0, dim_size=1)[
+                        None, :
+                    ]  # make a (1, 1) here
+                if t0.ndim > 2:
+                    raise ValueError(
+                        "It t0 is an Array, it represents n_cond"
+                        " imposed conditions and must be of shape (n_cond, 1)"
+                    )
+            else:
+                # in this case t0 clearly represents one (initial) condition
+                t0 = initial_condition_check(t0, dim_size=1)[
+                    None, :
+                ]  # make a (1, 1) here
+            if isinstance(u0, Array):
+                # at the end we want to end up with u0 of shape (:, dim) to account for
+                # possibly several data points
+                if not u0.shape:
+                    # in this case we assume u0 belongs to one (initial)
+                    # condition
+                    u0 = initial_condition_check(u0, dim_size=1)[
+                        None, :
+                    ]  # make a (1, 1) here
+                elif u0.ndim == 1:
+                    # in this case we assume u0 belongs to one (initial)
+                    # condition
+                    u0 = initial_condition_check(u0, dim_size=u0.shape[0])[
+                        None, :
+                    ]  # make a (1, dim) here
+                if u0.ndim > 2:
                     raise ValueError(
-                        f"Wrong t0 input (self.initial_condition[0]) It should be"
-                        f"a float or an array of shape (1,). Got shape: {t0.shape}"
+                        "It u0 is an Array, it represents n_cond "
+                        "imposed conditions and must be of shape (n_cond, dim)"
                     )
-            if isinstance(t0, float):  # e.g. user input: 0.
-                t0 = jnp.array([t0])
-            if isinstance(t0, int):  # e.g. user input: 0
-                t0 = jnp.array([float(t0)])
+            else:
+                # at the end we want to end up with u0 of shape (:, dim) to account for
+                # possibly several data points
+                u0 = initial_condition_check(u0, dim_size=None)[
+                    None, :
+                ]  # make a (1, 1) here
+
+            if t0.shape[0] != u0.shape[0] or t0.ndim != u0.ndim:
+                raise ValueError(
+                    "t0 and u0 must represent a same number of initial"
+                    " conditial conditions"
+                )
+
             self.initial_condition = (t0, u0)
 
         if self.obs_slice is None:
@@ -259,28 +309,42 @@ class LossODE(_LossODEAbstract):
 
         # initial condition
         if self.initial_condition is not None:
-            vmap_in_axes = (None,) + vmap_in_axes_params
-            if not jax.tree_util.tree_leaves(vmap_in_axes):
-                # test if only None in vmap_in_axes to avoid the value error:
-                # `vmap must have at least one non-None value in in_axes`
-                v_u = self.u
-            else:
-                v_u = vmap(self.u, (None,) + vmap_in_axes_params)
             t0, u0 = self.initial_condition
             u0 = jnp.array(u0)
-            initial_condition_fun = lambda p: jnp.mean(
-                jnp.sum(
-                    (
-                        v_u(
-                            t0,
-                            _set_derivatives(p, self.derivative_keys.initial_condition),  # type: ignore
-                        )
-                        - u0
+
+            # first construct the plain init loss no vmaping
+            initial_condition_fun__ = lambda t, u, p: jnp.sum(
+                (
+                    self.u(
+                        t,
+                        _set_derivatives(
+                            p,
+                            self.derivative_keys.initial_condition,  # type: ignore
+                        ),
                     )
-                    ** 2,
-                    axis=-1,
+                    - u
                 )
+                ** 2,
+                axis=0,
             )
+            # now vmap over the number of conditions (first dim of t0 and u0)
+            # and take the mean
+            initial_condition_fun_ = lambda p: jnp.mean(
+                vmap(initial_condition_fun__, (0, 0, None))(t0, u0, p)
+            )
+            # now vmap over the the possible batch of parameters and take the
+            # average. Note that we then finally have a cartesian product
+            # between the batch of parameters (if any) and the number of
+            # conditions (if any)
+            if not jax.tree_util.tree_leaves(vmap_in_axes_params):
+                # if there is no parameter batch to vmap over we cannot call
+                # vmap because calling vmap must be done with at least one non
+                # None in_axes or out_axes
+                initial_condition_fun = initial_condition_fun_
+            else:
+                initial_condition_fun = lambda p: jnp.mean(
+                    vmap(initial_condition_fun_, vmap_in_axes_params)(p)
+                )
         else:
             initial_condition_fun = None
 

jinns/loss/_LossPDE.py

@@ -21,6 +21,7 @@ from jinns.loss._loss_utils import (
     normalization_loss_apply,
     observations_loss_apply,
     initial_condition_apply,
+    initial_condition_check,
 )
 from jinns.parameters._params import (
     _get_vmap_in_axes_params,
@@ -700,22 +701,12 @@ class LossPDENonStatio(LossPDEStatio):
                 "case (e.g by. hardcoding it into the PINN output)."
             )
         # some checks for t0
-        if isinstance(self.t0, Array):
-            if not self.t0.shape:  # e.g. user input: jnp.array(0.)
-                self.t0 = jnp.array([self.t0])
-            elif self.t0.shape != (1,):
-                raise ValueError(
-                    f"Wrong self.t0 input. It should be"
-                    f"a float or an array of shape (1,). Got shape: {self.t0.shape}"
-                )
-        elif isinstance(self.t0, float):  # e.g. user input: 0.
-            self.t0 = jnp.array([self.t0])
-        elif isinstance(self.t0, int):  # e.g. user input: 0
-            self.t0 = jnp.array([float(self.t0)])
-        elif self.t0 is None:
-            self.t0 = jnp.array([0])
+        t0 = self.t0
+        if t0 is None:
+            t0 = jnp.array([0])
         else:
-            raise ValueError("Wrong value for t0")
+            t0 = initial_condition_check(t0, dim_size=1)
+        self.t0 = t0
 
         # witht the variables below we avoid memory overflow since a cartesian
         # product is taken

jinns/loss/_loss_utils.py

@@ -308,3 +308,26 @@ def initial_condition_apply(
     else:
         raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_initial_condition
+
+
+def initial_condition_check(x, dim_size=None):
+    """
+    Make a (dim_size,) jnp array from an int, a float or a 0D jnp array
+
+    """
+    if isinstance(x, Array):
+        if not x.shape:  # e.g. user input: jnp.array(0.)
+            x = jnp.array([x])
+        if dim_size is not None:  # we check for the required dims_ize
+            if x.shape != (dim_size,):
+                raise ValueError(
+                    f"Wrong dim_size. It should be({dim_size},). Got shape: {x.shape}"
+                )
+
+    elif isinstance(x, float):  # e.g. user input: 0.
+        x = jnp.array([x])
+    elif isinstance(x, int):  # e.g. user input: 0
+        x = jnp.array([float(x)])
+    else:
+        raise ValueError(f"Wrong value, expected Array, float or int, got {type(x)}")
+    return x

jinns/parameters/_params.py

@@ -10,6 +10,14 @@ from jaxtyping import Array, PyTree, Float
 T = TypeVar("T")  # the generic type for what is in the Params PyTree because we
 # have possibly Params of Arrays, boolean, ...
 
+### NOTE
+### We are taking derivatives with respect to Params eqx.Modules.
+### This has been shown to behave weirdly if some fields of eqx.Modules have
+### been set as `field(init=False)`, we then should never create such fields in
+### jinns' Params modules.
+### We currently have silenced the warning related to this (see jinns.__init__
+### see https://github.com/patrick-kidger/equinox/pull/1043/commits/f88e62ab809140334c2f987ed13eff0d80b8be13
+
 
 class Params(eqx.Module, Generic[T]):
     """

jinns/solver/_solve.py

@@ -179,6 +179,7 @@ def solve(
     best_val_params
         The best parameters according to the validation criterion
     """
+    initialization_time = time.time()
     if n_iter < 1:
         raise ValueError("Cannot run jinns.solve for n_iter<1")
 
@@ -225,11 +226,6 @@ def solve(
     # get_batch with device_put, the latter is not jittable
     get_batch = _get_get_batch(obs_batch_sharding)
 
-    # initialize the dict for stored parameter values
-    # we need to get a loss_term to init stuff
-    batch_ini, data, param_data, obs_data = get_batch(data, param_data, obs_data)
-    _, loss_terms = loss(init_params, batch_ini)
-
     # initialize parameter tracking
     if tracked_params is None:
         tracked_params = jax.tree.map(lambda p: None, init_params)
@@ -247,6 +243,13 @@ def solve(
         # being a complex data structure
     )
 
+    # initialize the dict for stored parameter values
+    # we need to get a loss_term to init stuff
+    # NOTE: we use jax.eval_shape to avoid FLOPS since we only need the tree
+    # structure
+    batch_ini, data, param_data, obs_data = get_batch(data, param_data, obs_data)
+    _, loss_terms = jax.eval_shape(loss, init_params, batch_ini)
+
     # initialize the PyTree for stored loss values
     stored_loss_terms = jax.tree_util.tree_map(
         lambda _: jnp.zeros((n_iter)), loss_terms
@@ -475,6 +478,9 @@ def solve(
             key,
         )
 
+    if verbose:
+        print("Initialization time:", time.time() - initialization_time)
+
     # Main optimization loop. We use the LAX while loop (fully jitted) version
     # if no mixing devices. Otherwise we use the standard while loop. Here devices only
     # concern obs_batch, but it could lead to more complex scheme in the future

{jinns-1.5.0.dist-info → jinns-1.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jinns
-Version: 1.5.0
+Version: 1.5.1
 Summary: Physics Informed Neural Network with JAX
 Author-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
 Maintainer-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-1.5.0.dist-info → jinns-1.5.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
-jinns/__init__.py,sha256=hyh3QKO2cQGK5cmvFYP0MrXb-tK_DM2T9CwLwO-sEX8,500
+jinns/__init__.py,sha256=Tjp5z0Mnd1nscvJaXnxZb4lEYbI0cpN6OPgCZ-Swo74,453
 jinns/data/_AbstractDataGenerator.py,sha256=O61TBOyeOFKwf1xqKzFD4KwCWRDnm2XgyJ-kKY9fmB4,557
 jinns/data/_Batchs.py,sha256=-DlD6Qag3zs5QbKtKAOvOzV7JOpNOqAm_P8cwo1dIZg,1574
-jinns/data/_CubicMeshPDENonStatio.py,sha256=c_8czJpxSoEvgZ8LDpL2sqtF9dcW4ELNO4juEFMOxog,16400
-jinns/data/_CubicMeshPDEStatio.py,sha256=stZ0Kbb7_VwFmWUSPs0P6a6qRj2Tu67p7sxEfb1Ajks,17865
+jinns/data/_CubicMeshPDENonStatio.py,sha256=4f21SgeNQsGJTz8-uehduU0X9TibRcs28Iq49Kv4nQQ,22250
+jinns/data/_CubicMeshPDEStatio.py,sha256=DVrP4qHVAJMu915EYH2PKyzwoG0nIMixAEKR2fz6C58,22525
 jinns/data/_DataGeneratorODE.py,sha256=5RzUbQFEsooAZsocDw4wRgA_w5lJmDMuY4M6u79K-1c,7260
 jinns/data/_DataGeneratorObservations.py,sha256=jknepLsJatSJHFq5lLMD-fFHkPGj5q286LEjE-vH24k,7738
 jinns/data/_DataGeneratorParameter.py,sha256=IedX3jcOj7ZDW_18IAcRR75KVzQzo85z9SICIKDBJl4,8539
@@ -12,13 +12,13 @@ jinns/experimental/__init__.py,sha256=DT9e57zbjfzPeRnXemGUqnGd--MhV77FspChT0z4Yr
 jinns/experimental/_diffrax_solver.py,sha256=upMr3kTTNrxEiSUO_oLvCXcjS9lPxSjvbB81h3qlhaU,6813
 jinns/loss/_DynamicLoss.py,sha256=4mb7OCP-cGZ_mG2MQ-AniddDcuBT78p4bQI7rZpwte4,22722
 jinns/loss/_DynamicLossAbstract.py,sha256=QhHRgvtcT-ifHlOxTyXbjDtHk9UfPN2Si8s3v9nEQm4,12672
-jinns/loss/_LossODE.py,sha256=DeejnU2ytgrOxUnwuVkQDWWRKJAgNQyjacTx-jT0xPA,13796
-jinns/loss/_LossPDE.py,sha256=ycjWJ99SuXe9DV5nROSWyq--xcp2JJ2PGWxsdWyZZog,36942
+jinns/loss/_LossODE.py,sha256=iVYDojaI6Co7S5CrU67_niopD4Bk7UBTuLzDiTHoWMc,16996
+jinns/loss/_LossPDE.py,sha256=VT56oQ_33fLq46lIch0slNsxu4d97eQBOgRAPeFESts,36401
 jinns/loss/__init__.py,sha256=z5xYgBipNFf66__5BqQc6R_8r4F6A3TXL60YjsM8Osk,1287
 jinns/loss/_abstract_loss.py,sha256=DMxn0SQe9PW-pq3p5Oqvb0YK3_ulLDOnoIXzK219GV4,4576
 jinns/loss/_boundary_conditions.py,sha256=9HGw1cGLfmEilP4V4B2T0zl0YP1kNtrtXVLQNiBmWgc,12464
 jinns/loss/_loss_components.py,sha256=MMzaGlaRqESPjRzT0j0WU9HAqWQSbIXpGAqM1xQCZHw,1106
-jinns/loss/_loss_utils.py,sha256=R6PffBAtg6z9M8x1DFXmmqZpC095b9gZ_DB1phQxSuY,11168
+jinns/loss/_loss_utils.py,sha256=eJ4JcBm396LHx7Tti88ZQrLcKqVL1oSfFGT23VNkytQ,11949
 jinns/loss/_loss_weight_updates.py,sha256=9Bwouh7shLyc_wrdzN6CYL0ZuQH81uEs-L6wCeiYFx8,6817
 jinns/loss/_loss_weights.py,sha256=kII5WddORgeommFTudT3CSvhICpo6nSe47LclUgu_78,2429
 jinns/loss/_operators.py,sha256=Ds5yRH7hu-jaGRp7PYbt821BgYuEvgWHufWhYgdMjw0,22909
@@ -34,12 +34,12 @@ jinns/nn/_spinn_mlp.py,sha256=uCL454sF0Tfj7KT-fdXPnvKJYRQOuq60N0r2b2VAB8Q,7606
 jinns/nn/_utils.py,sha256=9UXz73iHKHVQYPBPIEitrHYJzJ14dspRwPfLA8avx0c,1120
 jinns/parameters/__init__.py,sha256=O0n7y6R1LRmFzzugCxMFCMS2pgsuWSh-XHjfFViN_eg,265
 jinns/parameters/_derivative_keys.py,sha256=YlLDX49PfYhr2Tj--t3praiD8JOUTZU6PTmjbNZsbMc,19173
-jinns/parameters/_params.py,sha256=qn4IGMJhD9lDBqOWmGEMy4gXt5a6KHfirkYZwHO7Vwk,2633
+jinns/parameters/_params.py,sha256=nv0WScbgUdmuC0bSF15VbnKypJ58pl6wynZAcYfuF6M,3081
 jinns/plot/__init__.py,sha256=KPHX0Um4FbciZO1yD8kjZbkaT8tT964Y6SE2xCQ4eDU,135
 jinns/plot/_plot.py,sha256=-A5auNeElaz2_8UzVQJQE4143ZFg0zgMjStU7kwttEY,11565
 jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jinns/solver/_rar.py,sha256=vSVTnCGCusI1vTZCvIkP2_G8we44G_42yZHx2sOK9DE,10291
-jinns/solver/_solve.py,sha256=oVHnuc7Z0V2-ZYgZtCx7xdFd7TpB9w-6AwafX-kgBE4,28379
+jinns/solver/_solve.py,sha256=IsrmG2m48KkkYgvXYomlSbZ3hd1FySCj3rwlkovs-lI,28616
 jinns/solver/_utils.py,sha256=sM2UbVzYyjw24l4QSIR3IlynJTPGD_S08r8v0lXMxA8,5876
 jinns/utils/__init__.py,sha256=OEYWLCw8pKE7xoQREbd6SHvCjuw2QZHuVA6YwDcsBE8,53
 jinns/utils/_containers.py,sha256=YShcrPKfj5_I9mn3NMAS4Ea9MhhyL7fjv0e3MRbITHg,1837
@@ -47,9 +47,9 @@ jinns/utils/_types.py,sha256=jl_91HtcrtE6UHbdTrRI8iUmr2kBUL0oP0UNIKhAXYw,1170
 jinns/utils/_utils.py,sha256=M7NXX9ok-BkH5o_xo74PB1_Cc8XiDipSl51rq82dTH4,2821
 jinns/validation/__init__.py,sha256=FTyUO-v1b8Tv-FDSQsntrH7zl9E0ENexqKMT_dFRkYo,124
 jinns/validation/_validation.py,sha256=8p6sMKiBAvA6JNm65hjkMj0997LJ0BkyCREEh0AnPVE,4803
-jinns-1.5.0.dist-info/licenses/AUTHORS,sha256=7NwCj9nU-HNG1asvy4qhQ2w7oZHrn-Lk5_wK_Ve7a3M,80
-jinns-1.5.0.dist-info/licenses/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-1.5.0.dist-info/METADATA,sha256=jEp__DP39B1HiTYVhtVcWKPmzS22kSUD6jNVSmHFh8g,5314
-jinns-1.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-jinns-1.5.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-1.5.0.dist-info/RECORD,,
+jinns-1.5.1.dist-info/licenses/AUTHORS,sha256=7NwCj9nU-HNG1asvy4qhQ2w7oZHrn-Lk5_wK_Ve7a3M,80
+jinns-1.5.1.dist-info/licenses/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-1.5.1.dist-info/METADATA,sha256=K7Aii5ivFcczIwLlQtCPqzMJEfW86D1yW1q7qMtvWPE,5314
+jinns-1.5.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+jinns-1.5.1.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-1.5.1.dist-info/RECORD,,