How can I vectorize the evaluation of an arbitrary NumPy expression over specific indices?

Below is one common NumPy solution that avoids an explicit Python loop when your indices are known as arrays. For example, suppose that your indices are stored in an (N, 3) array (or as three one‑dimensional index arrays) so that each row gives a triple (i, j, k). In that case you can write:

PYTHONCopied
import numpy as np

# Example arrays; replace with your data
a = np.array([10, 20, 30, 40])
b = np.array([1, 2, 3, 4])
c = np.array([[100, 200, 300, 400],
              [110, 210, 310, 410],
              [120, 220, 320, 420],
              [130, 230, 330, 430]])

# Suppose indices is an (N, 3) array where each row is (i, j, k)
indices = np.array([
    [0, 1, 2],
    [2, 0, 3],
    [1, 3, 0]
])

# Unpack the indices along the proper axis
i_inds, j_inds, k_inds = indices.T  # Now i_inds, j_inds, k_inds are one-dimensional arrays

# Compute the expression (a[i] + b[j]) * c[k, i] for each triple
results = (a[i_inds] + b[j_inds]) * c[k_inds, i_inds]

# If you want the final scalar sum, simply:
total = results.sum()

print("Results per index:", results)
print("Total:", total)

Explanation

• Advanced Indexing:
  By using advanced indexing (e.g. a[i_inds]), NumPy creates new arrays that contain the elements corresponding to the selected indices. With this approach the whole expression is computed elementwise over all provided indices.

• Arbitrary Expressions:
  The technique naturally extends to arbitrary expressions that combine array elements using operations such as + and *. You simply index the arrays with your index arrays for each operand. In the example above, the expression was
  (a[i_inds] + b[j_inds]) * c[k_inds, i_inds].

• Performance:
  This method leverages vectorized operations rather than a Python-level loop which often leads to significant performance improvements, especially when the number of indices is large.

This solution provides a general pattern: provided your indices are available in a vectorized form, you can compute almost any elementwise expression without resorting to Python loops.