jinns 0.9.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

jinns/solver/_seq2seq.py

@@ -1,157 +0,0 @@
-"""
-Implements Seq2Seq training as described in “Characterizing possible
-failure modes in physics-informed neural networks”, A. S. Krishnapriyan,
-NeurIPS 2021.
-
-**Note:** we do not change tmin, we only let the interval grow longer.
-Indeed we noticed some unlearning happening.
-
-**Note:** using seq2seq might create some instability in training when
-interval changes. Some of this instability comes from the fact that Tmax in
-the dynamic loss rescaling must be the true and final (and potentially large and
-unstable) one from the beginning if we want to be able to catch the real dynamic.
-However it does offer some better results for learning on long time intervals.
-
-**Note:** As this is experimental some changes in the future might be:
-    - to dig deeper and try to attenuate the instability
-    - to try to attenuate the discrepancy with the real dynamic when we
-        also change Tmax in dynamic loss (this requires to treat the dynamic
-        loss as a dynamic attribute of a Loss class...).
-    - to investigate Tmax as input of the PINN
-
-"""
-
-import jax
-from jax import jit
-import jax.numpy as jnp
-from jinns.data._DataGenerators import (
-    DataGeneratorODE,
-    _reset_batch_idx_and_permute,
-)
-from jinns.loss._LossODE import SystemLossODE, LossODE
-from jinns.loss._LossPDE import LossPDENonStatio, LossPDEStatio, SystemLossPDE
-
-
-# @partial(jit, static_argnames=["_update_seq2seq_true", "_update_seq2seq_false"])
-@jit
-def trigger_seq2seq(
-    i,
-    loss,
-    params,
-    data,
-    opt_state,
-    curr_seq,
-    seq2seq,
-    # _update_seq2seq_true,
-    # _update_seq2seq_false,
-):
-    if seq2seq is not None:
-        curr_seq, loss, data, opt_state = jax.lax.cond(
-            curr_seq < jnp.sum(seq2seq["iter_steps"] < i),
-            # check if we fall in another time interval
-            _update_seq2seq_SystemLossODE,
-            # only SystemLoss are handled for now
-            _update_seq2seq_false,
-            (
-                loss,
-                seq2seq,
-                data,
-                params,
-                curr_seq,
-                opt_state,
-            ),
-        )
-    else:
-        # Do nothing if no seq2seq
-        curr_seq = -1
-
-    return loss, params, data, opt_state, curr_seq, seq2seq
-
-
-def initialize_seq2seq(loss, data, seq2seq, opt_state):
-    """
-    Helper function to set up Seq2Seq before going into scan or for loop
-    in `jinns.solve()`.
-    """
-    if isinstance(loss, SystemLossODE) and isinstance(data, DataGeneratorODE):
-        curr_seq = 0
-        # Note that boundaries for the first PINN are OK
-        # set new boundaries for the batch generator
-        data.tmax = seq2seq["time_steps"][curr_seq]
-
-        jax.debug.print(
-            "# -- Begin training on time segment [{tmin}, {tk}]",
-            tmin=data.tmin,
-            tk=seq2seq["time_steps"][curr_seq],
-        )
-        # and do not forget to regenerate the data
-        data.curr_omega_idx = 0
-        data.generate_time_data()
-        data._key, data.times, _ = _reset_batch_idx_and_permute(
-            (data._key, data.times, data.curr_omega_idx, None, data.p_times)
-        )
-        opt_state.hyperparams["learning_rate"] = seq2seq["learning_rate"][curr_seq]
-
-    elif isinstance(loss, (LossPDENonStatio, LossPDEStatio, SystemLossPDE)):
-        raise RuntimeError("Not implemented")
-
-    return data, opt_state
-
-
-def _update_seq2seq_SystemLossODE(operands):
-    """
-    Make all the necessary updates for a SystemLossODE in seq2seq learning mode
-
-    Parameters
-    ----------
-    operands
-        A tuple which comprises.
-
-        loss
-            A loss object (e.g. a LossODE, SystemLossODE, LossPDEStatio [...]
-            object). It must be jittable (e.g. implements via a pytree
-            registration)
-        seq2seq
-            A dictionary with keys 'times_steps'
-            and 'iter_steps' which mush have same length. The first represents
-            the time steps which represents the different time interval upon
-            which we perform the incremental learning. The second represents
-            the number of iteration we perform in each time interval.
-        data
-            A DataGenerator object which implements a `get_batch()`
-            method which returns a 3-tuple with (omega_grid, omega_border, time grid).
-            It must be jittable (e.g. implements via a pytree
-            registration)
-        params
-            The dictionary of parameters of the model.
-            Typically, it is a dictionary of
-            dictionaries: `eq_params` and `nn_params``, respectively the
-            differential equation parameters and the neural network parameter
-        curr_seq
-            A integer which represents which sequence we currently are in
-    """
-    loss, seq2seq, data, params, curr_seq, opt_state = operands
-    curr_seq += 1
-
-    jax.debug.print(
-        "# -- Entering training on time segment [{tmin}, {tk}]",
-        tmin=data.tmin,
-        tk=seq2seq["time_steps"][curr_seq],
-    )
-
-    # set new boundaries for the batch generator
-    data.tmax = seq2seq["time_steps"][curr_seq]
-    # and do not forget to regenerate the data
-    data.curr_omega_idx = 0
-    data.generate_time_data()
-    data._key, data.times, _ = _reset_batch_idx_and_permute(
-        (data._key, data.times, data.curr_omega_idx, None, data.p_times)
-    )
-
-    opt_state.hyperparams["learning_rate"] = seq2seq["learning_rate"][curr_seq]
-    return curr_seq, loss, data, opt_state
-
-
-def _update_seq2seq_false(operands):
-    # basically returns (curr_seq, loss, data, opt_state) in this order
-    return (operands[-2], operands[0], operands[2], operands[-1])

jinns/utils/_optim.py

@@ -1,147 +0,0 @@
-"""
-Implements utility functions for optimization
-"""
-
-import optax
-import jax
-import jax.numpy as jnp
-
-
-def alternate_optimizer(list_first_params, list_second_params, n_iter, evry, tx1, tx2):
-    """
-    Alternatively optimize on two sets of parameters for equal number of steps
-
-    Parameters
-    ----------
-    list_first_params
-        The first set of parameter to optimize on. A list of leaves from the `params` dict
-    list_second_params
-        The second set of parameter to optimize on. A list of leaves from the `params` dict
-    n_iter
-        The total number of iterations
-    evry
-        The number of iterations we spend optimizing a set of parameters before switching
-    tx1
-        An optax optimizer of the set of parameters 1
-    tx2
-        An optax optimizer of the set of parameters 1
-
-    Returns
-    -------
-    tx
-        An optax optimizer object
-    """
-
-    def map_nested_fn(fn):
-        """Recursively apply `fn` to the key-value pairs of a nested dict"""
-
-        def map_fn(nested_dict):
-            return {
-                k: (map_fn(v) if isinstance(v, dict) else fn(k, v))
-                for k, v in nested_dict.items()
-            }
-
-        return map_fn
-
-    label_fn = map_nested_fn(lambda k, _: k)
-
-    def should_update_1(step):
-        return jax.tree_util.tree_reduce(
-            lambda x, y: jnp.logical_or(x, y),
-            [
-                jnp.logical_and((step > i * evry), (step < (i + 1) * evry))
-                for i in range(1, n_iter // evry, 2)
-            ],
-        )
-
-    def should_update_2(step):
-        return jax.tree_util.tree_reduce(
-            lambda x, y: jnp.logical_or(x, y),
-            [
-                jnp.logical_and((step > i * evry), (step < (i + 1) * evry))
-                for i in range(
-                    0, n_iter // evry, 2
-                )  # starts at 0 since this one is blocked first
-            ],
-        )
-
-    first_adam = optax.chain(
-        tx1,
-        optax.maybe_update(
-            optax.scale(0.0), should_update_1
-        ),  # We add an update (a GradientTransform if should_update is True) i.e. we mult the update by 0.
-        # not to take a step
-    )
-    second_adam = optax.chain(
-        tx2,
-        optax.maybe_update(
-            optax.scale(0.0), should_update_2
-        ),  # We add an update (a GradientTransform if should_update is True) i.e. we mult the update by 0.
-        # not to take a step
-    )
-
-    return optax.multi_transform(
-        {k: first_adam for k in list_first_params}
-        | {
-            k: second_adam for k in list_second_params
-        },  # those gradient transforms must correspond to leaves of parameter pytree
-        label_fn,
-    )
-
-
-def delayed_optimizer(list_first_params, list_second_params, delay_steps, tx1, tx2):
-    """
-    Optimize on two sets of parameters, the optimization on the second set of
-    parameters start after `delay_steps` of freezing
-
-    Parameters
-    ----------
-    list_first_params
-        The first set of parameter to optimize on. A list of leaves from the `params` dict
-    list_second_params
-        The second set of parameter to optimize on. A list of leaves from the `params` dict
-    delay_steps
-        The number of steps we wait before starting the optimization on the
-        second set of parameters
-    tx1
-        An optax optimizer of the set of parameters 1
-    tx2
-        An optax optimizer of the set of parameters 1
-
-    Returns
-    -------
-    tx
-        An optax optimizer object
-    """
-
-    def map_nested_fn(fn):
-        """Recursively apply `fn` to the key-value pairs of a nested dict"""
-
-        def map_fn(nested_dict):
-            return {
-                k: (map_fn(v) if isinstance(v, dict) else fn(k, v))
-                for k, v in nested_dict.items()
-            }
-
-        return map_fn
-
-    label_fn = map_nested_fn(lambda k, _: k)
-
-    def should_update_2(step):
-        return step < delay_steps
-
-    delayed_tx2 = optax.chain(
-        tx2,
-        optax.maybe_update(
-            optax.scale(0.0), should_update_2
-        ),  # We add an update (a GradientTransform if should_update is True) i.e. we mult the update by 0.
-        # not to take a step
-    )
-
-    return optax.multi_transform(
-        {k: tx1 for k in list_first_params}
-        | {
-            k: delayed_tx2 for k in list_second_params
-        },  # those gradient transforms must correspond to leaves of parameter pytree
-        label_fn,
-    )