⚡️ Speed up function _get_mixed_actions by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for _get_mixed_actions in quantecon/game_theory/vertex_enumeration.py

⏱️ Runtime : 345 microseconds → 311 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves an 11% speedup by reducing overhead in the Numba-compiled function through three key strategies:

Key Optimizations

1. Localized tuple indexing
The original code repeatedly accessed equation_tup[0], equation_tup[1], trans_recips[0], and trans_recips[1] inside loops. The optimized version hoists these into local variables (eq0, eq1, tr0, tr1, last0, last1) at the start. In Numba's nopython mode, this eliminates redundant tuple indexing overhead on every iteration.

2. Loop unrolling and explicit normalization
Instead of iterating over a tuple of (start, stop, skip) parameters, the optimized code manually handles each player block separately. This eliminates the interpreter overhead of unpacking tuples in the loop. Additionally, normalization is changed from in-place slice division (out[start:stop] /= sum_) to an explicit loop with precomputed inverse (inv = 1.0 / sum_), which Numba can optimize more effectively and avoids potential slice operation overhead.

3. Reduced bit manipulation overhead
The code maintains a local copy of labeling_bits as lb and reuses a constant mask = np.uint64(1) instead of recreating it each iteration, reducing per-iteration constant creation overhead.

Impact Analysis

From the function_references, _get_mixed_actions is called within a generator that iterates over potentially many vertex pairs in game theory equilibrium computation. The function sits in the hot path of vertex enumeration, being called once per matching labeling pair. Given that:

• Tests show 5-14% speedups across various input sizes (most consistently 7-11%)
• The test_large_scale_many_iterations test (100 calls) shows 14.6% speedup (183μs → 160μs), confirming cumulative benefits
• Larger action spaces (m=50, n=50; m=10, n=90) maintain 6-8% gains

The optimization is particularly valuable when:

• The equilibrium enumeration involves many vertices (common in games with multiple actions)
• The function is called repeatedly in batch computations
• Players have moderate to large action spaces (n, m > 10), where the reduced per-iteration overhead compounds

The changes preserve exact numerical behavior (all tests pass) while delivering consistent performance gains across edge cases, including extreme coefficient ranges and various bit patterns.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import numpy as np  # used to construct numeric inputs and inspect outputs
# imports
import pytest  # used for our unit tests
from numba import \
    njit  # required because the function under test uses this decorator
from quantecon.game_theory.vertex_enumeration import _get_mixed_actions

# ---------------------------
# Helper to call the implementation
# ---------------------------
def _call_mixed(labeling_bits, equation_tup, trans_recips):
    """
    Call the real implementation. If Numba decorated function exposes the .py_func
    attribute (the pure-Python implementation), use that to avoid JIT overhead
    in the test runtime. Otherwise call the callable directly (it may be the
    compiled dispatcher).
    """
    func = getattr(_get_mixed_actions, "py_func", _get_mixed_actions)
    # Ensure labeling_bits is cast to np.uint64 exactly like the implementation expects.
    return func(np.uint64(labeling_bits), equation_tup, trans_recips)

# ---------------------------
# Auxiliary pure-Python reference implementation
# ---------------------------
def _expected_mixed_actions(labeling_bits, equation_tup, trans_recips):
    """
    Re-implement the algorithm in Python for generating expected outputs.
    This is intentionally a separate implementation (not calling the function
    under test) to detect behavioral changes.
    """
    # cast inputs similarly
    labeling_bits = np.uint64(labeling_bits)
    m = equation_tup[0].shape[0] - 1
    n = equation_tup[1].shape[0] - 1
    out = np.empty(m + n, dtype=np.float64)

    for pl in range(2):
        if pl == 0:
            start = 0
            stop = m
            skip = np.uint64(1)
        else:
            start = m
            stop = m + n
            skip = np.uint64(0)

        sum_ = 0.0
        for i in range(start, stop):
            if (labeling_bits & np.uint64(1)) == skip:
                out[i] = 0.0
            else:
                out[i] = equation_tup[pl][i - start] * trans_recips[pl] - equation_tup[pl][-1]
                sum_ += out[i]
            labeling_bits = labeling_bits >> np.uint64(1)
        if sum_ != 0.0:
            inv = 1.0 / sum_
            for k in range(start, stop):
                out[k] = out[k] * inv

    return out[:m].copy(), out[m:].copy()

def test_basic_all_zero_labeling():
    # Basic scenario: small m=2, n=2 and labeling_bits == 0 (all LSBs zero)
    # Expect player 0 entries to be computed and normalized; player 1 entries
    # should be set to zero due to skip logic.
    eq0 = np.array([2.0, 3.0, 1.0])  # last element is the constant term
    eq1 = np.array([4.0, 6.0, 2.0])
    trans_recips = (0.5, 0.25)  # reciprocals of translations

    # labeling_bits = 0 -> all bits zero
    a0, a1 = _call_mixed(0, (eq0, eq1), trans_recips)

def test_basic_all_one_labeling():
    # Basic scenario: small m=2, n=2 and labeling_bits == all ones in first m+n bits
    # Expect player 0 entries to be zeros (skipped) and player 1 computed and normalized.
    eq0 = np.array([1.0, 2.0, 0.5])
    eq1 = np.array([5.0, 7.0, 1.0])
    trans_recips = (2.0, 0.2)

    # Set first 4 bits to 1 -> labeling_bits = 0b1111 = 15
    a0, a1 = _call_mixed(15, (eq0, eq1), trans_recips)

def test_mixed_labeling_both_players_normalize():
    # Mixed labeling: ensure both players have some computed entries and each block
    # is normalized independently.
    eq0 = np.array([3.0, 1.0, 0.5])    # m = 2 (last element .5)
    eq1 = np.array([4.0, 2.0, 6.0, 1.0])  # n = 3 (last element 1.0)
    trans_recips = (1.0, 0.5)

    # Build labeling bits so that:
    # - For player 0 (m=2): bits LSB-> ... are [0,1] meaning first computed (0), second skipped (1)
    # - For player 1 (n=3): next bits LSB-> ... are [0,0,1] meaning first two computed, last skipped
    # Compose bits in LSB-first order: pl0 i=0 ->bit0, pl0 i=1 ->bit1, then pl1 bits.
    # So bitstring (from LSB) = [0,1,0,0,1] => binary 10010 (MSB->LSB) = reverse -> decimal:
    # We'll construct numerically by shifting.
    bits = 0
    pattern = [0, 1, 0, 0, 1]  # LSB-first
    for i, b in enumerate(pattern):
        if b:
            bits |= (1 << i)

    a0, a1 = _call_mixed(bits, (eq0, eq1), trans_recips)

    # Compute expected using the auxiliary expected implementation
    exp0, exp1 = _expected_mixed_actions(bits, (eq0, eq1), trans_recips)

    # Additionally ensure that when a block had non-zero sum, sum of probabilities is 1.0
    # Determine which blocks had sum != 0 by inspecting expected sums before normalization:
    # If exp sum > 0 then normalized -> sum approx 1.0, else sum approx 0.0
    sum0 = float(np.sum(exp0))
    sum1 = float(np.sum(exp1))
    if any(exp0):
        pass
    if any(exp1):
        pass

def test_sum_zero_no_normalization():
    # Edge case where the computed raw values for a block sum to zero: normalization should be skipped.
    # Construct eq0 such that eq0[i]*tr - eq0[-1] produce [1.0, -1.0] so sum zero.
    eq0 = np.array([3.0, 1.0, 2.0])  # yields [3*1 -2 =1, 1*1 -2 = -1]
    eq1 = np.array([2.0, 5.0, 1.0])  # another block for completeness
    trans_recips = (1.0, 0.5)

    # labeling_bits = 0 so that both entries in player 0 are computed
    a0, a1 = _call_mixed(0, (eq0, eq1), trans_recips)

    # For player 0, raw values are [1, -1] and sum is 0 -> no normalization applied.
    # So returned should be exactly the raw values.
    # Raw values:
    raw0 = np.array([eq0[0] * trans_recips[0] - eq0[-1],
                     eq0[1] * trans_recips[0] - eq0[-1]], dtype=float)

    # Player 1: compute raw values and possibly normalization; compare against expected implementation
    exp0, exp1 = _expected_mixed_actions(0, (eq0, eq1), trans_recips)

def test_zero_m_or_n_dimensions():
    # Edge case where m == 0 (first equation length is 1). The function should handle empty blocks.
    # Let eq0 length = 1 -> m = 0; eq1 length = 3 -> n = 2
    eq0 = np.array([0.0])  # no choices for player 0
    eq1 = np.array([1.0, 2.0, 1.5])  # two choices for player 1 (last is constant)
    trans_recips = (1.0, 0.5)

    # Various labeling bits; only the bits for player 1 matter.
    a0, a1 = _call_mixed(1, (eq0, eq1), trans_recips)

    # Compare against expected implementation for player 1 values
    _, exp1 = _expected_mixed_actions(1, (eq0, eq1), trans_recips)

import numpy as np
# imports
import pytest
from numba import njit
from quantecon.game_theory.vertex_enumeration import _get_mixed_actions

def test_basic_single_action_per_player():
    """
    Test with the simplest case: 1 action per player (m=1, n=1).
    Both players' mixed actions should be normalized probability distributions.
    """
    # Setup: 1 action for each player
    # equation_tup contains equations for both players' polar polytopes
    # Each equation array has shape (m+1,) for player 0 and (n+1,) for player 1
    equation_0 = np.array([1.0, 0.5], dtype=np.float64)  # m=1, shape (2,)
    equation_1 = np.array([1.0, 0.5], dtype=np.float64)  # n=1, shape (2,)
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (2.0, 2.0)
    
    # labeling_bits with no bits set (binary: 00)
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 9.04μs -> 8.55μs (5.67% faster)

def test_basic_two_actions_per_player():
    """
    Test with 2 actions per player (m=2, n=2).
    Verifies correct output shape and normalization.
    """
    # Setup: 2 actions for each player
    equation_0 = np.array([1.0, 1.0, 0.5], dtype=np.float64)  # m=2, shape (3,)
    equation_1 = np.array([1.0, 1.0, 0.5], dtype=np.float64)  # n=2, shape (3,)
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    # labeling_bits = 0 means all bits are 0
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.59μs -> 8.14μs (5.55% faster)

def test_basic_different_number_actions():
    """
    Test with different numbers of actions for each player (m=2, n=3).
    """
    # Setup: 2 actions for player 0, 3 actions for player 1
    equation_0 = np.array([2.0, 2.0, 1.0], dtype=np.float64)  # m=2, shape (3,)
    equation_1 = np.array([3.0, 2.0, 1.0, 0.5], dtype=np.float64)  # n=3, shape (4,)
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.5, 2.0)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.33μs -> 7.78μs (7.07% faster)

def test_pure_strategy_player_0():
    """
    Test case where player 0 plays a pure strategy (one action gets prob 1, others get 0).
    This happens when only one bit is not set in the labeling_bits.
    """
    # Setup: 3 actions for each player
    equation_0 = np.array([1.0, 0.5, 0.5, 0.2], dtype=np.float64)  # m=3
    equation_1 = np.array([1.0, 0.5, 0.5, 0.2], dtype=np.float64)  # n=3
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    # labeling_bits with specific pattern to result in pure strategy for player 0
    # Binary pattern: 000 (for 3 actions, first 3 bits all 0) -> one action active
    labeling_bits = np.uint64(0b0011111)  # First 3 bits are 0 for player 0
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.21μs -> 7.68μs (6.93% faster)

def test_edge_all_bits_set():
    """
    Test when all labeling bits are set to 1.
    This represents a specific labeling pattern.
    """
    # Setup: 2 actions for each player
    equation_0 = np.array([1.0, 1.0, 0.5], dtype=np.float64)
    equation_1 = np.array([1.0, 1.0, 0.5], dtype=np.float64)
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    # All bits set to 1
    labeling_bits = np.uint64(0xFFFFFFFFFFFFFFFF)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.53μs -> 7.95μs (7.31% faster)

def test_edge_zero_bits():
    """
    Test when all labeling bits are 0.
    This is the minimal case with all bits unset.
    """
    equation_0 = np.array([1.0, 1.0, 0.5], dtype=np.float64)
    equation_1 = np.array([1.0, 1.0, 0.5], dtype=np.float64)
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.26μs -> 7.67μs (7.74% faster)

def test_edge_very_small_trans_recips():
    """
    Test with very small trans_recips values (close to zero).
    Should still produce valid normalized distributions.
    """
    equation_0 = np.array([1.0, 1.0, 1.0], dtype=np.float64)
    equation_1 = np.array([1.0, 1.0, 1.0], dtype=np.float64)
    equation_tup = (equation_0, equation_1)
    
    # Very small reciprocals
    trans_recips = (1e-10, 1e-10)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.33μs -> 7.67μs (8.52% faster)

def test_edge_very_large_trans_recips():
    """
    Test with very large trans_recips values.
    Should still produce valid normalized distributions.
    """
    equation_0 = np.array([1.0, 1.0, 1.0], dtype=np.float64)
    equation_1 = np.array([1.0, 1.0, 1.0], dtype=np.float64)
    equation_tup = (equation_0, equation_1)
    
    # Very large reciprocals
    trans_recips = (1e10, 1e10)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.40μs -> 7.73μs (8.67% faster)

def test_edge_negative_equation_coefficients():
    """
    Test with negative coefficients in equations.
    The function should handle negative values correctly.
    """
    equation_0 = np.array([-1.0, 2.0, -0.5], dtype=np.float64)
    equation_1 = np.array([-1.0, 2.0, -0.5], dtype=np.float64)
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.41μs -> 7.80μs (7.77% faster)

def test_edge_single_large_action_count():
    """
    Test with a large number of actions for one player (m=1, n=10).
    """
    equation_0 = np.array([1.0, 0.5], dtype=np.float64)  # m=1
    equation_1 = np.array([1.0] * 10 + [0.5], dtype=np.float64)  # n=10
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.48μs -> 7.86μs (7.88% faster)

def test_edge_mixed_positive_negative_coefficients():
    """
    Test with mixed positive and negative coefficients in equations.
    """
    equation_0 = np.array([2.0, -1.0, 0.5], dtype=np.float64)
    equation_1 = np.array([-0.5, 1.0, 0.25], dtype=np.float64)
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.5, 2.0)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.34μs -> 7.83μs (6.54% faster)

def test_edge_zero_constant_term():
    """
    Test when the constant term (last element) of equations is zero.
    """
    equation_0 = np.array([1.0, 1.0, 0.0], dtype=np.float64)
    equation_1 = np.array([1.0, 1.0, 0.0], dtype=np.float64)
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.09μs -> 7.79μs (3.90% faster)

def test_edge_alternating_bits_pattern():
    """
    Test with an alternating bit pattern (0b0101... or 0b1010...).
    """
    equation_0 = np.array([1.0, 1.0, 1.0, 0.5], dtype=np.float64)  # m=3
    equation_1 = np.array([1.0, 1.0, 1.0, 0.5], dtype=np.float64)  # n=3
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    # Alternating pattern: 0b010101... (for 6 actions total)
    labeling_bits = np.uint64(0b010101)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.45μs -> 7.65μs (10.5% faster)

def test_large_scale_many_actions_per_player():
    """
    Test with a moderate number of actions for each player (m=50, n=50).
    This tests scalability and performance with larger problem sizes.
    """
    m = 50
    n = 50
    
    # Create equation arrays with random coefficients
    equation_0 = np.ones(m + 1, dtype=np.float64) * 0.5
    equation_0[-1] = 0.1  # constant term
    
    equation_1 = np.ones(n + 1, dtype=np.float64) * 0.5
    equation_1[-1] = 0.1  # constant term
    
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.62μs -> 7.97μs (8.10% faster)

def test_large_scale_many_actions_asymmetric():
    """
    Test with asymmetric action counts (m=10, n=90).
    One player has significantly more actions than the other.
    """
    m = 10
    n = 90
    
    equation_0 = np.linspace(0.1, 1.0, m + 1, dtype=np.float64)
    equation_1 = np.linspace(0.1, 1.0, n + 1, dtype=np.float64)
    
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.5, 2.0)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.48μs -> 7.97μs (6.40% faster)

def test_large_scale_varying_bit_patterns():
    """
    Test multiple different bit patterns with moderate-sized action spaces (m=20, n=20).
    """
    m = 20
    n = 20
    
    equation_0 = np.random.RandomState(42).uniform(0.1, 1.0, m + 1).astype(np.float64)
    equation_1 = np.random.RandomState(43).uniform(0.1, 1.0, n + 1).astype(np.float64)
    
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    # Test several different bit patterns
    bit_patterns = [
        np.uint64(0),
        np.uint64(1),
        np.uint64(0xAAAAAAAAAAAAAAAA),  # alternating 1s
        np.uint64(0x5555555555555555),  # alternating 0s
        np.uint64(0xFFFFFFFF00000000),  # first half 1s, second half 0s
    ]
    
    for labeling_bits in bit_patterns:
        codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 17.4μs -> 15.6μs (11.2% faster)

def test_large_scale_extreme_coefficient_ranges():
    """
    Test with coefficients spanning a wide range of magnitudes.
    This tests numerical stability with large-scale disparities.
    """
    m = 30
    n = 30
    
    # Create coefficients with extremely different magnitudes
    equation_0 = np.array(list(np.logspace(-5, 5, m)) + [1e-3], dtype=np.float64)
    equation_1 = np.array(list(np.logspace(-5, 5, n)) + [1e-3], dtype=np.float64)
    
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (1.0, 1.0)
    
    labeling_bits = np.uint64(0)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 9.28μs -> 8.78μs (5.63% faster)

def test_large_scale_many_iterations():
    """
    Test calling the function multiple times with different labelings
    to verify consistent behavior across many invocations.
    """
    m = 15
    n = 15
    
    equation_0 = np.ones(m + 1, dtype=np.float64) * 0.5
    equation_1 = np.ones(n + 1, dtype=np.float64) * 0.5
    
    equation_tup = (equation_0, equation_1)
    trans_recips = (1.0, 1.0)
    
    # Test with 100 different labelings
    for i in range(100):
        labeling_bits = np.uint64(i % (2 ** (m + n)))
        
        codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 183μs -> 160μs (14.6% faster)

def test_large_scale_high_dimensional_asymmetric():
    """
    Test with high-dimensional but asymmetric action spaces (m=5, n=95).
    """
    m = 5
    n = 95
    
    equation_0 = np.random.RandomState(100).normal(0.5, 0.2, m + 1).astype(np.float64)
    equation_1 = np.random.RandomState(101).normal(0.5, 0.2, n + 1).astype(np.float64)
    
    equation_tup = (equation_0, equation_1)
    
    trans_recips = (2.5, 1.8)
    
    labeling_bits = np.uint64(0xAAAAAAAAAAAAAAAA)
    
    codeflash_output = _get_mixed_actions(labeling_bits, equation_tup, trans_recips); result = codeflash_output # 8.85μs -> 8.23μs (7.52% 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-_get_mixed_actions-mkp8bi4n and push.