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

jinns/nn/_save_load.py

@@ -3,6 +3,7 @@ Implements save and load functions
 """
 
 from typing import Callable, Literal
+from dataclasses import fields
 import pickle
 import jax
 import equinox as eqx
@@ -130,8 +131,14 @@ def save_pinn(
         u = eqx.tree_at(lambda m: m.init_params, u, params)
     eqx.tree_serialise_leaves(filename + "-module.eqx", u)
 
+    # The class EqParams is malformed for pickling, hence we pickle it under
+    # its dictionary form
+    eq_params_as_dict = {
+        k.name: getattr(params.eq_params, k.name) for k in fields(params.eq_params)
+    }
+
     with open(filename + "-eq_params.pkl", "wb") as f:
-        pickle.dump(params.eq_params, f)
+        pickle.dump(eq_params_as_dict, f)
 
     kwargs_creation = kwargs_creation.copy()  # avoid side-effect that would be
     # very probably harmless anyway
@@ -187,9 +194,9 @@ def load_pinn(
     try:
         with open(filename + "-eq_params.pkl", "rb") as f:
             eq_params_reloaded = pickle.load(f)
-    except FileNotFoundError:
-        eq_params_reloaded = {}
-        print("No pickle file for equation parameters found!")
+    except FileNotFoundError as e:
+        raise e
+
     kwargs_reloaded["eqx_list"] = string_to_function(kwargs_reloaded["eqx_list"])
     if type_ == "pinn_mlp":
         # next line creates a shallow model, the jax arrays are just shapes and

jinns/nn/_spinn.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Union, Callable, Any, Literal, overload
+from typing import Union, Callable, Any, Literal
 from dataclasses import InitVar
 from jaxtyping import PyTree, Float, Array
 import jax
@@ -8,7 +8,6 @@ import equinox as eqx
 
 from jinns.parameters._params import Params
 from jinns.nn._abstract_pinn import AbstractPINN
-from jinns.nn._utils import _PyTree_to_Params
 
 
 class SPINN(AbstractPINN):
@@ -72,17 +71,6 @@ class SPINN(AbstractPINN):
             eqx_spinn_network, self.filter_spec
         )
 
-    @overload
-    @_PyTree_to_Params
-    def __call__(
-        self,
-        inputs: Float[Array, " input_dim"],
-        params: PyTree,
-        *args,
-        **kwargs,
-    ) -> Float[Array, " output_dim"]: ...
-
-    @_PyTree_to_Params
     def __call__(
         self,
         t_x: Float[Array, "  batch_size 1+dim"],
@@ -94,10 +82,7 @@ class SPINN(AbstractPINN):
         Note that that thanks to the decorator, params can also directly be the
         PyTree (SPINN, PINN_MLP, ...) that we get out of eqx.combine
         """
-        # try:
         spinn = eqx.combine(params.nn_params, self.static)
-        # except (KeyError, AttributeError, TypeError) as e:
-        #    spinn = eqx.combine(params, self.static)
         v_model = jax.vmap(spinn)
         res = v_model(t_x)  # type: ignore
 

jinns/nn/_spinn_mlp.py

@@ -8,7 +8,7 @@ from typing import Callable, Literal, Self, Union, Any, TypeGuard
 import jax
 import jax.numpy as jnp
 import equinox as eqx
-from jaxtyping import Key, Array, Float, PyTree
+from jaxtyping import PRNGKeyArray, Array, Float, PyTree
 
 from jinns.nn._mlp import MLP
 from jinns.nn._spinn import SPINN
@@ -20,7 +20,7 @@ class SMLP(eqx.Module):
 
     Parameters
     ----------
-    key : InitVar[Key]
+    key : InitVar[PRNGKeyArray]
         A jax random key for the layer initializations.
     d : int
         The number of dimensions to treat separately, including time `t` if
@@ -42,7 +42,7 @@ class SMLP(eqx.Module):
         )`.
     """
 
-    key: InitVar[Key] = eqx.field(kw_only=True)
+    key: InitVar[PRNGKeyArray] = eqx.field(kw_only=True)
     eqx_list: InitVar[tuple[tuple[Callable, int, int] | tuple[Callable], ...]] = (
         eqx.field(kw_only=True)
     )
@@ -74,7 +74,7 @@ class SPINN_MLP(SPINN):
     @classmethod
     def create(
         cls,
-        key: Key,
+        key: PRNGKeyArray,
         d: int,
         r: int,
         eqx_list: tuple[tuple[Callable, int, int] | tuple[Callable], ...],
@@ -93,7 +93,7 @@ class SPINN_MLP(SPINN):
 
         Parameters
         ----------
-        key : Key
+        key : PRNGKeyArray
             A JAX random key that will be used to initialize the network parameters
         d : int
             The number of dimensions to treat separately.

jinns/nn/_utils.py

@@ -1,38 +1,33 @@
-from typing import Any, ParamSpec, Callable, Concatenate
-from jaxtyping import PyTree, Array
-from jinns.parameters._params import Params
-
-
-P = ParamSpec("P")
-
-
-def _PyTree_to_Params(
-    call_fun: Callable[
-        Concatenate[Any, Any, PyTree | Params[Array], P],
-        Any,
-    ],
-) -> Callable[
-    Concatenate[Any, Any, PyTree | Params[Array], P],
-    Any,
-]:
-    """
-    Decorator to be used around __call__ functions of PINNs, SPINNs, etc. It
-    authorizes the __call__ with `params` being directly be the
-    PyTree (SPINN, PINN_MLP, ...) that we get out of `eqx.combine`
-
-    This generic approach enables to cleanly handle type hints, up to the small
-    effort required to understand type hints for decorators (ie ParamSpec).
-    """
-
-    def wrapper(
-        self: Any,
-        inputs: Any,
-        params: PyTree | Params[Array],
-        *args: P.args,
-        **kwargs: P.kwargs,
-    ):
-        if isinstance(params, PyTree) and not isinstance(params, Params):
-            params = Params(nn_params=params, eq_params={})
-        return call_fun(self, inputs, params, *args, **kwargs)
-
-    return wrapper
+# P = ParamSpec("P")
+#
+#
+# def _PyTree_to_Params(
+#    call_fun: Callable[
+#        Concatenate[Any, Any, PyTree | Params[Array], P],
+#        Any,
+#    ],
+# ) -> Callable[
+#    Concatenate[Any, Any, PyTree | Params[Array], P],
+#    Any,
+# ]:
+#    """
+#    Decorator to be used around __call__ functions of PINNs, SPINNs, etc. It
+#    authorizes the __call__ with `params` being directly be the
+#    PyTree (SPINN, PINN_MLP, ...) that we get out of `eqx.combine`
+#
+#    This generic approach enables to cleanly handle type hints, up to the small
+#    effort required to understand type hints for decorators (ie ParamSpec).
+#    """
+#
+#    def wrapper(
+#        self: Any,
+#        inputs: Any,
+#        params: PyTree | Params[Array],
+#        *args: P.args,
+#        **kwargs: P.kwargs,
+#    ):
+#        if isinstance(params, PyTree) and not isinstance(params, Params):
+#            params = Params(nn_params=params, eq_params={})
+#        return call_fun(self, inputs, params, *args, **kwargs)
+#
+#    return wrapper

jinns/parameters/__init__.py

@@ -1,4 +1,4 @@
-from ._params import Params
+from ._params import EqParams, Params, update_eq_params
 from ._derivative_keys import (
     DerivativeKeysODE,
     DerivativeKeysPDEStatio,
@@ -6,8 +6,10 @@ from ._derivative_keys import (
 )
 
 __all__ = [
+    "EqParams",
     "Params",
     "DerivativeKeysODE",
     "DerivativeKeysPDEStatio",
     "DerivativeKeysPDENonStatio",
+    "update_eq_params",
 ]

jinns/parameters/_derivative_keys.py

@@ -19,13 +19,10 @@ def _get_masked_parameters(
     """
     # start with a params object with True everywhere. We will update to False
     # for parameters wrt which we do want not to differentiate the loss
-    diff_params = jax.tree.map(
-        lambda x: True,
-        params,
-        is_leaf=lambda x: isinstance(x, eqx.Module)
-        and not isinstance(x, Params),  # do not travers nn_params, more
-        # granularity could be imagined here, in the future
-    )
+    diff_params = Params(
+        nn_params=True, eq_params=jax.tree.map(lambda _: True, params.eq_params)
+    )  # do not travers nn_params, more
+    # granularity could be imagined here, in the future
     if derivative_mask_str == "both":
         return diff_params
     if derivative_mask_str == "eq_params":
@@ -60,7 +57,7 @@ class DerivativeKeysODE(eqx.Module):
 
          1. For unspecified loss term, the default is to differentiate with
         respect to `"nn_params"` only.
-         2. No granularity inside `Params.nn_params` is currently supported.
+         2. No granularity inside `Params.nn_params` is currently supported. An easy way to do freeze part of a custom PINN module is to use `jax.lax.stop_gradient` as explained [here](https://docs.kidger.site/equinox/faq/#how-to-mark-arrays-as-non-trainable-like-pytorchs-buffers).
          3. Note that the main Params object of the problem is mandatory if initialization via `from_str()`.
 
     A typical specification is of the form:
@@ -95,38 +92,52 @@ class DerivativeKeysODE(eqx.Module):
         infer the content of `Params.eq_params`.
     """
 
-    dyn_loss: Params[bool] | None = eqx.field(kw_only=True, default=None)
-    observations: Params[bool] | None = eqx.field(kw_only=True, default=None)
-    initial_condition: Params[bool] | None = eqx.field(kw_only=True, default=None)
+    dyn_loss: Params[bool]
+    observations: Params[bool]
+    initial_condition: Params[bool]
 
-    params: InitVar[Params[Array] | None] = eqx.field(kw_only=True, default=None)
+    params: InitVar[Params[Array] | None]
 
-    def __post_init__(self, params: Params[Array] | None = None):
+    def __init__(
+        self,
+        *,
+        dyn_loss: Params[bool] | None = None,
+        observations: Params[bool] | None = None,
+        initial_condition: Params[bool] | None = None,
+        params: Params[Array] | None = None,
+    ):
+        super().__init__()
         if params is None and (
-            self.dyn_loss is None
-            or self.observations is None
-            or self.initial_condition is None
+            dyn_loss is None or observations is None or initial_condition is None
         ):
             raise ValueError(
                 "params cannot be None since at least one loss "
                 "term has an undefined derivative key Params PyTree"
             )
-        if self.dyn_loss is None:
+        if dyn_loss is None:
             if params is None:
                 raise ValueError("self.dyn_loss is None, hence params should be passed")
             self.dyn_loss = _get_masked_parameters("nn_params", params)
-        if self.observations is None:
+        else:
+            self.dyn_loss = dyn_loss
+
+        if observations is None:
             if params is None:
                 raise ValueError(
                     "self.observations is None, hence params should be passed"
                 )
             self.observations = _get_masked_parameters("nn_params", params)
-        if self.initial_condition is None:
+        else:
+            self.observations = observations
+
+        if initial_condition is None:
             if params is None:
                 raise ValueError(
                     "self.initial_condition is None, hence params should be passed"
                 )
             self.initial_condition = _get_masked_parameters("nn_params", params)
+        else:
+            self.initial_condition = initial_condition
 
     @classmethod
     def from_str(
@@ -216,36 +227,56 @@ class DerivativeKeysPDEStatio(eqx.Module):
         content of `Params.eq_params`.
     """
 
-    dyn_loss: Params[bool] | None = eqx.field(kw_only=True, default=None)
-    observations: Params[bool] | None = eqx.field(kw_only=True, default=None)
-    boundary_loss: Params[bool] | None = eqx.field(kw_only=True, default=None)
-    norm_loss: Params[bool] | None = eqx.field(kw_only=True, default=None)
+    dyn_loss: Params[bool] = eqx.field(kw_only=True, default=None)
+    observations: Params[bool] = eqx.field(kw_only=True, default=None)
+    boundary_loss: Params[bool] = eqx.field(kw_only=True, default=None)
+    norm_loss: Params[bool] = eqx.field(kw_only=True, default=None)
 
     params: InitVar[Params[Array] | None] = eqx.field(kw_only=True, default=None)
 
-    def __post_init__(self, params: Params[Array] | None = None):
-        if self.dyn_loss is None:
+    def __init__(
+        self,
+        *,
+        dyn_loss: Params[bool] | None = None,
+        observations: Params[bool] | None = None,
+        boundary_loss: Params[bool] | None = None,
+        norm_loss: Params[bool] | None = None,
+        params: Params[Array] | None = None,
+    ):
+        super().__init__()
+        if dyn_loss is None:
             if params is None:
                 raise ValueError("self.dyn_loss is None, hence params should be passed")
             self.dyn_loss = _get_masked_parameters("nn_params", params)
-        if self.observations is None:
+        else:
+            self.dyn_loss = dyn_loss
+
+        if observations is None:
             if params is None:
                 raise ValueError(
                     "self.observations is None, hence params should be passed"
                 )
             self.observations = _get_masked_parameters("nn_params", params)
-        if self.boundary_loss is None:
+        else:
+            self.observations = observations
+
+        if boundary_loss is None:
             if params is None:
                 raise ValueError(
                     "self.boundary_loss is None, hence params should be passed"
                 )
             self.boundary_loss = _get_masked_parameters("nn_params", params)
-        if self.norm_loss is None:
+        else:
+            self.boundary_loss = boundary_loss
+
+        if norm_loss is None:
             if params is None:
                 raise ValueError(
                     "self.norm_loss is None, hence params should be passed"
                 )
             self.norm_loss = _get_masked_parameters("nn_params", params)
+        else:
+            self.norm_loss = norm_loss
 
     @classmethod
     def from_str(
@@ -344,16 +375,33 @@ class DerivativeKeysPDENonStatio(DerivativeKeysPDEStatio):
         content of `Params.eq_params`.
     """
 
-    initial_condition: Params[bool] | None = eqx.field(kw_only=True, default=None)
-
-    def __post_init__(self, params: Params[Array] | None = None):
-        super().__post_init__(params=params)
-        if self.initial_condition is None:
+    initial_condition: Params[bool] = eqx.field(kw_only=True, default=None)
+
+    def __init__(
+        self,
+        *,
+        dyn_loss: Params[bool] | None = None,
+        observations: Params[bool] | None = None,
+        boundary_loss: Params[bool] | None = None,
+        norm_loss: Params[bool] | None = None,
+        initial_condition: Params[bool] | None = None,
+        params: Params[Array] | None = None,
+    ):
+        super().__init__(
+            dyn_loss=dyn_loss,
+            observations=observations,
+            boundary_loss=boundary_loss,
+            norm_loss=norm_loss,
+            params=params,
+        )
+        if initial_condition is None:
             if params is None:
                 raise ValueError(
                     "self.initial_condition is None, hence params should be passed"
                 )
             self.initial_condition = _get_masked_parameters("nn_params", params)
+        else:
+            self.initial_condition = initial_condition
 
     @classmethod
     def from_str(
@@ -432,7 +480,9 @@ class DerivativeKeysPDENonStatio(DerivativeKeysPDEStatio):
         )
 
 
-def _set_derivatives(params, derivative_keys):
+def _set_derivatives(
+    params: Params[Array], derivative_keys: Params[bool]
+) -> Params[Array]:
     """
     We construct an eqx.Module with the fields of derivative_keys, each field
     has a copy of the params with appropriate derivatives set
@@ -448,13 +498,21 @@ def _set_derivatives(params, derivative_keys):
         `Params(nn_params=True | False, eq_params={"alpha":True | False,
         "beta":True | False})`.
         """
-        return jax.tree.map(
-            lambda p, d: jax.lax.cond(d, lambda p: p, jax.lax.stop_gradient, p),
-            params_,
-            derivative_mask,
-            is_leaf=lambda x: isinstance(x, eqx.Module)
-            and not isinstance(x, Params),  # do not travers nn_params, more
-            # granularity could be imagined here, in the future
+
+        return Params(
+            nn_params=jax.lax.cond(
+                derivative_mask.nn_params,
+                lambda p: p,
+                jax.lax.stop_gradient,
+                params_.nn_params,
+            ),
+            eq_params=jax.tree.map(
+                lambda p, d: jax.lax.cond(d, lambda p: p, jax.lax.stop_gradient, p),
+                params_.eq_params,
+                derivative_mask.eq_params,
+            ),
         )
+        # NOTE that currently we do not travers nn_params, more
+        # granularity could be imagined here, in the future
 
     return _set_derivatives_(params, derivative_keys)

jinns/parameters/_params.py

@@ -2,14 +2,34 @@
 Formalize the data structure for the parameters
 """
 
+from dataclasses import fields
 from typing import Generic, TypeVar
-import jax
 import equinox as eqx
-from jaxtyping import Array, PyTree, Float
+from jaxtyping import Array, PyTree
+
+from jinns.utils._DictToModuleMeta import DictToModuleMeta
 
 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 EqParams(metaclass=DictToModuleMeta):
+    """
+    Note that this is exposed to the user for the particular case where the
+    user, during its work, wants to change the equation parameters. In this
+    case, the user must import EqParams and call `EqParams.clear()`
+    """
+
+    pass
+
 
 class Params(eqx.Module, Generic[T]):
     """
@@ -20,37 +40,47 @@ class Params(eqx.Module, Generic[T]):
     nn_params : PyTree[T]
         A PyTree of the non-static part of the PINN eqx.Module, i.e., the
         parameters of the PINN
-    eq_params : dict[str, T]
-        A dictionary of the equation parameters. Keys are the parameter name,
-        values are their corresponding value
+    eq_params : PyTree[T]
+        A PyTree of the equation parameters. For retrocompatibility it us
+        provided as a dictionary of the equation parameters where keys are the parameter names, and values are their corresponding values. Internally,
+        it will be transformed to a custom instance of `EqParams`.
     """
 
-    nn_params: PyTree[T] = eqx.field(kw_only=True, default=None)
-    eq_params: dict[str, T] = eqx.field(kw_only=True, default=None)
+    nn_params: PyTree[T]
+    eq_params: PyTree[T]
+
+    def __init__(
+        self,
+        nn_params: PyTree[T] | None = None,
+        eq_params: dict[str, T] | None = None,
+    ):
+        self.nn_params = nn_params
+        if isinstance(eq_params, dict):
+            self.eq_params = EqParams(eq_params, "EqParams")
+        else:
+            self.eq_params = eq_params
 
 
-def _update_eq_params_dict(
+def update_eq_params(
     params: Params[Array],
-    param_batch_dict: dict[str, Float[Array, " param_batch_size dim"]],
-) -> Params:
+    eq_param_batch: PyTree[Array] | None,
+) -> Params[Array]:
     """
     Update params.eq_params with a batch of eq_params for given key(s)
     """
 
-    # artificially "complete" `param_batch_dict` with None to match `params`
-    # PyTree  structure
-    param_batch_dict_ = param_batch_dict | {
-        k: None for k in set(params.eq_params.keys()) - set(param_batch_dict.keys())
-    }
+    if eq_param_batch is None:
+        return params
 
-    # Replace at non None leafs
+    param_names_to_update = tuple(f.name for f in fields(eq_param_batch))
     params = eqx.tree_at(
         lambda p: p.eq_params,
         params,
-        jax.tree_util.tree_map(
-            lambda p, q: q if q is not None else p,
+        eqx.tree_at(
+            lambda pt: tuple(getattr(pt, f) for f in param_names_to_update),
             params.eq_params,
-            param_batch_dict_,
+            tuple(getattr(eq_param_batch, f) for f in param_names_to_update),
+            is_leaf=lambda x: x is None or eqx.is_inexact_array(x),
         ),
     )
 
@@ -58,7 +88,7 @@ def _update_eq_params_dict(
 
 
 def _get_vmap_in_axes_params(
-    eq_params_batch_dict: dict[str, Array], params: Params[Array]
+    eq_param_batch: eqx.Module | None, params: Params[Array]
 ) -> tuple[Params[int | None] | None]:
     """
     Return the input vmap axes when there is batch(es) of parameters to vmap
@@ -69,19 +99,22 @@ def _get_vmap_in_axes_params(
     Note that we return a Params PyTree with an integer to designate the
     vmapped axis or None if there is not
     """
-    if eq_params_batch_dict is None:
+    if eq_param_batch is None:
         return (None,)
     # We use pytree indexing of vmapped axes and vmap on axis
     # 0 of the eq_parameters for which we have a batch
     # this is for a fine-grained vmaping
     # scheme over the params
+    param_names_to_vmap = tuple(f.name for f in fields(eq_param_batch))
+    vmap_axes_dict = {
+        k.name: (0 if k.name in param_names_to_vmap else None)
+        for k in fields(params.eq_params)
+    }
+    eq_param_vmap_axes = type(params.eq_params)(**vmap_axes_dict)
     vmap_in_axes_params = (
         Params(
             nn_params=None,
-            eq_params={
-                k: (0 if k in eq_params_batch_dict.keys() else None)
-                for k in params.eq_params.keys()
-            },
+            eq_params=eq_param_vmap_axes,
         ),
     )
     return vmap_in_axes_params

jinns/solver/_solve.py

@@ -14,7 +14,7 @@ import optax
 import jax
 from jax import jit
 import jax.numpy as jnp
-from jaxtyping import Float, Array, PyTree, Key
+from jaxtyping import Float, Array, PyTree, PRNGKeyArray
 import equinox as eqx
 from jinns.solver._rar import init_rar, trigger_rar
 from jinns.utils._utils import _check_nan_in_pytree
@@ -47,7 +47,7 @@ if TYPE_CHECKING:
         LossContainer,
         StoredObjectContainer,
         Float[Array, " n_iter"] | None,
-        Key | None,
+        PRNGKeyArray | None,
     ]
 
 
@@ -66,7 +66,7 @@ def solve(
     obs_batch_sharding: jax.sharding.Sharding | None = None,
     verbose: bool = True,
     ahead_of_time: bool = True,
-    key: Key = None,
+    key: PRNGKeyArray | None = None,
 ) -> tuple[
     Params[Array],
     Float[Array, " n_iter"],
@@ -179,6 +179,7 @@ def solve(
     best_val_params
         The best parameters according to the validation criterion
     """
+    initialization_time = time.time()
     if n_iter < 1:
         raise ValueError("Cannot run jinns.solve for n_iter<1")
 
@@ -225,11 +226,6 @@ def solve(
     # get_batch with device_put, the latter is not jittable
     get_batch = _get_get_batch(obs_batch_sharding)
 
-    # initialize the dict for stored parameter values
-    # we need to get a loss_term to init stuff
-    batch_ini, data, param_data, obs_data = get_batch(data, param_data, obs_data)
-    _, loss_terms = loss(init_params, batch_ini)
-
     # initialize parameter tracking
     if tracked_params is None:
         tracked_params = jax.tree.map(lambda p: None, init_params)
@@ -247,6 +243,13 @@ def solve(
         # being a complex data structure
     )
 
+    # initialize the dict for stored parameter values
+    # we need to get a loss_term to init stuff
+    # NOTE: we use jax.eval_shape to avoid FLOPS since we only need the tree
+    # structure
+    batch_ini, data, param_data, obs_data = get_batch(data, param_data, obs_data)
+    _, loss_terms = jax.eval_shape(loss, init_params, batch_ini)
+
     # initialize the PyTree for stored loss values
     stored_loss_terms = jax.tree_util.tree_map(
         lambda _: jnp.zeros((n_iter)), loss_terms
@@ -475,6 +478,9 @@ def solve(
             key,
         )
 
+    if verbose:
+        print("Initialization time:", time.time() - initialization_time)
+
     # Main optimization loop. We use the LAX while loop (fully jitted) version
     # if no mixing devices. Otherwise we use the standard while loop. Here devices only
     # concern obs_batch, but it could lead to more complex scheme in the future