jinns 0.8.1__py3-none-any.whl → 0.8.3__py3-none-any.whl

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/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.dist-info → jinns-0.8.3.dist-info}/METADATA

@@ -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.dist-info → jinns-0.8.3.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jinns/__init__.py,sha256=Nw5pdlmDhJwco3bXX3YttkeCF8czX_6m0poh8vu0lDQ,113
-jinns/data/_DataGenerators.py,sha256=qiGxuNlS2E1N1gaukA2N379iYn9Gw-s8AbSQN_JQbaI,60397
+jinns/data/_DataGenerators.py,sha256=N4-U4z3MG46UIzHCbKScv9Z7AN40w1wlLY_VsVNj2sI,62293
 jinns/data/__init__.py,sha256=yBOmoavSD-cABp4XcjQY1zsEVO0mDyIhi2MJ5WNp0l8,326
 jinns/data/_display.py,sha256=6renz4H7kHktutmLY7HM6PmxYH7cBfGHpC7GQa1Fnlk,7778
 jinns/experimental/__init__.py,sha256=3jCIy2R2i_0Erwxg-HwISdH79Nt1XCXhS9yY1F5awiY,208
@@ -17,15 +17,15 @@ jinns/solver/_rar.py,sha256=K-0y1-ofOAo1n_Ea3QShSGCGKVYTwiaE_Bz9-DZMJm8,14525
 jinns/solver/_seq2seq.py,sha256=FL-42hTgmVl7O3hHh1ccFVw2bT8bW82hvlDRz971Chk,5620
 jinns/solver/_solve.py,sha256=6kWFWpJ33uOUzZKn7gIOM7yQsVZUwSuorOWPojVeMQY,13721
 jinns/utils/__init__.py,sha256=44ms5UR6vMw3Nf6u4RCAzPFs4fom_YbBnH9mfne8m6k,313
-jinns/utils/_hyperpinn.py,sha256=nuy_V6qIXNxvLKvQAY6aZ_PLVroiiXQZ7RkHlF3GG60,11320
+jinns/utils/_hyperpinn.py,sha256=Mb5d6auzFfXcA81WgjiuhDBvAypAzVOENj_gUeqz6gI,11370
 jinns/utils/_optim.py,sha256=550kxH75TL30o1iKx1swJyP0KqyUPsJ7-imL1w65Qd0,4444
 jinns/utils/_pinn.py,sha256=N8LuB9Ql472O01USghkJkEOmx67DTjc279T8Lj-Lwd4,9722
 jinns/utils/_save_load.py,sha256=qgZ23nUcB8-B5IZ2guuUWC4M7r5Lxd_Ms3staScdyJo,5668
 jinns/utils/_spinn.py,sha256=aeIC3DBY7f_N8HABjvBNv375dMyjll3zt6KjY2bEIkM,8058
 jinns/utils/_utils.py,sha256=8dgvWXX9NT7_7-zltWp0C9tG45ZFNwXxueyxPBb4hjo,6740
 jinns/utils/_utils_uspinn.py,sha256=qcKcOw3zrwWSQyGVj6fD8c9GinHt_U6JWN_k0auTtXM,26039
-jinns-0.8.1.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.8.1.dist-info/METADATA,sha256=AKJ921rioUOvCh8PRFt-_KbvuUMPt73i-i9QV2Orrrg,2482
-jinns-0.8.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jinns-0.8.1.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.8.1.dist-info/RECORD,,
+jinns-0.8.3.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.8.3.dist-info/METADATA,sha256=vUEJOPFDFIZNoDAnUAzW6YbDbtKPNI5C7D5SAXCpXDg,2482
+jinns-0.8.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jinns-0.8.3.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.8.3.dist-info/RECORD,,