Temporal Evolution

Utilities to autoregressively evaluate steppers.

exponax.repeat

repeat(
    stepper_fn: Union[
        Callable[[PyTree], PyTree],
        Callable[[PyTree, PyTree], PyTree],
    ],
    n: int,
    *,
    takes_aux: bool = False,
    constant_aux: bool = True
)

Transform a stepper function into a function that autoregressively (i.e., recursively applied to its own output) applies the stepper n times and returns the final state.

Based on takes_aux, the stepper function is either fully automomous, just mapping state to state, or takes an additional auxiliary input. This can be a force/control or additional metadata (like physical parameters, or time for non-autonomous systems).

Args: - stepper_fn: The time stepper to transform. If takes_aux = False (default), expected signature is u_next = stepper_fn(u), else u_next = stepper_fn(u, aux). u and u_next need to be PyTrees of identical structure, in the easiest case just arrays of same shape. - n: The number of times to apply the stepper. - takes_aux: Whether the stepper function takes an additional PyTree as second argument. - constant_aux: Whether the auxiliary input is constant over the trajectory. If True, the auxiliary input is repeated n times, otherwise the leading axis in the PyTree arrays has to be of length n.

Returns: - repeated_stepper_fn: A function that takes an initial condition u_0 and an auxiliary input aux (if takes_aux = True) and produces the final state by autoregressively applying the stepper n times. Returns a PyTree of the same structure as the initial condition.

Source code in exponax/_utils.py

def repeat(
    stepper_fn: Union[Callable[[PyTree], PyTree], Callable[[PyTree, PyTree], PyTree]],
    n: int,
    *,
    takes_aux: bool = False,
    constant_aux: bool = True,
):
    """
    Transform a stepper function into a function that autoregressively (i.e.,
    recursively applied to its own output) applies the stepper `n` times and
    returns the final state.

    Based on `takes_aux`, the stepper function is either fully automomous, just
    mapping state to state, or takes an additional auxiliary input. This can be
    a force/control or additional metadata (like physical parameters, or time
    for non-autonomous systems).

    Args:
        - `stepper_fn`: The time stepper to transform. If `takes_aux = False`
            (default), expected signature is `u_next = stepper_fn(u)`, else
            `u_next = stepper_fn(u, aux)`. `u` and `u_next` need to be PyTrees
            of identical structure, in the easiest case just arrays of same
            shape.
        - `n`: The number of times to apply the stepper.
        - `takes_aux`: Whether the stepper function takes an additional PyTree
            as second argument.
        - `constant_aux`: Whether the auxiliary input is constant over the
            trajectory. If `True`, the auxiliary input is repeated `n` times,
            otherwise the leading axis in the PyTree arrays has to be of length
            `n`.

    Returns:
        - `repeated_stepper_fn`: A function that takes an initial condition
            `u_0` and an auxiliary input `aux` (if `takes_aux = True`) and
            produces the final state by autoregressively applying the stepper
            `n` times. Returns a PyTree of the same structure as the initial
            condition.
    """

    if takes_aux:

        def scan_fn(u, aux):
            u_next = stepper_fn(u, aux)
            return u_next, None

        def repeated_stepper_fn(u_0, aux):
            if constant_aux:
                aux = jtu.tree_map(
                    lambda x: jnp.repeat(jnp.expand_dims(x, axis=0), n, axis=0), aux
                )

            final, _ = jax.lax.scan(scan_fn, u_0, aux, length=n)
            return final

        return repeated_stepper_fn

    else:

        def scan_fn(u, _):
            u_next = stepper_fn(u)
            return u_next, None

        def repeated_stepper_fn(u_0):
            final, _ = jax.lax.scan(scan_fn, u_0, None, length=n)
            return final

        return repeated_stepper_fn

---

exponax.rollout

rollout(
    stepper_fn: Union[
        Callable[[PyTree], PyTree],
        Callable[[PyTree, PyTree], PyTree],
    ],
    n: int,
    *,
    include_init: bool = False,
    takes_aux: bool = False,
    constant_aux: bool = True
)

Transform a stepper function into a function that autoregressively (i.e., recursively applied to its own output) produces a trajectory of length n.

Based on takes_aux, the stepper function is either fully automomous, just mapping state to state, or takes an additional auxiliary input. This can be a force/control or additional metadata (like physical parameters, or time for non-autonomous systems).

Args: - stepper_fn: The time stepper to transform. If takes_aux = False (default), expected signature is u_next = stepper_fn(u), else u_next = stepper_fn(u, aux). u and u_next need to be PyTrees of identical structure, in the easiest case just arrays of same shape. - n: The number of time steps to rollout the trajectory into the future. If include_init = False (default) produces the n steps into the future. - include_init: Whether to include the initial condition in the trajectory. If True, the arrays in the returning PyTree have shape (n + 1, ...), else (n, ...). Default: False. - takes_aux: Whether the stepper function takes an additional PyTree as second argument. - constant_aux: Whether the auxiliary input is constant over the trajectory. If True, the auxiliary input is repeated n times, otherwise the leading axis in the PyTree arrays has to be of length n.

Returns: - rollout_stepper_fn: A function that takes an initial condition u_0 and an auxiliary input aux (if takes_aux = True) and produces the trajectory by autoregressively applying the stepper n times. If include_init = True, the trajectory has shape (n + 1, ...), else (n, ...). Returns a PyTree of the same structure as the initial condition, but with an additional leading axis of length n.

Source code in exponax/_utils.py

def rollout(
    stepper_fn: Union[Callable[[PyTree], PyTree], Callable[[PyTree, PyTree], PyTree]],
    n: int,
    *,
    include_init: bool = False,
    takes_aux: bool = False,
    constant_aux: bool = True,
):
    """
    Transform a stepper function into a function that autoregressively (i.e.,
    recursively applied to its own output) produces a trajectory of length `n`.

    Based on `takes_aux`, the stepper function is either fully automomous, just
    mapping state to state, or takes an additional auxiliary input. This can be
    a force/control or additional metadata (like physical parameters, or time
    for non-autonomous systems).

    Args:
        - `stepper_fn`: The time stepper to transform. If `takes_aux = False`
            (default), expected signature is `u_next = stepper_fn(u)`, else
            `u_next = stepper_fn(u, aux)`. `u` and `u_next` need to be PyTrees
            of identical structure, in the easiest case just arrays of same
            shape.
        - `n`: The number of time steps to rollout the trajectory into the
            future. If `include_init = False` (default) produces the `n` steps
            into the future.
        - `include_init`: Whether to include the initial condition in the
            trajectory. If `True`, the arrays in the returning PyTree have shape
            `(n + 1, ...)`, else `(n, ...)`. Default: `False`.
        - `takes_aux`: Whether the stepper function takes an additional PyTree
            as second argument.
        - `constant_aux`: Whether the auxiliary input is constant over the
            trajectory. If `True`, the auxiliary input is repeated `n` times,
            otherwise the leading axis in the PyTree arrays has to be of length
            `n`.

    Returns:
        - `rollout_stepper_fn`: A function that takes an initial condition `u_0`
            and an auxiliary input `aux` (if `takes_aux = True`) and produces
            the trajectory by autoregressively applying the stepper `n` times.
            If `include_init = True`, the trajectory has shape `(n + 1, ...)`,
            else `(n, ...)`. Returns a PyTree of the same structure as the
            initial condition, but with an additional leading axis of length
            `n`.
    """

    if takes_aux:

        def scan_fn(u, aux):
            u_next = stepper_fn(u, aux)
            return u_next, u_next

        def rollout_stepper_fn(u_0, aux):
            if constant_aux:
                aux = jtu.tree_map(
                    lambda x: jnp.repeat(jnp.expand_dims(x, axis=0), n, axis=0), aux
                )

            _, trj = jax.lax.scan(scan_fn, u_0, aux, length=n)

            if include_init:
                trj_with_init = jtu.tree_map(
                    lambda init, history: jnp.concatenate(
                        [jnp.expand_dims(init, axis=0), history], axis=0
                    ),
                    u_0,
                    trj,
                )
                return trj_with_init
            else:
                return trj

        return rollout_stepper_fn

    else:

        def scan_fn(u, _):
            u_next = stepper_fn(u)
            return u_next, u_next

        def rollout_stepper_fn(u_0):
            _, trj = jax.lax.scan(scan_fn, u_0, None, length=n)

            if include_init:
                trj_with_init = jtu.tree_map(
                    lambda init, history: jnp.concatenate(
                        [jnp.expand_dims(init, axis=0), history], axis=0
                    ),
                    u_0,
                    trj,
                )
                return trj_with_init
            else:
                return trj

        return rollout_stepper_fn

---

exponax.stack_sub_trajectories

stack_sub_trajectories(
    trj: PyTree[Float[Array, "n_timesteps ..."]],
    sub_len: int,
) -> PyTree[Float[Array, "n_stacks sub_len ..."]]

Slice a trajectory into subtrajectories of length n and stack them together. Useful for rollout training neural operators with temporal mixing.

!!! Note that this function can produce very large arrays.

Arguments: - trj: The trajectory to slice. Expected shape: (n_timesteps, ...). - sub_len: The length of the subtrajectories. If you want to perform rollout training with k steps, note that n=k+1 to also have an initial condition in the subtrajectories.

Returns: - sub_trjs: The stacked subtrajectories. Expected shape: (n_stacks, n, ...). n_stacks is the number of subtrajectories stacked together, i.e., n_timesteps - n + 1.

Source code in exponax/_utils.py

def stack_sub_trajectories(
    trj: PyTree[Float[Array, "n_timesteps ..."]],
    sub_len: int,
) -> PyTree[Float[Array, "n_stacks sub_len ..."]]:
    """
    Slice a trajectory into subtrajectories of length `n` and stack them
    together. Useful for rollout training neural operators with temporal mixing.

    !!! Note that this function can produce very large arrays.

    **Arguments:**
        - `trj`: The trajectory to slice. Expected shape: `(n_timesteps, ...)`.
        - `sub_len`: The length of the subtrajectories. If you want to perform rollout
            training with k steps, note that `n=k+1` to also have an initial
            condition in the subtrajectories.

    **Returns:**
        - `sub_trjs`: The stacked subtrajectories. Expected shape: `(n_stacks, n, ...)`.
           `n_stacks` is the number of subtrajectories stacked together, i.e.,
           `n_timesteps - n + 1`.
    """
    n_time_steps = [leaf.shape[0] for leaf in jtu.tree_leaves(trj)]

    if len(set(n_time_steps)) != 1:
        raise ValueError(
            "All arrays in trj must have the same number of time steps in the leading axis"
        )
    else:
        n_time_steps = n_time_steps[0]

    if sub_len > n_time_steps:
        raise ValueError(
            "n must be smaller than or equal to the number of time steps in trj"
        )

    n_sub_trjs = n_time_steps - sub_len + 1

    def scan_fn(_, i):
        sliced = jtu.tree_map(
            lambda leaf: jax.lax.dynamic_slice_in_dim(
                leaf,
                start_index=i,
                slice_size=sub_len,
                axis=0,
            ),
            trj,
        )
        return _, sliced

    _, sub_trjs = jax.lax.scan(scan_fn, None, jnp.arange(n_sub_trjs))

    return sub_trjs