jinns 0.8.4__py3-none-any.whl → 0.8.6__py3-none-any.whl

jinns/solver/_solve.py

@@ -218,6 +218,7 @@ def solve(
         batch, data, param_data, obs_data = get_batch(data, param_data, obs_data)
 
         (
+            loss,
             loss_val,
             loss_terms,
             params,
@@ -311,7 +312,7 @@ def solve(
     )
 
 
-@partial(jit, static_argnames=["loss", "optimizer"])
+@partial(jit, static_argnames=["optimizer"])
 def gradient_step(loss, optimizer, batch, params, opt_state, last_non_nan_params):
     """
     loss and optimizer cannot be jit-ted.
@@ -330,6 +331,7 @@ def gradient_step(loss, optimizer, batch, params, opt_state, last_non_nan_params
     )
 
     return (
+        loss,
         loss_val,
         loss_terms,
         params,

jinns/utils/_hyperpinn.py

@@ -216,13 +216,7 @@ def create_HYPERPINN(
 
     Returns
     -------
-    init_fn
-        A function which (re-)initializes the PINN parameters with the provided
-        jax random key
-    apply_fn
-        A function to apply the neural network on given inputs for given
-        parameters. A typical call will be of the form `u(t, params)` for
-        ODE or `u(t, x, params)` for nD PDEs (`x` being multidimensional)
+    `u`, a :class:`.HyperPINN` object which inherits from `eqx.Module` (hence callable).
 
     Raises
     ------
@@ -289,7 +283,6 @@ def create_HYPERPINN(
 
     if shared_pinn_outputs is not None:
         hyperpinns = []
-        static = None
         for output_slice in shared_pinn_outputs:
             hyperpinn = HYPERPINN(
                 mlp,
@@ -302,11 +295,6 @@ def create_HYPERPINN(
                 hypernet_input_size,
                 output_slice,
             )
-            # all the pinns are in fact the same so we share the same static
-            if static is None:
-                static = hyperpinn.static
-            else:
-                hyperpinn.static = static
             hyperpinns.append(hyperpinn)
         return hyperpinns
     hyperpinn = HYPERPINN(

jinns/utils/_pinn.py

@@ -200,13 +200,7 @@ def create_PINN(
 
     Returns
     -------
-    init_fn
-        A function which (re-)initializes the PINN parameters with the provided
-        jax random key
-    apply_fn
-        A function to apply the neural network on given inputs for given
-        parameters. A typical call will be of the form `u(t, params)` for
-        ODE or `u(t, x, params)` for nD PDEs (`x` being multidimensional)
+    `u`, a :class:`.PINN` object which inherits from `eqx.Module` (hence callable). This comes with a bound method :func:`u.init_params() <PINN.init_params>`. When `shared_pinn_ouput` is not None, a list of :class:`.PINN` with the same structure is returned, only differing by there final slicing of the network output.
 
     Raises
     ------
@@ -253,7 +247,6 @@ def create_PINN(
 
     if shared_pinn_outputs is not None:
         pinns = []
-        static = None
         for output_slice in shared_pinn_outputs:
             pinn = PINN(
                 mlp,
@@ -263,11 +256,6 @@ def create_PINN(
                 output_transform,
                 output_slice,
             )
-            # all the pinns are in fact the same so we share the same static
-            if static is None:
-                static = pinn.static
-            else:
-                pinn.static = static
             pinns.append(pinn)
         return pinns
     pinn = PINN(mlp, slice_solution, eq_type, input_transform, output_transform, None)

jinns/utils/_spinn.py

@@ -194,13 +194,7 @@ def create_SPINN(key, d, r, eqx_list, eq_type, m=1):
 
     Returns
     -------
-    init_fn
-        A function which (re-)initializes the SPINN parameters with the provided
-        jax random key
-    apply_fn
-        A function to apply the neural network on given inputs for given
-        parameters. A typical call will be of the form `u(t, params)` for
-        ODE or `u(t, x, params)` for nD PDEs (`x` being multidimensional)
+    `u`, a :class:`.SPINN` object which inherits from `eqx.Module` (hence callable).
 
     Raises
     ------

{jinns-0.8.4.dist-info → jinns-0.8.6.dist-info}/METADATA

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

@@ -15,17 +15,17 @@ jinns/loss/_operators.py,sha256=zDGJqYqeYH7xd-4dtGX9PS-pf0uSOpUUXGo5SVjIJ4o,1106
 jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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/solver/_solve.py,sha256=r4jn6hx7_t-Y2rBWA2npUmWWnDg4iRbgYBHZDNn9tmY,13745
 jinns/utils/__init__.py,sha256=44ms5UR6vMw3Nf6u4RCAzPFs4fom_YbBnH9mfne8m6k,313
-jinns/utils/_hyperpinn.py,sha256=Mb5d6auzFfXcA81WgjiuhDBvAypAzVOENj_gUeqz6gI,11370
+jinns/utils/_hyperpinn.py,sha256=93hbiATdp5W4l1cu9Oe6O2c45o-ZF_z2u6FzNLyjnm4,10878
 jinns/utils/_optim.py,sha256=550kxH75TL30o1iKx1swJyP0KqyUPsJ7-imL1w65Qd0,4444
-jinns/utils/_pinn.py,sha256=N8LuB9Ql472O01USghkJkEOmx67DTjc279T8Lj-Lwd4,9722
+jinns/utils/_pinn.py,sha256=mhA4-3PazyQTbWIx9oLaNwL0QDe8ZIBhbiy5J3kwa4I,9471
 jinns/utils/_save_load.py,sha256=qgZ23nUcB8-B5IZ2guuUWC4M7r5Lxd_Ms3staScdyJo,5668
-jinns/utils/_spinn.py,sha256=aeIC3DBY7f_N8HABjvBNv375dMyjll3zt6KjY2bEIkM,8058
+jinns/utils/_spinn.py,sha256=SzOUt1KHtB9QOpghpvitnXN-KEqXUXbvabC5k0TnKEo,7793
 jinns/utils/_utils.py,sha256=8dgvWXX9NT7_7-zltWp0C9tG45ZFNwXxueyxPBb4hjo,6740
 jinns/utils/_utils_uspinn.py,sha256=qcKcOw3zrwWSQyGVj6fD8c9GinHt_U6JWN_k0auTtXM,26039
-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,,
+jinns-0.8.6.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.8.6.dist-info/METADATA,sha256=3Ml6PCA-569v9-1FgyPDySX09RQas0zPOVEV_gqy9lk,2482
+jinns-0.8.6.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jinns-0.8.6.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.8.6.dist-info/RECORD,,