Vectorize expressibility sampling and batch fidelity computation

Author: ashutoshm1771

src/encoding_atlas/analysis/_utils.py

@@ -512,25 +512,26 @@ def simulate_encoding_statevectors_batch(
     encoding: BaseEncoding,
     X: NDArray[np.floating[Any]],
     backend: Literal["pennylane", "qiskit", "cirq"] = "pennylane",
-) -> list[StatevectorType]:
+) -> NDArray[np.complexfloating[Any, Any]]:
     """Simulate encoding circuits for multiple input vectors.
 
     This is a convenience function that applies :func:`simulate_encoding_statevector`
-    to each row of a 2D input array.
+    to each row of a 2D input array and returns a pre-allocated 2D array of
+    statevectors.
 
     Parameters
     ----------
     encoding : BaseEncoding
         The encoding instance to simulate.
     X : NDArray[np.floating]
         Input data array of shape ``(n_samples, n_features)``.
-    backend : {"pennylane", "qiskit"}, default="pennylane"
+    backend : {"pennylane", "qiskit", "cirq"}, default="pennylane"
         The quantum simulation backend to use.
 
     Returns
     -------
-    list[StatevectorType]
-        List of statevectors, one for each input sample.
+    NDArray[np.complexfloating], shape ``(n_samples, 2**n_qubits)``
+        2D array of statevectors, one row per input sample.
 
     Raises
     ------
@@ -554,25 +555,28 @@ def simulate_encoding_statevectors_batch(
     if X_array.ndim != 2:
         raise ValidationError(f"Input X must be 2D array, got shape {X_array.shape}")
 
+    n_samples = X_array.shape[0]
+    dim = 2**encoding.n_qubits
+
     _logger.debug(
         "Batch simulating %d samples for encoding %s",
-        X_array.shape[0],
+        n_samples,
         encoding.__class__.__name__,
     )
 
-    statevectors = []
+    states = np.zeros((n_samples, dim), dtype=np.complex128)
     for i, x in enumerate(X_array):
         try:
             state = simulate_encoding_statevector(encoding, x, backend)
-            statevectors.append(state)
+            states[i] = np.asarray(state, dtype=np.complex128).ravel()
         except SimulationError as e:
             raise SimulationError(
                 f"Simulation failed for sample {i}: {e}",
                 backend=backend,
                 details={"sample_index": i, "original_error": str(e)},
             ) from e
 
-    return statevectors
+    return states
 
 
 def _simulate_pennylane(
@@ -1469,6 +1473,53 @@ def compute_fidelity(
     return fidelity
 
 
+def _compute_fidelities_batch(
+    states1: NDArray[np.complexfloating[Any, Any]],
+    states2: NDArray[np.complexfloating[Any, Any]],
+) -> NDArray[np.floating[Any]]:
+    """Compute fidelities between pairs of statevectors in batch.
+
+    Vectorized version of :func:`compute_fidelity` for arrays of states.
+    Computes F_i = |⟨ψ₁ⁱ|ψ₂ⁱ⟩|² for each pair (i).
+
+    Parameters
+    ----------
+    states1 : NDArray[np.complexfloating], shape ``(n, d)``
+        First set of statevectors.
+    states2 : NDArray[np.complexfloating], shape ``(n, d)``
+        Second set of statevectors.
+
+    Returns
+    -------
+    NDArray[np.floating], shape ``(n,)``
+        Fidelity values, each in [0, 1].
+
+    Raises
+    ------
+    ValueError
+        If shapes of ``states1`` and ``states2`` do not match.
+    ValidationError
+        If any state contains NaN or infinite values.
+    """
+    if states1.shape != states2.shape:
+        raise ValueError(
+            f"States must have same shape: got {states1.shape} and {states2.shape}"
+        )
+    if states1.ndim != 2:
+        raise ValueError(f"States must be 2D arrays, got ndim={states1.ndim}")
+
+    if np.any(np.isnan(states1)) or np.any(np.isinf(states1)):
+        raise ValidationError("states1 contains NaN or infinite values")
+    if np.any(np.isnan(states2)) or np.any(np.isinf(states2)):
+        raise ValidationError("states2 contains NaN or infinite values")
+
+    overlaps = np.sum(np.conj(states1) * states2, axis=1)
+    fidelities = np.abs(overlaps) ** 2
+    fidelities = np.clip(fidelities, 0.0, 1.0).astype(np.float64)
+
+    return fidelities
+
+
 def compute_purity(
     density_matrix: DensityMatrixType,
 ) -> float:

src/encoding_atlas/analysis/expressibility.py

@@ -1040,11 +1040,9 @@ def _sample_fidelities(
 ) -> NDArray[np.floating[Any]]:
     """Sample fidelities between random input pairs.
 
-    This is the core sampling loop for expressibility computation.
-    For each of n_samples iterations:
-    1. Generate two random input vectors x1, x2
-    2. Simulate encoding to get states |ψ(x1)⟩, |ψ(x2)⟩
-    3. Compute fidelity F = |⟨ψ(x1)|ψ(x2)⟩|²
+    Generates random input pairs in batch, simulates encoding circuits,
+    and computes fidelities. RNG calls are vectorized into single batch
+    operations for reduced overhead.
 
     Parameters
     ----------
@@ -1073,24 +1071,20 @@ def _sample_fidelities(
     SimulationError
         If circuit simulation fails.
     """
-    # TODO: Vectorize this loop — batch rng.uniform calls and simulation
-    # calls to reduce per-sample overhead. This is the main performance
-    # bottleneck for large n_samples. See simulate_encoding_statevectors_batch
-    # in _utils.py for a possible starting point.
+    # Batch RNG generation — single call per input set
+    X1 = rng.uniform(input_range[0], input_range[1], size=(n_samples, n_features))
+    X2 = rng.uniform(input_range[0], input_range[1], size=(n_samples, n_features))
+
     fidelities = np.zeros(n_samples, dtype=np.float64)
 
     # Logging interval (log every 10% of progress)
     log_interval = max(1, n_samples // 10)
 
     for i in range(n_samples):
-        # Generate two random input vectors
-        x1 = rng.uniform(input_range[0], input_range[1], size=n_features)
-        x2 = rng.uniform(input_range[0], input_range[1], size=n_features)
-
         try:
             # Simulate encoding to get statevectors
-            state1 = simulate_encoding_statevector(encoding, x1, backend=backend)
-            state2 = simulate_encoding_statevector(encoding, x2, backend=backend)
+            state1 = simulate_encoding_statevector(encoding, X1[i], backend=backend)
+            state2 = simulate_encoding_statevector(encoding, X2[i], backend=backend)
 
             # Compute fidelity: F = |⟨ψ₁|ψ₂⟩|²
             fidelity = compute_fidelity(state1, state2)