优化FitzHugh-Nagumo方程求解效率：numba jitclass提升未达预期

Hey there! Let's break down why your numba jitclass only gave you a ~2x speedup for your FitzHugh-Nagumo ODE solver using Euler's method, and how you can squeeze more performance out of it:

1. Jitclass has inherent overhead (try pure function jitting instead) #

Jitclasses add a small but measurable overhead for attribute access and type checking, which can eat into gains when your core logic is a simple, tight loop like Euler's method. Instead of wrapping everything in a jitclass, extract the integration step into a standalone @njit-decorated function. This lets numba optimize the loop without the extra layer of class attribute lookups.

For example, instead of a jitclass with a step() method, write something like:

from numba import njit

@njit
def euler_fhn(u, v, a, b, tau, dt):
    n_cells = len(u)
    for i in range(n_cells):
        du = u[i] - (u[i]**3)/3 - v[i] + a
        dv = (u[i] - b - v[i]) / tau
        u[i] += du * dt
        v[i] += dv * dt

Then call this function repeatedly in your main loop—you’ll likely see a bigger speedup because numba can fully optimize the loop without class-related overhead.

2. Minimize attribute access inside hot loops #

If you stick with jitclass, avoid repeatedly accessing self.a, self.dt, etc., inside your loop. Pull these values into local variables before the loop starts to skip repeated type checks and attribute lookups:

from numba import jitclass, float64

spec = [
    ('u', float64[:]),
    ('v', float64[:]),
    ('a', float64),
    ('b', float64),
    ('tau', float64),
    ('dt', float64),
    ('n_cells', int)
]

@jitclass(spec)
class FitzHughNagumo:
    def __init__(self, u0, v0, a, b, tau, dt):
        self.u = u0
        self.v = v0
        self.a = a
        self.b = b
        self.tau = tau
        self.dt = dt
        self.n_cells = len(u0)
    
    def step(self):
        # Cache attributes to local variables first
        a = self.a
        b = self.b
        tau = self.tau
        dt = self.dt
        u = self.u
        v = self.v
        n_cells = self.n_cells
        
        for i in range(n_cells):
            du = u[i] - (u[i]**3)/3 - v[i] + a
            dv = (u[i] - b - v[i]) / tau
            u[i] += du * dt
            v[i] += dv * dt

This small change can reduce overhead in tight loops, especially when you’re iterating over thousands of myocardial cells.

3. Ensure numba can fully optimize your loop #

• Avoid Python operations in jitted code: Make sure you’re not calling un-jitted Python functions (like print(), or numpy higher-order functions) inside your loop. Replace numpy constructs like np.where with plain Python conditionals if possible—numba optimizes native loops far better.
• Use contiguous arrays: If your u and v arrays are non-contiguous (e.g., sliced from a larger array), convert them to contiguous arrays with np.ascontiguousarray() before passing to jitted code. Numba can vectorize operations more effectively on contiguous memory.
• Inline small helper functions: If you split the FHN right-hand side into a separate function, decorate it with @njit(inline='always') to let numba merge it into the main loop, eliminating function call overhead.

4. Consider switching to a higher-order ODE method (if feasible) #

Euler’s method is simple but requires small time steps for accuracy, which means more loop iterations. If your modeling requirements allow, try a higher-order method like RK4. Even though each iteration does more calculations, you can use a larger time step, reducing total iterations and potentially speeding up the overall simulation—especially when combined with numba optimizations.

After making these tweaks, you should see a much more significant speedup, potentially getting close to C-level performance for your solver.

内容的提问来源于stack exchange，提问作者ymmx