⚡️ Speed up function _is_data_matrix by 22%

[codeflash-ai]

📄 22% (0.22x) speedup for _is_data_matrix in pymde/preprocess/generic.py

⏱️ Runtime : 114 microseconds → 93.6 microseconds (best of 135 runs)

📝 Explanation and details

The optimization reorders the conditional checks to put the faster isinstance() check first, avoiding expensive sp.issparse() calls in the common case where data is a NumPy array or PyTorch tensor.

Key changes:

• Check reordering: Changed from sp.issparse(data) or isinstance(data, (np.ndarray, torch.Tensor)) to checking isinstance() first with early return
• Short-circuit optimization: When data is a NumPy array or PyTorch tensor (common cases), the function returns immediately without calling sp.issparse()

Why this is faster:

• isinstance() is a native Python operation that's very fast for built-in types like np.ndarray and torch.Tensor
• sp.issparse() is more expensive as it needs to check against multiple scipy sparse matrix types and their inheritance hierarchy
• The test results show this optimization is most effective for NumPy arrays and PyTorch tensors (60-120% faster), which are likely the most common input types

Performance characteristics:

• Excellent for dense matrices: NumPy arrays and PyTorch tensors see 75-120% speedup
• Slight regression for sparse matrices: Scipy sparse matrices are 45-55% slower due to the additional isinstance() check, but this is acceptable since they're likely less common
• Modest improvements for non-matrix types: Other data types see 5-22% improvements due to faster failure path

The 22% overall speedup suggests the workload is dominated by dense matrix inputs where this optimization provides the greatest benefit.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 80 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
import scipy.sparse as sp
import torch
from pymde.preprocess.generic import _is_data_matrix

# unit tests

# -------------------------------
# 1. Basic Test Cases
# -------------------------------

def test_numpy_array():
    # Test with a basic numpy array
    arr = np.array([[1, 2], [3, 4]])
    codeflash_output = _is_data_matrix(arr) # 1.58μs -> 803ns (96.6% faster)

def test_torch_tensor():
    # Test with a basic torch tensor
    tensor = torch.tensor([[1, 2], [3, 4]])
    codeflash_output = _is_data_matrix(tensor) # 1.64μs -> 888ns (84.3% faster)

def test_scipy_csr_matrix():
    # Test with a basic scipy CSR sparse matrix
    csr = sp.csr_matrix([[0, 1], [2, 0]])
    codeflash_output = _is_data_matrix(csr) # 679ns -> 1.51μs (54.9% slower)

def test_scipy_csc_matrix():
    # Test with a basic scipy CSC sparse matrix
    csc = sp.csc_matrix([[0, 1], [2, 0]])
    codeflash_output = _is_data_matrix(csc) # 641ns -> 1.38μs (53.4% slower)

def test_scipy_coo_matrix():
    # Test with a basic scipy COO sparse matrix
    coo = sp.coo_matrix([[0, 1], [2, 0]])
    codeflash_output = _is_data_matrix(coo) # 590ns -> 1.32μs (55.5% slower)

def test_empty_numpy_array():
    # Test with an empty numpy array
    arr = np.array([])
    codeflash_output = _is_data_matrix(arr) # 1.42μs -> 700ns (103% faster)

def test_empty_torch_tensor():
    # Test with an empty torch tensor
    tensor = torch.tensor([])
    codeflash_output = _is_data_matrix(tensor) # 1.60μs -> 832ns (92.5% faster)

def test_empty_scipy_matrix():
    # Test with an empty scipy sparse matrix
    csr = sp.csr_matrix((0, 0))
    codeflash_output = _is_data_matrix(csr) # 650ns -> 1.32μs (50.7% slower)

# -------------------------------
# 2. Edge Test Cases
# -------------------------------

def test_python_list():
    # Test with a regular Python list (should not be a data matrix)
    lst = [[1, 2], [3, 4]]
    codeflash_output = _is_data_matrix(lst) # 1.62μs -> 1.38μs (17.8% faster)

def test_python_tuple():
    # Test with a tuple (should not be a data matrix)
    tup = ((1, 2), (3, 4))
    codeflash_output = _is_data_matrix(tup) # 1.62μs -> 1.34μs (20.9% faster)

def test_python_dict():
    # Test with a dict (should not be a data matrix)
    dct = {0: [1, 2], 1: [3, 4]}
    codeflash_output = _is_data_matrix(dct) # 1.62μs -> 1.53μs (5.69% faster)

def test_integer():
    # Test with an integer (should not be a data matrix)
    codeflash_output = _is_data_matrix(42) # 1.46μs -> 1.28μs (13.8% faster)

def test_float():
    # Test with a float (should not be a data matrix)
    codeflash_output = _is_data_matrix(3.14) # 1.42μs -> 1.35μs (5.79% faster)

def test_string():
    # Test with a string (should not be a data matrix)
    codeflash_output = _is_data_matrix("not a matrix") # 1.45μs -> 1.38μs (5.29% faster)

def test_none():
    # Test with None (should not be a data matrix)
    codeflash_output = _is_data_matrix(None) # 1.39μs -> 1.34μs (3.35% faster)

def test_numpy_scalar():
    # Test with a numpy scalar (should not be a data matrix)
    scalar = np.float64(1.23)
    codeflash_output = _is_data_matrix(scalar) # 1.49μs -> 1.50μs (0.401% slower)

def test_torch_scalar():
    # Test with a torch scalar (should not be a data matrix)
    scalar = torch.tensor(1.23)
    # torch.tensor(1.23) is still a Tensor, so should be True
    codeflash_output = _is_data_matrix(scalar) # 1.59μs -> 893ns (77.9% faster)

def test_numpy_object_array():
    # Test with a numpy array of dtype object (still a numpy array)
    arr = np.array([{'a': 1}, {'b': 2}], dtype=object)
    codeflash_output = _is_data_matrix(arr) # 1.44μs -> 719ns (99.6% faster)

def test_numpy_matrix_class():
    # Test with numpy.matrix (deprecated, but still an ndarray subclass)
    mat = np.matrix([[1, 2], [3, 4]])
    codeflash_output = _is_data_matrix(mat) # 1.38μs -> 752ns (83.4% faster)

def test_scipy_matrix_subclass():
    # Test with a subclass of scipy.sparse matrix
    class MyCSR(sp.csr_matrix):
        pass
    mycsr = MyCSR([[1, 0], [0, 1]])
    codeflash_output = _is_data_matrix(mycsr) # 777ns -> 1.42μs (45.4% slower)

def test_numpy_array_subclass():
    # Test with a subclass of numpy.ndarray
    class MyArray(np.ndarray):
        pass
    arr = np.array([[1, 2], [3, 4]]).view(MyArray)
    codeflash_output = _is_data_matrix(arr) # 1.58μs -> 827ns (91.1% faster)

def test_torch_tensor_subclass():
    # Test with a subclass of torch.Tensor
    class MyTensor(torch.Tensor):
        pass
    # torch.Tensor cannot be directly instantiated, but we can check isinstance
    tensor = torch.tensor([[1, 2], [3, 4]])
    codeflash_output = _is_data_matrix(tensor) # 1.69μs -> 847ns (98.9% faster)

# -------------------------------
# 3. Large Scale Test Cases
# -------------------------------

def test_large_numpy_array():
    # Test with a large numpy array (1000x1000, ~8MB)
    arr = np.ones((1000, 1000), dtype=np.float64)
    codeflash_output = _is_data_matrix(arr) # 2.21μs -> 1.51μs (46.6% faster)

def test_large_torch_tensor():
    # Test with a large torch tensor (1000x1000, ~4MB for float32)
    tensor = torch.ones((1000, 1000), dtype=torch.float32)
    codeflash_output = _is_data_matrix(tensor) # 2.39μs -> 1.47μs (62.3% faster)

def test_large_scipy_csr_matrix():
    # Test with a large sparse matrix (1000x1000, 1000 nonzero elements)
    rows = np.arange(1000)
    cols = np.arange(1000)
    data = np.ones(1000)
    csr = sp.csr_matrix((data, (rows, cols)), shape=(1000, 1000))
    codeflash_output = _is_data_matrix(csr) # 815ns -> 1.49μs (45.3% slower)

def test_large_scipy_coo_matrix():
    # Test with a large COO matrix (1000x1000, 1000 nonzero elements)
    rows = np.arange(1000)
    cols = np.arange(1000)
    data = np.ones(1000)
    coo = sp.coo_matrix((data, (rows, cols)), shape=(1000, 1000))
    codeflash_output = _is_data_matrix(coo) # 813ns -> 1.49μs (45.5% slower)

def test_large_python_list():
    # Test with a large Python list (should not be a data matrix)
    lst = [[1]*1000 for _ in range(1000)]
    codeflash_output = _is_data_matrix(lst) # 3.07μs -> 2.56μs (19.9% faster)

def test_large_python_dict():
    # Test with a large Python dict (should not be a data matrix)
    dct = {i: [1]*10 for i in range(1000)}
    codeflash_output = _is_data_matrix(dct) # 2.31μs -> 2.07μs (11.8% faster)

# -------------------------------
# 4. Additional Edge Cases
# -------------------------------

def test_numpy_array_with_nan_inf():
    # Test with numpy array containing NaN and Inf
    arr = np.array([[np.nan, np.inf], [1, 2]])
    codeflash_output = _is_data_matrix(arr) # 1.41μs -> 738ns (91.1% faster)

def test_torch_tensor_with_nan_inf():
    # Test with torch tensor containing NaN and Inf
    tensor = torch.tensor([[float('nan'), float('inf')], [1, 2]])
    codeflash_output = _is_data_matrix(tensor) # 1.63μs -> 867ns (87.8% faster)

def test_scipy_matrix_with_zero_shape():
    # Test with scipy sparse matrix with zero shape
    csr = sp.csr_matrix((0, 0))
    codeflash_output = _is_data_matrix(csr) # 631ns -> 1.28μs (50.8% slower)

def test_numpy_array_with_zero_shape():
    # Test with numpy array with zero shape
    arr = np.empty((0, 0))
    codeflash_output = _is_data_matrix(arr) # 1.41μs -> 644ns (119% faster)

def test_torch_tensor_with_zero_shape():
    # Test with torch tensor with zero shape
    tensor = torch.empty((0, 0))
    codeflash_output = _is_data_matrix(tensor) # 1.55μs -> 805ns (93.2% faster)

# -------------------------------
# 5. Type Robustness
# -------------------------------

def test_numpy_array_with_str_dtype():
    # Test with numpy array of strings
    arr = np.array([["a", "b"], ["c", "d"]], dtype=str)
    codeflash_output = _is_data_matrix(arr) # 1.35μs -> 655ns (106% faster)

def test_numpy_array_with_bool_dtype():
    # Test with numpy array of bools
    arr = np.array([[True, False], [False, True]], dtype=bool)
    codeflash_output = _is_data_matrix(arr) # 1.22μs -> 648ns (88.3% faster)

def test_torch_tensor_with_bool_dtype():
    # Test with torch tensor of bools
    tensor = torch.tensor([[True, False], [False, True]], dtype=torch.bool)
    codeflash_output = _is_data_matrix(tensor) # 1.46μs -> 697ns (109% faster)

# -------------------------------
# 6. Negative Cases for Similar Types
# -------------------------------

def test_bytes_object():
    # Test with bytes object (should not be a data matrix)
    b = b"not a matrix"
    codeflash_output = _is_data_matrix(b) # 1.62μs -> 1.50μs (7.98% faster)

def test_set_object():
    # Test with set object (should not be a data matrix)
    s = {1, 2, 3}
    codeflash_output = _is_data_matrix(s) # 1.56μs -> 1.39μs (12.4% faster)


def test_function_object():
    # Test with function object (should not be a data matrix)
    def foo(): return 1
    codeflash_output = _is_data_matrix(foo) # 1.53μs -> 1.37μs (11.6% faster)

def test_module_object():
    # Test with a module object (should not be a data matrix)
    import math
    codeflash_output = _is_data_matrix(math) # 1.53μs -> 1.43μs (7.51% 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
import scipy.sparse as sp
import torch
from pymde.preprocess.generic import _is_data_matrix

# unit tests

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

def test_numpy_array_basic():
    # Test with a 1D numpy array
    arr = np.array([1, 2, 3])
    codeflash_output = _is_data_matrix(arr) # 1.15μs -> 655ns (75.3% faster)

def test_numpy_array_2d():
    # Test with a 2D numpy array
    arr = np.array([[1, 2], [3, 4]])
    codeflash_output = _is_data_matrix(arr) # 1.18μs -> 640ns (84.2% faster)

def test_torch_tensor_basic():
    # Test with a 1D torch tensor
    tensor = torch.tensor([1, 2, 3])
    codeflash_output = _is_data_matrix(tensor) # 1.55μs -> 851ns (82.5% faster)

def test_torch_tensor_2d():
    # Test with a 2D torch tensor
    tensor = torch.tensor([[1, 2], [3, 4]])
    codeflash_output = _is_data_matrix(tensor) # 1.49μs -> 811ns (83.4% faster)

def test_scipy_sparse_csr():
    # Test with a scipy CSR sparse matrix
    mat = sp.csr_matrix([[1, 0], [0, 2]])
    codeflash_output = _is_data_matrix(mat) # 648ns -> 1.40μs (53.8% slower)

def test_scipy_sparse_csc():
    # Test with a scipy CSC sparse matrix
    mat = sp.csc_matrix([[1, 0], [0, 2]])
    codeflash_output = _is_data_matrix(mat) # 636ns -> 1.35μs (52.9% slower)

def test_scipy_sparse_coo():
    # Test with a scipy COO sparse matrix
    mat = sp.coo_matrix([[1, 0], [0, 2]])
    codeflash_output = _is_data_matrix(mat) # 613ns -> 1.38μs (55.5% slower)

def test_python_list():
    # Test with a Python list (should not be considered a data matrix)
    lst = [1, 2, 3]
    codeflash_output = _is_data_matrix(lst) # 1.62μs -> 1.40μs (16.1% faster)

def test_python_tuple():
    # Test with a Python tuple (should not be considered a data matrix)
    tpl = (1, 2, 3)
    codeflash_output = _is_data_matrix(tpl) # 1.68μs -> 1.47μs (14.5% faster)

def test_python_dict():
    # Test with a Python dict (should not be considered a data matrix)
    dct = {'a': 1, 'b': 2}
    codeflash_output = _is_data_matrix(dct) # 1.59μs -> 1.50μs (5.65% faster)

def test_integer():
    # Test with a single integer
    val = 42
    codeflash_output = _is_data_matrix(val) # 1.67μs -> 1.39μs (20.1% faster)

def test_float():
    # Test with a single float
    val = 3.14
    codeflash_output = _is_data_matrix(val) # 1.52μs -> 1.31μs (15.6% faster)

def test_string():
    # Test with a string
    val = "not a matrix"
    codeflash_output = _is_data_matrix(val) # 1.47μs -> 1.33μs (9.98% faster)

def test_none():
    # Test with None
    codeflash_output = _is_data_matrix(None) # 1.52μs -> 1.24μs (22.0% faster)

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

def test_empty_numpy_array():
    # Test with an empty numpy array
    arr = np.array([])
    codeflash_output = _is_data_matrix(arr) # 1.23μs -> 642ns (91.3% faster)

def test_empty_torch_tensor():
    # Test with an empty torch tensor
    tensor = torch.tensor([])
    codeflash_output = _is_data_matrix(tensor) # 1.55μs -> 847ns (82.8% faster)

def test_empty_scipy_sparse():
    # Test with an empty scipy sparse matrix
    mat = sp.csr_matrix((0, 0))
    codeflash_output = _is_data_matrix(mat) # 667ns -> 1.30μs (48.7% slower)

def test_numpy_scalar():
    # Test with a numpy scalar (0-d array)
    scalar = np.array(5)
    codeflash_output = _is_data_matrix(scalar) # 1.17μs -> 650ns (80.0% faster)

def test_torch_scalar():
    # Test with a torch scalar (0-d tensor)
    scalar = torch.tensor(5)
    codeflash_output = _is_data_matrix(scalar) # 1.58μs -> 849ns (85.9% faster)

def test_numpy_object_dtype():
    # Test with a numpy array of object dtype
    arr = np.array([{'a': 1}, {'b': 2}], dtype=object)
    codeflash_output = _is_data_matrix(arr) # 1.39μs -> 720ns (93.2% faster)

def test_numpy_structured_dtype():
    # Test with a numpy structured array
    arr = np.array([(1, 2.0)], dtype=[('x', 'i4'), ('y', 'f4')])
    codeflash_output = _is_data_matrix(arr) # 1.30μs -> 671ns (94.0% faster)

def test_numpy_matrix_class():
    # Test with a numpy matrix (deprecated, but still a subclass of ndarray)
    mat = np.matrix([[1, 2], [3, 4]])
    codeflash_output = _is_data_matrix(mat) # 1.41μs -> 717ns (96.0% faster)

def test_torch_tensor_on_cuda():
    # Test with a torch tensor on CUDA (if available)
    if torch.cuda.is_available():
        tensor = torch.tensor([1, 2, 3]).cuda()
        codeflash_output = _is_data_matrix(tensor)

def test_scipy_sparse_with_zero_shape():
    # Test with a scipy sparse matrix with zero shape
    mat = sp.csr_matrix((0, 10))
    codeflash_output = _is_data_matrix(mat) # 695ns -> 1.39μs (49.9% slower)

def test_numpy_array_with_zero_shape():
    # Test with a numpy array with zero shape
    arr = np.empty((0, 10))
    codeflash_output = _is_data_matrix(arr) # 1.46μs -> 739ns (97.2% faster)

def test_torch_tensor_with_zero_shape():
    # Test with a torch tensor with zero shape
    tensor = torch.empty((0, 10))
    codeflash_output = _is_data_matrix(tensor) # 1.53μs -> 778ns (96.1% faster)

def test_numpy_array_subclass():
    # Test with a subclass of numpy ndarray
    class MyArray(np.ndarray):
        pass
    arr = np.array([1, 2, 3]).view(MyArray)
    codeflash_output = _is_data_matrix(arr) # 1.50μs -> 795ns (88.9% faster)

def test_torch_tensor_subclass():
    # Test with a subclass of torch.Tensor
    class MyTensor(torch.Tensor):
        pass
    tensor = torch.tensor([1, 2, 3])
    # torch.Tensor subclassing is tricky, but check base class
    codeflash_output = _is_data_matrix(tensor) # 1.55μs -> 824ns (87.9% faster)

def test_scipy_sparse_subclass():
    # Test with a subclass of scipy.sparse.csr_matrix
    class MyCSR(sp.csr_matrix):
        pass
    mat = MyCSR([[1, 0], [0, 2]])
    codeflash_output = _is_data_matrix(mat) # 696ns -> 1.40μs (50.4% slower)

def test_numpy_array_with_nan_inf():
    # Test with a numpy array containing NaN and Inf
    arr = np.array([np.nan, np.inf, -np.inf])
    codeflash_output = _is_data_matrix(arr) # 1.40μs -> 688ns (103% faster)

def test_torch_tensor_with_nan_inf():
    # Test with a torch tensor containing NaN and Inf
    tensor = torch.tensor([float('nan'), float('inf'), float('-inf')])
    codeflash_output = _is_data_matrix(tensor) # 1.64μs -> 820ns (99.6% faster)

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

def test_large_numpy_array():
    # Test with a large numpy array (1000 x 1000 floats, ~8MB)
    arr = np.ones((1000, 1000), dtype=np.float64)
    codeflash_output = _is_data_matrix(arr) # 2.38μs -> 1.47μs (61.5% faster)

def test_large_torch_tensor():
    # Test with a large torch tensor (1000 x 1000 floats, ~4MB)
    tensor = torch.ones((1000, 1000), dtype=torch.float32)
    codeflash_output = _is_data_matrix(tensor) # 2.44μs -> 1.46μs (66.8% faster)

def test_large_scipy_sparse_matrix():
    # Test with a large sparse matrix (1000 x 1000, 1% nonzero)
    rows = np.random.randint(0, 1000, size=10000)
    cols = np.random.randint(0, 1000, size=10000)
    data = np.ones(10000)
    mat = sp.coo_matrix((data, (rows, cols)), shape=(1000, 1000))
    codeflash_output = _is_data_matrix(mat) # 1.01μs -> 1.59μs (36.7% slower)

def test_large_python_list():
    # Test with a large Python list (should not be a data matrix)
    lst = [1] * 1000
    codeflash_output = _is_data_matrix(lst) # 1.83μs -> 1.61μs (13.9% faster)

def test_large_python_dict():
    # Test with a large Python dict (should not be a data matrix)
    dct = {i: i for i in range(1000)}
    codeflash_output = _is_data_matrix(dct) # 1.76μs -> 1.49μs (18.2% faster)

def test_large_string():
    # Test with a large string (should not be a data matrix)
    val = "x" * 10000
    codeflash_output = _is_data_matrix(val) # 1.55μs -> 1.43μs (8.54% faster)

# ---------------------- Negative/Mutation Cases ----------------------

def test_numpy_array_like_object():
    # Test with an object that looks like a numpy array but isn't
    class FakeArray:
        def __init__(self):
            self.shape = (10, 10)
    fake = FakeArray()
    codeflash_output = _is_data_matrix(fake) # 1.67μs -> 1.48μs (12.3% faster)

def test_torch_tensor_like_object():
    # Test with an object that looks like a torch tensor but isn't
    class FakeTensor:
        def __init__(self):
            self.size = (10, 10)
    fake = FakeTensor()
    codeflash_output = _is_data_matrix(fake) # 1.70μs -> 1.54μs (10.8% faster)

def test_scipy_sparse_like_object():
    # Test with an object that looks like a scipy sparse matrix but isn't
    class FakeSparse:
        def __init__(self):
            self.nnz = 0
    fake = FakeSparse()
    codeflash_output = _is_data_matrix(fake) # 1.52μs -> 1.42μs (7.70% 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-_is_data_matrix-mgsr5ql6 and push.