jinns 0.4.1__tar.gz → 0.4.2__tar.gz

{jinns-0.4.1 → jinns-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.4.1
+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.1 → jinns-0.4.2}/doc/source/index.rst

@@ -8,6 +8,10 @@ Welcome to jinn's documentation!
 
 Changelog:
 
+* v0.4.2:
+
+    - Critical bug correction concerning the manipulation of the optimizer's `opt_state` which caused weird failures in the optimization process
+
 * v0.4.1:
 
     - Generalize heterogeneity for the equation parameters. It can now be an arbitrary function provided by the user and thus depend on covariables. Update the corresponding notebook.

{jinns-0.4.1 → jinns-0.4.2}/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.1 → jinns-0.4.2}/jinns.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jinns
-Version: 0.4.1
+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.1 → jinns-0.4.2}/tests/solver_tests/test_Burger_x32.py

@@ -11,7 +11,6 @@ import jinns
 @pytest.fixture
 def train_Burger_init():
     jax.config.update("jax_enable_x64", False)
-    print(jax.config.FLAGS.jax_enable_x64)
     print(jax.devices())
     key = random.PRNGKey(2)
     eqx_list = [
@@ -111,4 +110,4 @@ def test_initial_loss_Burger(train_Burger_init):
 
 def test_10it_Burger(train_Burger_10it):
     total_loss_val = train_Burger_10it
-    assert jnp.round(total_loss_val, 5) == jnp.round(2.0262, 5)
+    assert jnp.round(total_loss_val, 5) == jnp.round(2.04182, 5)

{jinns-0.4.1 → jinns-0.4.2}/tests/solver_tests/test_Burger_x64.py

@@ -11,7 +11,6 @@ import jinns
 @pytest.fixture
 def train_Burger_init():
     jax.config.update("jax_enable_x64", True)
-    print(jax.config.FLAGS.jax_enable_x64)
     print(jax.devices())
     key = random.PRNGKey(2)
     eqx_list = [
@@ -110,4 +109,4 @@ def test_initial_loss_Burger(train_Burger_init):
 
 def test_10it_Burger(train_Burger_10it):
     total_loss_val = train_Burger_10it
-    assert jnp.round(total_loss_val, 5) == jnp.round(2.77797, 5)
+    assert jnp.round(total_loss_val, 5) == jnp.round(2.80653, 5)

{jinns-0.4.1 → jinns-0.4.2}/tests/solver_tests/test_Fisher_x32.py

@@ -12,7 +12,6 @@ import jinns
 @pytest.fixture
 def train_Fisher_init():
     jax.config.update("jax_enable_x64", False)
-    print(jax.config.FLAGS.jax_enable_x64)
     print(jax.devices())
     key = random.PRNGKey(2)
     eqx_list = [
@@ -142,4 +141,4 @@ def test_initial_loss_Fisher(train_Fisher_init):
 
 def test_10it_Fisher(train_Fisher_10it):
     total_loss_val = train_Fisher_10it
-    assert jnp.round(total_loss_val, 5) == jnp.round(10.7401, 5)
+    assert jnp.round(total_loss_val, 5) == jnp.round(10.79058, 5)

{jinns-0.4.1 → jinns-0.4.2}/tests/solver_tests/test_Fisher_x64.py

@@ -12,7 +12,6 @@ import jinns
 @pytest.fixture
 def train_Fisher_init():
     jax.config.update("jax_enable_x64", True)
-    print(jax.config.FLAGS.jax_enable_x64)
     print(jax.devices())
     key = random.PRNGKey(2)
     eqx_list = [
@@ -142,4 +141,4 @@ def test_initial_loss_Fisher(train_Fisher_init):
 
 def test_10it_Fisher(train_Fisher_10it):
     total_loss_val = train_Fisher_10it
-    assert jnp.round(total_loss_val, 5) == jnp.round(10.87023, 5)
+    assert jnp.round(total_loss_val, 5) == jnp.round(10.88394, 5)

{jinns-0.4.1 → jinns-0.4.2}/tests/solver_tests/test_GLV_x32.py

@@ -11,7 +11,6 @@ import jinns
 @pytest.fixture
 def train_GLV_init():
     jax.config.update("jax_enable_x64", False)
-    print(jax.config.FLAGS.jax_enable_x64)
     print(jax.devices())
     key = random.PRNGKey(2)
     key, subkey = random.split(key)
@@ -127,4 +126,4 @@ def test_initial_loss_GLV(train_GLV_init):
 
 def test_10it_GLV(train_GLV_10it):
     total_loss_val = train_GLV_10it
-    assert jnp.round(total_loss_val, 5) == jnp.round(4317.0625, 5)
+    assert jnp.round(total_loss_val, 5) == jnp.round(4318.117, 5)

{jinns-0.4.1 → jinns-0.4.2}/tests/solver_tests/test_GLV_x64.py

@@ -11,7 +11,6 @@ import jinns
 @pytest.fixture
 def train_GLV_init():
     jax.config.update("jax_enable_x64", True)
-    print(jax.config.FLAGS.jax_enable_x64)
     print(jax.devices())
     key = random.PRNGKey(2)
     key, subkey = random.split(key)
@@ -123,4 +122,4 @@ def test_initial_loss_GLV(train_GLV_init):
 
 def test_10it_GLV(train_GLV_10it):
     total_loss_val = train_GLV_10it
-    assert jnp.round(total_loss_val, 5) == jnp.round(3819.72582, 5)
+    assert jnp.round(total_loss_val, 5) == jnp.round(3819.15887, 5)

{jinns-0.4.1 → jinns-0.4.2}/tests/solver_tests/test_NSPipeFlow_x32.py

@@ -11,7 +11,6 @@ import jinns
 @pytest.fixture
 def train_NSPipeFlow_init():
     jax.config.update("jax_enable_x64", False)
-    print(jax.config.FLAGS.jax_enable_x64)
     print(jax.devices())
     key = random.PRNGKey(2)
 
@@ -138,4 +137,4 @@ def test_initial_loss_NSPipeFlow(train_NSPipeFlow_init):
 
 def test_10it_NSPipeFlow(train_NSPipeFlow_10it):
     total_loss_val = train_NSPipeFlow_10it
-    assert jnp.round(total_loss_val, 5) == jnp.round(0.00531, 5)
+    assert jnp.round(total_loss_val, 5) == jnp.round(0.00534, 5)

{jinns-0.4.1 → jinns-0.4.2}/tests/solver_tests/test_NSPipeFlow_x64.py

@@ -11,7 +11,6 @@ import jinns
 @pytest.fixture
 def train_NSPipeFlow_init():
     jax.config.update("jax_enable_x64", True)
-    print(jax.config.FLAGS.jax_enable_x64)
     print(jax.devices())
     key = random.PRNGKey(2)
 
@@ -138,4 +137,4 @@ def test_initial_loss_NSPipeFlow(train_NSPipeFlow_init):
 
 def test_10it_NSPipeFlow(train_NSPipeFlow_10it):
     total_loss_val = train_NSPipeFlow_10it
-    assert jnp.round(total_loss_val, 5) == jnp.round(0.00495, 5)
+    assert jnp.round(total_loss_val, 5) == jnp.round(0.00504, 5)