jinns 1.6.1__py3-none-any.whl → 1.7.0__py3-none-any.whl

jinns/__init__.py

@@ -7,8 +7,9 @@ from jinns import parameters as parameters
 from jinns import plot as plot
 from jinns import nn as nn
 from jinns.solver._solve import solve
+from jinns.solver._solve_alternate import solve_alternate
 
-__all__ = ["nn", "solve"]
+__all__ = ["nn", "solve", "solve_alternate"]
 
 import warnings
 

jinns/loss/_LossODE.py

@@ -47,7 +47,11 @@ if TYPE_CHECKING:
     )
 
 
-class LossODE(AbstractLoss[LossWeightsODE, ODEBatch, ODEComponents[Array | None]]):
+class LossODE(
+    AbstractLoss[
+        LossWeightsODE, ODEBatch, ODEComponents[Array | None], DerivativeKeysODE
+    ]
+):
     r"""Loss object for an ordinary differential equation
 
     $$
@@ -57,7 +61,6 @@ class LossODE(AbstractLoss[LossWeightsODE, ODEBatch, ODEComponents[Array | None]
     where $\mathcal{N}[\cdot]$ is a differential operator and the
     initial condition is $u(t_0)=u_0$.
 
-
     Parameters
     ----------
     u : eqx.Module
@@ -107,7 +110,6 @@ class LossODE(AbstractLoss[LossWeightsODE, ODEBatch, ODEComponents[Array | None]
     # (ie. jax.Array cannot be static) and that we do not expect to change
     u: AbstractPINN
     dynamic_loss: ODE | None
-    vmap_in_axes: tuple[int] = eqx.field(static=True)
     derivative_keys: DerivativeKeysODE
     loss_weights: LossWeightsODE
     initial_condition: InitialCondition | None
@@ -131,10 +133,6 @@ class LossODE(AbstractLoss[LossWeightsODE, ODEBatch, ODEComponents[Array | None]
         else:
             self.loss_weights = loss_weights
 
-        super().__init__(loss_weights=self.loss_weights, **kwargs)
-        self.u = u
-        self.dynamic_loss = dynamic_loss
-        self.vmap_in_axes = (0,)
         if derivative_keys is None:
             # by default we only take gradient wrt nn_params
             if params is None:
@@ -142,9 +140,27 @@ class LossODE(AbstractLoss[LossWeightsODE, ODEBatch, ODEComponents[Array | None]
                     "Problem at derivative_keys initialization "
                     f"received {derivative_keys=} and {params=}"
                 )
-            self.derivative_keys = DerivativeKeysODE(params=params)
+            derivative_keys = DerivativeKeysODE(params=params)
         else:
-            self.derivative_keys = derivative_keys
+            derivative_keys = derivative_keys
+
+        super().__init__(
+            loss_weights=self.loss_weights,
+            derivative_keys=derivative_keys,
+            vmap_in_axes=(0,),
+            **kwargs,
+        )
+        self.u = u
+        self.dynamic_loss = dynamic_loss
+        if self.update_weight_method is not None and jnp.any(
+            jnp.array(jax.tree.leaves(self.loss_weights)) == 0
+        ):
+            warnings.warn(
+                "self.update_weight_method is activated while some loss "
+                "weights are zero. The update weight method will likely "
+                "update the zero weight to some non-zero value. Check that "
+                "this is the desired behaviour."
+            )
 
         if initial_condition is None:
             warnings.warn(
@@ -220,7 +236,11 @@ class LossODE(AbstractLoss[LossWeightsODE, ODEBatch, ODEComponents[Array | None]
             self.obs_slice = obs_slice
 
     def evaluate_by_terms(
-        self, params: Params[Array], batch: ODEBatch
+        self,
+        opt_params: Params[Array],
+        batch: ODEBatch,
+        *,
+        non_opt_params: Params[Array] | None = None,
     ) -> tuple[
         ODEComponents[Float[Array, " "] | None], ODEComponents[Float[Array, " "] | None]
     ]:
@@ -231,15 +251,22 @@ class LossODE(AbstractLoss[LossWeightsODE, ODEBatch, ODEComponents[Array | None]
 
         Parameters
         ---------
-        params
-            Parameters at which the loss is evaluated
+        opt_params
+            Parameters, which are optimized, 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
+        non_opt_params
+            Parameters, which are not optimized, at which the loss is evaluated
         """
+        if non_opt_params is not None:
+            params = eqx.combine(opt_params, non_opt_params)
+        else:
+            params = opt_params
+
         temporal_batch = batch.temporal_batch
 
         # Retrieve the optional eq_params_batch

jinns/loss/_LossPDE.py

@@ -68,11 +68,14 @@ B = TypeVar("B", bound=PDEStatioBatch | PDENonStatioBatch)
 C = TypeVar(
     "C", bound=PDEStatioComponents[Array | None] | PDENonStatioComponents[Array | None]
 )
-D = TypeVar("D", bound=DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio)
+DKPDE = TypeVar("DKPDE", bound=DerivativeKeysPDEStatio | DerivativeKeysPDENonStatio)
 Y = TypeVar("Y", bound=PDEStatio | PDENonStatio | None)
 
 
-class _LossPDEAbstract(AbstractLoss[L, B, C], Generic[L, B, C, D, Y]):
+class _LossPDEAbstract(
+    AbstractLoss[L, B, C, DKPDE],
+    Generic[L, B, C, DKPDE, Y],
+):
     r"""
     Parameters
     ----------
@@ -158,9 +161,24 @@ class _LossPDEAbstract(AbstractLoss[L, B, C], Generic[L, B, C, D, Y]):
         norm_weights: Float[Array, " nb_norm_samples"] | float | int | None = None,
         obs_slice: EllipsisType | slice | None = None,
         key: PRNGKeyArray | None = None,
+        derivative_keys: DKPDE,
         **kwargs: Any,  # for arguments for super()
     ):
-        super().__init__(loss_weights=self.loss_weights, **kwargs)
+        super().__init__(
+            loss_weights=self.loss_weights,
+            derivative_keys=derivative_keys,
+            **kwargs,
+        )
+
+        if self.update_weight_method is not None and jnp.any(
+            jnp.array(jax.tree.leaves(self.loss_weights)) == 0
+        ):
+            warnings.warn(
+                "self.update_weight_method is activated while some loss "
+                "weights are zero. The update weight method will likely "
+                "update the zero weight to some non-zero value. Check that "
+                "this is the desired behaviour."
+            )
 
         if obs_slice is None:
             self.obs_slice = jnp.s_[...]
@@ -497,7 +515,6 @@ class LossPDEStatio(
     dynamic_loss: PDEStatio | None
     loss_weights: LossWeightsPDEStatio
     derivative_keys: DerivativeKeysPDEStatio
-    vmap_in_axes: tuple[int] = eqx.field(static=True)
 
     params: InitVar[Params[Array] | None]
 
@@ -516,25 +533,25 @@ class LossPDEStatio(
             self.loss_weights = LossWeightsPDEStatio()
         else:
             self.loss_weights = loss_weights
-        self.dynamic_loss = dynamic_loss
-
-        super().__init__(
-            **kwargs,
-        )
 
         if derivative_keys is None:
             # be default we only take gradient wrt nn_params
             try:
-                self.derivative_keys = DerivativeKeysPDEStatio(params=params)
+                derivative_keys = DerivativeKeysPDEStatio(params=params)
             except ValueError as exc:
                 raise ValueError(
                     "Problem at derivative_keys initialization "
                     f"received {derivative_keys=} and {params=}"
                 ) from exc
         else:
-            self.derivative_keys = derivative_keys
+            derivative_keys = derivative_keys
 
-        self.vmap_in_axes = (0,)
+        super().__init__(
+            derivative_keys=derivative_keys,
+            vmap_in_axes=(0,),
+            **kwargs,
+        )
+        self.dynamic_loss = dynamic_loss
 
     def _get_dynamic_loss_batch(
         self, batch: PDEStatioBatch
@@ -549,11 +566,12 @@ class LossPDEStatio(
     # we could have used typing.cast though
 
     def evaluate_by_terms(
-        self, params: Params[Array], batch: PDEStatioBatch
-    ) -> tuple[
-        PDEStatioComponents[Float[Array, ""] | None],
-        PDEStatioComponents[Float[Array, ""] | None],
-    ]:
+        self,
+        opt_params: Params[Array],
+        batch: PDEStatioBatch,
+        *,
+        non_opt_params: Params[Array] | None = None,
+    ) -> tuple[PDEStatioComponents[Array | None], PDEStatioComponents[Array | None]]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
@@ -561,15 +579,22 @@ class LossPDEStatio(
 
         Parameters
         ---------
-        params
-            Parameters at which the loss is evaluated
+        opt_params
+            Parameters, which are optimized, 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
+        non_opt_params
+            Parameters, which are non optimized, at which the loss is evaluated
         """
+        if non_opt_params is not None:
+            params = eqx.combine(opt_params, non_opt_params)
+        else:
+            params = opt_params
+
         # Retrieve the optional eq_params_batch
         # and update eq_params with the latter
         # and update vmap_in_axes
@@ -740,7 +765,6 @@ class LossPDENonStatio(
     initial_condition_fun: Callable[[Float[Array, " dimension"]], Array] | None = (
         eqx.field(static=True)
     )
-    vmap_in_axes: tuple[int] = eqx.field(static=True)
     max_norm_samples_omega: int = eqx.field(static=True)
     max_norm_time_slices: int = eqx.field(static=True)
 
@@ -766,25 +790,26 @@ class LossPDENonStatio(
             self.loss_weights = LossWeightsPDENonStatio()
         else:
             self.loss_weights = loss_weights
-        self.dynamic_loss = dynamic_loss
-
-        super().__init__(
-            **kwargs,
-        )
 
         if derivative_keys is None:
             # be default we only take gradient wrt nn_params
             try:
-                self.derivative_keys = DerivativeKeysPDENonStatio(params=params)
+                derivative_keys = DerivativeKeysPDENonStatio(params=params)
             except ValueError as exc:
                 raise ValueError(
                     "Problem at derivative_keys initialization "
                     f"received {derivative_keys=} and {params=}"
                 ) from exc
         else:
-            self.derivative_keys = derivative_keys
+            derivative_keys = derivative_keys
 
-        self.vmap_in_axes = (0,)  # for t_x
+        super().__init__(
+            derivative_keys=derivative_keys,
+            vmap_in_axes=(0,),  # for t_x
+            **kwargs,
+        )
+
+        self.dynamic_loss = dynamic_loss
 
         if initial_condition_fun is None:
             warnings.warn(
@@ -820,7 +845,11 @@ class LossPDENonStatio(
         )
 
     def evaluate_by_terms(
-        self, params: Params[Array], batch: PDENonStatioBatch
+        self,
+        opt_params: Params[Array],
+        batch: PDENonStatioBatch,
+        *,
+        non_opt_params: Params[Array] | None = None,
     ) -> tuple[
         PDENonStatioComponents[Array | None], PDENonStatioComponents[Array | None]
     ]:
@@ -831,15 +860,22 @@ class LossPDENonStatio(
 
         Parameters
         ---------
-        params
-            Parameters at which the loss is evaluated
+        opt_params
+            Parameters, which are optimized, 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
+        non_opt_params
+            Parameters, which are non optimized, at which the loss is evaluated
         """
+        if non_opt_params is not None:
+            params = eqx.combine(opt_params, non_opt_params)
+        else:
+            params = opt_params
+
         omega_initial_batch = batch.initial_batch
         assert omega_initial_batch is not None
 

jinns/loss/_abstract_loss.py

@@ -9,7 +9,12 @@ import jax.numpy as jnp
 import optax
 from jinns.parameters._params import Params
 from jinns.loss._loss_weight_updates import soft_adapt, lr_annealing, ReLoBRaLo
-from jinns.utils._types import AnyLossComponents, AnyBatch, AnyLossWeights
+from jinns.utils._types import (
+    AnyLossComponents,
+    AnyBatch,
+    AnyLossWeights,
+    AnyDerivativeKeys,
+)
 
 L = TypeVar(
     "L", bound=AnyLossWeights
@@ -25,31 +30,47 @@ C = TypeVar(
     "C", bound=AnyLossComponents[Array | None]
 )  # The above comment also works with Unions (https://docs.python.org/3/library/typing.html#typing.TypeVar)
 
+DK = TypeVar("DK", bound=AnyDerivativeKeys)
+
 # In the cases above, without the bound, we could not have covariance on
 # the type because it would break LSP. Note that covariance on the return type
 # is authorized in LSP hence we do not need the same TypeVar instruction for
 # the return types of evaluate_by_terms for example!
 
 
-class AbstractLoss(eqx.Module, Generic[L, B, C]):
+class AbstractLoss(eqx.Module, Generic[L, B, C, DK]):
     """
     About the call:
     https://github.com/patrick-kidger/equinox/issues/1002 + https://docs.kidger.site/equinox/pattern/
     """
 
+    derivative_keys: eqx.AbstractVar[DK]
     loss_weights: eqx.AbstractVar[L]
     update_weight_method: Literal["soft_adapt", "lr_annealing", "ReLoBRaLo"] | None = (
         eqx.field(kw_only=True, default=None, static=True)
     )
+    vmap_in_axes: tuple[int] = eqx.field(static=True)
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         return self.evaluate(*args, **kwargs)
 
     @abc.abstractmethod
-    def evaluate_by_terms(self, params: Params[Array], batch: B) -> tuple[C, C]:
+    def evaluate_by_terms(
+        self,
+        opt_params: Params[Array],
+        batch: B,
+        *,
+        non_opt_params: Params[Array] | None = None,
+    ) -> tuple[C, C]:
         pass
 
-    def evaluate(self, params: Params[Array], batch: B) -> tuple[Float[Array, " "], C]:
+    def evaluate(
+        self,
+        opt_params: Params[Array],
+        batch: B,
+        *,
+        non_opt_params: Params[Array] | None = None,
+    ) -> tuple[Float[Array, " "], C]:
         """
         Evaluate the loss function at a batch of points for given parameters.
 
@@ -57,16 +78,20 @@ class AbstractLoss(eqx.Module, Generic[L, B, C]):
 
         Parameters
         ---------
-        params
-            Parameters at which the loss is evaluated
+        opt_params
+            Parameters, which are optimized, 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
+        non_opt_params
+            Parameters, which are non optimized, at which the loss is evaluated
         """
-        loss_terms, _ = self.evaluate_by_terms(params, batch)
+        loss_terms, _ = self.evaluate_by_terms(
+            opt_params, batch, non_opt_params=non_opt_params
+        )
 
         loss_val = self.ponderate_and_sum_loss(loss_terms)
 
@@ -105,7 +130,7 @@ class AbstractLoss(eqx.Module, Generic[L, B, C]):
             f" got {len(weights)} and {len(terms_list)}"
         )
 
-    def ponderate_and_sum_gradient(self, terms: C) -> C:
+    def ponderate_and_sum_gradient(self, terms: C) -> Params[Array | None]:
         """
         Get total gradients from individual loss gradients and weights
         for each parameter

jinns/parameters/_derivative_keys.py

@@ -47,9 +47,9 @@ class DerivativeKeysODE(eqx.Module):
     [`DynamicLoss`][jinns.loss.DynamicLoss] should be differentiated both with
     respect to the neural network parameters *and* the equation parameters, or only some of them.
 
-    To do so, user can either use strings or a `Params` object
-    with PyTree structure matching the parameters of the problem at
-    hand, and booleans indicating if gradient is to be taken or not. Internally,
+    To do so, user can either use strings or a `Params[bool]` object
+    with PyTree structure matching the parameters of the problem (`Params[Array]`) at
+    hand, and leaves being booleans indicating if gradient is to be taken or not. Internally,
     a `jax.lax.stop_gradient()` is appropriately set to each `True` node when
     computing each loss term.
 
@@ -156,12 +156,12 @@ class DerivativeKeysODE(eqx.Module):
         """
         Construct the DerivativeKeysODE from strings. For each term of the
         loss, specify whether to differentiate wrt the neural network
-        parameters, the equation parameters or both. The `Params` object, which
+        parameters, the equation parameters or both. The `Params[Array]` object, which
         contains the actual array of parameters must be passed to
         construct the fields with the appropriate PyTree structure.
 
         !!! note
-            You can mix strings and `Params` if you need granularity.
+            You can mix strings and `Params[bool]` if you need granularity.
 
         Parameters
         ----------
@@ -498,7 +498,14 @@ def _set_derivatives(
         `Params(nn_params=True | False, eq_params={"alpha":True | False,
         "beta":True | False})`.
         """
-
+        assert jax.tree.structure(params_.eq_params) == jax.tree.structure(
+            derivative_mask.eq_params
+        ), (
+            "The derivative "
+            "mask for eq_params does not have the same tree structure as "
+            "Params.eq_params. This is often due to a wrong Params[bool] "
+            "passed when initializing the derivative key object."
+        )
         return Params(
             nn_params=jax.lax.cond(
                 derivative_mask.nn_params,

jinns/parameters/_params.py

@@ -2,6 +2,7 @@
 Formalize the data structure for the parameters
 """
 
+from __future__ import annotations
 from dataclasses import fields
 from typing import Generic, TypeVar
 import equinox as eqx
@@ -60,6 +61,15 @@ class Params(eqx.Module, Generic[T]):
         else:
             self.eq_params = eq_params
 
+    def partition(self, mask: Params[bool] | None):
+        """
+        following the boolean mask, partition into two Params
+        """
+        if mask is not None:
+            return eqx.partition(self, mask)
+        else:
+            return self, None
+
 
 def update_eq_params(
     params: Params[Array],