diff --git a/python/ffsim/contract/one_body.py b/python/ffsim/contract/one_body.py index 235ccd6a1..8f866bd2b 100644 --- a/python/ffsim/contract/one_body.py +++ b/python/ffsim/contract/one_body.py @@ -14,19 +14,14 @@ import numpy as np import scipy.sparse.linalg -from pyscf.fci import cistring from pyscf.fci.direct_nosym import contract_1e +from ffsim.cistring import gen_linkstr_index from ffsim.states import dim def contract_one_body( - vec: np.ndarray, - mat: np.ndarray, - norb: int, - nelec: tuple[int, int], - *, - link_index: tuple[np.ndarray, np.ndarray] | None = None, + vec: np.ndarray, mat: np.ndarray, norb: int, nelec: tuple[int, int] ) -> np.ndarray: r"""Contract a one-body tensor with a vector. @@ -42,14 +37,14 @@ def contract_one_body( mat: The one-body tensor. norb: The number of spatial orbitals. nelec: The number of alpha and beta electrons. - link_index: String index lookup tables for spin alpha and spin beta, as - generated by the function `pyscf.fci.cistring.gen_linkstr_index`_. Returns: A LinearOperator that implements the action of the one-body tensor. - - .. _pyscf.fci.cistring.gen_linkstr_index: https://pyscf.org/pyscf_api_docs/pyscf.fci.html#pyscf.fci.cistring.gen_linkstr_index """ + n_alpha, n_beta = nelec + link_index_a = gen_linkstr_index(range(norb), n_alpha) + link_index_b = gen_linkstr_index(range(norb), n_beta) + link_index = (link_index_a, link_index_b) result = contract_1e(mat.real, vec.real, norb, nelec, link_index=link_index).astype( complex ) @@ -81,20 +76,12 @@ def one_body_linop( A LinearOperator that implements the action of the one-body tensor. """ dim_ = dim(norb, nelec) - n_alpha, n_beta = nelec - link_index_a = cistring.gen_linkstr_index(range(norb), n_alpha) - link_index_b = cistring.gen_linkstr_index(range(norb), n_beta) - link_index = (link_index_a, link_index_b) def matvec(vec: np.ndarray): - return contract_one_body( - vec, mat, norb=norb, nelec=nelec, link_index=link_index - ) + return contract_one_body(vec, mat, norb=norb, nelec=nelec) def rmatvec(vec: np.ndarray): - return contract_one_body( - vec, mat.T.conj(), norb=norb, nelec=nelec, link_index=link_index - ) + return contract_one_body(vec, mat.T.conj(), norb=norb, nelec=nelec) return scipy.sparse.linalg.LinearOperator( shape=(dim_, dim_), matvec=matvec, rmatvec=rmatvec, dtype=complex diff --git a/python/ffsim/hamiltonians/molecular_hamiltonian.py b/python/ffsim/hamiltonians/molecular_hamiltonian.py index 5e1386450..1f868d540 100644 --- a/python/ffsim/hamiltonians/molecular_hamiltonian.py +++ b/python/ffsim/hamiltonians/molecular_hamiltonian.py @@ -15,12 +15,12 @@ import numpy as np import scipy.sparse.linalg -from pyscf.fci import cistring from pyscf.fci.direct_nosym import absorb_h1e, make_hdiag from pyscf.fci.fci_slow import contract_2e from scipy.sparse.linalg import LinearOperator from ffsim._lib import FermionOperator +from ffsim.cistring import gen_linkstr_index from ffsim.operators.fermion_action import cre_a, cre_b, des_a, des_b from ffsim.states import dim @@ -59,8 +59,8 @@ def norb(self) -> int: def _linear_operator_(self, norb: int, nelec: tuple[int, int]) -> LinearOperator: """Return a SciPy LinearOperator representing the object.""" n_alpha, n_beta = nelec - linkstr_index_a = cistring.gen_linkstr_index(range(norb), n_alpha) - linkstr_index_b = cistring.gen_linkstr_index(range(norb), n_beta) + linkstr_index_a = gen_linkstr_index(range(norb), n_alpha) + linkstr_index_b = gen_linkstr_index(range(norb), n_beta) link_index = (linkstr_index_a, linkstr_index_b) two_body = absorb_h1e( self.one_body_tensor, self.two_body_tensor, norb, nelec, 0.5 diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py index 3aaa22838..7ecc72fbf 100644 --- a/python/ffsim/states/states.py +++ b/python/ffsim/states/states.py @@ -27,6 +27,7 @@ ) from scipy.special import comb +from ffsim.cistring import gen_linkstr_index from ffsim.gates.orbital_rotation import apply_orbital_rotation @@ -207,7 +208,6 @@ def rdm( rank: int = 1, spin_summed: bool = True, reordered: bool = True, - link_index: tuple[np.ndarray, np.ndarray] | None = None, ) -> np.ndarray: """Return the reduced density matrix (RDM) of a state vector. @@ -250,14 +250,14 @@ def rdm( rank: The rank of the reduced density matrix. spin_summed: Whether to sum over the spin index. reordered: Whether to reorder the indices of the reduced density matrix. - link_index: String index lookup tables for spin alpha and spin beta, as - generated by the function `pyscf.fci.cistring.gen_linkstr_index`_. Returns: The reduced density matrix. - - .. _pyscf.fci.cistring.gen_linkstr_index: https://pyscf.org/pyscf_api_docs/pyscf.fci.html#pyscf.fci.cistring.gen_linkstr_index """ + n_alpha, n_beta = nelec + link_index_a = gen_linkstr_index(range(norb), n_alpha) + link_index_b = gen_linkstr_index(range(norb), n_beta) + link_index = (link_index_a, link_index_b) vec_real = np.real(vec) vec_imag = np.imag(vec) if rank == 1: diff --git a/python/ffsim/trotter/qdrift.py b/python/ffsim/trotter/qdrift.py index b2c245610..aba8ac1e6 100644 --- a/python/ffsim/trotter/qdrift.py +++ b/python/ffsim/trotter/qdrift.py @@ -90,7 +90,6 @@ def simulate_qdrift_double_factorized( hamiltonian.one_body_tensor ) step_time = time / n_steps - n_alpha, n_beta = nelec results = np.empty((n_samples, initial_state.shape[0]), dtype=complex) for i in range(n_samples): diff --git a/python/ffsim/trotter/trotter.py b/python/ffsim/trotter/trotter.py index 6594d6729..8092f1f9c 100644 --- a/python/ffsim/trotter/trotter.py +++ b/python/ffsim/trotter/trotter.py @@ -96,7 +96,6 @@ def simulate_trotter_double_factorized( hamiltonian.one_body_tensor ) step_time = time / n_steps - n_alpha, n_beta = nelec for _ in range(n_steps): vec = _simulate_trotter_step_double_factorized(