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

jinns/solver/_solve.py

@@ -8,56 +8,76 @@ from __future__ import (
 )  # https://docs.python.org/3/library/typing.html#constant
 
 import time
-from typing import TYPE_CHECKING, NamedTuple, Dict, Union
+from typing import TYPE_CHECKING, Any, TypeAlias, Callable
 from functools import partial
 import optax
 import jax
 from jax import jit
 import jax.numpy as jnp
-from jaxtyping import Int, Bool, Float, Array
+from jaxtyping import Float, Array, PyTree, Key
+import equinox as eqx
 from jinns.solver._rar import init_rar, trigger_rar
 from jinns.utils._utils import _check_nan_in_pytree
 from jinns.solver._utils import _check_batch_size
-from jinns.utils._containers import *
-from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-    append_obs_batch,
-    append_param_batch,
+from jinns.utils._containers import (
+    DataGeneratorContainer,
+    OptimizationContainer,
+    OptimizationExtraContainer,
+    LossContainer,
+    StoredObjectContainer,
 )
+from jinns.data._utils import append_param_batch, append_obs_batch
 
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.parameters._params import Params
+    from jinns.utils._types import AnyBatch
+    from jinns.loss._abstract_loss import AbstractLoss
+    from jinns.validation._validation import AbstractValidationModule
+    from jinns.data._DataGeneratorParameter import DataGeneratorParameter
+    from jinns.data._DataGeneratorObservations import DataGeneratorObservations
+    from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+
+    main_carry: TypeAlias = tuple[
+        int,
+        AbstractLoss,
+        OptimizationContainer,
+        OptimizationExtraContainer,
+        DataGeneratorContainer,
+        AbstractValidationModule | None,
+        LossContainer,
+        StoredObjectContainer,
+        Float[Array, " n_iter"] | None,
+        Key | None,
+    ]
 
 
 def solve(
-    n_iter: Int,
-    init_params: AnyParams,
-    data: AnyDataGenerator,
-    loss: AnyLoss,
+    n_iter: int,
+    init_params: Params[Array],
+    data: AbstractDataGenerator,
+    loss: AbstractLoss,
     optimizer: optax.GradientTransformation,
-    print_loss_every: Int = 1000,
-    opt_state: Union[NamedTuple, None] = None,
-    tracked_params: Params | ParamsDict | None = None,
+    print_loss_every: int = 1000,
+    opt_state: optax.OptState | None = None,
+    tracked_params: Params[Any | None] | None = None,
     param_data: DataGeneratorParameter | None = None,
-    obs_data: (
-        DataGeneratorObservations | DataGeneratorObservationsMultiPINNs | None
-    ) = None,
+    obs_data: DataGeneratorObservations | None = None,
     validation: AbstractValidationModule | None = None,
     obs_batch_sharding: jax.sharding.Sharding | None = None,
-    verbose: Bool = True,
-    ahead_of_time: Bool = True,
+    verbose: bool = True,
+    ahead_of_time: bool = True,
+    key: Key = None,
 ) -> tuple[
-    Params | ParamsDict,
-    Float[Array, "n_iter"],
-    Dict[str, Float[Array, "n_iter"]],
-    AnyDataGenerator,
-    AnyLoss,
-    NamedTuple,
-    AnyParams,
-    Float[Array, "n_iter"],
-    AnyParams,
+    Params[Array],
+    Float[Array, " n_iter"],
+    PyTree,
+    AbstractDataGenerator,
+    AbstractLoss,
+    optax.OptState,
+    Params[Array | None],
+    PyTree,
+    Float[Array, " n_iter"] | None,
+    Params[Array],
 ]:
     """
     Performs the optimization process via stochastic gradient descent
@@ -94,8 +114,7 @@ def solve(
         Default None. A DataGeneratorParameter object which can be used to
         sample equation parameters.
     obs_data
-        Default None. A DataGeneratorObservations or
-        DataGeneratorObservationsMultiPINNs
+        Default None. A DataGeneratorObservations
         object which can be used to sample minibatches of observations.
     validation
         Default None. Otherwise, a callable ``eqx.Module`` which implements a
@@ -127,6 +146,9 @@ def solve(
         transformed (see https://jax.readthedocs.io/en/latest/aot.html#aot-compiled-functions-cannot-be-transformed).
         When False, jinns does not provide any timing information (which would
         be nonsense in a JIT transformed `solve()` function).
+    key
+        Default None. A JAX random key that can be used for diverse purpose in
+        the main iteration loop.
 
     Returns
     -------
@@ -136,8 +158,8 @@ def solve(
     total_loss_values
         An array of the total loss term along the gradient steps
     stored_loss_terms
-        A dictionary. At each key an array of the values of a given loss
-        term is stored
+        A PyTree with attributes being arrays of all the values for each loss
+        term
     data
         The input data object
     loss
@@ -147,11 +169,19 @@ def solve(
     stored_params
         A Params objects with the stored values of the desired parameters (as
         signified in tracked_params argument)
+    stored_weights_terms
+        A PyTree with attributes being arrays of all the values for each loss
+        weight. Note that if Loss.update_weight_method is None, we return None,
+        because loss weights are never updated and we can then save some
+        computations
     validation_crit_values
         An array containing the validation criterion values of the training
     best_val_params
         The best parameters according to the validation criterion
     """
+    if n_iter < 1:
+        raise ValueError("Cannot run jinns.solve for n_iter<1")
+
     if param_data is not None:
         if param_data.param_batch_size is not None:
             # We need to check that batch sizes will all be compliant for
@@ -171,11 +201,21 @@ def solve(
             _check_batch_size(obs_data, param_data, "n")
 
     if opt_state is None:
-        opt_state = optimizer.init(init_params)
+        opt_state = optimizer.init(init_params)  # type: ignore
+        # our Params are eqx.Module (dataclass + PyTree), PyTree is
+        # compatible with optax transform but not dataclass, this leads to a
+        # type hint error: we could prevent this by ensuring with the eqx.filter that
+        # we have only floating points optimizable params given to optax
+        # see https://docs.kidger.site/equinox/faq/#optax-throwing-a-typeerror
+        # opt_state = optimizer.init(
+        #    eqx.filter(init_params, eqx.is_inexact_array)
+        # )
+        # but this seems like a hack and there is no better way
+        # https://github.com/google-deepmind/optax/issues/384
 
     # RAR sampling init (ouside scanned function to avoid dynamic slice error)
     # If RAR is not used the _rar_step_*() are juste None and data is unchanged
-    data, _rar_step_true, _rar_step_false = init_rar(data)
+    data, _rar_step_true, _rar_step_false = init_rar(data)  # type: ignore
 
     # Seq2seq
     curr_seq = 0
@@ -207,11 +247,38 @@ def solve(
         # being a complex data structure
     )
 
-    # initialize the dict for stored loss values
+    # initialize the PyTree for stored loss values
     stored_loss_terms = jax.tree_util.tree_map(
         lambda _: jnp.zeros((n_iter)), loss_terms
     )
 
+    # initialize the PyTree for stored loss weights values
+    if loss.update_weight_method is not None:
+        stored_weights_terms = eqx.tree_at(
+            lambda pt: jax.tree.leaves(
+                pt, is_leaf=lambda x: x is not None and eqx.is_inexact_array(x)
+            ),
+            loss.loss_weights,
+            tuple(
+                jnp.zeros((n_iter))
+                for n in range(
+                    len(
+                        jax.tree.leaves(
+                            loss.loss_weights,
+                            is_leaf=lambda x: x is not None and eqx.is_inexact_array(x),
+                        )
+                    )
+                )
+            ),
+        )
+    else:
+        stored_weights_terms = None
+    if loss.update_weight_method is not None and key is None:
+        raise ValueError(
+            "`key` argument must be passed to jinns.solve when"
+            " `loss.update_weight_method` is not None"
+        )
+
     train_data = DataGeneratorContainer(
         data=data, param_data=param_data, obs_data=obs_data
     )
@@ -228,6 +295,7 @@ def solve(
     )
     loss_container = LossContainer(
         stored_loss_terms=stored_loss_terms,
+        stored_weights_terms=stored_weights_terms,
         train_loss_values=train_loss_values,
     )
     stored_objects = StoredObjectContainer(
@@ -252,6 +320,7 @@ def solve(
         loss_container,
         stored_objects,
         validation_crit_values,
+        key,
     )
 
     def _one_iteration(carry: main_carry) -> main_carry:
@@ -265,24 +334,47 @@ def solve(
             loss_container,
             stored_objects,
             validation_crit_values,
+            key,
         ) = carry
 
         batch, data, param_data, obs_data = get_batch(
             train_data.data, train_data.param_data, train_data.obs_data
         )
 
-        # Gradient step
+        # ---------------------------------------------------------------------
+        # The following part is the equivalent of a
+        # > train_loss_value, grads = jax.values_and_grad(total_loss.evaluate)(params, ...)
+        # but it is decomposed on individual loss terms so that we can use it
+        # if needed for updating loss weights.
+        # Since the total loss is a weighted sum of individual loss terms, so
+        # are its total gradients.
+
+        # Compute individual losses and individual gradients
+        loss_terms, grad_terms = loss.evaluate_by_terms(optimization.params, batch)
+
+        if loss.update_weight_method is not None:
+            key, subkey = jax.random.split(key)  # type: ignore because key can
+            # still be None currently
+            # avoid computations of tree_at if no updates
+            loss = loss.update_weights(
+                i, loss_terms, loss_container.stored_loss_terms, grad_terms, subkey
+            )
+
+        # total grad
+        grads = loss.ponderate_and_sum_gradient(grad_terms)
+
+        # total loss
+        train_loss_value = loss.ponderate_and_sum_loss(loss_terms)
+        # ---------------------------------------------------------------------
+
+        # gradient step
         (
-            loss,
-            train_loss_value,
-            loss_terms,
             params,
             opt_state,
             last_non_nan_params,
         ) = _gradient_step(
-            loss,
+            grads,
             optimizer,
-            batch,
             optimization.params,
             optimization.opt_state,
             optimization.last_non_nan_params,
@@ -292,7 +384,7 @@ def solve(
         if verbose:
             _print_fn(i, train_loss_value, print_loss_every, prefix="[train] ")
 
-        if validation is not None:
+        if validation is not None and validation_crit_values is not None:
             # there is a jax.lax.cond because we do not necesarily call the
             # validation step every iteration
             (
@@ -306,7 +398,7 @@ def solve(
                 lambda operands: (
                     operands[0],
                     False,
-                    validation_crit_values[i - 1],
+                    validation_crit_values[i - 1],  # type: ignore don't know why it can still be None
                     False,
                 ),
                 (
@@ -350,14 +442,14 @@ def solve(
         )
 
         # save loss value and selected parameters
-        stored_params, stored_loss_terms, train_loss_values = _store_loss_and_params(
+        stored_objects, loss_container = _store_loss_and_params(
             i,
             params,
             stored_objects.stored_params,
-            loss_container.stored_loss_terms,
-            loss_container.train_loss_values,
+            loss_container,
             train_loss_value,
             loss_terms,
+            loss.loss_weights,
             tracked_params,
         )
 
@@ -377,9 +469,10 @@ def solve(
             ),
             DataGeneratorContainer(data, param_data, obs_data),
             validation,
-            LossContainer(stored_loss_terms, train_loss_values),
-            StoredObjectContainer(stored_params),
+            loss_container,
+            stored_objects,
             validation_crit_values,
+            key,
         )
 
     # Main optimization loop. We use the LAX while loop (fully jitted) version
@@ -419,6 +512,7 @@ def solve(
         loss_container,
         stored_objects,
         validation_crit_values,
+        key,
     ) = carry
 
     if verbose:
@@ -431,7 +525,7 @@ def solve(
     # get ready to return the parameters at last iteration...
     # (by default arbitrary choice, this could be None)
     validation_parameters = optimization.last_non_nan_params
-    if validation is not None:
+    if validation is not None and validation_crit_values is not None:
         jax.debug.print(
             "validation loss value = {validation_loss_val}",
             validation_loss_val=validation_crit_values[i - 1],
@@ -456,6 +550,7 @@ def solve(
         loss,  # return the Loss if needed (no-inplace modif)
         optimization.opt_state,
         stored_objects.stored_params,
+        loss_container.stored_weights_terms,
         validation_crit_values if validation is not None else None,
         validation_parameters,
     )
@@ -463,27 +558,26 @@ def solve(
 
 @partial(jit, static_argnames=["optimizer"])
 def _gradient_step(
-    loss: AnyLoss,
+    grads: Params[Array],
     optimizer: optax.GradientTransformation,
-    batch: AnyBatch,
-    params: AnyParams,
-    opt_state: NamedTuple,
-    last_non_nan_params: AnyParams,
+    params: Params[Array],
+    opt_state: optax.OptState,
+    last_non_nan_params: Params[Array],
 ) -> tuple[
-    AnyLoss,
-    float,
-    Dict[str, float],
-    AnyParams,
-    NamedTuple,
-    AnyParams,
+    Params[Array],
+    optax.OptState,
+    Params[Array],
 ]:
     """
     optimizer cannot be jit-ted.
     """
-    value_grad_loss = jax.value_and_grad(loss, has_aux=True)
-    (loss_val, loss_terms), grads = value_grad_loss(params, batch)
-    updates, opt_state = optimizer.update(grads, opt_state, params)
-    params = optax.apply_updates(params, updates)
+
+    updates, opt_state = optimizer.update(
+        grads,  # type: ignore
+        opt_state,
+        params,  # type: ignore
+    )  # see optimizer.init for explaination for the ignore(s) here
+    params = optax.apply_updates(params, updates)  # type: ignore
 
     # check if any of the parameters is NaN
     last_non_nan_params = jax.lax.cond(
@@ -494,9 +588,6 @@ def _gradient_step(
     )
 
     return (
-        loss,
-        loss_val,
-        loss_terms,
         params,
         opt_state,
         last_non_nan_params,
@@ -504,7 +595,7 @@ def _gradient_step(
 
 
 @partial(jit, static_argnames=["prefix"])
-def _print_fn(i: Int, loss_val: Float, print_loss_every: Int, prefix: str = ""):
+def _print_fn(i: int, loss_val: Float, print_loss_every: int, prefix: str = ""):
     # note that if the following is not jitted in the main lor loop, it is
     # super slow
     _ = jax.lax.cond(
@@ -521,17 +612,15 @@ def _print_fn(i: Int, loss_val: Float, print_loss_every: Int, prefix: str = ""):
 
 @jit
 def _store_loss_and_params(
-    i: Int,
-    params: AnyParams,
-    stored_params: AnyParams,
-    stored_loss_terms: Dict[str, Float[Array, "n_iter"]],
-    train_loss_values: Float[Array, "n_iter"],
+    i: int,
+    params: Params[Array],
+    stored_params: Params[Array | None],
+    loss_container: LossContainer,
     train_loss_val: float,
-    loss_terms: Dict[str, float],
-    tracked_params: AnyParams,
-) -> tuple[
-    Params | ParamsDict, Dict[str, Float[Array, "n_iter"]], Float[Array, "n_iter"]
-]:
+    loss_terms: PyTree[Array],
+    weight_terms: PyTree[Array],
+    tracked_params: Params,
+) -> tuple[StoredObjectContainer, LossContainer]:
     stored_params = jax.tree_util.tree_map(
         lambda stored_value, param, tracked_param: (
             None
@@ -550,15 +639,41 @@ def _store_loss_and_params(
     )
     stored_loss_terms = jax.tree_util.tree_map(
         lambda stored_term, loss_term: stored_term.at[i].set(loss_term),
-        stored_loss_terms,
+        loss_container.stored_loss_terms,
         loss_terms,
     )
 
-    train_loss_values = train_loss_values.at[i].set(train_loss_val)
-    return (stored_params, stored_loss_terms, train_loss_values)
+    if loss_container.stored_weights_terms is not None:
+        stored_weights_terms = jax.tree_util.tree_map(
+            lambda stored_term, weight_term: stored_term.at[i].set(weight_term),
+            jax.tree.leaves(
+                loss_container.stored_weights_terms,
+                is_leaf=lambda x: x is not None and eqx.is_inexact_array(x),
+            ),
+            jax.tree.leaves(
+                weight_terms,
+                is_leaf=lambda x: x is not None and eqx.is_inexact_array(x),
+            ),
+        )
+        stored_weights_terms = eqx.tree_at(
+            lambda pt: jax.tree.leaves(
+                pt, is_leaf=lambda x: x is not None and eqx.is_inexact_array(x)
+            ),
+            loss_container.stored_weights_terms,
+            stored_weights_terms,
+        )
+    else:
+        stored_weights_terms = None
+
+    train_loss_values = loss_container.train_loss_values.at[i].set(train_loss_val)
+    loss_container = LossContainer(
+        stored_loss_terms, stored_weights_terms, train_loss_values
+    )
+    stored_objects = StoredObjectContainer(stored_params)
+    return stored_objects, loss_container
 
 
-def _get_break_fun(n_iter: Int, verbose: Bool) -> Callable[[main_carry], Bool]:
+def _get_break_fun(n_iter: int, verbose: bool) -> Callable[[main_carry], bool]:
     """
     Wrapper to get the break_fun with appropriate `n_iter`.
     The verbose argument is here to control printing (or not) when exiting
@@ -586,7 +701,7 @@ def _get_break_fun(n_iter: Int, verbose: Bool) -> Callable[[main_carry], Bool]:
         def continue_while_loop(_):
             return True
 
-        (i, _, optimization, optimization_extra, _, _, _, _, _) = carry
+        (i, _, optimization, optimization_extra, _, _, _, _, _, _) = carry
 
         # Condition 1
         bool_max_iter = jax.lax.cond(
@@ -599,7 +714,7 @@ def _get_break_fun(n_iter: Int, verbose: Bool) -> Callable[[main_carry], Bool]:
         bool_nan_in_params = jax.lax.cond(
             _check_nan_in_pytree(optimization.params),
             lambda _: stop_while_loop(
-                "NaN values in parameters " "(returning last non NaN values)"
+                "NaN values in parameters (returning last non NaN values)"
             ),
             continue_while_loop,
             None,
@@ -622,18 +737,18 @@ def _get_break_fun(n_iter: Int, verbose: Bool) -> Callable[[main_carry], Bool]:
 
 
 def _get_get_batch(
-    obs_batch_sharding: jax.sharding.Sharding,
+    obs_batch_sharding: jax.sharding.Sharding | None,
 ) -> Callable[
     [
-        AnyDataGenerator,
+        AbstractDataGenerator,
         DataGeneratorParameter | None,
-        DataGeneratorObservations | DataGeneratorObservationsMultiPINNs | None,
+        DataGeneratorObservations | None,
     ],
     tuple[
         AnyBatch,
-        AnyDataGenerator,
+        AbstractDataGenerator,
         DataGeneratorParameter | None,
-        DataGeneratorObservations | DataGeneratorObservationsMultiPINNs | None,
+        DataGeneratorObservations | None,
     ],
 ]:
     """

jinns/solver/_utils.py

@@ -1,9 +1,7 @@
-from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    CubicMeshPDEStatio,
-    CubicMeshPDENonStatio,
-    DataGeneratorParameter,
-)
+from jinns.data._DataGeneratorODE import DataGeneratorODE
+from jinns.data._CubicMeshPDEStatio import CubicMeshPDEStatio
+from jinns.data._CubicMeshPDENonStatio import CubicMeshPDENonStatio
+from jinns.data._DataGeneratorParameter import DataGeneratorParameter
 
 
 def _check_batch_size(other_data, main_data, attr_name):

jinns/utils/__init__.py

@@ -1 +1,3 @@
 from ._utils import get_grid
+
+__all__ = ["get_grid"]

jinns/utils/_containers.py

@@ -6,28 +6,31 @@ from __future__ import (
     annotations,
 )  # https://docs.python.org/3/library/typing.html#constant
 
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING
 from jaxtyping import PyTree, Array, Float, Bool
 from optax import OptState
 import equinox as eqx
 
+from jinns.parameters._params import Params
+
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+    from jinns.data._DataGeneratorParameter import DataGeneratorParameter
+    from jinns.data._DataGeneratorObservations import DataGeneratorObservations
+    from jinns.utils._types import AnyLoss
 
 
 class DataGeneratorContainer(eqx.Module):
-    data: AnyDataGenerator
+    data: AbstractDataGenerator
     param_data: DataGeneratorParameter | None = None
-    obs_data: DataGeneratorObservations | DataGeneratorObservationsMultiPINNs | None = (
-        None
-    )
+    obs_data: DataGeneratorObservations | None = None
 
 
 class ValidationContainer(eqx.Module):
     loss: AnyLoss | None
     data: DataGeneratorContainer
     hyperparams: PyTree = None
-    loss_values: Float[Array, "n_iter"] | None = None
+    loss_values: Float[Array, " n_iter"] | None = None
 
 
 class OptimizationContainer(eqx.Module):
@@ -45,9 +48,12 @@ class OptimizationExtraContainer(eqx.Module):
 
 
 class LossContainer(eqx.Module):
-    stored_loss_terms: Dict[str, Float[Array, "n_iter"]]
-    train_loss_values: Float[Array, "n_iter"]
+    # PyTree below refers to ODEComponents or PDEStatioComponents or
+    # PDENonStatioComponents
+    stored_loss_terms: PyTree[Float[Array, " n_iter"]]
+    stored_weights_terms: PyTree[Float[Array, " n_iter"]]
+    train_loss_values: Float[Array, " n_iter"]
 
 
 class StoredObjectContainer(eqx.Module):
-    stored_params: list | None
+    stored_params: Params[Array | None]