基于numba优化scipy.optimize.root的excitingmixing算法问询

Hey there! Let's tackle your problem step by step. You're stuck because scipy.optimize.root(method="excitingmixing") works for your convergence needs, but Numba doesn't support Scipy in nopython mode. I'll break down two solid solutions for you: implementing the exciting mixing algorithm yourself (compatible with Numba) and optimizing the Scipy workflow with Numba-accelerated functions.

First, let's demystify what "exciting mixing" actually does under the hood in Scipy. It's a damped Newton method that uses a diagonal Jacobian approximation (instead of a full Jacobian) with adaptive step-size adjustment (the "exciting" part) to avoid divergence. Here's a functional implementation tailored to your use case:

import numpy as np
from numba import jit

@jit(nopython=True)
def func(x):
    a, b, c, d = x
    da = a*(1 - b)
    db = b*(1 - c)
    dc = c
    dd = d - 1  # Note: Fixed your example to have a valid solution (dd=0 when d=1)
    return np.array([da, db, dc, dd])

@jit(nopython=True)
def diag_jacobian(x):
    # Analytic diagonal Jacobian elements (∂F_i/∂x_i) for your function
    a, b, c, _ = x
    j11 = 1 - b  # ∂da/∂a
    j22 = 1 - c  # ∂db/∂b
    j33 = 1      # ∂dc/∂c
    j44 = 1      # ∂dd/∂d
    return np.array([j11, j22, j33, j44])

@jit(nopython=True)
def exciting_mixing_root(x0, tol=1e-8, max_iter=1000, alpha_init=1.0, alpha_adjust=0.5):
    x = x0.copy()
    alpha = alpha_init
    
    for _ in range(max_iter):
        F = func(x)
        # Check if residual meets tolerance
        if np.linalg.norm(F) < tol:
            return x
        
        # Compute inverse of diagonal Jacobian (avoid division by zero)
        J_diag = diag_jacobian(x)
        J_diag_inv = 1.0 / np.where(np.abs(J_diag) < 1e-10, 1e-10, J_diag)
        
        # Calculate Newton update direction
        delta = -F * J_diag_inv
        
        # Adaptive step size adjustment: shrink alpha until residual decreases
        x_new = x + alpha * delta
        F_new = func(x_new)
        while np.linalg.norm(F_new) >= np.linalg.norm(F) and alpha > 1e-10:
            alpha *= alpha_adjust
            x_new = x + alpha * delta
            F_new = func(x_new)
        
        if alpha <= 1e-10:
            # Step size too small to make progress, return current point
            return x
        
        # Update x and reset alpha for next iteration
        x = x_new
        alpha = min(alpha_init, alpha / alpha_adjust)
    
    # Return last point if max iterations reached
    return x

# Test the implementation
root = exciting_mixing_root(np.array([0.1, 0.1, 0.2, 0.4]))
print(root)

Key notes about this implementation:

• We use an analytic diagonal Jacobian for speed (you could also use finite differences if analytic derivatives aren't feasible)
• The adaptive step-size logic ensures we only accept updates that reduce the residual, which is the core of the "exciting" mixing behavior
• We handle edge cases like division by zero in the Jacobian inverse

You don't need to jit the entire getRoot function. Since the bottleneck is your func evaluation, you can just jit that and pass it to Scipy's root function directly. Scipy will happily call your Numba-accelerated func, giving you the best of both worlds: Scipy's robust algorithm and Numba's speed.

import numpy as np
from scipy import optimize
from numba import jit

@jit(nopython=True)
def func(x):
    a, b, c, d = x
    da = a*(1 - b)
    db = b*(1 - c)
    dc = c
    dd = d - 1  # Fixed example to have valid solution
    return [da, db, dc, dd]

# No need to jit this function—Scipy handles the algorithm, we just accelerate the core computation
def getRoot(x0):
    solution = optimize.root(func, x0, method="excitingmixing")
    return solution.x

root = getRoot([0.1, 0.1, 0.2, 0.4])
print(root)

If you want to match Scipy's exact implementation, you can look at its source code (the ExcitingMixing class in scipy/optimize/_root.py). The core logic mirrors what we implemented, with a few extra refinements:

• It uses finite differences for the diagonal Jacobian if no analytic Jacobian is provided
• It includes more sophisticated step-size adaptation logic
• It handles edge cases like non-positive diagonal Jacobian elements more rigorously

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