From bf860efbe08554d87b2c8b67cb4ff5067404d45f Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Thu, 6 Oct 2022 11:56:09 -0700
Subject: [PATCH 1/9] Work in progress on compute_rdms function

---
 .../helpers/circuits/measurement_basis.py     | 13 +++
 .../toolboxes/molecular_computation/rdms.py   | 94 +++++++++++++++++++
 .../molecular_computation/tests/test_rdms.py  | 44 ++++++---
 tangelo/toolboxes/operators/operators.py      |  9 +-
 4 files changed, 144 insertions(+), 16 deletions(-)

diff --git a/tangelo/linq/helpers/circuits/measurement_basis.py b/tangelo/linq/helpers/circuits/measurement_basis.py
index f7b9531f4..fac71d74a 100644
--- a/tangelo/linq/helpers/circuits/measurement_basis.py
+++ b/tangelo/linq/helpers/circuits/measurement_basis.py
@@ -41,6 +41,19 @@ def measurement_basis_gates(term):
     return gates
 
 
+def filter_bases(op, basis_list):
+    """Returns 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.
diff --git a/tangelo/toolboxes/molecular_computation/rdms.py b/tangelo/toolboxes/molecular_computation/rdms.py
index 76be298a2..0a6ba1c7b 100644
--- a/tangelo/toolboxes/molecular_computation/rdms.py
+++ b/tangelo/toolboxes/molecular_computation/rdms.py
@@ -14,9 +14,16 @@
 
 """Module containing functions to manipulate 1- and 2-RDMs."""
 
+import itertools as it
 import numpy as np
 
+from tangelo.linq.helpers import filter_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
 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 +48,93 @@ def matricize_2rdm(two_rdm, n_orbitals):
     return rho
 
 
+def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
+    """
+    Computes the 1- and 2-RDM and their spin-summed versions 
+    using a Molecule object and frequency data either in the form of a 
+    classical shadow or a dictionary of frequency histograms.
+
+    Args:
+        ferm_ham (FermionicOperator): Fermionic operator with n_spinorbitals and n_electrons defined
+        exp_data (ClassicalShadow or dict): classical shadow or a dictionary of expectation values of qubit terms
+        mapping: qubit mapping
+        up_then_down: spin ordering for the mapping
+
+    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)
+    onerdm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 2, dtype=complex)
+    twordm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 4, dtype=complex)
+
+    exp_vals = {}
+
+    # 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 type(exp_data) == ClassicalShadow:
+            for qterm, coeff in qubit_term.terms.items():
+                if coeff.real != 0:
+                    # Change depending on if it is randomized or not.
+                    eigenvalue += exp_data.get_term_observable(qterm, coeff)
+
+        if type(exp_data) == dict:
+            for qterm, coeff in qubit_term.terms.items():
+                if coeff.real != 0:
+
+                    try:
+                        exp_vals[qterm]
+                    except KeyError:
+                        if qterm:
+                            ps = pauli_of_to_string(qterm, ferm_ham.n_spinorbitals // 2)  # not sure about the number
+                            exp_vals[qterm] = aggregate_histograms(*[Histogram(exp_data[basis]) for basis in filter_bases(ps, exp_data.keys())]).get_expectation_value(qterm, 1.)
+                        else:
+                            continue
+
+                    exp_val = exp_vals[qterm] if qterm else 1.
+                    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
+
+    # 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
+
+
 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
diff --git a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
index 1f5ed73c2..6bb49fc44 100644
--- a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
+++ b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
@@ -14,15 +14,35 @@
 
 import os
 import unittest
+import json
+import numpy as np
 
 from openfermion.utils import load_operator
 
 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__
+ferm_op.n_spinorbitals = 4
+ferm_op.n_electrons = 2
+ferm_op.spin = 0
+
+cs_data = json.load(open("./data/H2_cs_data.dat", "r"))
+
+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]]]
+]
 
 class RDMsUtilitiesTest(unittest.TestCase):
 
@@ -30,23 +50,17 @@ 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)
         self.assertAlmostEqual(e_rdms, -1.1372701, delta=1e-5)
 
+    def test_compute_rdms(self):
+        """Load data and compute RDMs from frequency list"""
+        rdm1r, rdm2r, rdm1ssr, rdm2ssr = compute_rdms(ferm_op, cs_data, "scbk", False)
+
+        self.assertAlmostEqual((np.array(rdm1)-rdm1ssr).sum(), 0., delta=1e-5)
+        #self.assertAlmostEqual(np.array(rdm2), rdm2ssr, delta=1e-5)
+        #self.assertAlmostEqual(np.array(rdm1ss), rdm1ssr, delta=1e-5)
+        #self.assertAlmostEqual(np.array(rdm2ss), rdm2ssr, delta=1e-5)
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tangelo/toolboxes/operators/operators.py b/tangelo/toolboxes/operators/operators.py
index 854ff5c72..5556340ce 100644
--- a/tangelo/toolboxes/operators/operators.py
+++ b/tangelo/toolboxes/operators/operators.py
@@ -27,10 +27,17 @@
 
 
 class FermionOperator(openfermion.FermionOperator):
-    """Currently, this class is coming from openfermion. Can be later on be
+    """Currently, this class is coming from openfermion. Can later on be
     replaced by our own implementation.
     """
 
+    def __init__(self, term=None, coefficient=1., n_spinorbitals=None, n_electrons=None, spin=None):
+        super(openfermion.FermionOperator, self).__init__(term, coefficient)
+
+        self.n_spinorbitals = n_spinorbitals
+        self.n_electrons = n_electrons
+        self.spin = spin
+    
     def get_coeffs(self, coeff_threshold=1e-8):
         """Method to get the coefficient tensors from a fermion operator.
 

From 65e0b32b25cca18a72f7730bcf293afc9f117b51 Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Thu, 6 Oct 2022 16:33:54 -0700
Subject: [PATCH 2/9] First draft of compute_rdms

---
 .../toolboxes/molecular_computation/rdms.py   | 33 ++++++-----
 .../tests/data/H2_raw_exact.dat               |  1 +
 .../molecular_computation/tests/test_rdms.py  | 56 +++++++++++++------
 3 files changed, 58 insertions(+), 32 deletions(-)
 create mode 100644 tangelo/toolboxes/molecular_computation/tests/data/H2_raw_exact.dat

diff --git a/tangelo/toolboxes/molecular_computation/rdms.py b/tangelo/toolboxes/molecular_computation/rdms.py
index 0a6ba1c7b..30a8989a2 100644
--- a/tangelo/toolboxes/molecular_computation/rdms.py
+++ b/tangelo/toolboxes/molecular_computation/rdms.py
@@ -17,11 +17,11 @@
 import itertools as it
 import numpy as np
 
-from tangelo.linq.helpers import filter_bases
+from tangelo.linq.helpers import filter_bases, pauli_string_to_of
 from tangelo.toolboxes.operators import FermionOperator
-from tangelo.toolboxes.measurements import ClassicalShadow
+from tangelo.toolboxes.measurements import RandomizedClassicalShadow
 from tangelo.toolboxes.post_processing import Histogram, aggregate_histograms
-from tangelo.toolboxes.qubit_mappings.mapping_transform import fermion_to_qubit_mapping
+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
 
@@ -71,7 +71,13 @@ def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
     onerdm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 2, dtype=complex)
     twordm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 4, dtype=complex)
 
-    exp_vals = {}
+    # Check if the data dict is expectation values, or raw frequency data in the {basis: hist} format
+    if isinstance(exp_data, dict) and isinstance(list(exp_data.values())[0], float):
+        exp_vals = {pauli_string_to_of(term): exp_val for term,exp_val in exp_data.items()}
+    elif isinstance(list(exp_data.values())[0], dict):
+        exp_vals = {}
+
+    n_qubits = get_qubit_number(mapping, ferm_ham.n_spinorbitals)
 
     # Go over all terms in fermionic Hamiltonian
     for term in ferm_ham.terms:
@@ -87,33 +93,30 @@ def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
                                             n_spinorbitals = ferm_ham.n_spinorbitals,
                                             n_electrons = ferm_ham.n_electrons,
                                             mapping = mapping,
-                                            up_then_down = up_then_down,
-                                            spin = ferm_ham.spin)
+                                            up_then_down = up_then_down)
         qubit_term.compress()
 
         # Loop to go through all qubit terms.
         eigenvalue = 0.
         
-        if type(exp_data) == ClassicalShadow:
+        if isinstance(exp_data, RandomizedClassicalShadow):
             for qterm, coeff in qubit_term.terms.items():
                 if coeff.real != 0:
                     # Change depending on if it is randomized or not.
                     eigenvalue += exp_data.get_term_observable(qterm, coeff)
 
-        if type(exp_data) == dict:
+        if isinstance(exp_data, dict):
             for qterm, coeff in qubit_term.terms.items():
                 if coeff.real != 0:
 
                     try:
-                        exp_vals[qterm]
+                        exp_val = exp_vals[qterm] if qterm else 1.
                     except KeyError:
-                        if qterm:
-                            ps = pauli_of_to_string(qterm, ferm_ham.n_spinorbitals // 2)  # not sure about the number
-                            exp_vals[qterm] = aggregate_histograms(*[Histogram(exp_data[basis]) for basis in filter_bases(ps, exp_data.keys())]).get_expectation_value(qterm, 1.)
-                        else:
-                            continue
+                        ps = pauli_of_to_string(qterm, n_qubits)
+                        hist = aggregate_histograms(*[Histogram(exp_data[basis]) for basis in filter_bases(ps, exp_data.keys())])
+                        exp_val = hist.get_expectation_value(qterm, 1.)
+                        exp_vals[qterm] = exp_val
 
-                    exp_val = exp_vals[qterm] if qterm else 1.
                     eigenvalue += coeff * exp_val
 
         # Put the values in np arrays (differentiate 1- and 2-RDM)
diff --git a/tangelo/toolboxes/molecular_computation/tests/data/H2_raw_exact.dat b/tangelo/toolboxes/molecular_computation/tests/data/H2_raw_exact.dat
new file mode 100644
index 000000000..9738e3cb1
--- /dev/null
+++ b/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}}
\ No newline at end of file
diff --git a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
index 6bb49fc44..58b41bb2f 100644
--- a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
+++ b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
@@ -16,9 +16,11 @@
 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, compute_rdms
 
@@ -32,35 +34,55 @@
 ferm_op.n_electrons = 2
 ferm_op.spin = 0
 
-cs_data = json.load(open("./data/H2_cs_data.dat", "r"))
+exp_data = json.load(open("./data/H2_raw_exact.dat", "r"))
+n_shots = 10000
 
-rdm1 = [[1.97453997e+00, -7.05987336e-17],
-        [-7.05987336e-17, 2.54600303e-02]]
+# Construct ClassicalShadow
+bitstrings = []
+unitaries = []
 
-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]]]
-]
+for b, hist in exp_data.items():
+    for s,f in hist.items():
+        factor = round(f*n_shots)
+        bitstrings.extend([s]*factor)
+        unitaries.extend([b]*factor)
+
+cs_data = RandomizedClassicalShadow(unitaries=unitaries, bitstrings=bitstrings)
+
+
+rdm1 = np.array([[1.97454854+0.j, 0.+0.j],
+                 [0.+0.j, 0.02545146+0.j]])
+
+rdm2 = 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.
-        """
+        """Computes energy using known spin-summed 1RDM and 2RDM"""
         e_rdms = energy_from_rdms(ferm_op, rdm1, rdm2)
         self.assertAlmostEqual(e_rdms, -1.1372701, delta=1e-5)
 
-    def test_compute_rdms(self):
+    def test_compute_rdms_from_raw_data(self):
+        """Load data and compute RDMs from frequency list"""
+        rdm1r, rdm2r, rdm1ssr, rdm2ssr = compute_rdms(ferm_op, exp_data, "scbk", False)
+
+        assert_allclose(rdm1, rdm1ssr, rtol=1e-5)
+        assert_allclose(rdm2, rdm2ssr, rtol=1e-5)
+
+    def test_compute_rdms_from_classical_shadow(self):
         """Load data and compute RDMs from frequency list"""
         rdm1r, rdm2r, rdm1ssr, rdm2ssr = compute_rdms(ferm_op, cs_data, "scbk", False)
 
-        self.assertAlmostEqual((np.array(rdm1)-rdm1ssr).sum(), 0., delta=1e-5)
-        #self.assertAlmostEqual(np.array(rdm2), rdm2ssr, delta=1e-5)
-        #self.assertAlmostEqual(np.array(rdm1ss), rdm1ssr, delta=1e-5)
-        #self.assertAlmostEqual(np.array(rdm2ss), rdm2ssr, delta=1e-5)
+        assert_allclose(rdm1, rdm1ssr, rtol=1e-5)
+        assert_allclose(rdm2, rdm2ssr, rtol=1e-5)
 
 if __name__ == "__main__":
     unittest.main()

From d068f3caf4234a57ce46ac083a13959709f35cfd Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Thu, 6 Oct 2022 17:13:06 -0700
Subject: [PATCH 3/9] fixed tests

---
 tangelo/toolboxes/molecular_computation/rdms.py          | 4 ++--
 .../toolboxes/molecular_computation/tests/test_rdms.py   | 9 +++++----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/tangelo/toolboxes/molecular_computation/rdms.py b/tangelo/toolboxes/molecular_computation/rdms.py
index 30a8989a2..d9a692214 100644
--- a/tangelo/toolboxes/molecular_computation/rdms.py
+++ b/tangelo/toolboxes/molecular_computation/rdms.py
@@ -71,10 +71,10 @@ def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
     onerdm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 2, dtype=complex)
     twordm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 4, dtype=complex)
 
-    # Check if the data dict is expectation values, or raw frequency data in the {basis: hist} format
+    # Check if the data dict is expectation values or raw frequency data in the {basis: hist} format
     if isinstance(exp_data, dict) and isinstance(list(exp_data.values())[0], float):
         exp_vals = {pauli_string_to_of(term): exp_val for term,exp_val in exp_data.items()}
-    elif isinstance(list(exp_data.values())[0], dict):
+    elif isinstance(exp_data, dict) and isinstance(list(exp_data.values())[0], dict):
         exp_vals = {}
 
     n_qubits = get_qubit_number(mapping, ferm_ham.n_spinorbitals)
diff --git a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
index 58b41bb2f..e729e538c 100644
--- a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
+++ b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
@@ -34,7 +34,7 @@
 ferm_op.n_electrons = 2
 ferm_op.spin = 0
 
-exp_data = json.load(open("./data/H2_raw_exact.dat", "r"))
+exp_data = json.load(open(pwd_this_test + "/data/H2_raw_exact.dat", "r"))
 n_shots = 10000
 
 # Construct ClassicalShadow
@@ -78,11 +78,12 @@ def test_compute_rdms_from_raw_data(self):
         assert_allclose(rdm2, rdm2ssr, rtol=1e-5)
 
     def test_compute_rdms_from_classical_shadow(self):
-        """Load data and compute RDMs from frequency list"""
+        """Load data and compute RDMs from classical shadow"""
         rdm1r, rdm2r, rdm1ssr, rdm2ssr = compute_rdms(ferm_op, cs_data, "scbk", False)
 
-        assert_allclose(rdm1, rdm1ssr, rtol=1e-5)
-        assert_allclose(rdm2, rdm2ssr, rtol=1e-5)
+        # Have to adjust tolerance to account for classical shadow rounding to 10000 shots
+        assert_allclose(rdm1, rdm1ssr, atol=0.01)
+        assert_allclose(rdm2, rdm2ssr, atol=0.01)
 
 if __name__ == "__main__":
     unittest.main()

From 5ce88a15ccea1c69d59bf0bb0b211325370274c1 Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Thu, 6 Oct 2022 17:42:26 -0700
Subject: [PATCH 4/9] Fixed code style conformance

---
 .../helpers/circuits/measurement_basis.py     |  2 +-
 .../toolboxes/molecular_computation/rdms.py   | 20 +++++++++----------
 .../molecular_computation/tests/test_rdms.py  | 20 ++++++++++---------
 tangelo/toolboxes/operators/operators.py      |  2 +-
 4 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/tangelo/linq/helpers/circuits/measurement_basis.py b/tangelo/linq/helpers/circuits/measurement_basis.py
index fac71d74a..85d475363 100644
--- a/tangelo/linq/helpers/circuits/measurement_basis.py
+++ b/tangelo/linq/helpers/circuits/measurement_basis.py
@@ -51,7 +51,7 @@ def filter_bases(op, basis_list):
     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)])]
+    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):
diff --git a/tangelo/toolboxes/molecular_computation/rdms.py b/tangelo/toolboxes/molecular_computation/rdms.py
index d9a692214..45233929e 100644
--- a/tangelo/toolboxes/molecular_computation/rdms.py
+++ b/tangelo/toolboxes/molecular_computation/rdms.py
@@ -50,8 +50,8 @@ def matricize_2rdm(two_rdm, n_orbitals):
 
 def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
     """
-    Computes the 1- and 2-RDM and their spin-summed versions 
-    using a Molecule object and frequency data either in the form of a 
+    Computes the 1- and 2-RDM and their spin-summed versions
+    using a Molecule object and frequency data either in the form of a
     classical shadow or a dictionary of frequency histograms.
 
     Args:
@@ -73,7 +73,7 @@ def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
 
     # Check if the data dict is expectation values or raw frequency data in the {basis: hist} format
     if isinstance(exp_data, dict) and isinstance(list(exp_data.values())[0], float):
-        exp_vals = {pauli_string_to_of(term): exp_val for term,exp_val in exp_data.items()}
+        exp_vals = {pauli_string_to_of(term): exp_val for term, exp_val in exp_data.items()}
     elif isinstance(exp_data, dict) and isinstance(list(exp_data.values())[0], dict):
         exp_vals = {}
 
@@ -85,20 +85,20 @@ def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
 
         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)
+        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)
         qubit_term.compress()
 
         # Loop to go through all qubit terms.
         eigenvalue = 0.
-        
+
         if isinstance(exp_data, RandomizedClassicalShadow):
             for qterm, coeff in qubit_term.terms.items():
                 if coeff.real != 0:
diff --git a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
index e729e538c..50ca6f241 100644
--- a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
+++ b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
@@ -42,7 +42,7 @@
 unitaries = []
 
 for b, hist in exp_data.items():
-    for s,f in hist.items():
+    for s, f in hist.items():
         factor = round(f*n_shots)
         bitstrings.extend([s]*factor)
         unitaries.extend([b]*factor)
@@ -54,17 +54,18 @@
                  [0.+0.j, 0.02545146+0.j]])
 
 rdm2 = 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]]]])
+    [[[[ 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):
         """Computes energy using known spin-summed 1RDM and 2RDM"""
         e_rdms = energy_from_rdms(ferm_op, rdm1, rdm2)
@@ -85,5 +86,6 @@ def test_compute_rdms_from_classical_shadow(self):
         assert_allclose(rdm1, rdm1ssr, atol=0.01)
         assert_allclose(rdm2, rdm2ssr, atol=0.01)
 
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tangelo/toolboxes/operators/operators.py b/tangelo/toolboxes/operators/operators.py
index 5556340ce..1ef9c8bd0 100644
--- a/tangelo/toolboxes/operators/operators.py
+++ b/tangelo/toolboxes/operators/operators.py
@@ -37,7 +37,7 @@ def __init__(self, term=None, coefficient=1., n_spinorbitals=None, n_electrons=N
         self.n_spinorbitals = n_spinorbitals
         self.n_electrons = n_electrons
         self.spin = spin
-    
+
     def get_coeffs(self, coeff_threshold=1e-8):
         """Method to get the coefficient tensors from a fermion operator.
 

From 8a1a4b4c17ce275ad749df007491c301ad64051e Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Fri, 14 Oct 2022 15:16:07 -0700
Subject: [PATCH 5/9] Fixes after first PR review

---
 .../helpers/circuits/measurement_basis.py     |  2 +-
 .../toolboxes/molecular_computation/rdms.py   | 70 +++++++++++--------
 .../molecular_computation/tests/test_rdms.py  | 18 ++---
 tangelo/toolboxes/operators/operators.py      | 49 ++++++++++++-
 .../operators/tests/test_operators.py         | 27 ++++++-
 5 files changed, 123 insertions(+), 43 deletions(-)

diff --git a/tangelo/linq/helpers/circuits/measurement_basis.py b/tangelo/linq/helpers/circuits/measurement_basis.py
index 85d475363..24657ae7e 100644
--- a/tangelo/linq/helpers/circuits/measurement_basis.py
+++ b/tangelo/linq/helpers/circuits/measurement_basis.py
@@ -42,7 +42,7 @@ def measurement_basis_gates(term):
 
 
 def filter_bases(op, basis_list):
-    """Returns a list of measurement bases compatible with the given operator.
+    """Return a list of measurement bases compatible with the given operator.
 
     Args:
         op (str): Pauli string representing the operator to be measured
diff --git a/tangelo/toolboxes/molecular_computation/rdms.py b/tangelo/toolboxes/molecular_computation/rdms.py
index 45233929e..0c3861cdd 100644
--- a/tangelo/toolboxes/molecular_computation/rdms.py
+++ b/tangelo/toolboxes/molecular_computation/rdms.py
@@ -19,7 +19,7 @@
 
 from tangelo.linq.helpers import filter_bases, pauli_string_to_of
 from tangelo.toolboxes.operators import FermionOperator
-from tangelo.toolboxes.measurements import RandomizedClassicalShadow
+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
@@ -48,17 +48,22 @@ def matricize_2rdm(two_rdm, n_orbitals):
     return rho
 
 
-def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
+def compute_rdms(ferm_ham, mapping, up_then_down, exp_vals=None, exp_data=None, shadow=None, return_spinsum=True, **eval_args):
     """
-    Computes the 1- and 2-RDM and their spin-summed versions
+    Return the 1- and 2-RDM and their spin-summed versions
     using a Molecule object and frequency data either in the form of a
     classical shadow or a dictionary of frequency histograms.
 
     Args:
-        ferm_ham (FermionicOperator): Fermionic operator with n_spinorbitals and n_electrons defined
-        exp_data (ClassicalShadow or dict): classical shadow or a dictionary of expectation values of qubit terms
-        mapping: qubit mapping
-        up_then_down: spin ordering for the mapping
+        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): Dictionary of Pauli word expectation values
+        exp_data (dict): Dictionary of {basis: histogram} where basis is the measurement basis
+            and histogram is a {bitstring: frequency} dictionary
+        shadow (ClassicalShadow): A classical shadow object
+        return_spinsum (bool): If True, return also the spin-summed RDMs
+        eval_args: Optional arguments to pass to the ClassicalShadow object
 
     Returns:
         complex array: 1-RDM
@@ -68,14 +73,16 @@ def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
     """
     onerdm = np.zeros((ferm_ham.n_spinorbitals,) * 2, dtype=complex)
     twordm = np.zeros((ferm_ham.n_spinorbitals,) * 4, dtype=complex)
-    onerdm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 2, dtype=complex)
-    twordm_spinsum = np.zeros((ferm_ham.n_spinorbitals//2,) * 4, dtype=complex)
 
     # Check if the data dict is expectation values or raw frequency data in the {basis: hist} format
-    if isinstance(exp_data, dict) and isinstance(list(exp_data.values())[0], float):
+    if isinstance(exp_vals, dict) and (exp_data is None) and (shadow is None):
         exp_vals = {pauli_string_to_of(term): exp_val for term, exp_val in exp_data.items()}
-    elif isinstance(exp_data, dict) and isinstance(list(exp_data.values())[0], dict):
+    if (exp_vals is None) and isinstance(exp_data, dict) and (shadow is None):
         exp_vals = {}
+    if (exp_vals is None) and (exp_data is None) and isinstance(shadow, ClassicalShadow):
+        pass
+    else:
+        raise RuntimeError("Arguments exp_vals, exp_data and shadow are mutually exclusive. Provide only one of them.")
 
     n_qubits = get_qubit_number(mapping, ferm_ham.n_spinorbitals)
 
@@ -90,22 +97,20 @@ def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
         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)
+                                              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(exp_data, RandomizedClassicalShadow):
-            for qterm, coeff in qubit_term.terms.items():
-                if coeff.real != 0:
-                    # Change depending on if it is randomized or not.
-                    eigenvalue += exp_data.get_term_observable(qterm, coeff)
+        if isinstance(shadow, ClassicalShadow):
+            eigenvalue = shadow.get_observable(qubit_term, **eval_args)
 
-        if isinstance(exp_data, dict):
+        else:
             for qterm, coeff in qubit_term.terms.items():
                 if coeff.real != 0:
 
@@ -113,7 +118,8 @@ def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
                         exp_val = exp_vals[qterm] if qterm else 1.
                     except KeyError:
                         ps = pauli_of_to_string(qterm, n_qubits)
-                        hist = aggregate_histograms(*[Histogram(exp_data[basis]) for basis in filter_bases(ps, exp_data.keys())])
+                        hist = aggregate_histograms(*[Histogram(exp_data[basis])
+                                                      for basis in filter_bases(ps, list(exp_data.keys()))])
                         exp_val = hist.get_expectation_value(qterm, 1.)
                         exp_vals[qterm] = exp_val
 
@@ -127,15 +133,21 @@ def compute_rdms(ferm_ham, exp_data, mapping, up_then_down):
             iele, jele, kele, lele = (int(ele[0]) for ele in tuple(term[0:4]))
             twordm[iele, lele, jele, kele] += eigenvalue
 
-    # 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]
+    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]
+        # 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
+        return onerdm, twordm, onerdm_spinsum, twordm_spinsum
+    else:
+        return onerdm, twordm
 
 
 def energy_from_rdms(ferm_op, one_rdm, two_rdm):
diff --git a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
index 50ca6f241..982e7af7b 100644
--- a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
+++ b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
@@ -29,7 +29,7 @@
 
 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__
+ferm_op.__dict__ = ferm_op_of.__dict__.copy()
 ferm_op.n_spinorbitals = 4
 ferm_op.n_electrons = 2
 ferm_op.spin = 0
@@ -63,28 +63,28 @@
       [[ 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):
 
+class RDMsUtilitiesTest(unittest.TestCase):
 
     def test_energy_from_rdms(self):
-        """Computes energy using known spin-summed 1RDM and 2RDM"""
+        """Compute energy using known spin-summed 1RDM and 2RDM"""
         e_rdms = energy_from_rdms(ferm_op, rdm1, rdm2)
         self.assertAlmostEqual(e_rdms, -1.1372701, delta=1e-5)
 
     def test_compute_rdms_from_raw_data(self):
-        """Load data and compute RDMs from frequency list"""
-        rdm1r, rdm2r, rdm1ssr, rdm2ssr = compute_rdms(ferm_op, exp_data, "scbk", False)
+        """Compute RDMs from frequency list"""
+        rdm1r, rdm2r, rdm1ssr, rdm2ssr = compute_rdms(ferm_op, "scbk", True, exp_data=exp_data)
 
         assert_allclose(rdm1, rdm1ssr, rtol=1e-5)
         assert_allclose(rdm2, rdm2ssr, rtol=1e-5)
 
     def test_compute_rdms_from_classical_shadow(self):
-        """Load data and compute RDMs from classical shadow"""
-        rdm1r, rdm2r, rdm1ssr, rdm2ssr = compute_rdms(ferm_op, cs_data, "scbk", False)
+        """Compute RDMs from classical shadow"""
+        rdm1r, rdm2r, rdm1ssr, rdm2ssr = 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(rdm1, rdm1ssr, atol=0.01)
-        assert_allclose(rdm2, rdm2ssr, atol=0.01)
+        assert_allclose(rdm1, rdm1ssr, atol=0.05)
+        assert_allclose(rdm2, rdm2ssr, atol=0.05)
 
 
 if __name__ == "__main__":
diff --git a/tangelo/toolboxes/operators/operators.py b/tangelo/toolboxes/operators/operators.py
index 1ef9c8bd0..2700f3c50 100644
--- a/tangelo/toolboxes/operators/operators.py
+++ b/tangelo/toolboxes/operators/operators.py
@@ -18,6 +18,7 @@
 
 from math import sqrt
 from collections import OrderedDict
+from copy import copy, deepcopy
 
 import numpy as np
 from scipy.special import comb
@@ -27,17 +28,52 @@
 
 
 class FermionOperator(openfermion.FermionOperator):
-    """Currently, this class is coming from openfermion. Can later on be
-    replaced by our own implementation.
+    """Add n_spinorbitals, n_electrons, and spin parameters to the openfermion.FermionOperator class
     """
 
     def __init__(self, term=None, coefficient=1., n_spinorbitals=None, n_electrons=None, spin=None):
-        super(openfermion.FermionOperator, self).__init__(term, coefficient)
+        super(FermionOperator, self).__init__(term, coefficient)
 
         self.n_spinorbitals = n_spinorbitals
         self.n_electrons = n_electrons
         self.spin = spin
 
+    def __iadd__(self, other):
+        # Raise error if attributes are not the same across Operators.
+        if (self.n_spinorbitals is not None and other.n_spinorbitals is not None)\
+                or (self.n_electrons is not None and other.n_electrons is not None)\
+                or (self.spin is not None and other.spin is not None):
+
+            if self.n_spinorbitals != other.n_spinorbitals:
+                raise RuntimeError("n_spinorbitals must be the same for all FermionOperators.")
+            elif self.n_electrons != other.n_electrons:
+                raise RuntimeError("n_electrons must be the same for all FermionOperators.")
+            elif self.spin != other.spin:
+                raise RuntimeError("Spin must be the same for all FermionOperators.")
+
+        if isinstance(other, FermionOperator):
+            return super(FermionOperator, self).__iadd__(other)
+        elif isinstance(other, openfermion.FermionOperator):
+            if self.n_spinorbitals is not None or self.n_electrons is not None or self.spin is not 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)
+
+    def __add__(self, other):
+        return self.__iadd__(other)
+
+    def __radd__(self, other):
+        return self.__iadd__(other)
+
+    def __eq__(self, other):
+        # Additional checks for == operator.
+        if (self.n_spinorbitals != other.n_spinorbitals) or (self.n_electrons != other.n_electrons) or (self.spin != other.spin):
+            return False
+
+        return super(FermionOperator, self).__eq__(other)
+
     def get_coeffs(self, coeff_threshold=1e-8):
         """Method to get the coefficient tensors from a fermion operator.
 
@@ -80,6 +116,13 @@ 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 = openfermion.FermionOperator()
+        ferm_op.terms = self.terms.copy()
+
+        return ferm_op
+
 
 class QubitOperator(openfermion.QubitOperator):
     """Currently, this class is coming from openfermion. Can be later on be
diff --git a/tangelo/toolboxes/operators/tests/test_operators.py b/tangelo/toolboxes/operators/tests/test_operators.py
index 17feff853..62e2d3411 100644
--- a/tangelo/toolboxes/operators/tests/test_operators.py
+++ b/tangelo/toolboxes/operators/tests/test_operators.py
@@ -15,6 +15,7 @@
 import unittest
 
 import numpy as np
+import openfermion as of
 
 from tangelo.toolboxes.operators import QubitHamiltonian, FermionOperator, QubitOperator, count_qubits, qubitop_to_qubitham
 
@@ -73,6 +74,30 @@ def test_get_coeff(self):
         np.testing.assert_array_almost_equal(ref_one_body, one_body)
         np.testing.assert_array_almost_equal(ref_two_body, two_body)
 
+    def test_add_FermionOperator(self):
+        # addition between two compatible tangelo FermionOperator
+        fop = FermionOperator("0^ 0", 2.) + FermionOperator("0^ 1", 3.)
+
+        # addition between two incompatible tangelo FermionOperator
+        fop_1 = FermionOperator("0^ 0", 2.)
+        fop_2 = FermionOperator("0^ 1", 3.)
+        fop_1.spin = 1
+        fop_2.spin = 0
+        with self.assertRaises(RuntimeError):
+            fop_1 + fop_2
+
+        # addition between openfermion FermionOperator and Tangelo equivalent
+        fop_1 = FermionOperator("0^ 0", 2.) + of.FermionOperator("0^ 1", 3.)
+        self.assertTrue(isinstance(fop_1, FermionOperator))
+
+        # Reverse order addition test
+        fop_2 = of.FermionOperator("0^ 0", 2.) + FermionOperator("0^ 1", 3.)
+        self.assertEqual(fop_1, fop_2)
+
+        fop = FermionOperator("0^ 0", 2.)
+        fop += FermionOperator("0^ 1", 3.)
+        self.assertEqual(fop, fop_1)
+
 
 class QubitHamiltonianTest(unittest.TestCase):
 
@@ -85,7 +110,7 @@ def test_instantiate_QubitHamiltonian(self):
 
         self.assertDictEqual(qubit_ham.__dict__, reference_attributes)
 
-    def test_instantiate_QubitHamiltonian(self):
+    def test_add_QubitHamiltonian(self):
         """Test error raising when 2 incompatible QubitHamiltonian are summed up
         together.
         """

From 1b6a8c7fc6e3baa0fda47aae0e00bcf2d3c2e111 Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Fri, 14 Oct 2022 17:51:05 -0700
Subject: [PATCH 6/9] Fixed operator and molecule tests

---
 tangelo/toolboxes/molecular_computation/rdms.py | 4 ++--
 tangelo/toolboxes/operators/operators.py        | 4 +++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/tangelo/toolboxes/molecular_computation/rdms.py b/tangelo/toolboxes/molecular_computation/rdms.py
index 0c3861cdd..068b2f98c 100644
--- a/tangelo/toolboxes/molecular_computation/rdms.py
+++ b/tangelo/toolboxes/molecular_computation/rdms.py
@@ -77,9 +77,9 @@ def compute_rdms(ferm_ham, mapping, up_then_down, exp_vals=None, exp_data=None,
     # Check if the data dict is expectation values or raw frequency data in the {basis: hist} format
     if isinstance(exp_vals, dict) and (exp_data is None) and (shadow is None):
         exp_vals = {pauli_string_to_of(term): exp_val for term, exp_val in exp_data.items()}
-    if (exp_vals is None) and isinstance(exp_data, dict) and (shadow is None):
+    elif (exp_vals is None) and isinstance(exp_data, dict) and (shadow is None):
         exp_vals = {}
-    if (exp_vals is None) and (exp_data is None) and isinstance(shadow, ClassicalShadow):
+    elif (exp_vals is None) and (exp_data is None) and isinstance(shadow, ClassicalShadow):
         pass
     else:
         raise RuntimeError("Arguments exp_vals, exp_data and shadow are mutually exclusive. Provide only one of them.")
diff --git a/tangelo/toolboxes/operators/operators.py b/tangelo/toolboxes/operators/operators.py
index 2700f3c50..7ab26876d 100644
--- a/tangelo/toolboxes/operators/operators.py
+++ b/tangelo/toolboxes/operators/operators.py
@@ -69,7 +69,9 @@ def __radd__(self, other):
 
     def __eq__(self, other):
         # Additional checks for == operator.
-        if (self.n_spinorbitals != other.n_spinorbitals) or (self.n_electrons != other.n_electrons) or (self.spin != other.spin):
+        if isinstance(other, openfermion.FermionOperator):
+            return super(FermionOperator, self).__eq__(other)
+        elif (self.n_spinorbitals != other.n_spinorbitals) or (self.n_electrons != other.n_electrons) or (self.spin != other.spin):
             return False
 
         return super(FermionOperator, self).__eq__(other)

From 6e4bf370916fb7b919dd26b88dee70c47c030a8f Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Wed, 19 Oct 2022 14:40:56 -0700
Subject: [PATCH 7/9] Added __imul__, __isub__ and derived operators, and tests

---
 tangelo/toolboxes/operators/operators.py      | 58 ++++++++++++++++---
 .../operators/tests/test_operators.py         | 46 ++++++++++++++-
 2 files changed, 92 insertions(+), 12 deletions(-)

diff --git a/tangelo/toolboxes/operators/operators.py b/tangelo/toolboxes/operators/operators.py
index 7ab26876d..783656bf7 100644
--- a/tangelo/toolboxes/operators/operators.py
+++ b/tangelo/toolboxes/operators/operators.py
@@ -18,16 +18,16 @@
 
 from math import sqrt
 from collections import OrderedDict
-from copy import copy, deepcopy
 
 import numpy as np
+from openfermion.ops.operators.symbolic_operator import COEFFICIENT_TYPES
 from scipy.special import comb
 
 # Later on, if needed, we can extract the code for the operators themselves to remove the dependencies and customize
-import openfermion
+import openfermion as of
 
 
-class FermionOperator(openfermion.FermionOperator):
+class FermionOperator(of.FermionOperator):
     """Add n_spinorbitals, n_electrons, and spin parameters to the openfermion.FermionOperator class
     """
 
@@ -38,6 +38,37 @@ def __init__(self, term=None, coefficient=1., n_spinorbitals=None, n_electrons=N
         self.n_electrons = n_electrons
         self.spin = spin
 
+    def __imul__(self, other):
+        # Raise error if attributes are not the same across Operators.
+        if (self.n_spinorbitals is not None and other.n_spinorbitals is not None)\
+                or (self.n_electrons is not None and other.n_electrons is not None)\
+                or (self.spin is not None and other.spin is not None):
+
+            if self.n_spinorbitals != other.n_spinorbitals:
+                raise RuntimeError("n_spinorbitals must be the same for all FermionOperators.")
+            elif self.n_electrons != other.n_electrons:
+                raise RuntimeError("n_electrons must be the same for all FermionOperators.")
+            elif self.spin != other.spin:
+                raise RuntimeError("Spin must be the same for all FermionOperators.")
+
+        if isinstance(other, FermionOperator):
+            return super(FermionOperator, self).__imul__(other)
+        elif isinstance(other, of.FermionOperator):
+            if self.n_spinorbitals is not None or self.n_electrons is not None or self.spin is not 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 __rmul__(self, other):
+        return self.__imul__(other)
+
     def __iadd__(self, other):
         # Raise error if attributes are not the same across Operators.
         if (self.n_spinorbitals is not None and other.n_spinorbitals is not None)\
@@ -53,7 +84,7 @@ def __iadd__(self, other):
 
         if isinstance(other, FermionOperator):
             return super(FermionOperator, self).__iadd__(other)
-        elif isinstance(other, openfermion.FermionOperator):
+        elif isinstance(other, of.FermionOperator):
             if self.n_spinorbitals is not None or self.n_electrons is not None or self.spin is not None:
                 raise RuntimeError("openfermion FermionOperator did not define a necessary attribute")
             else:
@@ -67,9 +98,18 @@ def __add__(self, 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 self.__isub__(other)
+
     def __eq__(self, other):
         # Additional checks for == operator.
-        if isinstance(other, openfermion.FermionOperator):
+        if isinstance(other, of.FermionOperator):
             return super(FermionOperator, self).__eq__(other)
         elif (self.n_spinorbitals != other.n_spinorbitals) or (self.n_electrons != other.n_electrons) or (self.spin != other.spin):
             return False
@@ -88,7 +128,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)
@@ -120,13 +160,13 @@ def get_coeffs(self, coeff_threshold=1e-8):
 
     def to_openfermion(self):
         """Converts Tangelo FermionOperator to openfermion"""
-        ferm_op = openfermion.FermionOperator()
+        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.
     """
@@ -261,7 +301,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()
diff --git a/tangelo/toolboxes/operators/tests/test_operators.py b/tangelo/toolboxes/operators/tests/test_operators.py
index 62e2d3411..648739e1b 100644
--- a/tangelo/toolboxes/operators/tests/test_operators.py
+++ b/tangelo/toolboxes/operators/tests/test_operators.py
@@ -76,7 +76,7 @@ def test_get_coeff(self):
 
     def test_add_FermionOperator(self):
         # addition between two compatible tangelo FermionOperator
-        fop = FermionOperator("0^ 0", 2.) + FermionOperator("0^ 1", 3.)
+        FermionOperator("0^ 0", 2.) + FermionOperator("0^ 1", 3.)
 
         # addition between two incompatible tangelo FermionOperator
         fop_1 = FermionOperator("0^ 0", 2.)
@@ -94,10 +94,51 @@ def test_add_FermionOperator(self):
         fop_2 = of.FermionOperator("0^ 0", 2.) + FermionOperator("0^ 1", 3.)
         self.assertEqual(fop_1, fop_2)
 
+        # Test in-place addition
         fop = FermionOperator("0^ 0", 2.)
         fop += FermionOperator("0^ 1", 3.)
         self.assertEqual(fop, fop_1)
 
+    def test_mul_FermionOperator(self):
+        # Test in-place multiplication
+        fop_1 = FermionOperator("0^ 0", 2.)
+        fop_1 *= of.FermionOperator("0^ 1", 3.)
+        # Test multiplication
+        fop_2 = FermionOperator("0^ 0", 2.) * of.FermionOperator("0^ 1", 3.)
+        self.assertEqual(fop_1, fop_2)
+
+        # Test reverse multiplication
+        fop_3 = of.FermionOperator("0^ 1", 3.) * FermionOperator("0^ 0", 2.)
+        self.assertEqual(fop_2, fop_3)
+
+        # Test multiplication by number
+        fop_4 = 6. * FermionOperator("0^ 0 0^ 1", 1.)
+        self.assertEqual(fop_3, fop_4)
+
+    def test_sub_FermionOperator(self):
+        # Test in-place subtraction
+        fop_1 = FermionOperator("0^ 0", 2.)
+        fop_1 -= of.FermionOperator("0^ 1", 3.)
+        # Test subtraction
+        fop_2 = FermionOperator("0^ 0", 2.) - of.FermionOperator("0^ 1", 3.)
+        self.assertEqual(fop_1, fop_2)
+
+        # Test reverse subtraction
+        fop_3 = of.FermionOperator("0^ 1", 3.) - FermionOperator("0^ 0", 2.)
+        self.assertEqual(fop_2, fop_3)
+
+    def test_div_FermionOperator(self):
+        # Test in-place division
+        fop_1 = FermionOperator("0^ 0", 2.)
+        fop_1 /= 2.
+        # Test division
+        fop_2 = FermionOperator("0^ 0", 2.) / 2.
+        self.assertEqual(fop_1, fop_2)
+
+        # Test error for division by operator
+        with self.assertRaises(TypeError):
+            fop_1 / fop_2
+
 
 class QubitHamiltonianTest(unittest.TestCase):
 
@@ -128,8 +169,7 @@ def test_to_qubit_operator(self):
 
         qubit_ham = 1. * QubitHamiltonian("X0 Y1 Z2", mapping="JW",  up_then_down=True)
 
-        self.assertEqual(qubit_ham.to_qubitoperator(),
-            QubitOperator(term="X0 Y1 Z2", coefficient=1.))
+        self.assertEqual(qubit_ham.to_qubitoperator(), QubitOperator(term="X0 Y1 Z2", coefficient=1.))
 
     def test_get_operators(self):
         """Test get_operators methods, defined in QubitOperator class."""

From 3ae52d83a6556e4711a7b9c19820adaa8a4bfc81 Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Wed, 2 Nov 2022 16:41:49 +0100
Subject: [PATCH 8/9] Changes after Valentin's review

---
 .../helpers/circuits/measurement_basis.py     |  2 +-
 .../toolboxes/molecular_computation/rdms.py   | 25 ++++----
 .../molecular_computation/tests/test_rdms.py  | 46 +++++++------
 tangelo/toolboxes/operators/operators.py      | 64 ++++++++-----------
 .../operators/tests/test_operators.py         | 29 ++++++---
 5 files changed, 81 insertions(+), 85 deletions(-)

diff --git a/tangelo/linq/helpers/circuits/measurement_basis.py b/tangelo/linq/helpers/circuits/measurement_basis.py
index 24657ae7e..fd11b6857 100644
--- a/tangelo/linq/helpers/circuits/measurement_basis.py
+++ b/tangelo/linq/helpers/circuits/measurement_basis.py
@@ -41,7 +41,7 @@ def measurement_basis_gates(term):
     return gates
 
 
-def filter_bases(op, basis_list):
+def filter_measurement_bases(op, basis_list):
     """Return a list of measurement bases compatible with the given operator.
 
     Args:
diff --git a/tangelo/toolboxes/molecular_computation/rdms.py b/tangelo/toolboxes/molecular_computation/rdms.py
index 068b2f98c..e9a02ceb5 100644
--- a/tangelo/toolboxes/molecular_computation/rdms.py
+++ b/tangelo/toolboxes/molecular_computation/rdms.py
@@ -15,9 +15,10 @@
 """Module containing functions to manipulate 1- and 2-RDMs."""
 
 import itertools as it
+
 import numpy as np
 
-from tangelo.linq.helpers import filter_bases, pauli_string_to_of
+from tangelo.linq.helpers import pauli_string_to_of, filter_measurement_bases
 from tangelo.toolboxes.operators import FermionOperator
 from tangelo.toolboxes.measurements import ClassicalShadow
 from tangelo.toolboxes.post_processing import Histogram, aggregate_histograms
@@ -74,15 +75,13 @@ def compute_rdms(ferm_ham, mapping, up_then_down, exp_vals=None, exp_data=None,
     onerdm = np.zeros((ferm_ham.n_spinorbitals,) * 2, dtype=complex)
     twordm = np.zeros((ferm_ham.n_spinorbitals,) * 4, dtype=complex)
 
-    # Check if the data dict is expectation values or raw frequency data in the {basis: hist} format
-    if isinstance(exp_vals, dict) and (exp_data is None) and (shadow is None):
-        exp_vals = {pauli_string_to_of(term): exp_val for term, exp_val in exp_data.items()}
-    elif (exp_vals is None) and isinstance(exp_data, dict) and (shadow is None):
-        exp_vals = {}
-    elif (exp_vals is None) and (exp_data is None) and isinstance(shadow, ClassicalShadow):
-        pass
-    else:
-        raise RuntimeError("Arguments exp_vals, exp_data and shadow are mutually exclusive. Provide only one of them.")
+    # 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)
 
@@ -114,12 +113,12 @@ def compute_rdms(ferm_ham, mapping, up_then_down, exp_vals=None, exp_data=None,
             for qterm, coeff in qubit_term.terms.items():
                 if coeff.real != 0:
 
-                    try:
+                    if qterm in exp_vals:
                         exp_val = exp_vals[qterm] if qterm else 1.
-                    except KeyError:
+                    else:
                         ps = pauli_of_to_string(qterm, n_qubits)
                         hist = aggregate_histograms(*[Histogram(exp_data[basis])
-                                                      for basis in filter_bases(ps, list(exp_data.keys()))])
+                                                      for basis in filter_measurement_bases(ps, list(exp_data.keys()))])
                         exp_val = hist.get_expectation_value(qterm, 1.)
                         exp_vals[qterm] = exp_val
 
diff --git a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
index 982e7af7b..04ef90926 100644
--- a/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
+++ b/tangelo/toolboxes/molecular_computation/tests/test_rdms.py
@@ -15,9 +15,9 @@
 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
@@ -35,25 +35,11 @@
 ferm_op.spin = 0
 
 exp_data = json.load(open(pwd_this_test + "/data/H2_raw_exact.dat", "r"))
-n_shots = 10000
-
-# Construct ClassicalShadow
-bitstrings = []
-unitaries = []
-
-for b, hist in exp_data.items():
-    for s, f in hist.items():
-        factor = round(f*n_shots)
-        bitstrings.extend([s]*factor)
-        unitaries.extend([b]*factor)
-
-cs_data = RandomizedClassicalShadow(unitaries=unitaries, bitstrings=bitstrings)
 
-
-rdm1 = np.array([[1.97454854+0.j, 0.+0.j],
+rdm1ssr = np.array([[1.97454854+0.j, 0.+0.j],
                  [0.+0.j, 0.02545146+0.j]])
 
-rdm2 = np.array(
+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],
@@ -68,23 +54,35 @@ class RDMsUtilitiesTest(unittest.TestCase):
 
     def test_energy_from_rdms(self):
         """Compute energy using known spin-summed 1RDM and 2RDM"""
-        e_rdms = energy_from_rdms(ferm_op, rdm1, rdm2)
+        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"""
-        rdm1r, rdm2r, rdm1ssr, rdm2ssr = compute_rdms(ferm_op, "scbk", True, exp_data=exp_data)
+        rdm1, rdm2, rdm1ss, rdm2ss = compute_rdms(ferm_op, "scbk", True, exp_data=exp_data)
 
-        assert_allclose(rdm1, rdm1ssr, rtol=1e-5)
-        assert_allclose(rdm2, rdm2ssr, rtol=1e-5)
+        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"""
-        rdm1r, rdm2r, rdm1ssr, rdm2ssr = compute_rdms(ferm_op, "scbk", True, shadow=cs_data, k=5)
+        # 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(rdm1, rdm1ssr, atol=0.05)
-        assert_allclose(rdm2, rdm2ssr, atol=0.05)
+        assert_allclose(rdm1ssr, rdm1ss, atol=0.05)
+        assert_allclose(rdm2ssr, rdm2ss, atol=0.05)
 
 
 if __name__ == "__main__":
diff --git a/tangelo/toolboxes/operators/operators.py b/tangelo/toolboxes/operators/operators.py
index 783656bf7..7376c504b 100644
--- a/tangelo/toolboxes/operators/operators.py
+++ b/tangelo/toolboxes/operators/operators.py
@@ -20,15 +20,12 @@
 from collections import OrderedDict
 
 import numpy as np
-from openfermion.ops.operators.symbolic_operator import COEFFICIENT_TYPES
 from scipy.special import comb
-
-# Later on, if needed, we can extract the code for the operators themselves to remove the dependencies and customize
 import openfermion as of
 
 
 class FermionOperator(of.FermionOperator):
-    """Add n_spinorbitals, n_electrons, and spin parameters to the openfermion.FermionOperator class
+    """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):
@@ -39,27 +36,21 @@ def __init__(self, term=None, coefficient=1., n_spinorbitals=None, n_electrons=N
         self.spin = spin
 
     def __imul__(self, other):
-        # Raise error if attributes are not the same across Operators.
-        if (self.n_spinorbitals is not None and other.n_spinorbitals is not None)\
-                or (self.n_electrons is not None and other.n_electrons is not None)\
-                or (self.spin is not None and other.spin is not None):
-
-            if self.n_spinorbitals != other.n_spinorbitals:
-                raise RuntimeError("n_spinorbitals must be the same for all FermionOperators.")
-            elif self.n_electrons != other.n_electrons:
-                raise RuntimeError("n_electrons must be the same for all FermionOperators.")
-            elif self.spin != other.spin:
-                raise RuntimeError("Spin must be the same for all FermionOperators.")
-
         if isinstance(other, FermionOperator):
-            return super(FermionOperator, self).__imul__(other)
+            # 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 is not None or self.n_electrons is not None or self.spin is not None:
+            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)
 
@@ -70,28 +61,24 @@ def __rmul__(self, other):
         return self.__imul__(other)
 
     def __iadd__(self, other):
-        # Raise error if attributes are not the same across Operators.
-        if (self.n_spinorbitals is not None and other.n_spinorbitals is not None)\
-                or (self.n_electrons is not None and other.n_electrons is not None)\
-                or (self.spin is not None and other.spin is not None):
-
-            if self.n_spinorbitals != other.n_spinorbitals:
-                raise RuntimeError("n_spinorbitals must be the same for all FermionOperators.")
-            elif self.n_electrons != other.n_electrons:
-                raise RuntimeError("n_electrons must be the same for all FermionOperators.")
-            elif self.spin != other.spin:
-                raise RuntimeError("Spin must be the same for all FermionOperators.")
-
         if isinstance(other, FermionOperator):
-            return super(FermionOperator, self).__iadd__(other)
+            # 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 is not None or self.n_electrons is not None or self.spin is not None:
+            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)
 
@@ -105,16 +92,17 @@ def __sub__(self, other):
         return self.__isub__(other)
 
     def __rsub__(self, other):
-        return self.__isub__(other)
+        return -1 * self.__isub__(other)
 
     def __eq__(self, other):
         # Additional checks for == operator.
-        if isinstance(other, of.FermionOperator):
+        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)
-        elif (self.n_spinorbitals != other.n_spinorbitals) or (self.n_electrons != other.n_electrons) or (self.spin != other.spin):
-            return False
-
-        return super(FermionOperator, self).__eq__(other)
 
     def get_coeffs(self, coeff_threshold=1e-8):
         """Method to get the coefficient tensors from a fermion operator.
diff --git a/tangelo/toolboxes/operators/tests/test_operators.py b/tangelo/toolboxes/operators/tests/test_operators.py
index 648739e1b..611fd863d 100644
--- a/tangelo/toolboxes/operators/tests/test_operators.py
+++ b/tangelo/toolboxes/operators/tests/test_operators.py
@@ -74,15 +74,22 @@ def test_get_coeff(self):
         np.testing.assert_array_almost_equal(ref_one_body, one_body)
         np.testing.assert_array_almost_equal(ref_two_body, two_body)
 
-    def test_add_FermionOperator(self):
+    def test_eq(self):
+        fop_1 = FermionOperator("0^ 0", 2., spin=1)
+        fop_2 = of.FermionOperator("0^ 0", 2.)
+        fop_3 = FermionOperator("0^ 0", 2., spin=0)
+        fop_4 = FermionOperator("0^ 0", 1., spin=0)
+        self.assertEqual(fop_1, fop_2)
+        self.assertNotEqual(fop_1, fop_3)
+        self.assertNotEqual(fop_3, fop_4)
+
+    def test_add(self):
         # addition between two compatible tangelo FermionOperator
         FermionOperator("0^ 0", 2.) + FermionOperator("0^ 1", 3.)
 
         # addition between two incompatible tangelo FermionOperator
-        fop_1 = FermionOperator("0^ 0", 2.)
-        fop_2 = FermionOperator("0^ 1", 3.)
-        fop_1.spin = 1
-        fop_2.spin = 0
+        fop_1 = FermionOperator("0^ 0", 2., spin=1)
+        fop_2 = FermionOperator("0^ 1", 3., spin=0)
         with self.assertRaises(RuntimeError):
             fop_1 + fop_2
 
@@ -99,7 +106,11 @@ def test_add_FermionOperator(self):
         fop += FermionOperator("0^ 1", 3.)
         self.assertEqual(fop, fop_1)
 
-    def test_mul_FermionOperator(self):
+        # Test addition with non-compatible type
+        with self.assertRaises(RuntimeError):
+            fop + 1
+
+    def test_mul(self):
         # Test in-place multiplication
         fop_1 = FermionOperator("0^ 0", 2.)
         fop_1 *= of.FermionOperator("0^ 1", 3.)
@@ -115,7 +126,7 @@ def test_mul_FermionOperator(self):
         fop_4 = 6. * FermionOperator("0^ 0 0^ 1", 1.)
         self.assertEqual(fop_3, fop_4)
 
-    def test_sub_FermionOperator(self):
+    def test_sub(self):
         # Test in-place subtraction
         fop_1 = FermionOperator("0^ 0", 2.)
         fop_1 -= of.FermionOperator("0^ 1", 3.)
@@ -125,9 +136,9 @@ def test_sub_FermionOperator(self):
 
         # Test reverse subtraction
         fop_3 = of.FermionOperator("0^ 1", 3.) - FermionOperator("0^ 0", 2.)
-        self.assertEqual(fop_2, fop_3)
+        self.assertEqual(fop_2, -1*fop_3)
 
-    def test_div_FermionOperator(self):
+    def test_div(self):
         # Test in-place division
         fop_1 = FermionOperator("0^ 0", 2.)
         fop_1 /= 2.

From a3e0b31e82d65a003a3a5b5b00f51d42cc7024bf Mon Sep 17 00:00:00 2001
From: Krzysztof Bieniasz <krzysztof.bieniasz@1qbit.com>
Date: Tue, 8 Nov 2022 23:23:27 +0100
Subject: [PATCH 9/9] Resolve FermionOperator magic methods __rmul__ and __eq__

---
 .../helpers/circuits/measurement_basis.py     |  2 +-
 .../toolboxes/molecular_computation/rdms.py   | 24 ++++++++++++-------
 tangelo/toolboxes/operators/operators.py      |  3 ---
 .../operators/tests/test_operators.py         |  2 +-
 4 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/tangelo/linq/helpers/circuits/measurement_basis.py b/tangelo/linq/helpers/circuits/measurement_basis.py
index fd11b6857..826f125ce 100644
--- a/tangelo/linq/helpers/circuits/measurement_basis.py
+++ b/tangelo/linq/helpers/circuits/measurement_basis.py
@@ -41,7 +41,7 @@ def measurement_basis_gates(term):
     return gates
 
 
-def filter_measurement_bases(op, basis_list):
+def get_compatible_bases(op, basis_list):
     """Return a list of measurement bases compatible with the given operator.
 
     Args:
diff --git a/tangelo/toolboxes/molecular_computation/rdms.py b/tangelo/toolboxes/molecular_computation/rdms.py
index e9a02ceb5..5c462dbe6 100644
--- a/tangelo/toolboxes/molecular_computation/rdms.py
+++ b/tangelo/toolboxes/molecular_computation/rdms.py
@@ -18,7 +18,7 @@
 
 import numpy as np
 
-from tangelo.linq.helpers import pauli_string_to_of, filter_measurement_bases
+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
@@ -51,19 +51,21 @@ def matricize_2rdm(two_rdm, n_orbitals):
 
 def compute_rdms(ferm_ham, mapping, up_then_down, exp_vals=None, exp_data=None, shadow=None, return_spinsum=True, **eval_args):
     """
-    Return the 1- and 2-RDM and their spin-summed versions
-    using a Molecule object and frequency data either in the form of a
+    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): Dictionary of Pauli word expectation values
-        exp_data (dict): Dictionary of {basis: histogram} where basis is the measurement basis
+        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): A classical shadow object
-        return_spinsum (bool): If True, return also the spin-summed RDMs
+        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:
@@ -117,8 +119,12 @@ def compute_rdms(ferm_ham, mapping, up_then_down, exp_vals=None, exp_data=None,
                         exp_val = exp_vals[qterm] if qterm else 1.
                     else:
                         ps = pauli_of_to_string(qterm, n_qubits)
-                        hist = aggregate_histograms(*[Histogram(exp_data[basis])
-                                                      for basis in filter_measurement_bases(ps, list(exp_data.keys()))])
+                        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
 
diff --git a/tangelo/toolboxes/operators/operators.py b/tangelo/toolboxes/operators/operators.py
index 7376c504b..f6f4e99b2 100644
--- a/tangelo/toolboxes/operators/operators.py
+++ b/tangelo/toolboxes/operators/operators.py
@@ -57,9 +57,6 @@ def __imul__(self, other):
     def __mul__(self, other):
         return self.__imul__(other)
 
-    def __rmul__(self, other):
-        return self.__imul__(other)
-
     def __iadd__(self, other):
         if isinstance(other, FermionOperator):
             # Raise error if attributes are not the same across Operators.
diff --git a/tangelo/toolboxes/operators/tests/test_operators.py b/tangelo/toolboxes/operators/tests/test_operators.py
index 611fd863d..b112b7af9 100644
--- a/tangelo/toolboxes/operators/tests/test_operators.py
+++ b/tangelo/toolboxes/operators/tests/test_operators.py
@@ -119,7 +119,7 @@ def test_mul(self):
         self.assertEqual(fop_1, fop_2)
 
         # Test reverse multiplication
-        fop_3 = of.FermionOperator("0^ 1", 3.) * FermionOperator("0^ 0", 2.)
+        fop_3 = of.FermionOperator("0^ 0", 2.) * FermionOperator("0^ 1", 3.)
         self.assertEqual(fop_2, fop_3)
 
         # Test multiplication by number