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

jinns/loss/__init__.py

@@ -1,22 +1,20 @@
 from ._DynamicLossAbstract import DynamicLoss, ODE, PDEStatio, PDENonStatio
-from ._LossODE import LossODE, SystemLossODE
-from ._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
+from ._LossODE import LossODE
+from ._LossPDE import LossPDEStatio, LossPDENonStatio
 from ._DynamicLoss import (
     GeneralizedLotkaVolterra,
     BurgersEquation,
     FPENonStatioLoss2D,
     OU_FPENonStatioLoss2D,
     FisherKPP,
-    MassConservation2DStatio,
-    NavierStokes2DStatio,
+    NavierStokesMassConservation2DStatio,
 )
 from ._loss_weights import (
     LossWeightsODE,
-    LossWeightsODEDict,
     LossWeightsPDENonStatio,
     LossWeightsPDEStatio,
-    LossWeightsPDEDict,
 )
+from ._loss_weight_updates import soft_adapt, lr_annealing, ReLoBRaLo
 
 from ._operators import (
     divergence_fwd,
@@ -26,3 +24,31 @@ from ._operators import (
     vectorial_laplacian_fwd,
     vectorial_laplacian_rev,
 )
+
+__all__ = [
+    "DynamicLoss",
+    "ODE",
+    "PDEStatio",
+    "PDENonStatio",
+    "LossODE",
+    "LossPDEStatio",
+    "LossPDENonStatio",
+    "GeneralizedLotkaVolterra",
+    "BurgersEquation",
+    "FPENonStatioLoss2D",
+    "OU_FPENonStatioLoss2D",
+    "FisherKPP",
+    "NavierStokesMassConservation2DStatio",
+    "LossWeightsODE",
+    "LossWeightsPDEStatio",
+    "LossWeightsPDENonStatio",
+    "divergence_fwd",
+    "divergence_rev",
+    "laplacian_fwd",
+    "laplacian_rev",
+    "vectorial_laplacian_fwd",
+    "vectorial_laplacian_rev",
+    "soft_adapt",
+    "lr_annealing",
+    "ReLoBRaLo",
+]

jinns/loss/_abstract_loss.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import abc
+from typing import TYPE_CHECKING, Self, Literal, Callable
+from jaxtyping import Array, PyTree, Key
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import optax
+from jinns.loss._loss_weights import AbstractLossWeights
+from jinns.parameters._params import Params
+from jinns.loss._loss_weight_updates import soft_adapt, lr_annealing, ReLoBRaLo
+
+if TYPE_CHECKING:
+    from jinns.utils._types import AnyLossComponents, AnyBatch
+
+
+class AbstractLoss(eqx.Module):
+    """
+    About the call:
+    https://github.com/patrick-kidger/equinox/issues/1002 + https://docs.kidger.site/equinox/pattern/
+    """
+
+    loss_weights: AbstractLossWeights
+    update_weight_method: Literal["soft_adapt", "lr_annealing", "ReLoBRaLo"] | None = (
+        eqx.field(kw_only=True, default=None, static=True)
+    )
+
+    @abc.abstractmethod
+    def __call__(self, *_, **__) -> Array:
+        pass
+
+    @abc.abstractmethod
+    def evaluate_by_terms(
+        self, params: Params[Array], batch: AnyBatch
+    ) -> tuple[AnyLossComponents, AnyLossComponents]:
+        pass
+
+    def get_gradients(
+        self, fun: Callable[[Params[Array]], Array], params: Params[Array]
+    ) -> tuple[Array, Array]:
+        """
+        params already filtered with derivative keys here
+        """
+        if fun is None:
+            return None, None
+        value_grad_loss = jax.value_and_grad(fun)
+        loss_val, grads = value_grad_loss(params)
+        return loss_val, grads
+
+    def ponderate_and_sum_loss(self, terms):
+        """
+        Get total loss from individual loss terms and weights
+
+        tree.leaves is needed to get rid of None from non used loss terms
+        """
+        weights = jax.tree.leaves(
+            self.loss_weights,
+            is_leaf=lambda x: eqx.is_inexact_array(x) and x is not None,
+        )
+        terms = jax.tree.leaves(
+            terms, is_leaf=lambda x: eqx.is_inexact_array(x) and x is not None
+        )
+        if len(weights) == len(terms):
+            return jnp.sum(jnp.array(weights) * jnp.array(terms))
+        else:
+            raise ValueError(
+                "The numbers of declared loss weights and "
+                "declared loss terms do not concord "
+                f" got {len(weights)} and {len(terms)}"
+            )
+
+    def ponderate_and_sum_gradient(self, terms):
+        """
+        Get total gradients from individual loss gradients and weights
+        for each parameter
+
+        tree.leaves is needed to get rid of None from non used loss terms
+        """
+        weights = jax.tree.leaves(
+            self.loss_weights,
+            is_leaf=lambda x: eqx.is_inexact_array(x) and x is not None,
+        )
+        grads = jax.tree.leaves(terms, is_leaf=lambda x: isinstance(x, Params))
+        # gradient terms for each individual loss for each parameter (several
+        # Params structures)
+        weights_pytree = jax.tree.map(
+            lambda w: optax.tree_utils.tree_full_like(grads[0], w), weights
+        )  # We need several Params structures full of the weight scalar
+        weighted_grads = jax.tree.map(
+            lambda w, p: w * p, weights_pytree, grads, is_leaf=eqx.is_inexact_array
+        )  # Now we can multiply
+        return jax.tree.map(
+            lambda *grads: jnp.sum(jnp.array(grads), axis=0),
+            *weighted_grads,
+            is_leaf=eqx.is_inexact_array,
+        )
+
+    def update_weights(
+        self: Self,
+        iteration_nb: int,
+        loss_terms: PyTree,
+        stored_loss_terms: PyTree,
+        grad_terms: PyTree,
+        key: Key,
+    ) -> Self:
+        """
+        Update the loss weights according to a predefined scheme
+        """
+        if self.update_weight_method == "soft_adapt":
+            new_weights = soft_adapt(
+                self.loss_weights, iteration_nb, loss_terms, stored_loss_terms
+            )
+        elif self.update_weight_method == "lr_annealing":
+            new_weights = lr_annealing(self.loss_weights, grad_terms)
+        elif self.update_weight_method == "ReLoBRaLo":
+            new_weights = ReLoBRaLo(
+                self.loss_weights, iteration_nb, loss_terms, stored_loss_terms, key
+            )
+        else:
+            raise ValueError("update_weight_method for loss weights not implemented")
+
+        # Below we update the non None entry in the PyTree self.loss_weights
+        # we directly get the non None entries because None is not treated as a
+        # leaf
+        return eqx.tree_at(
+            lambda pt: jax.tree.leaves(pt.loss_weights), self, new_weights
+        )

jinns/loss/_boundary_conditions.py

@@ -7,31 +7,31 @@ from __future__ import (
 )  # https://docs.python.org/3/library/typing.html#constant
 
 from typing import TYPE_CHECKING, Callable
+from jaxtyping import Array, Float
 import jax
 import jax.numpy as jnp
 from jax import vmap, grad
-import equinox as eqx
 from jinns.utils._utils import get_grid, _subtract_with_check
-from jinns.data._Batchs import *
+from jinns.data._Batchs import PDEStatioBatch, PDENonStatioBatch
 from jinns.nn._pinn import PINN
 from jinns.nn._spinn import SPINN
 
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.parameters._params import Params
+    from jinns.utils._types import BoundaryConditionFun
+    from jinns.nn._abstract_pinn import AbstractPINN
 
 
 def _compute_boundary_loss(
     boundary_condition_type: str,
-    f: Callable[
-        [Float[Array, "dim"] | Float[Array, "dim + 1"]], Float[Array, "dim_solution"]
-    ],
+    f: BoundaryConditionFun,
     batch: PDEStatioBatch | PDENonStatioBatch,
-    u: eqx.Module,
-    params: AnyParams,
+    u: AbstractPINN,
+    params: Params[Array],
     facet: int,
     dim_to_apply: slice,
     vmap_in_axes: tuple,
-) -> float:
+) -> Float[Array, " "]:
     r"""A generic function that will compute the mini-batch MSE of a
     boundary condition in the stationary case, resp. non-stationary, given by:
 
@@ -67,7 +67,7 @@ def _compute_boundary_loss(
     u
         a PINN
     params
-        Params or ParamsDict
+        Params
     facet
         An integer which represents the id of the facet which is currently
         considered (in the order provided by the DataGenerator which is fixed)
@@ -96,15 +96,15 @@ def _compute_boundary_loss(
 
 def boundary_dirichlet(
     f: Callable[
-        [Float[Array, "dim"] | Float[Array, "dim + 1"]], Float[Array, "dim_solution"]
+        [Float[Array, " dim"] | Float[Array, " dim + 1"]], Float[Array, " dim_solution"]
     ],
     batch: PDEStatioBatch | PDENonStatioBatch,
-    u: eqx.Module,
-    params: Params | ParamsDict,
+    u: AbstractPINN,
+    params: Params[Array],
     facet: int,
     dim_to_apply: slice,
     vmap_in_axes: tuple,
-) -> float:
+) -> Float[Array, " "]:
     r"""
     This omega boundary condition enforces a solution that is equal to `f`
     at `times_batch` x `omega_border` (non stationary case) or at `omega_border`
@@ -135,6 +135,7 @@ def boundary_dirichlet(
     vmap_in_axes
         A tuple object which specifies the in_axes of the vmapping
     """
+    assert batch.border_batch is not None
     batch_array = batch.border_batch
     batch_array = batch_array[..., facet]
 
@@ -168,15 +169,15 @@ def boundary_dirichlet(
 
 def boundary_neumann(
     f: Callable[
-        [Float[Array, "dim"] | Float[Array, "dim + 1"]], Float[Array, "dim_solution"]
+        [Float[Array, " dim"] | Float[Array, " dim + 1"]], Float[Array, " dim_solution"]
     ],
     batch: PDEStatioBatch | PDENonStatioBatch,
-    u: eqx.Module,
-    params: Params | ParamsDict,
+    u: AbstractPINN,
+    params: Params[Array],
     facet: int,
     dim_to_apply: slice,
     vmap_in_axes: tuple,
-) -> float:
+) -> Float[Array, " "]:
     r"""
     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
@@ -208,6 +209,7 @@ def boundary_neumann(
     vmap_in_axes
         A tuple object which specifies the in_axes of the vmapping
     """
+    assert batch.border_batch is not None
     batch_array = batch.border_batch
     batch_array = batch_array[..., facet]
 
@@ -223,7 +225,6 @@ def boundary_neumann(
         n = jnp.array([[-1, 1, 0, 0], [0, 0, -1, 1]])
 
     if isinstance(u, PINN):
-
         u_ = lambda inputs, params: jnp.squeeze(u(inputs, params)[dim_to_apply])
 
         if u.eq_type == "statio_PDE":

jinns/loss/_loss_components.py

@@ -0,0 +1,43 @@
+from typing import TypeVar, Generic
+from dataclasses import fields
+import equinox as eqx
+
+T = TypeVar("T")
+
+
+class XDEComponentsAbstract(eqx.Module, Generic[T]):
+    """
+    Provides a template for ODE components with generic types.
+    One can inherit to specialize and add methods and attributes
+    We do not enforce keyword only to avoid being to verbose (this then can
+    work like a tuple)
+    """
+
+    def items(self):
+        """
+        For the dataclass to be iterated like a dictionary.
+        Practical and retrocompatible with old code when loss components were
+        dictionaries
+        """
+        return {
+            field.name: getattr(self, field.name)
+            for field in fields(self)
+            if getattr(self, field.name) is not None
+        }.items()
+
+
+class ODEComponents(XDEComponentsAbstract[T]):
+    dyn_loss: T
+    initial_condition: T
+    observations: T
+
+
+class PDEStatioComponents(XDEComponentsAbstract[T]):
+    dyn_loss: T
+    norm_loss: T
+    boundary_loss: T
+    observations: T
+
+
+class PDENonStatioComponents(PDEStatioComponents[T]):
+    initial_condition: T