Skip to content

Commit

Permalink
support interaction pairs for random ucj ops
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevinsung
kevinsung committed Oct 30, 2024
1 parent 6d96a0f commit f72a7cf
Show file tree
Hide file tree
Showing 8 changed files with 263 additions and 87 deletions.
108 changes: 105 additions & 3 deletions python/ffsim/random/random.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from typing_extensions import deprecated

from ffsim import hamiltonians, operators, variational
from ffsim.variational.util import validate_interaction_pairs


@deprecated(
Expand Down Expand Up @@ -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(
[
Expand All @@ -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,
Expand All @@ -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.
Expand All @@ -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(
[
Expand All @@ -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,
Expand All @@ -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)]
Expand All @@ -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,
Expand Down
28 changes: 5 additions & 23 deletions python/ffsim/variational/ucj_spin_balanced.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
66 changes: 29 additions & 37 deletions python/ffsim/variational/ucj_spin_unbalanced.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
Loading

0 comments on commit f72a7cf

Please sign in to comment.