From 8e2645af94eea2e56403e299a0fa164a422b3ef3 Mon Sep 17 00:00:00 2001 From: "Kevin J. Sung" Date: Tue, 29 Oct 2024 21:14:31 -0400 Subject: [PATCH] support interaction pairs for random ucj ops (#335) --- python/ffsim/random/random.py | 108 +++++++++++++++++- python/ffsim/variational/ucj_spin_balanced.py | 28 +---- .../ffsim/variational/ucj_spin_unbalanced.py | 66 +++++------ python/ffsim/variational/ucj_spinless.py | 24 +--- python/ffsim/variational/util.py | 19 +++ .../variational/ucj_spin_balanced_test.py | 35 +++++- .../variational/ucj_spin_unbalanced_test.py | 35 +++++- tests/python/variational/ucj_spinless_test.py | 35 +++++- 8 files changed, 263 insertions(+), 87 deletions(-) diff --git a/python/ffsim/random/random.py b/python/ffsim/random/random.py index a5d4c05bb..88697b275 100644 --- a/python/ffsim/random/random.py +++ b/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,21 +388,41 @@ 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. + r"""Sample a random spin-balanced unitary cluster Jastrow (UCJ) operator. Args: norb: The number of spatial orbitals. 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,14 +467,33 @@ 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: - """Sample a random spin-unbalanced unitary cluster Jastrow (UCJ) operator. + r"""Sample a random spin-unbalanced unitary cluster Jastrow (UCJ) operator. 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,22 +564,33 @@ 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: - """Sample a random spinless unitary cluster Jastrow (UCJ) operator. + r"""Sample a random spinless unitary cluster Jastrow (UCJ) operator. Args: norb: The number of orbitals. 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, diff --git a/python/ffsim/variational/ucj_spin_balanced.py b/python/ffsim/variational/ucj_spin_balanced.py index bcdd2c9ec..69d2807c4 100644 --- a/python/ffsim/variational/ucj_spin_balanced.py +++ b/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 diff --git a/python/ffsim/variational/ucj_spin_unbalanced.py b/python/ffsim/variational/ucj_spin_unbalanced.py index 4c4fd2e2a..2279e5c68 100644 --- a/python/ffsim/variational/ucj_spin_unbalanced.py +++ b/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) diff --git a/python/ffsim/variational/ucj_spinless.py b/python/ffsim/variational/ucj_spinless.py index 472a608ab..e274348a3 100644 --- a/python/ffsim/variational/ucj_spinless.py +++ b/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 diff --git a/python/ffsim/variational/util.py b/python/ffsim/variational/util.py index 21da0bbec..7f4ff27ca 100644 --- a/python/ffsim/variational/util.py +++ b/python/ffsim/variational/util.py @@ -18,6 +18,25 @@ import scipy.linalg +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 spinless, alpha-alpha or beta-beta " + "interaction pairs, you must provide only upper triangular pairs. " + f"Got {(i, j)}, which is a lower triangular pair." + ) + + def orbital_rotation_to_parameters( orbital_rotation: np.ndarray, real: bool = False ) -> np.ndarray: diff --git a/tests/python/variational/ucj_spin_balanced_test.py b/tests/python/variational/ucj_spin_balanced_test.py index f92efcb71..cbd839fd3 100644 --- a/tests/python/variational/ucj_spin_balanced_test.py +++ b/tests/python/variational/ucj_spin_balanced_test.py @@ -65,7 +65,7 @@ def test_n_params(): ) -def test_parameters_roundtrip(): +def test_parameters_roundtrip_all_to_all(): rng = np.random.default_rng() norb = 5 n_reps = 2 @@ -95,6 +95,39 @@ def test_parameters_roundtrip(): ) +def test_parameters_roundtrip_interaction_pairs(): + rng = np.random.default_rng() + norb = 5 + n_reps = 2 + interaction_pairs = ([(0, 1)], [(0, 1)]) + + for with_final_orbital_rotation in [False, True]: + operator = ffsim.random.random_ucj_op_spin_balanced( + norb, + n_reps=n_reps, + interaction_pairs=interaction_pairs, + with_final_orbital_rotation=with_final_orbital_rotation, + seed=rng, + ) + roundtripped = ffsim.UCJOpSpinBalanced.from_parameters( + operator.to_parameters(interaction_pairs=interaction_pairs), + norb=norb, + n_reps=n_reps, + interaction_pairs=interaction_pairs, + with_final_orbital_rotation=with_final_orbital_rotation, + ) + np.testing.assert_allclose( + roundtripped.diag_coulomb_mats, operator.diag_coulomb_mats + ) + np.testing.assert_allclose( + roundtripped.orbital_rotations, operator.orbital_rotations + ) + if with_final_orbital_rotation: + np.testing.assert_allclose( + roundtripped.final_orbital_rotation, operator.final_orbital_rotation + ) + + def test_t_amplitudes_energy(): mol = pyscf.gto.Mole() mol.build( diff --git a/tests/python/variational/ucj_spin_unbalanced_test.py b/tests/python/variational/ucj_spin_unbalanced_test.py index 4e08b530a..5045411ce 100644 --- a/tests/python/variational/ucj_spin_unbalanced_test.py +++ b/tests/python/variational/ucj_spin_unbalanced_test.py @@ -68,7 +68,7 @@ def test_n_params(): ) -def test_parameters_roundtrip(): +def test_parameters_roundtrip_all_to_all(): rng = np.random.default_rng() norb = 5 n_reps = 2 @@ -98,6 +98,39 @@ def test_parameters_roundtrip(): ) +def test_parameters_roundtrip_interaction_pairs(): + rng = np.random.default_rng() + norb = 5 + n_reps = 2 + interaction_pairs = ([(0, 1)], [(0, 1)], [(0, 1)]) + + for with_final_orbital_rotation in [False, True]: + operator = ffsim.random.random_ucj_op_spin_unbalanced( + norb, + n_reps=n_reps, + interaction_pairs=interaction_pairs, + with_final_orbital_rotation=with_final_orbital_rotation, + seed=rng, + ) + roundtripped = ffsim.UCJOpSpinUnbalanced.from_parameters( + operator.to_parameters(interaction_pairs=interaction_pairs), + norb=norb, + n_reps=n_reps, + interaction_pairs=interaction_pairs, + with_final_orbital_rotation=with_final_orbital_rotation, + ) + np.testing.assert_allclose( + roundtripped.diag_coulomb_mats, operator.diag_coulomb_mats + ) + np.testing.assert_allclose( + roundtripped.orbital_rotations, operator.orbital_rotations + ) + if with_final_orbital_rotation: + np.testing.assert_allclose( + roundtripped.final_orbital_rotation, operator.final_orbital_rotation + ) + + def test_t_amplitudes_energy(): mol = pyscf.gto.Mole() mol.build( diff --git a/tests/python/variational/ucj_spinless_test.py b/tests/python/variational/ucj_spinless_test.py index 2ec1ec785..0ed27c546 100644 --- a/tests/python/variational/ucj_spinless_test.py +++ b/tests/python/variational/ucj_spinless_test.py @@ -60,7 +60,7 @@ def test_n_params(): ) -def test_parameters_roundtrip(): +def test_parameters_roundtrip_all_to_all(): rng = np.random.default_rng() norb = 5 n_reps = 2 @@ -90,6 +90,39 @@ def test_parameters_roundtrip(): ) +def test_parameters_roundtrip_interaction_pairs(): + rng = np.random.default_rng() + norb = 5 + n_reps = 2 + interaction_pairs = [(0, 1)] + + for with_final_orbital_rotation in [False, True]: + operator = ffsim.random.random_ucj_op_spinless( + norb, + n_reps=n_reps, + interaction_pairs=interaction_pairs, + with_final_orbital_rotation=with_final_orbital_rotation, + seed=rng, + ) + roundtripped = ffsim.UCJOpSpinless.from_parameters( + operator.to_parameters(interaction_pairs=interaction_pairs), + norb=norb, + n_reps=n_reps, + interaction_pairs=interaction_pairs, + with_final_orbital_rotation=with_final_orbital_rotation, + ) + np.testing.assert_allclose( + roundtripped.diag_coulomb_mats, operator.diag_coulomb_mats + ) + np.testing.assert_allclose( + roundtripped.orbital_rotations, operator.orbital_rotations + ) + if with_final_orbital_rotation: + np.testing.assert_allclose( + roundtripped.final_orbital_rotation, operator.final_orbital_rotation + ) + + def test_t_amplitudes_energy(): mol = pyscf.gto.Mole() mol.build(