used cached link index (#69)

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

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

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:

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):

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(