compute_rdms function (#228)

* compute_rdms function
* Expanded FermionOperator attributes, defined logical and mathematical operations on it (add, mul, eq, etc)

tangelo/linq/helpers/circuits/measurement_basis.py

@@ -41,6 +41,19 @@ def measurement_basis_gates(term):
     return gates
 
 
+def get_compatible_bases(op, basis_list):
+    """Return a list of measurement bases compatible with the given operator.
+
+    Args:
+        op (str): Pauli string representing the operator to be measured
+        basis_list (list): List of Pauli strings for the measurement bases to be checked
+
+    Returns:
+        list: List of measurement bases compatible with the operator
+    """
+    return [b for b in basis_list if all([(o == p or o == "I") for o, p in zip(op, b)])]
+
+
 def pauli_string_to_of(pauli_string):
     """ Converts a string of I,X,Y,Z Pauli operators to an Openfermion-style
     representation.

tangelo/toolboxes/molecular_computation/rdms.py

@@ -14,9 +14,17 @@
 
 """Module containing functions to manipulate 1- and 2-RDMs."""
 
+import itertools as it
+
 import numpy as np
 
+from tangelo.linq.helpers import pauli_string_to_of, get_compatible_bases
+from tangelo.toolboxes.operators import FermionOperator
+from tangelo.toolboxes.measurements import ClassicalShadow
+from tangelo.toolboxes.post_processing import Histogram, aggregate_histograms
+from tangelo.toolboxes.qubit_mappings.mapping_transform import fermion_to_qubit_mapping, get_qubit_number
 from tangelo.toolboxes.molecular_computation.molecule import spatial_from_spinorb
+from tangelo.linq.helpers.circuits import pauli_of_to_string
 
 
 def matricize_2rdm(two_rdm, n_orbitals):
@@ -41,6 +49,112 @@ def matricize_2rdm(two_rdm, n_orbitals):
     return rho
 
 
+def compute_rdms(ferm_ham, mapping, up_then_down, exp_vals=None, exp_data=None, shadow=None, return_spinsum=True, **eval_args):
+    """
+    Compute the 1- and 2-RDM and their spin-summed versions
+    using a FermionOperator and experimental frequency data in the form of a
+    classical shadow or a dictionary of frequency histograms.
+    Exactly one of the following must be provided by the user:
+    exp_vals, exp_data or shadow
+
+    Args:
+        ferm_ham (FermionicOperator): Fermionic operator with n_spinorbitals, n_electrons, and spin defined
+        mapping (str): Qubit mapping
+        up_then_down (bool): Spin ordering for the mapping
+        exp_vals (dict): Optional, dictionary of Pauli word expectation values
+        exp_data (dict): Optional, dictionary of {basis: histogram} where basis is the measurement basis
+            and histogram is a {bitstring: frequency} dictionary
+        shadow (ClassicalShadow): Optional, a classical shadow object
+        return_spinsum (bool): Optional, if True, return also the spin-summed RDMs
+        eval_args: Optional arguments to pass to the ClassicalShadow object
+
+    Returns:
+        complex array: 1-RDM
+        complex array: 2-RDM
+        complex array: spin-summed 1-RDM
+        complex array: spin-summed 2-RDM
+    """
+    onerdm = np.zeros((ferm_ham.n_spinorbitals,) * 2, dtype=complex)
+    twordm = np.zeros((ferm_ham.n_spinorbitals,) * 4, dtype=complex)
+
+    # Optional arguments are mutually exclusive, return error if several of them have been passed
+    if [exp_vals, exp_data, shadow].count(None) != 2:
+        raise RuntimeError("Arguments exp_vals, exp_data and shadow are mutually exclusive. Provide exactly one of them.")
+
+    # Initialize exp_vals
+    exp_vals = {pauli_string_to_of(term): exp_val for term, exp_val in exp_data.items()} \
+        if isinstance(exp_vals, dict) else {}
+
+    n_qubits = get_qubit_number(mapping, ferm_ham.n_spinorbitals)
+
+    # Go over all terms in fermionic Hamiltonian
+    for term in ferm_ham.terms:
+        length = len(term)
+
+        if not term:
+            continue
+
+        # Fermionic term with a prefactor of 1.0.
+        fermionic_term = FermionOperator(term, 1.0)
+
+        qubit_term = fermion_to_qubit_mapping(fermion_operator=fermionic_term,
+                                              n_spinorbitals=ferm_ham.n_spinorbitals,
+                                              n_electrons=ferm_ham.n_electrons,
+                                              mapping=mapping,
+                                              up_then_down=up_then_down,
+                                              spin=ferm_ham.spin)
+        qubit_term.compress()
+
+        # Loop to go through all qubit terms.
+        eigenvalue = 0.
+
+        if isinstance(shadow, ClassicalShadow):
+            eigenvalue = shadow.get_observable(qubit_term, **eval_args)
+
+        else:
+            for qterm, coeff in qubit_term.terms.items():
+                if coeff.real != 0:
+
+                    if qterm in exp_vals:
+                        exp_val = exp_vals[qterm] if qterm else 1.
+                    else:
+                        ps = pauli_of_to_string(qterm, n_qubits)
+                        bases = get_compatible_bases(ps, list(exp_data.keys()))
+
+                        if not bases:
+                            raise RuntimeError(f"No experimental data for basis {ps}")
+
+                        hist = aggregate_histograms(*[Histogram(exp_data[basis]) for basis in bases])
+                        exp_val = hist.get_expectation_value(qterm, 1.)
+                        exp_vals[qterm] = exp_val
+
+                    eigenvalue += coeff * exp_val
+
+        # Put the values in np arrays (differentiate 1- and 2-RDM)
+        if length == 2:
+            iele, jele = (int(ele[0]) for ele in tuple(term[0:2]))
+            onerdm[iele, jele] += eigenvalue
+        elif length == 4:
+            iele, jele, kele, lele = (int(ele[0]) for ele in tuple(term[0:4]))
+            twordm[iele, lele, jele, kele] += eigenvalue
+
+    if return_spinsum:
+        onerdm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 2, dtype=complex)
+        twordm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 4, dtype=complex)
+
+        # Construct spin-summed 1-RDM.
+        for i, j in it.product(range(onerdm.shape[0]), repeat=2):
+            onerdm_spinsum[i//2, j//2] += onerdm[i, j]
+
+        # Construct spin-summed 2-RDM.
+        for i, j, k, l in it.product(range(twordm.shape[0]), repeat=4):
+            twordm_spinsum[i//2, j//2, k//2, l//2] += twordm[i, j, k, l]
+
+        return onerdm, twordm, onerdm_spinsum, twordm_spinsum
+    else:
+        return onerdm, twordm
+
+
 def energy_from_rdms(ferm_op, one_rdm, two_rdm):
     """Computes the molecular energy from one- and two-particle reduced
     density matrices (RDMs). Coefficients (integrals) are computed from the

tangelo/toolboxes/molecular_computation/tests/data/H2_raw_exact.dat

@@ -0,0 +1 @@
+{"ZZ": {"00": 0.01272572842689497, "10": 2.9605003812919443e-09, "01": 2.9605003812876728e-09, "11": 0.9872742656521027}, "XZ": {"00": 0.006356727736068299, "10": 0.006369003651327052, "01": 0.4936911974715798, "11": 0.49358307114102357}, "ZX": {"00": 0.006356727736068303, "10": 0.4936911974715798, "01": 0.006369003651327048, "11": 0.4935830711410234}, "XX": {"00": 0.19400378295380538, "10": 0.30604414225384274, "01": 0.30604414225384274, "11": 0.1939079325385078}, "YZ": {"00": 0.006362865693697675, "10": 0.006362865693697675, "01": 0.49363713430630146, "11": 0.49363713430630163}, "ZY": {"00": 0.006362865693697675, "10": 0.49363713430630146, "01": 0.006362865693697675, "11": 0.49363713430630163}, "YY": {"00": 0.30604414521434303, "10": 0.1939558547856562, "01": 0.1939558547856562, "11": 0.3060441452143432}, "XY": {"00": 0.25002396260382403, "10": 0.24997603739617527, "01": 0.25002396260382415, "11": 0.24997603739617533}, "YX": {"00": 0.25002396260382403, "10": 0.25002396260382415, "01": 0.24997603739617516, "11": 0.24997603739617527}}

tangelo/toolboxes/molecular_computation/tests/test_rdms.py

@@ -14,39 +14,76 @@
 
 import os
 import unittest
+import json
 
+import numpy as np
+from numpy.testing import assert_allclose
 from openfermion.utils import load_operator
 
+from tangelo.toolboxes.measurements import RandomizedClassicalShadow
 from tangelo.toolboxes.operators import FermionOperator
-from tangelo.toolboxes.molecular_computation.rdms import energy_from_rdms
+from tangelo.toolboxes.molecular_computation.rdms import energy_from_rdms, compute_rdms
 
 # For openfermion.load_operator function.
 pwd_this_test = os.path.dirname(os.path.abspath(__file__))
 
+ferm_op_of = load_operator("H2_ferm_op.data", data_directory=pwd_this_test + "/data", plain_text=True)
+ferm_op = FermionOperator()
+ferm_op.__dict__ = ferm_op_of.__dict__.copy()
+ferm_op.n_spinorbitals = 4
+ferm_op.n_electrons = 2
+ferm_op.spin = 0
+
+exp_data = json.load(open(pwd_this_test + "/data/H2_raw_exact.dat", "r"))
+
+rdm1ssr = np.array([[1.97454854+0.j, 0.+0.j],
+                 [0.+0.j, 0.02545146+0.j]])
+
+rdm2ssr = np.array(
+    [[[[ 1.97454853e+00+0.00000000e+00j,  0.00000000e+00+0.00000000e+00j],
+       [ 0.00000000e+00+0.00000000e+00j,  5.92100152e-09+0.00000000e+00j]],
+      [[ 0.00000000e+00+0.00000000e+00j, -2.24176575e-01+2.77555756e-17j],
+       [ 5.92100077e-09+0.00000000e+00j,  0.00000000e+00+0.00000000e+00j]]],
+     [[[ 0.00000000e+00+0.00000000e+00j,  5.92100077e-09+0.00000000e+00j],
+       [-2.24176575e-01-2.77555756e-17j,  0.00000000e+00+0.00000000e+00j]],
+      [[ 5.92100152e-09+0.00000000e+00j,  0.00000000e+00+0.00000000e+00j],
+       [ 0.00000000e+00+0.00000000e+00j,  2.54514569e-02+0.00000000e+00j]]]])
+
 
 class RDMsUtilitiesTest(unittest.TestCase):
 
     def test_energy_from_rdms(self):
-        """Same test as in test_molecule.SecondQuantizedMoleculeTest.test_energy_from_rdms,
-        but from a fermionic operator instead.
-        """
-        ferm_op_of = load_operator("H2_ferm_op.data", data_directory=pwd_this_test+"/data", plain_text=True)
-        ferm_op = FermionOperator()
-        ferm_op.__dict__ = ferm_op_of.__dict__
-
-        rdm1 = [[ 1.97453997e+00, -7.05987336e-17],
-                [-7.05987336e-17,  2.54600303e-02]]
-
-        rdm2 = [
-            [[[ 1.97453997e+00, -7.96423130e-17], [-7.96423130e-17,  3.21234218e-33]],
-             [[-7.96423130e-17, -2.24213843e-01], [ 0.00000000e+00,  9.04357944e-18]]],
-            [[[-7.96423130e-17,  0.00000000e+00], [-2.24213843e-01,  9.04357944e-18]],
-             [[ 3.21234218e-33,  9.04357944e-18], [ 9.04357944e-18,  2.54600303e-02]]]
-            ]
-
-        e_rdms = energy_from_rdms(ferm_op, rdm1, rdm2)
+        """Compute energy using known spin-summed 1RDM and 2RDM"""
+        e_rdms = energy_from_rdms(ferm_op, rdm1ssr, rdm2ssr)
         self.assertAlmostEqual(e_rdms, -1.1372701, delta=1e-5)
 
+    def test_compute_rdms_from_raw_data(self):
+        """Compute RDMs from frequency list"""
+        rdm1, rdm2, rdm1ss, rdm2ss = compute_rdms(ferm_op, "scbk", True, exp_data=exp_data)
+
+        assert_allclose(rdm1ssr, rdm1ss, rtol=1e-5)
+        assert_allclose(rdm2ssr, rdm2ss, rtol=1e-5)
+
+    def test_compute_rdms_from_classical_shadow(self):
+        """Compute RDMs from classical shadow"""
+        # Construct ClassicalShadow
+        bitstrings = []
+        unitaries = []
+
+        for b, hist in exp_data.items():
+            for s, f in hist.items():
+                factor = round(f * 10000)
+                bitstrings.extend([s] * factor)
+                unitaries.extend([b] * factor)
+
+        cs_data = RandomizedClassicalShadow(unitaries=unitaries, bitstrings=bitstrings)
+
+        rdm1, rdm2, rdm1ss, rdm2ss = compute_rdms(ferm_op, "scbk", True, shadow=cs_data, k=5)
+
+        # Have to adjust tolerance to account for classical shadow rounding to 10000 shots
+        assert_allclose(rdm1ssr, rdm1ss, atol=0.05)
+        assert_allclose(rdm2ssr, rdm2ss, atol=0.05)
+
 
 if __name__ == "__main__":
     unittest.main()

tangelo/toolboxes/operators/operators.py

@@ -21,16 +21,86 @@
 
 import numpy as np
 from scipy.special import comb
+import openfermion as of
 
-# Later on, if needed, we can extract the code for the operators themselves to remove the dependencies and customize
-import openfermion
 
-
-class FermionOperator(openfermion.FermionOperator):
-    """Currently, this class is coming from openfermion. Can be later on be
-    replaced by our own implementation.
+class FermionOperator(of.FermionOperator):
+    """Custom FermionOperator class. Based on openfermion's, with additional functionalities.
     """
 
+    def __init__(self, term=None, coefficient=1., n_spinorbitals=None, n_electrons=None, spin=None):
+        super(FermionOperator, self).__init__(term, coefficient)
+
+        self.n_spinorbitals = n_spinorbitals
+        self.n_electrons = n_electrons
+        self.spin = spin
+
+    def __imul__(self, other):
+        if isinstance(other, FermionOperator):
+            # Raise error if attributes are not the same across Operators.
+            if (self.n_spinorbitals, self.n_electrons, self.spin) != (other.n_spinorbitals, other.n_electrons, other.spin):
+                raise RuntimeError("n_spinorbitals, n_electrons and spin must be the same for all FermionOperators.")
+            else:
+                return super(FermionOperator, self).__imul__(other)
+
+        elif isinstance(other, of.FermionOperator):
+            if (self.n_spinorbitals, self.n_electrons, self.spin) != (None, None, None):
+                raise RuntimeError("openfermion FermionOperator did not define a necessary attribute")
+            else:
+                f_op = FermionOperator()
+                f_op.terms = other.terms.copy()
+                return super(FermionOperator, self).__imul__(f_op)
+
+        else:
+            return super(FermionOperator, self).__imul__(other)
+
+    def __mul__(self, other):
+        return self.__imul__(other)
+
+    def __iadd__(self, other):
+        if isinstance(other, FermionOperator):
+            # Raise error if attributes are not the same across Operators.
+            if (self.n_spinorbitals, self.n_electrons, self.spin) != (other.n_spinorbitals, other.n_electrons, other.spin):
+                raise RuntimeError("n_spinorbitals, n_electrons and spin must be the same for all FermionOperators.")
+            else:
+                return super(FermionOperator, self).__iadd__(other)
+
+        elif isinstance(other, of.FermionOperator):
+            if (self.n_spinorbitals, self.n_electrons, self.spin) != (None, None, None):
+                raise RuntimeError("openfermion FermionOperator did not define a necessary attribute")
+            else:
+                f_op = FermionOperator()
+                f_op.terms = other.terms.copy()
+                return super(FermionOperator, self).__iadd__(f_op)
+
+        else:
+            raise RuntimeError(f"You cannot add FermionOperator and {other.__class__}.")
+
+    def __add__(self, other):
+        return self.__iadd__(other)
+
+    def __radd__(self, other):
+        return self.__iadd__(other)
+
+    def __isub__(self, other):
+        return self.__iadd__(-1. * other)
+
+    def __sub__(self, other):
+        return self.__isub__(other)
+
+    def __rsub__(self, other):
+        return -1 * self.__isub__(other)
+
+    def __eq__(self, other):
+        # Additional checks for == operator.
+        if isinstance(other, FermionOperator):
+            if (self.n_spinorbitals, self.n_electrons, self.spin) == (other.n_spinorbitals, other.n_electrons, other.spin):
+                return super(FermionOperator, self).__eq__(other)
+            else:
+                return False
+        else:
+            return super(FermionOperator, self).__eq__(other)
+
     def get_coeffs(self, coeff_threshold=1e-8):
         """Method to get the coefficient tensors from a fermion operator.
 
@@ -43,7 +113,7 @@ def get_coeffs(self, coeff_threshold=1e-8):
                 two-body coefficient matrices (N*N*N*N), where N is the number
                 of spinorbitals.
         """
-        n_sos = openfermion.count_qubits(self)
+        n_sos = of.count_qubits(self)
 
         constant = 0.
         one_body = np.zeros((n_sos, n_sos), complex)
@@ -73,8 +143,15 @@ def get_coeffs(self, coeff_threshold=1e-8):
 
         return constant, one_body, two_body
 
+    def to_openfermion(self):
+        """Converts Tangelo FermionOperator to openfermion"""
+        ferm_op = of.FermionOperator()
+        ferm_op.terms = self.terms.copy()
+
+        return ferm_op
+
 
-class QubitOperator(openfermion.QubitOperator):
+class QubitOperator(of.QubitOperator):
     """Currently, this class is coming from openfermion. Can be later on be
     replaced by our own implementation.
     """
@@ -226,7 +303,7 @@ def normal_ordered(fe_op):
     """
 
     # Obtain normal ordered fermionic operator as list of terms
-    norm_ord_terms = openfermion.transforms.normal_ordered(fe_op).terms
+    norm_ord_terms = of.transforms.normal_ordered(fe_op).terms
 
     # Regeneratore full operator using class of tangelo.toolboxes.operators.FermionicOperator
     norm_ord_fe_op = FermionOperator()