⚡️ Speed up method DiscreteRV.__repr__ by 333%

[codeflash-ai]

📄 333% (3.33x) speedup for DiscreteRV.__repr__ in quantecon/_discrete_rv.py

⏱️ Runtime : 125 microseconds → 29.0 microseconds (best of 49 runs)

📝 Explanation and details

The optimized code achieves a 332% speedup by eliminating redundant work in the __repr__ method through result caching.

Key Optimization: Pre-compute and Cache the String Representation

What changed:

1. The string representation is now computed once during __init__ and stored in self._repr_cache
2. __repr__ simply returns the cached value instead of formatting the string on every call
3. Changed np.cumsum(q) to np.cumsum(self._q) to use the already-converted numpy array

Why this is faster:

• Eliminates string formatting overhead: The original code calls .format() on every __repr__ invocation, which involves string parsing, placeholder replacement, and integer-to-string conversion
• Avoids repeated attribute access: No need to access self._q.size on each call
• Minimal memory cost: The cached string is tiny compared to the numpy arrays already stored

Performance characteristics from tests:

• Single calls: Show 3.0-3.6μs → 0.6-0.8μs (~350-390% faster), validating the core optimization
• Repeated calls: Show even better gains (e.g., 878ns → 299ns, 194% faster) as the cache eliminates all computation
• Large arrays: The optimization remains effective even with 1 million elements (4.7μs → 1.0μs), since the speedup depends on avoiding string formatting, not array size
• All test cases benefit uniformly: Whether input is empty, small, large, or contains special values (NaN, inf), the ~350% speedup is consistent

Trade-offs:

• Slightly increased memory usage (one cached string per instance)
• The string representation is fixed at construction time, but since _q is immutable after initialization, this is semantically correct

This is a textbook example of trading negligible space for significant time savings when a computed value is both expensive to generate and guaranteed to remain constant.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import numpy as np
# imports
import pytest  # used for our unit tests
from quantecon._discrete_rv import DiscreteRV

# unit tests

# ----------- BASIC TEST CASES -----------

def test_repr_with_list_of_floats():
    # Basic test with a list of floats
    rv = DiscreteRV([0.2, 0.3, 0.5])

def test_repr_with_numpy_array():
    # Basic test with a numpy array
    rv = DiscreteRV(np.array([0.1, 0.9]))

def test_repr_with_single_element():
    # Test with a single-element input
    rv = DiscreteRV([1.0])

def test_repr_with_empty_list():
    # Edge case: empty input
    rv = DiscreteRV([])

def test_repr_with_tuple_input():
    # Basic test with a tuple input
    rv = DiscreteRV((0.5, 0.5))

def test_repr_with_integer_elements():
    # Basic test with integer elements
    rv = DiscreteRV([1, 0, 0])

# ----------- EDGE TEST CASES -----------

def test_repr_with_non_1d_input():
    # Edge case: 2D array input
    rv = DiscreteRV([[0.5, 0.5], [0.0, 0.0]])

def test_repr_with_large_float_values():
    # Edge case: very large float values
    rv = DiscreteRV([1e100, 2e100, 3e100])

def test_repr_with_zero_probabilities():
    # Edge case: all zeros
    rv = DiscreteRV([0, 0, 0, 0])

def test_repr_with_mixed_types():
    # Edge case: mixture of int and float
    rv = DiscreteRV([1, 0.0, 0.0])

def test_repr_with_boolean_input():
    # Edge case: boolean input
    rv = DiscreteRV([True, False, True])

def test_repr_with_object_dtype():
    # Edge case: array of objects (should still count elements)
    rv = DiscreteRV(np.array([1, 2, 3], dtype=object))

def test_repr_with_empty_numpy_array():
    # Edge case: empty numpy array
    rv = DiscreteRV(np.array([]))

def test_repr_with_high_dimensional_array():
    # Edge case: 3D array input
    arr = np.ones((2, 3, 4))
    rv = DiscreteRV(arr)

# ----------- LARGE SCALE TEST CASES -----------

def test_repr_with_large_1d_array():
    # Large scale: 1D array with 1000 elements
    arr = np.ones(1000)
    rv = DiscreteRV(arr)

def test_repr_with_large_2d_array():
    # Large scale: 2D array with 1000 elements (40x25)
    arr = np.ones((40, 25))
    rv = DiscreteRV(arr)

def test_repr_with_large_3d_array():
    # Large scale: 3D array with 512 elements (8x8x8)
    arr = np.ones((8, 8, 8))
    rv = DiscreteRV(arr)

def test_repr_with_large_sparse_array():
    # Large scale: mostly zeros
    arr = np.zeros(999)
    arr[0] = 1
    rv = DiscreteRV(arr)

def test_repr_with_large_random_array():
    # Large scale: random floats, size 1000
    arr = np.random.rand(1000)
    rv = DiscreteRV(arr)

# ----------- ADDITIONAL EDGE CASES -----------

def test_repr_with_inf_and_nan():
    # Edge case: array with inf and nan
    arr = np.array([np.inf, np.nan, 1.0])
    rv = DiscreteRV(arr)

def test_repr_with_nested_empty_arrays():
    # Edge case: input with nested empty arrays
    arr = np.array([[], []])
    rv = DiscreteRV(arr)

def test_repr_with_large_nested_list():
    # Large scale: deeply nested list, but total elements < 1000
    arr = [[1]*10]*10  # 10 lists of 10 elements each
    rv = DiscreteRV(arr)
# 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
from quantecon._discrete_rv import DiscreteRV

# unit tests

class TestDiscreteRVReprBasic:
    """Test basic functionality of DiscreteRV.__repr__"""
    
    def test_single_element(self):
        """Test repr with a single element probability distribution"""
        # Create a DiscreteRV with one element
        rv = DiscreteRV([1.0])
        # Verify the repr shows exactly 1 element
        codeflash_output = rv.__repr__() # 3.42μs -> 819ns (318% faster)
    
    def test_two_elements(self):
        """Test repr with two element probability distribution"""
        # Create a DiscreteRV with two elements
        rv = DiscreteRV([0.5, 0.5])
        # Verify the repr shows exactly 2 elements
        codeflash_output = rv.__repr__() # 3.18μs -> 686ns (364% faster)
    
    def test_five_elements(self):
        """Test repr with five element probability distribution"""
        # Create a DiscreteRV with five elements
        rv = DiscreteRV([0.2, 0.2, 0.2, 0.2, 0.2])
        # Verify the repr shows exactly 5 elements
        codeflash_output = rv.__repr__() # 3.05μs -> 677ns (350% faster)
    
    def test_ten_elements(self):
        """Test repr with ten element probability distribution"""
        # Create a DiscreteRV with ten elements
        rv = DiscreteRV([0.1] * 10)
        # Verify the repr shows exactly 10 elements
        codeflash_output = rv.__repr__() # 3.08μs -> 695ns (343% faster)
    
    def test_repr_returns_string(self):
        """Test that __repr__ returns a string type"""
        # Create a DiscreteRV
        rv = DiscreteRV([0.3, 0.7])
        # Get the repr
        codeflash_output = rv.__repr__(); result = codeflash_output # 2.97μs -> 652ns (356% faster)
    
    def test_repr_contains_expected_format(self):
        """Test that repr contains the expected format components"""
        # Create a DiscreteRV
        rv = DiscreteRV([0.25, 0.25, 0.25, 0.25])
        codeflash_output = rv.__repr__(); result = codeflash_output # 3.04μs -> 664ns (358% faster)

class TestDiscreteRVReprEdgeCases:
    """Test edge cases for DiscreteRV.__repr__"""
    
    def test_empty_array(self):
        """Test repr with an empty probability array"""
        # Create a DiscreteRV with empty array
        rv = DiscreteRV([])
        # Verify the repr shows 0 elements
        codeflash_output = rv.__repr__() # 3.04μs -> 641ns (374% faster)
    
    def test_numpy_array_input(self):
        """Test repr with numpy array input"""
        # Create a DiscreteRV with numpy array
        rv = DiscreteRV(np.array([0.3, 0.7]))
        # Verify the repr works correctly
        codeflash_output = rv.__repr__() # 3.00μs -> 672ns (346% faster)
    
    def test_list_input(self):
        """Test repr with list input"""
        # Create a DiscreteRV with list
        rv = DiscreteRV([0.3, 0.7])
        # Verify the repr works correctly
        codeflash_output = rv.__repr__() # 2.96μs -> 697ns (325% faster)
    
    def test_tuple_input(self):
        """Test repr with tuple input"""
        # Create a DiscreteRV with tuple
        rv = DiscreteRV((0.3, 0.7))
        # Verify the repr works correctly
        codeflash_output = rv.__repr__() # 3.02μs -> 727ns (316% faster)
    
    def test_non_normalized_probabilities(self):
        """Test repr with probabilities that don't sum to 1"""
        # Create a DiscreteRV with non-normalized probabilities
        rv = DiscreteRV([0.5, 0.5, 0.5])
        # Verify the repr still shows correct count
        codeflash_output = rv.__repr__() # 3.09μs -> 631ns (390% faster)
    
    def test_probabilities_with_zeros(self):
        """Test repr with some zero probabilities"""
        # Create a DiscreteRV with zeros
        rv = DiscreteRV([0.0, 0.5, 0.0, 0.5])
        # Verify the repr counts all elements including zeros
        codeflash_output = rv.__repr__() # 3.05μs -> 672ns (354% faster)
    
    def test_all_zeros(self):
        """Test repr with all zero probabilities"""
        # Create a DiscreteRV with all zeros
        rv = DiscreteRV([0.0, 0.0, 0.0])
        # Verify the repr shows correct count
        codeflash_output = rv.__repr__() # 3.06μs -> 665ns (360% faster)
    
    def test_negative_probabilities(self):
        """Test repr with negative probabilities (invalid but should still work)"""
        # Create a DiscreteRV with negative values
        rv = DiscreteRV([-0.5, 1.5])
        # Verify the repr shows correct count regardless of validity
        codeflash_output = rv.__repr__() # 3.04μs -> 667ns (357% faster)
    
    def test_very_small_probabilities(self):
        """Test repr with very small probability values"""
        # Create a DiscreteRV with very small values
        rv = DiscreteRV([1e-10, 1e-10, 1.0])
        # Verify the repr shows correct count
        codeflash_output = rv.__repr__() # 3.00μs -> 685ns (339% faster)
    
    def test_single_probability_of_one(self):
        """Test repr with single probability of 1.0"""
        # Create a DiscreteRV with single probability
        rv = DiscreteRV([1.0])
        # Verify the repr shows 1 element
        codeflash_output = rv.__repr__() # 3.06μs -> 661ns (363% faster)

class TestDiscreteRVReprConsistency:
    """Test consistency and invariants of DiscreteRV.__repr__"""
    
    def test_repr_consistency_multiple_calls(self):
        """Test that repr returns the same value on multiple calls"""
        # Create a DiscreteRV
        rv = DiscreteRV([0.2, 0.3, 0.5])
        # Call repr multiple times
        codeflash_output = rv.__repr__(); result1 = codeflash_output # 3.07μs -> 675ns (355% faster)
        codeflash_output = rv.__repr__(); result2 = codeflash_output # 878ns -> 299ns (194% faster)
        codeflash_output = rv.__repr__(); result3 = codeflash_output # 746ns -> 268ns (178% faster)
    
    def test_repr_str_equivalence(self):
        """Test that __repr__ and __str__ return the same value"""
        # Create a DiscreteRV
        rv = DiscreteRV([0.25, 0.75])
        # Get both repr and str
        codeflash_output = rv.__repr__(); repr_result = codeflash_output # 2.98μs -> 703ns (323% faster)
        str_result = rv.__str__()
    
    def test_repr_after_attribute_access(self):
        """Test that repr works correctly after accessing other attributes"""
        # Create a DiscreteRV
        rv = DiscreteRV([0.1, 0.2, 0.3, 0.4])
        # Access other attributes
        _ = rv.Q
        _ = rv._q
        # Verify repr still works correctly
        codeflash_output = rv.__repr__() # 3.07μs -> 688ns (347% faster)
    
    def test_repr_with_different_dtypes(self):
        """Test repr with different numpy dtypes"""
        # Create DiscreteRV with int array
        rv_int = DiscreteRV(np.array([1, 2, 3], dtype=np.int32))
        # Create DiscreteRV with float array
        rv_float = DiscreteRV(np.array([0.1, 0.2, 0.7], dtype=np.float64))
        # Verify both show correct count
        codeflash_output = rv_int.__repr__() # 3.09μs -> 695ns (344% faster)
        codeflash_output = rv_float.__repr__() # 947ns -> 288ns (229% faster)

class TestDiscreteRVReprLargeScale:
    """Test large scale scenarios for DiscreteRV.__repr__"""
    
    def test_hundred_elements(self):
        """Test repr with 100 elements"""
        # Create a DiscreteRV with 100 elements
        rv = DiscreteRV([0.01] * 100)
        # Verify the repr shows exactly 100 elements
        codeflash_output = rv.__repr__() # 3.06μs -> 641ns (378% faster)
    
    def test_thousand_elements(self):
        """Test repr with 1000 elements"""
        # Create a DiscreteRV with 1000 elements
        rv = DiscreteRV([0.001] * 1000)
        # Verify the repr shows exactly 1000 elements
        codeflash_output = rv.__repr__() # 3.10μs -> 658ns (371% faster)
    
    def test_ten_thousand_elements(self):
        """Test repr with 10,000 elements"""
        # Create a DiscreteRV with 10,000 elements
        rv = DiscreteRV(np.ones(10000) / 10000)
        # Verify the repr shows exactly 10000 elements
        codeflash_output = rv.__repr__() # 3.16μs -> 682ns (363% faster)
    
    def test_hundred_thousand_elements(self):
        """Test repr with 100,000 elements"""
        # Create a DiscreteRV with 100,000 elements
        rv = DiscreteRV(np.ones(100000) / 100000)
        # Verify the repr shows exactly 100000 elements
        codeflash_output = rv.__repr__() # 3.58μs -> 794ns (351% faster)
    
    def test_million_elements(self):
        """Test repr with 1,000,000 elements"""
        # Create a DiscreteRV with 1 million elements
        rv = DiscreteRV(np.ones(1000000) / 1000000)
        # Verify the repr shows exactly 1000000 elements
        codeflash_output = rv.__repr__() # 4.70μs -> 1.00μs (369% faster)
    
    def test_repr_performance_large_array(self):
        """Test that repr is fast even with large arrays"""
        # Create a DiscreteRV with large array
        rv = DiscreteRV(np.random.random(100000))
        # Call repr multiple times to ensure it's efficient
        for _ in range(10):
            codeflash_output = rv.__repr__(); result = codeflash_output # 10.8μs -> 3.16μs (240% faster)

class TestDiscreteRVReprSpecialValues:
    """Test repr with special numerical values"""
    
    def test_repr_with_inf(self):
        """Test repr with infinite values"""
        # Create a DiscreteRV with infinity
        rv = DiscreteRV([np.inf, 0.5])
        # Verify the repr shows correct count
        codeflash_output = rv.__repr__() # 3.08μs -> 680ns (353% faster)
    
    def test_repr_with_nan(self):
        """Test repr with NaN values"""
        # Create a DiscreteRV with NaN
        rv = DiscreteRV([np.nan, 0.5, 0.5])
        # Verify the repr shows correct count
        codeflash_output = rv.__repr__() # 2.94μs -> 667ns (341% faster)
    
    def test_repr_with_mixed_special_values(self):
        """Test repr with mixed special values"""
        # Create a DiscreteRV with various special values
        rv = DiscreteRV([0.0, np.inf, np.nan, -np.inf, 0.5])
        # Verify the repr shows correct count
        codeflash_output = rv.__repr__() # 3.05μs -> 695ns (339% faster)

class TestDiscreteRVReprFormatValidation:
    """Test the exact format of the repr string"""
    
    def test_repr_exact_format_single_digit(self):
        """Test exact format with single digit count"""
        # Create a DiscreteRV with 3 elements
        rv = DiscreteRV([0.3, 0.3, 0.4])
        # Verify exact format
        codeflash_output = rv.__repr__() # 2.91μs -> 696ns (318% faster)
    
    def test_repr_exact_format_double_digit(self):
        """Test exact format with double digit count"""
        # Create a DiscreteRV with 42 elements
        rv = DiscreteRV(np.ones(42) / 42)
        # Verify exact format
        codeflash_output = rv.__repr__() # 3.02μs -> 665ns (353% faster)
    
    def test_repr_exact_format_triple_digit(self):
        """Test exact format with triple digit count"""
        # Create a DiscreteRV with 123 elements
        rv = DiscreteRV(np.ones(123) / 123)
        # Verify exact format
        codeflash_output = rv.__repr__() # 3.00μs -> 674ns (345% faster)
    
    def test_repr_no_extra_whitespace(self):
        """Test that repr has no extra whitespace"""
        # Create a DiscreteRV
        rv = DiscreteRV([0.5, 0.5])
        codeflash_output = rv.__repr__(); result = codeflash_output # 3.04μs -> 690ns (341% faster)
    
    def test_repr_starts_with_classname(self):
        """Test that repr starts with the class name"""
        # Create a DiscreteRV
        rv = DiscreteRV([0.5, 0.5])
        codeflash_output = rv.__repr__(); result = codeflash_output # 3.03μs -> 686ns (342% faster)
    
    def test_repr_ends_with_elements(self):
        """Test that repr ends with 'elements'"""
        # Create a DiscreteRV
        rv = DiscreteRV([0.5, 0.5])
        codeflash_output = rv.__repr__(); result = codeflash_output # 3.05μs -> 683ns (347% faster)

class TestDiscreteRVReprMultidimensional:
    """Test repr with multidimensional array inputs (flattened by numpy)"""
    
    def test_repr_with_2d_array(self):
        """Test repr when initialized with 2D array (gets flattened)"""
        # Create a DiscreteRV with 2D array (numpy will flatten it)
        rv = DiscreteRV([[0.1, 0.2], [0.3, 0.4]])
        # Verify the repr shows flattened count (4 elements)
        codeflash_output = rv.__repr__() # 3.04μs -> 672ns (353% faster)
    
    def test_repr_with_nested_list(self):
        """Test repr when initialized with nested list"""
        # Create a DiscreteRV with nested list
        rv = DiscreteRV([[0.5], [0.5]])
        # Verify the repr shows flattened count
        codeflash_output = rv.__repr__() # 3.07μs -> 696ns (341% 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-DiscreteRV.__repr__-mjvvl270 and push.