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

jinns/loss/_loss_utils.py

@@ -40,7 +40,6 @@ def dynamic_loss_apply(
     ),
     params: Params[Array],
     vmap_axes: tuple[int, Params[int | None] | None],
-    loss_weight: float | Float[Array, " dyn_loss_dimension"],
     u_type: PINN | HyperPINN | None = None,
 ) -> Float[Array, " "]:
     """
@@ -58,10 +57,10 @@ def dynamic_loss_apply(
             0,
         )
         residuals = v_dyn_loss(batch, params)
-        mse_dyn_loss = jnp.mean(jnp.sum(loss_weight * residuals**2, axis=-1))
+        mse_dyn_loss = jnp.mean(jnp.sum(residuals**2, axis=-1))
     elif u_type == SPINN or isinstance(u, SPINN):
         residuals = dyn_loss(batch, u, params)
-        mse_dyn_loss = jnp.mean(jnp.sum(loss_weight * residuals**2, axis=-1))
+        mse_dyn_loss = jnp.mean(jnp.sum(residuals**2, axis=-1))
     else:
         raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
 
@@ -79,7 +78,6 @@ def normalization_loss_apply(
     params: Params[Array],
     vmap_axes_params: tuple[Params[int | None] | None],
     norm_weights: Float[Array, " nb_norm_samples"],
-    loss_weight: float,
 ) -> Float[Array, " "]:
     """
     Note the squeezing on each result. We expect unidimensional *PINN since
@@ -95,9 +93,7 @@ def normalization_loss_apply(
             res = v_u(*batches, params)
             assert res.shape[-1] == 1, "norm loss expects unidimensional *PINN"
             # Monte-Carlo integration using importance sampling
-            mse_norm_loss = loss_weight * (
-                jnp.abs(jnp.mean(res.squeeze() * norm_weights) - 1) ** 2
-            )
+            mse_norm_loss = jnp.abs(jnp.mean(res.squeeze() * norm_weights) - 1) ** 2
         else:
             # NOTE this cartesian product is costly
             batch_cart_prod = make_cartesian_product(
@@ -115,7 +111,7 @@ def normalization_loss_apply(
             assert res.shape[-1] == 1, "norm loss expects unidimensional *PINN"
             # For all times t, we perform an integration. Then we average the
             # losses over times.
-            mse_norm_loss = loss_weight * jnp.mean(
+            mse_norm_loss = jnp.mean(
                 jnp.abs(jnp.mean(res.squeeze() * norm_weights, axis=-1) - 1) ** 2
             )
     elif isinstance(u, SPINN):
@@ -123,8 +119,7 @@ def normalization_loss_apply(
             res = u(*batches, params)
             assert res.shape[-1] == 1, "norm loss expects unidimensional *SPINN"
             mse_norm_loss = (
-                loss_weight
-                * jnp.abs(
+                jnp.abs(
                     jnp.mean(
                         res.squeeze(),
                     )
@@ -144,7 +139,7 @@ def normalization_loss_apply(
             )
             assert res.shape[-1] == 1, "norm loss expects unidimensional *SPINN"
             # the outer mean() below is for the times stamps
-            mse_norm_loss = loss_weight * jnp.mean(
+            mse_norm_loss = jnp.mean(
                 jnp.abs(
                     jnp.mean(
                         res.squeeze(),
@@ -168,7 +163,6 @@ def boundary_condition_apply(
     omega_boundary_fun: BoundaryConditionFun | dict[str, BoundaryConditionFun],
     omega_boundary_condition: str | dict[str, str],
     omega_boundary_dim: slice | dict[str, slice],
-    loss_weight: float | Float[Array, " boundary_cond_dim"],
 ) -> Float[Array, " "]:
     assert batch.border_batch is not None
     vmap_in_axes = (0,) + _get_vmap_in_axes_params(batch.param_batch_dict, params)
@@ -205,10 +199,7 @@ def boundary_condition_apply(
                 None
                 if c is None
                 else jnp.mean(
-                    loss_weight
-                    * _compute_boundary_loss(
-                        c, f, batch, u, params, fa, d, vmap_in_axes
-                    )
+                    _compute_boundary_loss(c, f, batch, u, params, fa, d, vmap_in_axes)
                 )
             ),
             omega_boundary_dicts[0],  # omega_boundary_condition,
@@ -225,8 +216,7 @@ def boundary_condition_apply(
         facet_tuple = tuple(f for f in range(batch.border_batch.shape[-1]))
         b_losses_by_facet = jax.tree_util.tree_map(
             lambda fa: jnp.mean(
-                loss_weight
-                * _compute_boundary_loss(
+                _compute_boundary_loss(
                     omega_boundary_dicts[0],  # type: ignore -> need TypeIs from 3.13
                     omega_boundary_dicts[1],  # type: ignore -> need TypeIs from 3.13
                     batch,
@@ -251,7 +241,6 @@ def observations_loss_apply(
     params: Params[Array],
     vmap_axes: tuple[int, Params[int | None] | None],
     observed_values: Float[Array, " obs_batch_size observation_dim"],
-    loss_weight: float | Float[Array, " observation_dim"],
     obs_slice: EllipsisType | slice | None,
 ) -> Float[Array, " "]:
     if isinstance(u, (PINN, HyperPINN)):
@@ -263,8 +252,7 @@ def observations_loss_apply(
         val = v_u(batch, params)[:, obs_slice]
         mse_observation_loss = jnp.mean(
             jnp.sum(
-                loss_weight
-                * _subtract_with_check(
+                _subtract_with_check(
                     observed_values, val, cause="user defined observed_values"
                 )
                 ** 2,
@@ -285,7 +273,6 @@ def initial_condition_apply(
     vmap_axes: tuple[int, Params[int | None] | None],
     initial_condition_fun: Callable,
     t0: Float[Array, " 1"],
-    loss_weight: float | Float[Array, " initial_condition_dimension"],
 ) -> Float[Array, " "]:
     n = omega_batch.shape[0]
     t0_omega_batch = jnp.concatenate([t0 * jnp.ones((n, 1)), omega_batch], axis=1)
@@ -304,7 +291,7 @@ def initial_condition_apply(
         # dimension as params to be able to vmap.
         # Recall that by convention:
         # param_batch_dict = times_batch_size * omega_batch_size
-        mse_initial_condition = jnp.mean(jnp.sum(loss_weight * res**2, axis=-1))
+        mse_initial_condition = jnp.mean(jnp.sum(res**2, axis=-1))
     elif isinstance(u, SPINN):
         values = lambda t_x: u(
             t_x,
@@ -317,7 +304,30 @@ def initial_condition_apply(
             v_ini,
             cause="Output of initial_condition_fun",
         )
-        mse_initial_condition = jnp.mean(jnp.sum(loss_weight * res**2, axis=-1))
+        mse_initial_condition = jnp.mean(jnp.sum(res**2, axis=-1))
     else:
         raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_initial_condition
+
+
+def initial_condition_check(x, dim_size=None):
+    """
+    Make a (dim_size,) jnp array from an int, a float or a 0D jnp array
+
+    """
+    if isinstance(x, Array):
+        if not x.shape:  # e.g. user input: jnp.array(0.)
+            x = jnp.array([x])
+        if dim_size is not None:  # we check for the required dims_ize
+            if x.shape != (dim_size,):
+                raise ValueError(
+                    f"Wrong dim_size. It should be({dim_size},). Got shape: {x.shape}"
+                )
+
+    elif isinstance(x, float):  # e.g. user input: 0.
+        x = jnp.array([x])
+    elif isinstance(x, int):  # e.g. user input: 0
+        x = jnp.array([float(x)])
+    else:
+        raise ValueError(f"Wrong value, expected Array, float or int, got {type(x)}")
+    return x

jinns/loss/_loss_weight_updates.py

@@ -0,0 +1,202 @@
+"""
+A collection of specific weight update schemes in jinns
+"""
+
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from jaxtyping import Array, Key
+import jax.numpy as jnp
+import jax
+import equinox as eqx
+
+if TYPE_CHECKING:
+    from jinns.loss._loss_weights import AbstractLossWeights
+    from jinns.utils._types import AnyLossComponents
+
+
+def soft_adapt(
+    loss_weights: AbstractLossWeights,
+    iteration_nb: int,
+    loss_terms: AnyLossComponents,
+    stored_loss_terms: AnyLossComponents,
+) -> Array:
+    r"""
+    Implement the simple strategy given in
+    https://docs.nvidia.com/deeplearning/physicsnemo/physicsnemo-sym/user_guide/theory/advanced_schemes.html#softadapt
+
+    $$
+    w_j(i)= \frac{\exp(\frac{L_j(i)}{L_j(i-1)+\epsilon}-\mu(i))}
+    {\sum_{k=1}^{n_{loss}}\exp(\frac{L_k(i)}{L_k(i-1)+\epsilon}-\mu(i)}
+    $$
+
+    Note that since None is not treated as a leaf by jax tree.util functions,
+    we naturally avoid None components from loss_terms, stored_loss_terms etc.!
+    """
+
+    def do_nothing(loss_weights, _, __):
+        return jnp.array(
+            jax.tree.leaves(loss_weights, is_leaf=eqx.is_inexact_array), dtype=float
+        )
+
+    def soft_adapt_(_, loss_terms, stored_loss_terms):
+        ratio_pytree = jax.tree.map(
+            lambda lt, slt: lt / (slt[iteration_nb - 1] + 1e-6),
+            loss_terms,
+            stored_loss_terms,
+        )
+        mu = jax.tree.reduce(jnp.maximum, ratio_pytree, initializer=jnp.array(-jnp.inf))
+        ratio_pytree = jax.tree.map(lambda r: r - mu, ratio_pytree)
+        ratio_leaves = jax.tree.leaves(ratio_pytree)
+        return jax.nn.softmax(jnp.array(ratio_leaves))
+
+    return jax.lax.cond(
+        iteration_nb == 0,
+        lambda op: do_nothing(*op),
+        lambda op: soft_adapt_(*op),
+        (loss_weights, loss_terms, stored_loss_terms),
+    )
+
+
+def ReLoBRaLo(
+    loss_weights: AbstractLossWeights,
+    iteration_nb: int,
+    loss_terms: AnyLossComponents,
+    stored_loss_terms: AnyLossComponents,
+    key: Key,
+    decay_factor: float = 0.9,
+    tau: float = 1,  ## referred to as temperature in the article
+    p: float = 0.9,
+):
+    r"""
+    Implementing the extension of softadapt: Relative Loss Balancing with random LookBack
+    """
+    n_loss = len(jax.tree.leaves(loss_terms))  # number of loss terms
+    epsilon = 1e-6
+
+    def do_nothing(loss_weights, _):
+        return jnp.array(
+            jax.tree.leaves(loss_weights, is_leaf=eqx.is_inexact_array), dtype=float
+        )
+
+    def compute_softmax_weights(current, reference):
+        ratio_pytree = jax.tree.map(
+            lambda lt, ref: lt / (ref + epsilon),
+            current,
+            reference,
+        )
+        mu = jax.tree.reduce(jnp.maximum, ratio_pytree, initializer=-jnp.inf)
+        ratio_pytree = jax.tree.map(lambda r: r - mu, ratio_pytree)
+        ratio_leaves = jax.tree.leaves(ratio_pytree)
+        return jax.nn.softmax(jnp.array(ratio_leaves))
+
+    def soft_adapt_prev(stored_loss_terms):
+        # ω_j(i-1)
+        prev_terms = jax.tree.map(lambda slt: slt[iteration_nb - 1], stored_loss_terms)
+        prev_prev_terms = jax.tree.map(
+            lambda slt: slt[iteration_nb - 2], stored_loss_terms
+        )
+        return compute_softmax_weights(prev_terms, prev_prev_terms)
+
+    def look_back(loss_terms, stored_loss_terms):
+        # ω̂_j^(i,0)
+        initial_terms = jax.tree.map(lambda slt: tau * slt[0], stored_loss_terms)
+        weights = compute_softmax_weights(loss_terms, initial_terms)
+        return n_loss * weights
+
+    def soft_adapt_current(loss_terms, stored_loss_terms):
+        # ω_j(i)
+        prev_terms = jax.tree.map(lambda slt: slt[iteration_nb - 1], stored_loss_terms)
+        return compute_softmax_weights(loss_terms, prev_terms)
+
+    # Bernoulli variable for random lookback
+    rho = jax.random.bernoulli(key, p).astype(float)
+
+    # Base case for first iteration
+    def first_iter_case(_):
+        return do_nothing(loss_weights, None)
+
+    # Case for iteration >= 1
+    def subsequent_iter_case(_):
+        # Compute historical weights
+        def hist_weights_case1(_):
+            return soft_adapt_current(loss_terms, stored_loss_terms)
+
+        def hist_weights_case2(_):
+            return rho * soft_adapt_prev(stored_loss_terms) + (1 - rho) * look_back(
+                loss_terms, stored_loss_terms
+            )
+
+        loss_weights_hist = jax.lax.cond(
+            iteration_nb < 2,
+            hist_weights_case1,
+            hist_weights_case2,
+            None,
+        )
+
+        # Compute and return final weights
+        adaptive_weights = soft_adapt_current(loss_terms, stored_loss_terms)
+        return decay_factor * loss_weights_hist + (1 - decay_factor) * adaptive_weights
+
+    return jax.lax.cond(
+        iteration_nb == 0,
+        first_iter_case,
+        subsequent_iter_case,
+        None,
+    )
+
+
+def lr_annealing(
+    loss_weights: AbstractLossWeights,
+    grad_terms: AnyLossComponents,
+    decay_factor: float = 0.9,  # 0.9 is the recommended value from the article
+    eps: float = 1e-6,
+) -> Array:
+    r"""
+    Implementation of the Learning rate annealing
+    Algorithm 1 in the paper UNDERSTANDING AND MITIGATING GRADIENT PATHOLOGIES IN PHYSICS-INFORMED NEURAL NETWORKS
+
+    (a) Compute $\hat{\lambda}_i$ by
+    $$
+        \hat{\lambda}_i = \frac{\max_{\theta}\{|\nabla_\theta \mathcal{L}_r (\theta_n)|\}}{mean(|\nabla_\theta \mathcal{L}_i (\theta_n)|)}, \quad i=1,\dots, M,
+    $$
+
+    (b) Update the weights $\lambda_i$ using a moving average of the form
+    $$
+        \lambda_i = (1-\alpha) \lambda_{i-1} + \alpha \hat{\lambda}_i, \quad i=1, \dots, M.
+    $$
+
+    Note that since None is not treated as a leaf by jax tree.util functions,
+    we naturally avoid None components from loss_terms, stored_loss_terms etc.!
+
+    """
+    assert hasattr(grad_terms, "dyn_loss")
+    dyn_loss_grads = getattr(grad_terms, "dyn_loss")
+    data_fit_grads = [
+        getattr(grad_terms, att) if hasattr(grad_terms, att) else None
+        for att in ["norm_loss", "boundary_loss", "observations", "initial_condition"]
+    ]
+
+    dyn_loss_grads_leaves = jax.tree.leaves(
+        dyn_loss_grads,
+        is_leaf=eqx.is_inexact_array,
+    )
+
+    max_dyn_loss_grads = jnp.max(
+        jnp.stack([jnp.max(jnp.abs(g)) for g in dyn_loss_grads_leaves])
+    )
+
+    mean_gradients = [
+        jnp.mean(jnp.stack([jnp.abs(jnp.mean(g)) for g in jax.tree.leaves(t)]))
+        for t in data_fit_grads
+        if t is not None and jax.tree.leaves(t)
+    ]
+
+    lambda_hat = max_dyn_loss_grads / (jnp.array(mean_gradients) + eps)
+    old_weights = jnp.array(
+        jax.tree.leaves(
+            loss_weights,
+        )
+    )
+
+    new_weights = (1 - decay_factor) * old_weights[1:] + decay_factor * lambda_hat
+    return jnp.hstack([old_weights[0], new_weights])

jinns/loss/_loss_weights.py

@@ -2,26 +2,82 @@
 Formalize the loss weights data structure
 """
 
-from jaxtyping import Array, Float
+from __future__ import annotations
+from dataclasses import fields
+
+from jaxtyping import Array
+import jax.numpy as jnp
 import equinox as eqx
 
 
-class LossWeightsODE(eqx.Module):
-    dyn_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
-    initial_condition: Array | Float = eqx.field(kw_only=True, default=0.0)
-    observations: Array | Float = eqx.field(kw_only=True, default=0.0)
+def lw_converter(x):
+    if x is None:
+        return x
+    else:
+        return jnp.asarray(x)
+
+
+class AbstractLossWeights(eqx.Module):
+    """
+    An abstract class, currently only useful for type hints
+
+    TODO in the future maybe loss weights could be subclasses of
+    XDEComponentsAbstract?
+    """
+
+    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 LossWeightsODE(AbstractLossWeights):
+    dyn_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    initial_condition: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    observations: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
 
 
-class LossWeightsPDEStatio(eqx.Module):
-    dyn_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
-    norm_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
-    boundary_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
-    observations: Array | Float = eqx.field(kw_only=True, default=0.0)
+class LossWeightsPDEStatio(AbstractLossWeights):
+    dyn_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    norm_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    boundary_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    observations: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
 
 
-class LossWeightsPDENonStatio(eqx.Module):
-    dyn_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
-    norm_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
-    boundary_loss: Array | Float = eqx.field(kw_only=True, default=0.0)
-    observations: Array | Float = eqx.field(kw_only=True, default=0.0)
-    initial_condition: Array | Float = eqx.field(kw_only=True, default=0.0)
+class LossWeightsPDENonStatio(AbstractLossWeights):
+    dyn_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    norm_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    boundary_loss: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    observations: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )
+    initial_condition: Array | float | None = eqx.field(
+        kw_only=True, default=None, converter=lw_converter
+    )

jinns/parameters/_params.py

@@ -10,6 +10,14 @@ from jaxtyping import Array, PyTree, Float
 T = TypeVar("T")  # the generic type for what is in the Params PyTree because we
 # have possibly Params of Arrays, boolean, ...
 
+### NOTE
+### We are taking derivatives with respect to Params eqx.Modules.
+### This has been shown to behave weirdly if some fields of eqx.Modules have
+### been set as `field(init=False)`, we then should never create such fields in
+### jinns' Params modules.
+### We currently have silenced the warning related to this (see jinns.__init__
+### see https://github.com/patrick-kidger/equinox/pull/1043/commits/f88e62ab809140334c2f987ed13eff0d80b8be13
+
 
 class Params(eqx.Module, Generic[T]):
     """