⚡️ Speed up function hh_payoff_player by 72%

[codeflash-ai]

📄 72% (0.72x) speedup for hh_payoff_player in quantecon/game_theory/polymatrix_game.py

⏱️ Runtime : 518 milliseconds → 300 milliseconds (best of 13 runs)

📝 Explanation and details

The optimized code achieves a 72% speedup by eliminating the expensive nested list comprehension with np.vstack that dominated the original implementation (64.1% of runtime).

Key Optimizations:

1. Pre-allocated Arrays: Instead of building hh_actions_and_payoffs through repeated np.vstack and np.hstack calls (which create many intermediate arrays), the optimized version pre-allocates hh_actions and combined_payoffs arrays upfront based on computed dimensions.

2. Vectorized One-Hot Encoding: The original code used np.eye(nfg.nums_actions[p])[action_combination[p]] inside a list comprehension for each action combination. The optimized version uses advanced indexing hh_actions[np.arange(n_combinations), col_offset + actions_p] = 1.0 to set all one-hot values in a single vectorized operation per player.

3. Direct Array Construction: Converting product(*action_ranges) to a NumPy array upfront enables efficient column-wise slicing (action_combinations[:, p]) instead of repeatedly unpacking tuples in list comprehensions.

Performance Analysis:

The line profiler shows the critical bottleneck moved from array construction (64.1% → 0.8%) to the lstsq call (35% → 97% of remaining time), which is unavoidable. Test results show consistent speedups:

• Small games (2-3 players): 4-14% faster
• Large games (4-5 players, many actions): 172-196% faster
• Maximum scale test (1000 actions): 72% faster - matching the overall speedup

Impact on Workloads:

The function_references show this function is called in a hot path within PolymatrixGame.from_nf(), which iterates over all players and their actions (nested loops). For a game with N players and A actions each, this means N×A calls to hh_payoff_player. The optimization's impact scales quadratically with game size, making it particularly valuable for the library's typical use cases involving multi-player games with numerous actions.

The optimizations are most effective for games with many players and actions (see the 196% speedup for 5-player games), while maintaining correctness for edge cases like single-action scenarios.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from itertools import product

# function to test
import numpy as np
# imports
import pytest
from quantecon.game_theory.polymatrix_game import hh_payoff_player

class Player:
    """
    Minimal Player class for testing.
    """
    def __init__(self, payoff_array):
        self.payoff_array = payoff_array

class NormalFormGame:
    """
    Minimal NormalFormGame class for testing.
    """
    def __init__(self, payoff_arrays):
        self.N = len(payoff_arrays)  # number of players
        self.nums_actions = [arr.shape[0] for arr in payoff_arrays]
        self.players = [Player(arr) for arr in payoff_arrays]
from quantecon.game_theory.polymatrix_game import hh_payoff_player

# unit tests

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

def test_2p2a_identity_polymatrix():
    # 2 players, 2 actions each, simple additive payoffs (polymatrix)
    # Player 0 payoff: payoff_array[a0, a1] = a0 + a1
    payoff0 = np.fromfunction(lambda a0, a1: a0 + a1, (2, 2), dtype=int)
    payoff1 = np.fromfunction(lambda a0, a1: a0 - a1, (2, 2), dtype=int)
    nfg = NormalFormGame([payoff0, payoff1])

    # For player 0, action 1: payoff is always 1 + a1
    codeflash_output = hh_payoff_player(nfg, 0, 1); result = codeflash_output # 288μs -> 275μs (4.64% faster)

    # For player 1, action 0: payoff is a0 - 0 = a0
    codeflash_output = hh_payoff_player(nfg, 1, 0); result1 = codeflash_output # 160μs -> 149μs (7.21% faster)

def test_2p2a_nonpolymatrix_least_squares():
    # 2 players, 2 actions each, non-polymatrix game
    # Player 0 payoff: payoff_array[a0, a1] = a0 * a1
    payoff0 = np.fromfunction(lambda a0, a1: a0 * a1, (2,2), dtype=int)
    payoff1 = np.fromfunction(lambda a0, a1: a0 + 2*a1, (2,2), dtype=int)
    nfg = NormalFormGame([payoff0, payoff1])

    # Should not assert if is_polymatrix=False
    codeflash_output = hh_payoff_player(nfg, 0, 1, is_polymatrix=False); result = codeflash_output # 121μs -> 112μs (7.70% faster)

def test_3p2a_polymatrix():
    # 3 players, 2 actions each, polymatrix: payoff = sum of all actions
    payoff0 = np.fromfunction(lambda a0, a1, a2: a0 + a1 + a2, (2,2,2), dtype=int)
    payoff1 = np.fromfunction(lambda a0, a1, a2: a0 - a1 + a2, (2,2,2), dtype=int)
    payoff2 = np.fromfunction(lambda a0, a1, a2: a0 + a1 - a2, (2,2,2), dtype=int)
    nfg = NormalFormGame([payoff0, payoff1, payoff2])

    # For player 0, action 1: payoff = 1 + a1 + a2
    codeflash_output = hh_payoff_player(nfg, 0, 1); result = codeflash_output # 236μs -> 206μs (14.3% faster)
    # Should return {(1,0): 0.0, (1,1): 1.0, (2,0): 0.0, (2,1): 1.0}
    expected_keys = {(1,0), (1,1), (2,0), (2,1)}

# ---- Edge Test Cases ----

def test_one_action_per_player():
    # 2 players, each has only one action
    payoff0 = np.zeros((1,1))
    payoff1 = np.zeros((1,1))
    nfg = NormalFormGame([payoff0, payoff1])
    codeflash_output = hh_payoff_player(nfg, 0, 0); result = codeflash_output # 187μs -> 193μs (3.09% slower)

def test_invalid_player_number():
    # 2 players, 2 actions each
    payoff0 = np.zeros((2,2))
    payoff1 = np.zeros((2,2))
    nfg = NormalFormGame([payoff0, payoff1])
    # Player number out of bounds
    with pytest.raises(IndexError):
        hh_payoff_player(nfg, 2, 0) # 6.03μs -> 17.0μs (64.5% slower)

def test_invalid_action_number():
    # 2 players, 2 actions each
    payoff0 = np.zeros((2,2))
    payoff1 = np.zeros((2,2))
    nfg = NormalFormGame([payoff0, payoff1])
    # Action number out of bounds
    with pytest.raises(IndexError):
        hh_payoff_player(nfg, 0, 2) # 26.5μs -> 50.7μs (47.9% slower)

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

def test_large_polymatrix_game():
    # 5 players, 3 actions each, polymatrix: payoff = sum of all actions
    shape = (3,)*5
    payoff_arrays = [np.fromfunction(lambda *a: sum(a), shape, dtype=int) for _ in range(5)]
    nfg = NormalFormGame(payoff_arrays)
    # For player 0, action 2
    codeflash_output = hh_payoff_player(nfg, 0, 2); result = codeflash_output # 1.62ms -> 547μs (196% faster)
    # Check that the payoffs for each other player are 0, 1, 2
    for p in range(1,5):
        for a in range(3):
            key = (p,a) if p != 0 else (0,a)

def test_large_nonpolymatrix_game():
    # 4 players, 4 actions each, non-polymatrix (product of actions)
    shape = (4,)*4
    payoff0 = np.fromfunction(lambda a0, a1, a2, a3: a0*a1*a2*a3, shape, dtype=int)
    payoff1 = np.fromfunction(lambda a0, a1, a2, a3: a0+a1+a2+a3, shape, dtype=int)
    payoff2 = np.fromfunction(lambda a0, a1, a2, a3: a0*a1 + a2*a3, shape, dtype=int)
    payoff3 = np.fromfunction(lambda a0, a1, a2, a3: a0-a1+a2-a3, shape, dtype=int)
    nfg = NormalFormGame([payoff0, payoff1, payoff2, payoff3])
    # Should not raise error with is_polymatrix=False
    codeflash_output = hh_payoff_player(nfg, 0, 3, is_polymatrix=False); result = codeflash_output # 1.14ms -> 417μs (172% faster)

def test_maximum_actions():
    # 2 players, 1000 actions each, polymatrix: payoff = a0 + a1
    payoff0 = np.fromfunction(lambda a0, a1: a0 + a1, (1000,1000), dtype=int)
    payoff1 = np.fromfunction(lambda a0, a1: a0 - a1, (1000,1000), dtype=int)
    nfg = NormalFormGame([payoff0, payoff1])
    # For player 0, action 999
    codeflash_output = hh_payoff_player(nfg, 0, 999); result = codeflash_output # 514ms -> 298ms (72.3% faster)
    # The payoffs should be 0, 1, ..., 999
    for a in range(1000):
        pass
# 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-hh_payoff_player-mjw5cd12 and push.