factor orbital rotation from t1 amplitudes function

python/ffsim/variational/ucj_spin_balanced.py

@@ -17,11 +17,11 @@
 from typing import cast
 
 import numpy as np
-import scipy.linalg
 
 from ffsim import gates, linalg
 from ffsim.variational.util import (
     orbital_rotation_from_parameters,
+    orbital_rotation_from_t1_amplitudes,
     orbital_rotation_to_parameters,
     validate_interaction_pairs,
 )
@@ -451,10 +451,7 @@ def from_t_amplitudes(
 
         final_orbital_rotation = None
         if t1 is not None:
-            final_orbital_rotation_generator = np.zeros((norb, norb), dtype=complex)
-            final_orbital_rotation_generator[:nocc, nocc:] = t1
-            final_orbital_rotation_generator[nocc:, :nocc] = -t1.T
-            final_orbital_rotation = scipy.linalg.expm(final_orbital_rotation_generator)
+            final_orbital_rotation = orbital_rotation_from_t1_amplitudes(t1)
 
         # Zero out diagonal coulomb matrix entries if requested
         if pairs_aa is not None:

python/ffsim/variational/ucj_spin_unbalanced.py

@@ -17,11 +17,11 @@
 from typing import cast
 
 import numpy as np
-import scipy.linalg
 
 from ffsim import gates, linalg
 from ffsim.variational.util import (
     orbital_rotation_from_parameters,
+    orbital_rotation_from_t1_amplitudes,
     orbital_rotation_to_parameters,
     validate_interaction_pairs,
 )
@@ -579,20 +579,8 @@ def from_t_amplitudes(
         final_orbital_rotation = None
         if t1 is not None:
             t1a, t1b = t1
-
-            final_orbital_rotation_generator_a = np.zeros((norb, norb), dtype=complex)
-            final_orbital_rotation_generator_a[:nocc_a, nocc_a:] = t1a
-            final_orbital_rotation_generator_a[nocc_a:, :nocc_a] = -t1a.T
-            final_orbital_rotation_a = scipy.linalg.expm(
-                final_orbital_rotation_generator_a
-            )
-
-            final_orbital_rotation_generator_b = np.zeros((norb, norb), dtype=complex)
-            final_orbital_rotation_generator_b[:nocc_b, nocc_b:] = t1b
-            final_orbital_rotation_generator_b[nocc_b:, :nocc_b] = -t1b.T
-            final_orbital_rotation_b = scipy.linalg.expm(
-                final_orbital_rotation_generator_b
-            )
+            final_orbital_rotation_a = orbital_rotation_from_t1_amplitudes(t1a)
+            final_orbital_rotation_b = orbital_rotation_from_t1_amplitudes(t1b)
             final_orbital_rotation = np.stack(
                 [final_orbital_rotation_a, final_orbital_rotation_b]
             )

python/ffsim/variational/ucj_spinless.py

@@ -17,11 +17,11 @@
 from typing import cast
 
 import numpy as np
-import scipy.linalg
 
 from ffsim import gates, linalg
 from ffsim.variational.util import (
     orbital_rotation_from_parameters,
+    orbital_rotation_from_t1_amplitudes,
     orbital_rotation_to_parameters,
     validate_interaction_pairs,
 )
@@ -385,10 +385,7 @@ def from_t_amplitudes(
 
         final_orbital_rotation = None
         if t1 is not None:
-            final_orbital_rotation_generator = np.zeros((norb, norb), dtype=complex)
-            final_orbital_rotation_generator[:nocc, nocc:] = t1
-            final_orbital_rotation_generator[nocc:, :nocc] = -t1.T
-            final_orbital_rotation = scipy.linalg.expm(final_orbital_rotation_generator)
+            final_orbital_rotation = orbital_rotation_from_t1_amplitudes(t1)
 
         # Zero out diagonal coulomb matrix entries if requested
         if interaction_pairs is not None:

python/ffsim/variational/util.py

@@ -107,3 +107,22 @@ def orbital_rotation_from_parameters(
     generator[rows, cols] += vals
     generator[cols, rows] -= vals
     return scipy.linalg.expm(generator)
+
+
+def orbital_rotation_from_t1_amplitudes(t1: np.ndarray) -> np.ndarray:
+    """Construct an orbital rotation from t1 amplitudes.
+
+    Args:
+        t1: The t1 amplitudes.
+        norb: The number of spatial orbitals, which gives the width and height of the
+            orbital rotation matrix.
+
+    Returns:
+        The orbital rotation.
+    """
+    nocc, nvrt = t1.shape
+    norb = nocc + nvrt
+    generator = np.zeros((norb, norb), dtype=complex)
+    generator[:nocc, nocc:] = t1
+    generator[nocc:, :nocc] = -t1.T
+    return scipy.linalg.expm(generator)