jinns 0.8.2__tar.gz → 0.8.4__tar.gz

{jinns-0.8.2 → jinns-0.8.4}/.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.2 → jinns-0.8.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.2
+Version: 0.8.4
 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.2 → jinns-0.8.4}/doc/source/index.rst

@@ -8,9 +8,17 @@ Welcome to jinn's documentation!
 
 Changelog:
 
+* v0.8.4:
+
+    - Fix a bug: wrong argument in the wrapper function for heterogeneous parameter evaluation of a PDEStatio
+
+* 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
+    - Fix a bug: it was not possible to jit a reloaded HyperPINN model.
 
 * v0.8.1:
 

{jinns-0.8.2 → jinns-0.8.4}/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.2 → jinns-0.8.4}/jinns/loss/_DynamicLossAbstract.py

@@ -154,7 +154,7 @@ class PDEStatio(DynamicLoss):
             _params = {
                 "nn_params": params["nn_params"],
                 "eq_params": self.eval_heterogeneous_parameters(
-                    t, u, params, self.eq_params_heterogeneity
+                    x, u, params, self.eq_params_heterogeneity
                 ),
             }
             new_args = args[:-1] + (_params,)

{jinns-0.8.2 → jinns-0.8.4}/jinns.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.8.2
+Version: 0.8.4
 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.2 → jinns-0.8.4}/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.4/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,
+        )