jinns 0.8.1__tar.gz → 0.8.3__tar.gz

{jinns-0.8.1 → jinns-0.8.3}/.pre-commit-config.yaml

@@ -1,14 +1,14 @@
 default_stages: [commit]
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v4.6.0
     hooks:
     -   id: trailing-whitespace
         stages: [commit]
     -   id: end-of-file-fixer
         stages: [commit]
   - repo: https://github.com/psf/black
-    rev: 24.2.0
+    rev: 24.3.0
     hooks:
     -   id: black
         stages: [commit]

{jinns-0.8.1 → jinns-0.8.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.1
+Version: 0.8.3
 Summary: Physics Informed Neural Network with JAX
 Author-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
 Maintainer-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-0.8.1 → jinns-0.8.3}/doc/source/index.rst

@@ -8,6 +8,14 @@ Welcome to jinn's documentation!
 
 Changelog:
 
+* v0.8.3:
+
+    - Add the possibility to load user-provided tables of parameters in DataGeneratorParameter and not only to randomly sample them.
+
+* v0.8.2:
+
+    - Fix a bug: it was not possible to jit a reloaded HyperPINN model.
+
 * v0.8.1:
 
     - New feature: `save_pinn` and `load_pinn` in `jinns.utils` for pre-trained

{jinns-0.8.1 → jinns-0.8.3}/jinns/data/_DataGenerators.py

@@ -1101,8 +1101,9 @@ class DataGeneratorParameter:
         key,
         n,
         param_batch_size,
-        param_ranges,
+        param_ranges=None,
         method="grid",
+        user_data=None,
         data_exists=False,
     ):
         r"""
@@ -1140,25 +1141,47 @@ class DataGeneratorParameter:
             Must be left to `False` when created by the user. Avoids the
             regeneration of :math:`\Omega`, :math:`\partial\Omega` and
             time points at each pytree flattening and unflattening.
+        user_data
+            A dictionary containing user-provided data for parameters.
+            As for `param_ranges`, the key corresponds to the parameter name,
+            the keys must match the keys in `params["eq_params"]` and only
+            unidimensional arrays are supported. Therefore, the jnp arrays
+            found at `user_data[k]` must have shape `(n, 1)` or `(n,)`.
+            Note that if the same key appears in `param_ranges` and `user_data`
+            priority goes for the content in `user_data`.
+            Defaults to None.
         """
         self.data_exists = data_exists
         self.method = method
-        if not isinstance(key, dict):
-            self._keys = dict(
-                zip(param_ranges.keys(), jax.random.split(key, len(param_ranges)))
+
+        if n < param_batch_size:
+            raise ValueError(
+                f"Number of data points ({n}) is smaller than the"
+                f"number of batch points ({param_batch_size})."
             )
+
+        if user_data is None:
+            user_data = {}
+        if param_ranges is None:
+            param_ranges = {}
+        if not isinstance(key, dict):
+            all_keys = set().union(param_ranges, user_data)
+            self._keys = dict(zip(all_keys, jax.random.split(key, len(all_keys))))
         else:
             self._keys = key
         self.n = n
         self.param_batch_size = param_batch_size
         self.param_ranges = param_ranges
+        self.user_data = user_data
 
         if not self.data_exists:
             self.generate_data()
             # The previous call to self.generate_data() has created
-            # the dict self.param_n_samples
+            # the dict self.param_n_samples and then we will only use this one
+            # because it has merged the scattered data between `user_data` and
+            # `param_ranges`
             self.curr_param_idx = {}
-            for k in self.param_ranges.keys():
+            for k in self.param_n_samples.keys():
                 self.curr_param_idx[k] = 0
                 (
                     self._keys[k],
@@ -1167,22 +1190,40 @@ class DataGeneratorParameter:
                 ) = _reset_batch_idx_and_permute(self._get_param_operands(k))
 
     def generate_data(self):
-        # Generate param n samples
+        """
+        Generate parameter samples, either through generation
+        or using user-provided data.
+        """
         self.param_n_samples = {}
-        for k, e in self.param_ranges.items():
-            if self.method == "grid":
-                xmin, xmax = e[0], e[1]
-                self.partial = (xmax - xmin) / self.n
-                # shape (n, 1)
-                self.param_n_samples[k] = jnp.arange(xmin, xmax, self.partial)[:, None]
-            elif self.method == "uniform":
-                xmin, xmax = e[0], e[1]
-                self._keys[k], subkey = random.split(self._keys[k], 2)
-                self.param_n_samples[k] = random.uniform(
-                    subkey, shape=(self.n, 1), minval=xmin, maxval=xmax
-                )
+
+        all_keys = set().union(self.param_ranges, self.user_data)
+        for k in all_keys:
+            if self.user_data and k in self.user_data:
+                if self.user_data[k].shape == (self.n, 1):
+                    self.param_n_samples[k] = self.user_data[k]
+                if self.user_data[k].shape == (self.n,):
+                    self.param_n_samples[k] = self.user_data[k][:, None]
+                else:
+                    raise ValueError(
+                        "Wrong shape for user provided parameters"
+                        f" in user_data dictionary at key='{k}'"
+                    )
             else:
-                raise ValueError("Method " + self.method + " is not implemented.")
+                if self.method == "grid":
+                    xmin, xmax = self.param_ranges[k][0], self.param_ranges[k][1]
+                    self.partial = (xmax - xmin) / self.n
+                    # shape (n, 1)
+                    self.param_n_samples[k] = jnp.arange(xmin, xmax, self.partial)[
+                        :, None
+                    ]
+                elif self.method == "uniform":
+                    xmin, xmax = self.param_ranges[k][0], self.param_ranges[k][1]
+                    self._keys[k], subkey = random.split(self._keys[k], 2)
+                    self.param_n_samples[k] = random.uniform(
+                        subkey, shape=(self.n, 1), minval=xmin, maxval=xmax
+                    )
+                else:
+                    raise ValueError("Method " + self.method + " is not implemented.")
 
     def _get_param_operands(self, k):
         return (
@@ -1247,12 +1288,7 @@ class DataGeneratorParameter:
         )
         aux_data = {
             k: vars(self)[k]
-            for k in [
-                "n",
-                "param_batch_size",
-                "method",
-                "param_ranges",
-            ]
+            for k in ["n", "param_batch_size", "method", "param_ranges", "user_data"]
         }
         return (children, aux_data)
 

{jinns-0.8.1 → jinns-0.8.3}/jinns/utils/_hyperpinn.py

@@ -39,8 +39,8 @@ class HYPERPINN(PINN):
     static_hyper: eqx.Module
     hyperparams: list = eqx.field(static=True)
     hypernet_input_size: int
-    pinn_params_sum: ArrayLike
-    pinn_params_cumsum: ArrayLike
+    pinn_params_sum: ArrayLike = eqx.field(static=True)
+    pinn_params_cumsum: ArrayLike = eqx.field(static=True)
 
     def __init__(
         self,

{jinns-0.8.1 → jinns-0.8.3}/jinns.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.1
+Version: 0.8.3
 Summary: Physics Informed Neural Network with JAX
 Author-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>
 Maintainer-email: Hugo Gangloff <hugo.gangloff@inrae.fr>, Nicolas Jouvin <nicolas.jouvin@inrae.fr>

{jinns-0.8.1 → jinns-0.8.3}/jinns.egg-info/SOURCES.txt

@@ -82,6 +82,7 @@ tests/runtests.sh
 tests/dataGenerator_tests/test_CubicMeshPDENonStatio.py
 tests/dataGenerator_tests/test_CubicMeshPDEStatio.py
 tests/dataGenerator_tests/test_DataGeneratorODE.py
+tests/dataGenerator_tests/test_DataGeneratorParameter.py
 tests/save_load_tests/test_saving_loading_hyperpinn.py
 tests/save_load_tests/test_saving_loading_pinn.py
 tests/save_load_tests/test_saving_loading_spinn.py

jinns-0.8.3/tests/dataGenerator_tests/test_DataGeneratorParameter.py

@@ -0,0 +1,80 @@
+import pytest
+import jax
+import jax.numpy as jnp
+import jinns
+
+n = 64
+
+
+@pytest.fixture
+def create_DataGeneratorParameter():
+    key = jax.random.PRNGKey(2)
+    key, subkey = jax.random.split(key)
+
+    param_batch_size = 64
+    method = "uniform"
+    param_ranges = {"theta": (10.0, 11.0)}
+    user_data = {"nu": jnp.arange(n)}
+
+    return jinns.data.DataGeneratorParameter(
+        subkey,
+        n,
+        param_batch_size,
+        param_ranges=param_ranges,
+        method=method,
+        user_data=user_data,
+    )
+
+
+@pytest.fixture
+def create_DataGeneratorParameter_only_user_data():
+    key = jax.random.PRNGKey(2)
+    key, subkey = jax.random.split(key)
+
+    param_batch_size = 64
+    method = "uniform"
+    user_data = {"nu": jnp.arange(n)}
+
+    return jinns.data.DataGeneratorParameter(
+        subkey,
+        n,
+        param_batch_size,
+        method=method,
+        user_data=user_data,
+    )
+
+
+def test_get_batch(create_DataGeneratorParameter):
+    data_generator_parameters = create_DataGeneratorParameter
+    param_batch = data_generator_parameters.get_batch()
+    assert jnp.allclose(jnp.sort(jnp.unique(param_batch["nu"])), jnp.arange(n)) and (
+        jnp.all(param_batch["theta"] >= 10.0) and jnp.all(param_batch["theta"] <= 11.0)
+    )
+
+
+def test_get_batch_only_user_data(create_DataGeneratorParameter_only_user_data):
+    data_generator_parameters = create_DataGeneratorParameter_only_user_data
+    param_batch = data_generator_parameters.get_batch()
+    assert jnp.allclose(jnp.sort(jnp.unique(param_batch["nu"])), jnp.arange(n))
+
+
+def test_raise_error_with_wrong_shape_for_user_data():
+    key = jax.random.PRNGKey(2)
+    key, subkey = jax.random.split(key)
+
+    param_batch_size = 64
+    method = "uniform"
+    param_ranges = {"theta": (10.0, 11.0)}
+    # user_data is not (n,) or (n,1)
+    user_data = {"nu": jnp.ones((n, 1, 1))}
+
+    with pytest.raises(ValueError) as e_info:
+        # __init__ calls self.generate_data() that we are testing for
+        data_generator_parameters = jinns.data.DataGeneratorParameter(
+            subkey,
+            n,
+            param_batch_size,
+            param_ranges=param_ranges,
+            method=method,
+            user_data=user_data,
+        )

{jinns-0.8.1 → jinns-0.8.3}/tests/save_load_tests/test_saving_loading_hyperpinn.py

@@ -9,7 +9,7 @@ from jinns.utils import save_pinn, load_pinn
 
 
 @pytest.fixture
-def save_reload():
+def save_reload(tmpdir):
     jax.config.update("jax_enable_x64", False)
     key = random.PRNGKey(2)
 
@@ -73,7 +73,7 @@ def save_reload():
     }
 
     # Save
-    filename = "./test"
+    filename = str(tmpdir.join("test"))
     kwargs_creation = {
         "key": subkey,
         "eqx_list": eqx_list,
@@ -107,3 +107,24 @@ def test_equality_save_reload(save_reload):
         v_u_reloaded(test_points[:, 0:1], test_points[:, 1:3], params_reloaded),
         atol=1e-3,
     )
+
+
+def test_jitting_reloaded_hyperpinn(save_reload):
+    """
+    This test ensures that the reloaded hyperpinn is jit-able.
+    Some conversion of onp.array nodes can arise when reloading.
+    See this MR : https://gitlab.com/mia_jinns/jinns/-/merge_requests/32
+    jinns v0.8.2 uses eqx.field(static=True) to solve the problem.
+    This tests is here for testimony.
+    """
+
+    key, _, _, params_reloaded, u_reloaded = save_reload
+
+    v_u_reloaded = jax.vmap(
+        u_reloaded, (0, 0, {"nn_params": None, "eq_params": {"D": 0, "r": 0}})
+    )
+    v_u_reloaded_jitted = jax.jit(v_u_reloaded)
+
+    key, subkey = jax.random.split(key, 2)
+    test_points = jax.random.normal(subkey, shape=(10, 5))
+    v_u_reloaded_jitted(test_points[:, 0:1], test_points[:, 1:3], params_reloaded)

{jinns-0.8.1 → jinns-0.8.3}/tests/save_load_tests/test_saving_loading_pinn.py

@@ -9,7 +9,7 @@ from jinns.utils import save_pinn, load_pinn
 
 
 @pytest.fixture
-def save_reload():
+def save_reload(tmpdir):
     jax.config.update("jax_enable_x64", False)
     key = random.PRNGKey(2)
     eqx_list = [
@@ -28,7 +28,7 @@ def save_reload():
     params = {"nn_params": params, "eq_params": {}}
 
     # Save
-    filename = "./test"
+    filename = str(tmpdir.join("test"))
     kwargs_creation = {
         "key": subkey,
         "eqx_list": eqx_list,
@@ -57,3 +57,22 @@ def test_equality_save_reload(save_reload):
         v_u_reloaded(test_points[:, 0:1], test_points[:, 1:], params_reloaded),
         atol=1e-3,
     )
+
+
+def test_jitting_reloaded_pinn(save_reload):
+    """
+    This test ensures that the reloaded pinn is jit-able.
+    Some conversion of onp.array nodes can arise when reloading.
+    See this MR : https://gitlab.com/mia_jinns/jinns/-/merge_requests/32
+    jinns v0.8.2 uses eqx.field(static=True) to solve the problem.
+    This tests is here for testimony.
+    """
+
+    key, _, _, params_reloaded, u_reloaded = save_reload
+
+    key, subkey = jax.random.split(key, 2)
+    test_points = jax.random.normal(subkey, shape=(10, 2))
+    v_u_reloaded = jax.vmap(u_reloaded, (0, 0, None))
+    v_u_reloaded_jitted = jax.jit(v_u_reloaded)
+
+    v_u_reloaded_jitted(test_points[:, 0:1], test_points[:, 1:3], params_reloaded)

{jinns-0.8.1 → jinns-0.8.3}/tests/save_load_tests/test_saving_loading_spinn.py

@@ -9,7 +9,7 @@ from jinns.utils import save_pinn, load_pinn
 
 
 @pytest.fixture
-def save_reload():
+def save_reload(tmpdir):
     jax.config.update("jax_enable_x64", False)
     key = random.PRNGKey(2)
     d = 2
@@ -30,7 +30,7 @@ def save_reload():
     params = {"nn_params": params, "eq_params": {}}
 
     # Save
-    filename = "./test"
+    filename = str(tmpdir.join("test"))
     kwargs_creation = {
         "key": subkey,
         "d": d,
@@ -58,3 +58,21 @@ def test_equality_save_reload(save_reload):
         u_reloaded(test_points[:, 0:1], test_points[:, 1:], params_reloaded),
         atol=1e-3,
     )
+
+
+def test_jitting_reloaded_spinn(save_reload):
+    """
+    This test ensures that the reloaded spinn is jit-able.
+    Some conversion of onp.array nodes can arise when reloading.
+    See this MR : https://gitlab.com/mia_jinns/jinns/-/merge_requests/32
+    jinns v0.8.2 uses eqx.field(static=True) to solve the problem.
+    This tests is here for testimony.
+    """
+
+    key, _, _, params_reloaded, u_reloaded = save_reload
+
+    key, subkey = jax.random.split(key, 2)
+    test_points = jax.random.normal(subkey, shape=(10, 2))
+
+    u_reloaded_jitted = jax.jit(u_reloaded.__call__)
+    u_reloaded_jitted(test_points[:, 0:1], test_points[:, 1:], params_reloaded)