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

jinns/loss/_LossPDE.py

@@ -11,6 +11,7 @@ from dataclasses import InitVar
 from typing import TYPE_CHECKING, Callable, TypedDict
 from types import EllipsisType
 import warnings
+import jax
 import jax.numpy as jnp
 import equinox as eqx
 from jaxtyping import Float, Array, Key, Int
@@ -20,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,
@@ -31,16 +33,17 @@ from jinns.parameters._derivative_keys import (
     DerivativeKeysPDENonStatio,
 )
 from jinns.loss._abstract_loss import AbstractLoss
+from jinns.loss._loss_components import PDEStatioComponents, PDENonStatioComponents
 from jinns.loss._loss_weights import (
     LossWeightsPDEStatio,
     LossWeightsPDENonStatio,
 )
 from jinns.data._Batchs import PDEStatioBatch, PDENonStatioBatch
+from jinns.parameters._params import Params
 
 
 if TYPE_CHECKING:
     # imports for type hints only
-    from jinns.parameters._params import Params
     from jinns.nn._abstract_pinn import AbstractPINN
     from jinns.loss import PDENonStatio, PDEStatio
     from jinns.utils._types import BoundaryConditionFun
@@ -71,7 +74,10 @@ class _LossPDEAbstract(AbstractLoss):
         The loss weights for the differents term : dynamic loss,
         initial condition (if LossWeightsPDENonStatio), boundary conditions if
         any, normalization loss if any and observations if any.
-        All fields are set to 1.0 by default.
+        Can be updated according to a specific algorithm. See
+        `update_weight_method`
+    update_weight_method : Literal['soft_adapt', 'lr_annealing', 'ReLoBRaLo'], default=None
+        Default is None meaning no update for loss weights. Otherwise a string
     derivative_keys : DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio, default=None
         Specify which field of `params` should be differentiated for each
         composant of the total loss. Particularily useful for inverse problems.
@@ -336,7 +342,10 @@ class LossPDEStatio(_LossPDEAbstract):
         The loss weights for the differents term : dynamic loss,
         boundary conditions if any, normalization loss if any and
         observations if any.
-        All fields are set to 1.0 by default.
+        Can be updated according to a specific algorithm. See
+        `update_weight_method`
+    update_weight_method : Literal['soft_adapt', 'lr_annealing', 'ReLoBRaLo'], default=None
+        Default is None meaning no update for loss weights. Otherwise a string
     derivative_keys : DerivativeKeysPDEStatio, default=None
         Specify which field of `params` should be differentiated for each
         composant of the total loss. Particularily useful for inverse problems.
@@ -432,12 +441,13 @@ class LossPDEStatio(_LossPDEAbstract):
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
 
-    def evaluate(
+    def evaluate_by_terms(
         self, params: Params[Array], batch: PDEStatioBatch
-    ) -> tuple[Float[Array, " "], LossDictPDEStatio]:
+    ) -> tuple[PDEStatioComponents[Array | None], PDEStatioComponents[Array | None]]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
+        We retrieve two PyTrees with loss values and gradients for each term
 
         Parameters
         ---------
@@ -461,29 +471,27 @@ class LossPDEStatio(_LossPDEAbstract):
 
         # dynamic part
         if self.dynamic_loss is not None:
-            mse_dyn_loss = dynamic_loss_apply(
-                self.dynamic_loss.evaluate,
+            dyn_loss_fun = lambda p: dynamic_loss_apply(
+                self.dynamic_loss.evaluate,  # type: ignore
                 self.u,
                 self._get_dynamic_loss_batch(batch),
-                _set_derivatives(params, self.derivative_keys.dyn_loss),  # type: ignore
+                _set_derivatives(p, self.derivative_keys.dyn_loss),  # type: ignore
                 self.vmap_in_axes + vmap_in_axes_params,
-                self.loss_weights.dyn_loss,  # type: ignore
             )
         else:
-            mse_dyn_loss = jnp.array(0.0)
+            dyn_loss_fun = None
 
         # normalization part
         if self.norm_samples is not None:
-            mse_norm_loss = normalization_loss_apply(
+            norm_loss_fun = lambda p: normalization_loss_apply(
                 self.u,
                 self._get_normalization_loss_batch(batch),
-                _set_derivatives(params, self.derivative_keys.norm_loss),  # type: ignore
+                _set_derivatives(p, self.derivative_keys.norm_loss),  # type: ignore
                 vmap_in_axes_params,
                 self.norm_weights,  # type: ignore -> can't get the __post_init__ narrowing here
-                self.loss_weights.norm_loss,  # type: ignore
             )
         else:
-            mse_norm_loss = jnp.array(0.0)
+            norm_loss_fun = None
 
         # boundary part
         if (
@@ -491,47 +499,84 @@ class LossPDEStatio(_LossPDEAbstract):
             and self.omega_boundary_dim is not None
             and self.omega_boundary_fun is not None
         ):  # pyright cannot narrow down the three None otherwise as it is class attribute
-            mse_boundary_loss = boundary_condition_apply(
+            boundary_loss_fun = lambda p: boundary_condition_apply(
                 self.u,
                 batch,
-                _set_derivatives(params, self.derivative_keys.boundary_loss),  # type: ignore
-                self.omega_boundary_fun,
-                self.omega_boundary_condition,
-                self.omega_boundary_dim,
-                self.loss_weights.boundary_loss,  # type: ignore
+                _set_derivatives(p, self.derivative_keys.boundary_loss),  # type: ignore
+                self.omega_boundary_fun,  # type: ignore
+                self.omega_boundary_condition,  # type: ignore
+                self.omega_boundary_dim,  # type: ignore
             )
         else:
-            mse_boundary_loss = jnp.array(0.0)
+            boundary_loss_fun = None
 
         # Observation mse
         if batch.obs_batch_dict is not None:
             # update params with the batches of observed params
-            params = _update_eq_params_dict(params, batch.obs_batch_dict["eq_params"])
+            params_obs = _update_eq_params_dict(
+                params, batch.obs_batch_dict["eq_params"]
+            )
 
-            mse_observation_loss = observations_loss_apply(
+            obs_loss_fun = lambda po: observations_loss_apply(
                 self.u,
                 self._get_observations_loss_batch(batch),
-                _set_derivatives(params, self.derivative_keys.observations),  # type: ignore
+                _set_derivatives(po, self.derivative_keys.observations),  # type: ignore
                 self.vmap_in_axes + vmap_in_axes_params,
                 batch.obs_batch_dict["val"],
-                self.loss_weights.observations,  # type: ignore
                 self.obs_slice,
             )
         else:
-            mse_observation_loss = jnp.array(0.0)
+            params_obs = None
+            obs_loss_fun = None
 
-        # total loss
-        total_loss = (
-            mse_dyn_loss + mse_norm_loss + mse_boundary_loss + mse_observation_loss
+        # get the unweighted mses for each loss term as well as the gradients
+        all_funs: PDEStatioComponents[Callable[[Params[Array]], Array] | None] = (
+            PDEStatioComponents(
+                dyn_loss_fun, norm_loss_fun, boundary_loss_fun, obs_loss_fun
+            )
+        )
+        all_params: PDEStatioComponents[Params[Array] | None] = PDEStatioComponents(
+            params, params, params, params_obs
         )
-        return total_loss, (
-            {
-                "dyn_loss": mse_dyn_loss,
-                "norm_loss": mse_norm_loss,
-                "boundary_loss": mse_boundary_loss,
-                "observations": mse_observation_loss,
-            }
+        mses_grads = jax.tree.map(
+            lambda fun, params: self.get_gradients(fun, params),
+            all_funs,
+            all_params,
+            is_leaf=lambda x: x is None,
         )
+        mses = jax.tree.map(
+            lambda leaf: leaf[0], mses_grads, is_leaf=lambda x: isinstance(x, tuple)
+        )
+        grads = jax.tree.map(
+            lambda leaf: leaf[1], mses_grads, is_leaf=lambda x: isinstance(x, tuple)
+        )
+
+        return mses, grads
+
+    def evaluate(
+        self, params: Params[Array], batch: PDEStatioBatch
+    ) -> tuple[Float[Array, " "], PDEStatioComponents[Float[Array, " "] | None]]:
+        """
+        Evaluate the loss function at a batch of points for given parameters.
+
+        We retrieve the total value itself and a PyTree with loss values for each term
+
+        Parameters
+        ---------
+        params
+            Parameters at which the loss is evaluated
+        batch
+            Composed of a batch of points in the
+            domain, a batch of points in the domain
+            border and an optional additional batch of parameters (eg. for
+            metamodeling) and an optional additional batch of observed
+            inputs/outputs/parameters
+        """
+        loss_terms, _ = self.evaluate_by_terms(params, batch)
+
+        loss_val = self.ponderate_and_sum_loss(loss_terms)
+
+        return loss_val, loss_terms
 
 
 class LossPDENonStatio(LossPDEStatio):
@@ -569,7 +614,10 @@ class LossPDENonStatio(LossPDEStatio):
         The loss weights for the differents term : dynamic loss,
         boundary conditions if any, initial condition, normalization loss if any and
         observations if any.
-        All fields are set to 1.0 by default.
+        Can be updated according to a specific algorithm. See
+        `update_weight_method`
+    update_weight_method : Literal['soft_adapt', 'lr_annealing', 'ReLoBRaLo'], default=None
+        Default is None meaning no update for loss weights. Otherwise a string
     derivative_keys : DerivativeKeysPDENonStatio, default=None
         Specify which field of `params` should be differentiated for each
         composant of the total loss. Particularily useful for inverse problems.
@@ -653,20 +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 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
@@ -696,22 +736,25 @@ class LossPDENonStatio(LossPDEStatio):
     def __call__(self, *args, **kwargs):
         return self.evaluate(*args, **kwargs)
 
-    def evaluate(
+    def evaluate_by_terms(
         self, params: Params[Array], batch: PDENonStatioBatch
-    ) -> tuple[Float[Array, " "], LossDictPDENonStatio]:
+    ) -> tuple[
+        PDENonStatioComponents[Array | None], PDENonStatioComponents[Array | None]
+    ]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
+        We retrieve two PyTrees with loss values and gradients for each term
 
         Parameters
         ---------
         params
             Parameters at which the loss is evaluated
         batch
-            Composed of a batch of points in
-            the domain, a batch of points in the domain
-            border, a batch of time points and an optional additional batch
-            of parameters (eg. for metamodeling) and an optional additional batch of observed
+            Composed of a batch of points in the
+            domain, a batch of points in the domain
+            border and an optional additional batch of parameters (eg. for
+            metamodeling) and an optional additional batch of observed
             inputs/outputs/parameters
         """
         omega_batch = batch.initial_batch
@@ -728,27 +771,62 @@ class LossPDENonStatio(LossPDEStatio):
 
         # For mse_dyn_loss, mse_norm_loss, mse_boundary_loss,
         # mse_observation_loss we use the evaluate from parent class
-        partial_mse, partial_mse_terms = super().evaluate(params, batch)  # type: ignore
+        # As well as for their gradients
+        partial_mses, partial_grads = super().evaluate_by_terms(params, batch)  # type: ignore
         # ignore because batch is not PDEStatioBatch. We could use typing.cast though
 
         # initial condition
         if self.initial_condition_fun is not None:
-            mse_initial_condition = initial_condition_apply(
+            mse_initial_condition_fun = lambda p: initial_condition_apply(
                 self.u,
                 omega_batch,
-                _set_derivatives(params, self.derivative_keys.initial_condition),  # type: ignore
+                _set_derivatives(p, self.derivative_keys.initial_condition),  # type: ignore
                 (0,) + vmap_in_axes_params,
-                self.initial_condition_fun,
+                self.initial_condition_fun,  # type: ignore
                 self.t0,  # type: ignore can't get the narrowing in __post_init__
-                self.loss_weights.initial_condition,  # type: ignore
+            )
+            mse_initial_condition, grad_initial_condition = self.get_gradients(
+                mse_initial_condition_fun, params
             )
         else:
-            mse_initial_condition = jnp.array(0.0)
+            mse_initial_condition = None
+            grad_initial_condition = None
+
+        mses = PDENonStatioComponents(
+            partial_mses.dyn_loss,
+            partial_mses.norm_loss,
+            partial_mses.boundary_loss,
+            partial_mses.observations,
+            mse_initial_condition,
+        )
+
+        grads = PDENonStatioComponents(
+            partial_grads.dyn_loss,
+            partial_grads.norm_loss,
+            partial_grads.boundary_loss,
+            partial_grads.observations,
+            grad_initial_condition,
+        )
+
+        return mses, grads
+
+    def evaluate(
+        self, params: Params[Array], batch: PDENonStatioBatch
+    ) -> tuple[Float[Array, " "], PDENonStatioComponents[Float[Array, " "] | None]]:
+        """
+        Evaluate the loss function at a batch of points for given parameters.
+        We retrieve the total value itself and a PyTree with loss values for each term
 
-        # total loss
-        total_loss = partial_mse + mse_initial_condition
 
-        return total_loss, {
-            **partial_mse_terms,
-            "initial_condition": mse_initial_condition,
-        }
+        Parameters
+        ---------
+        params
+            Parameters at which the loss is evaluated
+        batch
+            Composed of a batch of points in
+            the domain, a batch of points in the domain
+            border, a batch of time points and an optional additional batch
+            of parameters (eg. for metamodeling) and an optional additional batch of observed
+            inputs/outputs/parameters
+        """
+        return super().evaluate(params, batch)  # type: ignore

jinns/loss/__init__.py

@@ -14,6 +14,7 @@ from ._loss_weights import (
     LossWeightsPDENonStatio,
     LossWeightsPDEStatio,
 )
+from ._loss_weight_updates import soft_adapt, lr_annealing, ReLoBRaLo
 
 from ._operators import (
     divergence_fwd,
@@ -47,4 +48,7 @@ __all__ = [
     "laplacian_rev",
     "vectorial_laplacian_fwd",
     "vectorial_laplacian_rev",
+    "soft_adapt",
+    "lr_annealing",
+    "ReLoBRaLo",
 ]

jinns/loss/_abstract_loss.py

@@ -1,15 +1,128 @@
+from __future__ import annotations
+
 import abc
-from jaxtyping import Array
+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):
     """
-    Basically just a way to add a __call__ to an eqx.Module.
-    The way to go for correct type hints apparently
+    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/_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