jinns 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

jinns/data/_DataGeneratorObservations.py

@@ -0,0 +1,189 @@
+"""
+Define the DataGenerators modules
+"""
+
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Key, Int, Array, Float
+from jinns.data._Batchs import ObsBatchDict
+from jinns.data._utils import _reset_or_increment
+from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+
+
+class DataGeneratorObservations(AbstractDataGenerator):
+    r"""
+    Despite the class name, it is rather a dataloader for user-provided
+    observations which will are used in the observations loss.
+
+    Parameters
+    ----------
+    key : Key
+        Jax random key to shuffle batches
+    obs_batch_size : int | None
+        The size of the batch of randomly selected points among
+        the `n` points. If None, no minibatch are used.
+    observed_pinn_in : Float[Array, " n_obs nb_pinn_in"]
+        Observed values corresponding to the input of the PINN
+        (eg. the time at which we recorded the observations). The first
+        dimension must corresponds to the number of observed_values.
+        The second dimension depends on the input dimension of the PINN,
+        that is `1` for ODE, `n_dim_x` for stationnary PDE and `n_dim_x + 1`
+        for non-stationnary PDE.
+    observed_values : Float[Array, " n_obs, nb_pinn_out"]
+        Observed values that the PINN should learn to fit. The first
+        dimension must be aligned with observed_pinn_in.
+    observed_eq_params : dict[str, Float[Array, " n_obs 1"]], default={}
+        A dict with keys corresponding to
+        the parameter name. The keys must match the keys in
+        `params["eq_params"]`. The values are jnp.array with 2 dimensions
+        with values corresponding to the parameter value for which we also
+        have observed_pinn_in and observed_values. Hence the first
+        dimension must be aligned with observed_pinn_in and observed_values.
+        Optional argument.
+    sharding_device : jax.sharding.Sharding, default=None
+        Default None. An optional sharding object to constraint the storage
+        of observed inputs, values and parameters. Typically, a
+        SingleDeviceSharding(cpu_device) to avoid loading on GPU huge
+        datasets of observations. Note that computations for **batches**
+        can still be performed on other devices (*e.g.* GPU, TPU or
+        any pre-defined Sharding) thanks to the `obs_batch_sharding`
+        arguments of `jinns.solve()`. Read `jinns.solve()` doc for more info.
+    """
+
+    key: Key
+    obs_batch_size: int | None = eqx.field(static=True)
+    observed_pinn_in: Float[Array, " n_obs nb_pinn_in"]
+    observed_values: Float[Array, " n_obs nb_pinn_out"]
+    observed_eq_params: dict[str, Float[Array, " n_obs 1"]] = eqx.field(
+        static=True, default_factory=lambda: {}
+    )
+    sharding_device: jax.sharding.Sharding = eqx.field(static=True, default=None)
+
+    n: int = eqx.field(init=False, static=True)
+    curr_idx: int = eqx.field(init=False)
+    indices: Array = eqx.field(init=False)
+
+    def __post_init__(self):
+        if self.observed_pinn_in.shape[0] != self.observed_values.shape[0]:
+            raise ValueError(
+                "self.observed_pinn_in and self.observed_values must have same first axis"
+            )
+        for _, v in self.observed_eq_params.items():
+            if v.shape[0] != self.observed_pinn_in.shape[0]:
+                raise ValueError(
+                    "self.observed_pinn_in and the values of"
+                    " self.observed_eq_params must have the same first axis"
+                )
+        if len(self.observed_pinn_in.shape) == 1:
+            self.observed_pinn_in = self.observed_pinn_in[:, None]
+        if self.observed_pinn_in.ndim > 2:
+            raise ValueError("self.observed_pinn_in must have 2 dimensions")
+        if len(self.observed_values.shape) == 1:
+            self.observed_values = self.observed_values[:, None]
+        if self.observed_values.ndim > 2:
+            raise ValueError("self.observed_values must have 2 dimensions")
+        for k, v in self.observed_eq_params.items():
+            if len(v.shape) == 1:
+                self.observed_eq_params[k] = v[:, None]
+            if len(v.shape) > 2:
+                raise ValueError(
+                    "Each value of observed_eq_params must have 2 dimensions"
+                )
+
+        self.n = self.observed_pinn_in.shape[0]
+
+        if self.sharding_device is not None:
+            self.observed_pinn_in = jax.lax.with_sharding_constraint(
+                self.observed_pinn_in, self.sharding_device
+            )
+            self.observed_values = jax.lax.with_sharding_constraint(
+                self.observed_values, self.sharding_device
+            )
+            self.observed_eq_params = jax.lax.with_sharding_constraint(
+                self.observed_eq_params, self.sharding_device
+            )
+
+        if self.obs_batch_size is not None:
+            self.curr_idx = self.n + self.obs_batch_size
+            # to be sure there is a shuffling at first get_batch()
+        else:
+            self.curr_idx = 0
+        # For speed and to avoid duplicating data what is really
+        # shuffled is a vector of indices
+        if self.sharding_device is not None:
+            self.indices = jax.lax.with_sharding_constraint(
+                jnp.arange(self.n), self.sharding_device
+            )
+        else:
+            self.indices = jnp.arange(self.n)
+
+        # recall post_init is the only place with _init_ where we can set
+        # self attribute in a in-place way
+        self.key, _ = jax.random.split(self.key, 2)  # to make it equivalent to
+        # the call to _reset_batch_idx_and_permute in legacy DG
+
+    def _get_operands(self) -> tuple[Key, Int[Array, " n"], int, int | None, None]:
+        return (
+            self.key,
+            self.indices,
+            self.curr_idx,
+            self.obs_batch_size,
+            None,
+        )
+
+    def obs_batch(
+        self,
+    ) -> tuple[DataGeneratorObservations, ObsBatchDict]:
+        """
+        Return an update DataGeneratorObservations instance and an ObsBatchDict
+        """
+        if self.obs_batch_size is None or self.obs_batch_size == self.n:
+            # Avoid unnecessary reshuffling
+            return self, {
+                "pinn_in": self.observed_pinn_in,
+                "val": self.observed_values,
+                "eq_params": self.observed_eq_params,
+            }
+
+        new_attributes = _reset_or_increment(
+            self.curr_idx + self.obs_batch_size,
+            self.n,
+            self._get_operands(),  # type: ignore
+            # ignore since the case self.obs_batch_size is None has been
+            # handled above
+        )
+        new = eqx.tree_at(
+            lambda m: (m.key, m.indices, m.curr_idx), self, new_attributes
+        )
+
+        minib_indices = jax.lax.dynamic_slice(
+            new.indices,
+            start_indices=(new.curr_idx,),
+            slice_sizes=(new.obs_batch_size,),
+        )
+
+        obs_batch: ObsBatchDict = {
+            "pinn_in": jnp.take(
+                new.observed_pinn_in, minib_indices, unique_indices=True, axis=0
+            ),
+            "val": jnp.take(
+                new.observed_values, minib_indices, unique_indices=True, axis=0
+            ),
+            "eq_params": jax.tree_util.tree_map(
+                lambda a: jnp.take(a, minib_indices, unique_indices=True, axis=0),
+                new.observed_eq_params,
+            ),
+        }
+        return new, obs_batch
+
+    def get_batch(
+        self,
+    ) -> tuple[DataGeneratorObservations, ObsBatchDict]:
+        """
+        Generic method to return a batch
+        """
+        return self.obs_batch()

jinns/data/_DataGeneratorParameter.py

@@ -0,0 +1,206 @@
+"""
+Define the DataGenerators modules
+"""
+
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Key, Array, Float
+from jinns.data._utils import _reset_or_increment
+from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+
+
+class DataGeneratorParameter(AbstractDataGenerator):
+    r"""
+    A data generator for additional unidimensional equation parameter(s).
+    Mostly useful for metamodeling where batch of `params.eq_params` are fed
+    to the network.
+
+    Parameters
+    ----------
+    keys : Key | dict[str, Key]
+        Jax random key to sample new time points and to shuffle batches
+        or a dict of Jax random keys with key entries from param_ranges
+    n : int
+        The number of total points that will be divided in
+        batches. Batches are made so that each data point is seen only
+        once during 1 epoch.
+    param_batch_size : int | None, default=None
+        The size of the batch of randomly selected points among
+        the `n` points.  **Important**: no check is performed but
+        `param_batch_size` must be the same as other collocation points
+        batch_size (time, space or timexspace depending on the context). This is because we vmap the network on all its axes at once to compute the MSE. Also, `param_batch_size` will be the same for all parameters. If None, no mini-batches are used.
+    param_ranges : dict[str, tuple[Float, Float] | None, default={}
+        A dict. A dict of tuples (min, max), which
+        reprensents the range of real numbers where to sample batches (of
+        length `param_batch_size` among `n` points).
+        The key corresponds to the parameter name. The keys must match the
+        keys in `params["eq_params"]`.
+        By providing several entries in this dictionary we can sample
+        an arbitrary number of parameters.
+        **Note** that we currently only support unidimensional parameters.
+        This argument can be None if we use `user_data`.
+    method : str, default="uniform"
+        Either `grid` or `uniform`, default is `uniform`. `grid` means
+        regularly spaced points over the domain. `uniform` means uniformly
+        sampled points over the domain
+    user_data : dict[str, Float[Array, " n"]] | None, default={}
+        A dictionary containing user-provided data for parameters.
+        The keys corresponds to the parameter name,
+        and must match the keys in `params["eq_params"]`. Only
+        unidimensional `jnp.array` are supported. Therefore, the array at
+        `user_data[k]` must have shape `(n, 1)` or `(n,)`.
+        Note that if the same key appears in `param_ranges` and `user_data`
+        priority goes for the content in `user_data`.
+        Defaults to None.
+    """
+
+    keys: Key | dict[str, Key]
+    n: int = eqx.field(static=True)
+    param_batch_size: int | None = eqx.field(static=True, default=None)
+    param_ranges: dict[str, tuple[Float, Float]] = eqx.field(
+        static=True, default_factory=lambda: {}
+    )
+    method: str = eqx.field(static=True, default="uniform")
+    user_data: dict[str, Float[Array, " n"]] | None = eqx.field(
+        default_factory=lambda: {}
+    )
+
+    curr_param_idx: dict[str, int] = eqx.field(init=False)
+    param_n_samples: dict[str, Array] = eqx.field(init=False)
+
+    def __post_init__(self):
+        if self.user_data is None:
+            self.user_data = {}
+        if self.param_ranges is None:
+            self.param_ranges = {}
+        if self.param_batch_size is not None and self.n < self.param_batch_size:
+            raise ValueError(
+                f"Number of data points ({self.n}) is smaller than the"
+                f"number of batch points ({self.param_batch_size})."
+            )
+        if not isinstance(self.keys, dict):
+            all_keys = set().union(self.param_ranges, self.user_data)
+            self.keys = dict(zip(all_keys, jax.random.split(self.keys, len(all_keys))))
+
+        if self.param_batch_size is None:
+            self.curr_param_idx = None  # type: ignore
+        else:
+            self.curr_param_idx = {}
+            for k in self.keys.keys():
+                self.curr_param_idx[k] = self.n + self.param_batch_size
+                # to be sure there is a shuffling at first get_batch()
+
+        # The call to self.generate_data() creates
+        # the dict self.param_n_samples and then we will only use this one
+        # because it merges the scattered data between `user_data` and
+        # `param_ranges`
+        self.keys, self.param_n_samples = self.generate_data(self.keys)
+
+    def generate_data(
+        self, keys: dict[str, Key]
+    ) -> tuple[dict[str, Key], dict[str, Float[Array, " n"]]]:
+        """
+        Generate parameter samples, either through generation
+        or using user-provided data.
+        """
+        param_n_samples = {}
+
+        all_keys = set().union(
+            self.param_ranges,
+            self.user_data,  # type: ignore this has been handled in post_init
+        )
+        for k in all_keys:
+            if self.user_data and k in self.user_data.keys():
+                if self.user_data[k].shape == (self.n, 1):
+                    param_n_samples[k] = self.user_data[k]
+                if self.user_data[k].shape == (self.n,):
+                    param_n_samples[k] = self.user_data[k][:, None]
+                else:
+                    raise ValueError(
+                        "Wrong shape for user provided parameters"
+                        f" in user_data dictionary at key='{k}'"
+                    )
+            else:
+                if self.method == "grid":
+                    xmin, xmax = self.param_ranges[k][0], self.param_ranges[k][1]
+                    partial = (xmax - xmin) / self.n
+                    # shape (n, 1)
+                    param_n_samples[k] = jnp.arange(xmin, xmax, partial)[:, None]
+                elif self.method == "uniform":
+                    xmin, xmax = self.param_ranges[k][0], self.param_ranges[k][1]
+                    keys[k], subkey = jax.random.split(keys[k], 2)
+                    param_n_samples[k] = jax.random.uniform(
+                        subkey, shape=(self.n, 1), minval=xmin, maxval=xmax
+                    )
+                else:
+                    raise ValueError("Method " + self.method + " is not implemented.")
+
+        return keys, param_n_samples
+
+    def _get_param_operands(
+        self, k: str
+    ) -> tuple[Key, Float[Array, " n"], int, int | None, None]:
+        return (
+            self.keys[k],
+            self.param_n_samples[k],
+            self.curr_param_idx[k],
+            self.param_batch_size,
+            None,
+        )
+
+    def param_batch(self):
+        """
+        Return a dictionary with batches of parameters
+        If all the batches have been seen, we reshuffle them,
+        otherwise we just return the next unseen batch.
+        """
+
+        if self.param_batch_size is None or self.param_batch_size == self.n:
+            return self, self.param_n_samples
+
+        def _reset_or_increment_wrapper(param_k, idx_k, key_k):
+            return _reset_or_increment(
+                idx_k + self.param_batch_size,
+                self.n,
+                (key_k, param_k, idx_k, self.param_batch_size, None),  # type: ignore
+                # ignore since the case self.param_batch_size is None has been
+                # handled above
+            )
+
+        res = jax.tree_util.tree_map(
+            _reset_or_increment_wrapper,
+            self.param_n_samples,
+            self.curr_param_idx,
+            self.keys,
+        )
+        # we must transpose the pytrees because keys are merged in res
+        # https://jax.readthedocs.io/en/latest/jax-101/05.1-pytrees.html#transposing-trees
+        new_attributes = jax.tree_util.tree_transpose(
+            jax.tree_util.tree_structure(self.keys),
+            jax.tree_util.tree_structure([0, 0, 0]),
+            res,
+        )
+
+        new = eqx.tree_at(
+            lambda m: (m.keys, m.param_n_samples, m.curr_param_idx),
+            self,
+            new_attributes,
+        )
+
+        return new, jax.tree_util.tree_map(
+            lambda p, q: jax.lax.dynamic_slice(
+                p, start_indices=(q, 0), slice_sizes=(new.param_batch_size, 1)
+            ),
+            new.param_n_samples,
+            new.curr_param_idx,
+        )
+
+    def get_batch(self):
+        """
+        Generic method to return a batch
+        """
+        return self.param_batch()

jinns/data/__init__.py

@@ -1,11 +1,21 @@
-from ._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-    DataGeneratorObservations,
-    DataGeneratorParameter,
-    DataGeneratorObservationsMultiPINNs,
-)
+from ._DataGeneratorODE import DataGeneratorODE
+from ._CubicMeshPDEStatio import CubicMeshPDEStatio
+from ._CubicMeshPDENonStatio import CubicMeshPDENonStatio
+from ._DataGeneratorObservations import DataGeneratorObservations
+from ._DataGeneratorParameter import DataGeneratorParameter
 from ._Batchs import ODEBatch, PDEStatioBatch, PDENonStatioBatch
 
-from ._DataGenerators import append_obs_batch, append_param_batch
+from ._utils import append_obs_batch, append_param_batch
+
+__all__ = [
+    "DataGeneratorODE",
+    "CubicMeshPDEStatio",
+    "CubicMeshPDENonStatio",
+    "DataGeneratorParameter",
+    "DataGeneratorObservations",
+    "ODEBatch",
+    "PDEStatioBatch",
+    "PDENonStatioBatch",
+    "append_obs_batch",
+    "append_param_batch",
+]

jinns/data/_utils.py

@@ -0,0 +1,149 @@
+"""
+Utility functions for DataGenerators
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Key, Array, Float
+
+if TYPE_CHECKING:
+    from jinns.utils._types import AnyBatch
+    from jinns.data._Batchs import ObsBatchDict
+
+
+def append_param_batch(batch: AnyBatch, param_batch_dict: dict[str, Array]) -> AnyBatch:
+    """
+    Utility function that fills the field `batch.param_batch_dict` of a batch object.
+    """
+    return eqx.tree_at(
+        lambda m: m.param_batch_dict,
+        batch,
+        param_batch_dict,
+        is_leaf=lambda x: x is None,
+    )
+
+
+def append_obs_batch(batch: AnyBatch, obs_batch_dict: ObsBatchDict) -> AnyBatch:
+    """
+    Utility function that fills the field `batch.obs_batch_dict` of a batch object
+    """
+    return eqx.tree_at(
+        lambda m: m.obs_batch_dict, batch, obs_batch_dict, is_leaf=lambda x: x is None
+    )
+
+
+def make_cartesian_product(
+    b1: Float[Array, " batch_size dim1"], b2: Float[Array, " batch_size dim2"]
+) -> Float[Array, " rows=batch_size*batch_size (dim1+dim2)"]:
+    # rows= serves to disable jaxtyping wish for runtime check since it does not like the star
+    # operator, we wish use not as expected
+    """
+    Create the cartesian product of a time and a border omega batches
+    by tiling and repeating
+    """
+    n1 = b1.shape[0]
+    n2 = b2.shape[0]
+    b1 = jnp.repeat(b1, n2, axis=0)
+    b2 = jnp.tile(b2, reps=(n1,) + tuple(1 for i in b2.shape[1:]))
+    return jnp.concatenate([b1, b2], axis=1)
+
+
+def _reset_batch_idx_and_permute(
+    operands: tuple[Key, Float[Array, " n dimension"], int, None, Float[Array, " n"]],
+) -> tuple[Key, Float[Array, " n dimension"], int]:
+    key, domain, curr_idx, _, p = operands
+    # resetting counter
+    curr_idx = 0
+    # reshuffling
+    key, subkey = jax.random.split(key)
+    if p is None:
+        domain = jax.random.permutation(subkey, domain, axis=0, independent=False)
+    else:
+        # otherwise p is used to avoid collocation points not in n_start
+        # NOTE that replace=True to avoid undefined behaviour but then, the
+        # domain.shape[0] does not really grow as in the original RAR. instead,
+        # it always comprises the same number of points, but the points are
+        # updated
+        domain = jax.random.choice(
+            subkey, domain, shape=(domain.shape[0],), replace=True, p=p
+        )
+
+    # return updated
+    return (key, domain, curr_idx)
+
+
+def _increment_batch_idx(
+    operands: tuple[Key, Float[Array, " n dimension"], int, int, Float[Array, " n"]],
+) -> tuple[Key, Float[Array, " n dimension"], int]:
+    key, domain, curr_idx, batch_size, _ = operands
+    # simply increases counter and get the batch
+    curr_idx += batch_size
+    return (key, domain, curr_idx)
+
+
+def _reset_or_increment(
+    bend: int,
+    n_eff: int,
+    operands: tuple[Key, Float[Array, " n dimension"], int, int, Float[Array, " n"]],
+) -> tuple[Key, Float[Array, " n dimension"], int]:
+    """
+    Factorize the code of the jax.lax.cond which checks if we have seen all the
+    batches in an epoch
+    If bend > n_eff (ie n when no RAR sampling) we reshuffle and start from 0
+    again. Otherwise, if bend < n_eff, this means there are still *_batch_size
+    samples at least that have not been seen and we can take a new batch
+
+    Parameters
+    ----------
+    bend
+        An integer. The new hypothetical index for the starting of the batch
+    n_eff
+        An integer. The number of points to see to complete an epoch
+    operands
+        A tuple. As passed to _reset_batch_idx_and_permute and
+        _increment_batch_idx
+
+    Returns
+    -------
+    res
+        A tuple as returned by _reset_batch_idx_and_permute or
+        _increment_batch_idx
+    """
+    return jax.lax.cond(
+        bend > n_eff, _reset_batch_idx_and_permute, _increment_batch_idx, operands
+    )
+
+
+def _check_and_set_rar_parameters(
+    rar_parameters: dict, n: int, n_start: int
+) -> tuple[int, Float[Array, " n"] | None, int | None, int | None]:
+    if rar_parameters is not None and n_start is None:
+        raise ValueError(
+            "n_start must be provided in the context of RAR sampling scheme"
+        )
+
+    if rar_parameters is not None:
+        # Default p is None. However, in the RAR sampling scheme we use 0
+        # probability to specify non-used collocation points (i.e. points
+        # above n_start). Thus, p is a vector of probability of shape (nt, 1).
+        p = jnp.zeros((n,))
+        p = p.at[:n_start].set(1 / n_start)
+        # set internal counter for the number of gradient steps since the
+        # last new collocation points have been added
+        # It is not 0 to ensure the first iteration of RAR happens just
+        # after start_iter. See the _proceed_to_rar() function in _rar.py
+        rar_iter_from_last_sampling = rar_parameters["update_every"] - 1
+        # set iternal counter for the number of times collocation points
+        # have been added
+        rar_iter_nb = 0
+    else:
+        n_start = n
+        p = None
+        rar_iter_from_last_sampling = None
+        rar_iter_nb = None
+
+    return n_start, p, rar_iter_from_last_sampling, rar_iter_nb

jinns/experimental/__init__.py

@@ -6,3 +6,12 @@ from ._diffrax_solver import (
     neumann_boundary_condition,
     plot_diffrax_solution,
 )
+
+__all__ = [
+    "SpatialDiscretisation",
+    "reaction_diffusion_2d_vector_field",
+    "laplacian",
+    "dirichlet_boundary_condition",
+    "neumann_boundary_condition",
+    "plot_diffrax_solution",
+]