Perf: 30% to 70% speedup for superimpose

benchmarks/structure/benchmark_superimpose.py

@@ -25,3 +25,22 @@ def benchmark_superimpose(method, atoms):
     Compute superimposition of two structures with the same number of atoms.
     """
     method(atoms[0], atoms[1])
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize(
+    "method",
+    [
+        struc.superimpose,
+        struc.superimpose_without_outliers,
+        struc.superimpose_homologs,
+    ],
+)
+def benchmark_superimpose_multi_model(method, atoms):
+    """
+    Superimpose every model of an ensemble onto a single reference model.
+
+    ``superimpose_structural_homologs`` is omitted, as it does not accept an
+    :class:`AtomArrayStack`.
+    """
+    method(atoms[0], atoms)

src/biotite/structure/superimpose.py

@@ -489,7 +489,7 @@ def _get_rotation_matrices(
     Both sets of coordinates must already be centered at origin.
     """
     # Calculate cross-covariance matrices
-    cov = np.sum(fixed[:, :, :, np.newaxis] * mobile[:, :, np.newaxis, :], axis=1)
+    cov = np.matmul(np.swapaxes(fixed, -1, -2), mobile)  # (M, XYZ, XYZ)
     v, s, w = np.linalg.svd(cov)
     # Remove possibility of reflected atom coordinates
     reflected_mask = np.linalg.det(v) * np.linalg.det(w) < 0