How can I optimize the nested-loop construction of a large NumPy matrix S1?

Below is one strategy to improve performance: try to replace the Python‐level nested loop with a vectorized operation. If you can cast your basis to a 3‑dimensional NumPy array (say, with shape (num, d, d)) then you can use tensor contractions (for example, via np.einsum) to compute S1 all at once.

For example, note that

  S1[i, j] = trace(b₁[i] · (b₁[j] · state))

can be rewritten as a double contraction

  S1[i, j] = ∑ₚ,₍q₎ (b₁[i])ₚ,₍q₎ · ( (b₁[j] · state) )₍q,p₎

If you let A be the tensor of all b₁ matrices with shape (num, d, d), then you can precompute

  P = np.matmul(A, state)

which gives an array of shape (num, d, d) so that P[j] = b₁[j]·state.

Then the whole S1 may be constructed as

PYTHONCopied
# A: a (num, d, d) array with A[i] = b1[i]
P = np.matmul(A, state)  # shape (num, d, d)
S1 = np.einsum('ipq,jqp->ij', A, P)

This np.einsum call contracts over the indices p and q so that for each i and j you obtain the trace of b1[i] · (b1[j] · state).

Additional Considerations

1. Memory usage:
   Packing all b₁ matrices into a single tensor A will use memory roughly equal to
     num × d × d × (sizeof(complex128)).
   For example, if d = 81 then num = 6561 and you need to store ≈ 6.5k × 81 × 81 ≈ 43 million numbers, which at 16 bytes each comes out to around 700 MB. If this is too heavy, you might have to work block‐wise.

2. Exploiting Sparsity:
   If b₁ and state are sparse, consider storing them in a sparse format (for example using SciPy’s sparse matrices). If many entries are zero then sparse matrix routines (or even custom loops that only iterate over nonzero entries) could reduce your computation time. However, note that converting many small d×d sparse matrices into a block–sparse structure may require some work.

3. Profiling before switching strategy:
   Often the best optimization is moving the costly innermost operations into vectorized NumPy calls. In our case, using np.einsum (or sometimes even np.tensordot) to get rid of the Python-level nested loops should provide a very significant speedup.

Summary

Using vectorized NumPy operations you can replace the double loop:

PYTHONCopied
# Original code:
S1 = np.zeros((num, num), dtype=np.complex128)
for i in range(num):
    for j in range(num):
        S1[i, j] = np.trace(np.dot(b1[i], np.dot(b1[j], state)))

with the vectorized version:

PYTHONCopied
# Convert list of d x d matrices to a single tensor (num, d, d)
A = np.array(b1)  # ensure b1 is a NumPy array with shape (num, d, d)
P = np.matmul(A, state)        # shape (num, d, d)
S1 = np.einsum('ipq,jqp->ij', A, P)

This approach exploits lower-level vectorized BLAS (and possibly the sparsity structure if you use sparse formats) to dramatically speed up the calculation.

By vectorizing the computation you reduce the overhead of nested Python loops and take full advantage of NumPy’s optimized routines.