GitHub

Workspace API

The workspace API provides a pre-allocated, reusable solving pattern via init and solve!. This avoids repeated memory allocation when solving the same type of problem many times, and is required for compatibility with Enzyme.jl reverse-mode AD.

DifferenceEquations.StateSpaceWorkspace — Type
StateSpaceWorkspace(prob, alg, output, cache)

Workspace for state-space problem solvers. Holds the problem, algorithm, pre-allocated output arrays, and scratch cache. Created by init and consumed by solve!.

Arguments

  • prob::AbstractStateSpaceProblem: Problem to solve.
  • alg: Difference-equation algorithm.
  • output: Pre-allocated output buffers owned by the workspace.
  • cache: Scratch buffers used by the solver.

Fields

  • prob: Problem solved by the workspace.
  • alg: Algorithm used by solve!.
  • output: Named tuple of pre-allocated state, observation, and covariance buffers.
  • cache: Scratch arrays and factorizations reused across solves.
  • save_everystep::Bool: Whether the full trajectory is retained.

Returns

  • StateSpaceWorkspace: A reusable workspace for repeated solves.

Interface

Construct workspaces with init unless an external AD or allocation pipeline already owns compatible output and cache buffers. Call solve! to overwrite the buffers and return a new StateSpaceSolution.

source

Creating and Using a Workspace

CommonSolve.init — Function
CommonSolve.init(prob::AbstractStateSpaceProblem, alg = default_alg(prob); save_everystep = true, kwargs...)

Create a StateSpaceWorkspace with pre-allocated output arrays and scratch cache. When save_everystep=false, allocate minimal two-element buffers containing only the initial and final states.

Arguments

  • prob::AbstractStateSpaceProblem: Problem to solve repeatedly.
  • alg: Algorithm; defaults to the problem-dependent algorithm selected by default_alg.

Keyword Arguments

  • save_everystep::Bool = true: Store the full trajectory when true; retain only endpoint buffers when false.
  • kwargs...: Additional allocation options forwarded to the workspace setup.

Returns

  • StateSpaceWorkspace: A reusable workspace consumed by solve!.
source
CommonSolve.solve! — Function
CommonSolve.solve!(ws::StateSpaceWorkspace; kwargs...)

Solve the state-space problem. Mutate ws.output arrays in place, then wrap them in a StateSpaceSolution and return it.

Arguments

  • ws::StateSpaceWorkspace: Workspace created by init.

Keyword Arguments

  • kwargs...: Solver options forwarded to the selected state-space algorithm.

Returns

  • StateSpaceSolution: The solution backed by the workspace's latest output.
source

Basic Usage

using DifferenceEquations, LinearAlgebra, Random
A = [0.95 6.2; 0.0 0.2]
B = [0.0; 0.01;;]
C = [0.09 0.67; 1.00 0.00]
u0 = zeros(2)
prob = LinearStateSpaceProblem(A, B, u0, (0, 5); C)
ws = init(prob, DirectIteration())
sol = solve!(ws)
sol.u[end]
2-element Vector{Float64}:
 -0.006400851860193932
 -0.001126390788464822

Cache Reuse

Calling solve!(ws) again on the same workspace reuses all previously allocated buffers. The solver fully overwrites all output arrays on each call, so no manual reset is needed between calls. This makes the workspace pattern ideal for tight loops:

ws = init(prob, DirectIteration())
for i in 1:1000
    sol = solve!(ws)
    # process sol...
end

You can also change the problem between calls for parameter sweeps using remake:

ws = init(prob, DirectIteration())
for a11 in [0.9, 0.95, 1.0]
    ws.prob = remake(ws.prob; A = [a11 6.2; 0.0 0.2])
    sol = solve!(ws)
    # process sol.logpdf...
end

Endpoints-Only Mode (save_everystep=false)

Pass save_everystep=false to init to allocate minimal 2-element buffers. The solver stores only the initial and final states, while still correctly accumulating logpdf:

ws_ep = init(prob, DirectIteration(); save_everystep=false)
sol_ep = solve!(ws_ep)
length(sol_ep.u)  # 2: [u_initial, u_final]
2

This is especially useful for ForwardDiff gradient computation, where reducing the number of dual-number allocations from O(T) to O(1) gives significant speedups (up to 7x with StaticArrays):

# ForwardDiff benefits from save_everystep=false
function neg_loglik(params)
    prob = make_problem(params)
    return -solve(prob, ConditionalLikelihood(); save_everystep=false).logpdf
end
ForwardDiff.gradient(neg_loglik, params0)

When to Use

The workspace API is useful in the following scenarios:

  • Enzyme AD: Enzyme requires pre-allocated buffers passed as Duplicated arguments. The workspace pattern via init/solve! is the recommended way to use Enzyme with DifferenceEquations.jl. See Enzyme AD for details.
  • Repeated solves in optimization loops: When solving the same problem structure many times (e.g., during parameter estimation), the workspace avoids allocating new arrays on every iteration.
  • ForwardDiff with save_everystep=false: Combining the workspace API with endpoints-only mode minimizes dual-number allocations, giving the best ForwardDiff performance.
  • Performance-critical code: Eliminating allocations reduces GC pressure and improves performance, especially for small to medium-sized problems.