jinns 0.8.2__py3-none-any.whl → 0.8.4__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/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.dist-info → jinns-0.8.4.dist-info}/METADATA

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

@@ -1,11 +1,11 @@
 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
 jinns/experimental/_diffrax_solver.py,sha256=sLT22byqh-6015_fhe1xtMWlFOYcCjzYKET4sLhA9R4,6818
 jinns/loss/_DynamicLoss.py,sha256=L4CVmmF0rTPbHntgqsLLHlnrlQgLHsetUocpJm7ZYag,27461
-jinns/loss/_DynamicLossAbstract.py,sha256=cXUJoRVphP607WPWZ5PGSHCKPCOtolB0Q6lC09PPeZY,8504
+jinns/loss/_DynamicLossAbstract.py,sha256=kTQlhLx7SBuH5dIDmYaE79sVHUZt1nUFa8LxPU5IHhM,8504
 jinns/loss/_LossODE.py,sha256=sxpgiDR6mfoREuc-qe0AkirOe5K_5oblaYCnodTNxoI,21912
 jinns/loss/_LossPDE.py,sha256=_yX3R-FrAScTn9_QfVC8PfDYRE4UQ5lnzITUYgNFitA,61766
 jinns/loss/_Losses.py,sha256=XOL3MFiKEd3ndsc78Qnpi1vbgR0B2HaAWOGGW2meDM8,11190
@@ -24,8 +24,8 @@ jinns/utils/_save_load.py,sha256=qgZ23nUcB8-B5IZ2guuUWC4M7r5Lxd_Ms3staScdyJo,566
 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.2.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.8.2.dist-info/METADATA,sha256=hTNXy_y_YGU0yGTnJWSsCB6Gj4P40v5m47tmnwyNDcU,2482
-jinns-0.8.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jinns-0.8.2.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.8.2.dist-info/RECORD,,
+jinns-0.8.4.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.8.4.dist-info/METADATA,sha256=QAq8dRIxqTZaBMb0YVOymae5X2kO5XqeHJXJHfe0380,2482
+jinns-0.8.4.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jinns-0.8.4.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.8.4.dist-info/RECORD,,