support interaction pairs for random ucj ops

python/ffsim/random/random.py

@@ -17,6 +17,7 @@
 from typing_extensions import deprecated
 
 from ffsim import hamiltonians, operators, variational
+from ffsim.variational.util import validate_interaction_pairs
 
 
 @deprecated(
@@ -387,6 +388,8 @@ def random_ucj_op_spin_balanced(
     *,
     n_reps: int = 1,
     with_final_orbital_rotation: bool = False,
+    interaction_pairs: tuple[list[tuple[int, int]] | None, list[tuple[int, int]] | None]
+    | None = None,
     seed=None,
 ) -> variational.UCJOpSpinBalanced:
     """Sample a random spin-balanced unitary cluster Jastrow (UCJ) operator.
@@ -396,12 +399,30 @@ def random_ucj_op_spin_balanced(
         n_reps: The number of ansatz repetitions.
         with_final_orbital_rotation: Whether to include a final orbital rotation
             in the operator.
+        interaction_pairs: Optional restrictions on allowed orbital interactions
+            for the diagonal Coulomb operators.
+            If specified, `interaction_pairs` should be a pair of lists,
+            for alpha-alpha and alpha-beta interactions, in that order.
+            Either list can be substituted with ``None`` to indicate no restrictions
+            on interactions.
+            Each list should contain pairs of integers representing the orbitals
+            that are allowed to interact. These pairs can also be interpreted as
+            indices of diagonal Coulomb matrix entries that are allowed to be
+            nonzero.
+            Each integer pair must be upper triangular, that is, of the form
+            :math:`(i, j)` where :math:`i \leq j`.
         seed: A seed to initialize the pseudorandom number generator.
             Should be a valid input to ``np.random.default_rng``.
 
     Returns:
         The sampled UCJ operator.
     """
+    if interaction_pairs is None:
+        interaction_pairs = (None, None)
+    pairs_aa, pairs_ab = interaction_pairs
+    validate_interaction_pairs(pairs_aa, ordered=False)
+    validate_interaction_pairs(pairs_ab, ordered=False)
+
     rng = np.random.default_rng(seed)
     diag_coulomb_mats = np.stack(
         [
@@ -420,6 +441,21 @@ def random_ucj_op_spin_balanced(
     final_orbital_rotation = None
     if with_final_orbital_rotation:
         final_orbital_rotation = random_unitary(norb, seed=rng)
+
+    # Zero out diagonal coulomb matrix entries if requested
+    if pairs_aa is not None:
+        mask = np.zeros((norb, norb), dtype=bool)
+        rows, cols = zip(*pairs_aa)
+        mask[rows, cols] = True
+        mask[cols, rows] = True
+        diag_coulomb_mats[:, 0] *= mask
+    if pairs_ab is not None:
+        mask = np.zeros((norb, norb), dtype=bool)
+        rows, cols = zip(*pairs_ab)
+        mask[rows, cols] = True
+        mask[cols, rows] = True
+        diag_coulomb_mats[:, 1] *= mask
+
     return variational.UCJOpSpinBalanced(
         diag_coulomb_mats=diag_coulomb_mats,
         orbital_rotations=orbital_rotations,
@@ -431,6 +467,12 @@ def random_ucj_op_spin_unbalanced(
     norb: int,
     *,
     n_reps: int = 1,
+    interaction_pairs: tuple[
+        list[tuple[int, int]] | None,
+        list[tuple[int, int]] | None,
+        list[tuple[int, int]] | None,
+    ]
+    | None = None,
     with_final_orbital_rotation: bool = False,
     seed=None,
 ) -> variational.UCJOpSpinUnbalanced:
@@ -439,6 +481,19 @@ def random_ucj_op_spin_unbalanced(
     Args:
         norb: The number of orbitals.
         n_reps: The number of ansatz repetitions.
+        interaction_pairs: Optional restrictions on allowed orbital interactions
+            for the diagonal Coulomb operators.
+            If specified, `interaction_pairs` should be a tuple of 3 lists,
+            for alpha-alpha, alpha-beta, and beta-beta interactions, in that order.
+            Any list can be substituted with ``None`` to indicate no restrictions
+            on interactions.
+            Each list should contain pairs of integers representing the orbitals
+            that are allowed to interact. These pairs can also be interpreted as
+            indices of diagonal Coulomb matrix entries that are allowed to be
+            nonzero.
+            For the alpha-alpha and beta-beta interactions, each integer
+            pair must be upper triangular, that is, of the form :math:`(i, j)` where
+            :math:`i \leq j`.
         with_final_orbital_rotation: Whether to include a final orbital rotation
             in the operator.
         seed: A seed to initialize the pseudorandom number generator.
@@ -447,6 +502,13 @@ def random_ucj_op_spin_unbalanced(
     Returns:
         The sampled UCJ operator.
     """
+    if interaction_pairs is None:
+        interaction_pairs = (None, None, None)
+    pairs_aa, pairs_ab, pairs_bb = interaction_pairs
+    validate_interaction_pairs(pairs_aa, ordered=False)
+    validate_interaction_pairs(pairs_bb, ordered=True)
+    validate_interaction_pairs(pairs_bb, ordered=False)
+
     rng = np.random.default_rng(seed)
     diag_coulomb_mats = np.stack(
         [
@@ -471,6 +533,26 @@ def random_ucj_op_spin_unbalanced(
         final_orbital_rotation = np.stack(
             [random_unitary(norb, seed=rng), random_unitary(norb, seed=rng)]
         )
+
+    # Zero out diagonal coulomb matrix entries if requested
+    if pairs_aa is not None:
+        mask = np.zeros((norb, norb), dtype=bool)
+        rows, cols = zip(*pairs_aa)
+        mask[rows, cols] = True
+        mask[cols, rows] = True
+        diag_coulomb_mats[:, 0] *= mask
+    if pairs_ab is not None:
+        mask = np.zeros((norb, norb), dtype=bool)
+        rows, cols = zip(*pairs_ab)
+        mask[rows, cols] = True
+        diag_coulomb_mats[:, 1] *= mask
+    if pairs_bb is not None:
+        mask = np.zeros((norb, norb), dtype=bool)
+        rows, cols = zip(*pairs_bb)
+        mask[rows, cols] = True
+        mask[cols, rows] = True
+        diag_coulomb_mats[:, 2] *= mask
+
     return variational.UCJOpSpinUnbalanced(
         diag_coulomb_mats=diag_coulomb_mats,
         orbital_rotations=orbital_rotations,
@@ -482,6 +564,7 @@ def random_ucj_op_spinless(
     norb: int,
     *,
     n_reps: int = 1,
+    interaction_pairs: list[tuple[int, int]] | None = None,
     with_final_orbital_rotation: bool = False,
     seed=None,
 ) -> variational.UCJOpSpinless:
@@ -492,12 +575,22 @@ def random_ucj_op_spinless(
         n_reps: The number of ansatz repetitions.
         with_final_orbital_rotation: Whether to include a final orbital rotation
             in the operator.
+        interaction_pairs: Optional restrictions on allowed orbital interactions
+            for the diagonal Coulomb operators.
+            If specified, `interaction_pairs` should be a list of integer pairs
+            representing the orbitals that are allowed to interact. These pairs
+            can also be interpreted as indices of diagonal Coulomb matrix entries
+            that are allowed to be nonzero.
+            Each integer pair must be upper triangular, that is, of the form
+            :math:`(i, j)` where :math:`i \leq j`.
         seed: A seed to initialize the pseudorandom number generator.
             Should be a valid input to ``np.random.default_rng``.
 
     Returns:
         The sampled UCJ operator.
     """
+    validate_interaction_pairs(interaction_pairs, ordered=False)
+
     rng = np.random.default_rng(seed)
     diag_coulomb_mats = np.stack(
         [random_real_symmetric_matrix(norb, seed=rng) for _ in range(n_reps)]
@@ -508,6 +601,15 @@ def random_ucj_op_spinless(
     final_orbital_rotation = None
     if with_final_orbital_rotation:
         final_orbital_rotation = random_unitary(norb, seed=rng)
+
+    # Zero out diagonal coulomb matrix entries if requested
+    if interaction_pairs is not None:
+        mask = np.zeros((norb, norb), dtype=bool)
+        rows, cols = zip(*interaction_pairs)
+        mask[rows, cols] = True
+        mask[cols, rows] = True
+        diag_coulomb_mats *= mask
+
     return variational.UCJOpSpinless(
         diag_coulomb_mats=diag_coulomb_mats,
         orbital_rotations=orbital_rotations,

python/ffsim/variational/ucj_spin_balanced.py

@@ -23,28 +23,10 @@
 from ffsim.variational.util import (
     orbital_rotation_from_parameters,
     orbital_rotation_to_parameters,
+    validate_interaction_pairs,
 )
 
 
-def _validate_interaction_pairs(
-    interaction_pairs: list[tuple[int, int]] | None, ordered: bool
-) -> None:
-    if interaction_pairs is None:
-        return
-    if len(set(interaction_pairs)) != len(interaction_pairs):
-        raise ValueError(
-            f"Duplicate interaction pairs encountered: {interaction_pairs}."
-        )
-    if not ordered:
-        for i, j in interaction_pairs:
-            if i > j:
-                raise ValueError(
-                    "When specifying alpha-alpha or beta-beta interaction pairs, "
-                    "you must provide only upper triangular pairs. "
-                    f"Got {(i, j)}, which is a lower triangular pair."
-                )
-
-
 @dataclass(frozen=True)
 class UCJOpSpinBalanced:
     r"""A spin-balanced unitary cluster Jastrow operator.
@@ -193,8 +175,8 @@ def n_params(
         if interaction_pairs is None:
             interaction_pairs = (None, None)
         pairs_aa, pairs_ab = interaction_pairs
-        _validate_interaction_pairs(pairs_aa, ordered=False)
-        _validate_interaction_pairs(pairs_ab, ordered=False)
+        validate_interaction_pairs(pairs_aa, ordered=False)
+        validate_interaction_pairs(pairs_ab, ordered=False)
         # Each diagonal Coulomb matrix has one parameter per upper triangular
         # entry unless indices are passed explicitly
         n_triu_indices = norb * (norb + 1) // 2
@@ -445,8 +427,8 @@ def from_t_amplitudes(
         if interaction_pairs is None:
             interaction_pairs = (None, None)
         pairs_aa, pairs_ab = interaction_pairs
-        _validate_interaction_pairs(pairs_aa, ordered=False)
-        _validate_interaction_pairs(pairs_ab, ordered=False)
+        validate_interaction_pairs(pairs_aa, ordered=False)
+        validate_interaction_pairs(pairs_ab, ordered=False)
 
         nocc, _, nvrt, _ = t2.shape
         norb = nocc + nvrt

python/ffsim/variational/ucj_spin_unbalanced.py

@@ -23,28 +23,10 @@
 from ffsim.variational.util import (
     orbital_rotation_from_parameters,
     orbital_rotation_to_parameters,
+    validate_interaction_pairs,
 )
 
 
-def _validate_interaction_pairs(
-    interaction_pairs: list[tuple[int, int]] | None, ordered: bool
-) -> None:
-    if interaction_pairs is None:
-        return
-    if len(set(interaction_pairs)) != len(interaction_pairs):
-        raise ValueError(
-            f"Duplicate interaction pairs encountered: {interaction_pairs}."
-        )
-    if not ordered:
-        for i, j in interaction_pairs:
-            if i > j:
-                raise ValueError(
-                    "When specifying alpha-alpha or beta-beta interaction pairs, "
-                    "you must provide only upper triangular pairs. "
-                    f"Got {(i, j)}, which is a lower triangular pair."
-                )
-
-
 @dataclass(frozen=True)
 class UCJOpSpinUnbalanced:
     r"""A spin-unbalanced unitary cluster Jastrow operator.
@@ -204,9 +186,9 @@ def n_params(
         if interaction_pairs is None:
             interaction_pairs = (None, None, None)
         pairs_aa, pairs_ab, pairs_bb = interaction_pairs
-        _validate_interaction_pairs(pairs_aa, ordered=False)
-        _validate_interaction_pairs(pairs_ab, ordered=True)
-        _validate_interaction_pairs(pairs_bb, ordered=False)
+        validate_interaction_pairs(pairs_aa, ordered=False)
+        validate_interaction_pairs(pairs_ab, ordered=True)
+        validate_interaction_pairs(pairs_bb, ordered=False)
         # Each same-spin diagonal Coulomb matrix has one parameter per upper triangular
         # entry unless indices are passed explicitly
         n_triu_indices = norb * (norb + 1) // 2
@@ -309,17 +291,28 @@ def from_parameters(
                 )
                 index += n_params
             # Diag Coulomb matrices
-            for indices, this_diag_coulomb_mat in zip(
-                (pairs_aa, pairs_ab, pairs_bb),
-                diag_coulomb_mat,
-            ):
-                if indices:
-                    n_params = len(indices)
-                    rows, cols = zip(*indices)
-                    vals = params[index : index + n_params]
-                    this_diag_coulomb_mat[cols, rows] = vals
-                    this_diag_coulomb_mat[rows, cols] = vals
-                    index += n_params
+            # for indices, this_diag_coulomb_mat in zip(
+            #     (pairs_aa, pairs_ab, pairs_bb),
+            #     diag_coulomb_mat,
+            # ):
+            #     if indices:
+            n_params = len(pairs_aa)
+            rows, cols = zip(*pairs_aa)
+            vals = params[index : index + n_params]
+            diag_coulomb_mat[0, cols, rows] = vals
+            diag_coulomb_mat[0, rows, cols] = vals
+            index += n_params
+            n_params = len(pairs_ab)
+            rows, cols = zip(*pairs_ab)
+            vals = params[index : index + n_params]
+            diag_coulomb_mat[1, rows, cols] = vals
+            index += n_params
+            n_params = len(pairs_bb)
+            rows, cols = zip(*pairs_bb)
+            vals = params[index : index + n_params]
+            diag_coulomb_mat[2, cols, rows] = vals
+            diag_coulomb_mat[2, rows, cols] = vals
+            index += n_params
         # Final orbital rotation
         final_orbital_rotation = None
         if with_final_orbital_rotation:
@@ -505,9 +498,9 @@ def from_t_amplitudes(
         if interaction_pairs is None:
             interaction_pairs = (None, None, None)
         pairs_aa, pairs_ab, pairs_bb = interaction_pairs
-        _validate_interaction_pairs(pairs_aa, ordered=False)
-        _validate_interaction_pairs(pairs_bb, ordered=True)
-        _validate_interaction_pairs(pairs_bb, ordered=False)
+        validate_interaction_pairs(pairs_aa, ordered=False)
+        validate_interaction_pairs(pairs_bb, ordered=True)
+        validate_interaction_pairs(pairs_bb, ordered=False)
 
         t2aa, t2ab, t2bb = t2
         nocc_a, nocc_b, nvrt_a, _ = t2ab.shape
@@ -615,7 +608,6 @@ def from_t_amplitudes(
             mask = np.zeros((norb, norb), dtype=bool)
             rows, cols = zip(*pairs_ab)
             mask[rows, cols] = True
-            mask[cols, rows] = True
             diag_coulomb_mats[:, 1] *= mask
         if pairs_bb is not None:
             mask = np.zeros((norb, norb), dtype=bool)

python/ffsim/variational/ucj_spinless.py

@@ -23,28 +23,10 @@
 from ffsim.variational.util import (
     orbital_rotation_from_parameters,
     orbital_rotation_to_parameters,
+    validate_interaction_pairs,
 )
 
 
-def _validate_interaction_pairs(
-    interaction_pairs: list[tuple[int, int]] | None, ordered: bool
-) -> None:
-    if interaction_pairs is None:
-        return
-    if len(set(interaction_pairs)) != len(interaction_pairs):
-        raise ValueError(
-            f"Duplicate interaction pairs encountered: {interaction_pairs}."
-        )
-    if not ordered:
-        for i, j in interaction_pairs:
-            if i > j:
-                raise ValueError(
-                    "When specifying interaction pairs, "
-                    "you must provide only upper triangular pairs. "
-                    f"Got {(i, j)}, which is a lower triangular pair."
-                )
-
-
 @dataclass(frozen=True)
 class UCJOpSpinless:
     r"""A spinless unitary cluster Jastrow operator.
@@ -171,7 +153,7 @@ def n_params(
             ValueError: Interaction pairs list contained duplicate interactions.
             ValueError: Interaction pairs list contained lower triangular pairs.
         """
-        _validate_interaction_pairs(interaction_pairs, ordered=False)
+        validate_interaction_pairs(interaction_pairs, ordered=False)
         # Each diagonal Coulomb matrix has one parameter per upper triangular
         # entry unless indices are passed explicitly
         n_triu_indices = norb * (norb + 1) // 2
@@ -381,7 +363,7 @@ def from_t_amplitudes(
             ValueError: Interaction pairs list contained duplicate interactions.
             ValueError: Interaction pairs list contained lower triangular pairs.
         """
-        _validate_interaction_pairs(interaction_pairs, ordered=False)
+        validate_interaction_pairs(interaction_pairs, ordered=False)
 
         nocc, _, nvrt, _ = t2.shape
         norb = nocc + nvrt