diff --git a/docs/source/theory.rst b/docs/source/theory.rst index 0fb932d1..a63e94d5 100644 --- a/docs/source/theory.rst +++ b/docs/source/theory.rst @@ -291,6 +291,44 @@ influences the convergence of the optimization procedure. Higher widths usually faster and easier. Thus one has to be careful when choosing the width. The width can be chosen on a per particle basis, but this is not recommended. +Returned orientation canonicalization (fundamental zone) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +The maximizing rotation is not unique. If :math:`q^\*` is an optimal quaternion for a +given symmetry target, then any left-multiplied equivalent +:math:`s \otimes q^\*` is also optimal for each proper symmetry rotation +:math:`s` in the relevant equivalence set. To make returned orientations +deterministic, ``spatula`` maps each result to a canonical representative. + +For a raw returned quaternion :math:`q` (stored as :math:`[w, x, y, z]`) and +equivalence set :math:`S`, we construct candidates: + +.. math:: + C = \left\{ \operatorname{norm}(s \otimes q)\;|\;s \in S \right\}, + +where :math:`\operatorname{norm}` normalizes to unit length and enforces the +double-cover sign convention by flipping sign when :math:`w < 0`. + +The reported quaternion is then selected as: + +.. math:: + q_{\mathrm{FZ}} = \underset{c \in C}{\arg\max}\; (c_w, c_x, c_y, c_z), + +using lexicographic ordering (with small floating-point tolerance in code). This +first maximizes :math:`w` (closest equivalent orientation to identity), then +deterministically breaks ties with :math:`x`, :math:`y`, and :math:`z`. + +The equivalence set :math:`S` is chosen as follows: + +* PGOP/BOOSOP group-level rotation outputs: all proper rotations + (determinant :math:`> 0`) from the tested point group. +* PGOP per-operator rotation outputs + (``compute_per_operator_values_for_final_orientation=True``): + proper-rotation subgroup generated by powers of the specific operator + (including identity). + +This canonicalization changes only the orientation representative and does not +change order parameter values. + Calculation of overlap ~~~~~~~~~~~~~~~~~~~~~~ To compute the overlap between two gaussians we use the Bhattacharyya diff --git a/spatula/_fundamental_zone.py b/spatula/_fundamental_zone.py new file mode 100644 index 00000000..bf89bc7d --- /dev/null +++ b/spatula/_fundamental_zone.py @@ -0,0 +1,156 @@ +# Copyright (c) 2021-2026 The Regents of the University of Michigan +# Part of spatula, released under the BSD 3-Clause License. + +"""Internal helpers for mapping quaternions to symmetry fundamental zones. + +Mapping rule used in this module +-------------------------------- +For each returned orientation quaternion ``q`` and associated proper-rotation +equivalence set ``S`` (a subset of ``SO(3)`` represented as quaternions), +we construct + +``C = { normalize(s ⊗ q) | s in S }``, + +where ``⊗`` is quaternion multiplication in ``[w, x, y, z]`` convention and +``normalize`` also enforces the double-cover convention ``w >= 0`` by flipping +sign when needed. The canonical representative is selected as: + +1. maximum ``w`` (closest equivalent rotation to the identity), +2. lexicographic tie-break on ``(x, y, z)`` with a small epsilon. + +This makes orientation outputs deterministic across all symmetry-equivalent +solutions. +""" + +from __future__ import annotations + +import numpy as np +from scipy.spatial.transform import Rotation + +_DET_EPS = 1e-6 +_EQ_EPS = 1e-7 +_MAX_OPERATOR_ORDER = 256 + + +def _normalize_wxyz(quaternions: np.ndarray) -> np.ndarray: + """Normalize quaternions and choose the ``w >= 0`` representative.""" + quaternions = np.asarray(quaternions, dtype=np.float32) + norms = np.linalg.norm(quaternions, axis=-1, keepdims=True) + norms = np.where(norms == 0, 1.0, norms) + normalized = quaternions / norms + normalized[normalized[..., 0] < 0] *= -1.0 + return normalized + + +def _quat_multiply_wxyz(lhs: np.ndarray, rhs: np.ndarray) -> np.ndarray: + """Multiply quaternions in ``[w, x, y, z]`` order.""" + lw, lx, ly, lz = np.moveaxis(lhs, -1, 0) + rw, rx, ry, rz = np.moveaxis(rhs, -1, 0) + return np.stack( + ( + lw * rw - lx * rx - ly * ry - lz * rz, + lw * rx + lx * rw + ly * rz - lz * ry, + lw * ry - lx * rz + ly * rw + lz * rx, + lw * rz + lx * ry - ly * rx + lz * rw, + ), + axis=-1, + ) + + +def _lexicographically_greater(lhs: np.ndarray, rhs: np.ndarray) -> np.ndarray: + """Return mask where ``lhs`` is lexicographically greater than ``rhs``.""" + delta = lhs[:, 0] - rhs[:, 0] + greater = delta > _EQ_EPS + tied = np.abs(delta) <= _EQ_EPS + for i in range(1, 4): + delta = lhs[:, i] - rhs[:, i] + greater = greater | (tied & (delta > _EQ_EPS)) + tied = tied & (np.abs(delta) <= _EQ_EPS) + return greater + + +def _choose_canonical_candidates(candidates: np.ndarray) -> np.ndarray: + """Select canonical quaternions from candidate set. + + Candidates are expected in shape ``(N_candidates, N_quaternions, 4)`` and must + already be normalized with ``w >= 0``. Selection is by maximal lexicographic + ``(w, x, y, z)`` order (with epsilon-aware comparisons). + """ + best = candidates[0].copy() + for i in range(1, candidates.shape[0]): + current = candidates[i] + update = _lexicographically_greater(current, best) + best[update] = current[update] + return best + + +def proper_rotation_quaternions(matrices: np.ndarray) -> np.ndarray: + """Return proper rotations from matrices as quaternions in ``[w, x, y, z]``. + + Matrices with determinant ``<= _DET_EPS`` are discarded so only proper + rotations are used for equivalence classes in orientation mapping. + """ + matrices = np.asarray(matrices, dtype=np.float64).reshape(-1, 3, 3) + proper = matrices[np.linalg.det(matrices) > _DET_EPS] + if proper.size == 0: + proper = np.eye(3, dtype=np.float64).reshape(1, 3, 3) + quats_xyzw = Rotation.from_matrix(proper).as_quat() + quats_wxyz = np.column_stack((quats_xyzw[:, 3], quats_xyzw[:, :3])) + return _normalize_wxyz(quats_wxyz) + + +def operator_generated_quaternions(operator: np.ndarray) -> np.ndarray: + """Return proper rotations generated by repeated powers of an operator. + + The sequence starts at identity, multiplies by the input operator until + returning to identity (or ``_MAX_OPERATOR_ORDER``), keeps only proper + rotations, deduplicates numerically, and converts to normalized + ``[w, x, y, z]`` quaternions. + """ + operator = np.asarray(operator, dtype=np.float64).reshape(3, 3) + current = np.eye(3, dtype=np.float64) + proper_powers = [current.copy()] + for _ in range(1, _MAX_OPERATOR_ORDER + 1): + current = current @ operator + if np.linalg.det(current) > _DET_EPS: + if not any(np.allclose(current, m, atol=1e-6) for m in proper_powers): + proper_powers.append(current.copy()) + if np.allclose(current, np.eye(3), atol=1e-6): + break + return proper_rotation_quaternions(np.asarray(proper_powers)) + + +def map_quaternions_to_fundamental_zone( + quaternions: np.ndarray, operator_sets: list[np.ndarray] +) -> np.ndarray: + """Map each quaternion to a deterministic representative in its fundamental zone. + + Parameters + ---------- + quaternions + Array shaped ``(N_particles, N_slots, 4)`` in ``[w, x, y, z]`` order. + operator_sets + ``operator_sets[i]`` provides the proper-rotation equivalence set for + slot ``i``. This function left-multiplies each operator in that set onto + the corresponding quaternion column and selects the canonical candidate + by the module-level rule (max ``w``, then lexicographic tie-break). + """ + mapped = _normalize_wxyz(np.asarray(quaternions, dtype=np.float32)) + if mapped.ndim != 3 or mapped.shape[-1] != 4: + raise ValueError("quaternions must have shape (N, M, 4)") + if mapped.shape[1] != len(operator_sets): + raise ValueError("operator_sets length must equal quaternion symmetry dimension") + + for i, operators in enumerate(operator_sets): + operators = _normalize_wxyz(np.asarray(operators, dtype=np.float32)) + if operators.shape[0] == 1: + continue + # Equivalent optima are generated by left-multiplying symmetry operators + # onto the returned orientation for the corresponding slot. + candidates = _quat_multiply_wxyz( + operators[:, np.newaxis, :], mapped[:, i, :][np.newaxis, :, :] + ) + candidates = _normalize_wxyz(candidates) + mapped[:, i, :] = _choose_canonical_candidates(candidates) + + return mapped diff --git a/spatula/boosop.py b/spatula/boosop.py index 2d62b37e..132f9462 100644 --- a/spatula/boosop.py +++ b/spatula/boosop.py @@ -14,6 +14,10 @@ import spatula._spatula_nb from . import freud, integrate, representations, sph_harm, util +from ._fundamental_zone import ( + map_quaternions_to_fundamental_zone, + proper_rotation_quaternions, +) def _get_neighbors( @@ -119,10 +123,17 @@ def __init__( else: raise ValueError(f"Distribution {dist} not supported.") matrices = [] + # Slot-aligned fundamental-zone operator sets for returned rotations: + # one proper-rotation equivalence set per symmetry column. + self._rotation_fz_operators = [] for point_group in self._symmetries: matrices.append( representations.WignerD(point_group, self._max_l).condensed_matrices ) + group_matrices = representations.CartesianRepMatrix(point_group).matrices + self._rotation_fz_operators.append( + proper_rotation_quaternions(np.asarray(group_matrices)) + ) D_ij = np.stack(matrices, axis=0) # noqa N806 self._cpp = cls_(D_ij, optimizer._cpp, dist_param) self._order = None @@ -216,6 +227,9 @@ def compute( self._rotations = np.asarray(self._rotations).reshape( neighbors.num_query_points, -1, 4 ) + self._rotations = map_quaternions_to_fundamental_zone( + self._rotations, self._rotation_fz_operators + ) if refine: quad_positions, quad_weights = integrate.gauss_legendre_quad_points( m=refine_m, weights=True, cartesian=True @@ -269,6 +283,12 @@ def rotations(self) -> np.ndarray: each query particle and each point group. Rotations are expressed as quaternions. Note that these use different convention to scipy! The convention used here is [w,x,y,z]. The scipy convention is [x,y,z,w]. + + Rotations are canonicalized into a deterministic fundamental-zone + representative by evaluating all equivalent orientations generated by + left-multiplication with proper symmetry rotations and selecting the + candidate with maximal ``w`` (then lexicographic tie-break on + ``[w, x, y, z]``). """ if self._rotations is None: raise ValueError("BOOSOP not computed, call compute first.") diff --git a/spatula/pgop.py b/spatula/pgop.py index ed5bc077..e9880c7c 100644 --- a/spatula/pgop.py +++ b/spatula/pgop.py @@ -12,6 +12,11 @@ import spatula._spatula_nb from . import freud, representations +from ._fundamental_zone import ( + map_quaternions_to_fundamental_zone, + operator_generated_quaternions, + proper_rotation_quaternions, +) from .boosop import _get_neighbors @@ -62,6 +67,9 @@ def __init__( PGOP value. Defaults to False. `order` values are in order point group symmetry, order for symmetry operators of this point group in order given by the representations.matrices, order for second point group symmetry, etc. + Returned rotations are mapped to a deterministic fundamental-zone + representative. For per-operator outputs this mapping uses the + proper-rotation subgroup generated by each specific operator. """ if isinstance(symmetries, str): @@ -70,13 +78,26 @@ def __init__( # computing the PGOP self._optimizer = optimizer matrices = [] + # Slot-aligned fundamental-zone operator sets for returned rotations: + # one set per output symmetry column, including optional per-operator slots. + self._rotation_fz_operators = [] for point_group in self._symmetries: pg = representations.CartesianRepMatrix(point_group) + group_matrices = np.asarray(pg.matrices, dtype=np.float32) + self._rotation_fz_operators.append( + proper_rotation_quaternions(group_matrices) + ) # skips E operator if group is not C1 if point_group == "C1": - matrices.append(pg.condensed_matrices.astype(np.float32)) + operator_matrices = group_matrices else: - matrices.append(pg.condensed_matrices.astype(np.float32)[9:]) + operator_matrices = group_matrices[1:] + matrices.append(np.asarray(operator_matrices, dtype=np.float32).reshape(-1)) + if compute_per_operator_values_for_final_orientation: + for operator in operator_matrices: + self._rotation_fz_operators.append( + operator_generated_quaternions(operator) + ) if mode == "full": m_mode = 0 elif mode == "boo": @@ -215,6 +236,9 @@ def compute( neighbors.neighbor_counts.astype(np.int32), sigmas.astype(np.float32), ) + self._rotations = map_quaternions_to_fundamental_zone( + self._rotations, self._rotation_fz_operators + ) @property def order(self) -> np.ndarray: @@ -235,6 +259,12 @@ def rotations(self) -> np.ndarray: each query particle and each point group. Rotations are expressed as quaternions. Note that these use different convention to scipy! The convention used here is [w,x,y,z]. The scipy convention is [x,y,z,w]. + + Rotations are canonicalized into a deterministic fundamental-zone + representative by evaluating all equivalent orientations generated by + left-multiplication with proper symmetry rotations and selecting the + candidate with maximal ``w`` (then lexicographic tie-break on + ``[w, x, y, z]``). """ if self._rotations is None: raise ValueError("PGOP not computed, call compute first.") diff --git a/tests/test_pgop.py b/tests/test_pgop.py index 2bc37e07..992bc77c 100644 --- a/tests/test_pgop.py +++ b/tests/test_pgop.py @@ -1762,6 +1762,65 @@ def test_increasing_number_of_symmetries(n, mode): ) +def _normalize_wxyz(quaternion): + quaternion = np.asarray(quaternion, dtype=float) + quaternion = quaternion / np.linalg.norm(quaternion) + if quaternion[0] < 0: + quaternion = -quaternion + return quaternion + + +def _quat_mul_wxyz(lhs, rhs): + lw, lx, ly, lz = lhs + rw, rx, ry, rz = rhs + return np.array( + [ + lw * rw - lx * rx - ly * ry - lz * rz, + lw * rx + lx * rw + ly * rz - lz * ry, + lw * ry - lx * rz + ly * rw + lz * rx, + lw * rz + lx * ry - ly * rx + lz * rw, + ] + ) + + +def _group_equivalent_quaternions(symmetry): + matrices = np.asarray(spatula.representations.CartesianRepMatrix(symmetry).matrices) + proper_matrices = matrices[np.linalg.det(matrices) > 0] + scipy_quats = scipy.spatial.transform.Rotation.from_matrix( + proper_matrices + ).as_quat() + quats = np.column_stack((scipy_quats[:, 3], scipy_quats[:, :3])) + return np.asarray([_normalize_wxyz(q) for q in quats]) + + +def _operator_generated_quaternions(operator): + operator = np.asarray(operator, dtype=float) + current = np.eye(3) + powers = [current.copy()] + for _ in range(1, 257): + current = current @ operator + if np.linalg.det(current) > 0 and not any( + np.allclose(current, prev, atol=1e-6) for prev in powers + ): + powers.append(current.copy()) + if np.allclose(current, np.eye(3), atol=1e-6): + break + scipy_quats = scipy.spatial.transform.Rotation.from_matrix( + np.asarray(powers) + ).as_quat() + quats = np.column_stack((scipy_quats[:, 3], scipy_quats[:, :3])) + return np.asarray([_normalize_wxyz(q) for q in quats]) + + +def _assert_rotation_is_fz_canonical(rotation, equivalent_ops): + rotation = _normalize_wxyz(rotation) + best_w = rotation[0] + for op in equivalent_ops: + candidate = _normalize_wxyz(_quat_mul_wxyz(op, rotation)) + best_w = max(best_w, candidate[0]) + assert rotation[0] >= best_w - 1e-6 + + @pytest.mark.parametrize("mode", modedict_types) @pytest.mark.parametrize("symmetries", [["T"], ["T", "Th"]]) def test_orientations(mode, symmetries): @@ -1787,6 +1846,46 @@ def test_orientations(mode, symmetries): np.testing.assert_allclose(order, op_no_opt.order[0], rtol=RTOL) +@pytest.mark.parametrize("mode", modedict_types) +def test_rotations_mapped_to_group_fundamental_zone(mode): + rot = scipy.spatial.transform.Rotation.random(random_state=RNG) + rotated_vertices = rot.apply(VERTICES_FOR_TESTING) + system, nlist = get_shape_sys_nlist(rotated_vertices) + symmetries = ["T", "D4h"] + op = compute_op_result( + symmetries, OPTIMIZER, mode, system, nlist, None, np.zeros((1, 3)) + ) + + for rotation, symmetry in zip(op.rotations[0], symmetries): + equivalent_ops = _group_equivalent_quaternions(symmetry) + _assert_rotation_is_fz_canonical(rotation, equivalent_ops) + + +def test_pgop_per_operator_rotations_mapped_to_operator_fundamental_zone(): + rot = scipy.spatial.transform.Rotation.random(random_state=RNG) + rotated_vertices = rot.apply(VERTICES_FOR_TESTING) + system, nlist = get_shape_sys_nlist(rotated_vertices) + + op_pg = spatula.PGOP( + ["T"], + OPTIMIZER, + mode="full", + compute_per_operator_values_for_final_orientation=True, + ) + op_pg.compute(system, None, nlist, np.zeros((1, 3))) + + group_matrices = np.asarray( + spatula.representations.CartesianRepMatrix("T").matrices + ) + _assert_rotation_is_fz_canonical( + op_pg.rotations[0, 0], _group_equivalent_quaternions("T") + ) + for rotation, operator in zip(op_pg.rotations[0, 1:], group_matrices[1:]): + _assert_rotation_is_fz_canonical( + rotation, _operator_generated_quaternions(operator) + ) + + OPTIMIZERS_TO_TEST = [ ( "Union_descent_random",