create cistring module with cached substitutes for pyscf functions (#68)

* create cistring module

* from __future__ import annotations

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

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",

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)

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)