Ode solver v2

[mjsyts]

Revert ODE solver files. Made sure to target dev branch this time. Tests are consolidated into one file that tests all three solvers.

[spluta]

You are doing something here that we need to avoid in real-time audio, which is constructing and destroying lists on every audio loop. This will create problems. But I don't actually see why you need to do that, since num_dims is known at compile time. All lists should be created as struct members and their size should be constant.

[mjsyts]

@spluta: ODE solvers operate on scalar Float64 state. Multi-channel parallelism should be handled by the caller (e.g. separate solver instances per channel), not inside the solver itself. Unless you can think of a counterexample?

[tedmoore]

This seems to still be initializing InlineArrays in the functions. Because they're structs the memory can be initialized in init.

Also are you sure that fn derivatives needs to be defined inside the next function?

[mjsyts]

@tedmoore solved both issues:

struct RK2[num_dims: Int, fn_deriv: fn(InlineArray[Float64, num_dims], mut InlineArray[Float64, num_dims]) capturing -> None](Copyable, Movable):
    """Runge-Kutta 2nd order ODE solver.

    Parameters:
        num_dims: Number of dimensions (state variables), e.g. 2 for position and velocity.
        fn_deriv: Function that computes derivatives given the current state.
    """

    var state: InlineArray[Float64, Self.num_dims]
    var world: World
    var dt: Float64
    var k1: InlineArray[Float64, Self.num_dims]
    var k2: InlineArray[Float64, Self.num_dims]
    var temp_state: InlineArray[Float64, Self.num_dims]

    fn __init__(out self, world: World):
        """Initialize the RK2 struct."""
        self.world = world
        self.dt = 1.0 / world[].sample_rate
        self.state = InlineArray[Float64, Self.num_dims](fill=0.0)
        self.k1 = InlineArray[Float64, Self.num_dims](fill=0.0)
        self.k2 = InlineArray[Float64, Self.num_dims](fill=0.0)
        self.temp_state = InlineArray[Float64, Self.num_dims](fill=0.0)

    fn step(mut self):
        """Perform a single RK2 integration step."""
        Self.fn_deriv(self.state, self.k1)
        for i in range(Self.num_dims):
            self.temp_state[i] = self.state[i] + self.k1[i] * (self.dt / 2.0)

        Self.fn_deriv(self.temp_state, self.k2)

        for i in range(Self.num_dims):
            self.state[i] = self.state[i] + self.k2[i] * self.dt

The problem is fn_deriv as a struct parameter can't capture runtime state like frequency. I'm thinking a Differentiable trait might be the right solution. Structs implement it so the derivative function lives alongside whatever runtime state it needs. Does that make sense as a design direction for MMMAudio, or do you have a better approach?

[mjsyts]

Actually... Mojo trait declarations do not support parameters yet. So I'll have to come up with something else in the mean time

[mjsyts]

@tedmoore good question!
I used this to check the allocation and can confirm they were stack-allocated:

fn rk4_step() -> Float64:
    var state = InlineArray[Float64, 4](fill=Float64(0.0))
    var temp  = InlineArray[Float64, 4](fill=Float64(0.0))
    var k1    = InlineArray[Float64, 4](fill=Float64(0.0))
    var k2    = InlineArray[Float64, 4](fill=Float64(0.0))
    var k3    = InlineArray[Float64, 4](fill=Float64(0.0))
    var k4    = InlineArray[Float64, 4](fill=Float64(0.0))

    for i in range(4):
        temp[i] = state[i] + k1[i] * 0.5
        k2[i]   = temp[i] * 2.0
        temp[i] = state[i] + k2[i] * 0.5
        k3[i]   = temp[i] * 2.0
        temp[i] = state[i] + k3[i]
        k4[i]   = temp[i] * 2.0
        state[i] = state[i] + (k1[i] + 2.0*k2[i] + 2.0*k3[i] + k4[i]) * (1.0/6.0)

    return state[0]

fn main():
    print(rk4_step())

I verified this by compiling to LLVM IR (--emit llvm) and confirming all allocations are alloca instructions with zero malloc calls. The i8, i64 2048 alloca entries are the arrays themselves (256 × 8 bytes), so the "allocation" is just a stack pointer bump that's essentially free.

Keeping them local to step() rather than as struct fields is also intentional design: they're intermediate scratch buffers that should be recomputed each step. Storing them as fields would waste struct memory and risk stale intermediate values between calls.

The cleaner long-term solution would be a Differentiable[num_dims: Int] parametric trait where the derivative function lives alongside the runtime state. That would completely eliminate the temporaries, but Mojo doesn't support this yet. #180