jinns 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

jinns/loss/_DynamicLossAbstract.py

@@ -29,10 +29,9 @@ class DynamicLoss:
             equation solution with as PINN.
         eq_params_heterogeneity
             Default None. A dict with the keys being the same as in eq_params
-            and the value being `time`, `space`, `both` or None which corresponds to
-            the heterogeneity of a given parameter. A value can be missing, in
-            this case there is no heterogeneity (=None). If
-            eq_params_heterogeneity is None this means there is no
+            and the value being either None (no heterogeneity) or a function which encodes for the spatio-temporal heterogeneity of the parameter. Such a function must be jittable and take three arguments `t`, `x` and `params["eq_params"]` even if one is not used. Therefore, one can introduce spatio-temporal covariates upon which a particular parameter can depend, e.g. in a GLM fashion. The effect of these covariables can themselves be estimated by being in `eq_params` too.
+            A value can be missing, in this case there is no heterogeneity (=None).
+            If eq_params_heterogeneity is None this means there is no
             heterogeneity for no parameters.
         """
         self.Tmax = Tmax
@@ -49,48 +48,18 @@ class DynamicLoss:
             return eq_params
         for k, p in eq_params.items():
             try:
-                eq_params_[k] = self._eval_heterogeneous_array_parameter(
-                    p, t, x, heterogeneity=eq_params_heterogeneity[k]
-                )
+                if eq_params_heterogeneity[k] is None:
+                    eq_params_[k] = p
+                else:
+                    eq_params_[k] = eq_params_heterogeneity[k](
+                        t, x, eq_params  # heterogeneity encoded through a function
+                    )
             except KeyError:
                 # we authorize missing eq_params_heterogeneity key
                 # is its heterogeneity is None anyway
                 eq_params_[k] = p
         return eq_params_
 
-    def _eval_heterogeneous_array_parameter(self, p, t, x, heterogeneity=None):
-        """
-        For time and/or space heterogeneous params defined by an n-dimensional
-        array `p` we return the value `p[t, x]` with discretization of the
-        collocation point
-
-        Parameters
-        ----------
-        p
-            The parameter
-        heterogeneity
-            A string. Either `time`, `space`, `both` or None to specify which
-            kind of heterogeneity we have. Default is None, is this case we do
-            not have heterogeneity.
-
-
-        **Note** t is assumed to be normalized in [0, 1] as well as x!
-        """
-        if heterogeneity is None:
-            return p
-        elif heterogeneity == "time":
-            return p[(t * len(p)).astype(int)]
-        elif heterogeneity == "space":
-            coords = (x * jnp.array(p.shape)).astype(int)
-            return jnp.take(p, jnp.ravel_multi_index(coords, p.shape, mode="clip"))
-        elif heterogeneity == "both":
-            coords = jnp.concatenate(
-                [(t * len(p))[:, None], x * jnp.array(p.shape)], axis=1
-            ).astype(int)
-            return jnp.take(p, jnp.ravel_multi_index(coords, p.shape, mode="clip"))
-        else:
-            raise ValueError("Wrong paramater value for parameter `heterogeneity`")
-
     def set_stop_gradient(self, params_dict):
         """
         Set the stop gradient operators in the dynamic loss `evaluate`

jinns/solver/_solve.py

@@ -196,7 +196,7 @@ def solve(
         if carry["param_data"] is not None:
             batch = append_param_batch(batch, carry["param_data"].get_batch())
         carry["params"], carry["opt_state"] = optimizer.update(
-            params=carry["params"], state=carry["state"], batch=batch
+            params=carry["params"], state=carry["opt_state"], batch=batch
         )
 
         # check if any of the parameters is NaN
@@ -260,7 +260,6 @@ def solve(
         {
             "params": init_params,
             "last_non_nan_params": init_params.copy(),
-            "state": opt_state,
             "data": data,
             "curr_seq": curr_seq,
             "seq2seq": seq2seq,
@@ -281,7 +280,7 @@ def solve(
 
     params = res["params"]
     last_non_nan_params = res["last_non_nan_params"]
-    opt_state = res["state"]
+    opt_state = res["opt_state"]
     data = res["data"]
     loss = res["loss"]
 

{jinns-0.4.0.dist-info → jinns-0.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.4.0
+Version: 0.4.2
 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.4.0.dist-info → jinns-0.4.2.dist-info}/RECORD

@@ -3,7 +3,7 @@ jinns/data/_DataGenerators.py,sha256=nIuKtkX4V4ckfT4-g0bjlY7BLkgcok5JbI9OzJn73mA
 jinns/data/__init__.py,sha256=S13J59Fxuph4uNJ542fP_Mj8U72ilhb5t_UQ-c1k3nY,232
 jinns/data/_display.py,sha256=Xnfo6_PH1g-ZFpWJcbF6CF6Pp12wJtNQb1W1bADuQrA,6134
 jinns/loss/_DynamicLoss.py,sha256=VyoyWdkoxRPeP2vs4ZZBK_T9xWgwkcDuaFrjUSid3Zo,52975
-jinns/loss/_DynamicLossAbstract.py,sha256=O_dxb7_0yBCxsoa7ornU5PvqOqDPlkOwK55CGgdlGSs,8578
+jinns/loss/_DynamicLossAbstract.py,sha256=V9NJvMmqnSC06yccu9bkFCWLF3cjyr8ze8qf0LRykjo,7718
 jinns/loss/_LossODE.py,sha256=FHTKQPqLSoMh18j_RYUoR7tUGg_ljd0JKL8xf2HiF5M,17541
 jinns/loss/_LossPDE.py,sha256=8wEEPnkibNOdqCXtcuU5nieRsPoffSyV8wcUMjcUkvg,57145
 jinns/loss/__init__.py,sha256=4JxMHHVMxTMsVZmV8mRSIyMVAEp3QIR8QtZvnwvj96Q,560
@@ -12,11 +12,11 @@ jinns/loss/_operators.py,sha256=HYGDq3K_K7lztDp6al88Q78F6o-hDRo3ODhw8UMRdC8,5690
 jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jinns/solver/_rar.py,sha256=V9y07F6objmP6rPA305dIJ82h7kwP4AWBAktZ68b-38,13894
 jinns/solver/_seq2seq.py,sha256=XNL9e0fBj85Q86XfGDzq9dzxIkPPMwoJF38C8doNYtM,6032
-jinns/solver/_solve.py,sha256=pCctVnrkfOptGd1wH2w9jDdFiOMmNqgatqaoYtyE_LM,10071
+jinns/solver/_solve.py,sha256=Yz_asD0ZuYN923E6ysxtPemN-36Zt-BtlSPkTX-BfA8,10047
 jinns/utils/__init__.py,sha256=-jDlwCjyEzWweswKdwLal3OhaUU3FVzK_Ge2S-7KHXs,149
 jinns/utils/_utils.py,sha256=bQm6z_xPKJj9BMCr2tXc44IA8JyGNN-PR5LNRhZ1fD8,20085
-jinns-0.4.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.4.0.dist-info/METADATA,sha256=iJyoQcOqKE4tgof5Xr_gfOgIKWkhD_VmfWGfHj9GJ14,1821
-jinns-0.4.0.dist-info/WHEEL,sha256=Xo9-1PvkuimrydujYJAjF7pCkriuXBpUPEjma1nZyJ0,92
-jinns-0.4.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.4.0.dist-info/RECORD,,
+jinns-0.4.2.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.4.2.dist-info/METADATA,sha256=maaTIojnCdHhTIPd2kCJa8QGfX36mWAt1DN7q8Pd_3o,1821
+jinns-0.4.2.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+jinns-0.4.2.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.4.2.dist-info/RECORD,,

{jinns-0.4.0.dist-info → jinns-0.4.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.3)
+Generator: bdist_wheel (0.42.0)
 Root-Is-Purelib: true
 Tag: py3-none-any