From b3dfe77ce67320f440a7304ffd781a853f21112e Mon Sep 17 00:00:00 2001 From: "Kevin J. Sung" Date: Sun, 5 Nov 2023 14:58:38 -0500 Subject: [PATCH] create cistring module with cached substitutes for pyscf functions (#68) * create cistring module * from __future__ import annotations --- benchmarks/gates.py | 11 +-- python/ffsim/__init__.py | 2 + python/ffsim/cistring.py | 107 ++++++++++++++++++++++ python/ffsim/contract/diag_coulomb.py | 69 ++------------ python/ffsim/contract/num_op_sum.py | 48 ++-------- python/ffsim/gates/diag_coulomb.py | 31 ++----- python/ffsim/gates/num_op_sum.py | 22 +---- python/ffsim/gates/orbital_rotation.py | 70 +------------- python/ffsim/trotter/qdrift.py | 54 ----------- python/ffsim/trotter/trotter.py | 26 ------ tests/_slow/contract/diag_coulomb_test.py | 10 +- tests/_slow/contract/num_op_sum_test.py | 6 +- tests/_slow/gates/diag_coulomb_test.py | 18 +--- tests/_slow/gates/num_op_sum_test.py | 6 +- 14 files changed, 148 insertions(+), 332 deletions(-) create mode 100644 python/ffsim/cistring.py diff --git a/benchmarks/gates.py b/benchmarks/gates.py index b8f164feb..27c48e6dc 100644 --- a/benchmarks/gates.py +++ b/benchmarks/gates.py @@ -9,7 +9,6 @@ # that they have been altered from the originals. import numpy as np -from pyscf.fci import cistring import ffsim @@ -30,7 +29,6 @@ def setup(self, norb: int, filling_fraction: float): self.norb = norb nocc = int(norb * filling_fraction) self.nelec = (nocc, nocc) - n_alpha, n_beta = self.nelec rng = np.random.default_rng() @@ -42,12 +40,7 @@ def setup(self, norb: int, filling_fraction: float): self.diag_coulomb_mat = ffsim.random.random_real_symmetric_matrix( self.norb, seed=rng ) - self.occupations_a = cistring.gen_occslst(range(self.norb), n_alpha).astype( - np.uint, copy=False - ) - self.occupations_b = cistring.gen_occslst(range(self.norb), n_beta).astype( - np.uint, copy=False - ) + ffsim.init_cache(self.norb, self.nelec) def time_apply_orbital_rotation_givens(self, *_): ffsim.apply_orbital_rotation( @@ -85,8 +78,6 @@ def time_apply_diag_coulomb_evolution(self, *_): time=1.0, norb=self.norb, nelec=self.nelec, - occupations_a=self.occupations_a, - occupations_b=self.occupations_b, copy=False, ) diff --git a/python/ffsim/__init__.py b/python/ffsim/__init__.py index dc20c867e..55ffa135c 100644 --- a/python/ffsim/__init__.py +++ b/python/ffsim/__init__.py @@ -11,6 +11,7 @@ """ffsim is a software library for fast simulation of fermionic quantum circuits.""" from ffsim import contract, linalg, random, testing +from ffsim.cistring import init_cache from ffsim.gates import ( apply_diag_coulomb_evolution, apply_fsim_gate, @@ -104,6 +105,7 @@ "fermion_operator", "hartree_fock_state", "indices_to_strings", + "init_cache", "linalg", "linear_operator", "multireference_state", diff --git a/python/ffsim/cistring.py b/python/ffsim/cistring.py new file mode 100644 index 000000000..08934152e --- /dev/null +++ b/python/ffsim/cistring.py @@ -0,0 +1,107 @@ +# (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. + +"""Tools for handling FCI strings.""" + +from __future__ import annotations + +from functools import lru_cache + +import numpy as np +from pyscf.fci import cistring +from scipy.special import comb + +from ffsim._lib import gen_orbital_rotation_index_in_place + + +@lru_cache(maxsize=None) +def make_strings(orbitals: range, nocc: int) -> np.ndarray: + """Cached version of pyscf.fci.cistring.make_strings.""" + return cistring.make_strings(orbitals, nocc) + + +@lru_cache(maxsize=None) +def gen_occslst(orbitals: range, nocc: int) -> np.ndarray: + """Cached version of pyscf.fci.cistring.gen_occslst.""" + return cistring.gen_occslst(orbitals, nocc).astype(np.uint, copy=False) + + +@lru_cache(maxsize=None) +def gen_linkstr_index(orbitals: range, nocc: int): + """Cached version of pyscf.fci.cistring.gen_linkstr_index.""" + return cistring.gen_linkstr_index(orbitals, nocc) + + +@lru_cache(maxsize=None) +def gen_orbital_rotation_index( + norb: int, nocc: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Generate string index used for performing orbital rotations. + + Returns a tuple (diag_strings, off_diag_strings, off_diag_index) + of three Numpy arrays. + + diag_strings is a norb x binom(norb - 1, nocc - 1) array. + The i-th row of this array contains all the strings with orbital i occupied. + + off_diag_strings is a norb x binom(norb - 1, nocc) array. + The i-th row of this array contains all the strings with orbital i unoccupied. + + off_diag_index is a norb x binom(norb - 1, nocc) x nocc x 3 array. + The first two axes of this array are in one-to-one correspondence with + off_diag_strings. For a fixed choice (i, str0) for the first two axes, + the last two axes form a nocc x 3 array. Each row of this array is a tuple + (j, str1, sign) where str1 is formed by annihilating orbital j in str0 and creating + orbital i, with sign giving the fermionic parity of this operation. + """ + if nocc == 0: + diag_strings = np.zeros((norb, 0), dtype=np.uint) + off_diag_strings = np.zeros((norb, 1), dtype=np.uint) + off_diag_index = np.zeros((norb, 1, 0, 3), dtype=np.int32) + return diag_strings, off_diag_strings, off_diag_index + + linkstr_index = gen_linkstr_index(range(norb), nocc) + dim_diag = comb(norb - 1, nocc - 1, exact=True) + dim_off_diag = comb(norb - 1, nocc, exact=True) + dim = dim_diag + dim_off_diag + diag_strings = np.empty((norb, dim_diag), dtype=np.uint) + off_diag_strings = np.empty((norb, dim_off_diag), dtype=np.uint) + # TODO should this be int64? pyscf uses int32 for linkstr_index though + off_diag_index = np.empty((norb, dim_off_diag, nocc, 3), dtype=np.int32) + off_diag_strings_index = np.empty((norb, dim), dtype=np.uint) + gen_orbital_rotation_index_in_place( + norb=norb, + nocc=nocc, + linkstr_index=linkstr_index, + diag_strings=diag_strings, + off_diag_strings=off_diag_strings, + off_diag_strings_index=off_diag_strings_index, + off_diag_index=off_diag_index, + ) + return diag_strings, off_diag_strings, off_diag_index + + +def init_cache(norb: int, nelec: tuple[int, int]) -> None: + """Initialize cached objects. + + Call this function to prepare ffsim for performing operations with given values + of `norb` and `nelec`. Typically there is no need to call this function, but it + should be called before benchmarking to avoid counting the cost of initializing + cached lookup tables. + + Args: + norb: The number of spatial orbitals. + nelec: The number of alpha and beta electrons. + """ + for nocc in nelec: + make_strings(range(norb), nocc) + gen_occslst(range(norb), nocc) + gen_linkstr_index(range(norb), nocc) + gen_orbital_rotation_index(norb, nocc) diff --git a/python/ffsim/contract/diag_coulomb.py b/python/ffsim/contract/diag_coulomb.py index 0c133edd8..d753030c3 100644 --- a/python/ffsim/contract/diag_coulomb.py +++ b/python/ffsim/contract/diag_coulomb.py @@ -14,17 +14,14 @@ import numpy as np import scipy.sparse.linalg -from pyscf.fci import cistring from scipy.special import comb from ffsim._lib import ( contract_diag_coulomb_into_buffer_num_rep, contract_diag_coulomb_into_buffer_z_rep, ) -from ffsim.gates.orbital_rotation import ( - apply_orbital_rotation, - gen_orbital_rotation_index, -) +from ffsim.cistring import gen_occslst, make_strings +from ffsim.gates.orbital_rotation import apply_orbital_rotation from ffsim.states import dim @@ -36,10 +33,6 @@ def contract_diag_coulomb( *, mat_alpha_beta: np.ndarray | None = None, z_representation: bool = False, - occupations_a: np.ndarray | None = None, - occupations_b: np.ndarray | None = None, - strings_a: np.ndarray | None = None, - strings_b: np.ndarray | None = None, ) -> np.ndarray: r"""Contract a diagonal Coulomb operator with a vector. @@ -62,10 +55,6 @@ def contract_diag_coulomb( mat_alpha_beta: A matrix of coefficients to use for interactions between orbitals with differing spin. z_representation: Whether the input matrices are in the "Z" representation. - occupations_a: List of occupied orbital lists for alpha strings. - occupations_b: List of occupied orbital lists for beta strings. - strings_a: List of alpha strings. - strings_b: List of beta strings. Returns: The result of applying the diagonal Coulomb operator on the input state vector. @@ -74,39 +63,21 @@ def contract_diag_coulomb( if mat_alpha_beta is None: mat_alpha_beta = mat - n_alpha, n_beta = nelec - if z_representation: - if strings_a is None: - strings_a = cistring.make_strings(range(norb), n_alpha) - if strings_b is None: - strings_b = cistring.make_strings(range(norb), n_beta) return _contract_diag_coulomb_z_rep( vec, mat, norb=norb, nelec=nelec, mat_alpha_beta=mat_alpha_beta, - strings_a=strings_a, - strings_b=strings_b, ) - if occupations_a is None: - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - if occupations_b is None: - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) return _contract_diag_coulomb_num_rep( vec, mat, norb=norb, nelec=nelec, mat_alpha_beta=mat_alpha_beta, - occupations_a=occupations_a, - occupations_b=occupations_b, ) @@ -117,13 +88,14 @@ def _contract_diag_coulomb_num_rep( nelec: tuple[int, int], *, mat_alpha_beta: np.ndarray, - occupations_a: np.ndarray, - occupations_b: np.ndarray, ) -> np.ndarray: n_alpha, n_beta = nelec dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) + occupations_a = gen_occslst(range(norb), n_alpha) + occupations_b = gen_occslst(range(norb), n_beta) + mat = mat.copy() mat[np.diag_indices(norb)] *= 0.5 vec = vec.reshape((dim_a, dim_b)) @@ -148,13 +120,14 @@ def _contract_diag_coulomb_z_rep( nelec: tuple[int, int], *, mat_alpha_beta: np.ndarray, - strings_a: np.ndarray, - strings_b: np.ndarray, ) -> np.ndarray: n_alpha, n_beta = nelec dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) + strings_a = make_strings(range(norb), n_alpha) + strings_b = make_strings(range(norb), n_beta) + vec = vec.reshape((dim_a, dim_b)) out = np.zeros_like(vec) contract_diag_coulomb_into_buffer_z_rep( @@ -209,26 +182,8 @@ def diag_coulomb_linop( """ if mat_alpha_beta is None: mat_alpha_beta = mat - n_alpha, n_beta = nelec dim_ = dim(norb, nelec) - occupations_a = None - occupations_b = None - strings_a = None - strings_b = None - if z_representation: - strings_a = cistring.make_strings(range(norb), n_alpha) - strings_b = cistring.make_strings(range(norb), n_beta) - else: - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) - orbital_rotation_index_a = gen_orbital_rotation_index(norb, n_alpha) - orbital_rotation_index_b = gen_orbital_rotation_index(norb, n_beta) - def matvec(vec): this_mat = mat this_mat_alpha_beta = mat_alpha_beta @@ -239,8 +194,6 @@ def matvec(vec): norb, nelec, allow_row_permutation=True, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, ) this_mat = perm0 @ mat @ perm0.T this_mat_alpha_beta = perm0 @ mat_alpha_beta @ perm0.T @@ -251,10 +204,6 @@ def matvec(vec): nelec=nelec, mat_alpha_beta=this_mat_alpha_beta, z_representation=z_representation, - occupations_a=occupations_a, - occupations_b=occupations_b, - strings_a=strings_a, - strings_b=strings_b, ) if orbital_rotation is not None: vec, perm1 = apply_orbital_rotation( @@ -263,8 +212,6 @@ def matvec(vec): norb, nelec, allow_col_permutation=True, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) np.testing.assert_allclose(perm0, perm1.T) diff --git a/python/ffsim/contract/num_op_sum.py b/python/ffsim/contract/num_op_sum.py index 83afa55f8..d48310112 100644 --- a/python/ffsim/contract/num_op_sum.py +++ b/python/ffsim/contract/num_op_sum.py @@ -14,27 +14,18 @@ import numpy as np import scipy.sparse.linalg -from pyscf.fci import cistring from scipy.special import comb from ffsim._lib import ( contract_num_op_sum_spin_into_buffer, ) -from ffsim.gates.orbital_rotation import ( - apply_orbital_rotation, - gen_orbital_rotation_index, -) +from ffsim.cistring import gen_occslst +from ffsim.gates.orbital_rotation import apply_orbital_rotation from ffsim.states import dim def contract_num_op_sum( - vec: np.ndarray, - coeffs: np.ndarray, - norb: int, - nelec: tuple[int, int], - *, - occupations_a: np.ndarray | None = None, - occupations_b: np.ndarray | None = None, + vec: np.ndarray, coeffs: np.ndarray, norb: int, nelec: tuple[int, int] ): r"""Contract a linear combination of number operators with a vector. @@ -52,8 +43,6 @@ def contract_num_op_sum( coeffs: The coefficients of the linear combination. norb: The number of spatial orbitals. nelec: The number of alpha and beta electrons. - occupations_a: List of occupied orbital lists for alpha strings. - occupations_b: List of occupied orbital lists for beta strings. Returns: The result of applying the linear combination of number operators on the input @@ -62,14 +51,8 @@ def contract_num_op_sum( vec = vec.astype(complex, copy=False) n_alpha, n_beta = nelec - if occupations_a is None: - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - if occupations_b is None: - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) + occupations_a = gen_occslst(range(norb), n_alpha) + occupations_b = gen_occslst(range(norb), n_beta) dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) @@ -122,14 +105,6 @@ def num_op_sum_linop( """ n_alpha, n_beta = nelec dim_ = dim(norb, nelec) - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) - orbital_rotation_index_a = gen_orbital_rotation_index(norb, n_alpha) - orbital_rotation_index_b = gen_orbital_rotation_index(norb, n_beta) def matvec(vec): these_coeffs = coeffs @@ -140,18 +115,9 @@ def matvec(vec): norb, nelec, allow_row_permutation=True, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, ) these_coeffs = perm0 @ these_coeffs - vec = contract_num_op_sum( - vec, - these_coeffs, - norb=norb, - nelec=nelec, - occupations_a=occupations_a, - occupations_b=occupations_b, - ) + vec = contract_num_op_sum(vec, these_coeffs, norb=norb, nelec=nelec) if orbital_rotation is not None: vec, perm1 = apply_orbital_rotation( vec, @@ -159,8 +125,6 @@ def matvec(vec): norb, nelec, allow_col_permutation=True, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) np.testing.assert_allclose(perm0, perm1.T) diff --git a/python/ffsim/gates/diag_coulomb.py b/python/ffsim/gates/diag_coulomb.py index b335876bd..55e25b25b 100644 --- a/python/ffsim/gates/diag_coulomb.py +++ b/python/ffsim/gates/diag_coulomb.py @@ -15,13 +15,13 @@ from typing import cast import numpy as np -from pyscf.fci import cistring from scipy.special import comb from ffsim._lib import ( apply_diag_coulomb_evolution_in_place_num_rep, apply_diag_coulomb_evolution_in_place_z_rep, ) +from ffsim.cistring import gen_occslst, make_strings from ffsim.gates.orbital_rotation import apply_orbital_rotation @@ -35,12 +35,6 @@ def apply_diag_coulomb_evolution( *, mat_alpha_beta: np.ndarray | None = None, z_representation: bool = False, - occupations_a: np.ndarray | None = None, - occupations_b: np.ndarray | None = None, - strings_a: np.ndarray | None = None, - strings_b: np.ndarray | None = None, - orbital_rotation_index_a: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, - orbital_rotation_index_b: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, copy: bool = True, ) -> np.ndarray: r"""Apply time evolution by a (rotated) diagonal Coulomb operator. @@ -69,10 +63,7 @@ def apply_diag_coulomb_evolution( orbital_rotation: A unitary matrix describing the optional orbital rotation. mat_alpha_beta: A matrix of coefficients to use for interactions between orbitals with differing spin. - occupations_a: List of occupied orbital lists for alpha strings. - occupations_b: List of occupied orbital lists for beta strings. - orbital_rotation_index_a: The orbital rotation index for alpha strings. - orbital_rotation_index_b: The orbital rotation index for beta strings. + z_representation: Whether the input matrices are in the "Z" representation. copy: Whether to copy the vector before operating on it. - If ``copy=True`` then this function always returns a newly allocated vector and the original vector is left untouched. @@ -100,8 +91,6 @@ def apply_diag_coulomb_evolution( norb, nelec, allow_row_permutation=True, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) mat = perm0 @ mat @ perm0.T @@ -118,10 +107,8 @@ def apply_diag_coulomb_evolution( vec = vec.reshape((dim_a, dim_b)) if z_representation: - if strings_a is None: - strings_a = cistring.make_strings(range(norb), n_alpha) - if strings_b is None: - strings_b = cistring.make_strings(range(norb), n_beta) + strings_a = make_strings(range(norb), n_alpha) + strings_b = make_strings(range(norb), n_beta) apply_diag_coulomb_evolution_in_place_z_rep( vec, mat_exp, @@ -133,12 +120,8 @@ def apply_diag_coulomb_evolution( strings_b=strings_b, ) else: - if occupations_a is None: - occupations_a = cistring.gen_occslst(range(norb), n_alpha) - if occupations_b is None: - occupations_b = cistring.gen_occslst(range(norb), n_beta) - occupations_a = occupations_a.astype(np.uint, copy=False) - occupations_b = occupations_b.astype(np.uint, copy=False) + occupations_a = gen_occslst(range(norb), n_alpha) + occupations_b = gen_occslst(range(norb), n_beta) apply_diag_coulomb_evolution_in_place_num_rep( vec, mat_exp, @@ -156,8 +139,6 @@ def apply_diag_coulomb_evolution( norb, nelec, allow_col_permutation=True, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) np.testing.assert_allclose(perm0, perm1.T) diff --git a/python/ffsim/gates/num_op_sum.py b/python/ffsim/gates/num_op_sum.py index 5c5096a65..2d01f4a0a 100644 --- a/python/ffsim/gates/num_op_sum.py +++ b/python/ffsim/gates/num_op_sum.py @@ -13,10 +13,10 @@ from __future__ import annotations import numpy as np -from pyscf.fci import cistring from scipy.special import comb from ffsim._lib import apply_num_op_sum_evolution_in_place +from ffsim.cistring import gen_occslst from ffsim.gates.orbital_rotation import apply_orbital_rotation @@ -28,10 +28,6 @@ def apply_num_op_sum_evolution( nelec: tuple[int, int], orbital_rotation: np.ndarray | None = None, *, - occupations_a: np.ndarray | None = None, - occupations_b: np.ndarray | None = None, - orbital_rotation_index_a: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, - orbital_rotation_index_b: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, copy: bool = True, ): r"""Apply time evolution by a (rotated) linear combination of number operators. @@ -55,10 +51,6 @@ def apply_num_op_sum_evolution( norb: The number of spatial orbitals. nelec: The number of alpha and beta electrons. orbital_rotation: A unitary matrix describing the optional orbital rotation. - occupations_a: List of occupied orbital lists for alpha strings. - occupations_b: List of occupied orbital lists for beta strings. - orbital_rotation_index_a: The orbital rotation index for alpha strings. - orbital_rotation_index_b: The orbital rotation index for beta strings. copy: Whether to copy the vector before operating on it. - If ``copy=True`` then this function always returns a newly allocated vector and the original vector is left untouched. @@ -82,12 +74,8 @@ def apply_num_op_sum_evolution( n_alpha, n_beta = nelec dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) - if occupations_a is None: - occupations_a = cistring.gen_occslst(range(norb), n_alpha) - if occupations_b is None: - occupations_b = cistring.gen_occslst(range(norb), n_beta) - occupations_a = occupations_a.astype(np.uint, copy=False) - occupations_b = occupations_b.astype(np.uint, copy=False) + occupations_a = gen_occslst(range(norb), n_alpha) + occupations_b = gen_occslst(range(norb), n_beta) if orbital_rotation is not None: vec, perm0 = apply_orbital_rotation( @@ -96,8 +84,6 @@ def apply_num_op_sum_evolution( norb, nelec, allow_row_permutation=True, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) coeffs = perm0 @ coeffs @@ -118,8 +104,6 @@ def apply_num_op_sum_evolution( norb, nelec, allow_col_permutation=True, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) np.testing.assert_allclose(perm1.T, perm0) diff --git a/python/ffsim/gates/orbital_rotation.py b/python/ffsim/gates/orbital_rotation.py index da6313613..fdda2da44 100644 --- a/python/ffsim/gates/orbital_rotation.py +++ b/python/ffsim/gates/orbital_rotation.py @@ -24,60 +24,11 @@ apply_givens_rotation_in_place, apply_phase_shift_in_place, apply_single_column_transformation_in_place, - gen_orbital_rotation_index_in_place, ) +from ffsim.cistring import gen_orbital_rotation_index from ffsim.linalg import givens_decomposition, lup -def gen_orbital_rotation_index( - norb: int, nocc: int, linkstr_index: np.ndarray | None = None -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Generate string index used for performing orbital rotations. - - Returns a tuple (diag_strings, off_diag_strings, off_diag_index) - of three Numpy arrays. - - diag_strings is a norb x binom(norb - 1, nocc - 1) array. - The i-th row of this array contains all the strings with orbital i occupied. - - off_diag_strings is a norb x binom(norb - 1, nocc) array. - The i-th row of this array contains all the strings with orbital i unoccupied. - - off_diag_index is a norb x binom(norb - 1, nocc) x nocc x 3 array. - The first two axes of this array are in one-to-one correspondence with - off_diag_strings. For a fixed choice (i, str0) for the first two axes, - the last two axes form a nocc x 3 array. Each row of this array is a tuple - (j, str1, sign) where str1 is formed by annihilating orbital j in str0 and creating - orbital i, with sign giving the fermionic parity of this operation. - """ - if nocc == 0: - diag_strings = np.zeros((norb, 0), dtype=np.uint) - off_diag_strings = np.zeros((norb, 1), dtype=np.uint) - off_diag_index = np.zeros((norb, 1, 0, 3), dtype=np.int32) - return diag_strings, off_diag_strings, off_diag_index - - if linkstr_index is None: - linkstr_index = cistring.gen_linkstr_index(range(norb), nocc) - dim_diag = comb(norb - 1, nocc - 1, exact=True) - dim_off_diag = comb(norb - 1, nocc, exact=True) - dim = dim_diag + dim_off_diag - diag_strings = np.empty((norb, dim_diag), dtype=np.uint) - off_diag_strings = np.empty((norb, dim_off_diag), dtype=np.uint) - # TODO should this be int64? pyscf uses int32 for linkstr_index though - off_diag_index = np.empty((norb, dim_off_diag, nocc, 3), dtype=np.int32) - off_diag_strings_index = np.empty((norb, dim), dtype=np.uint) - gen_orbital_rotation_index_in_place( - norb=norb, - nocc=nocc, - linkstr_index=linkstr_index, - diag_strings=diag_strings, - off_diag_strings=off_diag_strings, - off_diag_strings_index=off_diag_strings_index, - off_diag_index=off_diag_index, - ) - return diag_strings, off_diag_strings, off_diag_index - - @overload def apply_orbital_rotation( vec: np.ndarray, @@ -98,8 +49,6 @@ def apply_orbital_rotation( nelec: tuple[int, int], *, allow_row_permutation: Literal[True], - orbital_rotation_index_a: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, - orbital_rotation_index_b: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, copy: bool = True, ) -> tuple[np.ndarray, np.ndarray]: ... @@ -113,8 +62,6 @@ def apply_orbital_rotation( nelec: tuple[int, int], *, allow_col_permutation: Literal[True], - orbital_rotation_index_a: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, - orbital_rotation_index_b: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, copy: bool = True, ) -> tuple[np.ndarray, np.ndarray]: ... @@ -128,8 +75,6 @@ def apply_orbital_rotation( *, allow_row_permutation: bool = False, allow_col_permutation: bool = False, - orbital_rotation_index_a: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, - orbital_rotation_index_b: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None, copy: bool = True, ) -> np.ndarray | tuple[np.ndarray, np.ndarray]: r"""Apply an orbital rotation to a vector. @@ -157,8 +102,6 @@ def apply_orbital_rotation( of the orbital rotation matrix. allow_col_permutation: Whether to allow a permutation of the columns of the orbital rotation matrix. - orbital_rotation_index_a: The orbital rotation index for alpha strings. - orbital_rotation_index_b: The orbital rotation index for beta strings. copy: Whether to copy the vector before operating on it. - If ``copy=True`` then this function always returns a newly allocated vector and the original vector is left untouched. @@ -183,19 +126,12 @@ def apply_orbital_rotation( "You can choose to allow either row or column permutations, but not both." ) if allow_row_permutation or allow_col_permutation: - n_alpha, n_beta = nelec - if orbital_rotation_index_a is None: - orbital_rotation_index_a = gen_orbital_rotation_index(norb, n_alpha) - if orbital_rotation_index_b is None: - orbital_rotation_index_b = gen_orbital_rotation_index(norb, n_beta) return _apply_orbital_rotation_lu( vec, mat, norb, nelec, permute_rows=allow_row_permutation, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=copy, ) return _apply_orbital_rotation_givens(vec, mat, norb, nelec, copy=copy) @@ -207,8 +143,6 @@ def _apply_orbital_rotation_lu( norb: int, nelec: tuple[int, int], permute_rows: bool, - orbital_rotation_index_a: tuple[np.ndarray, np.ndarray, np.ndarray], - orbital_rotation_index_b: tuple[np.ndarray, np.ndarray, np.ndarray], copy: bool, ) -> tuple[np.ndarray, np.ndarray]: if copy: @@ -220,6 +154,8 @@ def _apply_orbital_rotation_lu( eye = np.eye(norb, dtype=complex) transformation_mat = eye - lower + scipy.linalg.solve_triangular(upper, eye) n_alpha, n_beta = nelec + orbital_rotation_index_a = gen_orbital_rotation_index(norb, n_alpha) + orbital_rotation_index_b = gen_orbital_rotation_index(norb, n_beta) dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) vec = vec.reshape((dim_a, dim_b)) diff --git a/python/ffsim/trotter/qdrift.py b/python/ffsim/trotter/qdrift.py index 7f377fa53..b2c245610 100644 --- a/python/ffsim/trotter/qdrift.py +++ b/python/ffsim/trotter/qdrift.py @@ -14,10 +14,8 @@ import numpy as np import scipy.linalg -from pyscf.fci import cistring from ffsim.gates import apply_diag_coulomb_evolution, apply_num_op_sum_evolution -from ffsim.gates.orbital_rotation import gen_orbital_rotation_index from ffsim.hamiltonians import DoubleFactorizedHamiltonian from ffsim.states.wick import expectation_one_body_power, expectation_one_body_product @@ -93,14 +91,6 @@ def simulate_qdrift_double_factorized( ) step_time = time / n_steps n_alpha, n_beta = nelec - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) - orbital_rotation_index_a = gen_orbital_rotation_index(norb, n_alpha) - orbital_rotation_index_b = gen_orbital_rotation_index(norb, n_beta) results = np.empty((n_samples, initial_state.shape[0]), dtype=complex) for i in range(n_samples): @@ -122,10 +112,6 @@ def simulate_qdrift_double_factorized( nelec=nelec, probabilities=probabilities, term_indices=term_indices, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, ) else: vec = _simulate_qdrift_step_double_factorized( @@ -140,10 +126,6 @@ def simulate_qdrift_double_factorized( nelec=nelec, probabilities=probabilities, term_indices=term_indices, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, ) results[i] = vec @@ -164,10 +146,6 @@ def _simulate_qdrift_step_double_factorized( nelec: tuple[int, int], probabilities: np.ndarray, term_indices: np.ndarray, - occupations_a: np.ndarray, - occupations_b: np.ndarray, - orbital_rotation_index_a: tuple[np.ndarray, np.ndarray, np.ndarray], - orbital_rotation_index_b: tuple[np.ndarray, np.ndarray, np.ndarray], ) -> np.ndarray: for term_index in term_indices: if term_index == 0: @@ -178,10 +156,6 @@ def _simulate_qdrift_step_double_factorized( norb=norb, nelec=nelec, orbital_rotation=one_body_basis_change, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) else: @@ -193,10 +167,6 @@ def _simulate_qdrift_step_double_factorized( nelec=nelec, orbital_rotation=orbital_rotations[term_index - 1], z_representation=z_representation, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) return vec @@ -214,10 +184,6 @@ def _simulate_qdrift_step_double_factorized_symmetric( nelec: tuple[int, int], probabilities: np.ndarray, term_indices: np.ndarray, - occupations_a: np.ndarray, - occupations_b: np.ndarray, - orbital_rotation_index_a: tuple[np.ndarray, np.ndarray, np.ndarray], - orbital_rotation_index_b: tuple[np.ndarray, np.ndarray, np.ndarray], ) -> np.ndarray: # simulate the one-body term for half the time vec = apply_num_op_sum_evolution( @@ -227,10 +193,6 @@ def _simulate_qdrift_step_double_factorized_symmetric( norb=norb, nelec=nelec, orbital_rotation=one_body_basis_change, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) # simulate the first sampled term @@ -243,10 +205,6 @@ def _simulate_qdrift_step_double_factorized_symmetric( nelec=nelec, orbital_rotation=orbital_rotations[term_index - 1], z_representation=z_representation, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) # simulate the remaining sampled terms @@ -258,10 +216,6 @@ def _simulate_qdrift_step_double_factorized_symmetric( norb=norb, nelec=nelec, orbital_rotation=one_body_basis_change, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) vec = apply_diag_coulomb_evolution( @@ -272,10 +226,6 @@ def _simulate_qdrift_step_double_factorized_symmetric( nelec=nelec, orbital_rotation=orbital_rotations[term_index - 1], z_representation=z_representation, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) # simulate the one-body term for half the time @@ -286,10 +236,6 @@ def _simulate_qdrift_step_double_factorized_symmetric( norb=norb, nelec=nelec, orbital_rotation=one_body_basis_change, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) diff --git a/python/ffsim/trotter/trotter.py b/python/ffsim/trotter/trotter.py index 0e370355b..6594d6729 100644 --- a/python/ffsim/trotter/trotter.py +++ b/python/ffsim/trotter/trotter.py @@ -13,10 +13,8 @@ from collections.abc import Iterator import numpy as np -from pyscf.fci import cistring from ffsim.gates import apply_diag_coulomb_evolution, apply_num_op_sum_evolution -from ffsim.gates.orbital_rotation import gen_orbital_rotation_index from ffsim.hamiltonians import DoubleFactorizedHamiltonian @@ -99,14 +97,6 @@ def simulate_trotter_double_factorized( ) step_time = time / n_steps n_alpha, n_beta = nelec - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) - orbital_rotation_index_a = gen_orbital_rotation_index(norb, n_alpha) - orbital_rotation_index_b = gen_orbital_rotation_index(norb, n_beta) for _ in range(n_steps): vec = _simulate_trotter_step_double_factorized( @@ -120,10 +110,6 @@ def simulate_trotter_double_factorized( nelec=nelec, order=order, z_representation=hamiltonian.z_representation, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, ) return vec @@ -140,10 +126,6 @@ def _simulate_trotter_step_double_factorized( nelec: tuple[int, int], order: int, z_representation: bool, - occupations_a: np.ndarray, - occupations_b: np.ndarray, - orbital_rotation_index_a: tuple[np.ndarray, np.ndarray, np.ndarray], - orbital_rotation_index_b: tuple[np.ndarray, np.ndarray, np.ndarray], ) -> np.ndarray: for term_index, time in _simulate_trotter_step_iterator( 1 + len(diag_coulomb_mats), time, order @@ -156,10 +138,6 @@ def _simulate_trotter_step_double_factorized( norb=norb, nelec=nelec, orbital_rotation=one_body_basis_change, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) else: @@ -171,10 +149,6 @@ def _simulate_trotter_step_double_factorized( nelec=nelec, orbital_rotation=orbital_rotations[term_index - 1], z_representation=z_representation, - occupations_a=occupations_a, - occupations_b=occupations_b, - orbital_rotation_index_a=orbital_rotation_index_a, - orbital_rotation_index_b=orbital_rotation_index_b, copy=False, ) return vec diff --git a/tests/_slow/contract/diag_coulomb_test.py b/tests/_slow/contract/diag_coulomb_test.py index 21dcc96b8..c5141efd9 100644 --- a/tests/_slow/contract/diag_coulomb_test.py +++ b/tests/_slow/contract/diag_coulomb_test.py @@ -13,10 +13,10 @@ from __future__ import annotations import numpy as np -from pyscf.fci import cistring from scipy.special import comb import ffsim +from ffsim import cistring from ffsim._lib import ( contract_diag_coulomb_into_buffer_num_rep, contract_diag_coulomb_into_buffer_z_rep, @@ -36,12 +36,8 @@ def test_contract_diag_coulomb_into_buffer_num_rep_slow(): n_beta = rng.integers(1, norb + 1) dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) + occupations_a = cistring.gen_occslst(range(norb), n_alpha) + occupations_b = cistring.gen_occslst(range(norb), n_beta) mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) vec = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( diff --git a/tests/_slow/contract/num_op_sum_test.py b/tests/_slow/contract/num_op_sum_test.py index a59ce546f..925b2a915 100644 --- a/tests/_slow/contract/num_op_sum_test.py +++ b/tests/_slow/contract/num_op_sum_test.py @@ -13,10 +13,10 @@ from __future__ import annotations import numpy as np -from pyscf.fci import cistring from scipy.special import comb import ffsim +from ffsim import cistring from ffsim._lib import contract_num_op_sum_spin_into_buffer from ffsim._slow.contract.num_op_sum import contract_num_op_sum_spin_into_buffer_slow @@ -30,9 +30,7 @@ def test_contract_num_op_sum_spin_into_buffer_slow(): n_beta = rng.integers(1, norb + 1) dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) - occupations = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) + occupations = cistring.gen_occslst(range(norb), n_alpha) coeffs = np.random.uniform(size=norb) vec = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( (dim_a, dim_b) diff --git a/tests/_slow/gates/diag_coulomb_test.py b/tests/_slow/gates/diag_coulomb_test.py index 04c6dff80..92b661722 100644 --- a/tests/_slow/gates/diag_coulomb_test.py +++ b/tests/_slow/gates/diag_coulomb_test.py @@ -11,10 +11,10 @@ from __future__ import annotations import numpy as np -from pyscf.fci import cistring from scipy.special import comb import ffsim +from ffsim import cistring from ffsim._lib import ( apply_diag_coulomb_evolution_in_place_num_rep, apply_diag_coulomb_evolution_in_place_z_rep, @@ -35,12 +35,8 @@ def test_apply_diag_coulomb_evolution_num_rep_slow(): n_beta = rng.integers(1, norb + 1) dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) + occupations_a = cistring.gen_occslst(range(norb), n_alpha) + occupations_b = cistring.gen_occslst(range(norb), n_beta) time = 0.6 mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) mat_exp = np.exp(-1j * time * mat) @@ -121,12 +117,8 @@ def test_apply_diag_coulomb_evolution_num_rep_numpy(): n_beta = rng.integers(1, norb + 1) dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) + occupations_a = cistring.gen_occslst(range(norb), n_alpha) + occupations_b = cistring.gen_occslst(range(norb), n_beta) time = 0.6 mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) mat_exp = np.exp(-1j * time * mat) diff --git a/tests/_slow/gates/num_op_sum_test.py b/tests/_slow/gates/num_op_sum_test.py index 365b0e501..679a3be66 100644 --- a/tests/_slow/gates/num_op_sum_test.py +++ b/tests/_slow/gates/num_op_sum_test.py @@ -11,10 +11,10 @@ from __future__ import annotations import numpy as np -from pyscf.fci import cistring from scipy.special import comb import ffsim +from ffsim import cistring from ffsim._lib import apply_num_op_sum_evolution_in_place from ffsim._slow.gates.num_op_sum import apply_num_op_sum_evolution_in_place_slow @@ -28,9 +28,7 @@ def test_apply_num_op_sum_evolution_in_place_slow(): n_beta = rng.integers(1, norb + 1) dim_a = comb(norb, n_alpha, exact=True) dim_b = comb(norb, n_beta, exact=True) - occupations = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) + occupations = cistring.gen_occslst(range(norb), n_alpha) exponents = np.random.uniform(0, 2 * np.pi, size=norb) phases = np.exp(1j * exponents) vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape(