⚡️ Speed up method DiGraph.subgraph by 54%

[codeflash-ai]

📄 54% (0.54x) speedup for DiGraph.subgraph in quantecon/_graph_tools.py

⏱️ Runtime : 15.4 milliseconds → 10.0 milliseconds (best of 89 runs)

📝 Explanation and details

The optimized code achieves a 54% speedup through two key optimizations:

1. Avoiding Unnecessary Copy in __init__ (copy=False)

The original code calls sparse.csr_matrix(adj_matrix, dtype=dtype), which by default creates a copy of the input data. The optimization adds copy=False to avoid this unnecessary allocation when adj_matrix is already a compatible sparse matrix. This reduces memory allocation overhead during graph construction.

2. Replacing np.ix_() with Direct Slicing in subgraph

The original uses self.csgraph[np.ix_(nodes, nodes)] to extract a subgraph, which creates intermediate index arrays for fancy indexing. The optimized version:

• Converts nodes to a numpy array once: nodes = np.asarray(nodes)
• Uses chained slicing: self.csgraph[nodes][:, nodes]

This approach is faster because:

• np.ix_() constructs broadcasting-compatible index arrays, adding overhead
• Direct slicing [nodes][:, nodes] performs row selection then column selection sequentially, which is more efficient for CSR matrices (row-major format)
• For sparse matrices, row slicing is particularly fast in CSR format

Impact on Test Cases

The optimization excels with:

• Large subgraphs: 358-416% speedup for 250+ node extractions (test_large_scale_subgraph_preserves_sparsity_and_counts, test_subgraph_performance_large_extraction)
• Sparse graphs: 66-127% speedup on chain-structured graphs where CSR format shines
• Random node selections: 224% speedup on random graphs due to reduced indexing overhead
• Small subgraphs: 10-16% consistent improvement across basic cases

The optimizations are particularly effective when subgraph() is called repeatedly or on large graphs, making them valuable for algorithms that perform frequent graph decompositions or analyses on substructures.

✅ Correctness verification report:

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

⚙️ Click to see Existing Unit Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_graph_tools.py::test_node_labels_subgraph	181μs	155μs	17.1%✅
test_graph_tools.py::test_subgraph	177μs	151μs	17.1%✅
test_graph_tools.py::test_subgraph_weighted	178μs	153μs	16.4%✅

🌀 Click to see Generated Regression Tests

import numpy as np  # used to build adjacency matrices and node index arrays
# imports
import pytest  # used for our unit tests
from quantecon._graph_tools import DiGraph
from scipy import \
    sparse  # used to construct sparse matrices for large-scale tests

def test_init_raises_for_non_square_matrix():
    # Create a non-square adjacency matrix (2 x 3)
    adj = np.array([[1, 0, 0],
                    [0, 1, 0]])
    # Expect ValueError complaining about non-square input
    with pytest.raises(ValueError) as excinfo:
        DiGraph(adj)  # constructor should validate squareness

def test_subgraph_basic_unweighted_no_labels():
    # Build a small 3x3 unweighted adjacency (bool) matrix
    adj = np.array([[0, 1, 0],
                    [0, 0, 1],
                    [1, 0, 0]], dtype=bool)
    # Instantiate DiGraph without labels and unweighted (default)
    G = DiGraph(adj)

    # Take a subgraph consisting of nodes 0 and 2
    nodes = [0, 2]
    codeflash_output = G.subgraph(nodes); H = codeflash_output # 179μs -> 153μs (16.9% faster)

    # The adjacency of the subgraph should equal the sliced adjacency from original
    expected = G.csgraph[np.ix_(nodes, nodes)].toarray().tolist()
    actual = H.csgraph.toarray().tolist()

def test_subgraph_empty_nodes_returns_empty_graph():
    # Create a 2x2 adjacency matrix
    adj = np.array([[1, 0],
                    [0, 1]], dtype=bool)
    G = DiGraph(adj)

    # Provide an empty list of nodes (edge case)
    nodes = np.array([], dtype=int)
    codeflash_output = G.subgraph(nodes); H = codeflash_output # 105μs -> 105μs (0.199% faster)

def test_subgraph_out_of_range_node_index_raises_index_error():
    # Create a small 2x2 graph
    adj = np.array([[0, 1],
                    [1, 0]], dtype=bool)
    G = DiGraph(adj)

    # Choose a node index that is out of bounds (e.g., 5)
    nodes = [0, 5]
    # Expect an IndexError when slicing with out-of-range indices
    with pytest.raises(IndexError):
        G.subgraph(nodes) # 30.8μs -> 17.9μs (72.2% faster)

def test_subgraph_non_integer_indices_raise_index_error():
    # Create a 3x3 graph
    adj = np.eye(3, dtype=int)
    G = DiGraph(adj, weighted=True)

    # Provide non-integer indices (floats) - indexing must fail
    nodes = [0.0, 1.0]
    # Numpy / sparse indexing with float indices raises IndexError
    with pytest.raises(IndexError):
        G.subgraph(nodes)

def test_large_scale_subgraph_preserves_sparsity_and_counts():
    # Build a moderately large sparse graph (size 500 x 500) under the 1000 limit
    n = 500  # total nodes
    # Create a sparse adjacency that has ones on the diagonal and on the first upper diagonal
    # Use scipy.sparse to avoid creating a huge dense matrix
    main_diag = np.ones(n, dtype=int)
    upper_diag = np.ones(n - 1, dtype=int)
    A = sparse.diags([main_diag, upper_diag], [0, 1], shape=(n, n), format='csr')

    # Instantiate DiGraph using the sparse matrix directly; mark weighted=True to preserve dtype
    G = DiGraph(A, weighted=True)

    # Select every other node to form a medium-sized subgraph (~250 nodes)
    nodes = np.arange(0, n, 2, dtype=int)  # deterministic selection
    codeflash_output = G.subgraph(nodes); H = codeflash_output # 839μs -> 162μs (416% faster)

    # The number of stored non-zeros in the subgraph should match slicing the original
    expected_slice = G.csgraph[np.ix_(nodes, nodes)]

    # Sample a small deterministic property: diagonal ones are preserved for the selected indices
    # Check that every selected node corresponds to a 1 on the diagonal in the subgraph
    diag = H.csgraph.diagonal()
# 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
import pytest
from quantecon._graph_tools import DiGraph
from scipy import sparse

class TestDiGraphSubgraphBasic:
    """Basic test cases for DiGraph.subgraph functionality."""

    def test_subgraph_single_node(self):
        """Test subgraph with a single node."""
        adj_matrix = np.array([[1, 1, 0],
                               [0, 1, 1],
                               [1, 0, 1]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph containing only node 0
        codeflash_output = g.subgraph(np.array([0])); subgraph = codeflash_output # 163μs -> 146μs (11.9% faster)

    def test_subgraph_two_nodes(self):
        """Test subgraph with two nodes."""
        adj_matrix = np.array([[1, 1, 0],
                               [0, 1, 1],
                               [1, 0, 1]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph containing nodes 0 and 1
        codeflash_output = g.subgraph(np.array([0, 1])); subgraph = codeflash_output # 171μs -> 149μs (14.7% faster)

    def test_subgraph_all_nodes(self):
        """Test subgraph containing all nodes (should be equivalent to original)."""
        adj_matrix = np.array([[1, 1, 0],
                               [0, 1, 1],
                               [1, 0, 1]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph containing all nodes
        codeflash_output = g.subgraph(np.array([0, 1, 2])); subgraph = codeflash_output # 172μs -> 150μs (14.9% faster)

    def test_subgraph_preserves_edge_structure(self):
        """Test that subgraph preserves the edge structure correctly."""
        adj_matrix = np.array([[0, 1, 0, 1],
                               [1, 0, 1, 0],
                               [0, 1, 0, 1],
                               [1, 0, 1, 0]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with nodes 1 and 3
        codeflash_output = g.subgraph(np.array([1, 3])); subgraph = codeflash_output # 161μs -> 151μs (6.77% faster)

    def test_subgraph_with_node_labels(self):
        """Test that subgraph correctly preserves node labels."""
        adj_matrix = np.array([[1, 1, 0],
                               [0, 1, 1],
                               [1, 0, 1]])
        node_labels = np.array(['a', 'b', 'c'])
        g = DiGraph(adj_matrix, node_labels=node_labels)
        
        # Extract subgraph containing nodes 0 and 2
        codeflash_output = g.subgraph(np.array([0, 2])); subgraph = codeflash_output # 175μs -> 155μs (13.1% faster)

    def test_subgraph_with_numeric_node_labels(self):
        """Test subgraph with numeric node labels."""
        adj_matrix = np.array([[0, 1, 1],
                               [1, 0, 1],
                               [1, 1, 0]])
        node_labels = np.array([10, 20, 30])
        g = DiGraph(adj_matrix, node_labels=node_labels)
        
        # Extract subgraph containing nodes 1 and 2
        codeflash_output = g.subgraph(np.array([1, 2])); subgraph = codeflash_output # 174μs -> 153μs (13.9% faster)

    def test_subgraph_without_node_labels(self):
        """Test that subgraph without node labels returns None for node_labels."""
        adj_matrix = np.array([[1, 1, 0],
                               [0, 1, 1],
                               [1, 0, 1]])
        g = DiGraph(adj_matrix, node_labels=None)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([0, 1])); subgraph = codeflash_output # 171μs -> 147μs (15.9% faster)

    def test_subgraph_returns_digraph(self):
        """Test that subgraph returns a DiGraph instance."""
        adj_matrix = np.array([[1, 0], [0, 1]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([0, 1])); subgraph = codeflash_output # 170μs -> 147μs (15.1% faster)

    def test_subgraph_weighted_preservation(self):
        """Test that subgraph preserves weighted adjacency values."""
        adj_matrix = np.array([[1.5, 2.0, 0.0],
                               [0.0, 3.5, 1.5],
                               [2.5, 0.0, 0.5]])
        g = DiGraph(adj_matrix, weighted=True)
        
        # Extract subgraph containing nodes 0 and 2
        codeflash_output = g.subgraph(np.array([0, 2])); subgraph = codeflash_output # 172μs -> 149μs (15.8% faster)

class TestDiGraphSubgraphEdgeCases:
    """Edge case tests for DiGraph.subgraph functionality."""

    def test_subgraph_single_isolated_node(self):
        """Test subgraph with a single isolated node (no edges)."""
        adj_matrix = np.array([[0, 1, 0],
                               [1, 0, 0],
                               [0, 0, 0]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph containing only the isolated node
        codeflash_output = g.subgraph(np.array([2])); subgraph = codeflash_output # 164μs -> 146μs (12.4% faster)

    def test_subgraph_disconnected_nodes(self):
        """Test subgraph with disconnected nodes."""
        adj_matrix = np.array([[1, 0, 0, 0],
                               [0, 1, 0, 0],
                               [0, 0, 1, 0],
                               [0, 0, 0, 1]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with nodes 0 and 2 (disconnected)
        codeflash_output = g.subgraph(np.array([0, 2])); subgraph = codeflash_output # 170μs -> 149μs (14.5% faster)

    def test_subgraph_with_non_sequential_indices(self):
        """Test subgraph with non-sequential node indices."""
        adj_matrix = np.array([[1, 1, 1, 1, 1],
                               [1, 1, 1, 1, 1],
                               [1, 1, 1, 1, 1],
                               [1, 1, 1, 1, 1],
                               [1, 1, 1, 1, 1]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with non-sequential indices: 0, 2, 4
        codeflash_output = g.subgraph(np.array([0, 2, 4])); subgraph = codeflash_output # 171μs -> 148μs (15.3% faster)

    def test_subgraph_with_reversed_indices(self):
        """Test subgraph with reversed node indices."""
        adj_matrix = np.array([[0, 1, 0],
                               [0, 0, 1],
                               [1, 0, 0]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with reversed indices
        codeflash_output = g.subgraph(np.array([2, 1, 0])); subgraph = codeflash_output # 171μs -> 152μs (13.1% faster)

    def test_subgraph_with_duplicated_indices(self):
        """Test subgraph behavior with duplicated node indices."""
        adj_matrix = np.array([[1, 1, 0],
                               [0, 1, 1],
                               [1, 0, 1]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with duplicated indices [0, 1, 0]
        # np.ix_ will handle this and create a subgraph using the specified indices
        codeflash_output = g.subgraph(np.array([0, 1, 0])); subgraph = codeflash_output # 172μs -> 151μs (13.9% faster)

    def test_subgraph_from_sparse_matrix(self):
        """Test subgraph when original graph uses sparse matrix."""
        # Create a sparse adjacency matrix
        adj_matrix = sparse.csr_matrix(np.array([[0, 1, 0, 0],
                                                  [0, 0, 1, 0],
                                                  [0, 0, 0, 1],
                                                  [1, 0, 0, 0]]))
        g = DiGraph(adj_matrix)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([0, 2])); subgraph = codeflash_output # 159μs -> 149μs (7.03% faster)

    def test_subgraph_empty_after_selection(self):
        """Test subgraph when selecting nodes results in no edges."""
        adj_matrix = np.array([[0, 1, 0, 0],
                               [1, 0, 0, 0],
                               [0, 0, 0, 1],
                               [0, 0, 1, 0]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with two isolated nodes from different components
        codeflash_output = g.subgraph(np.array([0, 2])); subgraph = codeflash_output # 160μs -> 149μs (7.06% faster)

    def test_subgraph_boolean_vs_weighted(self):
        """Test that subgraph correctly handles weighted flag."""
        adj_matrix = np.array([[0, 1, 2],
                               [3, 0, 4],
                               [5, 6, 0]])
        g = DiGraph(adj_matrix, weighted=True)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([0, 1])); subgraph = codeflash_output # 173μs -> 150μs (15.3% faster)

    def test_subgraph_with_zero_weights(self):
        """Test subgraph with zero weights in weighted matrix."""
        adj_matrix = np.array([[0.0, 1.5, 0.0],
                               [0.0, 0.0, 2.5],
                               [3.5, 0.0, 0.0]])
        g = DiGraph(adj_matrix, weighted=True)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([0, 2])); subgraph = codeflash_output # 172μs -> 149μs (15.3% faster)

    def test_subgraph_preserves_symmetry(self):
        """Test that subgraph preserves edge symmetry."""
        # Create a symmetric adjacency matrix
        adj_matrix = np.array([[0, 1, 1],
                               [1, 0, 1],
                               [1, 1, 0]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([0, 1])); subgraph = codeflash_output # 171μs -> 147μs (16.3% faster)

    def test_subgraph_large_index_value(self):
        """Test subgraph with large node indices."""
        # Create a 10x10 adjacency matrix
        adj_matrix = np.eye(10, dtype=int)
        g = DiGraph(adj_matrix)
        
        # Extract subgraph using largest index
        codeflash_output = g.subgraph(np.array([9])); subgraph = codeflash_output # 161μs -> 146μs (10.8% faster)

class TestDiGraphSubgraphDataTypes:
    """Test cases for different data types and formats."""

    def test_subgraph_with_float_weights(self):
        """Test subgraph with floating-point weighted adjacency matrix."""
        adj_matrix = np.array([[0.1, 0.5, 0.3],
                               [0.2, 0.4, 0.6],
                               [0.7, 0.1, 0.2]], dtype=float)
        g = DiGraph(adj_matrix, weighted=True)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([0, 2])); subgraph = codeflash_output # 172μs -> 150μs (15.1% faster)

    def test_subgraph_with_int_weights(self):
        """Test subgraph with integer weighted adjacency matrix."""
        adj_matrix = np.array([[1, 2, 3],
                               [4, 5, 6],
                               [7, 8, 9]], dtype=int)
        g = DiGraph(adj_matrix, weighted=True)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([0, 1])); subgraph = codeflash_output # 172μs -> 149μs (15.5% faster)

    def test_subgraph_with_string_labels(self):
        """Test subgraph with string node labels."""
        adj_matrix = np.array([[1, 1, 0],
                               [0, 1, 1],
                               [1, 0, 1]])
        node_labels = np.array(['node_a', 'node_b', 'node_c'])
        g = DiGraph(adj_matrix, node_labels=node_labels)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([1, 2])); subgraph = codeflash_output # 175μs -> 153μs (14.2% faster)

    def test_subgraph_preserves_sparsity(self):
        """Test that subgraph preserves sparsity structure."""
        # Create a sparse adjacency matrix
        adj_matrix = np.array([[0, 0, 0, 1],
                               [0, 0, 1, 0],
                               [0, 1, 0, 0],
                               [1, 0, 0, 0]])
        g = DiGraph(adj_matrix)
        
        # Extract subgraph
        codeflash_output = g.subgraph(np.array([0, 1, 2, 3])); subgraph = codeflash_output # 172μs -> 148μs (15.6% faster)

class TestDiGraphSubgraphLargeScale:
    """Large scale test cases for DiGraph.subgraph functionality."""

    def test_subgraph_large_complete_graph(self):
        """Test subgraph extraction from a large complete graph."""
        # Create a complete graph with 100 nodes
        n = 100
        adj_matrix = np.ones((n, n), dtype=bool)
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with 50 nodes
        nodes = np.arange(0, 50)
        codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 262μs -> 172μs (52.4% faster)

    def test_subgraph_large_sparse_graph(self):
        """Test subgraph extraction from a large sparse graph."""
        # Create a large sparse adjacency matrix (chain structure)
        n = 500
        adj_matrix = np.zeros((n, n), dtype=int)
        for i in range(n - 1):
            adj_matrix[i, i + 1] = 1
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with 100 nodes in the middle
        nodes = np.arange(200, 300)
        codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 268μs -> 161μs (66.3% faster)

    def test_subgraph_large_random_graph(self):
        """Test subgraph extraction from a large random graph."""
        # Create a large random adjacency matrix
        n = 300
        np.random.seed(42)
        adj_matrix = np.random.randint(0, 2, size=(n, n))
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with 100 random nodes
        nodes = np.sort(np.random.choice(n, 100, replace=False))
        codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 863μs -> 266μs (224% faster)

    def test_subgraph_large_with_labels(self):
        """Test subgraph extraction preserves large number of labels."""
        # Create a large graph with labels
        n = 200
        adj_matrix = np.eye(n, dtype=int)
        node_labels = np.array([f'node_{i}' for i in range(n)])
        g = DiGraph(adj_matrix, node_labels=node_labels)
        
        # Extract subgraph with 50 nodes
        nodes = np.arange(0, 50)
        codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 208μs -> 161μs (29.1% faster)

    def test_subgraph_sequential_extractions(self):
        """Test multiple sequential subgraph extractions."""
        # Create a large graph
        n = 100
        adj_matrix = np.random.randint(0, 2, size=(n, n))
        g = DiGraph(adj_matrix)
        
        # Extract subgraph multiple times
        for i in range(10):
            nodes = np.arange(i * 5, (i + 1) * 5)
            codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 1.46ms -> 1.22ms (19.7% faster)

    def test_subgraph_large_weighted_matrix(self):
        """Test subgraph extraction from a large weighted matrix."""
        # Create a large weighted adjacency matrix
        n = 150
        np.random.seed(42)
        adj_matrix = np.random.rand(n, n)
        g = DiGraph(adj_matrix, weighted=True)
        
        # Extract subgraph with 75 nodes
        nodes = np.arange(0, 75)
        codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 392μs -> 211μs (86.0% faster)
        # All weights should be preserved
        for i in range(75):
            for j in range(75):
                pass

    def test_subgraph_many_isolated_nodes(self):
        """Test subgraph with many isolated nodes."""
        # Create a graph with mostly isolated nodes
        n = 250
        adj_matrix = np.zeros((n, n), dtype=int)
        # Add edges only for first 10 nodes
        for i in range(10):
            for j in range(10):
                adj_matrix[i, j] = 1
        g = DiGraph(adj_matrix)
        
        # Extract subgraph with isolated nodes
        nodes = np.arange(100, 200)
        codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 231μs -> 155μs (49.1% faster)

    def test_subgraph_performance_large_extraction(self):
        """Test performance of extracting a large subgraph."""
        # Create a large complete subgraph
        n = 500
        adj_matrix = np.ones((n, n), dtype=bool)
        g = DiGraph(adj_matrix)
        
        # Extract a large subgraph (250 nodes)
        nodes = np.arange(0, 250)
        codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 2.60ms -> 568μs (358% faster)

    def test_subgraph_many_extractions_memory_efficient(self):
        """Test that multiple subgraph extractions are memory efficient."""
        # Create a moderately large graph
        n = 200
        adj_matrix = np.random.randint(0, 2, size=(n, n))
        g = DiGraph(adj_matrix)
        
        # Extract many subgraphs
        subgraphs = []
        for i in range(20):
            start = i * 5
            end = start + 10
            if end > n:
                break
            nodes = np.arange(start, min(end, n))
            codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 3.00ms -> 2.44ms (23.1% faster)
            subgraphs.append(subgraph)
        for subgraph in subgraphs:
            pass

    def test_subgraph_sparse_connectivity(self):
        """Test subgraph with very sparse connectivity."""
        # Create a graph with extremely sparse connections
        n = 300
        adj_matrix = np.zeros((n, n), dtype=int)
        # Add only n edges (very sparse)
        for i in range(n - 1):
            adj_matrix[i, i + 1] = 1
        g = DiGraph(adj_matrix)
        
        # Extract a large subgraph
        nodes = np.arange(0, 150)
        codeflash_output = g.subgraph(nodes); subgraph = codeflash_output # 368μs -> 162μs (127% 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-DiGraph.subgraph-mkp3k41x and push.