⚡️ Speed up function close_normal_form_games by 371%

[codeflash-ai]

📄 371% (3.71x) speedup for close_normal_form_games in quantecon/game_theory/tests/test_polymatrix_game.py

⏱️ Runtime : 2.95 milliseconds → 627 microseconds (best of 164 runs)

📝 Explanation and details

The optimized version achieves a 371% speedup by replacing Python loops and NumPy's allclose function with faster Numba-compiled alternatives for the most computationally expensive operations.

Key Optimizations Applied:

1. Numba-compiled action comparison (_nums_actions_equal): Replaces the Python loop that compares nums_actions arrays with a JIT-compiled function, eliminating Python interpreter overhead for this comparison.

2. Custom Numba array comparison (_allclose_ndarray): Replaces NumPy's allclose function with a specialized Numba implementation that directly iterates over flattened arrays, avoiding NumPy's overhead for tolerance checking.

3. Smart fallback strategy: The code maintains full compatibility by falling back to the original NumPy logic for edge cases (empty games, shape mismatches) while using the optimized path for common scenarios.

Why This Creates Significant Speedup:

• Elimination of Python loops: The line profiler shows that allclose calls consumed 98% of the original runtime. The Numba version compiles to native machine code, removing Python's interpreted loop overhead.
• Reduced function call overhead: Instead of calling NumPy's general-purpose allclose, the optimized version uses a specialized comparison that's tailored for this specific use case.
• Cache benefits: The @njit(cache=True) decorator ensures the compiled functions are cached, making subsequent calls even faster.

Performance Impact on Workloads:

Based on the function references, close_normal_form_games is called in test scenarios that involve:

• Converting between different game representations (PolymatrixGame.from_nf, polymg.to_nfg)
• Validating game equality in unit tests
• Checking approximation quality

The test results show consistent 2-4x speedups across various scenarios, with the largest gains (up to 16x) occurring when arrays differ early in the comparison, allowing for fast rejection.

Best Performance Gains For:

• Large games with many players/actions (355-420% speedups)
• Games with identical payoffs (314-400% speedups)
• Games where differences are detected early (up to 1603% speedup)

The optimization is particularly valuable in iterative algorithms or batch processing scenarios where game comparisons are performed repeatedly.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 34 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  # used for our unit tests
# function to test
from numpy import allclose
from quantecon.game_theory.tests.test_polymatrix_game import \
    close_normal_form_games

# Minimal mock for NormalFormGame and Player to be self-contained for testing
class Player:
    def __init__(self, payoff_array):
        self.payoff_array = np.array(payoff_array)

class NormalFormGame:
    def __init__(self, payoff_arrays):
        # payoff_arrays: list of n arrays, one per player
        self.N = len(payoff_arrays)
        self.nums_actions = [arr.shape[0] if arr.ndim > 0 else 1 for arr in payoff_arrays]
        self.players = [Player(arr) for arr in payoff_arrays]
from quantecon.game_theory.tests.test_polymatrix_game import \
    close_normal_form_games

# unit tests

# --- Basic Test Cases ---

def test_identical_games():
    # Two identical 2-player games with 2 actions each
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([[4, 3], [2, 1]])
    nf1 = NormalFormGame([arr1, arr2])
    nf2 = NormalFormGame([arr1.copy(), arr2.copy()])
    codeflash_output = close_normal_form_games(nf1, nf2) # 51.9μs -> 10.4μs (400% faster)

def test_games_different_payoffs():
    # Games with different payoffs should not be close
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([[4, 3], [2, 1]])
    arr2_diff = np.array([[4, 3], [2, 0]])
    nf1 = NormalFormGame([arr1, arr2])
    nf2 = NormalFormGame([arr1, arr2_diff])
    codeflash_output = not close_normal_form_games(nf1, nf2) # 28.3μs -> 6.92μs (310% faster)

def test_games_within_atol():
    # Games with payoffs within atol should be close
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([[4, 3], [2, 1]])
    arr1_perturbed = arr1 + 1e-5
    arr2_perturbed = arr2 - 1e-5
    nf1 = NormalFormGame([arr1, arr2])
    nf2 = NormalFormGame([arr1_perturbed, arr2_perturbed])
    codeflash_output = close_normal_form_games(nf1, nf2) # 26.8μs -> 7.92μs (238% faster)

def test_games_outside_atol():
    # Games with payoffs outside atol should not be close
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([[4, 3], [2, 1]])
    arr1_perturbed = arr1 + 1e-2
    arr2_perturbed = arr2 - 1e-2
    nf1 = NormalFormGame([arr1, arr2])
    nf2 = NormalFormGame([arr1_perturbed, arr2_perturbed])
    codeflash_output = not close_normal_form_games(nf1, nf2, atol=1e-4) # 15.4μs -> 6.54μs (136% faster)

def test_games_different_num_players():
    # Games with different numbers of players should not be close
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([[4, 3], [2, 1]])
    arr3 = np.array([[0, 0], [0, 0]])
    nf1 = NormalFormGame([arr1, arr2])
    nf2 = NormalFormGame([arr1, arr2, arr3])
    codeflash_output = not close_normal_form_games(nf1, nf2) # 125ns -> 166ns (24.7% slower)

def test_games_different_num_actions():
    # Games with different numbers of actions for a player should not be close
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([[4, 3], [2, 1]])
    arr1_short = np.array([[1, 2]])
    nf1 = NormalFormGame([arr1, arr2])
    nf2 = NormalFormGame([arr1_short, arr2])
    codeflash_output = not close_normal_form_games(nf1, nf2) # 458ns -> 5.42μs (91.5% slower)

# --- Edge Test Cases ---

def test_zero_players():
    # Games with zero players should be close
    nf1 = NormalFormGame([])
    nf2 = NormalFormGame([])
    codeflash_output = close_normal_form_games(nf1, nf2) # 458ns -> 500ns (8.40% slower)

def test_one_player_game():
    # Single player games with identical payoffs
    arr1 = np.array([1, 2, 3])
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr1.copy()])
    codeflash_output = close_normal_form_games(nf1, nf2) # 16.6μs -> 6.58μs (153% faster)

def test_one_player_game_different_length():
    # Single player games, different number of actions
    arr1 = np.array([1, 2, 3])
    arr2 = np.array([1, 2])
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr2])
    codeflash_output = not close_normal_form_games(nf1, nf2) # 417ns -> 4.58μs (90.9% slower)

def test_empty_payoff_arrays():
    # Games with empty payoff arrays
    arr1 = np.array([])
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr1.copy()])
    codeflash_output = close_normal_form_games(nf1, nf2) # 14.8μs -> 6.46μs (129% faster)

def test_scalar_payoff_arrays():
    # Scalar payoff arrays (0-dimensional)
    arr1 = np.array(5)
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr1.copy()])
    codeflash_output = close_normal_form_games(nf1, nf2) # 17.5μs -> 7.21μs (143% faster)

def test_mixed_shape_payoff_arrays():
    # Games where payoff arrays have different shapes (should not be close)
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([1, 2])
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr2])
    codeflash_output = not close_normal_form_games(nf1, nf2) # 19.2μs -> 45.6μs (58.0% slower)

def test_large_atol_accepts_difference():
    # Large atol should accept bigger differences
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = arr1 + 0.5
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr2])
    codeflash_output = close_normal_form_games(nf1, nf2, atol=1) # 16.6μs -> 6.58μs (152% faster)

def test_negative_payoffs():
    # Negative payoffs should be handled correctly
    arr1 = np.array([[-1, -2], [-3, -4]])
    arr2 = np.array([[-1, -2], [-3, -4]])
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr2])
    codeflash_output = close_normal_form_games(nf1, nf2) # 16.4μs -> 5.54μs (196% faster)

# --- Large Scale Test Cases ---

def test_large_number_of_actions():
    # Test with 2 players, each with 100 actions, identical payoffs
    arr1 = np.ones((100, 100))
    arr2 = np.ones((100, 100))
    nf1 = NormalFormGame([arr1, arr2])
    nf2 = NormalFormGame([arr1.copy(), arr2.copy()])
    codeflash_output = close_normal_form_games(nf1, nf2) # 85.8μs -> 18.9μs (355% faster)

def test_large_but_different_payoffs():
    # Large games, but one payoff array differs by more than atol
    arr1 = np.zeros((50, 50))
    arr2 = np.zeros((50, 50))
    arr2[0, 0] = 1  # difference
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr2])
    codeflash_output = not close_normal_form_games(nf1, nf2) # 22.2μs -> 6.04μs (267% faster)

def test_large_games_within_atol():
    # Large games, all payoffs within atol
    arr1 = np.random.rand(20, 20)
    arr2 = arr1 + np.random.uniform(-1e-5, 1e-5, size=(20, 20))
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr2])
    codeflash_output = close_normal_form_games(nf1, nf2, atol=1e-4) # 16.8μs -> 6.04μs (179% faster)

def test_large_games_outside_atol():
    # Large games, some payoffs outside atol
    arr1 = np.random.rand(20, 20)
    arr2 = arr1.copy()
    arr2[0, 0] += 1e-2  # outside default atol
    nf1 = NormalFormGame([arr1])
    nf2 = NormalFormGame([arr2])
    codeflash_output = not close_normal_form_games(nf1, nf2, atol=1e-4) # 16.3μs -> 5.67μs (187% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import numpy as np
# imports
import pytest  # used for our unit tests
# function to test
from numpy import allclose
from quantecon.game_theory.tests.test_polymatrix_game import \
    close_normal_form_games

# Minimal NormalFormGame and Player classes for testing
class Player:
    def __init__(self, payoff_array):
        self.payoff_array = np.array(payoff_array)

class NormalFormGame:
    def __init__(self, payoff_arrays):
        # payoff_arrays: list of arrays, one for each player
        self.N = len(payoff_arrays)
        self.nums_actions = [arr.shape[0] if arr.ndim == 1 else arr.shape[0] for arr in payoff_arrays]
        self.players = [Player(arr) for arr in payoff_arrays]
from quantecon.game_theory.tests.test_polymatrix_game import \
    close_normal_form_games

# unit tests

# ------------------------
# Basic Test Cases
# ------------------------

def test_identical_games():
    # Two identical games (2 players, 2 actions each)
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([[4, 3], [2, 1]])
    g1 = NormalFormGame([arr1, arr2])
    g2 = NormalFormGame([arr1.copy(), arr2.copy()])
    # Should be considered close
    codeflash_output = close_normal_form_games(g1, g2) # 31.1μs -> 7.50μs (314% faster)

def test_games_within_tolerance():
    # Payoffs differ by less than atol
    arr1 = np.array([[1.00001, 2.00001], [3.00001, 4.00001]])
    arr2 = np.array([[4.00001, 3.00001], [2.00001, 1.00001]])
    arr1b = np.array([[1.00002, 2.00002], [3.00002, 4.00002]])
    arr2b = np.array([[4.00002, 3.00002], [2.00002, 1.00002]])
    g1 = NormalFormGame([arr1, arr2])
    g2 = NormalFormGame([arr1b, arr2b])
    # Should be considered close with default atol
    codeflash_output = close_normal_form_games(g1, g2) # 24.8μs -> 6.79μs (266% faster)

def test_games_outside_tolerance():
    # Payoffs differ by more than atol
    arr1 = np.array([[1.0, 2.0], [3.0, 4.0]])
    arr2 = np.array([[4.0, 3.0], [2.0, 1.0]])
    arr1b = np.array([[1.1, 2.0], [3.0, 4.0]])
    arr2b = np.array([[4.0, 3.0], [2.0, 1.0]])
    g1 = NormalFormGame([arr1, arr2])
    g2 = NormalFormGame([arr1b, arr2b])
    # Should NOT be considered close (1.0 vs 1.1 diff > 1e-4)
    codeflash_output = not close_normal_form_games(g1, g2) # 14.8μs -> 6.33μs (134% faster)

def test_games_with_different_number_of_players():
    # One game has 2 players, the other has 3
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([[4, 3], [2, 1]])
    arr3 = np.array([[0, 0], [0, 0]])
    g1 = NormalFormGame([arr1, arr2])
    g2 = NormalFormGame([arr1, arr2, arr3])
    # Should NOT be considered close
    codeflash_output = not close_normal_form_games(g1, g2) # 208ns -> 167ns (24.6% faster)

def test_empty_games():
    # Both games have zero players
    g1 = NormalFormGame([])
    g2 = NormalFormGame([])
    # Should be considered close
    codeflash_output = close_normal_form_games(g1, g2) # 541ns -> 667ns (18.9% slower)

def test_one_empty_one_nonempty():
    # One game is empty, the other is not
    arr1 = np.array([[1, 2], [3, 4]])
    arr2 = np.array([[4, 3], [2, 1]])
    g1 = NormalFormGame([])
    g2 = NormalFormGame([arr1, arr2])
    # Should NOT be considered close
    codeflash_output = not close_normal_form_games(g1, g2) # 166ns -> 166ns (0.000% faster)

def test_single_player_games():
    # Games with a single player (degenerate case)
    arr1 = np.array([1, 2, 3])
    arr2 = np.array([1, 2, 3])
    g1 = NormalFormGame([arr1])
    g2 = NormalFormGame([arr2])
    # Should be considered close
    codeflash_output = close_normal_form_games(g1, g2) # 27.2μs -> 11.5μs (136% faster)

def test_single_player_games_different_actions():
    # Single player, but different number of actions
    arr1 = np.array([1, 2, 3])
    arr2 = np.array([1, 2])
    g1 = NormalFormGame([arr1])
    g2 = NormalFormGame([arr2])
    # Should NOT be considered close
    codeflash_output = not close_normal_form_games(g1, g2) # 458ns -> 4.79μs (90.4% slower)

def test_games_with_nan_payoff():
    # Games with np.nan in payoffs (should not be close)
    arr1 = np.array([[1, 2], [3, np.nan]])
    arr2 = np.array([[1, 2], [3, 4]])
    g1 = NormalFormGame([arr1])
    g2 = NormalFormGame([arr2])
    codeflash_output = not close_normal_form_games(g1, g2) # 17.8μs -> 6.62μs (168% faster)

def test_games_with_inf_payoff():
    # Games with np.inf in payoffs (should not be close)
    arr1 = np.array([[1, 2], [3, np.inf]])
    arr2 = np.array([[1, 2], [3, 4]])
    g1 = NormalFormGame([arr1])
    g2 = NormalFormGame([arr2])
    codeflash_output = not close_normal_form_games(g1, g2) # 16.0μs -> 5.62μs (184% faster)

def test_games_with_different_shapes_but_same_num_actions():
    # 2D vs 1D payoff arrays for single player (should be close if values match)
    arr1 = np.array([1, 2, 3])
    arr2 = np.array([[1, 2, 3]])
    g1 = NormalFormGame([arr1])
    g2 = NormalFormGame([arr2])
    # Should NOT be considered close as shapes differ (allclose will still work)
    codeflash_output = close_normal_form_games(g1, g2) # 417ns -> 4.54μs (90.8% slower)

# ------------------------
# Large Scale Test Cases
# ------------------------

def test_large_games_identical():
    # Large 2-player game, each with 500 actions
    arr1 = np.random.rand(500, 500)
    arr2 = np.random.rand(500, 500)
    g1 = NormalFormGame([arr1, arr2])
    g2 = NormalFormGame([arr1.copy(), arr2.copy()])
    # Should be considered close
    codeflash_output = close_normal_form_games(g1, g2) # 988μs -> 190μs (420% faster)

def test_large_games_with_small_noise():
    # Large 2-player game, each with 500 actions, small noise within atol
    arr1 = np.random.rand(500, 500)
    arr2 = np.random.rand(500, 500)
    noise1 = arr1 + np.random.uniform(-1e-5, 1e-5, arr1.shape)
    noise2 = arr2 + np.random.uniform(-1e-5, 1e-5, arr2.shape)
    g1 = NormalFormGame([arr1, arr2])
    g2 = NormalFormGame([noise1, noise2])
    # Should be considered close
    codeflash_output = close_normal_form_games(g1, g2) # 971μs -> 187μs (417% faster)

def test_large_games_with_large_noise():
    # Large 2-player game, each with 500 actions, some noise outside atol
    arr1 = np.random.rand(500, 500)
    arr2 = np.random.rand(500, 500)
    noise1 = arr1.copy()
    noise2 = arr2.copy()
    # Introduce a large difference at a random spot
    noise1[100, 100] += 1e-2  # much larger than default atol
    g1 = NormalFormGame([arr1, arr2])
    g2 = NormalFormGame([noise1, noise2])
    # Should NOT be considered close
    codeflash_output = not close_normal_form_games(g1, g2) # 493μs -> 29.0μs (1603% faster)

def test_large_games_different_num_actions():
    # Large games, but one player has a different number of actions
    arr1 = np.random.rand(500, 500)
    arr2 = np.random.rand(500, 500)
    arr1b = np.random.rand(499, 500)
    arr2b = np.random.rand(500, 500)
    g1 = NormalFormGame([arr1, arr2])
    g2 = NormalFormGame([arr1b, arr2b])
    # Should NOT be considered close
    codeflash_output = not close_normal_form_games(g1, g2) # 833ns -> 8.25μs (89.9% slower)

def test_large_games_different_number_of_players():
    # Large games, but different number of players
    arr1 = np.random.rand(500, 500)
    arr2 = np.random.rand(500, 500)
    arr3 = np.random.rand(500, 500)
    g1 = NormalFormGame([arr1, arr2])
    g2 = NormalFormGame([arr1, arr2, arr3])
    # Should NOT be considered close
    codeflash_output = not close_normal_form_games(g1, g2) # 208ns -> 208ns (0.000% faster)
# 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-close_normal_form_games-mj9prtk2 and push.