Merge branch 'direct_Slater_sampler' of github.com:jrm874/ffsim_slate…

…r_sampler into direct_Slater_sampler

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "maturin"
 [project]
 name = "ffsim"
 requires-python = ">=3.9"
-version = "0.0.42.dev"
+version = "0.0.44.dev"
 description = "Faster simulations of fermionic quantum circuits."
 readme = "README.md"
 license = { file = "LICENSE" }

python/ffsim/hamiltonians/diagonal_coulomb_hamiltonian.py

@@ -15,6 +15,7 @@
 
 import numpy as np
 import scipy.linalg
+from pyscf.fci.direct_uhf import make_hdiag
 from scipy.sparse.linalg import LinearOperator
 
 from ffsim._lib import FermionOperator
@@ -174,3 +175,23 @@ def from_fermion_operator(op: FermionOperator) -> DiagonalCoulombHamiltonian:
             diag_coulomb_mats=diag_coulomb_mats,
             constant=constant,
         )
+
+    def _diag_(self, norb: int, nelec: tuple[int, int]) -> np.ndarray:
+        """Return the diagonal entries of the Hamiltonian."""
+        if np.iscomplexobj(self.one_body_tensor):
+            raise NotImplementedError(
+                "Computing diagonal of complex diagonal Coulomb Hamiltonian is not yet "
+                "supported."
+            )
+        two_body_tensor_aa = np.zeros((self.norb, self.norb, self.norb, self.norb))
+        two_body_tensor_ab = np.zeros((self.norb, self.norb, self.norb, self.norb))
+        diag_coulomb_mat_aa, diag_coulomb_mat_ab = self.diag_coulomb_mats
+        for p, q in itertools.product(range(self.norb), repeat=2):
+            two_body_tensor_aa[p, p, q, q] = diag_coulomb_mat_aa[p, q]
+            two_body_tensor_ab[p, p, q, q] = diag_coulomb_mat_ab[p, q]
+        one_body_tensor = self.one_body_tensor + 0.5 * np.einsum(
+            "prqr", two_body_tensor_aa
+        )
+        h1e = (one_body_tensor, one_body_tensor)
+        h2e = (two_body_tensor_aa, two_body_tensor_ab, two_body_tensor_aa)
+        return make_hdiag(h1e, h2e, norb=norb, nelec=nelec) + self.constant

python/ffsim/hamiltonians/molecular_hamiltonian.py

@@ -158,8 +158,9 @@ def _diag_(self, norb: int, nelec: tuple[int, int]) -> np.ndarray:
         if np.iscomplexobj(self.two_body_tensor) or np.iscomplexobj(
             self.one_body_tensor
         ):
-            raise ValueError(
-                "Computing diagonal of complex molecular Hamiltonian is not supported."
+            raise NotImplementedError(
+                "Computing diagonal of complex molecular Hamiltonian is not yet "
+                "supported."
             )
         return (
             make_hdiag(self.one_body_tensor, self.two_body_tensor, norb, nelec)

python/ffsim/molecular_data.py

@@ -26,6 +26,7 @@
 import pyscf
 import pyscf.cc
 import pyscf.ci
+import pyscf.fci
 import pyscf.mcscf
 import pyscf.mp
 import pyscf.symm
@@ -68,6 +69,9 @@ class MolecularData:
             CCSD t2 amplitudes.
         cisd_energy (float | None): The CISD energy.
         cisd_vec (np.ndarray | None): The CISD state vector.
+        sci_energy (float | None): The SCI energy.
+        sci_vec (tuple[np.ndarray, np.ndarray, np.ndarray] | None): The SCI state
+            vector coefficients, spin alpha strings, and spin beta strings.
         fci_energy (float | None): The FCI energy.
         fci_vec (np.ndarray | None): The FCI state vector.
         dipole_integrals (np.ndarray | None): The dipole integrals.
@@ -104,6 +108,9 @@ class MolecularData:
     # CISD data
     cisd_energy: float | None = None
     cisd_vec: np.ndarray | None = None
+    # SCI data
+    sci_energy: float | None = None
+    sci_vec: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None
     # FCI data
     fci_energy: float | None = None
     fci_vec: np.ndarray | None = None
@@ -220,6 +227,19 @@ def run_cisd(self, *, store_cisd_vec: bool = False) -> None:
         if store_cisd_vec:
             self.cisd_vec = cisd_vec
 
+    def run_sci(self, *, store_sci_vec: bool = False) -> None:
+        """Run SCI and store results."""
+        sci = pyscf.fci.SCI(self.scf.run())
+        sci_energy, sci_vec = sci.kernel(
+            self.one_body_integrals,
+            self.two_body_integrals,
+            norb=self.norb,
+            nelec=self.nelec,
+        )
+        self.sci_energy = sci_energy + self.core_energy
+        if store_sci_vec:
+            self.sci_vec = (sci_vec, *sci_vec._strs)
+
     def run_fci(self, *, store_fci_vec: bool = False) -> None:
         """Run FCI and store results."""
         cas = pyscf.mcscf.CASCI(self.scf.run(), ncas=self.norb, nelecas=self.nelec)
@@ -343,6 +363,8 @@ def as_array_tuple_or_none(val):
             ccsd_t2=arrays_func(data.get("ccsd_t2")),
             cisd_energy=data.get("cisd_energy"),
             cisd_vec=as_array_or_none(data.get("cisd_vec")),
+            sci_energy=data.get("sci_energy"),
+            sci_vec=as_array_tuple_or_none(data.get("sci_vec")),
             fci_energy=data.get("fci_energy"),
             fci_vec=as_array_or_none(data.get("fci_vec")),
             dipole_integrals=as_array_or_none(data.get("dipole_integrals")),

python/ffsim/states/slater.py

@@ -162,7 +162,6 @@ def _generate_marginals(
     """
     marginals = np.zeros(len(empty_orbitals), dtype=float)
     for i, orbital in enumerate(empty_orbitals):
-        new_sample = sample.copy()
         new_sample = sample + [orbital]
         rest_rdm = rdm[np.ix_(new_sample, new_sample)]
         marginals[i] = np.linalg.det(rest_rdm).real
@@ -189,14 +188,13 @@ def _autoregressive_slater(
         A Numpy array with the position of the sampled electrons.
     """
     rng = np.random.default_rng(seed)
-
     probs = np.diag(rdm).real / nelec
     sample = [rng.choice(norb, p=probs)]
     marginal = [probs[sample[0]]]
     all_orbs = set(range(norb))
     empty_orbitals = list(all_orbs.difference(sample))
     for k in range(nelec - 1):
-        marginals = _generate_marginals(rdm, sample, empty_orbitals)
+        marginals = _generate_marginals(rdm, sample, empty_orbitals, marginal[-1])
         conditionals = marginals / marginal[-1]
         conditionals /= np.sum(conditionals)
         index = rng.choice(len(empty_orbitals), p=conditionals)
@@ -233,12 +231,10 @@ def _sample_spinless_direct(
         Each row is a sample.
     """
     norb, _ = rdm.shape
-
     if nelec == 0:
         return [0] * shots
     if nelec == norb:
         return [(1 << norb) - 1] * shots
     rng = np.random.default_rng(seed)
-    samples = []
     samples = [_autoregressive_slater(rdm, norb, nelec, rng) for _ in range(shots)]
     return [sum(1 << orb for orb in sample) for sample in samples]

python/ffsim/variational/ucj_spin_balanced.py

@@ -410,7 +410,12 @@ def from_t_amplitudes(
         Args:
             t2: The t2 amplitudes.
             t1: The t1 amplitudes.
-            n_reps: The number of ansatz repetitions.
+            n_reps: The number of ansatz repetitions. If not specified, then it is set
+                to the number of terms resulting from the double-factorization of the
+                t2 amplitudes. If the specified number of repetitions is larger than the
+                number of terms resulting from the double-factorization, then the ansatz
+                is padded with additional identity operators up to the specified number
+                of repetitions.
             interaction_pairs: Optional restrictions on allowed orbital interactions
                 for the diagonal Coulomb operators.
                 If specified, `interaction_pairs` should be a pair of lists,
@@ -450,6 +455,17 @@ def from_t_amplitudes(
         diag_coulomb_mats = np.stack([diag_coulomb_mats, diag_coulomb_mats], axis=1)
         orbital_rotations = orbital_rotations.reshape(-1, norb, norb)[:n_reps]
 
+        n_vecs, _, _, _ = diag_coulomb_mats.shape
+        if n_reps is not None and n_vecs < n_reps:
+            # Pad with no-ops to the requested number of repetitions
+            diag_coulomb_mats = np.concatenate(
+                [diag_coulomb_mats, np.zeros((n_reps - n_vecs, 2, norb, norb))]
+            )
+            eye = np.eye(norb)
+            orbital_rotations = np.concatenate(
+                [orbital_rotations, np.stack([eye for _ in range(n_reps - n_vecs)])]
+            )
+
         final_orbital_rotation = None
         if t1 is not None:
             final_orbital_rotation_generator = np.zeros((norb, norb), dtype=complex)

python/ffsim/variational/ucj_spin_unbalanced.py

@@ -468,6 +468,12 @@ def from_t_amplitudes(
                 the number of terms to use from the alpha-beta t2 amplitudes,
                 and the second integer specifies the number of terms to use from the
                 alpha-alpha and beta-beta t2 amplitudes.
+                If not specified, then it is set
+                to the number of terms resulting from the double-factorization of the
+                t2 amplitudes. If the specified number of repetitions is larger than the
+                number of terms resulting from the double-factorization, then the ansatz
+                is padded with additional identity operators up to the specified number
+                of 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,
@@ -524,23 +530,27 @@ def from_t_amplitudes(
         diag_coulomb_mats_bb = diag_coulomb_mats_bb.reshape(-1, norb, norb)
         orbital_rotations_bb = orbital_rotations_bb.reshape(-1, norb, norb)
         zero = np.zeros((norb, norb))
-        diag_coulomb_mats_same_spin = np.stack(
-            [
-                [mat_aa, zero, mat_bb]
-                for mat_aa, mat_bb in itertools.zip_longest(
-                    diag_coulomb_mats_aa, diag_coulomb_mats_bb, fillvalue=zero
-                )
-            ]
-        )
-        eye = np.eye(norb)
-        orbital_rotations_same_spin = np.stack(
-            [
-                [orbital_rotation_aa, orbital_rotation_bb]
-                for orbital_rotation_aa, orbital_rotation_bb in itertools.zip_longest(
-                    orbital_rotations_aa, orbital_rotations_bb, fillvalue=eye
-                )
-            ]
-        )
+        if len(diag_coulomb_mats_aa) or len(diag_coulomb_mats_bb):
+            diag_coulomb_mats_same_spin = np.stack(
+                [
+                    [mat_aa, zero, mat_bb]
+                    for mat_aa, mat_bb in itertools.zip_longest(
+                        diag_coulomb_mats_aa, diag_coulomb_mats_bb, fillvalue=zero
+                    )
+                ]
+            )
+            eye = np.eye(norb)
+            orbital_rotations_same_spin = np.stack(
+                [
+                    [orb_rot_aa, orb_rot_bb]
+                    for orb_rot_aa, orb_rot_bb in itertools.zip_longest(
+                        orbital_rotations_aa, orbital_rotations_bb, fillvalue=eye
+                    )
+                ]
+            )
+        else:
+            diag_coulomb_mats_same_spin = np.empty((0, 3, norb, norb))
+            orbital_rotations_same_spin = np.empty((0, 2, norb, norb))
         # concatenate
         if n_reps is None or isinstance(n_reps, int):
             diag_coulomb_mats = np.concatenate(
@@ -549,19 +559,27 @@ def from_t_amplitudes(
             orbital_rotations = np.concatenate(
                 [orbital_rotations_ab, orbital_rotations_same_spin]
             )[:n_reps]
+            diag_coulomb_mats = _pad_diag_coulomb_mats(diag_coulomb_mats, n_reps)
+            orbital_rotations = _pad_orbital_rotations(orbital_rotations, n_reps)
         else:
             n_reps_ab, n_reps_same_spin = n_reps
+            diag_coulomb_mats_ab = _pad_diag_coulomb_mats(
+                diag_coulomb_mats_ab[:n_reps_ab], n_reps_ab
+            )
+            orbital_rotations_ab = _pad_orbital_rotations(
+                orbital_rotations_ab[:n_reps_ab], n_reps_ab
+            )
+            diag_coulomb_mats_same_spin = _pad_diag_coulomb_mats(
+                diag_coulomb_mats_same_spin[:n_reps_same_spin], n_reps_same_spin
+            )
+            orbital_rotations_same_spin = _pad_orbital_rotations(
+                orbital_rotations_same_spin[:n_reps_same_spin], n_reps_same_spin
+            )
             diag_coulomb_mats = np.concatenate(
-                [
-                    diag_coulomb_mats_ab[:n_reps_ab],
-                    diag_coulomb_mats_same_spin[:n_reps_same_spin],
-                ]
+                [diag_coulomb_mats_ab, diag_coulomb_mats_same_spin]
             )
             orbital_rotations = np.concatenate(
-                [
-                    orbital_rotations_ab[:n_reps_ab],
-                    orbital_rotations_same_spin[:n_reps_same_spin],
-                ]
+                [orbital_rotations_ab, orbital_rotations_same_spin]
             )
 
         final_orbital_rotation = None
@@ -676,3 +694,28 @@ def _approx_eq_(self, other, rtol: float, atol: float) -> bool:
                 )
             return True
         return NotImplemented
+
+
+def _pad_diag_coulomb_mats(diag_coulomb_mats: np.ndarray, n_reps: int | None):
+    n_vecs, _, norb, _ = diag_coulomb_mats.shape
+    if n_reps is not None and n_vecs < n_reps:
+        diag_coulomb_mats = np.concatenate(
+            [
+                diag_coulomb_mats,
+                np.zeros((n_reps - n_vecs, 3, norb, norb)),
+            ]
+        )
+    return diag_coulomb_mats
+
+
+def _pad_orbital_rotations(orbital_rotations: np.ndarray, n_reps: int | None):
+    n_vecs, _, norb, _ = orbital_rotations.shape
+    if n_reps is not None and n_vecs < n_reps:
+        eye = np.eye(norb)
+        orbital_rotations = np.concatenate(
+            [
+                orbital_rotations,
+                np.stack([np.stack([eye, eye]) for _ in range(n_reps - n_vecs)]),
+            ]
+        )
+    return orbital_rotations

python/ffsim/variational/ucj_spinless.py

@@ -353,7 +353,12 @@ def from_t_amplitudes(
         Args:
             t2: The t2 amplitudes.
             t1: The t1 amplitudes.
-            n_reps: The number of ansatz repetitions.
+            n_reps: The number of ansatz repetitions. If not specified, then it is set
+                to the number of terms resulting from the double-factorization of the
+                t2 amplitudes. If the specified number of repetitions is larger than the
+                number of terms resulting from the double-factorization, then the ansatz
+                is padded with additional identity operators up to the specified number
+                of repetitions.
             interaction_pairs: Optional restrictions on allowed orbital interactions
                 for the diagonal Coulomb operators.
                 If specified, `interaction_pairs` should be a list of integer pairs
@@ -384,6 +389,17 @@ def from_t_amplitudes(
         diag_coulomb_mats = diag_coulomb_mats.reshape(-1, norb, norb)[:n_reps]
         orbital_rotations = orbital_rotations.reshape(-1, norb, norb)[:n_reps]
 
+        n_vecs, _, _ = diag_coulomb_mats.shape
+        if n_reps is not None and n_vecs < n_reps:
+            # Pad with no-ops to the requested number of repetitions
+            diag_coulomb_mats = np.concatenate(
+                [diag_coulomb_mats, np.zeros((n_reps - n_vecs, norb, norb))]
+            )
+            eye = np.eye(norb)
+            orbital_rotations = np.concatenate(
+                [orbital_rotations, np.stack([eye for _ in range(n_reps - n_vecs)])]
+            )
+
         final_orbital_rotation = None
         if t1 is not None:
             final_orbital_rotation_generator = np.zeros((norb, norb), dtype=complex)

tests/python/hamiltonians/diagonal_coulomb_hamiltonian_test.py

@@ -273,3 +273,24 @@ def test_from_fermion_operator_fermi_hubbard_2d(norb: int, nelec: tuple[int, int
     actual_periodic = actual_linop_periodic @ vec
     expected_periodic = expected_linop_periodic @ vec
     np.testing.assert_allclose(actual_periodic, expected_periodic)
+
+
+def test_diag():
+    """Test computing diagonal."""
+    rng = np.random.default_rng(2222)
+    norb = 5
+    nelec = (3, 2)
+    # TODO test complex one-body after adding support for it
+    one_body_tensor = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+    diag_coulomb_mat_a = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+    diag_coulomb_mat_b = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+    diag_coulomb_mats = np.stack([diag_coulomb_mat_a, diag_coulomb_mat_b])
+    constant = rng.standard_normal()
+    hamiltonian = ffsim.DiagonalCoulombHamiltonian(
+        one_body_tensor, diag_coulomb_mats, constant=constant
+    )
+    linop = ffsim.linear_operator(hamiltonian, norb=norb, nelec=nelec)
+    hamiltonian_dense = linop @ np.eye(ffsim.dim(norb, nelec))
+    np.testing.assert_allclose(
+        ffsim.diag(hamiltonian, norb=norb, nelec=nelec), np.diag(hamiltonian_dense)
+    )

tests/python/molecular_data_test.py

@@ -31,6 +31,7 @@ def _assert_mol_data_equal(
             "np.ndarray | None",
             "np.ndarray | tuple[np.ndarray, np.ndarray] | None",
             "np.ndarray | tuple[np.ndarray, np.ndarray, np.ndarray] | None",
+            "tuple[np.ndarray, np.ndarray, np.ndarray] | None",
         ]:
             if actual is not None:
                 if isinstance(actual, tuple):
@@ -85,11 +86,13 @@ def test_molecular_data_run_methods():
     mol_data.run_mp2()
     mol_data.run_ccsd()
     mol_data.run_cisd()
+    mol_data.run_sci()
     mol_data.run_fci()
 
     np.testing.assert_allclose(mol_data.mp2_energy, -108.58852784026)
     np.testing.assert_allclose(mol_data.ccsd_energy, -108.5933309085008)
     np.testing.assert_allclose(mol_data.cisd_energy, -108.5878344909782)
+    np.testing.assert_allclose(mol_data.sci_energy, -108.59598682615388)
     np.testing.assert_allclose(mol_data.fci_energy, -108.595987350986)
 
 
@@ -108,6 +111,7 @@ def test_json_closed_shell(tmp_path: pathlib.Path):
     mol_data.run_mp2(store_t2=True)
     mol_data.run_ccsd(store_t1=True, store_t2=True)
     mol_data.run_cisd(store_cisd_vec=True)
+    mol_data.run_sci(store_sci_vec=True)
     mol_data.run_fci(store_fci_vec=True)
 
     for compression in [None, "gzip", "bz2", "lzma"]:
@@ -132,6 +136,7 @@ def test_json_open_shell(tmp_path: pathlib.Path):
     mol_data.run_mp2(store_t2=True)
     mol_data.run_ccsd(store_t1=True, store_t2=True)
     mol_data.run_cisd(store_cisd_vec=True)
+    mol_data.run_sci(store_sci_vec=True)
     mol_data.run_fci(store_fci_vec=True)
 
     for compression in [None, "gzip", "bz2", "lzma"]: