jinns 1.0.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

jinns/solver/_solve.py

@@ -7,6 +7,7 @@ from __future__ import (
     annotations,
 )  # https://docs.python.org/3/library/typing.html#constant
 
+import time
 from typing import TYPE_CHECKING, NamedTuple, Dict, Union
 from functools import partial
 import optax
@@ -16,6 +17,7 @@ import jax.numpy as jnp
 from jaxtyping import Int, Bool, Float, Array
 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,
@@ -29,31 +31,6 @@ if TYPE_CHECKING:
     from jinns.utils._types import *
 
 
-def _check_batch_size(other_data, main_data, attr_name):
-    if (
-        (
-            isinstance(main_data, DataGeneratorODE)
-            and getattr(other_data, attr_name) != main_data.temporal_batch_size
-        )
-        or (
-            isinstance(main_data, CubicMeshPDEStatio)
-            and not isinstance(main_data, CubicMeshPDENonStatio)
-            and getattr(other_data, attr_name) != main_data.omega_batch_size
-        )
-        or (
-            isinstance(main_data, CubicMeshPDENonStatio)
-            and getattr(other_data, attr_name)
-            != main_data.omega_batch_size * main_data.temporal_batch_size
-        )
-    ):
-        raise ValueError(
-            "Optional other_data.param_batch_size must be"
-            " equal to main_data.temporal_batch_size or main_data.omega_batch_size or"
-            " the product of both dependeing on the type of the main"
-            " datagenerator"
-        )
-
-
 def solve(
     n_iter: Int,
     init_params: AnyParams,
@@ -167,10 +144,22 @@ def solve(
         The best parameters according to the validation criterion
     """
     if param_data is not None:
-        _check_batch_size(param_data, data, "param_batch_size")
-
-    if obs_data is not None:
-        _check_batch_size(obs_data, data, "obs_batch_size")
+        if param_data.param_batch_size is not None:
+            # We need to check that batch sizes will all be compliant for
+            # correct vectorization
+            _check_batch_size(param_data, data, "param_batch_size")
+        else:
+            # If DataGeneratorParameter does not have a batch size we will
+            # vectorization using `n`, and the same checks must be done
+            _check_batch_size(param_data, data, "n")
+
+    if obs_data is not None and param_data is not None:
+        # obs_data batch dimensions need only to be aligned with param_data
+        # batch dimensions if the latter exist
+        if obs_data.obs_batch_size is not None:
+            _check_batch_size(obs_data, param_data, "obs_batch_size")
+        else:
+            _check_batch_size(obs_data, param_data, "n")
 
     if opt_state is None:
         opt_state = optimizer.init(init_params)
@@ -224,6 +213,8 @@ def solve(
     )
     optimization_extra = OptimizationExtraContainer(
         curr_seq=curr_seq,
+        best_iter_id=0,
+        best_val_criterion=jnp.nan,
         best_val_params=init_params,
     )
     loss_container = LossContainer(
@@ -323,16 +314,26 @@ def solve(
                 validation_criterion
             )
 
-            # update best_val_params w.r.t val_loss if needed
-            best_val_params = jax.lax.cond(
+            # update best_val_params and best_val_criterion w.r.t val_loss if needed
+            (best_val_params, best_val_criterion, best_iter_id) = jax.lax.cond(
                 update_best_params,
-                lambda _: params,  # update with current value
-                lambda operands: operands[0].best_val_params,  # unchanged
+                lambda operands: (
+                    params,
+                    validation_criterion,
+                    i,
+                ),  # update with current value
+                lambda operands: (
+                    operands[0].best_val_params,
+                    operands[0].best_val_criterion,
+                    operands[0].best_iter_id,
+                ),  # unchanged
                 (optimization_extra,),
             )
         else:
             early_stopping = False
+            best_iter_id = 0
             best_val_params = params
+            best_val_criterion = jnp.nan
 
         # Trigger RAR
         loss, params, data = trigger_rar(
@@ -358,7 +359,13 @@ def solve(
             i,
             loss,
             OptimizationContainer(params, last_non_nan_params, opt_state),
-            OptimizationExtraContainer(curr_seq, best_val_params, early_stopping),
+            OptimizationExtraContainer(
+                curr_seq,
+                best_iter_id,
+                best_val_criterion,
+                best_val_params,
+                early_stopping,
+            ),
             DataGeneratorContainer(data, param_data, obs_data),
             validation,
             LossContainer(stored_loss_terms, train_loss_values),
@@ -373,7 +380,20 @@ def solve(
         while break_fun(carry):
             carry = _one_iteration(carry)
     else:
-        carry = jax.lax.while_loop(break_fun, _one_iteration, carry)
+
+        def train_fun(carry):
+            return jax.lax.while_loop(break_fun, _one_iteration, carry)
+
+        start = time.time()
+        compiled_train_fun = jax.jit(train_fun).lower(carry).compile()
+        end = time.time()
+        print("\nCompilation took\n", end - start, "\n")
+
+        start = time.time()
+        carry = compiled_train_fun(carry)
+        jax.block_until_ready(carry)
+        end = time.time()
+        print("\nTraining took\n", end - start, "\n")
 
     (
         i,
@@ -389,15 +409,30 @@ def solve(
 
     if verbose:
         jax.debug.print(
-            "Final iteration {i}: train loss value = {train_loss_val}",
+            "\nFinal iteration {i}: train loss value = {train_loss_val}",
             i=i,
             train_loss_val=loss_container.train_loss_values[i - 1],
         )
+
+    # 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:
         jax.debug.print(
             "validation loss value = {validation_loss_val}",
             validation_loss_val=validation_crit_values[i - 1],
         )
+        if optimization_extra.early_stopping:
+            jax.debug.print(
+                "\n Returning a set of best parameters from early stopping"
+                " as last argument!\n"
+                " Best parameters from iteration {best_iter_id}"
+                " with validation loss criterion = {best_val_criterion}",
+                best_iter_id=optimization_extra.best_iter_id,
+                best_val_criterion=optimization_extra.best_val_criterion,
+            )
+            # ...but if early stopping, return the parameters at the best_iter_id
+            validation_parameters = optimization_extra.best_val_params
 
     return (
         optimization.last_non_nan_params,
@@ -408,7 +443,7 @@ def solve(
         optimization.opt_state,
         stored_objects.stored_params,
         validation_crit_values if validation is not None else None,
-        optimization_extra.best_val_params if validation is not None else None,
+        validation_parameters,
     )
 
 
@@ -531,7 +566,7 @@ def _get_break_fun(n_iter: Int, verbose: Bool) -> Callable[[main_carry], Bool]:
             string is not a valid JAX type that can be fed into the operands
             """
             if verbose:
-                jax.debug.print(f"Stopping main optimization loop, cause: {msg}")
+                jax.debug.print(f"\nStopping main optimization loop, cause: {msg}")
             return False
 
         def continue_while_loop(_):

jinns/solver/_utils.py

@@ -0,0 +1,122 @@
+from jinns.data._DataGenerators import (
+    DataGeneratorODE,
+    CubicMeshPDEStatio,
+    CubicMeshPDENonStatio,
+    DataGeneratorParameter,
+)
+
+
+def _check_batch_size(other_data, main_data, attr_name):
+    if isinstance(main_data, DataGeneratorODE):
+        if main_data.temporal_batch_size is not None:
+            if getattr(other_data, attr_name) != main_data.temporal_batch_size:
+                raise ValueError(
+                    f"{other_data.__class__}.{attr_name} must be equal"
+                    f" to {main_data.__class__}.temporal_batch_size for correct"
+                    " vectorization"
+                )
+        else:
+            if main_data.nt is not None:
+                if getattr(other_data, attr_name) != main_data.nt:
+                    raise ValueError(
+                        f"{other_data.__class__}.{attr_name} must be equal"
+                        f" to {main_data.__class__}.nt for correct"
+                        " vectorization"
+                    )
+    if isinstance(main_data, CubicMeshPDEStatio) and not isinstance(
+        main_data, CubicMeshPDENonStatio
+    ):
+        if main_data.omega_batch_size is not None:
+            if getattr(other_data, attr_name) != main_data.omega_batch_size:
+                raise ValueError(
+                    f"{other_data.__class__}.{attr_name} must be equal"
+                    f" to {main_data.__class__}.omega_batch_size for correct"
+                    " vectorization"
+                )
+        else:
+            if main_data.n is not None:
+                if getattr(other_data, attr_name) != main_data.n:
+                    raise ValueError(
+                        f"{other_data.__class__}.{attr_name} must be equal"
+                        f" to {main_data.__class__}.n for correct"
+                        " vectorization"
+                    )
+        if main_data.omega_border_batch_size is not None:
+            if getattr(other_data, attr_name) != main_data.omega_border_batch_size:
+                raise ValueError(
+                    f"{other_data.__class__}.{attr_name} must be equal"
+                    f" to {main_data.__class__}.omega_border_batch_size for correct"
+                    " vectorization"
+                )
+        else:
+            if main_data.nb is not None:
+                if getattr(other_data, attr_name) != main_data.nb:
+                    raise ValueError(
+                        f"{other_data.__class__}.{attr_name} must be equal"
+                        f" to {main_data.__class__}.nb for correct"
+                        " vectorization"
+                    )
+    if isinstance(main_data, CubicMeshPDENonStatio):
+        if main_data.domain_batch_size is not None:
+            if getattr(other_data, attr_name) != main_data.domain_batch_size:
+                raise ValueError(
+                    f"{other_data.__class__}.{attr_name} must be equal"
+                    f" to {main_data.__class__}.domain_batch_size for correct"
+                    " vectorization"
+                )
+        else:
+            if main_data.n is not None:
+                if getattr(other_data, attr_name) != main_data.n:
+                    raise ValueError(
+                        f"{other_data.__class__}.{attr_name} must be equal"
+                        f" to {main_data.__class__}.n for correct"
+                        " vectorization"
+                    )
+        if main_data.border_batch_size is not None:
+            if getattr(other_data, attr_name) != main_data.border_batch_size:
+                raise ValueError(
+                    f"{other_data.__class__}.{attr_name} must be equal"
+                    f" to {main_data.__class__}.border_batch_size for correct"
+                    " vectorization"
+                )
+        else:
+            if main_data.nb is not None:
+                if main_data.dim > 1 and getattr(other_data, attr_name) != (
+                    main_data.nb // 2**main_data.dim
+                ):
+                    raise ValueError(
+                        f"{other_data.__class__}.{attr_name} must be equal"
+                        f" to ({main_data.__class__}.nb // 2**{main_data.__class__}.dim)"
+                        " for correct vectorization"
+                    )
+        if main_data.initial_batch_size is not None:
+            if getattr(other_data, attr_name) != main_data.initial_batch_size:
+                raise ValueError(
+                    f"{other_data.__class__}.{attr_name} must be equal"
+                    f" to {main_data.__class__}.initial_batch_size for correct"
+                    " vectorization"
+                )
+        else:
+            if main_data.ni is not None:
+                if getattr(other_data, attr_name) != main_data.ni:
+                    raise ValueError(
+                        f"{other_data.__class__}.{attr_name} must be equal"
+                        f" to {main_data.__class__}.ni for correct"
+                        " vectorization"
+                    )
+    if isinstance(main_data, DataGeneratorParameter):
+        if main_data.param_batch_size is not None:
+            if getattr(other_data, attr_name) != main_data.param_batch_size:
+                raise ValueError(
+                    f"{other_data.__class__}.{attr_name} must be equal"
+                    f" to {main_data.__class__}.param_batch_size for correct"
+                    " vectorization"
+                )
+        else:
+            if main_data.n is not None:
+                if getattr(other_data, attr_name) != main_data.n:
+                    raise ValueError(
+                        f"{other_data.__class__}.{attr_name} must be equal"
+                        f" to {main_data.__class__}.n for correct"
+                        " vectorization"
+                    )

jinns/utils/__init__.py

@@ -1,4 +1,6 @@
 from ._pinn import create_PINN, PINN
+from ._ppinn import create_PPINN, PPINN
 from ._spinn import create_SPINN, SPINN
 from ._hyperpinn import create_HYPERPINN, HYPERPINN
 from ._save_load import save_pinn, load_pinn
+from ._utils import get_grid

jinns/utils/_containers.py

@@ -38,7 +38,9 @@ class OptimizationContainer(eqx.Module):
 
 class OptimizationExtraContainer(eqx.Module):
     curr_seq: int
-    best_val_params: Params
+    best_iter_id: int  # the best iteration number (that which achieves best_val_params and best_val_params)
+    best_val_criterion: float  # the best validation criterion at early stopping
+    best_val_params: Params  # the best parameter values at early stopping
     early_stopping: Bool = False
 
 

jinns/utils/_hyperpinn.py

@@ -16,7 +16,7 @@ import equinox as eqx
 import numpy as onp
 
 from jinns.utils._pinn import PINN, _MLP
-from jinns.parameters._params import Params
+from jinns.parameters._params import Params, ParamsDict
 
 
 def _get_param_nb(
@@ -114,6 +114,7 @@ class HYPERPINN(PINN):
         )
         self.pinn_params_sum, self.pinn_params_cumsum = _get_param_nb(self.params)
 
+    @property
     def init_params(self) -> Params:
         """
         Returns an initial set of parameters
@@ -138,14 +139,20 @@ class HYPERPINN(PINN):
             is_leaf=lambda x: isinstance(x, jnp.ndarray),
         )
 
-    def eval_nn(
+    def __call__(
         self,
         inputs: Float[Array, "input_dim"],
-        params: Params | PyTree,
+        params: Params | ParamsDict | PyTree,
     ) -> Float[Array, "output_dim"]:
         """
-        Evaluate the HYPERPINN on some inputs with some params.
+        Evaluate the HyperPINN on some inputs with some params.
         """
+        if len(inputs.shape) == 0:
+            # This can happen often when the user directly provides some
+            # collocation points (eg for plotting, whithout using
+            # DataGenerators)
+            inputs = inputs[None]
+
         try:
             hyper = eqx.combine(params.nn_params, self.static_hyper)
         except (KeyError, AttributeError, TypeError) as e:  # give more flexibility
@@ -190,7 +197,7 @@ def create_HYPERPINN(
     slice_solution: slice = None,
     shared_pinn_outputs: slice = None,
     eqx_list_hyper: tuple[tuple[Callable, int, int] | Callable, ...] = None,
-) -> HYPERPINN | list[HYPERPINN]:
+) -> tuple[HYPERPINN | list[HYPERPINN], PyTree | list[PyTree]]:
     r"""
     Utility function to create a standard PINN neural network with the equinox
     library.
@@ -274,6 +281,9 @@ def create_HYPERPINN(
         A HYPERPINN instance or, when `shared_pinn_ouput` is not None,
         a list of HYPERPINN instances with the same structure is returned,
         only differing by there final slicing of the network output.
+    hyperpinn.init_params
+        The initial set of parameters for the HyperPINN or a list of the latter
+        when `shared_pinn_ouput` is not None.
 
 
     Raises
@@ -389,7 +399,7 @@ def create_HYPERPINN(
                     output_slice=output_slice,
                 )
             hyperpinns.append(hyperpinn)
-        return hyperpinns
+        return hyperpinns, [h.init_params for h in hyperpinns]
     with warnings.catch_warnings():
         # Catch the equinox warning because we put the number of
         # parameters as static while being jnp.Array. This this time
@@ -407,4 +417,4 @@ def create_HYPERPINN(
             hypernet_input_size=hypernet_input_size,
             output_slice=None,
         )
-    return hyperpinn
+    return hyperpinn, hyperpinn.init_params

jinns/utils/_pinn.py

@@ -10,7 +10,7 @@ import equinox as eqx
 
 from jaxtyping import Array, Key, PyTree, Float
 
-from jinns.parameters._params import Params
+from jinns.parameters._params import Params, ParamsDict
 
 
 class _MLP(eqx.Module):
@@ -128,40 +128,27 @@ class PINN(eqx.Module):
     def __post_init__(self, mlp):
         self.params, self.static = eqx.partition(mlp, eqx.is_inexact_array)
 
+    @property
     def init_params(self) -> PyTree:
         """
         Returns an initial set of parameters
         """
         return self.params
 
-    def __call__(self, *args) -> Float[Array, "output_dim"]:
-        """
-        Calls `eval_nn` with rearranged arguments
-        """
-        if self.eq_type == "ODE":
-            (t, params) = args
-            if len(t.shape) == 0:
-                t = t[..., None]  #  Add mandatory dimension which can be lacking
-                # (eg. for the ODE batches) but this dimension can already
-                # exists (eg. for user provided observation times)
-            return self.eval_nn(t, params)
-        if self.eq_type == "statio_PDE":
-            (x, params) = args
-            return self.eval_nn(x, params)
-        if self.eq_type == "nonstatio_PDE":
-            (t, x, params) = args
-            t_x = jnp.concatenate([t, x], axis=-1)
-            return self.eval_nn(t_x, params)
-        raise ValueError("Wrong value for self.eq_type")
-
-    def eval_nn(
+    def __call__(
         self,
-        inputs: Float[Array, "input_dim"],
-        params: Params | PyTree,
+        inputs: Float[Array, "1"] | Float[Array, "dim"] | Float[Array, "1+dim"],
+        params: Params | ParamsDict | PyTree,
     ) -> Float[Array, "output_dim"]:
         """
         Evaluate the PINN on some inputs with some params.
         """
+        if len(inputs.shape) == 0:
+            # This can happen often when the user directly provides some
+            # collocation points (eg for plotting, whithout using
+            # DataGenerators)
+            inputs = inputs[None]
+
         try:
             model = eqx.combine(params.nn_params, self.static)
         except (KeyError, AttributeError, TypeError) as e:  # give more flexibility
@@ -193,7 +180,7 @@ def create_PINN(
     ] = None,
     shared_pinn_outputs: tuple[slice] = None,
     slice_solution: slice = None,
-) -> PINN | list[PINN]:
+) -> tuple[PINN | list[PINN], PyTree | list[PyTree]]:
     r"""
     Utility function to create a standard PINN neural network with the equinox
     library.
@@ -266,6 +253,9 @@ def create_PINN(
         A PINN instance or, when `shared_pinn_ouput` is not None,
         a list of PINN instances with the same structure is returned,
         only differing by there final slicing of the network output.
+    pinn.init_params
+        An initial set of parameters for the PINN or a list of the latter
+        when `shared_pinn_ouput` is not None.
 
     Raises
     ------
@@ -322,7 +312,7 @@ def create_PINN(
                 output_slice=output_slice,
             )
             pinns.append(pinn)
-        return pinns
+        return pinns, [p.init_params for p in pinns]
     pinn = PINN(
         mlp=mlp,
         slice_solution=slice_solution,
@@ -331,4 +321,4 @@ def create_PINN(
         output_transform=output_transform,
         output_slice=None,
     )
-    return pinn
+    return pinn, pinn.init_params