jinns 0.8.10__py3-none-any.whl → 1.0.0__py3-none-any.whl

jinns/loss/_loss_weights.py

@@ -0,0 +1,59 @@
+"""
+Formalize the loss weights data structure
+"""
+
+from typing import Dict
+from jaxtyping import Array, Float
+import equinox as eqx
+
+
+class LossWeightsODE(eqx.Module):
+
+    dyn_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+    initial_condition: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+    observations: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+
+
+class LossWeightsODEDict(eqx.Module):
+
+    dyn_loss: Dict[str, Array | Float | None] = eqx.field(kw_only=True, default=None)
+    initial_condition: Dict[str, Array | Float | None] = eqx.field(
+        kw_only=True, default=None
+    )
+    observations: Dict[str, Array | Float | None] = eqx.field(
+        kw_only=True, default=None
+    )
+
+
+class LossWeightsPDEStatio(eqx.Module):
+
+    dyn_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+    norm_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+    boundary_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+    observations: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+
+
+class LossWeightsPDENonStatio(eqx.Module):
+
+    dyn_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+    norm_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+    boundary_loss: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+    observations: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+    initial_condition: Array | Float | None = eqx.field(kw_only=True, default=1.0)
+
+
+class LossWeightsPDEDict(eqx.Module):
+    """
+    Only one type of LossWeights data structure for the SystemLossPDE:
+    Include the initial condition always for the code to be more generic
+    """
+
+    dyn_loss: Dict[str, Array | Float | None] = eqx.field(kw_only=True, default=1.0)
+    norm_loss: Dict[str, Array | Float | None] = eqx.field(kw_only=True, default=1.0)
+    boundary_loss: Dict[str, Array | Float | None] = eqx.field(
+        kw_only=True, default=1.0
+    )
+    observations: Dict[str, Array | Float | None] = eqx.field(kw_only=True, default=1.0)
+    initial_condition: Dict[str, Array | Float | None] = eqx.field(
+        kw_only=True, default=1.0
+    )

jinns/loss/_operators.py

@@ -5,15 +5,20 @@ Implements diverse operators for dynamic losses
 import jax
 import jax.numpy as jnp
 from jax import grad
+import equinox as eqx
+from jaxtyping import Float, Array
 from jinns.utils._pinn import PINN
 from jinns.utils._spinn import SPINN
+from jinns.parameters._params import Params
 
 
-def _div_rev(t, x, u, params):
+def _div_rev(
+    t: Float[Array, "1"], x: Float[Array, "dimension"], u: eqx.Module, params: Params
+) -> float:
     r"""
-    Compute the divergence of a vector field :math:`\mathbf{u}`, i.e.,
-    :math:`\nabla \cdot \mathbf{u}(\mathbf{x})` with :math:`\mathbf{u}` a vector
-    field from :math:`\mathbb{R}^d` to :math:`\mathbb{R}^d`.
+    Compute the divergence of a vector field $\mathbf{u}$, i.e.,
+    $\nabla \cdot \mathbf{u}(\mathbf{x})$ with $\mathbf{u}$ a vector
+    field from $\mathbb{R}^d$ to $\mathbb{R}^d$.
     The computation is done using backward AD
     """
 
@@ -28,15 +33,21 @@ def _div_rev(t, x, u, params):
     return jnp.sum(accu)
 
 
-def _div_fwd(t, x, u, params):
+def _div_fwd(
+    t: Float[Array, "1"], x: Float[Array, "dimension"], u: eqx.Module, params: Params
+) -> float:
     r"""
-    Compute the divergence of a **batched** vector field :math:`\mathbf{u}`, i.e.,
-    :math:`\nabla \cdot \mathbf{u}(\mathbf{x})` with :math:`\mathbf{u}` a vector
-    field from :math:`\mathbb{R}^{b \times d}` to :math:`\mathbb{R}^{b \times b
-    \times d}`. The result is then in :math:`\mathbb{R}^{b\times b}`.
+    Compute the divergence of a **batched** vector field $\mathbf{u}$, i.e.,
+    $\nabla \cdot \mathbf{u}(\mathbf{x})$ with $\mathbf{u}$ a vector
+    field from $\mathbb{R}^{b \times d}$ to $\mathbb{R}^{b \times b
+    \times d}$. The result is then in $\mathbb{R}^{b\times b}$.
     Because of the embedding that happens in SPINNs the
-    computation is most efficient with forward AD. This is the idea behind Separable PINNs.
-    This function is to be used in the context of SPINNs only.
+    computation is most efficient with forward AD. This is the idea behind
+    Separable PINNs.
+
+    !!! warning "Warning"
+
+        This function is to be used in the context of SPINNs only.
     """
 
     def scan_fun(_, i):
@@ -55,11 +66,13 @@ def _div_fwd(t, x, u, params):
     return jnp.sum(accu, axis=0)
 
 
-def _laplacian_rev(t, x, u, params):
+def _laplacian_rev(
+    t: Float[Array, "1"], x: Float[Array, "dimension"], u: eqx.Module, params: Params
+) -> float:
     r"""
-    Compute the Laplacian of a scalar field :math:`u` (from :math:`\mathbb{R}^d`
-    to :math:`\mathbb{R}`) for :math:`\mathbf{x}` of arbitrary dimension, i.e.,
-    :math:`\Delta u(\mathbf{x})=\nabla\cdot\nabla u(\mathbf{x})`.
+    Compute the Laplacian of a scalar field $u$ (from $\mathbb{R}^d$
+    to $\mathbb{R}$) for $\mathbf{x}$ of arbitrary dimension, i.e.,
+    $\Delta u(\mathbf{x})=\nabla\cdot\nabla u(\mathbf{x})$.
     The computation is done using backward AD.
     """
 
@@ -98,15 +111,24 @@ def _laplacian_rev(t, x, u, params):
     # return jnp.sum(trace_hessian)
 
 
-def _laplacian_fwd(t, x, u, params):
+def _laplacian_fwd(
+    t: Float[Array, "batch_size 1"],
+    x: Float[Array, "batch_size dimension"],
+    u: eqx.Module,
+    params: Params,
+) -> Float[Array, "batch_size batch_size"]:
     r"""
-    Compute the Laplacian of a **batched** scalar field :math:`u`
-    (from :math:`\mathbb{R}^{b\times d}` to :math:`\mathbb{R}^{b\times b}`)
-    for :math:`\mathbf{x}` of arbitrary dimension :math:`d` with batch
-    dimension :math:`b`.
+    Compute the Laplacian of a **batched** scalar field $u$
+    (from $\mathbb{R}^{b\times d}$ to $\mathbb{R}^{b\times b}$)
+    for $\mathbf{x}$ of arbitrary dimension $d$ with batch
+    dimension $b$.
     Because of the embedding that happens in SPINNs the
-    computation is most efficient with forward AD. This is the idea behind Separable PINNs.
-    This function is to be used in the context of SPINNs only.
+    computation is most efficient with forward AD. This is the idea behind
+    Separable PINNs.
+
+    !!! warning "Warning"
+
+        This function is to be used in the context of SPINNs only.
     """
 
     def scan_fun(_, i):
@@ -134,22 +156,30 @@ def _laplacian_fwd(t, x, u, params):
     return jnp.sum(trace_hessian, axis=0)
 
 
-def _vectorial_laplacian(t, x, u, params, u_vec_ndim=None):
+def _vectorial_laplacian(
+    t: Float[Array, "1"] | Float[Array, "batch_size 1"],
+    x: Float[Array, "dimension_in"] | Float[Array, "batch_size dimension"],
+    u: eqx.Module,
+    params: Params,
+    u_vec_ndim: int = None,
+) -> (
+    Float[Array, "dimension_out"] | Float[Array, "batch_size batch_size dimension_out"]
+):
     r"""
-    Compute the vectorial Laplacian of a vector field :math:`\mathbf{u}` (from
-    :math:`\mathbb{R}^d`
-    to :math:`\mathbb{R}^n`) for :math:`\mathbf{x}` of arbitrary dimension, i.e.,
-    :math:`\Delta \mathbf{u}(\mathbf{x})=\nabla\cdot\nabla
-    \mathbf{u}(\mathbf{x})`.
+    Compute the vectorial Laplacian of a vector field $\mathbf{u}$ (from
+    $\mathbb{R}^d$
+    to $\mathbb{R}^n$) for $\mathbf{x}$ of arbitrary dimension, i.e.,
+    $\Delta \mathbf{u}(\mathbf{x})=\nabla\cdot\nabla
+    \mathbf{u}(\mathbf{x})$.
 
     **Note:** We need to provide `u_vec_ndim` the dimension of the vector
-    :math:`\mathbf{u}(\mathbf{x})` if it is different than that of
-    :math:`\mathbf{x}`.
+    $\mathbf{u}(\mathbf{x})$ if it is different than that of
+    $\mathbf{x}$.
 
     **Note:** `u` can be a SPINN, in this case, it corresponds to a vector
-    field from (from :math:`\mathbb{R}^{b\times d}` to
-    :math:`\mathbb{R}^{b\times b\times n}`) and forward mode AD is used.
-    Technically, the return is of dimension :math:`n\times b \times b`.
+    field from (from $\mathbb{R}^{b\times d}$ to
+    $\mathbb{R}^{b\times b\times n}$) and forward mode AD is used.
+    Technically, the return is of dimension $n\times b \times b$.
     """
     if u_vec_ndim is None:
         u_vec_ndim = x.shape[0]
@@ -172,6 +202,8 @@ def _vectorial_laplacian(t, x, u, params, u_vec_ndim=None):
                     u(t, x, params)[..., j], axis=-1
                 )
             lap_on_j = _laplacian_fwd(t, x, uj, params)
+        else:
+            raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
 
         return _, lap_on_j
 
@@ -179,12 +211,14 @@ def _vectorial_laplacian(t, x, u, params, u_vec_ndim=None):
     return vec_lap
 
 
-def _u_dot_nabla_times_u_rev(t, x, u, params):
+def _u_dot_nabla_times_u_rev(
+    t: Float[Array, "1"], x: Float[Array, "2"], u: eqx.Module, params: Params
+) -> Float[Array, "2"]:
     r"""
-    Implement :math:`((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})` for
-    :math:`\mathbf{x}` of arbitrary
-    dimension. :math:`\mathbf{u}` is a vector field from :math:`\mathbb{R}^n`
-    to :math:`\mathbb{R}^n`. **Currently for** `x.ndim=2` **only**.
+    Implement $((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})$ for
+    $\mathbf{x}$ of arbitrary
+    dimension. $\mathbf{u}$ is a vector field from $\mathbb{R}^n$
+    to $\mathbb{R}^n$. **Currently for** `x.ndim=2` **only**.
     The computation is done using backward AD.
     We do not use loops but code explicitly the expression to avoid
     computing twice some terms
@@ -224,7 +258,12 @@ def _u_dot_nabla_times_u_rev(t, x, u, params):
     raise NotImplementedError("x.ndim must be 2")
 
 
-def _u_dot_nabla_times_u_fwd(t, x, u, params):
+def _u_dot_nabla_times_u_fwd(
+    t: Float[Array, "batch_size 1"],
+    x: Float[Array, "batch_size 2"],
+    u: eqx.Module,
+    params: Params,
+) -> Float[Array, "batch_size batch_size 2"]:
     r"""
     Implement :math:`((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})` for
     :math:`\mathbf{x}` of arbitrary dimension **with a batch dimension**.
@@ -264,36 +303,3 @@ def _u_dot_nabla_times_u_fwd(t, x, u, params):
             axis=-1,
         )
     raise NotImplementedError("x.ndim must be 2")
-
-
-def _sobolev(u, m, statio=True):
-    r"""
-    Compute the Sobolev regularization of order :math:`m`
-    of a scalar field :math:`u` (from :math:`\mathbb{R}^{d}` to :math:`\mathbb{R}`)
-    for :math:`\mathbf{x}` of arbitrary dimension :math:`d`, i.e.,
-    :math:`\frac{1}{n_l}\sum_{l=1}^{n_l}\sum_{|\alpha|=1}^{m+1} ||\partial^{\alpha} u(x_l)||_2^2` where
-    :math:`m\geq\max(d_1 // 2, K)` with :math:`K` the order of the differential
-    operator.
-
-    This regularization is proposed in *Convergence and error analysis of
-    PINNs*, Doumeche et al., 2023, https://arxiv.org/pdf/2305.01240.pdf
-    """
-
-    def jac_recursive(u, order, start):
-        # Compute the derivative of order `start`
-        if order == 0:
-            return u
-        if start == 0:
-            return jac_recursive(jax.jacrev(u), order - 1, start + 1)
-        return jac_recursive(jax.jacfwd(u), order - 1, start + 1)
-
-    if statio:
-        return lambda x, params: jnp.sum(
-            jac_recursive(lambda x: u(x, params), m + 1, 0)(x) ** 2
-        )
-    return lambda t, x, params: jnp.sum(
-        jac_recursive(lambda tx: u(tx[0:1], tx[1:], params), m + 1, 0)(
-            jnp.concatenate([t, x], axis=0)
-        )
-        ** 2
-    )

jinns/parameters/__init__.py

@@ -0,0 +1,6 @@
+from ._params import Params, ParamsDict
+from ._derivative_keys import (
+    DerivativeKeysODE,
+    DerivativeKeysPDEStatio,
+    DerivativeKeysPDENonStatio,
+)

jinns/parameters/_derivative_keys.py

@@ -0,0 +1,94 @@
+"""
+Formalize the data structure for the derivative keys
+"""
+
+from dataclasses import fields
+from typing import Literal
+import jax
+import equinox as eqx
+
+from jinns.parameters._params import Params
+
+
+class DerivativeKeysODE(eqx.Module):
+    # we use static = True because all fields are string, hence should be
+    # invisible by JAX transforms (JIT, etc.)
+    dyn_loss: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
+        kw_only=True, default="nn_params", static=True
+    )
+    observations: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
+        kw_only=True, default="nn_params", static=True
+    )
+    initial_condition: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
+        kw_only=True, default="nn_params", static=True
+    )
+
+
+class DerivativeKeysPDEStatio(eqx.Module):
+
+    dyn_loss: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
+        kw_only=True, default="nn_params", static=True
+    )
+    observations: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
+        kw_only=True, default="nn_params", static=True
+    )
+    boundary_loss: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
+        kw_only=True, default="nn_params", static=True
+    )
+    norm_loss: Literal["nn_params", "eq_params", "both"] | None = eqx.field(
+        kw_only=True, default="nn_params", static=True
+    )
+
+
+class DerivativeKeysPDENonStatio(DerivativeKeysPDEStatio):
+
+    initial_condition: Literal["nn_params", "eq_params", "both"] = eqx.field(
+        kw_only=True, default="nn_params", static=True
+    )
+
+
+def _set_derivatives(params, derivative_keys):
+    """
+    We construct an eqx.Module with the fields of derivative_keys, each field
+    has a copy of the params with appropriate derivatives set
+    """
+
+    def _set_derivatives_(loss_term_derivative):
+        if loss_term_derivative == "both":
+            return params
+        # the next line put a stop_gradient around the fields that do not
+        # appear in loss_term_derivative. Currently there are only two possible
+        # values nn_params and eq_params but there might be more in the future
+        return eqx.tree_at(
+            lambda p: tuple(
+                getattr(p, f.name)
+                for f in fields(Params)
+                if f.name != loss_term_derivative
+            ),
+            params,
+            replace_fn=jax.lax.stop_gradient,
+        )
+
+    def _set_derivatives_dict(loss_term_derivative):
+        if loss_term_derivative == "both":
+            return params
+        # the next line put a stop_gradient around the fields that do not
+        # appear in loss_term_derivative. Currently there are only two possible
+        # values nn_params and eq_params but there might be more in the future
+        return {
+            k: eqx.tree_at(
+                lambda p: tuple(
+                    getattr(p, f.name)
+                    for f in fields(Params)
+                    if f.name != loss_term_derivative
+                ),
+                params_,
+                replace_fn=jax.lax.stop_gradient,
+            )
+            for k, params_ in params
+        }
+
+    if not isinstance(params, dict):
+        return _set_derivatives_(derivative_keys)
+    else:
+        return _set_derivatives_dict(derivative_keys)

jinns/parameters/_params.py

@@ -0,0 +1,115 @@
+"""
+Formalize the data structure for the parameters
+"""
+
+import jax
+import equinox as eqx
+from typing import Dict
+from jaxtyping import Array, PyTree
+
+
+class Params(eqx.Module):
+    """
+    The equinox module for the parameters
+
+    Parameters
+    ----------
+    nn_params : Pytree
+        A PyTree of the non-static part of the PINN eqx.Module, i.e., the
+        parameters of the PINN
+    eq_params : Dict[str, Array]
+        A dictionary of the equation parameters. Keys are the parameter name,
+        values are their corresponding value
+    """
+
+    nn_params: PyTree = eqx.field(kw_only=True, default=None)
+    eq_params: Dict[str, Array] = eqx.field(kw_only=True, default=None)
+
+
+class ParamsDict(eqx.Module):
+    """
+    The equinox module for a dictionnary of parameters with different keys
+    corresponding to different equations.
+
+    Parameters
+    ----------
+    nn_params : Dict[str, PyTree]
+        The neural network's parameters. Most of the time, it will be the
+        Array part of an `eqx.Module` obtained by
+        `eqx.partition(module, eqx.is_inexact_array)`.
+    eq_params : Dict[str, Array]
+        A dictionary of the equation parameters. Dict keys are the parameter name as defined your custom loss.
+    """
+
+    nn_params: Dict[str, PyTree] = eqx.field(kw_only=True, default=None)
+    eq_params: Dict[str, Array] = eqx.field(kw_only=True, default=None)
+
+    def extract_params(self, nn_key: str) -> Params:
+        """
+        Extract the corresponding `nn_params` and `eq_params` for `nn_key` and
+        return them in the form of a `Params` object.
+        """
+        try:
+            return Params(
+                nn_params=self.nn_params[nn_key],
+                eq_params=self.eq_params[nn_key],
+            )
+        except (KeyError, IndexError) as e:
+            return Params(
+                nn_params=self.nn_params[nn_key],
+                eq_params=self.eq_params,
+            )
+
+
+def _update_eq_params_dict(
+    params: Params, param_batch_dict: Dict[str, Array]
+) -> Params:
+    """
+    Update params.eq_params with a batch of eq_params for given key(s)
+    """
+
+    # artificially "complete" `param_batch_dict` with None to match `params`
+    # PyTree  structure
+    param_batch_dict_ = param_batch_dict | {
+        k: None for k in set(params.eq_params.keys()) - set(param_batch_dict.keys())
+    }
+
+    # Replace at non None leafs
+    params = eqx.tree_at(
+        lambda p: p.eq_params,
+        params,
+        jax.tree_util.tree_map(
+            lambda p, q: q if q is not None else p,
+            params.eq_params,
+            param_batch_dict_,
+        ),
+    )
+
+    return params
+
+
+def _get_vmap_in_axes_params(
+    eq_params_batch_dict: Dict[str, Array], params: Params | ParamsDict
+) -> tuple[Params]:
+    """
+    Return the input vmap axes when there is batch(es) of parameters to vmap
+    over. The latter are designated by keys in eq_params_batch_dict.
+    If eq_params_batch_dict is None (i.e. no additional parameter batch), we
+    return (None,).
+    """
+    if eq_params_batch_dict is None:
+        return (None,)
+    # We use pytree indexing of vmapped axes and vmap on axis
+    # 0 of the eq_parameters for which we have a batch
+    # this is for a fine-grained vmaping
+    # scheme over the params
+    vmap_in_axes_params = (
+        type(params)(
+            nn_params=None,
+            eq_params={
+                k: (0 if k in eq_params_batch_dict.keys() else None)
+                for k in params.eq_params.keys()
+            },
+        ),
+    )
+    return vmap_in_axes_params

jinns/plot/__init__.py

@@ -0,0 +1,5 @@
+from jinns.plot._plot import (
+    plot2d,
+    plot1d_slice,
+    plot1d_image,
+)