jinns 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

jinns/solver/_solve.py

@@ -1,8 +1,8 @@
 from functools import partial
-from jaxopt import OptaxSolver
+from jaxopt import OptaxSolver, LBFGS
+from optax import GradientTransformation
 from jax_tqdm import scan_tqdm
 import jax
-from jax import jit
 import jax.numpy as jnp
 from jinns.solver._seq2seq import (
     _initialize_seq2seq,
@@ -24,7 +24,8 @@ def solve(
     init_params,
     data,
     loss,
-    optax_solver,
+    optimizer,
+    print_loss_every=1000,
     opt_state=None,
     seq2seq=None,
     tracked_params_key_list=None,
@@ -54,11 +55,19 @@ def solve(
         A loss object (e.g. a LossODE, SystemLossODE, LossPDEStatio [...]
         object). It must be jittable (e.g. implements via a pytree
         registration)
-    optax_solver
-        An optax solver (e.g. adam with a given step-size)
+    optimizer
+        Can be an `optax` optimizer (e.g. `optax.adam`).
+        In such case, it is wrapped in the `jaxopt.OptaxSolver` wrapper.
+        Can be a `jaxopt` optimizer (e.g. `jaxopt.BFGS`) which supports the
+        methods `init_state` and `update`.
+        Can be a string (currently only `bfgs`), in such case a `jaxopt`
+        optimizer is created with default parameters.
+    print_loss_every
+        Integer. Default 100. The rate at which we print the loss value in the
+        gradient step loop.
     opt_state
         Default None. Provide an optional initial optional state to the
-        optimizer
+        optimizer. Not valid for all optimizers.
     seq2seq
         Default None. A dictionary with keys 'times_steps'
         and 'iter_steps' which mush have same length. The first represents
@@ -99,13 +108,16 @@ def solve(
     """
     params = init_params
 
-    # Wrap the optax solver with jaxopt
-    optax_solver = OptaxSolver(
-        opt=optax_solver,
-        fun=loss,
-        has_aux=True,  # because the objective has aux output
-        maxiter=n_iter,
-    )
+    if isinstance(optimizer, GradientTransformation):
+        optimizer = OptaxSolver(
+            opt=optimizer,
+            fun=loss,
+            has_aux=True,
+            maxiter=n_iter,
+        )
+    elif optimizer == "lbfgs":
+        optimizer = LBFGS(fun=loss, has_aux=True, maxiter=n_iter)
+    # else, we trust that the user has given a valid jaxopt optimizer
 
     if param_data is not None:
         if (
@@ -134,7 +146,7 @@ def solve(
         batch = data.get_batch()
         if param_data is not None:
             batch = append_param_batch(batch, param_data.get_batch())
-        opt_state = optax_solver.init_state(params, batch=batch)
+        opt_state = optimizer.init_state(params, batch=batch)
 
     curr_seq = 0
     if seq2seq is not None:
@@ -183,7 +195,7 @@ def solve(
         batch = carry["data"].get_batch()
         if carry["param_data"] is not None:
             batch = append_param_batch(batch, carry["param_data"].get_batch())
-        carry["params"], carry["opt_state"] = optax_solver.update(
+        carry["params"], carry["opt_state"] = optimizer.update(
             params=carry["params"], state=carry["state"], batch=batch
         )
 
@@ -197,6 +209,18 @@ def solve(
 
         total_loss_val, loss_terms = loss(carry["params"], batch)
 
+        # Print loss during optimization
+        _ = jax.lax.cond(
+            i % print_loss_every == 0,
+            lambda _: jax.debug.print(
+                "Iteration {i}: loss value = " "{total_loss_val}",
+                i=i,
+                total_loss_val=total_loss_val,
+            ),
+            lambda _: None,
+            (None,),
+        )
+
         # optionnal seq2seq
         if seq2seq is not None:
             carry = _seq2seq_triggerer(
@@ -249,6 +273,12 @@ def solve(
         jnp.arange(n_iter),
     )
 
+    jax.debug.print(
+        "Iteration {i}: loss value = " "{total_loss_val}",
+        i=n_iter,
+        total_loss_val=accu[-1][-1],
+    )
+
     params = res["params"]
     last_non_nan_params = res["last_non_nan_params"]
     opt_state = res["state"]

{jinns-0.3.2.dist-info → jinns-0.3.3.dist-info}/METADATA

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

@@ -12,11 +12,11 @@ jinns/loss/_operators.py,sha256=nza4gfBn3Ppx7BN9Eben7KXPZ4cDQDyl6pfC8sHnAlI,4380
 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=XaY8I9dYGhFArB5GFWjcoe1D1nQOoXxpsUEPhrqaW0U,8896
+jinns/solver/_solve.py,sha256=pCctVnrkfOptGd1wH2w9jDdFiOMmNqgatqaoYtyE_LM,10071
 jinns/utils/__init__.py,sha256=-jDlwCjyEzWweswKdwLal3OhaUU3FVzK_Ge2S-7KHXs,149
 jinns/utils/_utils.py,sha256=bQm6z_xPKJj9BMCr2tXc44IA8JyGNN-PR5LNRhZ1fD8,20085
-jinns-0.3.2.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-0.3.2.dist-info/METADATA,sha256=YAstgf9edypib3O8bdTN7VUG5e811vhD3IaLjfuK7JM,1821
-jinns-0.3.2.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-jinns-0.3.2.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-0.3.2.dist-info/RECORD,,
+jinns-0.3.3.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
+jinns-0.3.3.dist-info/METADATA,sha256=GL7VeN4nQ5QuFqnKlnvagQiES6AyWsw9eeRVxnbu4vc,1821
+jinns-0.3.3.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+jinns-0.3.3.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
+jinns-0.3.3.dist-info/RECORD,,