Julia vs NumPy broadcasting

I think you can do mapslices instead of loop.

And I’m seconding @gdalle on using StaticArrays.jl for the arrays of known dimension

function step_analytical(A::AbstractMatrix, B::AbstractArray, F::Number, n::Integer)
    # Eigendecomposition: A = V D V_inv
    eigvals, V = eigen(A)
    # Precompute constant term
    BF_transformed = V \ (B .* F) # (4, 1), must be the same as inv(V) * (B .* F)
    # Time steps (n,)
    t = 1:n
    # Compute (exp - 1) / lambda
    diag_terms = expm1.(t' .* eigvals) ./ eigvals # (4, n) Must be more numerically stable than exp(x) - 1
    # Scale V columns by diag_terms
    V_scaled = reshape(V, 4, 4, :) .* reshape(diag_terms, 4, :, n)
    # Multiply by BF in transformed space
    x = mapslices(V_scaled; dims=(1, 2)) do v # (4, n)
        v * BF_transformed
    end
    return x
end

For static arrays:

using LinearAlgebra
using StaticArrays

function step_analytical(A::SMatrix, B::SArray, F::Number, n::Integer)
    # Eigendecomposition: A = V D V_inv
    eigvals, V = eigen(A)
    # Precompute constant term
    BF_transformed = V \ (B * F) # (4, 1), must be the same as inv(V) * (B .* F)
    # Compute (exp - 1) / lambda
    diag_terms = [expm1.(t .* eigvals) ./ eigvals for t in 1:n] # (4, n) Must be more numerically stable than exp(x) - 1
    # Scale V columns by diag_terms
    V_scaled = V .* Diagonal.(diag_terms)
    # Multiply by BF in transformed space
    x = V_scaled .* BF_transformed
    return x
end
2 Likes