Upgrade HF state for different mappings (#5472)

**Context:** Upgrade hf_state function to work with different mappings **Description of the Change:** hf_state function works with parity and Bravyi-Kitaev mappings. **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:** --------- Co-authored-by: soranjh <[email protected]> Co-authored-by: Austin Huang <[email protected]> Co-authored-by: Thomas R. Bromley <[email protected]> Co-authored-by: soranjh <[email protected]>
PennyLaneAI · Apr 18, 2024 · c2de427 · c2de427
1 parent e2529d8
commit c2de427
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 10 deletions.
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
@@ -177,7 +177,6 @@
   >>> op.error(method="commutator")
   SpectralNormError(6.166666666666668e-06)
   ```
-
 <h3>Improvements 🛠</h3>
 
 * `qml.ops.Conditional` now stores the `data`, `num_params`, and `ndim_param` attributes of
@@ -234,6 +233,9 @@
 * `qml.transforms.hamiltonian_expand` can now handle multi-term observables with a constant offset.
   [(#5414)](https://github.com/PennyLaneAI/pennylane/pull/5414)
 
+* The `qml.qchem.hf_state` function is upgraded to be compatible with the parity and Bravyi-Kitaev bases.
+  [(#5472)](https://github.com/PennyLaneAI/pennylane/pull/5472)
+
 <h4>Community contributions 🥳</h4>
 
 * Functions `measure_with_samples` and `sample_state` have been added to the new `qutrit_mixed` module found in

diff --git a/pennylane/qchem/structure.py b/pennylane/qchem/structure.py
@@ -282,8 +282,30 @@ def excitations(electrons, orbitals, delta_sz=0):
     return singles, doubles
 
 
-def hf_state(electrons, orbitals):
-    r"""Generate the occupation-number vector representing the Hartree-Fock state.
+def _beta_matrix(orbitals):
+    r"""Generate the conversion matrix to transform the occupation number basis to the Bravyi-Kitaev basis.
+
+    Args:
+        orbitals (int): Number of spin orbitals. If an active space is defined,
+            this is the number of active spin orbitals.
+
+    Returns:
+        (array): The transformation matrix
+    """
+
+    bin_range = int(np.ceil(np.log2(orbitals)))
+
+    beta = np.array([[1]])
+
+    for i in range(bin_range):
+        beta = np.kron(np.eye(2), beta)
+        beta[-1, : 2**i] = 1
+
+    return beta[:orbitals, :orbitals]
+
+
+def hf_state(electrons, orbitals, basis="occupation_number"):
+    r"""Generate the Hartree-Fock statevector with respect to a chosen basis.
 
     The many-particle wave function in the Hartree-Fock (HF) approximation is a `Slater determinant
     <https://en.wikipedia.org/wiki/Slater_determinant>`_. In Fock space, a Slater determinant
@@ -296,11 +318,18 @@ def hf_state(electrons, orbitals):
 
     where :math:`n_i` indicates the occupation of the :math:`i`-th orbital.
 
+    The Hartree-Fock state can also be generated in the parity basis, where each qubit stores the parity of
+    the spin orbital, and in the Bravyi-Kitaev basis, where a qubit :math:`j` stores the occupation state of orbital
+    :math:`j` if :math:`j` is even and stores partial sum of the occupation state of a set of orbitals of indices
+    less than :math:`j` if :math:`j` is odd [`Tranter et al. Int. J. Quantum Chem. 115, 1431 (2015)
+    <https://doi.org/10.1002/qua.24969>`_].
+
     Args:
         electrons (int): Number of electrons. If an active space is defined, this
             is the number of active electrons.
         orbitals (int): Number of *spin* orbitals. If an active space is defined,
             this is the number of active spin-orbitals.
+        basis (string): Basis in which the HF state is represented. Options are ``occupation_number``, ``parity`` and ``bravyi_kitaev``.
 
     Returns:
         array: NumPy array containing the vector :math:`\vert {\bf n} \rangle`
@@ -310,6 +339,15 @@ def hf_state(electrons, orbitals):
     >>> state = hf_state(2, 6)
     >>> print(state)
     [1 1 0 0 0 0]
+
+    >>> state = hf_state(2, 6, basis="parity")
+    >>> print(state)
+    [1 0 0 0 0 0]
+
+    >>> state = hf_state(2, 6, basis="bravyi_kitaev")
+    >>> print(state)
+    [1 0 0 0 0 0]
+
     """
 
     if electrons <= 0:
@@ -326,7 +364,15 @@ def hf_state(electrons, orbitals):
 
     state = np.where(np.arange(orbitals) < electrons, 1, 0)
 
-    return np.array(state)
+    if basis == "parity":
+        pi_matrix = np.tril(np.ones((orbitals, orbitals)))
+        return (np.matmul(pi_matrix, state) % 2).astype(int)
+
+    if basis == "bravyi_kitaev":
+        beta_matrix = _beta_matrix(orbitals)
+        return (np.matmul(beta_matrix, state) % 2).astype(int)
+
+    return state
 
 
 def excitations_to_wires(singles, doubles, wires=None):

diff --git a/tests/qchem/test_structure.py b/tests/qchem/test_structure.py
@@ -216,22 +216,72 @@ def circuit(weights):
 
 
 @pytest.mark.parametrize(
-    ("electrons", "orbitals", "exp_state"),
+    ("electrons", "orbitals", "basis", "exp_state"),
+    # Obtained manually using Eqs (10, 14) of
+    # [`Tranter et al. Int. J. Quantum Chem. 115, 1431 (2015)
+    # <https://doi.org/10.1002/qua.24969>`_]
     [
-        (2, 5, np.array([1, 1, 0, 0, 0])),
-        (1, 5, np.array([1, 0, 0, 0, 0])),
-        (5, 5, np.array([1, 1, 1, 1, 1])),
+        (1, 1, "occupation_number", np.array([1])),
+        (2, 5, "occupation_number", np.array([1, 1, 0, 0, 0])),
+        (1, 5, "occupation_number", np.array([1, 0, 0, 0, 0])),
+        (5, 5, "occupation_number", np.array([1, 1, 1, 1, 1])),
+        (1, 1, "parity", np.array([1])),
+        (2, 5, "parity", np.array([1, 0, 0, 0, 0])),
+        (1, 5, "parity", np.array([1, 1, 1, 1, 1])),
+        (5, 5, "parity", np.array([1, 0, 1, 0, 1])),
+        (1, 1, "bravyi_kitaev", np.array([1])),
+        (2, 5, "bravyi_kitaev", np.array([1, 0, 0, 0, 0])),
+        (1, 5, "bravyi_kitaev", np.array([1, 1, 0, 1, 0])),
+        (5, 5, "bravyi_kitaev", np.array([1, 0, 1, 0, 1])),
     ],
 )
-def test_hf_state(electrons, orbitals, exp_state):
+def test_hf_state(electrons, orbitals, basis, exp_state):
     r"""Test the correctness of the generated occupation-number vector"""
 
-    res_state = qchem.hf_state(electrons, orbitals)
+    res_state = qchem.hf_state(electrons, orbitals, basis)
 
     assert len(res_state) == len(exp_state)
     assert np.allclose(res_state, exp_state)
 
 
+@pytest.mark.parametrize(
+    ("electrons", "symbols", "geometry", "charge"),
+    [
+        (2, ["H", "H"], np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.4]], requires_grad=False), 0),
+        (
+            2,
+            ["H", "H", "H"],
+            np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0]], requires_grad=False),
+            1,
+        ),
+    ],
+)
+def test_hf_state_basis(electrons, symbols, geometry, charge):
+    r"""Test the correctness of the generated HF state in a circuit."""
+
+    mol = qml.qchem.Molecule(symbols, geometry, charge)
+    h_ferm = qchem.fermionic_hamiltonian(mol)()
+    qubits = len(h_ferm.wires)
+
+    state_occ = qchem.hf_state(electrons, qubits, basis="occupation_number")
+    state_parity = qchem.hf_state(electrons, qubits, basis="parity")
+    state_bk = qchem.hf_state(electrons, qubits, basis="bravyi_kitaev")
+
+    h_occ = qml.jordan_wigner(h_ferm, ps=True, tol=1e-16).hamiltonian()
+    h_parity = qml.parity_transform(h_ferm, qubits, ps=True, tol=1e-16).hamiltonian()
+    h_bk = qml.bravyi_kitaev(h_ferm, qubits, ps=True, tol=1e-16).hamiltonian()
+
+    dev = qml.device("default.qubit", wires=qubits)
+
+    @qml.qnode(dev)
+    def circuit(hf_state, h):
+        qml.BasisState(hf_state, wires=range(qubits))
+        return qml.expval(h)
+
+    assert circuit(state_occ, h_occ) == circuit(state_parity, h_parity)
+    assert circuit(state_occ, h_occ) == circuit(state_bk, h_bk)
+
+
 @pytest.mark.parametrize(
     ("electrons", "orbitals", "msg_match"),
     [