Code for jellium Hartree-Fock state in plane wave and dual basis (#120)

* Code for jellium Hartree-Fock state in plane wave and dual basis

* Trying to figure out where Python3 is making a float

Author: idk3

src/fermilib/utils/_jellium_hf_state.py

@@ -0,0 +1,93 @@
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+"""This module constructs the uniform electron gas' Hartree-Fock state."""
+from __future__ import absolute_import
+
+from fermilib.ops import FermionOperator, normal_ordered
+from fermilib.utils import jellium_model, inverse_fourier_transform
+
+from scipy.sparse import csr_matrix
+
+import numpy
+
+
+def hartree_fock_state_jellium(grid, n_electrons, spinless=True,
+                               plane_wave=False):
+    """Give the Hartree-Fock state of jellium.
+
+    Args:
+        grid (Grid): The discretization to use.
+        n_electrons (int): Number of electrons in the system.
+        spinless (bool): Whether to use the spinless model or not.
+        plane_wave (bool): Whether to return the Hartree-Fock state in
+                           the plane wave (True) or dual basis (False).
+
+    Notes:
+        The jellium model is built up by filling the lowest-energy
+        single-particle states in the plane-wave Hamiltonian until
+        n_electrons states are filled.
+    """
+    # Get the jellium Hamiltonian in the plane wave basis.
+    hamiltonian = jellium_model(grid, spinless, plane_wave=True)
+    hamiltonian = normal_ordered(hamiltonian)
+    hamiltonian.compress()
+
+    # Enumerate the single-particle states.
+    n_single_particle_states = (grid.length ** grid.dimensions)
+    if not spinless:
+        n_single_particle_states *= 2
+
+    # Compute the energies for each of the single-particle states.
+    single_particle_energies = numpy.zeros(n_single_particle_states,
+                                           dtype=float)
+    for i in range(n_single_particle_states):
+        single_particle_energies[i] = hamiltonian.terms.get(
+            ((i, 1), (i, 0)), 0.0)
+
+    # The number of occupied single-particle states is the number of electrons.
+    # Those states with the lowest single-particle energies are occupied first.
+    occupied_states = single_particle_energies.argsort()[:n_electrons]
+
+    if plane_wave:
+        # In the plane wave basis the HF state is a single determinant.
+        hartree_fock_state_index = numpy.sum(2 ** occupied_states)
+        hartree_fock_state = csr_matrix(
+            ([1.0], ([hartree_fock_state_index], [0])),
+            shape=(2 ** n_single_particle_states, 1))
+
+    else:
+        # Inverse Fourier transform the creation operators for the state to get
+        # to the dual basis state, then use that to get the dual basis state.
+        hartree_fock_state_creation_operator = FermionOperator.identity()
+        for state in occupied_states[::-1]:
+            hartree_fock_state_creation_operator *= (
+                FermionOperator(((int(state), 1),)))
+        dual_basis_hf_creation_operator = inverse_fourier_transform(
+            hartree_fock_state_creation_operator, grid, spinless)
+
+        dual_basis_hf_creation = normal_ordered(
+            dual_basis_hf_creation_operator)
+
+        # Initialize the HF state as a sparse matrix.
+        hartree_fock_state = csr_matrix(
+            ([], ([], [])), shape=(2 ** n_single_particle_states, 1),
+            dtype=complex)
+
+        # Populate the elements of the HF state in the dual basis.
+        for term in dual_basis_hf_creation.terms:
+            index = 0
+            for operator in term:
+                index += 2 ** operator[0]
+            hartree_fock_state[index, 0] = dual_basis_hf_creation.terms[term]
+
+    return hartree_fock_state

src/fermilib/utils/_jellium_hf_state_test.py

@@ -0,0 +1,99 @@
+"""Tests for _jellium_hf_state.py."""
+from __future__ import absolute_import
+
+from itertools import permutations
+from scipy.sparse import csr_matrix
+
+import numpy
+import unittest
+
+from fermilib.transforms import get_sparse_operator
+from fermilib.utils import expectation, get_ground_state, Grid, jellium_model
+from fermilib.utils._plane_wave_hamiltonian import wigner_seitz_length_scale
+from fermilib.utils._sparse_tools import jw_number_restrict_operator
+
+from fermilib.utils._jellium_hf_state import hartree_fock_state_jellium
+
+
+class JelliumHartreeFockStateTest(unittest.TestCase):
+    def test_hf_state_energy_close_to_ground_energy_at_high_density(self):
+        grid_length = 8
+        dimension = 1
+        spinless = True
+        n_particles = grid_length ** dimension // 2
+
+        # High density -> small length_scale.
+        length_scale = 0.25
+
+        grid = Grid(dimension, grid_length, length_scale)
+        hamiltonian = jellium_model(grid, spinless)
+        hamiltonian_sparse = get_sparse_operator(hamiltonian)
+
+        hf_state = hartree_fock_state_jellium(grid, n_particles,
+                                              spinless, plane_wave=True)
+
+        restricted_hamiltonian = jw_number_restrict_operator(
+            hamiltonian_sparse, n_particles)
+
+        E_g = get_ground_state(restricted_hamiltonian)[0]
+        E_HF_plane_wave = expectation(hamiltonian_sparse, hf_state)
+
+        self.assertAlmostEqual(E_g, E_HF_plane_wave, places=5)
+
+    def test_hf_state_energy_same_in_plane_wave_and_dual_basis(self):
+        grid_length = 4
+        dimension = 1
+        wigner_seitz_radius = 10.0
+        spinless = False
+
+        n_orbitals = grid_length ** dimension
+        if not spinless:
+            n_orbitals *= 2
+        n_particles = n_orbitals // 2
+
+        length_scale = wigner_seitz_length_scale(
+            wigner_seitz_radius, n_particles, dimension)
+
+        grid = Grid(dimension, grid_length, length_scale)
+        hamiltonian = jellium_model(grid, spinless)
+        hamiltonian_dual_basis = jellium_model(
+            grid, spinless, plane_wave=False)
+
+        # Get the Hamiltonians as sparse operators.
+        hamiltonian_sparse = get_sparse_operator(hamiltonian)
+        hamiltonian_dual_sparse = get_sparse_operator(hamiltonian_dual_basis)
+
+        hf_state = hartree_fock_state_jellium(
+            grid, n_particles, spinless, plane_wave=True)
+        hf_state_dual = hartree_fock_state_jellium(
+            grid, n_particles, spinless, plane_wave=False)
+
+        E_HF_plane_wave = expectation(hamiltonian_sparse, hf_state)
+        E_HF_dual = expectation(hamiltonian_dual_sparse, hf_state_dual)
+
+        self.assertAlmostEqual(E_HF_dual, E_HF_plane_wave)
+
+    def test_hf_state_plane_wave_basis_lowest_single_determinant_state(self):
+        grid_length = 7
+        dimension = 1
+        spinless = True
+        n_particles = 4
+        length_scale = 2.0
+
+        grid = Grid(dimension, grid_length, length_scale)
+        hamiltonian = jellium_model(grid, spinless)
+        hamiltonian_sparse = get_sparse_operator(hamiltonian)
+
+        hf_state = hartree_fock_state_jellium(grid, n_particles,
+                                              spinless, plane_wave=True)
+
+        HF_energy = expectation(hamiltonian_sparse, hf_state)
+
+        for occupied_orbitals in permutations(
+                [1] * n_particles + [0] * (grid_length - n_particles)):
+            state_index = numpy.sum(2 ** numpy.array(occupied_orbitals))
+            HF_competitor = csr_matrix(([1.0], ([state_index], [0])),
+                                       shape=(2 ** grid_length, 1))
+
+            self.assertLessEqual(
+                HF_energy, expectation(hamiltonian_sparse, HF_competitor))