add assert_allclose_up_to_global_phase testing utility function

python/ffsim/linalg/__init__.py

@@ -16,7 +16,12 @@
     modified_cholesky,
 )
 from ffsim.linalg.givens import apply_matrix_to_slices, givens_decomposition
-from ffsim.linalg.linalg import expm_multiply_taylor, lup, reduced_matrix
+from ffsim.linalg.linalg import (
+    expm_multiply_taylor,
+    lup,
+    match_global_phase,
+    reduced_matrix,
+)
 from ffsim.linalg.predicates import (
     is_antihermitian,
     is_hermitian,
@@ -39,6 +44,7 @@
     "is_special_orthogonal",
     "is_unitary",
     "lup",
+    "match_global_phase",
     "modified_cholesky",
     "reduced_matrix",
 ]

python/ffsim/linalg/linalg.py

@@ -12,6 +12,7 @@
 
 from __future__ import annotations
 
+import cmath
 from collections.abc import Sequence
 
 import numpy as np
@@ -66,3 +67,23 @@ def reduced_matrix(
         for i, state_i in enumerate(vecs):
             result[i, j] = np.vdot(state_i, mat_state_j)
     return result
+
+
+def match_global_phase(a: np.ndarray, b: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Phase the given arrays so that their phases match at one entry.
+
+    Args:
+        a: A Numpy array.
+        b: Another Numpy array.
+
+    Returns:
+        A pair of arrays (a', b') that are equal if and only if a == b * exp(i phi)
+        for some real number phi.
+    """
+    if a.shape != b.shape:
+        return a, b
+    # use the largest entry of one of the matrices to maximize precision
+    index = max(np.ndindex(*a.shape), key=lambda i: abs(b[i]))
+    phase_a = cmath.phase(a[index])
+    phase_b = cmath.phase(b[index])
+    return a * cmath.rect(1, -phase_a), b * cmath.rect(1, -phase_b)

python/ffsim/testing/__init__.py

@@ -10,9 +10,14 @@
 
 """Testing utilities."""
 
-from ffsim.testing.testing import random_nelec, random_occupied_orbitals
+from ffsim.testing.testing import (
+    assert_allclose_up_to_global_phase,
+    random_nelec,
+    random_occupied_orbitals,
+)
 
 __all__ = [
+    "assert_allclose_up_to_global_phase",
     "random_nelec",
     "random_occupied_orbitals",
 ]

python/ffsim/testing/testing.py

@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import numpy as np
+from ffsim.linalg import match_global_phase
 
 
 def random_nelec(norb: int, *, seed=None) -> tuple[int, int]:
@@ -50,3 +51,39 @@ def random_occupied_orbitals(
     occ_a = list(rng.choice(norb, n_alpha, replace=False))
     occ_b = list(rng.choice(norb, n_beta, replace=False))
     return (occ_a, occ_b)
+
+
+def assert_allclose_up_to_global_phase(
+    actual: np.ndarray,
+    desired: np.ndarray,
+    rtol: float = 1e-7,
+    atol: float = 0,
+    equal_nan: bool = True,
+    err_msg: str = "",
+    verbose: bool = True,
+):
+    """Check if a == b * exp(i phi) for some real number phi.
+
+    Args:
+        actual: A Numpy array.
+        desired: Another Numpy array.
+        rtol: Relative tolerance.
+        atol: Absolute tolerance.
+        equal_nan: If True, NaNs will compare equal.
+        err_msg: The error message to be printed in case of failure.
+        verbose: If True, the conflicting values are appended to the error message.
+
+    Raises:
+        AssertionError: If a and b are not equal up to global phase, up to the
+            specified precision.
+    """
+    actual, desired = match_global_phase(actual, desired)
+    np.testing.assert_allclose(
+        actual,
+        desired,
+        rtol=rtol,
+        atol=atol,
+        equal_nan=equal_nan,
+        err_msg=err_msg,
+        verbose=verbose,
+    )

tests/linalg/linalg_test.py

@@ -43,3 +43,15 @@ def test_reduced_matrix():
         actual = reduced_mat[i, j]
         expected = np.vdot(vecs[i], mat @ vecs[j])
         np.testing.assert_allclose(actual, expected)
+
+
+def test_match_global_phase():
+    a = np.array([[1, 2, 3], [4, 5, 6]])
+    b = 1j * a
+    c, d = ffsim.linalg.match_global_phase(a, b)
+    np.testing.assert_allclose(c, d)
+
+    a = np.array([[1, 2, 3], [4, 5, 6]])
+    b = 2j * a
+    c, d = ffsim.linalg.match_global_phase(a, b)
+    np.testing.assert_allclose(2 * c, d)

tests/testing_test.py

@@ -0,0 +1,32 @@
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Tests for testing utilities."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import ffsim
+
+
+def test_assert_allclose_up_to_global_phase():
+    rng = np.random.default_rng()
+    shape = (2, 3, 4)
+    a = rng.standard_normal(shape).astype(complex)
+    a += 1j + rng.standard_normal(shape)
+    b = a * (1j) ** 1.5
+
+    ffsim.testing.assert_allclose_up_to_global_phase(a, b)
+
+    b[0, 0, 0] *= 1.1
+    with pytest.raises(AssertionError):
+        ffsim.testing.assert_allclose_up_to_global_phase(a, b)