⚡️ Speed up function _equilibrium_payoffs_abreu_sannikov by 20%

[codeflash-ai]

📄 20% (0.20x) speedup for _equilibrium_payoffs_abreu_sannikov in quantecon/game_theory/repeated_game.py

⏱️ Runtime : 68.1 milliseconds → 57.0 milliseconds (best of 30 runs)

📝 Explanation and details

The optimized code achieves a 19% speedup by replacing non-JIT compiled functions with Numba-optimized versions, which is particularly beneficial for the computationally intensive loops in this repeated game algorithm.

Key Optimizations Applied:

1. JIT-compiled helper functions: Created _compute_best_dev_gain_numba(), _R_numba(), and _update_u_numba() with @njit(cache=True) decorators to eliminate Python interpreter overhead in hot loops.

2. Manual loop optimization in _R_numba(): Replaced NumPy's vectorized operations like (action_profile_payoff >= IC).all() and (equations @ extended_payoff <= tol).all() with explicit loops that Numba can heavily optimize.

3. Eliminated tuple comprehension in _best_dev_gains(): The original code used a generator expression with np.max() calls that couldn't be JIT-compiled. The optimized version uses a dedicated Numba function with manual loops for maximum performance.

Why These Optimizations Work:

• Numba's strength in tight loops: The _R() function contains nested loops over all action pairs, making it the dominant bottleneck (89.5% of runtime in the original). Numba eliminates Python overhead and enables CPU-level optimizations.
• Cache benefits: The cache=True parameter avoids recompilation overhead across multiple calls, important since this function is called from equilibrium_payoffs() which may be invoked repeatedly.
• Manual vectorization control: Explicit loops give Numba better optimization opportunities than complex NumPy operations involving boolean indexing and matrix operations.

Impact on Workloads:

Based on the function references, _equilibrium_payoffs_abreu_sannikov() is called from the main equilibrium_payoffs() method, making it a core computational path. The test results show consistent improvements across different game sizes:

• Large games benefit most: 20x20 games see 212% speedup, 10x10 games see 11.7% speedup
• Convergence scenarios: When algorithms converge quickly (high tolerance), speedups reach 58.1%
• Small games: Even 2x2 games show modest improvements, indicating the optimization overhead is minimal

The optimization is particularly valuable for researchers running multiple equilibrium computations or analyzing larger game matrices, where the cumulative time savings become significant.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 11 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest
from numba import njit
from quantecon.game_theory.repeated_game import \
    _equilibrium_payoffs_abreu_sannikov
from scipy.spatial import ConvexHull

# --- Minimal RepeatedGame and StageGame classes for testing ---

class StageGame:
    def __init__(self, payoff_arrays):
        self.payoff_arrays = payoff_arrays
        self.N = len(payoff_arrays)
        self.nums_actions = tuple(arr.shape[0] for arr in payoff_arrays)
        # payoff_profile_array: shape (a1, a2, 2)
        self.payoff_profile_array = np.empty(self.nums_actions + (2,))
        for a0 in range(self.nums_actions[0]):
            for a1 in range(self.nums_actions[1]):
                self.payoff_profile_array[a0, a1, 0] = payoff_arrays[0][a0, a1]
                self.payoff_profile_array[a0, a1, 1] = payoff_arrays[1][a1, a0]

class RepeatedGame:
    def __init__(self, payoff_arrays, delta):
        self.sg = StageGame(payoff_arrays)
        self.delta = delta
from quantecon.game_theory.repeated_game import \
    _equilibrium_payoffs_abreu_sannikov

# --- Basic Test Cases ---

def test_battle_of_the_sexes_basic():
    # Battle of the Sexes
    A = np.array([[2,0],[0,1]])
    B = np.array([[1,0],[0,2]])
    delta = 0.8
    rpg = RepeatedGame([A, B], delta)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 5.23ms -> 145μs (3495% faster)
    # Should contain (2,1), (1,2), (0,0)
    points = hull.points[hull.vertices]
    expected = np.array([[2,1],[1,2],[0,0]])
    for pt in expected:
        pass

def test_one_player_game_raises():
    # Only one player: should raise NotImplementedError
    A = np.array([[1,2],[3,4]])
    class FakeStageGame:
        def __init__(self):
            self.N = 1
            self.nums_actions = (2,)
            self.payoff_arrays = [A]
            self.payoff_profile_array = np.zeros((2,1,2))
    class FakeRepeatedGame:
        def __init__(self):
            self.sg = FakeStageGame()
            self.delta = 0.9
    rpg = FakeRepeatedGame()
    with pytest.raises(NotImplementedError):
        _equilibrium_payoffs_abreu_sannikov(rpg) # 791ns -> 750ns (5.47% faster)

def test_non_square_payoff_matrices():
    # Player 0 has 2 actions, player 1 has 3
    A = np.array([[1,2,3],[4,5,6]])
    B = np.array([[6,5],[4,3],[2,1]])
    rpg = RepeatedGame([A, B], 0.8)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 158μs -> 165μs (4.27% slower)

def test_performance_on_20x20_game():
    # 20x20 payoff matrices, random
    np.random.seed(123)
    A = np.random.randint(-5, 5, size=(20,20))
    B = np.random.randint(-5, 5, size=(20,20))
    rpg = RepeatedGame([A, B], 0.8)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg, max_iter=10); hull = codeflash_output # 2.33ms -> 746μs (212% faster)

# --- Determinism test ---

import numpy as np
# imports
import pytest
from numba import njit
from quantecon.game_theory.repeated_game import \
    _equilibrium_payoffs_abreu_sannikov
from scipy.spatial import ConvexHull

# --- Minimal RepeatedGame and StageGame classes for testing ---
class StageGame:
    def __init__(self, payoff_arrays):
        self.payoff_arrays = payoff_arrays
        self.N = len(payoff_arrays)
        self.nums_actions = tuple(arr.shape[0] for arr in payoff_arrays)
        # payoff_profile_array: shape (A1, A2, N)
        self.payoff_profile_array = np.stack(payoff_arrays, axis=-1)

class RepeatedGame:
    def __init__(self, payoff_arrays, delta):
        self.sg = StageGame(payoff_arrays)
        self.delta = delta
from quantecon.game_theory.repeated_game import \
    _equilibrium_payoffs_abreu_sannikov

# --- Unit tests ---

# 1. Basic Test Cases

def test_basic_asymmetric_payoff():
    # 2x2 game, asymmetric payoffs
    payoff_arrays = (
        np.array([[2, 0], [0, 1]], dtype=float),
        np.array([[1, 0], [0, 3]], dtype=float)
    )
    rpg = RepeatedGame(payoff_arrays, delta=0.8)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 3.89ms -> 3.89ms (0.145% faster)
    pts = hull.points[hull.vertices]
    # The vertices should be the four possible payoff pairs
    expected = sorted([[2,1],[0,0],[0,3],[1,0]])
    result = sorted([list(map(float,pt)) for pt in pts])
    for e, r in zip(expected, result):
        pass

def test_basic_high_discount():
    # 2x2 game, delta near 1
    payoff_arrays = (
        np.array([[5, 0], [0, 2]], dtype=float),
        np.array([[2, 0], [0, 5]], dtype=float)
    )
    rpg = RepeatedGame(payoff_arrays, delta=0.99)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 15.2ms -> 15.1ms (0.434% faster)
    pts = hull.points[hull.vertices]
    # Vertices should be the four possible payoff pairs
    expected = sorted([[5,2],[0,0],[0,5],[2,0]])
    result = sorted([list(map(float,pt)) for pt in pts])
    for e, r in zip(expected, result):
        pass

def test_edge_non_two_player_game():
    # 3-player game should raise NotImplementedError
    payoff_arrays = (
        np.ones((2,2), dtype=float),
        np.ones((2,2), dtype=float),
        np.ones((2,2), dtype=float)
    )
    class StageGame3:
        def __init__(self, payoff_arrays):
            self.payoff_arrays = payoff_arrays
            self.N = len(payoff_arrays)
            self.nums_actions = tuple(arr.shape[0] for arr in payoff_arrays)
            self.payoff_profile_array = np.stack(payoff_arrays, axis=-1)
    class RepeatedGame3:
        def __init__(self, payoff_arrays, delta):
            self.sg = StageGame3(payoff_arrays)
            self.delta = delta
    rpg = RepeatedGame3(payoff_arrays, delta=0.9)
    with pytest.raises(NotImplementedError):
        _equilibrium_payoffs_abreu_sannikov(rpg) # 750ns -> 792ns (5.30% slower)

def test_large_scale_10x10_game():
    # 10x10 game, random payoffs
    np.random.seed(42)
    payoff_arrays = (
        np.random.uniform(-10, 10, size=(10,10)),
        np.random.uniform(-10, 10, size=(10,10))
    )
    rpg = RepeatedGame(payoff_arrays, delta=0.95)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 40.2ms -> 36.0ms (11.7% faster)
    pts = hull.points[hull.vertices]
    # All points should be within the min/max of payoffs
    min0, max0 = payoff_arrays[0].min(), payoff_arrays[0].max()
    min1, max1 = payoff_arrays[1].min(), payoff_arrays[1].max()

def test_large_scale_sparse_payoff():
    # 30x30 game, mostly zeros, some high values
    payoff_arrays = (
        np.zeros((30,30), dtype=float),
        np.zeros((30,30), dtype=float)
    )
    payoff_arrays[0][0,0] = 100.0
    payoff_arrays[1][29,29] = 200.0
    rpg = RepeatedGame(payoff_arrays, delta=0.9)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 663μs -> 576μs (15.2% faster)
    pts = hull.points[hull.vertices]

def test_large_scale_max_iter():
    # 5x5 game, set max_iter to a small value
    np.random.seed(123)
    payoff_arrays = (
        np.random.uniform(0, 10, size=(5,5)),
        np.random.uniform(0, 10, size=(5,5))
    )
    rpg = RepeatedGame(payoff_arrays, delta=0.9)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg, max_iter=1); hull = codeflash_output # 98.9μs -> 92.9μs (6.46% faster)
    pts = hull.points[hull.vertices]

def test_large_scale_tol():
    # 5x5 game, set tol very large so it converges immediately
    np.random.seed(456)
    payoff_arrays = (
        np.random.uniform(0, 10, size=(5,5)),
        np.random.uniform(0, 10, size=(5,5))
    )
    rpg = RepeatedGame(payoff_arrays, delta=0.9)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg, tol=1e5); hull = codeflash_output # 360μs -> 227μs (58.1% faster)
    pts = hull.points[hull.vertices]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-_equilibrium_payoffs_abreu_sannikov-mja8w7c1 and push.