jinns 0.9.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

jinns/loss/__init__.py

@@ -1,11 +1,27 @@
-from ._DynamicLossAbstract import ODE, PDEStatio, PDENonStatio
+from ._DynamicLossAbstract import DynamicLoss, ODE, PDEStatio, PDENonStatio
+from ._LossODE import LossODE, SystemLossODE
+from ._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
 from ._DynamicLoss import (
-    FisherKPP,
-    BurgerEquation,
     GeneralizedLotkaVolterra,
+    BurgerEquation,
+    FPENonStatioLoss2D,
     OU_FPENonStatioLoss2D,
+    FisherKPP,
     MassConservation2DStatio,
     NavierStokes2DStatio,
 )
-from ._LossPDE import LossPDENonStatio, LossPDEStatio, SystemLossPDE
-from ._LossODE import LossODE, SystemLossODE
+from ._loss_weights import (
+    LossWeightsODE,
+    LossWeightsODEDict,
+    LossWeightsPDENonStatio,
+    LossWeightsPDEStatio,
+    LossWeightsPDEDict,
+)
+
+from ._operators import (
+    _div_fwd,
+    _div_rev,
+    _laplacian_fwd,
+    _laplacian_rev,
+    _vectorial_laplacian,
+)

jinns/loss/_boundary_conditions.py

@@ -2,38 +2,53 @@
 Implements the main boundary conditions for all kinds of losses in jinns
 """
 
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
+
+from typing import TYPE_CHECKING, Callable
 import jax
 import jax.numpy as jnp
 from jax import vmap, grad
+import equinox as eqx
 from jinns.utils._utils import (
     _get_grid,
     _check_user_func_return,
-    _get_vmap_in_axes_params,
 )
-from jinns.data._DataGenerators import PDEStatioBatch, PDENonStatioBatch
+from jinns.parameters._params import _get_vmap_in_axes_params
+from jinns.data._Batchs import *
 from jinns.utils._pinn import PINN
 from jinns.utils._spinn import SPINN
 
+if TYPE_CHECKING:
+    from jinns.utils._types import *
+
 
 def _compute_boundary_loss(
-    boundary_condition_type, f, batch, u, params, facet, dim_to_apply
-):
+    boundary_condition_type: str,
+    f: Callable,
+    batch: PDEStatioBatch | PDENonStatioBatch,
+    u: eqx.Module,
+    params: Params | ParamsDict,
+    facet: int,
+    dim_to_apply: slice,
+) -> float:
     r"""A generic function that will compute the mini-batch MSE of a
     boundary condition in the stationary case, resp. non-stationary, given by:
 
-    .. math::
+    $$
         D[u](\partial x) = f(\partial x), \forall \partial x \in \partial \Omega
-
+    $$
     resp.,
 
-    .. math::
+    $$
         D[u](t, \partial x) = f(\partial x), \forall t \in I, \forall \partial
         x \in \partial \Omega
+    $$
 
+    Where $D[\cdot]$ is a differential operator, possibly identity.
 
-    Where :math:`D[\cdot]` is a differential operator, possibly identity.
-
-    __Note__: if using a batch.param_batch_dict, we need to resolve the
+    **Note**: if using a batch.param_batch_dict, we need to resolve the
     vmapping axes in the boundary functions,  however params["eq_params"]
     has already been fed with the batch in the `evaluate()` of `LossPDEStatio`,
     resp. `LossPDENonStatio`.
@@ -41,27 +56,24 @@ def _compute_boundary_loss(
     Parameters
     ----------
     boundary_condition_type
-        a string defining the differential operator :math:`D[\cdot]`.
-        Currently implements one of "Dirichlet" (:math:`D = Id`) and Von
-        Neuman (:math:`D[\cdot] = \nabla \cdot n`) where :math:`n` is the
-        unitary outgoing vector normal to :math:`\partial\Omega`
+        a string defining the differential operator $D[\cdot]$.
+        Currently implements one of "Dirichlet" ($D = Id$) and Von
+        Neuman ($D[u] = \nabla u \cdot n$) where $n$ is the
+        unitary outgoing vector normal to $\partial\Omega$
     f
         the function to be matched in the boundary condition. It should have
-        one argument only (other are ignored).
+        one or two arguments only (other are ignored).
     batch
-        a PDEStatioBatch object or PDENonStatioBatch
+        a PDEStatioBatch or PDENonStatioBatch
     u
         a PINN
     params
-        The dictionary of parameters of the model.
-        Typically, it is a dictionary of
-        dictionaries: `eq_params` and `nn_params``, respectively the
-        differential equation parameters and the neural network parameter
-    facet:
+        Params or ParamsDict
+    facet
         An integer which represents the id of the facet which is currently
         considered (in the order provided by the DataGenerator which is fixed)
     dim_to_apply
-        A jnp.s\_ object which indicates which dimension(s) of u will be forced
+        A `jnp.s_` object which indicates which dimension(s) of u will be forced
         to match the boundary condition
 
     Returns
@@ -91,7 +103,14 @@ def _compute_boundary_loss(
     return mse
 
 
-def boundary_dirichlet_statio(f, batch, u, params, facet, dim_to_apply):
+def boundary_dirichlet_statio(
+    f: Callable,
+    batch: PDEStatioBatch,
+    u: eqx.Module,
+    params: Params | ParamsDict,
+    facet: int,
+    dim_to_apply: slice,
+) -> float:
     r"""
     This omega boundary condition enforces a solution that is equal to f on
     border batch.
@@ -102,17 +121,14 @@ def boundary_dirichlet_statio(f, batch, u, params, facet, dim_to_apply):
 
     Parameters
     ----------
-    f:
+    f
         the constraint function
     batch
         A PDEStatioBatch object.
     u
         The PINN
     params
-        The dictionary of parameters of the model.
-        Typically, it is a dictionary of
-        dictionaries: `eq_params` and `nn_params``, respectively the
-        differential equation parameters and the neural network parameter
+        Params or ParamsDict
     dim_to_apply
         A jnp.s\_ object. The dimension of u on which to apply the boundary condition
     """
@@ -139,14 +155,23 @@ def boundary_dirichlet_statio(f, batch, u, params, facet, dim_to_apply):
             res**2,
             axis=-1,
         )
+    else:
+        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_u_boundary
 
 
-def boundary_neumann_statio(f, batch, u, params, facet, dim_to_apply):
+def boundary_neumann_statio(
+    f: Callable,
+    batch: PDEStatioBatch,
+    u: eqx.Module,
+    params: Params | ParamsDict,
+    facet: int,
+    dim_to_apply: slice,
+) -> float:
     r"""
-    This omega boundary condition enforces a solution where :math:`\nabla u\cdot
-    n` is equal to `f` on omega borders. :math:`n` is the unitary
-    outgoing vector normal at border :math:`\partial\Omega`.
+    This omega boundary condition enforces a solution where $\nabla u\cdot
+    n$ is equal to `f` on omega borders. $n$ is the unitary
+    outgoing vector normal at border $\partial\Omega$.
 
     __Note__: if using a batch.param_batch_dict, we need to resolve the
     vmapping axes here however params["eq_params"] has already been fed with
@@ -165,7 +190,7 @@ def boundary_neumann_statio(f, batch, u, params, facet, dim_to_apply):
         Typically, it is a dictionary of
         dictionaries: `eq_params` and `nn_params``, respectively the
         differential equation parameters and the neural network parameter
-    facet:
+    facet
         An integer which represents the id of the facet which is currently
         considered (in the order provided wy the DataGenerator which is fixed)
     dim_to_apply
@@ -248,13 +273,22 @@ def boundary_neumann_statio(f, batch, u, params, facet, dim_to_apply):
         boundaries = _check_user_func_return(f(x_grid), values.shape)
         res = values - boundaries
         mse_u_boundary = jnp.sum(res**2, axis=-1)
+    else:
+        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_u_boundary
 
 
-def boundary_dirichlet_nonstatio(f, batch, u, params, facet, dim_to_apply):
+def boundary_dirichlet_nonstatio(
+    f: Callable,
+    batch: PDENonStatioBatch,
+    u: eqx.Module,
+    params: Params | ParamsDict,
+    facet: int,
+    dim_to_apply: slice,
+) -> float:
     r"""
-    This omega boundary condition enforces a solution that is equal to f
-    at times_batch x omega borders
+    This omega boundary condition enforces a solution that is equal to `f`
+    at `times_batch` x `omega borders`
 
     __Note__: if using a batch.param_batch_dict, we need to resolve the
     vmapping axes here however params["eq_params"] has already been fed with
@@ -271,7 +305,7 @@ def boundary_dirichlet_nonstatio(f, batch, u, params, facet, dim_to_apply):
     params
         The dictionary of parameters of the model.
         Typically, it is a dictionary of
-        dictionaries: `eq_params` and `nn_params``, respectively the
+        dictionaries: `eq_params` and `nn_params`, respectively the
         differential equation parameters and the neural network parameter
     facet:
         An integer which represents the id of the facet which is currently
@@ -309,14 +343,24 @@ def boundary_dirichlet_nonstatio(f, batch, u, params, facet, dim_to_apply):
         )
         res = values - boundaries
         mse_u_boundary = jnp.sum(res**2, axis=-1)
+    else:
+        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_u_boundary
 
 
-def boundary_neumann_nonstatio(f, batch, u, params, facet, dim_to_apply):
+def boundary_neumann_nonstatio(
+    f: Callable,
+    batch: PDENonStatioBatch,
+    u: eqx.Module,
+    params: Params | ParamsDict,
+    facet: int,
+    dim_to_apply: slice,
+) -> float:
     r"""
-    This omega boundary condition enforces a solution where :math:`\nabla u\cdot
-    n` is equal to `f` at time_batch x omega borders. :math:`n` is the unitary
-    outgoing vector normal at border :math:`\partial\Omega`.
+    This omega boundary condition enforces a solution where $\nabla u\cdot
+    n$ is equal to `f` at the cartesian product of `time_batch` x `omega
+    borders`. $n$ is the unitary outgoing vector normal at border
+    $\partial\Omega$.
 
     __Note__: if using a batch.param_batch_dict, we need to resolve the
     vmapping axes here however params["eq_params"] has already been fed with
@@ -424,4 +468,6 @@ def boundary_neumann_nonstatio(f, batch, u, params, facet, dim_to_apply):
             res**2,
             axis=-1,
         )
+    else:
+        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_u_boundary

jinns/loss/{_Losses.py → _loss_utils.py}

@@ -2,22 +2,43 @@
 Interface for diverse loss functions to factorize code
 """
 
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
+
+from typing import TYPE_CHECKING, Callable, Dict
 import jax
 import jax.numpy as jnp
 from jax import vmap
+import equinox as eqx
+from jaxtyping import Float, Array, PyTree
 
-from jinns.utils._pinn import PINN
-from jinns.utils._spinn import SPINN
-from jinns.utils._hyperpinn import HYPERPINN
 from jinns.loss._boundary_conditions import (
     _compute_boundary_loss,
 )
 from jinns.utils._utils import _check_user_func_return, _get_grid
+from jinns.data._DataGenerators import (
+    append_obs_batch,
+)
+from jinns.utils._pinn import PINN
+from jinns.utils._spinn import SPINN
+from jinns.utils._hyperpinn import HYPERPINN
+from jinns.data._Batchs import *
+from jinns.parameters._params import Params, ParamsDict
+
+if TYPE_CHECKING:
+    from jinns.utils._types import *
 
 
 def dynamic_loss_apply(
-    dyn_loss, u, batches, params, vmap_axes, loss_weight, u_type=None
-):
+    dyn_loss: DynamicLoss,
+    u: eqx.Module,
+    batches: ODEBatch | PDEStatioBatch | PDENonStatioBatch,
+    params: Params | ParamsDict,
+    vmap_axes: tuple[int | None, ...],
+    loss_weight: float | Float[Array, "dyn_loss_dimension"],
+    u_type: PINN | HYPERPINN | None = None,
+) -> float:
     """
     Sometimes when u is a lambda function a or dict we do not have access to
     its type here, hence the last argument
@@ -35,11 +56,20 @@ def dynamic_loss_apply(
     elif u_type == SPINN or isinstance(u, SPINN):
         residuals = dyn_loss(*batches, u, params)
         mse_dyn_loss = jnp.mean(jnp.sum(loss_weight * residuals**2, axis=-1))
+    else:
+        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
 
     return mse_dyn_loss
 
 
-def normalization_loss_apply(u, batches, params, vmap_axes, int_length, loss_weight):
+def normalization_loss_apply(
+    u: eqx.Module,
+    batches: ODEBatch | PDEStatioBatch | PDENonStatioBatch,
+    params: Params | ParamsDict,
+    vmap_axes: tuple[int | None, ...],
+    int_length: int,
+    loss_weight: float,
+) -> float:
     # TODO merge stationary and non stationary cases
     if isinstance(u, (PINN, HYPERPINN)):
         if len(batches) == 1:
@@ -95,26 +125,38 @@ def normalization_loss_apply(u, batches, params, vmap_axes, int_length, loss_wei
                 )
                 ** 2
             )
+    else:
+        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
 
     return mse_norm_loss
 
 
 def boundary_condition_apply(
-    u,
-    batch,
-    params,
-    omega_boundary_fun,
-    omega_boundary_condition,
-    omega_boundary_dim,
-    loss_weight,
-):
+    u: eqx.Module,
+    batch: PDEStatioBatch | PDENonStatioBatch,
+    params: Params | ParamsDict,
+    omega_boundary_fun: Callable,
+    omega_boundary_condition: str,
+    omega_boundary_dim: int,
+    loss_weight: float | Float[Array, "boundary_cond_dim"],
+) -> float:
     if isinstance(omega_boundary_fun, dict):
         # We must create the facet tree dictionary as we do not have the
         # enumerate from the for loop to pass the id integer
-        if batch[1].shape[-1] == 2:
+        if (
+            isinstance(batch, PDEStatioBatch) and batch.border_batch.shape[-1] == 2
+        ) or (
+            isinstance(batch, PDENonStatioBatch)
+            and batch.times_x_border_batch.shape[-1] == 2
+        ):
             # 1D
             facet_tree = {"xmin": 0, "xmax": 1}
-        elif batch[1].shape[-1] == 4:
+        elif (
+            isinstance(batch, PDEStatioBatch) and batch.border_batch.shape[-1] == 4
+        ) or (
+            isinstance(batch, PDENonStatioBatch)
+            and batch.times_x_border_batch.shape[-1] == 4
+        ):
             # 2D
             facet_tree = {"xmin": 0, "xmax": 1, "ymin": 2, "ymax": 3}
         else:
@@ -138,7 +180,10 @@ def boundary_condition_apply(
         # Note that to keep the behaviour given in the comment above we neede
         # to specify is_leaf according to the note in the release of 0.4.29
     else:
-        facet_tuple = tuple(f for f in range(batch[1].shape[-1]))
+        if isinstance(batch, PDEStatioBatch):
+            facet_tuple = tuple(f for f in range(batch.border_batch.shape[-1]))
+        else:
+            facet_tuple = tuple(f for f in range(batch.times_x_border_batch.shape[-1]))
         b_losses_by_facet = jax.tree_util.tree_map(
             lambda fa: jnp.mean(
                 loss_weight
@@ -161,8 +206,14 @@ def boundary_condition_apply(
 
 
 def observations_loss_apply(
-    u, batches, params, vmap_axes, observed_values, loss_weight, obs_slice
-):
+    u: eqx.Module,
+    batches: ODEBatch | PDEStatioBatch | PDENonStatioBatch,
+    params: Params | ParamsDict,
+    vmap_axes: tuple[int | None, ...],
+    observed_values: Float[Array, "batch_size observation_dim"],
+    loss_weight: float | Float[Array, "observation_dim"],
+    obs_slice: slice,
+) -> float:
     # TODO implement for SPINN
     if isinstance(u, (PINN, HYPERPINN)):
         v_u = vmap(
@@ -181,12 +232,20 @@ def observations_loss_apply(
         )
     elif isinstance(u, SPINN):
         raise RuntimeError("observation loss term not yet implemented for SPINNs")
+    else:
+        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_observation_loss
 
 
 def initial_condition_apply(
-    u, omega_batch, params, vmap_axes, initial_condition_fun, n, loss_weight
-):
+    u: eqx.Module,
+    omega_batch: Float[Array, "dimension"],
+    params: Params | ParamsDict,
+    vmap_axes: tuple[int | None, ...],
+    initial_condition_fun: Callable,
+    n: int,
+    loss_weight: float | Float[Array, "initial_condition_dimension"],
+) -> float:
     if isinstance(u, (PINN, HYPERPINN)):
         v_u_t0 = vmap(
             lambda x, params: initial_condition_fun(x) - u(jnp.zeros((1,)), x, params),
@@ -212,25 +271,17 @@ def initial_condition_apply(
         )
         res = ini - v_ini
         mse_initial_condition = jnp.mean(jnp.sum(loss_weight * res**2, axis=-1))
+    else:
+        raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_initial_condition
 
 
-def sobolev_reg_apply(u, batches, params, vmap_axes, sobolev_reg, loss_weight):
-    # TODO implement for SPINN
-    if isinstance(u, (PINN, HYPERPINN)):
-        v_sob_reg = vmap(
-            lambda *args: sobolev_reg(*args),  # pylint: disable=E1121
-            vmap_axes,
-            0,
-        )
-        mse_sobolev_loss = loss_weight * jnp.mean(v_sob_reg(*batches, params))
-    elif isinstance(u, SPINN):
-        raise RuntimeError("Sobolev loss term not yet implemented for SPINNs")
-    return mse_sobolev_loss
-
-
 def constraints_system_loss_apply(
-    u_constraints_dict, batch, params_dict, loss_weights, loss_weight_struct
+    u_constraints_dict: Dict,
+    batch: ODEBatch | PDEStatioBatch | PDENonStatioBatch,
+    params_dict: ParamsDict,
+    loss_weights: Dict[str, float | Array],
+    loss_weight_struct: PyTree,
 ):
     """
     Same function for systemlossODE and systemlossPDE!
@@ -243,17 +294,17 @@ def constraints_system_loss_apply(
         loss_weights,
     )
 
-    if isinstance(params_dict["nn_params"], dict):
+    if isinstance(params_dict.nn_params, dict):
 
         def apply_u_constraint(
-            u_constraint, nn_params, loss_weights_for_u, obs_batch_u
+            u_constraint, nn_params, eq_params, loss_weights_for_u, obs_batch_u
         ):
             res_dict_for_u = u_constraint.evaluate(
-                {
-                    "nn_params": nn_params,
-                    "eq_params": params_dict["eq_params"],
-                },
-                batch._replace(obs_batch_dict=obs_batch_u),
+                Params(
+                    nn_params=nn_params,
+                    eq_params=eq_params,
+                ),
+                append_obs_batch(batch, obs_batch_u),
             )[1]
             res_dict_ponderated = jax.tree_util.tree_map(
                 lambda w, l: w * l, res_dict_for_u, loss_weights_for_u
@@ -267,7 +318,12 @@ def constraints_system_loss_apply(
         res_dict = jax.tree_util.tree_map(
             apply_u_constraint,
             u_constraints_dict,
-            params_dict["nn_params"],
+            params_dict.nn_params,
+            (
+                params_dict.eq_params
+                if params_dict.eq_params.keys() == params_dict.nn_params.keys()
+                else {k: params_dict.eq_params for k in params_dict.nn_params.keys()}
+            ),  # this manipulation is needed since we authorize eq_params not to have the same structure as nn_params in ParamsDict
             loss_weights_T,
             batch.obs_batch_dict,
             is_leaf=lambda x: (
@@ -283,7 +339,7 @@ def constraints_system_loss_apply(
         def apply_u_constraint(u_constraint, loss_weights_for_u, obs_batch_u):
             res_dict_for_u = u_constraint.evaluate(
                 params_dict,
-                batch._replace(obs_batch_dict=obs_batch_u),
+                append_obs_batch(batch, obs_batch_u),
             )[1]
             res_dict_ponderated = jax.tree_util.tree_map(
                 lambda w, l: w * l, res_dict_for_u, loss_weights_for_u

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
+    )