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

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,10 +2,12 @@
 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, ...
@@ -19,6 +21,16 @@ T = TypeVar("T")  # the generic type for what is in the Params PyTree because we
 ### 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]):
     """
     The equinox module for the parameters
@@ -28,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),
         ),
     )
 
@@ -66,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
@@ -77,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"],

jinns/utils/_DictToModuleMeta.py

@@ -0,0 +1,66 @@
+from typing import Any
+import equinox as eqx
+
+
+class DictToModuleMeta(type):
+    """
+    A Metaclass based solution to handle the fact that we only
+    want one type to be created for EqParams.
+    If we were to create a new **class type** (despite same name) each time we
+    create a new Params object, nothing would be broadcastable in terms of jax
+    tree utils operations and this would be useless. The difficulty comes from
+    the fact that we need to instanciate from this same class at different
+    moments of the jinns workflow eg: parameter creation, derivative keys
+    creations, tracked parameter designation, etc. (ie. each time a Params
+    class is instanciated whatever its usage, we need the same EqParams class
+    to be instanciated)
+
+    This is inspired by the Singleton pattern in Python
+    (https://stackoverflow.com/a/10362179)
+
+    Here we need the call of a metaclass because as explained in
+     https://stackoverflow.com/a/45536640). To quote from the answer
+    Metaclasses implement how the class will behave (not the instance). So when you look at the instance creation:
+    `x = Foo()`
+    This literally "calls" the class Foo. That's why __call__ of the metaclass
+    is invoked before the __new__ and
+    __init__ methods of your class initialize the instance.
+    Other viewpoint: Metaclasses,as well as classes making use of those
+    metaclasses, are created when the lines of code containing
+    the class statement body is executed
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(DictToModuleMeta, self).__init__(*args, **kwargs)
+        self._class = None
+
+    def __call__(self, d: dict[str, Any], class_name: str | None = None) -> eqx.Module:
+        """
+        Notably, once the class template is registered (after the first call to
+        EqParams()), all calls with different keys in `d` will fail.
+        """
+        if self._class is None and class_name is not None:
+            self._class = type(
+                class_name,
+                (eqx.Module,),
+                {"__annotations__": {k: type(v) for k, v in d.items()}},
+            )
+        try:
+            return self._class(**d)  # type: ignore
+        except TypeError as _:
+            print(
+                "DictToModuleMeta has been created with the fields"
+                f"{tuple(k for k in self._class.__annotations__.keys())}"
+                f" but an instanciation is resquested with fields={tuple(k for k in d.keys())}"
+                " which results in an error"
+            )
+            raise ValueError
+
+    def clear(cls) -> None:
+        """
+        The current Metaclass implementation freezes the list of equation parameters inside a Python session;
+        only one EqParams annotation can exist at a given time. Use `EqParams.clear()`  to reset.
+        Also useful for pytest where stuff is not complety reset after tests
+        Taken from https://stackoverflow.com/a/50065732
+        """
+        cls._class = None

jinns/utils/_ItemizableModule.py

@@ -0,0 +1,19 @@
+from dataclasses import fields
+from typing import Any, ItemsView
+import equinox as eqx
+
+
+class ItemizableModule(eqx.Module):
+    def items(self) -> ItemsView[str, Any]:
+        """
+        For the dataclass to be iterated like a dictionary.
+        Practical and retrocompatible with old code when loss components were
+        dictionaries
+
+        About the type hint: https://stackoverflow.com/questions/73022688/type-annotation-for-dict-items
+        """
+        return {
+            field.name: getattr(self, field.name)
+            for field in fields(self)
+            if getattr(self, field.name) is not None
+        }.items()

jinns/utils/__init__.py

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