Data API Reference

This subpackage contains data generation functions for testing and examples.

Data Subpackage

data

Implementations of data generation and handling functions.

KS_1D

KS_1D(tN: Float, u0: Array = None, dt: Float = 0.25, domain: tuple[Float, Float] = (0, 22), Nx: int = 64) -> tuple[Float, Float]

Solve the Kuramoto-Sivashinsky equation in 1D with periodic boundary conditions.

The KS PDE solved is: u_t + u*u_x + u_xx + u_xxxx = 0

The solver uses a fixed time-step ETDRK4 (Kassam & Trefethen 2005) method for handling the stiffness of the PDE. Dealiasing (2/3 rule) is applied.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the PDE to.	required
u0	ndarray	The initial condition for the PDE (shape (Nx,)). Default is None, which initializes u0 to sin((32/domain[1])*pi*x).	None
dt	float	The time step size for the solution. Default is 0.25.	0.25
domain	tuple[float, float]	The spatial domain (x_min, x_max). Default is (0, 22).	(0, 22)
Nx	int	The number of spatial grid points. Default is 64.	64

Returns:

Name	Type	Description
U	ndarray	The solution array with shape (Nt, Nx+1), where Nt is the number of time steps. Includes the periodic boundary point.
t	ndarray	The time vector corresponding to the solution steps.

colpitts

colpitts(tN: Float, dt: Float, u0: Float[Array, 3] = (1.0, -1.0, 1.0), alpha: Float = 5.0, gamma: Float = 0.0797, q: Float = 0.6898, eta: Float = 6.2723, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Colpitts oscillator system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is (1.0, -1.0, 1.0).	(1.0, -1.0, 1.0)
alpha	float	The alpha parameter for the Colpitts oscillator system. Default is 5.0 (Platt, 2020).	5.0
gamma	float	The gamma parameter for the Colpitts oscillator system. Default is 0.0797 (Platt, 2020).	0.0797
q	float	The q parameter for the Colpitts oscillator system. Default is 0.6898 (Platt, 2020).	0.6898
eta	float	The eta parameter for the Colpitts oscillator system. Default is 6.2723 (Platt, 2020).	6.2723
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, 3), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

double_pendulum

double_pendulum(tN: Float, dt: Float, u0: Float[Array, 4] = (pi / 4, -1.0, pi / 2, 1.0), m1: Float = 1.0, m2: Float = 1.0, L1: Float = 1.0, L2: Float = 1.0, g: Float = 9.81, damping: Float = 0.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the equations of motion for a damped double pendulum.

The state u is represented as a 4-tuple (theta1, omega1, theta2, omega2) where: - theta1 is the angle of the first pendulum from vertical (in radians). - omega1 is the angular velocity of the first pendulum (in radians/s). - theta2 is the angle of the second pendulum from vertical (in radians). - omega2 is the angular velocity of the second pendulum (in radians/s).

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is (jnp.pi/4, -1.0, jnp.pi/2, 1.0).	(pi / 4, -1.0, pi / 2, 1.0)
m1	float	The mass of the first pendulum bob. Default is 1.0.	1.0
m2	float	The mass of the second pendulum bob. Default is 1.0.	1.0
L1	float	The length of the first pendulum rod. Default is 1.0.	1.0
L2	float	The length of the second pendulum rod. Default is 1.0.	1.0
g	float	The acceleration due to gravity. Default is 9.81.	9.81
damping	float	The damping coefficient for the pendulum system. Default is 0.0 (no damping).	0.0
diffeqsolve_kwargs	(dict, optional)	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, 4), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

hyper_lorenz63

hyper_lorenz63(tN: Float, dt: Float, u0: Float[Array, 4] = (-10.0, 6.0, 0.0, 10.0), a: Float = 10.0, b: Float = 28.0, c: Float = 8.0 / 3.0, d: Float = -1.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Hyper-Lorenz 63 system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is (-10.0, 6.0, 0.0, 10.0).	(-10.0, 6.0, 0.0, 10.0)
a	float	The a parameter for the Hyper-Lorenz system. Default is 10.0.	10.0
b	float	The b parameter for the Hyper-Lorenz system. Default is 28.0.	28.0
c	float	The c parameter for the Hyper-Lorenz system. Default is 8.0/3.0.	8.0 / 3.0
d	float	The d parameter for the Hyper-Lorenz system. Default is -1.0.	-1.0
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, 4), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

hyper_xu

hyper_xu(tN: Float, dt: Float, u0: Float[Array, 4] = (-2.0, -1.0, -2.0, -10.0), a: Float = 10.0, b: Float = 40.0, c: Float = 2.5, d: Float = 2.0, e: Float = 16.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Hyper-Xu system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is (-2.0, -1.0, -2.0, -10.0).	(-2.0, -1.0, -2.0, -10.0)
a	float	The a parameter for the Hyper-Xu system. Default is 10.0.	10.0
b	float	The b parameter for the Hyper-Xu system. Default is 40.0.	40.0
c	float	The c parameter for the Hyper-Xu system. Default is 2.5.	2.5
d	float	The d parameter for the Hyper-Xu system. Default is 2.0.	2.0
e	float	The e parameter for the Hyper-Xu system. Default is 16.0.	16.0
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, 4), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

lorenz63

lorenz63(tN: Float, dt: Float, u0: Float[Array, 3] = (-10.0, 1.0, 10.0), rho: Float = 28.0, sigma: Float = 10.0, beta: Float = 8.0 / 3.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Lorenz 63 system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is (-10, 1, 10).	(-10.0, 1.0, 10.0)
rho	float	The rho parameter for the Lorenz system. Default is 28.0.	28.0
sigma	float	The sigma parameter for the Lorenz system. Default is 10.0.	10.0
beta	float	The beta parameter for the Lorenz system. Default is 8.0/3.0.	8.0 / 3.0
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, 3), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

lorenz96

lorenz96(tN: Float, dt: Float, u0: Array = None, N: Float = 10, F: Float = 8.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Lorenz 96 system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is None, which initializes y0 to jnp.sin(jnp.arange(N)).	None
N	int	The number of variables in the Lorenz 96 system. Default is 10.	10
F	float	The forcing parameter for the Lorenz 96 system. Default is 8.0.	8.0
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, N), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

rossler

rossler(tN: Float, dt: Float, u0: Float[Array, 3] = (1.0, 1.0, 1.0), a: Float = 0.1, b: Float = 0.1, c: Float = 14.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Rossler system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is (1.0, 1.0, 1.0).	(1.0, 1.0, 1.0)
a	float	The a parameter for the Rössler system. Default is 0.1.	0.1
b	float	The b parameter for the Rössler system. Default is 0.1.	0.1
c	float	The c parameter for the Rössler system. Default is 14.0.	14.0
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, 3), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

sakaraya

sakaraya(tN: Float, dt: Float, u0: Float[Array, 3] = (-2.8976045, 3.8877978, 3.07465), a: Float = 1.0, b: Float = 1.0, m: Float = 1.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Sakaraya system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is (-2.8976045, 3.8877978, 3.07465).	(-2.8976045, 3.8877978, 3.07465)
a	float	The a parameter for the Sakaraya system. Default is 1.0.	1.0
b	float	The b parameter for the Sakaraya system. Default is 1.0.	1.0
m	float	The m parameter for the Sakaraya system. Default is 1.0.	1.0
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, 3), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

Lorenz63 System

lorenz63

lorenz63(tN: Float, dt: Float, u0: Float[Array, 3] = (-10.0, 1.0, 10.0), rho: Float = 28.0, sigma: Float = 10.0, beta: Float = 8.0 / 3.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Lorenz 63 system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is (-10, 1, 10).	(-10.0, 1.0, 10.0)
rho	float	The rho parameter for the Lorenz system. Default is 28.0.	28.0
sigma	float	The sigma parameter for the Lorenz system. Default is 10.0.	10.0
beta	float	The beta parameter for the Lorenz system. Default is 8.0/3.0.	8.0 / 3.0
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, 3), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

Source code in src/orc/data/integrators.py

def lorenz63(
    tN: Float,
    dt: Float,
    u0: Float[Array, "3"] = (-10.0, 1.0, 10.0),
    rho: Float = 28.0,
    sigma: Float = 10.0,
    beta: Float = 8.0 / 3.0,
    **diffeqsolve_kwargs,
) -> tuple[Float, Float]:
    """Solve the Lorenz 63 system of ODEs.

    Parameters
    ----------
    tN : float
        The final time to solve the ODEs to.
    dt : float
        The time step size for the interpolated solution. Will be overridden if
        `diffeqsolve_kwargs` contains a saveat argument with a different time grid.
    u0 : jnp.ndarray, optional
        The initial conditions for the ODEs. Default is (-10, 1, 10).
    rho : float, optional
        The rho parameter for the Lorenz system. Default is 28.0.
    sigma : float, optional
        The sigma parameter for the Lorenz system. Default is 10.0.
    beta : float, optional
        The beta parameter for the Lorenz system. Default is 8.0/3.0.
    diffeqsolve_kwargs : dict, optional
        Additional keyword arguments to pass to the `diffrax.diffeqsolve` function.
        Default solver is `diffrax.Tsit5()`, saveat is set to a grid of times from 0 to
        tN with step size dt, and stepsize_controller is set to
        `diffrax.PIDController(rtol=1e-7, atol=1e-9)`. dt0 is set to dt. max_steps is
        set to None.

    Returns
    -------
    us : jnp.ndarray
        The solution array with shape (Nt, 3), where Nt is the number of time steps.
    ts : jnp.ndarray
        The time vector corresponding to the solution steps.
    """
    # set kwarg defaults
    diffeqsolve_kwargs.setdefault("solver", diffrax.Tsit5())
    diffeqsolve_kwargs.setdefault("saveat", diffrax.SaveAt(ts=jnp.arange(0, tN, dt)))
    diffeqsolve_kwargs.setdefault(
        "stepsize_controller", diffrax.PIDController(rtol=1e-7, atol=1e-9)
    )
    diffeqsolve_kwargs.setdefault("dt0", dt)
    diffeqsolve_kwargs.setdefault("max_steps", None)

    # solve
    u0 = jnp.array(u0)
    term = diffrax.ODETerm(_lorenz63_f)
    args = (rho, sigma, beta)
    sol = diffrax.diffeqsolve(term, t0=0, t1=tN, y0=u0, args=args, **diffeqsolve_kwargs)
    us = sol.ys
    return us, sol.ts

Kuramoto-Sivashinsky Equation

KS_1D

KS_1D(tN: Float, u0: Array = None, dt: Float = 0.25, domain: tuple[Float, Float] = (0, 22), Nx: int = 64) -> tuple[Float, Float]

Solve the Kuramoto-Sivashinsky equation in 1D with periodic boundary conditions.

The KS PDE solved is: u_t + u*u_x + u_xx + u_xxxx = 0

The solver uses a fixed time-step ETDRK4 (Kassam & Trefethen 2005) method for handling the stiffness of the PDE. Dealiasing (2/3 rule) is applied.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the PDE to.	required
u0	ndarray	The initial condition for the PDE (shape (Nx,)). Default is None, which initializes u0 to sin((32/domain[1])*pi*x).	None
dt	float	The time step size for the solution. Default is 0.25.	0.25
domain	tuple[float, float]	The spatial domain (x_min, x_max). Default is (0, 22).	(0, 22)
Nx	int	The number of spatial grid points. Default is 64.	64

Returns:

Name	Type	Description
U	ndarray	The solution array with shape (Nt, Nx+1), where Nt is the number of time steps. Includes the periodic boundary point.
t	ndarray	The time vector corresponding to the solution steps.

Source code in src/orc/data/integrators.py

@functools.partial(jax.jit, static_argnames=["tN", "dt", "Nx"])
def KS_1D(
    tN: Float,
    u0: Array = None,
    dt: Float = 0.25,
    domain: tuple[Float, Float] = (0, 22),
    Nx: int = 64,
) -> tuple[Float, Float]:
    """Solve the Kuramoto-Sivashinsky equation in 1D with periodic boundary conditions.

    The KS PDE solved is:
        u_t + u*u_x + u_xx + u_xxxx = 0

    The solver uses a fixed time-step ETDRK4 (Kassam & Trefethen 2005) method
    for handling the stiffness of the PDE. Dealiasing (2/3 rule) is applied.

    Parameters
    ----------
    tN : float
        The final time to solve the PDE to.
    u0 : jnp.ndarray, optional
        The initial condition for the PDE (shape (Nx,)). Default is None, which
        initializes u0 to `sin((32/domain[1])*pi*x)`.
    dt : float, optional
        The time step size for the solution. Default is 0.25.
    domain : tuple[float, float], optional
        The spatial domain (x_min, x_max). Default is (0, 22).
    Nx : int, optional
        The number of spatial grid points. Default is 64.

    Returns
    -------
    U : jnp.ndarray
        The solution array with shape (Nt, Nx+1), where Nt is the number of time steps.
        Includes the periodic boundary point.
    t : jnp.ndarray
        The time vector corresponding to the solution steps.
    """
    # Setup spatial grid
    if u0 is None:
        x0 = jnp.linspace(domain[0], domain[1], Nx, endpoint=True)
        u0 = jnp.sin((32 / domain[1]) * jnp.pi * x0)  # TODO find a better default
    u0 = u0[:-1]
    Nx = Nx - 1  # remove duplicate periodic point
    x = jnp.linspace(domain[0], domain[1], Nx, endpoint=False)
    dx = x[1] - x[0]

    Nt = int(tN / dt)
    U = jnp.zeros((Nx, Nt))
    U = U.at[:, 0].set(u0)

    # Wavenumbers
    k = jnp.fft.fftfreq(Nx, d=dx) * 2 * jnp.pi
    k2 = k**2
    k4 = k**4
    L_op = k2 - k4

    # Dealiasing (2/3 rule)
    def dealias(f_hat):
        cutoff = Nx // 3
        f_hat = f_hat.at[cutoff:-cutoff].set(0)
        return f_hat

    # nonlinear operators on u and u_hat
    def N_op_u(u):
        return dealias(1j * k * jnp.fft.fft(-0.5 * u**2))

    def N_op_uhat(u_hat):
        return dealias(1j * k * jnp.fft.fft(-0.5 * jnp.real(jnp.fft.ifft(u_hat)) ** 2))

    # ETDRK4 coefficients (Kassam & Trefethen 2005)
    E1 = jnp.exp(L_op * dt)
    E2 = jnp.exp(L_op * dt / 2)
    M = 16
    r = jnp.exp(1j * jnp.pi * (jnp.arange(1, M + 1) - 0.5) / M)
    LR = dt * jnp.column_stack([L_op] * M) + jnp.vstack([r] * Nx)
    Q = dt * jnp.mean((jnp.exp(LR / 2) - 1) / LR, axis=1)
    f1 = dt * jnp.mean((-4 - LR + jnp.exp(LR) * (4 - 3 * LR + LR**2)) / LR**3, axis=1)
    f2 = dt * jnp.mean((2 + LR + jnp.exp(LR) * (-2 + LR)) / LR**3, axis=1)
    f3 = dt * jnp.mean((-4 - 3 * LR - LR**2 + jnp.exp(LR) * (4 - LR)) / LR**3, axis=1)

    def _KS_ETDRK4_step(carry, _):
        u, E1, E2, Q, f1, f2, f3 = carry

        u_hat = jnp.fft.fft(u)

        a = E2 * u_hat + Q * N_op_u(u)
        b = E2 * u_hat + Q * N_op_uhat(a)
        c = E2 * a + Q * (2 * N_op_uhat(b) - N_op_u(u))

        u_hat = (
            E1 * u_hat
            + f1 * N_op_u(u)
            + f2 * (N_op_uhat(a) + N_op_uhat(b))
            + f3 * N_op_uhat(c)
        )

        # Enforce conservation by zeroing the mean mode
        u_hat = u_hat.at[0].set(0.0)

        u_next = jnp.real(jnp.fft.ifft(u_hat, n=Nx))
        carry_next = (u_next, E1, E2, Q, f1, f2, f3)
        return carry_next, u_next

    _, u_vals = jax.lax.scan(
        _KS_ETDRK4_step, (u0, E1, E2, Q, f1, f2, f3), length=Nt - 1
    )

    # add back in the initial point and boundary points
    U = jnp.concatenate([u0[None, :], u_vals], axis=0)
    U = jnp.concatenate((U, U[:, 0:1]), axis=1)

    # create time vector for output
    t = jnp.arange(0, tN, dt)

    return U, t

Rössler System

rossler

rossler(tN: Float, dt: Float, u0: Float[Array, 3] = (1.0, 1.0, 1.0), a: Float = 0.1, b: Float = 0.1, c: Float = 14.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Rossler system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is (1.0, 1.0, 1.0).	(1.0, 1.0, 1.0)
a	float	The a parameter for the Rössler system. Default is 0.1.	0.1
b	float	The b parameter for the Rössler system. Default is 0.1.	0.1
c	float	The c parameter for the Rössler system. Default is 14.0.	14.0
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, 3), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

Source code in src/orc/data/integrators.py

def rossler(
    tN: Float,
    dt: Float,
    u0: Float[Array, "3"] = (1.0, 1.0, 1.0),
    a: Float = 0.1,
    b: Float = 0.1,
    c: Float = 14.0,
    **diffeqsolve_kwargs,
) -> tuple[Float, Float]:
    """Solve the Rossler system of ODEs.

    Parameters
    ----------
    tN : float
        The final time to solve the ODEs to.
    dt : float
        The time step size for the interpolated solution. Will be overridden if
        `diffeqsolve_kwargs` contains a saveat argument with a different time grid.
    u0 : jnp.ndarray, optional
        The initial conditions for the ODEs. Default is (1.0, 1.0, 1.0).
    a : float, optional
        The a parameter for the Rössler system. Default is 0.1.
    b : float, optional
        The b parameter for the Rössler system. Default is 0.1.
    c : float, optional
        The c parameter for the Rössler system. Default is 14.0.
    diffeqsolve_kwargs : dict, optional
        Additional keyword arguments to pass to the `diffrax.diffeqsolve` function.
        Default solver is `diffrax.Tsit5()`, saveat is set to a grid of times from 0 to
        tN with step size dt, and stepsize_controller is set to
        `diffrax.PIDController(rtol=1e-7, atol=1e-9)`. dt0 is set to dt. max_steps is
        set to None.

    Returns
    -------
    us : jnp.ndarray
        The solution array with shape (Nt, 3), where Nt is the number of time steps.
    ts : jnp.ndarray
        The time vector corresponding to the solution steps.
    """
    # set kwarg defaults
    diffeqsolve_kwargs.setdefault("solver", diffrax.Tsit5())
    diffeqsolve_kwargs.setdefault("saveat", diffrax.SaveAt(ts=jnp.arange(0, tN, dt)))
    diffeqsolve_kwargs.setdefault(
        "stepsize_controller", diffrax.PIDController(rtol=1e-7, atol=1e-9)
    )
    diffeqsolve_kwargs.setdefault("dt0", dt)
    diffeqsolve_kwargs.setdefault("max_steps", None)

    # solve
    u0 = jnp.array(u0)
    term = diffrax.ODETerm(_rossler_f)
    args = (a, b, c)
    sol = diffrax.diffeqsolve(term, t0=0, t1=tN, y0=u0, args=args, **diffeqsolve_kwargs)
    us = sol.ys
    return us, sol.ts

Lorenz96 System

lorenz96

lorenz96(tN: Float, dt: Float, u0: Array = None, N: Float = 10, F: Float = 8.0, **diffeqsolve_kwargs) -> tuple[Float, Float]

Solve the Lorenz 96 system of ODEs.

Parameters:

Name	Type	Description	Default
tN	float	The final time to solve the ODEs to.	required
dt	float	The time step size for the interpolated solution. Will be overridden if diffeqsolve_kwargs contains a saveat argument with a different time grid.	required
u0	ndarray	The initial conditions for the ODEs. Default is None, which initializes y0 to jnp.sin(jnp.arange(N)).	None
N	int	The number of variables in the Lorenz 96 system. Default is 10.	10
F	float	The forcing parameter for the Lorenz 96 system. Default is 8.0.	8.0
diffeqsolve_kwargs	dict	Additional keyword arguments to pass to the diffrax.diffeqsolve function. Default solver is diffrax.Tsit5(), saveat is set to a grid of times from 0 to tN with step size dt, and stepsize_controller is set to diffrax.PIDController(rtol=1e-7, atol=1e-9). dt0 is set to dt. max_steps is set to None.	{}

Returns:

Name	Type	Description
us	ndarray	The solution array with shape (Nt, N), where Nt is the number of time steps.
ts	ndarray	The time vector corresponding to the solution steps.

Source code in src/orc/data/integrators.py

def lorenz96(
    tN: Float,
    dt: Float,
    u0: Array = None,
    N: Float = 10,
    F: Float = 8.0,
    **diffeqsolve_kwargs,
) -> tuple[Float, Float]:
    """Solve the Lorenz 96 system of ODEs.

    Parameters
    ----------
    tN : float
        The final time to solve the ODEs to.
    dt : float
        The time step size for the interpolated solution. Will be overridden if
        `diffeqsolve_kwargs` contains a saveat argument with a different time grid.
    u0 : jnp.ndarray, optional
        The initial conditions for the ODEs. Default is None, which initializes y0
        to `jnp.sin(jnp.arange(N))`.
    N : int, optional
        The number of variables in the Lorenz 96 system. Default is 10.
    F : float, optional
        The forcing parameter for the Lorenz 96 system. Default is 8.0.
    diffeqsolve_kwargs : dict, optional
        Additional keyword arguments to pass to the `diffrax.diffeqsolve` function.
        Default solver is `diffrax.Tsit5()`, saveat is set to a grid of times from 0 to
        tN with step size dt, and stepsize_controller is set to
        `diffrax.PIDController(rtol=1e-7, atol=1e-9)`. dt0 is set to dt. max_steps is
        set to None.

    Returns
    -------
    us : jnp.ndarray
        The solution array with shape (Nt, N), where Nt is the number of time steps.
    ts : jnp.ndarray
        The time vector corresponding to the solution steps.
    """
    # set kwarg defaults
    diffeqsolve_kwargs.setdefault("solver", diffrax.Tsit5())
    diffeqsolve_kwargs.setdefault("saveat", diffrax.SaveAt(ts=jnp.arange(0, tN, dt)))
    diffeqsolve_kwargs.setdefault(
        "stepsize_controller", diffrax.PIDController(rtol=1e-7, atol=1e-9)
    )
    diffeqsolve_kwargs.setdefault("dt0", dt)
    diffeqsolve_kwargs.setdefault("max_steps", None)

    # solve
    if u0 is None:
        u0 = jnp.sin(jnp.arange(N))
    u0 = jnp.array(u0)
    term = diffrax.ODETerm(_lorenz96_f)
    args = (N, F)
    sol = diffrax.diffeqsolve(term, t0=0, t1=tN, y0=u0, args=args, **diffeqsolve_kwargs)
    us = sol.ys
    return us, sol.ts