add hartree_fock_state function (#60)

docs/tutorials/03-double-factorized.ipynb

@@ -238,10 +238,7 @@
    "outputs": [],
    "source": [
     "# Construct the Hartree-Fock state\n",
-    "n_alpha, n_beta = nelec\n",
-    "initial_state = ffsim.slater_determinant(\n",
-    "    norb=norb, occupied_orbitals=(range(n_alpha), range(n_beta))\n",
-    ")\n",
+    "initial_state = ffsim.hartree_fock_state(norb, nelec)\n",
     "\n",
     "# Get the Hamiltonian as a LinearOperator\n",
     "hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)\n",
@@ -403,7 +400,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.10.12"
   },
   "orig_nbformat": 4
  },

docs/tutorials/04-lucj.ipynb

@@ -104,10 +104,7 @@
     "operator = ffsim.UCJOperator.from_t_amplitudes(t2, n_reps=n_reps)\n",
     "\n",
     "# Construct the Hartree-Fock state to use as the reference state\n",
-    "n_alpha, n_beta = nelec\n",
-    "reference_state = ffsim.slater_determinant(\n",
-    "    norb=norb, occupied_orbitals=(range(n_alpha), range(n_beta))\n",
-    ")\n",
+    "reference_state = ffsim.hartree_fock_state(norb, nelec)\n",
     "\n",
     "# Apply the operator to the reference state\n",
     "ansatz_state = ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)\n",
@@ -124,7 +121,7 @@
    "source": [
     "To facilitate variational optimization of the ansatz, `UCJOperator` implements methods for conversion to and from a vector of real-valued parameters. The precise relation between a parameter vector and the matrices of the UCJ operator is somewhat complicated. In short, the parameter vector stores the entries of the UCJ matrices in a non-redundant way (for the orbital rotations, the parameter vector actually stores the entries of their generators.)\n",
     "\n",
-    "The following code cell shows how one can define an objective function that takes as input a parameter vector and outputs the energy of the associated ansatz state, and then optimize this objective function using `scipy.optimize.minimize`."
+    "The following code cell shows how one can define an objective function that takes as input a parameter vector and outputs the energy of the associated ansatz state, and then optimize this objective function using `scipy.optimize.minimize`. Here, we set a small limit on the number of function evaluations; increase the value if you would like to run it to convergence."
    ]
   },
   {
@@ -140,15 +137,13 @@
     "    # Initialize the ansatz operator from the parameter vector\n",
     "    operator = ffsim.UCJOperator.from_parameters(x, norb=norb, n_reps=n_reps)\n",
     "    # Apply the ansatz operator to the reference state\n",
-    "    final_state = ffsim.apply_unitary(\n",
-    "        reference_state, operator, norb=norb, nelec=(n_alpha, n_beta)\n",
-    "    )\n",
+    "    final_state = ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)\n",
     "    # Return the energy ⟨ψ|H|ψ⟩ of the ansatz state\n",
     "    return np.real(np.vdot(final_state, hamiltonian @ final_state))\n",
     "\n",
     "\n",
     "result = scipy.optimize.minimize(\n",
-    "    fun, x0=operator.to_parameters(), method=\"L-BFGS-B\", options=dict(maxfun=50000)\n",
+    "    fun, x0=operator.to_parameters(), method=\"L-BFGS-B\", options=dict(maxfun=1000)\n",
     ")\n",
     "\n",
     "print(f\"Number of parameters: {len(result.x)}\")\n",
@@ -195,9 +190,7 @@
     "        alpha_alpha_indices=alpha_alpha_indices,\n",
     "        alpha_beta_indices=alpha_beta_indices,\n",
     "    )\n",
-    "    final_state = ffsim.apply_unitary(\n",
-    "        reference_state, operator, norb=norb, nelec=(n_alpha, n_beta)\n",
-    "    )\n",
+    "    final_state = ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)\n",
     "    return np.real(np.vdot(final_state, hamiltonian @ final_state))\n",
     "\n",
     "\n",
@@ -207,7 +200,7 @@
     "        alpha_alpha_indices=alpha_alpha_indices, alpha_beta_indices=alpha_beta_indices\n",
     "    ),\n",
     "    method=\"L-BFGS-B\",\n",
-    "    options=dict(maxfun=50000),\n",
+    "    options=dict(maxfun=1000),\n",
     ")\n",
     "print(f\"Number of parameters: {len(result.x)}\")\n",
     "print(result)"
@@ -230,7 +223,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,

python/ffsim/__init__.py

@@ -50,6 +50,7 @@
     dims,
     expectation_one_body_power,
     expectation_one_body_product,
+    hartree_fock_state,
     indices_to_strings,
     one_hot,
     slater_determinant,
@@ -96,6 +97,7 @@
     "dims",
     "expectation_one_body_power",
     "expectation_one_body_product",
+    "hartree_fock_state",
     "indices_to_strings",
     "linalg",
     "linear_operator",

python/ffsim/states/__init__.py

@@ -13,6 +13,7 @@
 from ffsim.states.states import (
     dim,
     dims,
+    hartree_fock_state,
     indices_to_strings,
     one_hot,
     slater_determinant,
@@ -25,6 +26,7 @@
     "dims",
     "expectation_one_body_power",
     "expectation_one_body_product",
+    "hartree_fock_state",
     "indices_to_strings",
     "one_hot",
     "slater_determinant",

python/ffsim/states/states.py

@@ -80,10 +80,9 @@ def slater_determinant(
         orbital_rotation: An optional orbital rotation to apply to the
             electron configuration. In other words, this is a unitary matrix that
             describes the orbitals of the Slater determinant.
-        dtype:
 
     Returns:
-        The Slater determinant.
+        The Slater determinant as a statevector.
     """
     alpha_orbitals, beta_orbitals = occupied_orbitals
     n_alpha = len(alpha_orbitals)
@@ -104,6 +103,22 @@ def slater_determinant(
     return vec
 
 
+def hartree_fock_state(norb: int, nelec: tuple[int, int]) -> np.ndarray:
+    """Return the Hartree-Fock state.
+
+    Args:
+        norb: The number of spatial orbitals.
+        nelec: The number of alpha and beta electrons.
+
+    Returns:
+        The Hartree-Fock state as a statevector.
+    """
+    n_alpha, n_beta = nelec
+    return slater_determinant(
+        norb=norb, occupied_orbitals=(range(n_alpha), range(n_beta))
+    )
+
+
 def slater_determinant_one_rdm(
     norb: int, occupied_orbitals: tuple[Sequence[int], Sequence[int]]
 ) -> np.ndarray:

tests/states_test.py

@@ -11,6 +11,7 @@
 """Test states."""
 
 import numpy as np
+import pyscf
 
 import ffsim
 
@@ -33,6 +34,27 @@ def test_slater_determinant():
     np.testing.assert_allclose(hamiltonian @ state, eig * state)
 
 
+def test_hartree_fock_state():
+    """Test Hartree-Fock state."""
+    mol = pyscf.gto.Mole()
+    mol.build(
+        atom=[["H", (0, 0, 0)], ["H", (0, 0, 1.8)]],
+        basis="sto-6g",
+    )
+    hartree_fock = pyscf.scf.RHF(mol)
+    hartree_fock_energy = hartree_fock.kernel()
+
+    mol_data = ffsim.MolecularData.from_hartree_fock(hartree_fock)
+
+    vec = ffsim.hartree_fock_state(mol_data.norb, mol_data.nelec)
+    hamiltonian = ffsim.linear_operator(
+        mol_data.hamiltonian, norb=mol_data.norb, nelec=mol_data.nelec
+    )
+    energy = np.vdot(vec, hamiltonian @ vec)
+
+    np.testing.assert_allclose(energy, hartree_fock_energy)
+
+
 def test_indices_to_strings():
     """Test converting statevector indices to strings."""
     norb = 3