Capturing Extra Output with CyRK’s ODE Solvers

It can be advantageous to capture additional output parameters during the integration of systems of ordinary differential equations. Particularly when the differential equations in question are computationally expensive. For example, consider the following simplified differential equation.


def diffeq(t, y, a, b):
    
    # Some complicated function of time and y_1; this may be calling many other inner functions.
    f = calculate_f(t, y[1])
    
    # Calculate dy/dt
    dy_0 = f * a * y[0] + b * y[1]
    dy_1 = b * y[1] 
    
    return np.asarray([dy_0, dy_1])

Using CyRK it is easy to integrate this function to find \(y_0\) and \(y_1\) as a function of time. But suppose that in our final analysis we also want to see how \(f\) is changing with time and with \(y\)’s. We could extract it by using y_0 and dy_0 (which we would have to calculate), but the more complicated the relationship the harder this is to do and will require additional computations and code. An easier solution would be to make separate calls to calculate_f with the \(y\) and \(t\) found from integration, but again this requires additional computations which may be expensive. In this case calculate_f will be called twice for each value in the time domain (once for the initial integration and once to find the value of f afterwards).

Instead, it would be ideal if we could just store the value of f during integration.

CyRK offers this functionality at the expense of a small hit on performance when using this feature. The values of these extra parameters are not analyzed by the solver when determining the adaptive step size or numerical error handling.

Limitations

Current limitations of this feature as of v0.17.0:

  • Only numerical parameters can be used as extra outputs (no strings, booleans, other structs).

  • All extra outputs must have the same type as the input ys. (for pysolve_ivp and cysolve_ivp extra outputs can only be double floating point numbers).

How to use with CyRK.nbsolve_ivp (Numba-based)

Define a differential equation function with the additional extra outputs put after the y-derivatives. It is a good idea to turn the output into numpy ndarray before the final return.

import numpy as np
from numba import njit

@njit
def diffeq(t, y, arg_0, arg_1):
    
    extra_parameter_0 = arg_0 * y[0] * np.sin(arg_1 * t)
    extra_parameter_1 = arg_0 * y[0] * np.cos(arg_1 * t - y[1])
    dy_0 = extra_parameter_0
    dy_1 = np.acos(extra_parameter_1)
    
    return np.asarray([dy_0, dy_1, extra_parameter_0, extra_parameter_1], dtype=y.dtype)

Call the nbsolve_ivp solver in the standard form except now with the capture_extra flag to True (by default it is False).

from CyRK import nbsolve_ivp
time_domain, all_output, success, message = \
    nbsolve_ivp(diffeq, time_span, initial_conds, rk_method=1, rtol=rtol, atol=atol, capture_extra=True)

The output of the solver (if successful) will now contain both the dependent y values over time as well as the extra parameters.

y_0 = all_output[0, :]
y_1 = all_output[1, :]
extra_parameter_0 = all_output[2, :]
extra_parameter_1 = all_output[3, :]

How to use with CyRK.pysolve_ivp

The process for capturing extra outputs with the pysolve_ivp solver follows similar steps to the numba-based approach where extra outputs are stored in the dy/dt output array. Note that the extra outputs must be stored at the end of the output array, after all dy/dts.

If you are using a returned dy/dt the format would look like:

def diffeq_cy(t, y, arg_0, arg_1):
    
    dy_size = y.size + 2 # 2 extra spots for the extra output
    dy = np.empty(dy_size, dtype=np.float64)
    extra_0 = (1. - arg_0 * y[1])
    extra_1 = (arg_1 * y[0] - 1.)
    # The differential equations
    dy[0] = extra_0 * y[0]
    dy[1] = extra_1 * y[1]

    # The extra output
    dy[2] = extra_0
    dy[3] = extra_1

Or if you pass in the differential equation as the first argument of the diffeq (recommended for improved performance!):

def diffeq_cy(dy, t, y, arg_0, arg_1):
    
    extra_0 = (1. - arg_0 * y[1])
    extra_1 = (arg_1 * y[0] - 1.)
    # The differential equations
    dy[0] = extra_0 * y[0]
    dy[1] = extra_1 * y[1]

    # The extra output
    dy[2] = extra_0
    dy[3] = extra_1

The next change is that the solver needs to know how many extra outputs it should expect. In our example there are two additional outputs. num_extra defaults to 0.

from CyRK import pysolve_ivp
result = \
    cyrk_ode(diffeq_cy, time_span, initial_conds, rk_method=1, rtol=rtol, atol=atol, num_extra=2)

The output matches the output described for the numba-based approach,

y_0 = result.y[0, :]
y_1 = result.y[1, :]
extra_parameter_0 = result.y[2, :]
extra_parameter_1 = result.y[3, :]

Interpolating Extra Outputs

Interpolating for cysolve_ivp and pysolve_ivp

When using either dense_output or t_eval, interpolations must be performed. For dependent \(y\) variables, CyRK will utilize the user-specified differential equation method’s approach to interpolation. However, this only works for dependent variables. If you are capturing extra output they will not be included in this interpolation. Instead, both cysolve_ivp and pysolve_ivp will perform said dependent-variable interpolation and then use the results to make additional calls (one per interpolation, i.e., once per value in t_eval) to the differential equation to find the values of the extra outputs at the requested time steps.

This approach is similar to “option 1” discussed in the next section for the numba-based solver.

Note

This means that collecting extra outputs when interpolation is one will result in additional calls to the ODE’s diffeq. If it is a computationally expensive function then this could cause a noticeable impact on performance. Read more about managing the performance of CyRK solvers here.

Interpolating for numba-based nbsolve_ivp

By setting the t_eval argument for either the nbsolve_ivp solver, an interpolation will occur at the end of integration. As of CyRK 0.17.0, nbsolve_ivp does not support dense output interpolators like the kind discussed in the previous section. Instead it must rely on (less accurate and less efficient) general purpose, linear interpolators. These use the solved \(y\)’s and \(t\) to find a new reduced y_reduced at t_eval intervals using a method similar to numpy.interp function. Since we are potentially storing extra parameters during integration, we need to tell the solver how to handle any potential interpolation on these additional parameters.

  • Option 1: Solve for the extra parameters at the new interpolated y values

    • How to set: Set t_eval to the desired numpy array, set capture_extra=True, and interpolate_extra=False.

    • Pros:

      • Less information is stored during integration which can reduce memory usage, especially when an integration requires many steps.

      • Results are more accurate since all the interpolation error is only in \(y\), extra parameters are as correct as \(y\) is (no compounding interpolation error).

    • Cons:

      • Additional calls to the diffeq are made after integration is completed (number of additional calls = len(t_eval)). If this function is computationally expensive then this may have a noticeable performance hit.

  • Option 2: Store extra parameters during integration and then perform interpolation on both y and all extra parameters.

    • How to set: Set t_eval to the desired numpy array, set capture_extra=True, and interpolate_extra=True.

    • Pros:

      • No additional calls to diffeq are made. This could improve performance if this is a particularly expensive function.

    • Mixed:

      • The same amount of information is stored during integration as if t_eval were not set in the first place. This can have a negative impact on memory and performance depending on the number of extra parameters and the number of integration steps.

    • Cons:

      • Results will be less accurate because the extra parameters are interpolated independent of y’s interpolation. Parameters that depend on \(y\) may no longer match expectations, particularly when \(y\) is changing quickly. For example, if an extra output is defined as x = y[0] * 2.0, then at some time \(t\) where y[0] was interpolated to be 5.0, \(x\) may differ from its definition (we expect it to be equal to 10.0): x(t) = 9.8.

      • More call(s) to interp will be made (once for each extra parameter) this may have a negative impact on performance, particularly for many extra parameters.

Tip

It is almost always better to use interpolate_extra=False. The results will be more accurate and even if there is a hit to performance, it should be minimal. A future release of CyRK will wrap the cysolve_ivp method to work with numba, eliminating these considerations.

Extra Outputs with Events

Extra outputs work with events as well. Events store y-data as an array each time an event is triggered. If extra outputs are used, then they will be stored with with each dependent y-value for each event that is triggered.