jinns 0.8.4__tar.gz → 0.8.6__tar.gz

{jinns-0.8.4 → jinns-0.8.6}/PKG-INFO

@@ -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 → jinns-0.8.6}/doc/source/index.rst

@@ -8,6 +8,14 @@ Welcome to jinn's documentation!
 
 Changelog:
 
+* v0.8.6:
+
+    - Merge `[!37] <https://gitlab.com/mia_jinns/jinns/-/merge_requests/37>`_
+
+* v0.8.5:
+
+    - Merge `[!36] <https://gitlab.com/mia_jinns/jinns/-/merge_requests/36>`_
+
 * v0.8.4:
 
     - Fix a bug: wrong argument in the wrapper function for heterogeneous parameter evaluation of a PDEStatio

{jinns-0.8.4 → jinns-0.8.6}/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-0.8.4 → jinns-0.8.6}/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-0.8.4 → jinns-0.8.6}/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-0.8.4 → jinns-0.8.6}/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 → jinns-0.8.6}/jinns.egg-info/PKG-INFO

@@ -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 → jinns-0.8.6}/jinns.egg-info/SOURCES.txt

@@ -102,4 +102,7 @@ tests/solver_tests_spinn/test_Burger_x32.py
 tests/solver_tests_spinn/test_Fisher_x32.py
 tests/solver_tests_spinn/test_NSPipeFlow_x32_spinn.py
 tests/solver_tests_spinn/test_OU2D_x32.py
-tests/solver_tests_spinn/test_ReactionDiffusion_nonhomo_x64.py
+tests/solver_tests_spinn/test_ReactionDiffusion_nonhomo_x64.py
+tests/utils_tests/test_hyperpinns.py
+tests/utils_tests/test_pinn.py
+tests/utils_tests/test_spinn.py

{jinns-0.8.4 → jinns-0.8.6}/tests/runtests.sh

@@ -16,4 +16,8 @@ if [ $? -ne 0 ]; then
    exit $?
 fi
 pytest solver_tests_spinn/*
+if [ $? -ne 0 ]; then
+   exit $?
+fi
+pytest utils_tests/*
 exit $?

jinns-0.8.6/tests/utils_tests/test_hyperpinns.py

@@ -0,0 +1,138 @@
+"""
+Test script for custom PINN eqx.Module
+"""
+
+import pytest
+import jax
+import jax.random as random
+import jax.numpy as jnp
+import equinox as eqx
+
+import jinns
+from jinns.utils import create_PINN
+import jinns.utils
+
+key = random.PRNGKey(2)
+key, subkey = random.split(key)
+
+d = 5
+n_param = 42
+hyperparams = [f"param {i}" for i in range(n_param)]
+
+EQX_LIST = [
+    [jax.nn.swish],
+    [eqx.nn.Linear, 16, 16],
+    [jax.nn.swish],
+    [eqx.nn.Linear, 16, 16],
+]
+
+eqx_list_hyper = [
+    [eqx.nn.Linear, n_param, 32],  # input is of size 42
+    [jax.nn.tanh],
+    [eqx.nn.Linear, 32, 16],
+    [jax.nn.tanh],
+    [
+        eqx.nn.Linear,
+        16,
+        1000,
+    ],  # 1000 is a random guess, it will automatically be filled with the correct value
+]
+
+
+@pytest.fixture
+def create_pinn_ode():
+    eqx_list = [[eqx.nn.Linear, 1, 16]] + EQX_LIST
+    u_ode = jinns.utils.create_HYPERPINN(
+        subkey, eqx_list, "ODE", hyperparams, n_param, 0, eqx_list_hyper
+    )
+    return u_ode
+
+
+@pytest.fixture
+def create_pinn_statio():
+    eqx_list = [[eqx.nn.Linear, d, 16]] + EQX_LIST
+    u_statio = jinns.utils.create_HYPERPINN(
+        subkey, eqx_list, "statio_PDE", hyperparams, n_param, d, eqx_list_hyper
+    )
+    return u_statio
+
+
+@pytest.fixture
+def create_pinn_nonstatio():
+    eqx_list = [[eqx.nn.Linear, d + 1, 16]] + EQX_LIST
+    u_nonstatio = jinns.utils.create_HYPERPINN(
+        subkey, eqx_list, "nonstatio_PDE", hyperparams, n_param, d + 1, eqx_list_hyper
+    )
+
+    return u_nonstatio
+
+
+@pytest.fixture
+def create_pinn_nonstatio_shared_output():
+
+    # specific argument since we want to have u1 and u2 as separate nns
+    shared_pinn_output = (
+        jnp.s_[:2],
+        jnp.s_[2],
+    )
+
+    eqx_list = [[eqx.nn.Linear, d, 16]] + EQX_LIST
+    u1, u2 = jinns.utils.create_HYPERPINN(
+        subkey,
+        eqx_list,
+        "nonstatio_PDE",
+        hyperparams,
+        n_param,
+        d,
+        eqx_list_hyper,
+        shared_pinn_outputs=shared_pinn_output,
+    )
+    return u1, u2, shared_pinn_output
+
+
+def test_ode_pinn_struct(create_pinn_ode):
+
+    u_ode = create_pinn_ode
+    assert isinstance(u_ode, jinns.utils._hyperpinn.HYPERPINN)
+    assert u_ode.eq_type == "ODE"
+    assert u_ode.output_slice is None
+    assert isinstance(u_ode.slice_solution, slice)
+    _ = u_ode.init_params()
+
+
+def test_statio_pinn_struct(create_pinn_statio):
+
+    u_statio = create_pinn_statio
+    assert u_statio.eq_type == "statio_PDE"
+    assert isinstance(u_statio, jinns.utils._hyperpinn.HYPERPINN)
+    assert u_statio.output_slice is None
+    assert isinstance(u_statio.slice_solution, slice)
+    _ = u_statio.init_params()
+
+
+def test_nonstatio_pinn_struct(create_pinn_nonstatio):
+
+    u_nonstatio = create_pinn_nonstatio
+    assert u_nonstatio.eq_type == "nonstatio_PDE"
+    assert isinstance(u_nonstatio, jinns.utils._hyperpinn.HYPERPINN)
+    assert u_nonstatio.output_slice is None
+    assert isinstance(u_nonstatio.slice_solution, slice)
+    _ = u_nonstatio.init_params()
+
+
+def test_nonstatio_pinn_shared_output(create_pinn_nonstatio_shared_output):
+
+    u1, u2, shared_pinn_ouput = create_pinn_nonstatio_shared_output
+    assert u1.eq_type == "nonstatio_PDE"
+    assert u1.output_slice == shared_pinn_ouput[0]
+    assert isinstance(u1.slice_solution, slice)
+
+    assert u2.eq_type == "nonstatio_PDE"
+    assert u2.output_slice == shared_pinn_ouput[1]
+    assert isinstance(u2.slice_solution, slice)
+
+    param1 = u1.init_params()
+    param2 = u2.init_params()
+
+    # the init parameters for the 2 nns should be the same PyTree
+    assert eqx.tree_equal(param1, param2)

jinns-0.8.6/tests/utils_tests/test_pinn.py

@@ -0,0 +1,131 @@
+"""
+Test script for custom PINN eqx.Module
+"""
+
+import pytest
+import jax
+import jax.random as random
+import jax.numpy as jnp
+import equinox as eqx
+
+import jinns
+from jinns.utils import create_PINN
+
+
+d = 5
+
+
+@pytest.fixture
+def create_pinn_ode():
+    key = random.PRNGKey(2)
+    eqx_list = [
+        [eqx.nn.Linear, 1, 128],
+        [jax.nn.tanh],
+        [eqx.nn.Linear, 128, 1],
+    ]
+    key, subkey = random.split(key)
+    u_statio = create_PINN(subkey, eqx_list, "ODE")
+
+    return u_statio
+
+
+@pytest.fixture
+def create_pinn_statio():
+    key = random.PRNGKey(2)
+    eqx_list = [
+        [eqx.nn.Linear, d, 128],
+        [jax.nn.tanh],
+        [eqx.nn.Linear, 128, 1],
+    ]
+    key, subkey = random.split(key)
+    u_statio = create_PINN(subkey, eqx_list, "statio_PDE", d)
+
+    return u_statio
+
+
+@pytest.fixture
+def create_pinn_nonstatio():
+    key = random.PRNGKey(2)
+    eqx_list = [
+        [eqx.nn.Linear, d, 128],
+        [jax.nn.tanh],
+        [eqx.nn.Linear, 128, 1],
+    ]
+    key, subkey = random.split(key)
+    u_nonstatio = create_PINN(subkey, eqx_list, "nonstatio_PDE", d)
+
+    return u_nonstatio
+
+
+@pytest.fixture
+def create_pinn_nonstatio_shared_output():
+    key = random.PRNGKey(2)
+    eqx_list = [
+        [eqx.nn.Linear, d, 128],
+        [jax.nn.tanh],
+        [eqx.nn.Linear, 128, 1],
+    ]
+    key, subkey = random.split(key)
+    # specific argument since we want to have u1 and u2 as separate nns
+    shared_pinn_output = (
+        jnp.s_[:2],
+        jnp.s_[2],
+    )
+
+    u1, u2 = jinns.utils.create_PINN(
+        subkey,
+        eqx_list,
+        "nonstatio_PDE",
+        d,
+        shared_pinn_outputs=shared_pinn_output,
+    )
+    return u1, u2, shared_pinn_output
+
+
+def test_ode_pinn_struct(create_pinn_ode):
+
+    u_ode = create_pinn_ode
+    assert u_ode.eq_type == "ODE"
+    assert isinstance(u_ode, jinns.utils._pinn.PINN)
+    assert u_ode.output_slice is None
+    assert isinstance(u_ode.slice_solution, slice)
+    _ = u_ode.init_params()
+
+
+def test_statio_pinn_struct(create_pinn_statio):
+
+    u_statio = create_pinn_statio
+    assert u_statio.eq_type == "statio_PDE"
+    assert isinstance(u_statio, jinns.utils._pinn.PINN)
+
+    assert u_statio.output_slice is None
+    assert isinstance(u_statio.slice_solution, slice)
+    _ = u_statio.init_params()
+
+
+def test_nonstatio_pinn_struct(create_pinn_nonstatio):
+
+    u_nonstatio = create_pinn_nonstatio
+    assert u_nonstatio.eq_type == "nonstatio_PDE"
+    assert isinstance(u_nonstatio, jinns.utils._pinn.PINN)
+    assert u_nonstatio.output_slice is None
+    assert isinstance(u_nonstatio.slice_solution, slice)
+    _ = u_nonstatio.init_params()
+
+
+def test_nonstatio_pinn_shared_output(create_pinn_nonstatio_shared_output):
+
+    u1, u2, shared_pinn_ouput = create_pinn_nonstatio_shared_output
+    assert u1.eq_type == "nonstatio_PDE"
+    assert u1.output_slice == shared_pinn_ouput[0]
+    assert isinstance(u1.slice_solution, slice)
+
+    assert u2.eq_type == "nonstatio_PDE"
+    assert u2.output_slice == shared_pinn_ouput[1]
+    assert isinstance(u2.slice_solution, slice)
+
+    param1 = u1.init_params()
+    param2 = u2.init_params()
+
+    # the init parameters for the 2 nns should be the same PyTree
+    assert eqx.tree_equal(param1, param2)

jinns-0.8.6/tests/utils_tests/test_spinn.py

@@ -0,0 +1,110 @@
+"""
+Test script for custom PINN eqx.Module
+"""
+
+import pytest
+import jax
+import jax.random as random
+import jax.numpy as jnp
+import equinox as eqx
+
+import jinns
+from jinns.utils import create_SPINN
+import jinns.utils
+
+
+d = 5
+r = 100  # embedding dim
+m = 1  # output dim
+eqx_list = [
+    [eqx.nn.Linear, 1, 128],
+    [jax.nn.tanh],
+    [eqx.nn.Linear, 128, 128],
+    [jax.nn.tanh],
+    [eqx.nn.Linear, 128, 128],
+    [jax.nn.tanh],
+    [eqx.nn.Linear, 128, r * m],
+]
+
+
+def _assert_attr_equal(u):
+    assert u.m == m
+    assert u.r == r
+
+
+@pytest.fixture
+def create_SPINN_ode():
+    key = random.PRNGKey(2)
+    key, subkey = random.split(key)
+    u_statio = create_SPINN(subkey, 1, r, eqx_list, "ODE", m)
+
+    return u_statio
+
+
+@pytest.fixture
+def create_SPINN_statio():
+    key = random.PRNGKey(2)
+    key, subkey = random.split(key)
+    u_statio = create_SPINN(subkey, d, r, eqx_list, "statio_PDE", m)
+
+    return u_statio
+
+
+@pytest.fixture
+def create_SPINN_nonstatio():
+    key = random.PRNGKey(2)
+    key, subkey = random.split(key)
+    u_nonstatio = create_SPINN(subkey, d, r, eqx_list, "nonstatio_PDE", m)
+
+    return u_nonstatio
+
+
+def test_ode_pinn_struct(create_SPINN_ode):
+
+    u_ode = create_SPINN_ode
+    assert u_ode.eq_type == "ODE"
+    assert isinstance(u_ode, jinns.utils._spinn.SPINN)
+    assert u_ode.d == 1
+    _assert_attr_equal(u_ode)
+    _ = u_ode.init_params()
+
+
+def test_statio_pinn_struct(create_SPINN_statio):
+
+    u_statio = create_SPINN_statio
+    assert u_statio.eq_type == "statio_PDE"
+    assert isinstance(u_statio, jinns.utils._spinn.SPINN)
+
+    assert u_statio.d == d
+    _assert_attr_equal(u_statio)
+    _ = u_statio.init_params()
+
+
+def test_nonstatio_pinn_struct(create_SPINN_nonstatio):
+
+    u_nonstatio = create_SPINN_nonstatio
+    assert u_nonstatio.eq_type == "nonstatio_PDE"
+    assert isinstance(u_nonstatio, jinns.utils._spinn.SPINN)
+    assert u_nonstatio.d == d  # in non-statio SPINN user should include `t` in `d`
+    _assert_attr_equal(u_nonstatio)
+    _ = u_nonstatio.init_params()
+
+
+def test_raising_error_init_SPINN():
+
+    # output_dim != r*m
+    with pytest.raises(ValueError) as e:
+        wrong_eqx_list = [
+            [eqx.nn.Linear, 1, 128],
+            [jax.nn.tanh],
+            [eqx.nn.Linear, 128, r * m + 1],  # output_dim != r*m
+        ]
+        _ = create_SPINN(random.PRNGKey(1), d, r, wrong_eqx_list, "nonstatio_PDE", m)
+
+    # d > 24
+    with pytest.raises(ValueError) as e:
+        _ = create_SPINN(random.PRNGKey(1), 24, r, wrong_eqx_list, "nonstatio_PDE", m)
+
+    # d > 24
+    with pytest.raises(ValueError) as e:
+        _ = create_SPINN(random.PRNGKey(1), 24, r, wrong_eqx_list, "nonstatio_PDE", m)