Optimise RotationAnalyser permutation loop

CHANGELOG.md

@@ -6,6 +6,15 @@
 
 - Cached `symmetry_measure` property on `CoordinationPolyhedron` to avoid redundant CSM calculations on repeated access.
 - `edge_graph` and `symmetry_measure` now return defensive copies to prevent callers from corrupting the internal cache.
+- Replaced N! permutation loop in `RotationAnalyser` with symmetry-reduced approach using bsym, giving ~50x speedup for octahedral geometries.
+
+### Bug fixes
+
+- Fixed `OrientationDict.all_rotational_distances` type annotation (`np.ndarray`, not `float`).
+
+### Other changes
+
+- Removed redundant `assert isinstance` checks across the codebase.
 
 ## 0.4.0
 

README.md

@@ -4,9 +4,7 @@
 [![Coverage Status](https://coveralls.io/repos/github/bjmorgan/polyhedral-analysis/badge.svg?branch=main)](https://coveralls.io/github/bjmorgan/polyhedral-analysis?branch=main)
 [![Documentation Status](https://readthedocs.org/projects/polyhedral-analysis/badge/?version=latest)](http://polyhedral-analysis.readthedocs.io/en/latest/?badge=latest)
 
-`polyhedral-analysis` is a Python module for analysing coordination polyhedra in crystal structures and molecular dynamics trajectories. It identifies coordination environments around central atoms and computes geometric properties such as distortion measures, bond lengths, angles, volumes, and connectivity.
-
-Built on top of [pymatgen](https://pymatgen.org), the package works with any structure that pymatgen can read.
+`polyhedral-analysis` is a Python module for analysing coordination polyhedra in crystal structures and molecular dynamics trajectories. Built on top of [pymatgen](https://pymatgen.org), it works with any structure that pymatgen can read.
 
 ## Features
 
@@ -16,7 +14,7 @@ Built on top of [pymatgen](https://pymatgen.org), the package works with any str
 - Bond lengths, angles, volumes, and edge connectivity graphs
 - Corner-, edge-, and face-sharing neighbour analysis
 - Off-centre displacement and radial distortion parameters
-- Vertex vector orientation analysis with stereographic projection plotting
+- Vertex vector orientation analysis with orientation distribution plotting
 - Trajectory analysis for tracking polyhedral distortions over molecular dynamics runs
 
 ## Installation

polyhedral_analysis/rotation_analyser.py

@@ -1,27 +1,105 @@
+from __future__ import annotations
+
 from polyhedral_analysis.csm import continuous_symmetry_measure
-from itertools import permutations
-import math
 import numpy as np
 from typing import TypedDict
 from polyhedral_analysis.coordination_polyhedron import CoordinationPolyhedron
 
 
+def _analyse_point_group(reference_points: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Analyse the point group of a reference geometry.
+
+    Performs a single symmetry analysis and returns both the
+    symmetry-reduced permutation indices and the proper rotation matrices.
+
+    Args:
+        reference_points: An Nx3 array of reference vertex coordinates.
+
+    Returns:
+        A tuple of (reduced_permutations, proper_rotations) where:
+
+        - reduced_permutations is an MxN integer array of symmetry-inequivalent
+          permutation representatives.
+        - proper_rotations is a Kx3x3 array of proper (det > 0) rotation
+          matrices of the point group.
+    """
+    from pymatgen.core import Molecule
+    from pymatgen.symmetry.analyzer import PointGroupAnalyzer
+    from pymatgen.util.coord import coord_list_mapping
+    from bsym import ConfigurationSpace  # type: ignore[import]
+    from bsym import PointGroup as BsymPointGroup  # type: ignore[import]
+    from bsym.symmetry_operation import SymmetryOperation  # type: ignore[import]
+
+    n = len(reference_points)
+    mol = Molecule(['H'] * n, reference_points - reference_points.mean(axis=0))
+    pga = PointGroupAnalyzer(mol)
+    symm_ops = pga.get_symmetry_operations()
+
+    # Extract proper 3x3 rotation matrices
+    proper_rotations = np.array([
+        op.rotation_matrix for op in symm_ops
+        if np.linalg.det(op.rotation_matrix) > 0.5
+    ])
+    if len(proper_rotations) == 0:
+        raise RuntimeError(
+            'No proper rotations found for the reference geometry.')
+
+    # Convert symmetry operations to permutation representations for bsym
+    coords = mol.cart_coords
+    seen: set[tuple[int, ...]] = set()
+    mappings: list[list[int]] = []
+    for op in symm_ops:
+        mapped_coords = op.operate_multi(coords)
+        new_mol = Molecule(mol.species, mapped_coords)
+        mapping_zero = coord_list_mapping(new_mol.cart_coords, coords)
+        if len(set(mapping_zero)) != n:
+            raise RuntimeError(
+                'coord_list_mapping did not produce a bijection; '
+                'symmetry operation maps multiple vertices to the same target.')
+        mapping = [x + 1 for x in mapping_zero]
+        key = tuple(mapping)
+        if key not in seen:
+            seen.add(key)
+            mappings.append(mapping)
+
+    sym_ops = [SymmetryOperation.from_vector(m) for m in mappings]
+    point_group = BsymPointGroup(symmetry_operations=sym_ops)
+    config_space = ConfigurationSpace(
+        objects=list(range(n)),
+        symmetry_group=point_group,
+    )
+    unique_configs = config_space.unique_configurations(
+        site_distribution={i: 1 for i in range(n)},
+    )
+    reduced_permutations = np.array(
+        [c.tolist() for c in unique_configs],
+        dtype=np.intp,
+    )
+
+    return reduced_permutations, proper_rotations
+
+
 class OrientationDict(TypedDict):
     orientation_index: int
     reference_geometry_index: int
     rotational_distance: float
     symmetry_measure: float
-    all_rotational_distances: float
+    all_rotational_distances: np.ndarray
 
-class RotationAnalyser(object):
+
+class RotationAnalyser:
     """Class for analysing rotational orientation of polyhedra.
 
+    All reference orientations must be rotated copies of the same geometry
+    (i.e. share the same point group). The symmetry analysis is performed
+    once from the first reference orientation and reused for all others.
+
     Attributes:
-        reference_points (np.array): A Mx3xN numpy array of points around the
-            origin ([0.0, 0.0, 0.0]) that define sets of centre-vertex
+        reference_points (np.ndarray): An MxNx3 numpy array of points around
+            the origin ([0.0, 0.0, 0.0]) that define sets of centre-vertex
             vectors for specific reference orientations, e.g. for a single
-            orientation of a tetrahedron, reference points will be a size 1x3x4
-            array, e.g.::
+            orientation of a tetrahedron, reference points will be a size
+            1x4x3 array, e.g.::
 
                 [[[ 1.0, -1.0,  1.0],
                   [-1.0, -1.0, -1.0],
@@ -35,88 +113,158 @@ def __init__(self,
         """Initialise a RotationAnalyser instance.
 
         Args:
-            reference_points (np.array): Either a 3xN or Mx3xN numpy array
-            of points around the origin ([0.0, 0.0, 0.0]) that define sets of 
-            centre-vertex vectors, and are used to classify the rotational
-            orientation of polyhedra.
-
-            If a 3xN array is passed in, this will be converted to a 1x3xN array.
-
-        Returns:
-            None
+            reference_points: Either an Nx3 or MxNx3 numpy array of points
+                around the origin ([0.0, 0.0, 0.0]) that define sets of
+                centre-vertex vectors, and are used to classify the rotational
+                orientation of polyhedra.
 
+                If an Nx3 array is passed in, this will be converted to a
+                1xNx3 array.
         """
         if len(reference_points.shape) == 2:
             reference_points = np.array([reference_points])
         if len(reference_points.shape) == 1:
             raise ValueError(
                 'Reference points must all contain the same number of coordinates')
         self.reference_points = reference_points
+        self._symmetry_data: tuple[np.ndarray, np.ndarray] | None = None
+
+    def _ensure_symmetry_analysis(self) -> tuple[np.ndarray, np.ndarray]:
+        """Run point group analysis if not already done.
+
+        Returns:
+            A tuple of (reduced_permutations, proper_rotations) from
+            a single symmetry analysis.
+        """
+        if self._symmetry_data is None:
+            self._symmetry_data = _analyse_point_group(self.reference_points[0])
+        return self._symmetry_data
+
+    @property
+    def reduced_permutations(self) -> np.ndarray:
+        """Symmetry-inequivalent permutation indices for the reference geometry.
+
+        Computed lazily on first access.
+        """
+        return self._ensure_symmetry_analysis()[0]
+
+    @property
+    def proper_rotations(self) -> np.ndarray:
+        """Proper rotation matrices (Kx3x3) of the reference geometry's point group.
+
+        Computed lazily on first access.
+        """
+        return self._ensure_symmetry_analysis()[1]
 
     def discrete_orientation(self,
                              points: np.ndarray) -> OrientationDict:
         """Find the discrete "closest orientation" for an input polyhedron of points.
 
-        For example, a tetrahedon has 12 pure rotation symmetry operations. A distorted
-        tetrahedron with vertices approximately aligned along the unordered vectors::
+        For example, a tetrahedron has 12 proper rotation symmetry operations.
+        A distorted tetrahedron with vertices approximately aligned along the
+        unordered vectors::
 
             [[ 1.0, -1.0,  1.0],
              [-1.0, -1.0, -1.0],
              [ 1.0,  1.0, -1.0],
              [-1.0,  1.0,  1.0]]
 
-        can be in one of 12 orientations. This method compares the input points to
-        all permutations of self.reference_points that can be generated by proper rotations, and returns
-        an index for orientation that minimises the rotational distance between the reordered
-        reference points and the input points.
+        can be in one of 12 orientations. This method compares the input
+        points to the reference orientations and returns the orientation that
+        minimises the rotational distance.
+
+        The algorithm uses a two-phase approach to avoid the N! cost of
+        evaluating all vertex permutations:
 
-        The algorithm is:
-        1. Perform continuous symmetry measure analysis for all permutations of reference points.
-        2. Collect operations that have the minimum continous symmetry measure value.
-        These give the set of all propert and improper rotations.
-        3. Find all proper rotations using det(M_rot)>0.
-        4. Calculate the rotational distance.
+        1. Find the best vertex alignment using symmetry-reduced permutations
+           of the reference points (N!/|G| CSM evaluations instead of N!).
+        2. Compose the best-fit Procrustes rotation R_best with all proper
+           rotation matrices S_i of the reference geometry's point group to
+           give the rotation matrix for each orientation: R_i = S_i @ R_best.
+        3. Calculate rotational distances from the composed matrices.
+        4. Return the orientation with the minimum rotational distance.
+
+        The composed rotation matrices in step 2 are exact for undistorted
+        geometries. For distorted geometries they are approximate, since the
+        true Procrustes rotation for each symmetry-equivalent permutation
+        differs slightly from the composed matrix. In practice the results
+        are numerically indistinguishable from the brute-force approach for
+        moderate distortions.
+
+        If exact rotation matrices are needed for highly distorted
+        geometries, an alternative is a two-phase orbit-expansion approach:
+        use reduced permutations to identify the minimum-CSM orbit (phase 1),
+        then expand that orbit to all |G| equivalent permutations and compute
+        CSM for each to obtain exact rotation matrices (phase 2). This costs
+        N!/|G| + |G| CSM evaluations per reference orientation rather than
+        N!/|G|, but avoids the composition approximation.
 
         Args:
-            points (np.array(float)): Nx3 numpy array of points describing the coordinates of the
-                input polyhedron, centered around zero.
+            points: Nx3 numpy array of points describing the coordinates of
+                the input polyhedron, centred around zero.
 
         Returns:
-            (dict): Dictionary describing the orientation, with keys:
-                - ``orientation_index`` (int): Index of this particular orientation.
-                - ``reference_geometry_index`` (int): Index of the reference geometry
-                    the closest discrete orientation is equivalent to.
-                - ``rotational_distance`` (float): Angle of rotation from the relevant
-                    reference orientation, in radians.
-                - ``all_rotational_distances`` (np.array(float)): Array of angles of rotation
-                    from each reference orientation.
-                - ``symmetry_measure`` (float): The continuous symmetry measure (CSM)
-                  for this polyhedron (J. Am. Chem. Soc., 114, 7843-7851 (1992)) 
+            Dictionary describing the orientation, with keys:
 
+            - ``orientation_index`` (int): Index of this particular orientation.
+            - ``reference_geometry_index`` (int): Index of the reference geometry
+              the closest discrete orientation is equivalent to.
+            - ``rotational_distance`` (float): Angle of rotation from the
+              relevant reference orientation, in radians.
+            - ``all_rotational_distances`` (np.ndarray): Array of angles of
+              rotation from each reference orientation.
+            - ``symmetry_measure`` (float): The continuous symmetry measure
+              (CSM) for this polyhedron.
         """
-        points -= np.mean(points, axis=0, dtype=float)
-        sm = []
-        for rp in self.reference_points:
-            sm.extend([continuous_symmetry_measure(points, np.array(p)) for p in permutations(rp)])
-        min_sm = min(s.symmetry_measure for s in sm)
-        pure_rot_sm = [s for s in sm if math.isclose(
-            s.symmetry_measure, min_sm)]
-        proper_rot_sm = [s for s in pure_rot_sm if np.linalg.det(
-            s.rotation_matrix) > 0]
-        proper_rot_matrices = np.array(
-            [s.rotation_matrix for s in proper_rot_sm])
-        trace = np.trace(proper_rot_matrices, axis1=1, axis2=2)
-        rot_distance = np.arccos(np.clip((trace - 1.0) / 2.0, -1.0, 1.0))
-        index = int(np.argmin(rot_distance))
-        reference_geometry_index = index // int(
-            len(proper_rot_sm) / len(self.reference_points))
-        return OrientationDict(orientation_index=index,
-                               reference_geometry_index=reference_geometry_index,
-                               rotational_distance=rot_distance[index],
-                               symmetry_measure=proper_rot_sm[index].symmetry_measure,
-                               all_rotational_distances=rot_distance)
+        expected_n = self.reference_points.shape[1]
+        if points.shape != (expected_n, 3):
+            raise ValueError(
+                f'Expected points with shape ({expected_n}, 3), '
+                f'got {points.shape}.')
+        points = points - np.mean(points, axis=0, dtype=float)
+        n_refs = len(self.reference_points)
+        n_proper = len(self.proper_rotations)
+        # Phase 1: find best alignment per reference orientation
+        best_csm_per_ref = np.full(n_refs, float('inf'))
+        best_rotation_per_ref: list[np.ndarray] = [np.eye(3) for _ in range(n_refs)]
+        for i, rp in enumerate(self.reference_points):
+            for perm in self.reduced_permutations:
+                result = continuous_symmetry_measure(points, rp[perm])
+                if result.symmetry_measure < best_csm_per_ref[i]:
+                    best_csm_per_ref[i] = result.symmetry_measure
+                    best_rotation_per_ref[i] = result.rotation_matrix
+        # Phase 2: compose best-fit rotation with proper symmetry operations
+        all_rot_distances = np.empty(n_refs * n_proper)
+        for i in range(n_refs):
+            composed = self.proper_rotations @ best_rotation_per_ref[i]
+            traces = np.trace(composed, axis1=1, axis2=2)
+            start = i * n_proper
+            all_rot_distances[start:start + n_proper] = np.arccos(
+                np.clip((traces - 1.0) / 2.0, -1.0, 1.0))
+        index = int(np.argmin(all_rot_distances))
+        reference_geometry_index = index // n_proper
+        return OrientationDict(
+            orientation_index=index,
+            reference_geometry_index=reference_geometry_index,
+            rotational_distance=all_rot_distances[index],
+            symmetry_measure=float(best_csm_per_ref[reference_geometry_index]),
+            all_rotational_distances=all_rot_distances,
+        )
 
     def polyhedron_orientation(self,
                                polyhedron: CoordinationPolyhedron) -> OrientationDict:
+        """Find the discrete orientation of a coordination polyhedron.
+
+        Convenience wrapper around :meth:`discrete_orientation` that
+        extracts vertex vectors from the polyhedron relative to its
+        central atom.
+
+        Args:
+            polyhedron: The coordination polyhedron to analyse.
+
+        Returns:
+            Dictionary describing the orientation.
+            See :meth:`discrete_orientation` for details.
+        """
         points = polyhedron.vertex_vectors(reference='central_atom')
         return self.discrete_orientation(points)