From 8e400fd7d56cf5d070bc3c5d2a31fa5f9e998ac4 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Sun, 11 Sep 2022 23:24:28 -0500
Subject: [PATCH 01/24] adding noise simulator

---
 .../noise_simulator/ComparisonSimulator.py    | 246 ++++++++++++++++++
 qtensor/noise_simulator/NoiseChannels.py      | 112 ++++++++
 qtensor/noise_simulator/NoiseModel.py         |  73 ++++++
 .../NoiseSimComparisonResult.py               | 148 +++++++++++
 qtensor/noise_simulator/NoiseSimulator.py     | 116 +++++++++
 qtensor/noise_simulator/QAOANoiseSimulator.py | 127 +++++++++
 qtensor/noise_simulator/__init__.py           |   3 +
 qtensor/noise_simulator/comparison.py         |  64 +++++
 qtensor/noise_simulator/comparison_mpi.py     |  65 +++++
 qtensor/noise_simulator/helper_functions.py   |  83 ++++++
 qtensor/noise_simulator/noisy_qaoa.py         |  21 ++
 11 files changed, 1058 insertions(+)
 create mode 100644 qtensor/noise_simulator/ComparisonSimulator.py
 create mode 100644 qtensor/noise_simulator/NoiseChannels.py
 create mode 100644 qtensor/noise_simulator/NoiseModel.py
 create mode 100644 qtensor/noise_simulator/NoiseSimComparisonResult.py
 create mode 100644 qtensor/noise_simulator/NoiseSimulator.py
 create mode 100644 qtensor/noise_simulator/QAOANoiseSimulator.py
 create mode 100644 qtensor/noise_simulator/__init__.py
 create mode 100644 qtensor/noise_simulator/comparison.py
 create mode 100644 qtensor/noise_simulator/comparison_mpi.py
 create mode 100644 qtensor/noise_simulator/helper_functions.py
 create mode 100644 qtensor/noise_simulator/noisy_qaoa.py

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
new file mode 100644
index 00000000..8adc44bd
--- /dev/null
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -0,0 +1,246 @@
+from NoiseChannels import *
+from NoiseSimulator import * 
+from NoiseModel import *
+from NoiseSimComparisonResult import *
+from helper_functions import *
+from qiskit import execute, Aer
+from qtensor.tests.test_composers import *
+from qtensor import QiskitQAOAComposer
+from qtree.operators import from_qiskit_circuit
+from qtensor import tools
+from itertools import repeat
+
+class ComparisonSimulator:
+    def __init__(self):
+        pass
+
+class QAOAComparisonSimulator(ComparisonSimulator):
+    def __init__(self, n: int, p: int, d: int, noise_model_qiskit: noise.NoiseModel, noise_model_qtensor: NoiseModel, num_circs_list: list):
+        self.n = n
+        self.p = p
+        self.d = d
+        self.noise_model_qiskit = noise_model_qiskit
+        self.noise_model_qtensor = noise_model_qtensor
+        self.num_circs_list = num_circs_list
+        self.num_circs_simulated = []
+        self.results = []
+
+        # We can save on simulation time by not recomputing previous work. This is done by 
+        # using the results of a previous, smaller ensemble simulation: 
+        # i.e. if we ran 1,000 circuits in the previous simulation, and now we want to run 10,000 
+        # circuits using the exact same circuit and parameters, we actually only need to run
+        # 9,000 circuits, add the previous results, and renormalize. 
+        # The downside of this approach is that it requires the storage of an extra state 
+        # in memory. Therefore it should only be done for simulations that do not require 
+        # all of the available memory
+        self.recompute_previous_ensemble: bool
+
+        self._check_params()
+        self.num_circs_list.sort()
+
+    def qtensor_qiskit_noisy_qaoa(self, recompute_previous_ensemble: bool = False, print_stats: bool = False):
+        self.recompute_previous_ensemble = recompute_previous_ensemble
+        # Prepare circuits, simulator
+        G, gamma, beta = get_qaoa_params(n = self.n, p = self.p, d = self.d)
+        self._get_circuits(G, gamma, beta)
+        noise_sim = NoiseSimulator(self.noise_model_qtensor)
+
+        # Run simulation
+        for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
+            result = NoiseSimComparisonResult(self.qiskit_circ, self.qtensor_circ, self.noise_model_qiskit, 
+                self.noise_model_qtensor, self.n, self.p, self.d)
+
+            if i == 0 or recompute_previous_ensemble == True:
+                self.num_circs_simulated.append(num_circs)
+                noise_sim.simulate_batch_ensemble(self.qtensor_circ, num_circs, self.num_qubits)
+                #print("num_circs: {}".format(num_circs))
+                qtensor_probs = noise_sim.normalized_ensemble_probs 
+            else: 
+                actual_num_circs = num_circs - self.num_circs_list[i - 1]
+                self.num_circs_simulated.append(actual_num_circs)
+                #print("num_circs: {}, actual_num_circs: {}".format(num_circs, actual_num_circs_to_sim))
+                noise_sim.simulate_batch_ensemble(self.qtensor_circ, actual_num_circs, self.num_qubits)
+                prev_qtensor_probs = self.prev_probs
+                curr_qtensor_probs = noise_sim.normalized_ensemble_probs
+                qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
+            
+            if recompute_previous_ensemble == False:
+                self.prev_probs = qtensor_probs
+
+            qtensor_sim_time = noise_sim.time_taken
+            self.simulate_qiskit_density_matrix(self.qiskit_circ, self.noise_model_qiskit)    
+
+            # Save results
+            result.save_result(self.qiskit_probs, qtensor_probs, num_circs, 
+                self.num_circs_simulated[i], G, gamma, beta, qtensor_sim_time, self.qiskit_sim_time)
+            self.results.append(result.data)
+            if print_stats:
+                result.print_result()
+
+    def _mpi_parallel_unit(self, args):
+        noise_sim, qtensor_circ, num_circs, num_qubits = args
+        #print("this workers num_circs: ", num_circs)
+        noise_sim.simulate_batch_ensemble(qtensor_circ, num_circs, num_qubits)
+        fraction_of_qtensor_probs = noise_sim.normalized_ensemble_probs
+        return fraction_of_qtensor_probs
+
+    def qtensor_qiskit_noisy_qaoa_mpi(self, num_nodes: int, num_jobs_per_node: int, recompute_previous_ensemble: bool = False, print_stats: bool = True, pbar: bool = True):
+        self.num_nodes = num_nodes 
+        self.num_jobs_per_node = num_jobs_per_node
+        self.recompute_previous_ensemble = recompute_previous_ensemble
+
+        # Prepare circuit, simulator, and area to save results, 
+        G, gamma, beta = get_qaoa_params(n = self.n, p = self.p, d = self.d)
+        self._get_circuits(G, gamma, beta)
+        self.noise_sim = NoiseSimulator(self.noise_model_qtensor)
+
+        for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
+            result = NoiseSimComparisonResult(self.qiskit_circ, self.qtensor_circ, self.noise_model_qiskit, 
+                self.noise_model_qtensor, self.n, self.p, self.d)
+            self._get_args(i)
+            qtensor_probs_list = tools.mpi.mpi_map(self._mpi_parallel_unit, self._arggen, pbar=pbar, total=num_nodes)
+            self._check_correct_num_circs_simulated(i)
+            if qtensor_probs_list:
+                if i == 0 or recompute_previous_ensemble == True: 
+                    qtensor_probs = sum(qtensor_probs_list) / self._total_jobs
+                else:
+                    #prev_qtensor_probs = self.results[i - 1].qtensor_probs
+                    prev_qtensor_probs = self.prev_probs
+                    curr_qtensor_probs = sum(qtensor_probs_list) / self._total_jobs
+                    qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
+                qtensor_sim_time = self.noise_sim.time_taken
+                if recompute_previous_ensemble == False:
+                    self.prev_probs = qtensor_probs
+
+                self.simulate_qiskit_density_matrix(self.qiskit_circ, self.noise_model_qiskit)
+
+                # Save results 
+                result.save_result(self.qiskit_probs, qtensor_probs, num_circs,
+                    self.num_circs_simulated[i], G, gamma, beta, qtensor_sim_time, self.qiskit_sim_time)
+                self.results.append(result.data) 
+                if print_stats:
+                    tools.mpi.print_stats()
+                    result.print_result()
+
+    def qtensor_qiskit_noisy_qaoa_density(self, recompute_previous_ensemble: bool = True):
+        self.recompute_previous_ensemble = recompute_previous_ensemble
+        # Prepare circuits, simulator, and area to save results
+        G, gamma, beta = get_qaoa_params(n = self.n, p = self.p, d = self.d)
+        self._get_circuits(G, gamma, beta)
+        noise_sim = NoiseSimulator(self.noise_model_qtensor)
+        
+        # Simulate
+        for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
+            result = NoiseSimComparisonResult(self.qiskit_circ, self.qtensor_circ, self.noise_model_qiskit, 
+                self.noise_model_qtensor, self.n, self.p, self.d)
+            if i == 0 or recompute_previous_ensemble == True: 
+                self.num_circs_simulated.append(num_circs)
+                noise_sim.simulate_batch_ensemble_density(self.qtensor_circ, num_circs, self.n)
+                qtensor_density_matrix = noise_sim.normalized_ensemble_density_matrix
+            else: 
+                actual_num_circs = num_circs - self.num_circs_list[i - 1]
+                self.num_circs_simulated.append(actual_num_circs)
+                noise_sim.simulate_batch_ensemble_density(self.qtensor_circ, actual_num_circs, self.n)
+                prev_density_matrix = self.prev_qtensor_density_matrix
+                curr_density_matrix = noise_sim.normalized_ensemble_probs
+                qtensor_density_matrix = (curr_density_matrix + prev_density_matrix) / 2
+            qtensor_sim_time = noise_sim.time_taken
+            self.simulate_qiskit_density_matrix(self.qiskit_circ, self.noise_model_qiskit, take_trace = False)
+
+            if recompute_previous_ensemble == False: 
+                self.prev_qtensor_density_matrix = qtensor_density_matrix
+            # Save results
+            result.save_results_density(self.qiskit_density_matrix, qtensor_density_matrix, num_circs, 
+                self.num_circs_simulated[i], G, gamma, beta, qtensor_sim_time, self.qiskit_sim_time)
+            self.results.append(result.data)
+
+
+    # Prepare arguments to be sent to each unit of work 
+    def _get_args(self, i):
+        if i == 0 or self.recompute_previous_ensemble == True:
+            num_circs = self.num_circs_list[i]
+        else:
+            num_circs = self.num_circs_list[i] - self.num_circs_list[i - 1]
+        ## We make sure that each job has a minimum number of circuits. We do this because 
+        ## if there are too few circuits for each unit of work, the overhead from 
+        ## parallelization removes any advantage gained 
+        ## TODO: determine a better minimum number of circuits. Currently 10 is chosen arbitrarily
+        min_circs_per_job = min(10, num_circs)
+        if self.num_nodes * self.num_jobs_per_node > num_circs / min_circs_per_job:
+            num_circs_per_job = min_circs_per_job
+            total_jobs = int(np.ceil(num_circs / num_circs_per_job))
+        else: 
+            total_jobs = self.num_nodes * self.num_jobs_per_node
+            num_circs_per_job = int(np.ceil(num_circs / total_jobs))
+
+        ## We make sure that regardless of how many nodes and jobs per node we have, we always 
+        ## simulate the exact number of circuits in the ensemble specified. 
+        if num_circs_per_job * total_jobs == num_circs:
+            self._arggen = list(zip(repeat(self.noise_sim, total_jobs), repeat(self.qtensor_circ, total_jobs), 
+                repeat(num_circs_per_job, total_jobs), repeat(self.n, total_jobs)))
+            self._total_jobs = total_jobs
+            self.num_circs_simulated.append(num_circs)
+            #print("num_circs: {}, actual num_circs simulated on this iteration: {}".format(self.num_circs_list[i], num_circs))
+        else: 
+            self._arggen = list(zip(repeat(self.noise_sim, total_jobs - 1), repeat(self.qtensor_circ, total_jobs - 1), 
+                repeat(num_circs_per_job, total_jobs - 1), repeat(self.n, total_jobs - 1)))
+            num_circs_in_last_job = num_circs % num_circs_per_job
+            self._arggen.append((self.noise_sim, self.qtensor_circ, num_circs_in_last_job, self.n))
+            self._total_jobs = total_jobs
+            actual_num_circs = (total_jobs - 1) * num_circs_per_job + num_circs_in_last_job
+            assert num_circs == actual_num_circs
+            self.num_circs_simulated.append(actual_num_circs)
+            #print("num_circs: {}, actual num_circs simulated on this iteration: {}, total jobs: {}".format(self.num_circs_list[i], actual_num_circs, total_jobs))
+
+
+    def simulate_qiskit_density_matrix(self, circuit, noise_model_qiskit, take_trace = True):
+        start = time.time_ns() / (10 ** 9)
+        result = execute(circuit, Aer.get_backend('aer_simulator_density_matrix'), shots=1, noise_model=noise_model_qiskit).result()
+        if take_trace:
+            self.qiskit_probs = np.diagonal(result.results[0].data.density_matrix.real)
+            end = time.time_ns() / (10 ** 9)
+            self.qiskit_sim_time = end - start
+        else:
+            self.qiskit_density_matrix = result.results[0].data.density_matrix.real
+            end = time.time_ns() / (10 ** 9)
+            self.qiskit_sim_time = end - start
+
+    def _get_circuits(self, G, gamma, beta):
+        # Create Qiskit circuit 
+        qiskit_com = QiskitQAOAComposer(graph=G, gamma=gamma, beta=beta)
+        qiskit_com.ansatz_state()
+
+        # Convert Qiskit circuit to Qtree circuit
+        self.num_qubits, self.qtensor_circ = from_qiskit_circuit(qiskit_com.circuit)
+        
+        # Finish building remaining portion of Qiskit circuit used only in an Aer simulation 
+        qiskit_com.circuit = qiskit_com.circuit.reverse_bits()
+        qiskit_com.circuit.save_density_matrix()
+        qiskit_com.circuit.measure_all(add_bits = False)
+        self.qiskit_circ = qiskit_com.circuit
+
+    def _check_params(self):
+        if not isinstance(self.n, int):
+            raise Exception("n must an integer.")
+        if not isinstance(self.p, int):
+            raise Exception("p must an integer.")
+        if not isinstance(self.d, int):
+            raise Exception("d must an integer.")
+        if not isinstance(self.noise_model_qiskit, noise.NoiseModel):
+            raise Exception("Qiskit noise model must be of type 'noisel.NoiseModel'")
+        if not isinstance(self.noise_model_qtensor, NoiseModel):
+            raise Exception("QTensor noise model must be of type NoiseModel.NoiseModel")
+        if not isinstance(self.num_circs_list, list):
+            raise Exception("The number of circuits must be given as a list. I.e. if num_circs = 10, the argument should be [10].")
+        if any(not isinstance(y, int) for y in self.num_circs_list):
+            raise Exception("The number of circuits specified must a list of integers.")
+        if (self.n * self.d) % 2 != 0:
+            raise Exception("n * d must be even. This was not satisfied for the given values d: {}, n: {}".format(self.d, self.n))
+        if not 0 <= self.d < self.n:
+            raise Exception("The inequality 0 <= d < n was not satisfied for the given values d: {}, n: {}".format(self.d, self.n))
+
+    def _check_correct_num_circs_simulated(self, i):
+        if i > 0:
+            assert self.num_circs_list[i] == self.num_circs_list[i - 1] + self.num_circs_simulated[i]
+
+
diff --git a/qtensor/noise_simulator/NoiseChannels.py b/qtensor/noise_simulator/NoiseChannels.py
new file mode 100644
index 00000000..78261508
--- /dev/null
+++ b/qtensor/noise_simulator/NoiseChannels.py
@@ -0,0 +1,112 @@
+import itertools as it
+
+# In a depolarizing channel, there is a depolarizing parameter lambda.
+# which is given by λ = 4^n / (4^n - 1) , where n is the number of qubits
+# when we run this function, we pass in a parameter term param, where 
+# 0 <= param <= λ. The probability of an error occuring is related to param and λ by
+# prob_depol_error = param/λ, and the probability of an particular pauli error 
+# from the depolarizing channel is given by param / 4^n 
+
+class DepolarizingChannel: 
+    def __init__(self, param, num_qubits):
+
+        num_terms = 4**num_qubits
+        lam = num_terms / (num_terms - 1)
+        
+        if param < 0:
+            raise ValueError("Error param is too small. It cannot be less than 0")
+        elif param > lam:
+            raise ValueError("Error param  is too large. It cannot be greater than {} when the channel has {} qubits").format(lam, num_qubits)
+
+        self.name = 'depolarizing'
+        self.param = param
+        self.num_qubits = num_qubits
+        self.num_terms = num_terms
+        self.lam = lam
+        self.kraus_ops = []
+        self.qubits = []
+        self._get_channel()
+
+    def _get_channel(self):
+        
+        pauli_error_prob = self.param / self.num_terms
+        identity_prob = 1 - self.param / self.lam
+        depol_channel_probs = [identity_prob] + (self.num_terms - 1) * [pauli_error_prob]
+
+        # it.product creates a set of tuples, where each element is a 
+        # cartensian product of the num_qubit and each of the pauli operators. 
+        # We use list(tup) for tup in beforehand to return a list instead of a tuple
+        depol_channel_paulis = (''.join(tup) for tup in it.product('IXYZ', repeat=self.num_qubits))
+        depol_channel = list(zip(depol_channel_paulis, depol_channel_probs))
+        self.kraus_ops.extend(depol_channel)
+
+    def add_qubits(self, qubits = ['all']):
+        if qubits is not None and not isinstance(qubits, list):
+            raise TypeError("qubits parameter must either be a list or left blank")
+
+        self.qubits.extend(qubits)
+
+    
+class BitFlipChannel:
+    def __init__(self, param):
+        
+        if param < 0:
+            raise ValueError("Error param is too small. It cannot be less than 0")
+        elif param > 1:
+            raise ValueError("Error param is too large. It cannot be greater than 1")
+        
+        self.name = 'bit_flip'
+        self.param = param
+        self.num_qubits = 1
+        self.kraus_ops = []
+        self.qubits = []
+        self._get_channel()
+
+    def _get_channel(self):
+        bit_flip_prob = self.param 
+        identity_prob = 1 - self.param
+        bit_flip_channel = [('I', identity_prob), ('X', bit_flip_prob)]
+        self.kraus_ops.extend(bit_flip_channel)
+    
+    def add_qubits(self, qubits = ['all']):
+        if qubits is not None and not isinstance(qubits, list):
+            raise TypeError("qubits parameter must either be a list or left blank")
+        if len(qubits) > 1:
+            raise Exception("Bit Flip channel is only defined for a single qubit")
+
+        self.qubits.extend(qubits)
+
+class PhaseFlipChannel:
+    def __init__(self, param):
+        
+        if param < 0:
+            raise ValueError("Error param is too small. It cannot be less than 0")
+        elif param > 1:
+            raise ValueError("Error param is too large. It cannot be greater than 1")
+        
+        self.name = 'phase_flip'
+        self.param = param
+        self.num_qubits = 1
+        self.kraus_ops = []
+        self.qubits = []
+        self._get_channel()
+
+    def _get_channel(self):
+        phase_flip_prob = self.param 
+        identity_prob = 1 - self.param
+        phase_flip_channel = [('I', identity_prob), ('Z', phase_flip_prob)]
+        self.kraus_ops.extend(phase_flip_channel)
+        
+
+    def add_qubits(self, qubits = ['all']):
+        if qubits is not None and not isinstance(qubits, list):
+            raise TypeError("qubits parameter must either be a list or left blank")
+        if len(qubits) > 1:
+            raise Exception("Phase Flip channel is only defined for a single qubit")
+
+        self.qubits.extend(qubits)
+
+class KrausOperator:
+    def __init__(self, name, prob):
+        self.name = name
+        self.prob = prob
\ No newline at end of file
diff --git a/qtensor/noise_simulator/NoiseModel.py b/qtensor/noise_simulator/NoiseModel.py
new file mode 100644
index 00000000..d6fb68ca
--- /dev/null
+++ b/qtensor/noise_simulator/NoiseModel.py
@@ -0,0 +1,73 @@
+#from qtensor.noise_simulator import NoiseGate
+# from NoiseGate import *
+# import attr
+
+class NoiseModel:
+    _1qubit_gate = set([
+        'M', 'I', 'H', 'Z', 'T', 'Tdag', 'S', 'Sdag', 'X_1_2', 'Y_1_2', 'W_1_2'
+        'XPhase', 'YPhase', 'ZPhase', 'U', 
+
+    ])
+    _2qubit_gate = set([
+        'cX', 'cY', 'cZ', 'SWAP', 'fSim'
+    ])
+    _3qubit_gate = set([
+        'ccX'
+    ])
+
+    def __init__(self):
+        self.noise_gates = {}
+        # not used currently. Just here in case it needs to be referenced for some later use. Channels is accessed in NoiseGate class
+        self.channels = []
+
+    # applies noise to all qubits that a gate is acting on
+    def add_channel_to_all_qubits(self, channel, gates):
+        for gate in gates:
+            if gate in self.noise_gates:
+                self._check_gate_with_channel(gate, channel)
+                channel.add_qubits()
+                self.noise_gates[gate].add_channel(channel)
+            else: 
+                self._check_gate_with_channel(gate, channel)
+                channel.add_qubits()
+                noise_gate = NoiseGate(gate)
+                noise_gate.add_channel(channel)
+                self.noise_gates[gate] = noise_gate
+                self.channels.append(channel)
+
+    # applies noise to only specific qubits that a gate is acting on. Only matters
+    # for multi qubit gates. Eg: if a gate is 'cx' applied to qubits (0, 1), but the
+    # list qubits = [1], then noise from channel will only be applied to qubit 1
+    def add_channel_to_specific_qubits(self, channel, gates, qubits):
+        for gate in gates:
+            if gate in self.noise_gates:
+                self._check_gate_with_channel(gate, channel)
+                channel.add_qubits(qubits)
+                self.noise_gates[gate].add_channel(channel)
+            else: 
+                self._check_gate_with_channel(gate, channel)
+                channel.add_qubits(qubits)
+                noise_gate = NoiseGate(gate)
+                noise_gate.add_channel(channel)
+                self.noise_gates[gate] = noise_gate
+                self.channels.append(channel)
+                
+    def _check_gate_with_channel(self, gate, channel):
+        if gate in self._1qubit_gate and channel.num_qubits != 1:
+            raise Exception("{} qubit channel ".format(channel.num_qubits) + \
+                            "cannot be applied to 1-qubit gate {}".format(gate))
+        if gate in self._2qubit_gate and channel.num_qubits != 2:
+            raise Exception("{} qubit channel ".format(channel.num_qubits) + \
+                            "cannot be applied to 2-qubit gate {}".format(gate))
+        if gate in self._3qubit_gate and channel.num_qubits != 3:
+            raise Exception("{} qubit channel ".format(channel.num_qubits) + \
+                            "cannot be applied to 3-qubit gate {}".format(gate))
+    ## TODO: Add a warning or error if a bad gate name is given    
+
+class NoiseGate: 
+    def __init__(self, name):
+        self.name = name
+        self.channels = []
+
+    def add_channel(self, channel):
+        self.channels.append(channel)
\ No newline at end of file
diff --git a/qtensor/noise_simulator/NoiseSimComparisonResult.py b/qtensor/noise_simulator/NoiseSimComparisonResult.py
new file mode 100644
index 00000000..9b9d0d80
--- /dev/null
+++ b/qtensor/noise_simulator/NoiseSimComparisonResult.py
@@ -0,0 +1,148 @@
+from helper_functions import *
+
+from datetime import datetime
+import time
+import json
+import jsbeautifier
+import jsonpickle
+
+from qtensor import QiskitQAOAComposer
+import qiskit.providers.aer.noise as noise
+
+from NoiseModel import NoiseModel
+
+class NoiseSimComparisonResult:
+    def __init__(self, qiskit_circ, qtensor_circ, qiskit_noise_model, qtensor_noise_model, n, p, d, name = ""):
+        self.name = name
+        self.experiment_date =  datetime.now().isoformat()
+        self.experiment_start_time = time.time_ns() / (10 ** 9) # convert to floating-point seconds
+
+        self.qiskit_circuit: QiskitQAOAComposer.circuit
+        self.qtensor_circuit: list
+        self.n: int 
+        self.p: int
+        self.d: int
+        self.qiskit_noise_model: noise.NoiseModel
+        self.qtensor_noise_model: NoiseModel
+
+        self.qiskit_circuit = qiskit_circ
+        self.qtensor_circuit = qtensor_circ
+        self.n = n
+        self.p = p
+        self.d = d
+        self.qiskit_noise_model = qiskit_noise_model
+        self.qtensor_noise_model = qtensor_noise_model
+        self.data = {}
+        self._get_noise_model_string()
+        self.pickled_qtensor_noise_model = jsonpickle.encode(self.qtensor_noise_model)
+
+    def save_results_density(self, qiskit_density_matrix, qtensor_density_matrix, num_circs, num_circs_simulated, G, gamma, beta, qtensor_time_total, qiskit_time_total):
+        self.qiskit_density_matrix = qiskit_density_matrix
+        self.qtensor_density_matrix = qtensor_density_matrix
+        self.num_circs = num_circs
+        self.num_circs_simulated = num_circs_simulated
+        self.graph = G
+        self.gamma = gamma
+        self.beta = beta
+        self.qtensor_time_taken = qtensor_time_total
+        self.qiskit_time_taken = qiskit_time_total
+        self.experiment_end_time =  time.time_ns() / (10 ** 9) 
+        self.total_time_taken = self.experiment_end_time - self.experiment_start_time
+        self._calc_similarity()
+
+    def save_result(self, qiskit_probs, qtensor_probs, num_circs, num_circs_simulated, G, gamma, beta, qtensor_time_total, qiskit_time_total):
+        self.qiskit_probs = qiskit_probs
+        self.qtensor_probs = qtensor_probs
+        self.num_circs = num_circs
+        self.num_circs_simulated = num_circs_simulated
+        self.graph = G
+        self.gamma = gamma
+        self.beta = beta
+        self.qtensor_time_taken = qtensor_time_total
+        self.qiskit_time_taken = qiskit_time_total
+        self.experiment_end_time =  time.time_ns() / (10 ** 9) 
+        self.experiment_time_taken = self.experiment_end_time - self.experiment_start_time
+        self._calc_similarity_of_probs()
+        self._to_dict()
+
+    def print_result(self):
+        print("\nExperiment with n = {}, p = {}, d = {}, num_circs = {}, actual_num_circs = {}".format(self.n, self.p, self.d, self.num_circs, self.num_circs_simulated))
+        if type(self.fidelity) is str:
+            print(
+                f"{'Cosine Similarity:':<20}{np.round(self.cos_sim.real, 7):<10}",
+                f"{'Total time taken:':<20}{self.experiment_time_taken:<10}")
+        else:
+            print(
+                f"{'Cosine Similarity:':<20}{np.round(self.cos_sim.real, 7):<10}",
+                f"\n{'Fidelity:':<20}{np.round(self.fidelity.real, 7):<10}",
+                f"{'Total time taken:':<20}{self.experiment_time_taken:<10}")
+
+    def print_noise_model(self):
+        print(self.noise_model_str)
+
+    def save_experiment_to_file(self, outfile):
+        options = jsbeautifier.default_options()
+        options.indent_size = 2
+        with open(outfile, 'a') as f:
+            f.write(jsbeautifier.beautify(json.dumps(self.experiment_dict), options))
+            f.write(', ')
+
+    def _calc_similarity(self):
+        qiskit_probs_root = np.sqrt(np.diagonal(self.qiskit_density_matrix))
+        qtensor_probs_root = np.sqrt(np.diagonal(self.qtensor_density_matrix))
+        
+        self.cos_sim = cosine_similarity(qiskit_probs_root, qtensor_probs_root)
+        self.fidelity = fidelity(self.qiskit_density_matrix.data, self.qtensor_density_matrix)
+
+    def _calc_similarity_of_probs(self):
+        qiskit_probs_root = np.sqrt(self.qiskit_probs)
+        qtensor_probs_root = np.sqrt(self.qtensor_probs)
+
+        self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
+        # we don't need to take the conjugate, as both population density vectors are strictly 
+        # real by the time they have made it here. 
+        self.fidelity = (np.inner(qiskit_probs_root, qtensor_probs_root))**2
+
+    def _to_dict(self):
+        #self.experiment_dict['name'] = self.name
+        self.data['num_qubits'] = self.n
+        self.data['depth'] = self.p
+        self.data['degree'] = self.d
+        self.data['num_circs'] = self.num_circs
+        self.data['actual_num_circs_simulated'] = self.num_circs_simulated
+        self.data['gamma'] = self.gamma
+        self.data['beta'] = self.beta
+        self.data['cosine_similarity'] = self.cos_sim.real
+        self.data['fidelity'] = self.fidelity
+        self.data['noise_model_string'] = self.noise_model_str_condensed
+        self.data['noise_model_pickle'] = self.pickled_qtensor_noise_model
+        self.data['experiment_date'] = self.experiment_date
+        self.data['qtensor_time_taken'] = self.qtensor_time_taken
+        self.data['qiskit_time_taken'] = self.qiskit_time_taken
+        self.data['total_time_taken'] = self.experiment_time_taken
+
+    def _get_noise_model_string(self):
+        gates = self.qtensor_noise_model.noise_gates
+        noise_model_str = 'Noise Model:'
+        noise_model_str_condensed = ''
+        for gate in gates:
+            num_channels = len(self.qtensor_noise_model.noise_gates[gate].channels)
+            if num_channels == 1:
+                noise_model_str += '\nThe gate {} has 1 channel:\n'.format(gate)
+                noise_model_str_condensed += "Gate: {}, Channel: ".format(gate)
+            else:
+                noise_model_str += '\nThe gate {} has {} channels:\n'.format(gate, num_channels)
+                noise_model_str_condensed += "Gate: {}, {} Channels: ".format(gate, num_channels)
+            all_channels_info = ''
+            all_channels_info_short = ''
+            for i in range(num_channels):
+                name = self.qtensor_noise_model.noise_gates[gate].channels[i].name
+                param = self.qtensor_noise_model.noise_gates[gate].channels[i].param 
+                channel_info = 'Channel name: {}, with parameter: {}.\n'.format(name, param)
+                channel_info_short = 'name: {}, param: {}. '.format(name, param)
+                all_channels_info += channel_info
+                all_channels_info_short += channel_info_short
+            noise_model_str += all_channels_info
+            noise_model_str_condensed += all_channels_info_short
+        self.noise_model_str = noise_model_str
+        self.noise_model_str_condensed = noise_model_str_condensed
\ No newline at end of file
diff --git a/qtensor/noise_simulator/NoiseSimulator.py b/qtensor/noise_simulator/NoiseSimulator.py
new file mode 100644
index 00000000..af58aaef
--- /dev/null
+++ b/qtensor/noise_simulator/NoiseSimulator.py
@@ -0,0 +1,116 @@
+from qtensor.Simulate import *
+from qtensor.OpFactory import *
+from qtree import *
+from NoiseModel import *
+import time
+
+class NoiseSimulator(QtreeSimulator):
+    def __init__(self, noise_model, bucket_backend=NumpyBackend(), optimizer=None, max_tw=None):
+        super().__init__(bucket_backend, optimizer, max_tw)
+        if not isinstance(noise_model, NoiseModel):
+            raise ValueError("Error: noise_model value must be of type NoiseModel")
+        self.noise_model = noise_model
+
+    # If all you want is the probabiltiies, and not the amplitudes, this is a better function to call. It does not create the density matrix, 
+    # only a state vector, so it will take up much less memeory 
+    # The vector returned is dim(2^n), where n = number of qubits
+    def simulate_batch_ensemble(self, qc, num_circs, batch_vars=0, peo=None):
+        start = time.time_ns() / (10 ** 9)
+        if num_circs < 0 or not isinstance(num_circs, int):
+            raise Exception("Error: The argument num_circs must be a positive integer")
+        
+        unnormalized_ensemble_probs = [0] * 2**batch_vars
+        for _ in range(num_circs):
+            noisy_state_amps = self.simulate_batch(qc, batch_vars, peo)
+            noisy_state_probs = np.square(np.absolute(noisy_state_amps))
+            unnormalized_ensemble_probs += noisy_state_probs
+        
+        normalized_ensemble_probs = unnormalized_ensemble_probs / num_circs
+        self.normalized_ensemble_probs = normalized_ensemble_probs
+        end = time.time_ns() / (10 ** 9)
+        self.time_taken = end - start
+        #return normalized_ensemble_probs
+
+    # This returns a density matrix, which contains all of the amplitudes of the final state. 
+    # If we want the probabilities, we can take the trace of that matrix
+    # The density matrix returned will be dimension m x m where m = 2^n and n = number of qubits  
+    def simulate_batch_ensemble_density(self, qc, num_circs, batch_vars=0, peo=None):
+        start = time.time_ns() / (10 ** 9)
+        if num_circs < 0 or not isinstance(num_circs, int):
+            raise Exception("Error: The argument num_circs must be a positive integer")
+        
+        unnormalized_ensemble_density_matrix = np.zeros(shape=(2**batch_vars, 2**batch_vars), dtype=complex)
+        for _ in range(num_circs):
+            noisy_state_amps = self.simulate_batch(qc, batch_vars, peo)
+            conj_noisy_state_amps = noisy_state_amps.conj()
+            noisy_state_density_matrix = np.outer(conj_noisy_state_amps, noisy_state_amps)
+            unnormalized_ensemble_density_matrix += noisy_state_density_matrix
+        
+        self.normalized_ensemble_density_matrix = np.divide(unnormalized_ensemble_density_matrix, num_circs)
+        end = time.time_ns() / (10 ** 9)
+        self.time_taken = end - start
+        #normalized_ensemble_density_matrix = np.divide(unnormalized_ensemble_density_matrix, num_circs)
+        #return normalized_ensemble_density_matrix
+    
+    # Simulates and returns only the first amplitude of the ensemble  
+    def simulate_ensemble(self, qc, num_circs):
+        return self.simulate_state(qc, num_circs)
+
+    def simulate_state_ensemble(self, qc, num_circs, peo=None):
+        return self.simulate_batch(qc, num_circs, peo=peo, batch_vars=0)
+
+    #-- Internal helpers
+    def _new_circuit(self, ideal_qc):
+        self.all_gates = []
+        for gate in ideal_qc:
+            self.all_gates.append(gate)
+            if gate.name in self.noise_model.noise_gates:
+                self._apply_channel(gate)
+        
+    def _apply_channel(self, gate):
+        for i in range(len(self.noise_model.noise_gates[gate.name].channels)):
+            error_name = self.noise_model.noise_gates[gate.name].channels[i].name
+            if error_name == 'depolarizing':
+                # The first tuple in the list of Kraus operators is always  II....I, n times, where n is the number of qubits
+                # kraus_ops[0][1] refers to the second element of the first tuple in the list of kraus operators. 
+                # The second element of the tuple contains the probability of that operator 
+                prob_no_noise = self.noise_model.noise_gates[gate.name].channels[i].kraus_ops[0][1]
+                p = np.random.uniform()
+                if p >= prob_no_noise:
+                    # In a depolarizing channel, the probability for any non II...I Kraus operator is always the same, 
+                    # so we can just randomly choose any of the Kraus operators besides the first one in our noise_info list 
+                    randIndex = np.random.randint(1, len(self.noise_model.noise_gates[gate.name].channels[i].kraus_ops))
+                    # Get string of pauli(s), turn them into an iterable object, then add each pauli one by one to the circuit
+                    kraus_op_name = self.noise_model.noise_gates[gate.name].channels[i].kraus_ops[randIndex][0]
+                    paulis = iter(kraus_op_name)
+                    for qubit in gate.qubits:
+                        pauli = next(paulis)
+                        if pauli != 'I':
+                            self.all_gates.append(getattr(qtree.operators, pauli)(qubit))
+                # else: p < prob_no_noise, so no noise is applied 
+            
+            elif error_name == 'bit_flip':
+                if len(gate.qubits) > 1:
+                    raise ValueError("Bit flip noise can only be applied to single qubit gates")
+
+                prob_no_noise = self.noise_model.noise_gates[gate.name].channels[i].kraus_ops[0][1]
+                p = np.random.uniform()
+                if p >= prob_no_noise: 
+                    self.all_gates.append(qtree.operators.X(gate.qubits[0]))
+            
+            elif error_name == 'phase_flip':
+                if len(gate.qubits) > 1:
+                    raise ValueError("Phase flip noise can only be applied to single qubit gates")
+
+                prob_no_noise = self.noise_model.noise_gates[gate.name].channels[i].kraus_ops[0][1]
+                p = np.random.uniform()
+                if p >= prob_no_noise:
+                    self.all_gates.append(qtree.operators.Z(gate.qubits[0]))
+
+
+
+
+
+
+
+
diff --git a/qtensor/noise_simulator/QAOANoiseSimulator.py b/qtensor/noise_simulator/QAOANoiseSimulator.py
new file mode 100644
index 00000000..6abb8066
--- /dev/null
+++ b/qtensor/noise_simulator/QAOANoiseSimulator.py
@@ -0,0 +1,127 @@
+from xml.dom.minicompat import NodeList
+from NoiseChannels import *
+from NoiseSimulator import * 
+from NoiseModel import *
+from helper_functions import *
+from qtensor.tests.test_composers import *
+from qtensor.tools import mpi
+from itertools import repeat
+
+## TODO
+# Save data somewhere (and figure out what exactly to save)
+# potentially refactor comparison simulator to inherit this and/or call this. remove duplicate code
+
+class QAOANoiseSimulator:
+    def __init__(self):
+        pass
+
+    def noisy_qaoa_sim(self, n: int, p: int, d: int, num_circs: int, noise_model: NoiseModel):
+        self._set_params(n=n, p=p, d=d, num_circs=num_circs, noise_model=noise_model)
+        G, gamma, beta = get_qaoa_params(n=self.n, p=self.p, d=self.d)
+        composer = QtreeQAOAComposer(graph=G, gamma=gamma, beta=beta)
+        composer.ansatz_state()
+        self.circuit = composer.circuit
+        noise_sim = NoiseSimulator(self.noise_model)
+        noise_sim.simulate_batch_ensemble(self.circuit, self.num_circs, self.num_qubits)
+        self.probs = noise_sim.normalized_ensemble_probs
+        self.sim_time = noise_sim.time_taken
+
+    def _mpi_parallel_unit(self, args):
+        noise_sim, qtensor_circ, num_circs, num_qubits = args
+        noise_sim.simulate_batch_ensemble(qtensor_circ, num_circs, num_qubits)
+        fraction_of_qtensor_probs = noise_sim.normalized_ensemble_probs
+        return fraction_of_qtensor_probs
+
+    def noisy_qaoa_sim_mpi(self, n: int, p: int, d: int, num_circs: int, noise_model: NoiseModel, 
+            num_nodes: int, num_jobs_per_node: int, print_stats: bool=False, pbar: bool=True):
+            
+        self._set_params(mpi_sim=True, n=n, p=p, d=d, num_circs=num_circs, noise_model=noise_model, 
+            num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node, print_stats=print_stats, pbar=pbar)
+        G, gamma, beta = get_qaoa_params(n=self.n, p=self.p, d=self.d)
+        composer = QtreeQAOAComposer(graph=G, gamma=gamma, beta=beta)
+        composer.ansatz_state()
+        self.circuit = composer.circuit
+        self.noise_sim = NoiseSimulator(self.noise_model)
+        arggen = self._get_args(self.num_circs, self.num_nodes, self.num_jobs_per_node)
+        qaoa_probs_list = mpi.mpi_map(self._mpi_parallel_unit, arggen, pbar=pbar, total=self.num_nodes)
+
+        if qaoa_probs_list:
+            self.probs = sum(qaoa_probs_list) / self.total_jobs
+            print(self.probs)
+            #self.mpi_sim_time = mpi.get_stats()
+            if print_stats:
+                mpi.print_stats()
+
+    def _get_args(self, num_circs, num_nodes, num_jobs_per_node):
+        total_jobs, num_circs_per_job = self._get_total_jobs(num_circs, num_nodes, num_jobs_per_node)
+        """
+        We make sure that regardless of how many nodes and jobs per node we have, 
+        we always simulate the exact number of circuits in the ensemble specified. 
+        """
+        if num_circs_per_job * total_jobs == num_circs:
+            arggen = list(zip(repeat(self.noise_sim, total_jobs), repeat(self.circuit, total_jobs), 
+                repeat(num_circs_per_job, total_jobs), repeat(self.num_qubits, total_jobs)))
+        else: 
+            arggen = list(zip(repeat(self.noise_sim, total_jobs - 1), repeat(self.circuit, total_jobs - 1), 
+                repeat(num_circs_per_job, total_jobs - 1), repeat(self.n, total_jobs - 1)))
+            num_circs_in_last_job = num_circs % num_circs_per_job
+            arggen.append((self.noise_sim, self.circuit, num_circs_in_last_job, self.n))
+            actual_num_circs = (total_jobs - 1) * num_circs_per_job + num_circs_in_last_job
+            assert num_circs == actual_num_circs
+        self.total_jobs = total_jobs
+        return arggen
+
+    def _get_total_jobs(self, num_circs, num_nodes, num_jobs_per_node):
+        """
+        We make sure that each job has a minimum number of circuits. 
+        We do this because if there are too few circuits for each unit of work, 
+        the overhead from parallelization removes any advantage gained
+
+        TODO: determine a better minimum number of circuits. Currently 10 is chosen arbitrarily
+        """
+        min_circs_per_job = min(10, num_circs)
+        if num_nodes * num_jobs_per_node > num_circs / min_circs_per_job:
+            num_circs_per_job = min_circs_per_job
+            total_jobs = int(np.ceil(num_circs / num_circs_per_job))
+        else: 
+            total_jobs = num_nodes * num_jobs_per_node
+            num_circs_per_job = int(np.ceil(num_circs / total_jobs))
+        return total_jobs, num_circs_per_job
+
+    def _set_params(self, mpi_sim = False, **params):
+        self.n = params['n']
+        self.p = params['p']
+        self.d = params['d']
+        self.num_circs = params['num_circs']
+        self.noise_model = params['noise_model']
+        self.num_qubits = params['n']
+        if mpi_sim:
+            self.num_nodes = params['num_nodes']
+            self.num_jobs_per_node = params['num_jobs_per_node']
+        self._check_params(mpi_sim, params)
+
+    def _check_params(self, mpi_sim, params):
+        if not isinstance(self.n, int):
+            raise Exception("n must be an integer.")
+        if not isinstance(self.p, int):
+            raise Exception("p must be an integer.")
+        if not isinstance(self.d, int):
+            raise Exception("d must be an integer.")
+        if not isinstance(self.num_circs, int):
+            raise Exception("num_circs must be an integer")
+        if not isinstance(self.noise_model, NoiseModel):
+            raise Exception("noise_model must be of type NoiseModel.NoiseModel")
+        if (self.n * self.d) % 2 != 0:
+            raise Exception("n * d must be even. This was not satisfied for the given values d: {}, n: {}".format(self.d, self.n))
+        if not 0 <= self.d < self.n:
+            raise Exception("The inequality 0 <= d < n was not satisfied for the given values d: {}, n: {}".format(self.d, self.n))
+        if mpi_sim:
+            if not isinstance(self.num_nodes, int):
+                raise Exception("num_nodes must be an integer")
+            if not isinstance(self.num_jobs_per_node, int):
+                raise Exception("num_jobs_per_node must be an integer")
+            if not isinstance(params['print_stats'], bool):
+                raise Exception("print_stats must be a bool")
+            if not isinstance(params['pbar'], bool):
+                raise Exception("pbar must be a bool")
+        
diff --git a/qtensor/noise_simulator/__init__.py b/qtensor/noise_simulator/__init__.py
new file mode 100644
index 00000000..9d2d0e4a
--- /dev/null
+++ b/qtensor/noise_simulator/__init__.py
@@ -0,0 +1,3 @@
+from . import NoiseChannels
+from . import NoiseGate
+from . import NoiseModel
\ No newline at end of file
diff --git a/qtensor/noise_simulator/comparison.py b/qtensor/noise_simulator/comparison.py
new file mode 100644
index 00000000..882dc3ae
--- /dev/null
+++ b/qtensor/noise_simulator/comparison.py
@@ -0,0 +1,64 @@
+from NoiseChannels import *
+from NoiseSimulator import * 
+from NoiseModel import *
+from NoiseSimComparisonResult import *
+from helper_functions import *
+import qiskit.providers.aer.noise as noise
+from ComparisonSimulator import QAOAComparisonSimulator
+from qtensor.tests.test_composers import *
+
+
+prob_1 = 0.01
+prob_2 = 0.1
+
+# Qiskit Noise Model
+depol_chan_qiskit_1Q = noise.depolarizing_error(prob_1, 1)
+depol_chan_qiskit_2Q = noise.depolarizing_error(prob_2, 2)
+
+noise_model_qiskit = noise.NoiseModel()
+noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_1Q, ['h', 'rx', 'rz'])
+noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_2Q, ['cx'])
+
+# QTensor Noise Model
+depol_chan_qtensor_1Q = DepolarizingChannel(prob_1, 1)
+depol_chan_qtensor_2Q = DepolarizingChannel(prob_2, 2)
+
+noise_model_qtensor = NoiseModel()
+noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['H', 'XPhase', 'ZPhase'])
+noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX'])
+
+min_n = 5
+max_n = 7
+
+min_p = 2
+max_p = 3
+
+min_d = 2
+max_d = 3
+
+num_samples = 1
+num_circs_list = [10, 18, 32, 100, 178, 316]
+results = []
+
+outfile_name = '{}.json'.format(datetime.now().isoformat())
+
+
+for n in range(min_n, max_n):
+    for p in range(min_p, max_p):
+        for d in range(min_d, max_d):
+            # n * d must be even
+            if (n * d) % 2 != 0:
+                if d != max_d:
+                    d += 1
+                else:  
+                    break
+            if not 0 <= d < n:
+                break
+
+            for _  in range(num_samples):
+                print("\n\nnum_circs_list:", num_circs_list)
+                comparison = QAOAComparisonSimulator(n, p, d, noise_model_qiskit, noise_model_qtensor, num_circs_list)
+                comparison.qtensor_qiskit_noisy_qaoa(print_stats = True)
+                results.extend(comparison.results)
+
+save_dict_to_file(results, outfile_name)                 
diff --git a/qtensor/noise_simulator/comparison_mpi.py b/qtensor/noise_simulator/comparison_mpi.py
new file mode 100644
index 00000000..28b46507
--- /dev/null
+++ b/qtensor/noise_simulator/comparison_mpi.py
@@ -0,0 +1,65 @@
+from NoiseChannels import *
+from NoiseSimulator import * 
+from NoiseModel import *
+from NoiseSimComparisonResult import *
+from helper_functions import *
+import qiskit.providers.aer.noise as noise
+from qtensor.tests.test_composers import *
+
+
+### Noise model, simulation, and results of a QAOA algorithmm  ###
+prob_1 = 0.01
+prob_2 = 0.1
+
+# Qiskit Noise Model
+depol_chan_qiskit_1Q = noise.depolarizing_error(prob_1, 1)
+depol_chan_qiskit_2Q = noise.depolarizing_error(prob_2, 2)
+
+noise_model_qiskit = noise.NoiseModel()
+noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_1Q, ['h', 'rx', 'rz'])
+noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_2Q, ['cx'])
+
+# QTensor Noise Model
+depol_chan_qtensor_1Q = DepolarizingChannel(prob_1, 1)
+depol_chan_qtensor_2Q = DepolarizingChannel(prob_2, 2)
+
+noise_model_qtensor = NoiseModel()
+noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['H', 'XPhase', 'ZPhase'])
+noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX'])
+
+min_n = 5
+max_n = 7
+
+min_p = 2
+max_p = 3
+
+min_d = 2
+max_d = 3
+
+num_samples = 3
+num_circs_list = [10, 18, 32, 100, 178, 316]
+num_nodes = 1
+num_jobs_per_node = 3
+
+outfile_name = '{}.json'.format(datetime.now().isoformat())
+results = []
+
+for n in range(min_n, max_n):
+    for p in range(min_p, max_p):
+        for d in range(min_d, max_d):
+            # n * d must be even
+            if (n * d) % 2 != 0:
+                if d != max_d:
+                    d += 1
+                else:  
+                    break
+            if not 0 <= d < n:
+                break
+
+            for _ in range(num_samples):
+                print("\n\nnum_circs_list:", num_circs_list)
+                comparison = QAOAComparisonSimulator(n, p, d, noise_model_qiskit, noise_model_qtensor, num_circs_list)
+                comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node)
+                results.extend(comparison.results)
+
+save_dict_to_file(results, name = outfile_name) 
\ No newline at end of file
diff --git a/qtensor/noise_simulator/helper_functions.py b/qtensor/noise_simulator/helper_functions.py
new file mode 100644
index 00000000..aac630cb
--- /dev/null
+++ b/qtensor/noise_simulator/helper_functions.py
@@ -0,0 +1,83 @@
+from logging import raiseExceptions
+from qtree.operators import *
+
+from collections import OrderedDict
+import numpy as np
+import numpy as np
+from numpy import inner
+from numpy.linalg import norm
+import json
+import jsbeautifier
+import networkx as nx
+from datetime import datetime
+from os import SEEK_SET
+
+def decimal_to_binary(n):
+# converting decimal to binary
+# and removing the prefix(0b)
+    return bin(n).replace("0b", "")
+
+def create_counts_dict(num_qubits, big_endian):
+    # get the binary length of num_qubits
+    length = int(np.ceil(np.log(2**num_qubits + 1)/np.log(2)) - 1)
+    counts = OrderedDict()
+    for i in range(2**num_qubits):
+        # convert to binary and then pad the right side with 0s
+        # length is used so that the length of each key is the same
+        if big_endian == True:
+            key_i = str(decimal_to_binary(i).zfill(length))
+            key_i[::-1]
+            counts[key_i] = 0
+        else:
+            key_i = str(decimal_to_binary(i).zfill(length))
+            counts[key_i] = 0
+    #print(counts)
+    return counts
+
+# big_endian refers to the order that qubits are displayed 
+# Use big_endian = False if you want the qubit ordering that Qiskit uses 
+def attach_qubit_names(probs_list, big_endian = True):
+    num_qubits = int(np.log2(len(probs_list)))
+    probs_dict = create_counts_dict(num_qubits, big_endian)
+    for key, prob in zip(probs_dict, probs_list):
+        probs_dict[key] = prob
+    return probs_dict
+
+def fidelity(A, B):
+    return (np.sum(np.diagonal(np.sqrt(np.sqrt(A)*B*np.sqrt(A)))))**2
+
+def cosine_similarity(A, B):
+    return inner(A, B) / (norm(A)* norm(B)) 
+
+def get_qaoa_params(n, p, d):
+    if (n * d) % 2 != 0:
+        raise ValueError("n * d must be even.")
+    if not 0 <= d < n:
+        raise ValueError("Bad value for d. The inequality 0 <= d < n must be satisfied.")
+
+    print('Circuit params: n: {}, p: {}, d: {}'.format(n, p, d))
+    G = nx.random_regular_graph(d, n)
+    gamma = [np.random.uniform(0, 2*np.pi) for _ in range(p)]
+    beta = [np.random.uniform(0, np.pi) for _ in range(p)]
+    return G, gamma, beta
+
+def save_dict_to_file(dict, name):
+    options = jsbeautifier.default_options()
+    options.indent_size = 2
+    with open(name, 'a') as f:
+        f.write(jsbeautifier.beautify(json.dumps(dict), options))
+
+def get_total_jobs(num_circs: int, num_nodes: int, num_jobs_per_node: int, min_circs_per_job = 10):
+    ## We make sure that each job has a minimum number of circuits. We do this because 
+    ## if there are too few circuits for each unit of work, the overhead from 
+    ## parallelization removes any advantage gained 
+    ## TODO: determine a better minimum number of circuits. Currently 10 is chosen arbitrarily
+    min_circs_per_job = min(min_circs_per_job, num_circs)
+    if num_nodes * num_jobs_per_node > num_circs / min_circs_per_job:
+        num_circs_per_job = min_circs_per_job
+        total_jobs = int(np.ceil(num_circs / num_circs_per_job))
+    else: 
+        total_jobs = num_nodes * num_jobs_per_node
+        num_circs_per_job = int(np.ceil(num_circs / total_jobs))
+    return total_jobs, num_circs_per_job
+
diff --git a/qtensor/noise_simulator/noisy_qaoa.py b/qtensor/noise_simulator/noisy_qaoa.py
new file mode 100644
index 00000000..cd1ddec7
--- /dev/null
+++ b/qtensor/noise_simulator/noisy_qaoa.py
@@ -0,0 +1,21 @@
+from NoiseChannels import *
+from NoiseSimulator import * 
+from NoiseModel import *
+from NoiseSimComparisonResult import *
+from helper_functions import *
+from QAOANoiseSimulator import *
+
+prob_1 = 0.01
+prob_2 = 0.1
+
+# QTensor Noise Model
+depol_chan_1Q = DepolarizingChannel(prob_1, 1)
+depol_chan_2Q = DepolarizingChannel(prob_2, 2)
+
+noise_model = NoiseModel()
+noise_model.add_channel_to_all_qubits(depol_chan_1Q, ['H', 'XPhase', 'ZPhase'])
+noise_model.add_channel_to_all_qubits(depol_chan_2Q, ['cX'])
+
+qaoa_sim = QAOANoiseSimulator()
+qaoa_sim.noisy_qaoa_sim(n = 6, p = 1, d = 1, num_circs = 100, noise_model = noise_model)
+#print(qaoa_sim.probs)
\ No newline at end of file

From 8f07128933541596068f3774e73c39c4d5cc8be8 Mon Sep 17 00:00:00 2001
From: Danil Lykov <lkv97dn@gmail.com>
Date: Sat, 17 Sep 2022 19:33:07 -0600
Subject: [PATCH 02/24] add generic circ sim and review on noise

---
 qtensor/noise_simulator/Review.md             |  25 ++++
 .../comparison_generic_circ.py                |  87 ++++++++++++
 qtensor/noise_simulator/generate_cz_rnd.py    | 129 ++++++++++++++++++
 3 files changed, 241 insertions(+)
 create mode 100644 qtensor/noise_simulator/Review.md
 create mode 100644 qtensor/noise_simulator/comparison_generic_circ.py
 create mode 100644 qtensor/noise_simulator/generate_cz_rnd.py

diff --git a/qtensor/noise_simulator/Review.md b/qtensor/noise_simulator/Review.md
new file mode 100644
index 00000000..b9cf8ed8
--- /dev/null
+++ b/qtensor/noise_simulator/Review.md
@@ -0,0 +1,25 @@
+1. Recompute previous ensemble
+   This feature makes different runs with different K correlated: if our K=100 was particularly good at estimating the result, our K=500 might be biased. It is better to statistically decouple different runs for more reliable results across the runs.
+
+
+2. Naming is good
+   The long names make it easy to visually browse the code and intuitively understand the purpose of it.
+
+3. `from ... import *`
+   This is not a good practice. It makes hard to understand what exactly gets imported from the module. This makes it harder to read the code.
+
+4. NoisySimulator design
+   Why do methods for noisy simulation not return the vectors/density matrices, but instead save the result to `.self`?
+
+   The call stack is following: `NoisySimulator.simulate_batch_ensemble()` -> `QtreeSimulator.simulate_batch()` ->
+   `NoisySimulator._new_circuit()` -> ...
+   This is confusing, since by looking at just NoisySimulator it's hard to understand that `_new_circuit()` is called. This is probably a problem in design of Qtreesimulator.
+   However, it makes `NoisySimulator.simulate_batch()` behave appropriatly: to return a stochastic sample of the noisy circuit, which is a good thing.
+
+   TBD: Make `NoisySimulator` to have a reference to `Simulator` and initialization would be `NoisySimulator(noise_model, simulator)`. Then in `NoisySimulator.simulate_ensemble` calls `NoisySimulator.simulate_batch()`, which in turn explicitly applies channels, then calls `self.simulator.simulate_batch(modified_circ, ...)`. Will have to explicitly write the `simulate_batch` method.  
+
+5. Code formatting
+   Docstrings in python come after method definition and in triple quotes: """Docstring""". See "Google python styleguide".
+
+6. TODO:
+   Noise Channel should contain information on which kraus operator to apply (instead of doing a check by name of channel via `error_name`). In this way, we can allow user to specify custom Kraus Operators.
\ No newline at end of file
diff --git a/qtensor/noise_simulator/comparison_generic_circ.py b/qtensor/noise_simulator/comparison_generic_circ.py
new file mode 100644
index 00000000..74d9ce95
--- /dev/null
+++ b/qtensor/noise_simulator/comparison_generic_circ.py
@@ -0,0 +1,87 @@
+from generate_cz_rnd import get_cz_circ, CZBrickworkComposer, QiskitCZBrickworkComposer
+import numpy as np
+from qtree.operators import from_qiskit_circuit
+from qiskit import execute, Aer
+from qtensor import QiskitQAOAComposer
+import networkx as nx
+import qiskit.providers.aer.noise as noise
+import time
+
+def get_circs(S, d):
+    G = nx.random_regular_graph(3, S**2)
+
+    comp = QiskitCZBrickworkComposer(S)
+    comp.two_qubit_rnd(layers=d)
+
+    ###
+    #gamma = [0.3] * d
+    #beta = [0.3] * d
+    #comp = QiskitQAOAComposer(G, gamma=gamma, beta=beta)
+    #comp.ansatz_state()
+    n, qtensor_circ = from_qiskit_circuit(comp.circuit)
+
+    comp.circuit = comp.circuit.reverse_bits()
+    comp.circuit.save_density_matrix()
+    comp.circuit.measure_all(add_bits = False)
+    return qtensor_circ, comp.circuit
+
+def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = True):
+    start = time.time_ns() / (10 ** 9)
+    result = execute(circuit, Aer.get_backend('aer_simulator_density_matrix'), shots=1, noise_model=noise_model_qiskit).result()
+    if take_trace:
+        qiskit_probs = np.diagonal(result.results[0].data.density_matrix.real)
+        end = time.time_ns() / (10 ** 9)
+        qiskit_sim_time = end - start
+        return qiskit_probs, qiskit_sim_time
+    else:
+        qiskit_density_matrix = result.results[0].data.density_matrix.real
+        end = time.time_ns() / (10 ** 9)
+        qiskit_sim_time = end - start
+        return qiskit_density_matrix, qiskit_sim_time
+
+
+if __name__=="__main__":
+    from NoiseChannels import DepolarizingChannel
+    from NoiseModel import NoiseModel
+    from NoiseSimulator import NoiseSimulator
+
+    prob_1 = 0.01
+    prob_2 = 0.1
+
+    S = 2
+    d = 3
+    qtensor_circ, qiskit_circ = get_circs(S, d)
+    num_qubits = S**2
+    # Qiskit Noise Model
+    depol_chan_qiskit_1Q = noise.depolarizing_error(prob_1, 1)
+    depol_chan_qiskit_2Q = noise.depolarizing_error(prob_2, 2)
+
+    noise_model_qiskit = noise.NoiseModel()
+    noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_1Q, ['x', 'y', 'z'])
+    noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_2Q, ['cx'])
+
+    # QTensor Noise Model
+    depol_chan_qtensor_1Q = DepolarizingChannel(prob_1, 1)
+    depol_chan_qtensor_2Q = DepolarizingChannel(prob_2, 2)
+
+    noise_model_qtensor = NoiseModel()
+    noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['X', 'Y', 'Z'])
+    noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX'])
+
+    noise_sim = NoiseSimulator(noise_model_qtensor)
+
+
+    for num_circs in [100, 300, 1000, 3000]:
+        start = time.time_ns() / (10 ** 9)
+        noise_sim.simulate_batch_ensemble(sum(qtensor_circ, []), num_circs, num_qubits)
+        qtensor_probs = noise_sim.normalized_ensemble_probs 
+        end = time.time_ns() / (10 ** 9)
+
+        qiskit_probs, qiskit_time = simulate_qiskit_density_matrix(qiskit_circ, noise_model_qiskit, take_trace=True)
+        #fidelity = np.sum(np.square(np.sqrt(qtensor_probs) * np.sqrt(qiskit_probs)))
+        a = np.sqrt(qtensor_probs)
+        b = np.sqrt(qiskit_probs)
+        cos_sim = np.dot(a, b)/np.linalg.norm(a)/np.linalg.norm(b)
+        print("probs norm", print(sum(qtensor_probs)), print(sum(qiskit_probs)))
+        fidelity = np.dot(a, b)**2
+        print(f"num_circs: {num_circs}, Fidelity: {fidelity}, Cossim: {cos_sim}, TimeQiskit: {qiskit_time}, TimeQTensor: {end - start}")
diff --git a/qtensor/noise_simulator/generate_cz_rnd.py b/qtensor/noise_simulator/generate_cz_rnd.py
new file mode 100644
index 00000000..5b52262b
--- /dev/null
+++ b/qtensor/noise_simulator/generate_cz_rnd.py
@@ -0,0 +1,129 @@
+import numpy as np
+import qtree
+import qtensor
+
+def register_qiskit_gate():
+    import qiskit
+    qtensor.OpFactory.QiskitFactory.PauliGate = lambda x, alpha, beta, gamma: qiskit.extensions.UnitaryGate(
+        qtensor.tools.unitary.random_pauli(alpha, beta, gamma),
+        label='RandomPauliGate'
+    )
+
+def pauli_unitary(a, b, c):
+    """
+    Returns exp(-i*c*pauli.m)
+    where m = (sin(a)cos(b), sin(a)sin(b), cos(b))
+    """
+    v = np.array([
+        np.sin(a)*np.cos(b),
+        np.sin(a)*np.sin(b),
+        np.cos(a)
+    ])
+    paulis = np.array([
+        np.array([[0, 1], [1, 0]], dtype=complex),
+        np.array([[0, 1j], [-1j, 0]], dtype=complex),
+        np.array([[1, 0], [0, -1]], dtype=complex),
+    ])
+    R = np.sum([paulis[i]*v[i] for i in range(3)], axis=0)
+    return np.cos(c)*np.eye(2) - 1j*np.sin(c)*R
+
+class PauliGate(qtree.operators.ParametricGate):
+    name = "Pauli"
+    def __init__(self, *args, **kwargs):
+        self.ar = pauli_unitary(kwargs['alpha'], kwargs['beta'], kwargs['gamma'])
+        super().__init__(*args, **kwargs)
+        # used by data_key in global storage
+
+    def gen_tensor(self):
+        return self.ar
+
+    _changes_qubits = (0,)
+
+
+def get_cx_random_circuit(N, d, angle_sequence=None):
+    if angle_sequence is None:
+        angle_sequence = np.random.uniform(0, 2*np.pi, size=(N*d*2, 3))
+    circuit = []
+    k = 0
+    for l in range(d*2):
+        layer = []
+        for i in range(N):
+            gate = PauliGate(
+                i,
+                alpha=angle_sequence[k, 0],
+                beta=angle_sequence[k, 1],
+                gamma=angle_sequence[k, 2]
+            )
+            layer.append(gate)
+            k += 1
+        s = l % 2
+        #for s in range(2):
+        for i in range(N//2-s):
+            u, v = (s+2*i, s+2*i+1)
+            g = qtree.operators.cZ(u, v)
+            layer.append(g)
+        circuit.append(layer)
+    return circuit
+
+
+class CZBrickworkComposer(qtensor.CircuitComposer):
+    def _get_builder(self):
+        return qtensor.QtreeBuilder(self.n_qubits)
+
+    def __init__(self, S, *args, **kwargs):
+        self.S = S
+        self.n_qubits = S*S
+        super().__init__(self.n_qubits, *args, **kwargs)
+
+    def apply_layer(self, mode):
+        switch = 'n' in mode
+        start = lambda s: 0 if s else 1
+        end_i = self.S if 'v' in mode else (self.S - 1)
+        end_j = self.S if 'h' in mode else self.S - 1
+        for j in range(0, end_j, 1):
+            for i in range(start(switch), end_i, 2):
+                A = i+self.S*j
+                if 'v' in mode:
+                    # vertical
+                    B = i + self.S*(j+1)
+                else:
+                    # horisontal
+                    B = i + 1 + self.S*j
+                self.two_qubit(A, B)
+            switch = not switch
+
+    def two_qubit(self, u, v):
+        self.apply_gate(self.builder.operators.cZ, u, v)
+
+    def one_qubit(self, u, params):
+        options = [self.builder.operators.X, self.builder.operators.Y, self.builder.operators.Z]
+        gate = np.random.choice(options)
+        self.apply_gate(gate, u)
+
+    def one_qubit_layer(self, params=None):
+        if params is None:
+            params = 2*np.pi*np.random.rand(self.n_qubits, 3)
+        for i in range(self.n_qubits):
+            self.one_qubit(i, params[i])
+
+    def two_qubit_rnd(self, layers=10):
+        modes = ['vp', 'vn', 'hp', 'hn']
+        for i in range(layers):
+            for mode in modes:
+                self.one_qubit_layer()
+                self.apply_layer(mode)
+
+
+def get_cz_circ(S, d):
+    comp = CZBrickworkComposer(S)
+    comp.two_qubit_rnd(layers=d)
+    return comp.circuit
+
+class QiskitCZBrickworkComposer(CZBrickworkComposer):
+    def _get_builder(self):
+        return qtensor.QiskitBuilder(self.n_qubits)
+
+
+if __name__ == "__main__":
+    circ = get_cz_circ(3, 5)
+    print(circ)

From 2a6add69d5261c6a2fc1aa4c59e9022155b9c639 Mon Sep 17 00:00:00 2001
From: Danil Lykov <lkv97dn@gmail.com>
Date: Thu, 29 Sep 2022 16:53:35 -0500
Subject: [PATCH 03/24] add x and y to qiskit opfactory

---
 qtensor/OpFactory.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/qtensor/OpFactory.py b/qtensor/OpFactory.py
index 75bde0dc..d167c336 100644
--- a/qtensor/OpFactory.py
+++ b/qtensor/OpFactory.py
@@ -181,6 +181,14 @@ def cX(cls):
     def cZ(cls):
         return qiskit_lib.CZGate
 
+    @property
+    def X(cls):
+        return qiskit_lib.XGate
+
+    @property
+    def Y(cls):
+        return qiskit_lib.YGate
+
     @property
     def Z(cls):
         return qiskit_lib.ZGate

From 6643a0d09052ea193b4d7be54cd30ba4e1d67f6a Mon Sep 17 00:00:00 2001
From: Dan Lykov <lkv97dn@gmail.com>
Date: Thu, 6 Oct 2022 16:45:52 -0600
Subject: [PATCH 04/24] fix the set up of qiskit noisy simulation

---
 .../comparison_generic_circ.py                | 23 +++++++++++++------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/qtensor/noise_simulator/comparison_generic_circ.py b/qtensor/noise_simulator/comparison_generic_circ.py
index 74d9ce95..8fc7324a 100644
--- a/qtensor/noise_simulator/comparison_generic_circ.py
+++ b/qtensor/noise_simulator/comparison_generic_circ.py
@@ -3,8 +3,10 @@
 from qtree.operators import from_qiskit_circuit
 from qiskit import execute, Aer
 from qtensor import QiskitQAOAComposer
+import qtensor
 import networkx as nx
 import qiskit.providers.aer.noise as noise
+from qiskit.providers.aer import AerSimulator
 import time
 
 def get_circs(S, d):
@@ -13,9 +15,11 @@ def get_circs(S, d):
     comp = QiskitCZBrickworkComposer(S)
     comp.two_qubit_rnd(layers=d)
 
-    ###
-    #gamma = [0.3] * d
-    #beta = [0.3] * d
+    ##
+    gammabeta = np.array(qtensor.tools.BETHE_QAOA_VALUES[str(d)]['angles'])
+
+    gamma = -gammabeta[:d]
+    beta = gammabeta[d:]
     #comp = QiskitQAOAComposer(G, gamma=gamma, beta=beta)
     #comp.ansatz_state()
     n, qtensor_circ = from_qiskit_circuit(comp.circuit)
@@ -27,7 +31,10 @@ def get_circs(S, d):
 
 def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = True):
     start = time.time_ns() / (10 ** 9)
-    result = execute(circuit, Aer.get_backend('aer_simulator_density_matrix'), shots=1, noise_model=noise_model_qiskit).result()
+    backend = AerSimulator(method='density_matrix', noise_model=noise_model_qiskit)
+
+    result = execute(circuit, backend, shots=1).result()
+    #result = execute(circuit, Aer.get_backend('aer_simulator_density_matrix'), shots=1, noise_model=noise_model_qiskit).result()
     if take_trace:
         qiskit_probs = np.diagonal(result.results[0].data.density_matrix.real)
         end = time.time_ns() / (10 ** 9)
@@ -49,7 +56,7 @@ def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = Tru
     prob_2 = 0.1
 
     S = 2
-    d = 3
+    d = 1
     qtensor_circ, qiskit_circ = get_circs(S, d)
     num_qubits = S**2
     # Qiskit Noise Model
@@ -58,7 +65,8 @@ def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = Tru
 
     noise_model_qiskit = noise.NoiseModel()
     noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_1Q, ['x', 'y', 'z'])
-    noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_2Q, ['cx'])
+    noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_1Q, ['rz', 'ry', 'rx', 'h'])
+    noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_2Q, ['cx', 'cz'])
 
     # QTensor Noise Model
     depol_chan_qtensor_1Q = DepolarizingChannel(prob_1, 1)
@@ -66,7 +74,8 @@ def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = Tru
 
     noise_model_qtensor = NoiseModel()
     noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['X', 'Y', 'Z'])
-    noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX'])
+    #noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['XPhase', 'YPhase', 'ZPhase'])
+    noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX',])
 
     noise_sim = NoiseSimulator(noise_model_qtensor)
 

From 1357658f7e8c1ae8555bb4282f7c44bc453e01a0 Mon Sep 17 00:00:00 2001
From: Dan Lykov <lkv97dn@gmail.com>
Date: Thu, 6 Oct 2022 17:34:41 -0600
Subject: [PATCH 05/24] add readme and add H error to qtensor noise model

---
 qtensor/noise_simulator/README.md             | 38 +++++++++++++++++++
 .../comparison_generic_circ.py                |  2 +-
 2 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 qtensor/noise_simulator/README.md

diff --git a/qtensor/noise_simulator/README.md b/qtensor/noise_simulator/README.md
new file mode 100644
index 00000000..5f0cc8a5
--- /dev/null
+++ b/qtensor/noise_simulator/README.md
@@ -0,0 +1,38 @@
+# Noisy simulations using QTensor
+
+This folder uses stochastic simulation of quantum circuits
+and benchmarks it against qiskit density matrix simulations.
+
+
+## Observations
+
+Two types of circuits are considered:
+
+1. QAOA MaxCut with CZ two-qubit gate
+2. Randomized brickwork circuit with XYZ gates and CX gate
+
+The noise model is depolarizing noise applied to every quantum gate.
+
+Two types of qiskit runs were considered: with passing noise_model to `execute` and to `backend`. These modes are named E and B below.
+Mode B is the correct one.
+
+The qiskit simulation results are then compared to qtensor simulation results and the discrepancy between the two is evaluated.
+
+### 1. QAOA circuits
+
+When running with mode E, we see that there is very small error and it reduces with number of circuits k. The same was with mode B.
+
+### 2. Generic circuits
+
+When running in mode E, we found that the error does not converge with number of circuits K. 
+In the mode B we found that the error is larger but converges better (by how much?) K, which indicates that we compare against the correct noisy output.
+
+When running in mode E, removing two-qubit gate error channels reduced the simulation error.
+
+When running in mode E, adding a layer of Hadamards in the begining of the circuit reduced the error significantly.
+
+When running in both modes E and B, increasing error rate in channel reduces the simulation error.
+
+In some cases hovewer, there was convergence of error in mode E
+
+These results are preliminary.
diff --git a/qtensor/noise_simulator/comparison_generic_circ.py b/qtensor/noise_simulator/comparison_generic_circ.py
index 8fc7324a..2fc607e1 100644
--- a/qtensor/noise_simulator/comparison_generic_circ.py
+++ b/qtensor/noise_simulator/comparison_generic_circ.py
@@ -74,7 +74,7 @@ def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = Tru
 
     noise_model_qtensor = NoiseModel()
     noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['X', 'Y', 'Z'])
-    #noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['XPhase', 'YPhase', 'ZPhase'])
+    noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['XPhase', 'YPhase', 'ZPhase', 'H'])
     noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX',])
 
     noise_sim = NoiseSimulator(noise_model_qtensor)

From 8f7724402624bfc09aa597a817ec3b21af987393 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Sat, 8 Oct 2022 15:46:27 -0500
Subject: [PATCH 06/24] fixed qtensor noise model so it matches the qiskit
 noise model. Ran some tests and recorded results in README.md

---
 qtensor/noise_simulator/README.md                  | 6 ++++++
 qtensor/noise_simulator/comparison_generic_circ.py | 2 +-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/qtensor/noise_simulator/README.md b/qtensor/noise_simulator/README.md
index 5f0cc8a5..772f459b 100644
--- a/qtensor/noise_simulator/README.md
+++ b/qtensor/noise_simulator/README.md
@@ -36,3 +36,9 @@ When running in both modes E and B, increasing error rate in channel reduces the
 In some cases hovewer, there was convergence of error in mode E
 
 These results are preliminary.
+
+When running in mode E with a layer of Hadamards at the beginning and end of the circuit, a depolazing error on 0.01 on single qubit gates but no error on two-qubit gates had significant error. 
+
+When running in mode E with a layer of Hadamards at the beginning and end of the circuit, a depolazing error on 0.01 on single qubit gates and 0.1 on two-qubit gates significantly reduced the error. 
+
+When running in mode E with a layer of Hadamards at the beginning and end of the circuit, a depolazing error on 0.25 on single qubit gates and 0.1 on two-qubit gates caused the error to return. 
\ No newline at end of file
diff --git a/qtensor/noise_simulator/comparison_generic_circ.py b/qtensor/noise_simulator/comparison_generic_circ.py
index 2fc607e1..b20e0f11 100644
--- a/qtensor/noise_simulator/comparison_generic_circ.py
+++ b/qtensor/noise_simulator/comparison_generic_circ.py
@@ -75,7 +75,7 @@ def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = Tru
     noise_model_qtensor = NoiseModel()
     noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['X', 'Y', 'Z'])
     noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['XPhase', 'YPhase', 'ZPhase', 'H'])
-    noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX',])
+    noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX', 'cZ'])
 
     noise_sim = NoiseSimulator(noise_model_qtensor)
 

From 721c9c90fc96f4eaefd92640e6376818a36c93b9 Mon Sep 17 00:00:00 2001
From: Dan Lykov <lkv97dn@gmail.com>
Date: Sun, 9 Oct 2022 15:44:21 -0600
Subject: [PATCH 07/24] Fix qiskit simulation even one more time

---
 .../comparison_generic_circ.py                | 42 ++++++++++++++-----
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/qtensor/noise_simulator/comparison_generic_circ.py b/qtensor/noise_simulator/comparison_generic_circ.py
index b20e0f11..cfb520f0 100644
--- a/qtensor/noise_simulator/comparison_generic_circ.py
+++ b/qtensor/noise_simulator/comparison_generic_circ.py
@@ -2,7 +2,8 @@
 import numpy as np
 from qtree.operators import from_qiskit_circuit
 from qiskit import execute, Aer
-from qtensor import QiskitQAOAComposer
+from qtensor import QiskitQAOAComposer, QtreeSimulator
+from qtensor import QiskitBuilder
 import qtensor
 import networkx as nx
 import qiskit.providers.aer.noise as noise
@@ -12,28 +13,37 @@
 def get_circs(S, d):
     G = nx.random_regular_graph(3, S**2)
 
-    comp = QiskitCZBrickworkComposer(S)
-    comp.two_qubit_rnd(layers=d)
+    #comp = QiskitCZBrickworkComposer(S)
+    #comp.two_qubit_rnd(layers=d)
 
     ##
     gammabeta = np.array(qtensor.tools.BETHE_QAOA_VALUES[str(d)]['angles'])
 
     gamma = -gammabeta[:d]
     beta = gammabeta[d:]
-    #comp = QiskitQAOAComposer(G, gamma=gamma, beta=beta)
-    #comp.ansatz_state()
+    comp = QiskitQAOAComposer(G, gamma=gamma, beta=beta)
+    comp.ansatz_state()
+    #builder = QiskitBuilder(1)
+    #for i in range(10):
+    #    builder.apply_gate(builder.operators.X, 0)
+    #    builder.apply_gate(builder.operators.Z, 0)
+    #comp.circuit = builder.circuit
+
     n, qtensor_circ = from_qiskit_circuit(comp.circuit)
 
     comp.circuit = comp.circuit.reverse_bits()
     comp.circuit.save_density_matrix()
     comp.circuit.measure_all(add_bits = False)
+    #print("Circuit:")
+    #print(comp.circuit)
     return qtensor_circ, comp.circuit
 
 def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = True):
     start = time.time_ns() / (10 ** 9)
-    backend = AerSimulator(method='density_matrix', noise_model=noise_model_qiskit)
+    backend = AerSimulator(method='density_matrix', noise_model=noise_model_qiskit, fusion_enable=False, fusion_verbose=True)
 
     result = execute(circuit, backend, shots=1).result()
+    result = backend.run(circuit).result()
     #result = execute(circuit, Aer.get_backend('aer_simulator_density_matrix'), shots=1, noise_model=noise_model_qiskit).result()
     if take_trace:
         qiskit_probs = np.diagonal(result.results[0].data.density_matrix.real)
@@ -56,9 +66,10 @@ def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = Tru
     prob_2 = 0.1
 
     S = 2
-    d = 1
+    d = 4
     qtensor_circ, qiskit_circ = get_circs(S, d)
     num_qubits = S**2
+    #num_qubits = 1
     # Qiskit Noise Model
     depol_chan_qiskit_1Q = noise.depolarizing_error(prob_1, 1)
     depol_chan_qiskit_2Q = noise.depolarizing_error(prob_2, 2)
@@ -67,6 +78,7 @@ def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = Tru
     noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_1Q, ['x', 'y', 'z'])
     noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_1Q, ['rz', 'ry', 'rx', 'h'])
     noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_2Q, ['cx', 'cz'])
+    #print("Qiskit noise model", noise_model_qiskit.to_dict())
 
     # QTensor Noise Model
     depol_chan_qtensor_1Q = DepolarizingChannel(prob_1, 1)
@@ -78,19 +90,29 @@ def simulate_qiskit_density_matrix(circuit, noise_model_qiskit, take_trace = Tru
     noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX', 'cZ'])
 
     noise_sim = NoiseSimulator(noise_model_qtensor)
+    exact_sim = QtreeSimulator()
 
 
     for num_circs in [100, 300, 1000, 3000]:
         start = time.time_ns() / (10 ** 9)
         noise_sim.simulate_batch_ensemble(sum(qtensor_circ, []), num_circs, num_qubits)
-        qtensor_probs = noise_sim.normalized_ensemble_probs 
+        qtensor_probs = noise_sim.normalized_ensemble_probs
         end = time.time_ns() / (10 ** 9)
 
+        amps = exact_sim.simulate_batch(sum(qtensor_circ, []), batch_vars=num_qubits)
+        probs_uniform = np.ones(2**num_qubits)/2**num_qubits
+        probs = np.abs(amps)**2
+        c = np.sqrt(probs)
+        d = np.sqrt(probs_uniform)
+
         qiskit_probs, qiskit_time = simulate_qiskit_density_matrix(qiskit_circ, noise_model_qiskit, take_trace=True)
         #fidelity = np.sum(np.square(np.sqrt(qtensor_probs) * np.sqrt(qiskit_probs)))
+        print(f"Qiskit_probs: {qiskit_probs[:10]}, QTensor_noiseless_probs: {probs[:10]}, Qtensor_stochastic_probs {qtensor_probs[:10]}")
         a = np.sqrt(qtensor_probs)
         b = np.sqrt(qiskit_probs)
         cos_sim = np.dot(a, b)/np.linalg.norm(a)/np.linalg.norm(b)
-        print("probs norm", print(sum(qtensor_probs)), print(sum(qiskit_probs)))
+        ref_noiseless = np.dot(a, c)**2
+        print("probs norm", sum(qtensor_probs), sum(qiskit_probs))
         fidelity = np.dot(a, b)**2
-        print(f"num_circs: {num_circs}, Fidelity: {fidelity}, Cossim: {cos_sim}, TimeQiskit: {qiskit_time}, TimeQTensor: {end - start}")
+        fidelity_uniform = np.dot(b, d)**2
+        print(f"num_circs: {num_circs}, Fidelity: {fidelity}, Cossim: {cos_sim}, Fid_noiseless : {ref_noiseless}, Fid_uniform: {fidelity_uniform}, TimeQiskit: {qiskit_time}, TimeQTensor: {end - start}")

From 6ef5419db8d3e136df0cbb7b685ebe2dcf698a9e Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Sun, 9 Oct 2022 23:41:53 -0500
Subject: [PATCH 08/24] fixed qiskit density matrix simulator function calls so
 they are used correctly. Added more fidelity measures for uniform and
 noiseless states

---
 .../noise_simulator/ComparisonSimulator.py    | 54 +++++++++++--------
 .../NoiseSimComparisonResult.py               | 44 ++++++++++-----
 qtensor/noise_simulator/comparison.py         | 24 ++++-----
 qtensor/noise_simulator/comparison_mpi.py     | 16 ++++--
 qtensor/noise_simulator/helper_functions.py   | 12 ++++-
 5 files changed, 96 insertions(+), 54 deletions(-)

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
index 8adc44bd..7dba8734 100644
--- a/qtensor/noise_simulator/ComparisonSimulator.py
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -9,6 +9,8 @@
 from qtree.operators import from_qiskit_circuit
 from qtensor import tools
 from itertools import repeat
+from qiskit import execute, Aer
+from qiskit.providers.aer import AerSimulator
 
 class ComparisonSimulator:
     def __init__(self):
@@ -44,34 +46,36 @@ def qtensor_qiskit_noisy_qaoa(self, recompute_previous_ensemble: bool = False, p
         G, gamma, beta = get_qaoa_params(n = self.n, p = self.p, d = self.d)
         self._get_circuits(G, gamma, beta)
         noise_sim = NoiseSimulator(self.noise_model_qtensor)
+        exact_sim = QtreeSimulator()
 
         # Run simulation
         for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
             result = NoiseSimComparisonResult(self.qiskit_circ, self.qtensor_circ, self.noise_model_qiskit, 
                 self.noise_model_qtensor, self.n, self.p, self.d)
 
-            if i == 0 or recompute_previous_ensemble == True:
+            if i == 0 or recompute_previous_ensemble == False:
                 self.num_circs_simulated.append(num_circs)
-                noise_sim.simulate_batch_ensemble(self.qtensor_circ, num_circs, self.num_qubits)
+                noise_sim.simulate_batch_ensemble(sum(self.qtensor_circ, []), num_circs, self.num_qubits)
                 #print("num_circs: {}".format(num_circs))
-                qtensor_probs = noise_sim.normalized_ensemble_probs 
+                self.qtensor_probs = noise_sim.normalized_ensemble_probs
             else: 
                 actual_num_circs = num_circs - self.num_circs_list[i - 1]
                 self.num_circs_simulated.append(actual_num_circs)
                 #print("num_circs: {}, actual_num_circs: {}".format(num_circs, actual_num_circs_to_sim))
-                noise_sim.simulate_batch_ensemble(self.qtensor_circ, actual_num_circs, self.num_qubits)
+                noise_sim.simulate_batch_ensemble(sum(self.qtensor_circ, []), actual_num_circs, self.num_qubits)
                 prev_qtensor_probs = self.prev_probs
                 curr_qtensor_probs = noise_sim.normalized_ensemble_probs
-                qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
+                self.qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
             
-            if recompute_previous_ensemble == False:
-                self.prev_probs = qtensor_probs
+            if recompute_previous_ensemble == True:
+                self.prev_probs = self.qtensor_probs
 
             qtensor_sim_time = noise_sim.time_taken
             self.simulate_qiskit_density_matrix(self.qiskit_circ, self.noise_model_qiskit)    
+            self.exact_qtensor_amps = exact_sim.simulate_batch(sum(self.qtensor_circ, []), batch_vars=self.num_qubits)
 
             # Save results
-            result.save_result(self.qiskit_probs, qtensor_probs, num_circs, 
+            result.save_result(self.qiskit_probs, self.qtensor_probs, self.exact_qtensor_amps, num_circs,
                 self.num_circs_simulated[i], G, gamma, beta, qtensor_sim_time, self.qiskit_sim_time)
             self.results.append(result.data)
             if print_stats:
@@ -80,7 +84,7 @@ def qtensor_qiskit_noisy_qaoa(self, recompute_previous_ensemble: bool = False, p
     def _mpi_parallel_unit(self, args):
         noise_sim, qtensor_circ, num_circs, num_qubits = args
         #print("this workers num_circs: ", num_circs)
-        noise_sim.simulate_batch_ensemble(qtensor_circ, num_circs, num_qubits)
+        noise_sim.simulate_batch_ensemble(sum(qtensor_circ, []), num_circs, num_qubits)
         fraction_of_qtensor_probs = noise_sim.normalized_ensemble_probs
         return fraction_of_qtensor_probs
 
@@ -93,6 +97,7 @@ def qtensor_qiskit_noisy_qaoa_mpi(self, num_nodes: int, num_jobs_per_node: int,
         G, gamma, beta = get_qaoa_params(n = self.n, p = self.p, d = self.d)
         self._get_circuits(G, gamma, beta)
         self.noise_sim = NoiseSimulator(self.noise_model_qtensor)
+        exact_sim = QtreeSimulator()
 
         for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
             result = NoiseSimComparisonResult(self.qiskit_circ, self.qtensor_circ, self.noise_model_qiskit, 
@@ -101,28 +106,28 @@ def qtensor_qiskit_noisy_qaoa_mpi(self, num_nodes: int, num_jobs_per_node: int,
             qtensor_probs_list = tools.mpi.mpi_map(self._mpi_parallel_unit, self._arggen, pbar=pbar, total=num_nodes)
             self._check_correct_num_circs_simulated(i)
             if qtensor_probs_list:
-                if i == 0 or recompute_previous_ensemble == True: 
-                    qtensor_probs = sum(qtensor_probs_list) / self._total_jobs
+                if i == 0 or recompute_previous_ensemble == False: 
+                    self.qtensor_probs = sum(qtensor_probs_list) / self._total_jobs
                 else:
                     #prev_qtensor_probs = self.results[i - 1].qtensor_probs
                     prev_qtensor_probs = self.prev_probs
                     curr_qtensor_probs = sum(qtensor_probs_list) / self._total_jobs
-                    qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
+                    self.qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
                 qtensor_sim_time = self.noise_sim.time_taken
-                if recompute_previous_ensemble == False:
-                    self.prev_probs = qtensor_probs
+                if recompute_previous_ensemble == True:
+                    self.prev_probs = self.qtensor_probs
 
                 self.simulate_qiskit_density_matrix(self.qiskit_circ, self.noise_model_qiskit)
-
+                self.exact_qtensor_amps = exact_sim.simulate_batch(sum(self.qtensor_circ, []), batch_vars=self.num_qubits)
                 # Save results 
-                result.save_result(self.qiskit_probs, qtensor_probs, num_circs,
+                result.save_result(self.qiskit_probs, self.qtensor_probs, self.exact_qtensor_amps, num_circs,
                     self.num_circs_simulated[i], G, gamma, beta, qtensor_sim_time, self.qiskit_sim_time)
                 self.results.append(result.data) 
                 if print_stats:
                     tools.mpi.print_stats()
                     result.print_result()
 
-    def qtensor_qiskit_noisy_qaoa_density(self, recompute_previous_ensemble: bool = True):
+    def qtensor_qiskit_noisy_qaoa_density(self, recompute_previous_ensemble: bool = False):
         self.recompute_previous_ensemble = recompute_previous_ensemble
         # Prepare circuits, simulator, and area to save results
         G, gamma, beta = get_qaoa_params(n = self.n, p = self.p, d = self.d)
@@ -136,28 +141,28 @@ def qtensor_qiskit_noisy_qaoa_density(self, recompute_previous_ensemble: bool =
             if i == 0 or recompute_previous_ensemble == True: 
                 self.num_circs_simulated.append(num_circs)
                 noise_sim.simulate_batch_ensemble_density(self.qtensor_circ, num_circs, self.n)
-                qtensor_density_matrix = noise_sim.normalized_ensemble_density_matrix
+                self.qtensor_density_matrix = noise_sim.normalized_ensemble_density_matrix
             else: 
                 actual_num_circs = num_circs - self.num_circs_list[i - 1]
                 self.num_circs_simulated.append(actual_num_circs)
                 noise_sim.simulate_batch_ensemble_density(self.qtensor_circ, actual_num_circs, self.n)
                 prev_density_matrix = self.prev_qtensor_density_matrix
                 curr_density_matrix = noise_sim.normalized_ensemble_probs
-                qtensor_density_matrix = (curr_density_matrix + prev_density_matrix) / 2
+                self.qtensor_density_matrix = (curr_density_matrix + prev_density_matrix) / 2
             qtensor_sim_time = noise_sim.time_taken
             self.simulate_qiskit_density_matrix(self.qiskit_circ, self.noise_model_qiskit, take_trace = False)
 
             if recompute_previous_ensemble == False: 
-                self.prev_qtensor_density_matrix = qtensor_density_matrix
+                self.prev_qtensor_density_matrix = self.qtensor_density_matrix
             # Save results
-            result.save_results_density(self.qiskit_density_matrix, qtensor_density_matrix, num_circs, 
+            result.save_results_density(self.qiskit_density_matrix, self.qtensor_density_matrix, num_circs, 
                 self.num_circs_simulated[i], G, gamma, beta, qtensor_sim_time, self.qiskit_sim_time)
             self.results.append(result.data)
 
 
     # Prepare arguments to be sent to each unit of work 
     def _get_args(self, i):
-        if i == 0 or self.recompute_previous_ensemble == True:
+        if i == 0 or self.recompute_previous_ensemble == False:
             num_circs = self.num_circs_list[i]
         else:
             num_circs = self.num_circs_list[i] - self.num_circs_list[i - 1]
@@ -195,7 +200,10 @@ def _get_args(self, i):
 
     def simulate_qiskit_density_matrix(self, circuit, noise_model_qiskit, take_trace = True):
         start = time.time_ns() / (10 ** 9)
-        result = execute(circuit, Aer.get_backend('aer_simulator_density_matrix'), shots=1, noise_model=noise_model_qiskit).result()
+        backend = AerSimulator(method='density_matrix', noise_model=noise_model_qiskit, fusion_enable=False, fusion_verbose=True)
+        result = execute(circuit, backend, shots=1).result()
+        result = backend.run(circuit).result()
+        #result = execute(circuit, Aer.get_backend('aer_simulator_density_matrix'), shots=1, noise_model=noise_model_qiskit).result()
         if take_trace:
             self.qiskit_probs = np.diagonal(result.results[0].data.density_matrix.real)
             end = time.time_ns() / (10 ** 9)
diff --git a/qtensor/noise_simulator/NoiseSimComparisonResult.py b/qtensor/noise_simulator/NoiseSimComparisonResult.py
index 9b9d0d80..6c418cbe 100644
--- a/qtensor/noise_simulator/NoiseSimComparisonResult.py
+++ b/qtensor/noise_simulator/NoiseSimComparisonResult.py
@@ -50,7 +50,7 @@ def save_results_density(self, qiskit_density_matrix, qtensor_density_matrix, nu
         self.total_time_taken = self.experiment_end_time - self.experiment_start_time
         self._calc_similarity()
 
-    def save_result(self, qiskit_probs, qtensor_probs, num_circs, num_circs_simulated, G, gamma, beta, qtensor_time_total, qiskit_time_total):
+    def save_result(self, qiskit_probs, qtensor_probs, exact_qtensor_amps, num_circs, num_circs_simulated, G, gamma, beta, qtensor_time_total, qiskit_time_total):
         self.qiskit_probs = qiskit_probs
         self.qtensor_probs = qtensor_probs
         self.num_circs = num_circs
@@ -62,20 +62,24 @@ def save_result(self, qiskit_probs, qtensor_probs, num_circs, num_circs_simulate
         self.qiskit_time_taken = qiskit_time_total
         self.experiment_end_time =  time.time_ns() / (10 ** 9) 
         self.experiment_time_taken = self.experiment_end_time - self.experiment_start_time
+        self.exact_qtensor_amps = exact_qtensor_amps
         self._calc_similarity_of_probs()
         self._to_dict()
 
     def print_result(self):
         print("\nExperiment with n = {}, p = {}, d = {}, num_circs = {}, actual_num_circs = {}".format(self.n, self.p, self.d, self.num_circs, self.num_circs_simulated))
-        if type(self.fidelity) is str:
-            print(
+        # if type(self.fidelity) is str:
+        #     print(
+        #         f"{'Cosine Similarity:':<20}{np.round(self.cos_sim.real, 7):<10}",
+        #         f"{'Total time taken:':<20}{self.experiment_time_taken:<10}")
+        # else:
+        print(
                 f"{'Cosine Similarity:':<20}{np.round(self.cos_sim.real, 7):<10}",
-                f"{'Total time taken:':<20}{self.experiment_time_taken:<10}")
-        else:
-            print(
-                f"{'Cosine Similarity:':<20}{np.round(self.cos_sim.real, 7):<10}",
-                f"\n{'Fidelity:':<20}{np.round(self.fidelity.real, 7):<10}",
-                f"{'Total time taken:':<20}{self.experiment_time_taken:<10}")
+                f"\n{'Noisy Fidelity:':<20}{np.round(self.noisy_fidelity.real, 7):<10}",
+                f"\n{'Noiseless Fidelity:':<20}{np.round(self.noiseless_fidelity.real, 7):<10}",
+                f"\n{'Uniform Qiskit Fidelity: ':<20}{np.round(self.uniform_qiskit_fidelity.real, 7):<10}",
+                f"\n{'Uniform QTensor Fidelity: ':<20}{np.round(self.uniform_qtensor_fidelity.real, 7):<10}",
+                f"\n{'Total time taken:':<20}{self.experiment_time_taken:<10}")
 
     def print_noise_model(self):
         print(self.noise_model_str)
@@ -97,11 +101,24 @@ def _calc_similarity(self):
     def _calc_similarity_of_probs(self):
         qiskit_probs_root = np.sqrt(self.qiskit_probs)
         qtensor_probs_root = np.sqrt(self.qtensor_probs)
+        noiseless_qtensor_probs_root = np.sqrt(np.abs(self.exact_qtensor_amps)**2)
+        uniform_probs_root = np.sqrt(np.ones(2**self.n)/2**self.n)
 
-        self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
         # we don't need to take the conjugate, as both population density vectors are strictly 
         # real by the time they have made it here. 
-        self.fidelity = (np.inner(qiskit_probs_root, qtensor_probs_root))**2
+
+        """Measures the fidelity between the qiskit density matrix noisy state and the qtensoir stochastic noisy state"""
+        self.noisy_fidelity = (np.inner(qiskit_probs_root, qtensor_probs_root))**2
+
+        """Measures the fidelity between a noisy qtensor state and the noiseless version of the same state"""
+        self.noiseless_fidelity = (np.inner(noiseless_qtensor_probs_root, qtensor_probs_root))**2
+
+        """Measures the fidelity between a qiskit noisy state and a uniform distribution state"""
+        self.uniform_qiskit_fidelity = (np.inner(uniform_probs_root, qiskit_probs_root))**2
+
+        """Measures the fidelity between a qtensor noisy state and a uniform distribution state"""
+        self.uniform_qtensor_fidelity = (np.inner(uniform_probs_root, qtensor_probs_root))**2
+        self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
 
     def _to_dict(self):
         #self.experiment_dict['name'] = self.name
@@ -113,7 +130,10 @@ def _to_dict(self):
         self.data['gamma'] = self.gamma
         self.data['beta'] = self.beta
         self.data['cosine_similarity'] = self.cos_sim.real
-        self.data['fidelity'] = self.fidelity
+        self.data['noisy_fidelity'] = self.noisy_fidelity
+        self.data['noiseless_fidelity'] = self.noiseless_fidelity
+        self.data['uniform_qiskit_fidelity'] = self.uniform_qiskit_fidelity
+        self.data['uniform_qtensor_fidelity'] = self.uniform_qtensor_fidelity
         self.data['noise_model_string'] = self.noise_model_str_condensed
         self.data['noise_model_pickle'] = self.pickled_qtensor_noise_model
         self.data['experiment_date'] = self.experiment_date
diff --git a/qtensor/noise_simulator/comparison.py b/qtensor/noise_simulator/comparison.py
index 882dc3ae..3229708a 100644
--- a/qtensor/noise_simulator/comparison.py
+++ b/qtensor/noise_simulator/comparison.py
@@ -8,8 +8,8 @@
 from qtensor.tests.test_composers import *
 
 
-prob_1 = 0.01
-prob_2 = 0.1
+prob_1 = 0.003
+prob_2 = 0.03
 
 # Qiskit Noise Model
 depol_chan_qiskit_1Q = noise.depolarizing_error(prob_1, 1)
@@ -27,21 +27,21 @@
 noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['H', 'XPhase', 'ZPhase'])
 noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX'])
 
-min_n = 5
-max_n = 7
-
-min_p = 2
-max_p = 3
-
-min_d = 2
-max_d = 3
 
-num_samples = 1
-num_circs_list = [10, 18, 32, 100, 178, 316]
+num_samples = 30
+num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
 results = []
 
 outfile_name = '{}.json'.format(datetime.now().isoformat())
 
+min_n = 5
+max_n = 16
+
+min_p = 2
+max_p = 3
+
+min_d = 4
+max_d = 5
 
 for n in range(min_n, max_n):
     for p in range(min_p, max_p):
diff --git a/qtensor/noise_simulator/comparison_mpi.py b/qtensor/noise_simulator/comparison_mpi.py
index 28b46507..b57d2a50 100644
--- a/qtensor/noise_simulator/comparison_mpi.py
+++ b/qtensor/noise_simulator/comparison_mpi.py
@@ -5,11 +5,12 @@
 from helper_functions import *
 import qiskit.providers.aer.noise as noise
 from qtensor.tests.test_composers import *
+from ComparisonSimulator import QAOAComparisonSimulator
 
 
 ### Noise model, simulation, and results of a QAOA algorithmm  ###
-prob_1 = 0.01
-prob_2 = 0.1
+prob_1 = 0.003
+prob_2 = 0.03
 
 # Qiskit Noise Model
 depol_chan_qiskit_1Q = noise.depolarizing_error(prob_1, 1)
@@ -36,10 +37,15 @@
 min_d = 2
 max_d = 3
 
+# num samples determines how many samples of a particular n, p, d, and num_circs we take
+# generally we want to get >30 to reduce sampling noise 
 num_samples = 3
-num_circs_list = [10, 18, 32, 100, 178, 316]
-num_nodes = 1
-num_jobs_per_node = 3
+
+num_nodes = 2
+# num jobs per node is equivalent to the number of cores per node
+num_jobs_per_node = 2
+
+num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
 
 outfile_name = '{}.json'.format(datetime.now().isoformat())
 results = []
diff --git a/qtensor/noise_simulator/helper_functions.py b/qtensor/noise_simulator/helper_functions.py
index aac630cb..84e9750e 100644
--- a/qtensor/noise_simulator/helper_functions.py
+++ b/qtensor/noise_simulator/helper_functions.py
@@ -12,6 +12,8 @@
 from datetime import datetime
 from os import SEEK_SET
 
+import qtensor
+
 def decimal_to_binary(n):
 # converting decimal to binary
 # and removing the prefix(0b)
@@ -57,8 +59,14 @@ def get_qaoa_params(n, p, d):
 
     print('Circuit params: n: {}, p: {}, d: {}'.format(n, p, d))
     G = nx.random_regular_graph(d, n)
-    gamma = [np.random.uniform(0, 2*np.pi) for _ in range(p)]
-    beta = [np.random.uniform(0, np.pi) for _ in range(p)]
+    gammabeta = np.array(qtensor.tools.BETHE_QAOA_VALUES[str(d)]['angles'])
+    #gammabeta = gammabeta.tolist()
+    gamma = gammabeta[:d]
+    beta = gammabeta[d:]
+    gamma = gamma.tolist()
+    beta = beta.tolist()
+    #gamma = [np.random.uniform(0, 2*np.pi) for _ in range(p)]
+    #beta = [np.random.uniform(0, np.pi) for _ in range(p)]
     return G, gamma, beta
 
 def save_dict_to_file(dict, name):

From 53fea743aaea5efcf3d9277142c080a720d9af73 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Mon, 10 Oct 2022 01:09:55 -0500
Subject: [PATCH 09/24] minor fixes

---
 qtensor/noise_simulator/comparison_mpi.py | 29 +++++++++++++----------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/qtensor/noise_simulator/comparison_mpi.py b/qtensor/noise_simulator/comparison_mpi.py
index b57d2a50..482f6d7b 100644
--- a/qtensor/noise_simulator/comparison_mpi.py
+++ b/qtensor/noise_simulator/comparison_mpi.py
@@ -28,6 +28,22 @@
 noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_1Q, ['H', 'XPhase', 'ZPhase'])
 noise_model_qtensor.add_channel_to_all_qubits(depol_chan_qtensor_2Q, ['cX'])
 
+"""
+num samples determines how many samples of a particular n, p, d, and num_circs we take
+generally we want to get >30 to reduce sampling noise 
+"""
+num_samples = 30
+
+num_nodes = 10
+
+"""num jobs per node is equivalent to the number of cores per node"""
+num_jobs_per_node = 56
+
+num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
+
+outfile_name = '{}.json'.format(datetime.now().isoformat())
+results = []
+
 min_n = 5
 max_n = 7
 
@@ -37,19 +53,6 @@
 min_d = 2
 max_d = 3
 
-# num samples determines how many samples of a particular n, p, d, and num_circs we take
-# generally we want to get >30 to reduce sampling noise 
-num_samples = 3
-
-num_nodes = 2
-# num jobs per node is equivalent to the number of cores per node
-num_jobs_per_node = 2
-
-num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
-
-outfile_name = '{}.json'.format(datetime.now().isoformat())
-results = []
-
 for n in range(min_n, max_n):
     for p in range(min_p, max_p):
         for d in range(min_d, max_d):

From 9d1dee26568ac0a8e7fa12e46563b1aca7028d09 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Tue, 11 Oct 2022 09:56:52 -0500
Subject: [PATCH 10/24] fixed correct number of circuit check. Only checks now
 if we choose to recompute previous results, instead of all the time.

---
 qtensor/noise_simulator/ComparisonSimulator.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
index 7dba8734..26d42ad1 100644
--- a/qtensor/noise_simulator/ComparisonSimulator.py
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -68,6 +68,7 @@ def qtensor_qiskit_noisy_qaoa(self, recompute_previous_ensemble: bool = False, p
                 self.qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
             
             if recompute_previous_ensemble == True:
+                #self._check_correct_num_circs_simulated(i)
                 self.prev_probs = self.qtensor_probs
 
             qtensor_sim_time = noise_sim.time_taken
@@ -104,7 +105,6 @@ def qtensor_qiskit_noisy_qaoa_mpi(self, num_nodes: int, num_jobs_per_node: int,
                 self.noise_model_qtensor, self.n, self.p, self.d)
             self._get_args(i)
             qtensor_probs_list = tools.mpi.mpi_map(self._mpi_parallel_unit, self._arggen, pbar=pbar, total=num_nodes)
-            self._check_correct_num_circs_simulated(i)
             if qtensor_probs_list:
                 if i == 0 or recompute_previous_ensemble == False: 
                     self.qtensor_probs = sum(qtensor_probs_list) / self._total_jobs
@@ -115,6 +115,7 @@ def qtensor_qiskit_noisy_qaoa_mpi(self, num_nodes: int, num_jobs_per_node: int,
                     self.qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
                 qtensor_sim_time = self.noise_sim.time_taken
                 if recompute_previous_ensemble == True:
+                    self._check_correct_num_circs_simulated(i)
                     self.prev_probs = self.qtensor_probs
 
                 self.simulate_qiskit_density_matrix(self.qiskit_circ, self.noise_model_qiskit)
@@ -138,7 +139,7 @@ def qtensor_qiskit_noisy_qaoa_density(self, recompute_previous_ensemble: bool =
         for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
             result = NoiseSimComparisonResult(self.qiskit_circ, self.qtensor_circ, self.noise_model_qiskit, 
                 self.noise_model_qtensor, self.n, self.p, self.d)
-            if i == 0 or recompute_previous_ensemble == True: 
+            if i == 0 or recompute_previous_ensemble == False: 
                 self.num_circs_simulated.append(num_circs)
                 noise_sim.simulate_batch_ensemble_density(self.qtensor_circ, num_circs, self.n)
                 self.qtensor_density_matrix = noise_sim.normalized_ensemble_density_matrix

From 17bbc13b76441af47eeed4fe8525e59343190095 Mon Sep 17 00:00:00 2001
From: Danil <mail@example.com>
Date: Tue, 11 Oct 2022 17:16:47 +0000
Subject: [PATCH 11/24] add hpc scripts

---
 qtensor/noise_simulator/hpc/entry.sh | 4 ++++
 qtensor/noise_simulator/hpc/run.sh   | 2 ++
 2 files changed, 6 insertions(+)
 create mode 100755 qtensor/noise_simulator/hpc/entry.sh
 create mode 100755 qtensor/noise_simulator/hpc/run.sh

diff --git a/qtensor/noise_simulator/hpc/entry.sh b/qtensor/noise_simulator/hpc/entry.sh
new file mode 100755
index 00000000..16b766a2
--- /dev/null
+++ b/qtensor/noise_simulator/hpc/entry.sh
@@ -0,0 +1,4 @@
+#!/bin/bash
+source $HOME/.profile
+mpiexec -n 250 python $HOME/anl/Qensor/qtensor/noise_simulator/comparison_mpi.py
+
diff --git a/qtensor/noise_simulator/hpc/run.sh b/qtensor/noise_simulator/hpc/run.sh
new file mode 100755
index 00000000..2e84d13f
--- /dev/null
+++ b/qtensor/noise_simulator/hpc/run.sh
@@ -0,0 +1,2 @@
+#!/bin/bash
+qsub -t 300 -n 5 -q skylake_8180 --cwd ~ $PWD/entry.sh

From 6f792dd41829d158f17149abd0d0dbdcf78e5838 Mon Sep 17 00:00:00 2001
From: Danil <mail@example.com>
Date: Tue, 11 Oct 2022 17:40:59 +0000
Subject: [PATCH 12/24] change to n=5

---
 qtensor/noise_simulator/comparison_mpi.py | 4 ++--
 qtensor/noise_simulator/hpc/entry.sh      | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/qtensor/noise_simulator/comparison_mpi.py b/qtensor/noise_simulator/comparison_mpi.py
index 482f6d7b..fa20447b 100644
--- a/qtensor/noise_simulator/comparison_mpi.py
+++ b/qtensor/noise_simulator/comparison_mpi.py
@@ -34,7 +34,7 @@
 """
 num_samples = 30
 
-num_nodes = 10
+num_nodes = 5
 
 """num jobs per node is equivalent to the number of cores per node"""
 num_jobs_per_node = 56
@@ -71,4 +71,4 @@
                 comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node)
                 results.extend(comparison.results)
 
-save_dict_to_file(results, name = outfile_name) 
\ No newline at end of file
+save_dict_to_file(results, name = outfile_name) 
diff --git a/qtensor/noise_simulator/hpc/entry.sh b/qtensor/noise_simulator/hpc/entry.sh
index 16b766a2..957c35bb 100755
--- a/qtensor/noise_simulator/hpc/entry.sh
+++ b/qtensor/noise_simulator/hpc/entry.sh
@@ -1,4 +1,4 @@
 #!/bin/bash
 source $HOME/.profile
-mpiexec -n 250 python $HOME/anl/Qensor/qtensor/noise_simulator/comparison_mpi.py
+mpiexec -n 5 python $HOME/anl/Qensor/qtensor/noise_simulator/comparison_mpi.py
 

From 3dbbadba814fb66324915c2e2225f8ff06495f78 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Tue, 11 Oct 2022 14:14:13 -0500
Subject: [PATCH 13/24] temporary fix for _get_args(). Needs a full rewrite
 eventually

---
 qtensor/noise_simulator/ComparisonSimulator.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
index 26d42ad1..2191dcda 100644
--- a/qtensor/noise_simulator/ComparisonSimulator.py
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -162,6 +162,9 @@ def qtensor_qiskit_noisy_qaoa_density(self, recompute_previous_ensemble: bool =
 
 
     # Prepare arguments to be sent to each unit of work 
+    """Function needs a rewrite. If mod(num_circs/num_circs_per_job) !=, the job 
+    is not bounded above in the number of circuits it can simulate. The proper way for this to run
+    is for num_circuits_in_last_job <= num_circuits_per_job """
     def _get_args(self, i):
         if i == 0 or self.recompute_previous_ensemble == False:
             num_circs = self.num_circs_list[i]
@@ -188,12 +191,21 @@ def _get_args(self, i):
             self.num_circs_simulated.append(num_circs)
             #print("num_circs: {}, actual num_circs simulated on this iteration: {}".format(self.num_circs_list[i], num_circs))
         else: 
+            # self._arggen = list(zip(repeat(self.noise_sim, total_jobs - 1), repeat(self.qtensor_circ, total_jobs - 1), 
+            #     repeat(num_circs_per_job, total_jobs - 1), repeat(self.n, total_jobs - 1)))
+            # num_circs_in_last_job = num_circs % num_circs_per_job
             self._arggen = list(zip(repeat(self.noise_sim, total_jobs - 1), repeat(self.qtensor_circ, total_jobs - 1), 
-                repeat(num_circs_per_job, total_jobs - 1), repeat(self.n, total_jobs - 1)))
-            num_circs_in_last_job = num_circs % num_circs_per_job
+                repeat(num_circs_per_job - 1, total_jobs - 1), repeat(self.n, total_jobs - 1)))
+            if num_circs_per_job == min_circs_per_job:
+                num_circs_in_last_job = num_circs % num_circs_per_job
+                actual_num_circs = (total_jobs - 1) * num_circs_per_job + num_circs_in_last_job
+            else: 
+                num_circs_in_last_job = num_circs - (total_jobs - 1) * (num_circs_per_job - 1)
+                actual_num_circs = (total_jobs - 1) * (num_circs_per_job - 1) + num_circs_in_last_job
             self._arggen.append((self.noise_sim, self.qtensor_circ, num_circs_in_last_job, self.n))
             self._total_jobs = total_jobs
-            actual_num_circs = (total_jobs - 1) * num_circs_per_job + num_circs_in_last_job
+            print("actual_num_circs: {}, num_circs: {}, num_nodes: {}, num_jobs_per_node: {}, total_jobs: {}, num_circs_per_job: {}, num_circs_in_last_job: {}".format(
+                actual_num_circs, num_circs, self.num_nodes, self.num_jobs_per_node, self._total_jobs, num_circs_per_job -1 , num_circs_in_last_job))
             assert num_circs == actual_num_circs
             self.num_circs_simulated.append(actual_num_circs)
             #print("num_circs: {}, actual num_circs simulated on this iteration: {}, total jobs: {}".format(self.num_circs_list[i], actual_num_circs, total_jobs))

From 4aeb6b57d6b510f716cbd2385013a2f593ba0394 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Tue, 11 Oct 2022 15:04:28 -0500
Subject: [PATCH 14/24] added a better fix to _get_args()

---
 .../noise_simulator/ComparisonSimulator.py    | 28 +++++++++----------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
index 2191dcda..463dc384 100644
--- a/qtensor/noise_simulator/ComparisonSimulator.py
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -162,9 +162,6 @@ def qtensor_qiskit_noisy_qaoa_density(self, recompute_previous_ensemble: bool =
 
 
     # Prepare arguments to be sent to each unit of work 
-    """Function needs a rewrite. If mod(num_circs/num_circs_per_job) !=, the job 
-    is not bounded above in the number of circuits it can simulate. The proper way for this to run
-    is for num_circuits_in_last_job <= num_circuits_per_job """
     def _get_args(self, i):
         if i == 0 or self.recompute_previous_ensemble == False:
             num_circs = self.num_circs_list[i]
@@ -180,7 +177,7 @@ def _get_args(self, i):
             total_jobs = int(np.ceil(num_circs / num_circs_per_job))
         else: 
             total_jobs = self.num_nodes * self.num_jobs_per_node
-            num_circs_per_job = int(np.ceil(num_circs / total_jobs))
+            num_circs_per_job = int(np.floor(num_circs / total_jobs))
 
         ## We make sure that regardless of how many nodes and jobs per node we have, we always 
         ## simulate the exact number of circuits in the ensemble specified. 
@@ -191,21 +188,24 @@ def _get_args(self, i):
             self.num_circs_simulated.append(num_circs)
             #print("num_circs: {}, actual num_circs simulated on this iteration: {}".format(self.num_circs_list[i], num_circs))
         else: 
-            # self._arggen = list(zip(repeat(self.noise_sim, total_jobs - 1), repeat(self.qtensor_circ, total_jobs - 1), 
-            #     repeat(num_circs_per_job, total_jobs - 1), repeat(self.n, total_jobs - 1)))
-            # num_circs_in_last_job = num_circs % num_circs_per_job
-            self._arggen = list(zip(repeat(self.noise_sim, total_jobs - 1), repeat(self.qtensor_circ, total_jobs - 1), 
-                repeat(num_circs_per_job - 1, total_jobs - 1), repeat(self.n, total_jobs - 1)))
             if num_circs_per_job == min_circs_per_job:
+                self._arggen = list(zip(repeat(self.noise_sim, total_jobs - 1), repeat(self.qtensor_circ, total_jobs - 1), 
+                    repeat(num_circs_per_job, total_jobs - 1), repeat(self.n, total_jobs - 1)))
                 num_circs_in_last_job = num_circs % num_circs_per_job
                 actual_num_circs = (total_jobs - 1) * num_circs_per_job + num_circs_in_last_job
+                self._arggen.append((self.noise_sim, self.qtensor_circ, num_circs_in_last_job, self.n))
             else: 
-                num_circs_in_last_job = num_circs - (total_jobs - 1) * (num_circs_per_job - 1)
-                actual_num_circs = (total_jobs - 1) * (num_circs_per_job - 1) + num_circs_in_last_job
-            self._arggen.append((self.noise_sim, self.qtensor_circ, num_circs_in_last_job, self.n))
+                first_set_of_jobs = num_circs - (total_jobs * num_circs_per_job)
+                second_set_of_jobs = total_jobs - first_set_of_jobs
+                print("first_set_of_jobs: {}, secomd_set_of_jobs: {}".format(first_set_of_jobs, second_set_of_jobs))
+                self._arggen = list(zip(repeat(self.noise_sim, first_set_of_jobs), repeat(self.qtensor_circ,  first_set_of_jobs), 
+                    repeat(num_circs_per_job + 1,  first_set_of_jobs), repeat(self.n, first_set_of_jobs)))
+                self._arggen.extend(list(zip(repeat(self.noise_sim, second_set_of_jobs), repeat(self.qtensor_circ,  second_set_of_jobs), 
+                    repeat(num_circs_per_job,  second_set_of_jobs), repeat(self.n, second_set_of_jobs))))  
+                actual_num_circs = first_set_of_jobs * (num_circs_per_job + 1) + second_set_of_jobs * num_circs_per_job
             self._total_jobs = total_jobs
-            print("actual_num_circs: {}, num_circs: {}, num_nodes: {}, num_jobs_per_node: {}, total_jobs: {}, num_circs_per_job: {}, num_circs_in_last_job: {}".format(
-                actual_num_circs, num_circs, self.num_nodes, self.num_jobs_per_node, self._total_jobs, num_circs_per_job -1 , num_circs_in_last_job))
+            print("actual_num_circs: {}, num_circs: {}, num_nodes: {}, num_jobs_per_node: {}, total_jobs: {}".format(
+                actual_num_circs, num_circs, self.num_nodes, self.num_jobs_per_node, self._total_jobs))
             assert num_circs == actual_num_circs
             self.num_circs_simulated.append(actual_num_circs)
             #print("num_circs: {}, actual num_circs simulated on this iteration: {}, total jobs: {}".format(self.num_circs_list[i], actual_num_circs, total_jobs))

From 089223f4119ce99e7e2d9540be7b21b055bf5e78 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Fri, 20 Jan 2023 13:12:02 -0600
Subject: [PATCH 15/24] adding parallelized tensor slicing file. Can be used on
 multi gpu nodes. Mostly for learning

---
 .../node_parallel/node_parallelism_mpi.py     | 243 ++++++++++++++++++
 1 file changed, 243 insertions(+)
 create mode 100644 scratchpad/node_parallel/node_parallelism_mpi.py

diff --git a/scratchpad/node_parallel/node_parallelism_mpi.py b/scratchpad/node_parallel/node_parallelism_mpi.py
new file mode 100644
index 00000000..9cdcc040
--- /dev/null
+++ b/scratchpad/node_parallel/node_parallelism_mpi.py
@@ -0,0 +1,243 @@
+import pyrofiler as pyrof
+from pyrofiler.pyrofiler import Profiler
+from pyrofiler import callbacks
+import matplotlib.pyplot as plt
+import matplotlib.pyplot as plt
+import seaborn as sns
+import numpy as np
+import sys
+from multiprocessing import Pool, Array
+from multiprocessing.dummy import Pool as ThreadPool
+import os
+from qtensor import tools
+from itertools import repeat
+import subprocess
+
+sns.set_style('whitegrid')
+np.random.seed(42)
+# pool = ThreadPool(processes=2**7)
+
+
+
+def _none_slice():
+    return slice(None)
+
+def _get_idx(x, idxs, slice_idx, shapes=None):
+    if shapes is None:
+        shapes = [2]*len(idxs)
+    point = np.unravel_index(slice_idx, shapes)
+    get_point = {i:p for i,p in zip(idxs, point)}
+    if x in idxs:
+        p = get_point[x]
+        return slice(p,p+1)
+    else:
+        return _none_slice()
+
+def _slices_for_idxs(idxs, *args, shapes=None, slice_idx=0):
+    """Return array of slices along idxs"""
+    slices = []
+    for indexes in args:
+        _slice = [_get_idx(x, idxs, slice_idx, shapes) for x in indexes ]
+        slices.append(tuple(_slice))
+    return slices
+        
+def get_example_task(A=8, B=10, C=7, dim1=0):
+    shape1 = [2]*(A+B)
+    shape2 = [2]*(A+C)
+    for i in range(dim1):
+        shape1[-i] = 1
+        shape2[-i] = 1
+
+    T1 = np.random.randn(*shape1)
+    T2 = np.random.randn(*shape2)
+    common = list(range(A))
+    idxs1 = common + list(range(A, A+B))
+    idxs2 = common + list(range(A+B, A+B+C))
+    return (T1, idxs1), (T2, idxs2)
+
+def contract(A, B, output = None):
+    a, idxa = A
+    b, idxb = B
+    contract_idx = set(idxa) & set(idxb)
+    result_idx = set(idxa + idxb)
+    if output is not None:
+        f'{output}\n contract result idx: {result_idx}'
+    else: 
+        f'contract result idx: {result_idx}'
+    C = np.einsum(a,idxa, b,idxb, result_idx)
+    return C
+
+def sliced_contract(x, y, idxs, num, output = None):
+    slices = _slices_for_idxs(idxs, x[1], y[1], slice_idx=num)
+    a = x[0][slices[0]]
+    b = y[0][slices[1]]
+    if output is not None:
+        with pyrof.timing(f'{output}\n \tcontract sliced {num}'):
+            C = contract((a, x[1]), (b, y[1]), output)
+    else: 
+        with pyrof.timing(f'\tcontract sliced {num}'):
+            C = contract((a, x[1]), (b, y[1]))
+    return C
+
+def target_slice(result_idx, idxs, num):
+    slices = _slices_for_idxs(idxs, result_idx, slice_idx=num)
+    return slices
+
+def __contract_bound(A, B):
+    a, idxa = A
+    b, idxb = B
+    contract_idx = set(idxa) & set(idxb)
+    def glue_first(shape):
+        sh = [shape[0] * shape[1]] + list(shape[2:])
+        return sh
+    
+    result_idx = set(idxa + idxb)
+    _map_a = {k:v for k,v in zip(idxa, a.shape)}
+    _map_b = {k:v for k,v in zip(idxb, b.shape)}
+    _map = {**_map_a, **_map_b}
+    print(_map)
+    result_idx = sorted(tuple(_map.keys()))
+    target_shape = tuple([_map[i] for i in result_idx])
+    
+    
+            
+    _dimlen = len(result_idx)
+    _maxdims = 22
+    print('dimlen',_dimlen)
+    new_a, new_b = a.shape, b.shape
+    if _dimlen>_maxdims:
+        _contr_dim = _dimlen - _maxdims
+        print(len(new_a), len(new_b))
+        for i in range(_contr_dim):
+            idxa = idxa[1:]
+            idxb = idxb[1:]
+                    
+            new_a = glue_first(new_a)
+            new_b = glue_first(new_b)
+            
+    _map_a = {k:v for k,v in zip(idxa, a.shape)}
+    _map_b = {k:v for k,v in zip(idxb, b.shape)}
+    _map = {**_map_a, **_map_b}
+    print(_map)
+    result_idx = sorted(tuple(_map.keys()))
+    print(len(new_a), len(new_b))
+    a = a.reshape(new_a)
+    b = b.reshape(new_b)
+    print(a.shape, b.shape)
+    print(idxa, idxb)
+    print('btsh',result_idx, target_shape)
+        
+        
+    C = np.einsum(a,idxa, b,idxb, result_idx)
+    
+    return C.reshape(*target_shape)
+
+def __add_dims(x, dims, ofs):
+    arr, idxs = x
+    arr = arr.reshape(list(arr.shape) + [1]*dims)
+    md = max(idxs)
+    return arr, idxs + list(range(md+ofs, ofs+md+dims)) 
+
+# def _mpi_unit(args):
+#     i,x,y, par_vars, result_idx, target_shape =  args
+#     patch = sliced_contract(x, y, par_vars, i)
+#     sl = target_slice(result_idx, par_vars, i)
+#     os.global_C[sl[0]] = patch
+
+def _mpi_unit(args):
+    output = subprocess.getoutput('echo echo')
+    #output = subprocess.getoutput('echo “RANK= ${PMI_RANK} LOCAL_RANK= ${PMI_LOCAL_RANK} gpu= $((${PMI_LOCAL_RANK} % $(nvidia-smi -L | wc -l)))”')
+    i,x,y, par_vars, result_idx, target_shape =  args
+    patch = sliced_contract(x, y, par_vars, i, output)
+    sl = target_slice(result_idx, par_vars, i)
+    C_par = np.empty(target_shape)
+    C_par[sl[0]] = patch
+    return C_par
+
+def parallel_contract(x, y, num_jobs, pbar = False, print_stats = True):
+    # x, y = get_example_task(A=10)
+
+    prof_seq = Profiler()
+    prof_seq.use_append()
+
+    contract_idx = set(x[1]) & set(y[1])
+    result_idx = set(x[1] + y[1])
+
+    arggen = i , 
+    for i in range(1):
+        _ = contract(x,y)
+    for rank in range(1,7):
+        with prof_seq.timing('Single thread'):
+            C = contract(x,y)
+        
+        par_vars = list(range(rank))
+        print(par_vars)
+        target_shape = C.shape
+
+        with prof_seq.timing('One patch: total'):
+            i = 0
+            with prof_seq.timing('One patch: compute'):
+                patch = sliced_contract(x, y, par_vars, i)
+            C_par = np.empty(target_shape)
+            with prof_seq.timing('One patch: assign'):
+                _slice = target_slice(result_idx, par_vars, i)
+                C_par[_slice[0]] = patch
+
+def parallel_contract_mpi_old(x, y, num_jobs, pbar = False, print_stats = True):
+    # x, y = get_example_task(A=10)
+
+    prof_seq = Profiler()
+    prof_seq.use_append()
+
+    contract_idx = set(x[1]) & set(y[1])
+    result_idx = set(x[1] + y[1])
+
+    i = 0
+    for rank in range(1,7):
+        with prof_seq.timing('Single thread'):
+            C = contract(x,y)
+        
+        par_vars = list(range(rank))
+        target_shape = C.shape
+        os.global_C = np.empty(target_shape, dtype=x[0].dtype)
+        # arggen = [i,x,y, par_vars, result_idx]
+        # arggen = (i,x,y,par_vars,result_idx) for i in range(num_jobs)
+        arggen = list(zip(repeat(i, num_jobs), repeat(x, num_jobs), repeat(y, num_jobs), repeat(par_vars, num_jobs), repeat(result_idx, num_jobs), repeat(target_shape, num_jobs)))
+        result = tools.mpi.mpi_map(_mpi_unit, arggen, pbar=pbar, total=num_jobs)
+        if result: 
+            pass
+            # result = tools.mpi.mpi_map(_mpi_unit, arggen, pbar=pbar, total=num_jobs)
+            # if print_stats:
+            #     tools.mpi.print_stats()
+            # print(os.global_C)
+            # return os.global_C
+            # print(result)
+            # return result
+
+def parallel_contract_mpi(x, y, num_jobs, pbar = False, print_stats = True):
+    """Get args for slicing"""
+    result_idx = set(x[1] + y[1])
+    #result_idx = set([i for i in x[1] if i not in set(y[1])])
+    print(x[1], "\n", y[1], '\n', result_idx)
+    idx_shape_pairs = (x[1] + y[1], x[0].shape + y[0].shape)
+    shape_dict = {i:s for i, s in zip(*idx_shape_pairs)}
+    target_shape = [shape_dict[i] for i in result_idx]
+    par_vars = np.arange(1,7)
+    i = 0
+
+    """Create a set of the args to send out to workers"""
+    arggen = list(zip(repeat(i, num_jobs), repeat(x, num_jobs), repeat(y, num_jobs), repeat(par_vars, num_jobs), repeat(result_idx, num_jobs), repeat(target_shape, num_jobs)))
+    
+    """Send out slicing jobs to workers"""
+    result = tools.mpi.mpi_map(_mpi_unit, arggen, pbar=pbar, total=num_jobs)
+
+    if result: 
+        pass
+
+
+if __name__ == '__main__':
+    x, y = get_example_task(A=8, B = 9)
+    num_nodes = 2
+    num_jobs_per_node = 4
+    num_jobs = num_nodes * num_jobs_per_node
+    parallel_contract_mpi(x, y, num_jobs)

From a7b8c9473e22d2e069903ef5da5134e101de3850 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Fri, 20 Jan 2023 13:51:41 -0600
Subject: [PATCH 16/24] small updates

---
 .../noise_simulator/ComparisonSimulator.py    |  4 +++-
 .../NoiseSimComparisonResult.py               |  8 ++++----
 qtensor/noise_simulator/comparison_mpi.py     | 19 ++++++++++---------
 qtensor/noise_simulator/helper_functions.py   |  3 +++
 4 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
index 463dc384..c36ed2d8 100644
--- a/qtensor/noise_simulator/ComparisonSimulator.py
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -87,6 +87,7 @@ def _mpi_parallel_unit(self, args):
         #print("this workers num_circs: ", num_circs)
         noise_sim.simulate_batch_ensemble(sum(qtensor_circ, []), num_circs, num_qubits)
         fraction_of_qtensor_probs = noise_sim.normalized_ensemble_probs
+        # fraction_of_qtensor_probs = [0] * 2**self.n
         return fraction_of_qtensor_probs
 
     def qtensor_qiskit_noisy_qaoa_mpi(self, num_nodes: int, num_jobs_per_node: int, recompute_previous_ensemble: bool = False, print_stats: bool = True, pbar: bool = True):
@@ -155,7 +156,7 @@ def qtensor_qiskit_noisy_qaoa_density(self, recompute_previous_ensemble: bool =
 
             if recompute_previous_ensemble == False: 
                 self.prev_qtensor_density_matrix = self.qtensor_density_matrix
-            # Save results
+            #Save results
             result.save_results_density(self.qiskit_density_matrix, self.qtensor_density_matrix, num_circs, 
                 self.num_circs_simulated[i], G, gamma, beta, qtensor_sim_time, self.qiskit_sim_time)
             self.results.append(result.data)
@@ -209,6 +210,7 @@ def _get_args(self, i):
             assert num_circs == actual_num_circs
             self.num_circs_simulated.append(actual_num_circs)
             #print("num_circs: {}, actual num_circs simulated on this iteration: {}, total jobs: {}".format(self.num_circs_list[i], actual_num_circs, total_jobs))
+        print("total_jobs: ", total_jobs)
 
 
     def simulate_qiskit_density_matrix(self, circuit, noise_model_qiskit, take_trace = True):
diff --git a/qtensor/noise_simulator/NoiseSimComparisonResult.py b/qtensor/noise_simulator/NoiseSimComparisonResult.py
index 6c418cbe..31e63bdd 100644
--- a/qtensor/noise_simulator/NoiseSimComparisonResult.py
+++ b/qtensor/noise_simulator/NoiseSimComparisonResult.py
@@ -108,16 +108,16 @@ def _calc_similarity_of_probs(self):
         # real by the time they have made it here. 
 
         """Measures the fidelity between the qiskit density matrix noisy state and the qtensoir stochastic noisy state"""
-        self.noisy_fidelity = (np.inner(qiskit_probs_root, qtensor_probs_root))**2
+        self.noisy_fidelity = np.abs((np.inner(qiskit_probs_root, qtensor_probs_root)))**2
 
         """Measures the fidelity between a noisy qtensor state and the noiseless version of the same state"""
-        self.noiseless_fidelity = (np.inner(noiseless_qtensor_probs_root, qtensor_probs_root))**2
+        self.noiseless_fidelity = np.abs((np.inner(noiseless_qtensor_probs_root, qtensor_probs_root)))**2
 
         """Measures the fidelity between a qiskit noisy state and a uniform distribution state"""
-        self.uniform_qiskit_fidelity = (np.inner(uniform_probs_root, qiskit_probs_root))**2
+        self.uniform_qiskit_fidelity = np.abs((np.inner(uniform_probs_root, qiskit_probs_root)))**2
 
         """Measures the fidelity between a qtensor noisy state and a uniform distribution state"""
-        self.uniform_qtensor_fidelity = (np.inner(uniform_probs_root, qtensor_probs_root))**2
+        self.uniform_qtensor_fidelity = np.abs((np.inner(uniform_probs_root, qtensor_probs_root)))**2
         self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
 
     def _to_dict(self):
diff --git a/qtensor/noise_simulator/comparison_mpi.py b/qtensor/noise_simulator/comparison_mpi.py
index fa20447b..3c0baefa 100644
--- a/qtensor/noise_simulator/comparison_mpi.py
+++ b/qtensor/noise_simulator/comparison_mpi.py
@@ -32,26 +32,27 @@
 num samples determines how many samples of a particular n, p, d, and num_circs we take
 generally we want to get >30 to reduce sampling noise 
 """
-num_samples = 30
+num_samples = 1
 
-num_nodes = 5
+num_nodes = 2
 
 """num jobs per node is equivalent to the number of cores per node"""
-num_jobs_per_node = 56
+num_jobs_per_node = 1
 
-num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
+# num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
+num_circs_list = [1000]
 
 outfile_name = '{}.json'.format(datetime.now().isoformat())
 results = []
 
 min_n = 5
-max_n = 7
+max_n = 6
 
 min_p = 2
 max_p = 3
 
-min_d = 2
-max_d = 3
+min_d = 4
+max_d = 5
 
 for n in range(min_n, max_n):
     for p in range(min_p, max_p):
@@ -70,5 +71,5 @@
                 comparison = QAOAComparisonSimulator(n, p, d, noise_model_qiskit, noise_model_qtensor, num_circs_list)
                 comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node)
                 results.extend(comparison.results)
-
-save_dict_to_file(results, name = outfile_name) 
+if results:
+    save_dict_to_file(results, name = outfile_name) 
diff --git a/qtensor/noise_simulator/helper_functions.py b/qtensor/noise_simulator/helper_functions.py
index 84e9750e..a98586e3 100644
--- a/qtensor/noise_simulator/helper_functions.py
+++ b/qtensor/noise_simulator/helper_functions.py
@@ -89,3 +89,6 @@ def get_total_jobs(num_circs: int, num_nodes: int, num_jobs_per_node: int, min_c
         num_circs_per_job = int(np.ceil(num_circs / total_jobs))
     return total_jobs, num_circs_per_job
 
+
+if __name__ == '__main__':
+    G, gamma, beta = get_qaoa_params(3,2,2)
\ No newline at end of file

From d0c05c42ba4fb5c934b631e01086175eed6c0469 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Fri, 20 Jan 2023 17:03:17 -0600
Subject: [PATCH 17/24] fixed imports

---
 .../noise_simulator/ComparisonSimulator.py    | 20 ++++++++++---------
 qtensor/noise_simulator/NoiseModel.py         |  4 ----
 .../NoiseSimComparisonResult.py               |  4 +++-
 qtensor/noise_simulator/NoiseSimulator.py     |  9 +++++----
 qtensor/noise_simulator/comparison_mpi.py     | 13 +++++-------
 qtensor/noise_simulator/helper_functions.py   |  5 -----
 6 files changed, 24 insertions(+), 31 deletions(-)

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
index c36ed2d8..1df9be7d 100644
--- a/qtensor/noise_simulator/ComparisonSimulator.py
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -1,16 +1,18 @@
-from NoiseChannels import *
-from NoiseSimulator import * 
-from NoiseModel import *
-from NoiseSimComparisonResult import *
-from helper_functions import *
-from qiskit import execute, Aer
-from qtensor.tests.test_composers import *
+from NoiseSimulator import NoiseSimulator
+from NoiseModel import NoiseModel 
+from NoiseSimComparisonResult import NoiseSimComparisonResult
+from helper_functions import get_qaoa_params
+from qtensor.tests.test_composers import QtreeSimulator
 from qtensor import QiskitQAOAComposer
 from qtree.operators import from_qiskit_circuit
 from qtensor import tools
+
+from qiskit import execute
+from qiskit.providers.aer import AerSimulator, noise
+
+import numpy as np
 from itertools import repeat
-from qiskit import execute, Aer
-from qiskit.providers.aer import AerSimulator
+import time
 
 class ComparisonSimulator:
     def __init__(self):
diff --git a/qtensor/noise_simulator/NoiseModel.py b/qtensor/noise_simulator/NoiseModel.py
index d6fb68ca..242829ba 100644
--- a/qtensor/noise_simulator/NoiseModel.py
+++ b/qtensor/noise_simulator/NoiseModel.py
@@ -1,7 +1,3 @@
-#from qtensor.noise_simulator import NoiseGate
-# from NoiseGate import *
-# import attr
-
 class NoiseModel:
     _1qubit_gate = set([
         'M', 'I', 'H', 'Z', 'T', 'Tdag', 'S', 'Sdag', 'X_1_2', 'Y_1_2', 'W_1_2'
diff --git a/qtensor/noise_simulator/NoiseSimComparisonResult.py b/qtensor/noise_simulator/NoiseSimComparisonResult.py
index 31e63bdd..db2ac9e1 100644
--- a/qtensor/noise_simulator/NoiseSimComparisonResult.py
+++ b/qtensor/noise_simulator/NoiseSimComparisonResult.py
@@ -1,10 +1,12 @@
-from helper_functions import *
+# from helper_functions import *
+from helper_functions import fidelity, cosine_similarity
 
 from datetime import datetime
 import time
 import json
 import jsbeautifier
 import jsonpickle
+import numpy as np
 
 from qtensor import QiskitQAOAComposer
 import qiskit.providers.aer.noise as noise
diff --git a/qtensor/noise_simulator/NoiseSimulator.py b/qtensor/noise_simulator/NoiseSimulator.py
index af58aaef..3e4dc872 100644
--- a/qtensor/noise_simulator/NoiseSimulator.py
+++ b/qtensor/noise_simulator/NoiseSimulator.py
@@ -1,7 +1,8 @@
-from qtensor.Simulate import *
-from qtensor.OpFactory import *
-from qtree import *
-from NoiseModel import *
+from qtensor.Simulate import QtreeSimulator, NumpyBackend
+import qtree
+from NoiseModel import NoiseModel
+
+import numpy as np
 import time
 
 class NoiseSimulator(QtreeSimulator):
diff --git a/qtensor/noise_simulator/comparison_mpi.py b/qtensor/noise_simulator/comparison_mpi.py
index 3c0baefa..fb3211ca 100644
--- a/qtensor/noise_simulator/comparison_mpi.py
+++ b/qtensor/noise_simulator/comparison_mpi.py
@@ -1,12 +1,9 @@
-from NoiseChannels import *
-from NoiseSimulator import * 
-from NoiseModel import *
-from NoiseSimComparisonResult import *
-from helper_functions import *
+from NoiseChannels import DepolarizingChannel
+from NoiseModel import NoiseModel
 import qiskit.providers.aer.noise as noise
-from qtensor.tests.test_composers import *
 from ComparisonSimulator import QAOAComparisonSimulator
-
+from datetime import datetime
+from helper_functions import save_dict_to_file
 
 ### Noise model, simulation, and results of a QAOA algorithmm  ###
 prob_1 = 0.003
@@ -40,7 +37,7 @@
 num_jobs_per_node = 1
 
 # num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
-num_circs_list = [1000]
+num_circs_list = [10]
 
 outfile_name = '{}.json'.format(datetime.now().isoformat())
 results = []
diff --git a/qtensor/noise_simulator/helper_functions.py b/qtensor/noise_simulator/helper_functions.py
index a98586e3..221b3182 100644
--- a/qtensor/noise_simulator/helper_functions.py
+++ b/qtensor/noise_simulator/helper_functions.py
@@ -1,6 +1,3 @@
-from logging import raiseExceptions
-from qtree.operators import *
-
 from collections import OrderedDict
 import numpy as np
 import numpy as np
@@ -9,8 +6,6 @@
 import json
 import jsbeautifier
 import networkx as nx
-from datetime import datetime
-from os import SEEK_SET
 
 import qtensor
 

From 57b9fd1eca0869c8d4522f5ed09d6926273888dd Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Sat, 21 Jan 2023 21:19:19 -0600
Subject: [PATCH 18/24] fixed some documentation

---
 qtensor/noise_simulator/NoiseChannels.py    | 17 ++++----
 qtensor/noise_simulator/NoiseModel.py       | 15 ++++---
 qtensor/noise_simulator/NoiseSimulator.py   | 23 +++++++---
 qtensor/noise_simulator/helper_functions.py | 48 +++++++++++++--------
 4 files changed, 65 insertions(+), 38 deletions(-)

diff --git a/qtensor/noise_simulator/NoiseChannels.py b/qtensor/noise_simulator/NoiseChannels.py
index 78261508..7a8f65b9 100644
--- a/qtensor/noise_simulator/NoiseChannels.py
+++ b/qtensor/noise_simulator/NoiseChannels.py
@@ -1,13 +1,15 @@
 import itertools as it
 
-# In a depolarizing channel, there is a depolarizing parameter lambda.
-# which is given by λ = 4^n / (4^n - 1) , where n is the number of qubits
-# when we run this function, we pass in a parameter term param, where 
-# 0 <= param <= λ. The probability of an error occuring is related to param and λ by
-# prob_depol_error = param/λ, and the probability of an particular pauli error 
-# from the depolarizing channel is given by param / 4^n 
-
 class DepolarizingChannel: 
+    """
+    In a depolarizing channel, there is a depolarizing parameter lambda
+    which is given by λ = 4^n / (4^n - 1) , where n is the number of qubits
+    when we run this function, we pass in a parameter term param, where 
+    0 <= param <= λ. The probability of an error occuring is related to param and λ by
+    prob_depol_error = param/λ, and the probability of an particular pauli error 
+    from the depolarizing channel is given by param / 4^n
+    """
+
     def __init__(self, param, num_qubits):
 
         num_terms = 4**num_qubits
@@ -28,7 +30,6 @@ def __init__(self, param, num_qubits):
         self._get_channel()
 
     def _get_channel(self):
-        
         pauli_error_prob = self.param / self.num_terms
         identity_prob = 1 - self.param / self.lam
         depol_channel_probs = [identity_prob] + (self.num_terms - 1) * [pauli_error_prob]
diff --git a/qtensor/noise_simulator/NoiseModel.py b/qtensor/noise_simulator/NoiseModel.py
index 242829ba..b3d8586c 100644
--- a/qtensor/noise_simulator/NoiseModel.py
+++ b/qtensor/noise_simulator/NoiseModel.py
@@ -16,8 +16,9 @@ def __init__(self):
         # not used currently. Just here in case it needs to be referenced for some later use. Channels is accessed in NoiseGate class
         self.channels = []
 
-    # applies noise to all qubits that a gate is acting on
     def add_channel_to_all_qubits(self, channel, gates):
+        """Applies noise to all qubits that a gate is acting on"""
+
         for gate in gates:
             if gate in self.noise_gates:
                 self._check_gate_with_channel(gate, channel)
@@ -31,10 +32,12 @@ def add_channel_to_all_qubits(self, channel, gates):
                 self.noise_gates[gate] = noise_gate
                 self.channels.append(channel)
 
-    # applies noise to only specific qubits that a gate is acting on. Only matters
-    # for multi qubit gates. Eg: if a gate is 'cx' applied to qubits (0, 1), but the
-    # list qubits = [1], then noise from channel will only be applied to qubit 1
     def add_channel_to_specific_qubits(self, channel, gates, qubits):
+        """Applies noise to only specific qubits that a gate is acting on. 
+        
+        Only matters for multi qubit gates. Eg: if a gate is 'cx' applied to qubits (0, 1), but the
+        list qubits = [1], then noise from channel will only be applied to qubit 1
+        """
         for gate in gates:
             if gate in self.noise_gates:
                 self._check_gate_with_channel(gate, channel)
@@ -49,6 +52,8 @@ def add_channel_to_specific_qubits(self, channel, gates, qubits):
                 self.channels.append(channel)
                 
     def _check_gate_with_channel(self, gate, channel):
+        """TODO: Add a warning or error if a bad gate name is given"""
+        
         if gate in self._1qubit_gate and channel.num_qubits != 1:
             raise Exception("{} qubit channel ".format(channel.num_qubits) + \
                             "cannot be applied to 1-qubit gate {}".format(gate))
@@ -58,7 +63,7 @@ def _check_gate_with_channel(self, gate, channel):
         if gate in self._3qubit_gate and channel.num_qubits != 3:
             raise Exception("{} qubit channel ".format(channel.num_qubits) + \
                             "cannot be applied to 3-qubit gate {}".format(gate))
-    ## TODO: Add a warning or error if a bad gate name is given    
+
 
 class NoiseGate: 
     def __init__(self, name):
diff --git a/qtensor/noise_simulator/NoiseSimulator.py b/qtensor/noise_simulator/NoiseSimulator.py
index 3e4dc872..6b4e5b4a 100644
--- a/qtensor/noise_simulator/NoiseSimulator.py
+++ b/qtensor/noise_simulator/NoiseSimulator.py
@@ -12,10 +12,15 @@ def __init__(self, noise_model, bucket_backend=NumpyBackend(), optimizer=None, m
             raise ValueError("Error: noise_model value must be of type NoiseModel")
         self.noise_model = noise_model
 
-    # If all you want is the probabiltiies, and not the amplitudes, this is a better function to call. It does not create the density matrix, 
-    # only a state vector, so it will take up much less memeory 
-    # The vector returned is dim(2^n), where n = number of qubits
+
     def simulate_batch_ensemble(self, qc, num_circs, batch_vars=0, peo=None):
+        """Stoachastic noise simulator using statevector approach
+        
+        This method uses significantly less memory than the density matrix approach (2^n vs 4^n). However, one should note that 
+        while this uses the statevector approach to simulate, the user never gets the amplitudes, only probabilities. 
+        This is because we must take the modulus squared of each statevector in the ensemble to conserve probability. Thus the user will 
+        receive a probability density vector, not a probability amplitude vector.  
+        """
         start = time.time_ns() / (10 ** 9)
         if num_circs < 0 or not isinstance(num_circs, int):
             raise Exception("Error: The argument num_circs must be a positive integer")
@@ -32,10 +37,12 @@ def simulate_batch_ensemble(self, qc, num_circs, batch_vars=0, peo=None):
         self.time_taken = end - start
         #return normalized_ensemble_probs
 
-    # This returns a density matrix, which contains all of the amplitudes of the final state. 
-    # If we want the probabilities, we can take the trace of that matrix
-    # The density matrix returned will be dimension m x m where m = 2^n and n = number of qubits  
     def simulate_batch_ensemble_density(self, qc, num_circs, batch_vars=0, peo=None):
+        """Stochastic noise simulator using density matrix apporach.
+        
+        Uses significantly more memory than the statevector approach, but gives more information. 
+        The user will get the relative phaseses as well as the probabilities.
+        """
         start = time.time_ns() / (10 ** 9)
         if num_circs < 0 or not isinstance(num_circs, int):
             raise Exception("Error: The argument num_circs must be a positive integer")
@@ -53,8 +60,8 @@ def simulate_batch_ensemble_density(self, qc, num_circs, batch_vars=0, peo=None)
         #normalized_ensemble_density_matrix = np.divide(unnormalized_ensemble_density_matrix, num_circs)
         #return normalized_ensemble_density_matrix
     
-    # Simulates and returns only the first amplitude of the ensemble  
     def simulate_ensemble(self, qc, num_circs):
+        """Simulates and returns only the first amplitude of the ensemble"""
         return self.simulate_state(qc, num_circs)
 
     def simulate_state_ensemble(self, qc, num_circs, peo=None):
@@ -69,6 +76,8 @@ def _new_circuit(self, ideal_qc):
                 self._apply_channel(gate)
         
     def _apply_channel(self, gate):
+        """A noisy gate has all of the proper noise channels applied to it"""
+        
         for i in range(len(self.noise_model.noise_gates[gate.name].channels)):
             error_name = self.noise_model.noise_gates[gate.name].channels[i].name
             if error_name == 'depolarizing':
diff --git a/qtensor/noise_simulator/helper_functions.py b/qtensor/noise_simulator/helper_functions.py
index 221b3182..192b57ad 100644
--- a/qtensor/noise_simulator/helper_functions.py
+++ b/qtensor/noise_simulator/helper_functions.py
@@ -10,17 +10,17 @@
 import qtensor
 
 def decimal_to_binary(n):
-# converting decimal to binary
-# and removing the prefix(0b)
+    """converts decimal to binary and removes the prefix (0b)"""
     return bin(n).replace("0b", "")
 
-def create_counts_dict(num_qubits, big_endian):
-    # get the binary length of num_qubits
+def create_counts_dict(num_qubits: int, big_endian: bool):
+    """Hxelper function for attach_qubit_names(). We create the dictionary that will eventually hold the probabilties."""
+
+    # get the binary length of num_qubits. length is used so that the length of each key is the same
     length = int(np.ceil(np.log(2**num_qubits + 1)/np.log(2)) - 1)
     counts = OrderedDict()
     for i in range(2**num_qubits):
         # convert to binary and then pad the right side with 0s
-        # length is used so that the length of each key is the same
         if big_endian == True:
             key_i = str(decimal_to_binary(i).zfill(length))
             key_i[::-1]
@@ -28,12 +28,21 @@ def create_counts_dict(num_qubits, big_endian):
         else:
             key_i = str(decimal_to_binary(i).zfill(length))
             counts[key_i] = 0
-    #print(counts)
     return counts
 
-# big_endian refers to the order that qubits are displayed 
-# Use big_endian = False if you want the qubit ordering that Qiskit uses 
-def attach_qubit_names(probs_list, big_endian = True):
+
+def attach_qubit_names(probs_list: list, big_endian: bool = True):
+    """Creates a dictionary of qubit names and probabilities from a probability list. 
+
+    Args:
+        probs_list (list): the list probabilities we are turning into a dictionary
+        big_endian (bool): the order that qubits are displayed.  
+                           use big_endian = False if you want the qubit ordering that Qiskit uses
+    
+    Returns: 
+        probs_dict (dict): a dictionary of qubit name (int): probability (float)
+    """
+
     num_qubits = int(np.log2(len(probs_list)))
     probs_dict = create_counts_dict(num_qubits, big_endian)
     for key, prob in zip(probs_dict, probs_list):
@@ -46,7 +55,7 @@ def fidelity(A, B):
 def cosine_similarity(A, B):
     return inner(A, B) / (norm(A)* norm(B)) 
 
-def get_qaoa_params(n, p, d):
+def get_qaoa_params(n: int, p: int, d: int):
     if (n * d) % 2 != 0:
         raise ValueError("n * d must be even.")
     if not 0 <= d < n:
@@ -55,13 +64,10 @@ def get_qaoa_params(n, p, d):
     print('Circuit params: n: {}, p: {}, d: {}'.format(n, p, d))
     G = nx.random_regular_graph(d, n)
     gammabeta = np.array(qtensor.tools.BETHE_QAOA_VALUES[str(d)]['angles'])
-    #gammabeta = gammabeta.tolist()
     gamma = gammabeta[:d]
     beta = gammabeta[d:]
     gamma = gamma.tolist()
     beta = beta.tolist()
-    #gamma = [np.random.uniform(0, 2*np.pi) for _ in range(p)]
-    #beta = [np.random.uniform(0, np.pi) for _ in range(p)]
     return G, gamma, beta
 
 def save_dict_to_file(dict, name):
@@ -70,11 +76,17 @@ def save_dict_to_file(dict, name):
     with open(name, 'a') as f:
         f.write(jsbeautifier.beautify(json.dumps(dict), options))
 
-def get_total_jobs(num_circs: int, num_nodes: int, num_jobs_per_node: int, min_circs_per_job = 10):
-    ## We make sure that each job has a minimum number of circuits. We do this because 
-    ## if there are too few circuits for each unit of work, the overhead from 
-    ## parallelization removes any advantage gained 
-    ## TODO: determine a better minimum number of circuits. Currently 10 is chosen arbitrarily
+def get_total_jobs(num_circs: int, num_nodes: int, num_jobs_per_node: int, min_circs_per_job: int = 10):
+    """Gets the total number of jobs for a simulation.
+    
+    We make sure that each job has a minimum number of circuits. We do this because 
+    if there are too few circuits for each unit of work, the overhead from 
+    parallelization removes any advantage gained 
+    
+    TODO: 
+        determine a better minimum number of circuits. Currently 10 is chosen arbitrarily
+    """
+
     min_circs_per_job = min(min_circs_per_job, num_circs)
     if num_nodes * num_jobs_per_node > num_circs / min_circs_per_job:
         num_circs_per_job = min_circs_per_job

From b6cf5c550ac1e967ecfbe32c9cdeb918ecdeaaf3 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Sun, 22 Jan 2023 12:39:15 -0600
Subject: [PATCH 19/24] added CupyBackend support for noisy simulations

---
 .../noise_simulator/ComparisonSimulator.py    | 21 +++++++++++--------
 qtensor/noise_simulator/NoiseSimulator.py     |  2 +-
 qtensor/noise_simulator/comparison_mpi.py     | 16 ++++++++------
 3 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
index 1df9be7d..93729361 100644
--- a/qtensor/noise_simulator/ComparisonSimulator.py
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -2,11 +2,14 @@
 from NoiseModel import NoiseModel 
 from NoiseSimComparisonResult import NoiseSimComparisonResult
 from helper_functions import get_qaoa_params
-from qtensor.tests.test_composers import QtreeSimulator
+# from qtensor.tests.test_composers import QtreeSimulator, NumpyBackend
+from qtensor.Simulate import QtreeSimulator, NumpyBackend
+from qtensor.contraction_backends import CuPyBackend
 from qtensor import QiskitQAOAComposer
 from qtree.operators import from_qiskit_circuit
 from qtensor import tools
 
+
 from qiskit import execute
 from qiskit.providers.aer import AerSimulator, noise
 
@@ -42,13 +45,13 @@ def __init__(self, n: int, p: int, d: int, noise_model_qiskit: noise.NoiseModel,
         self._check_params()
         self.num_circs_list.sort()
 
-    def qtensor_qiskit_noisy_qaoa(self, recompute_previous_ensemble: bool = False, print_stats: bool = False):
+    def qtensor_qiskit_noisy_qaoa(self, backend = NumpyBackend(), recompute_previous_ensemble: bool = False, print_stats: bool = False):
         self.recompute_previous_ensemble = recompute_previous_ensemble
         # Prepare circuits, simulator
         G, gamma, beta = get_qaoa_params(n = self.n, p = self.p, d = self.d)
         self._get_circuits(G, gamma, beta)
-        noise_sim = NoiseSimulator(self.noise_model_qtensor)
-        exact_sim = QtreeSimulator()
+        noise_sim = NoiseSimulator(self.noise_model_qtensor, bucket_backend=backend)
+        exact_sim = QtreeSimulator(backend=backend)
 
         # Run simulation
         for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
@@ -92,7 +95,7 @@ def _mpi_parallel_unit(self, args):
         # fraction_of_qtensor_probs = [0] * 2**self.n
         return fraction_of_qtensor_probs
 
-    def qtensor_qiskit_noisy_qaoa_mpi(self, num_nodes: int, num_jobs_per_node: int, recompute_previous_ensemble: bool = False, print_stats: bool = True, pbar: bool = True):
+    def qtensor_qiskit_noisy_qaoa_mpi(self,  num_nodes: int, num_jobs_per_node: int, backend = NumpyBackend(), recompute_previous_ensemble: bool = False, print_stats: bool = True, pbar: bool = True):
         self.num_nodes = num_nodes 
         self.num_jobs_per_node = num_jobs_per_node
         self.recompute_previous_ensemble = recompute_previous_ensemble
@@ -100,8 +103,8 @@ def qtensor_qiskit_noisy_qaoa_mpi(self, num_nodes: int, num_jobs_per_node: int,
         # Prepare circuit, simulator, and area to save results, 
         G, gamma, beta = get_qaoa_params(n = self.n, p = self.p, d = self.d)
         self._get_circuits(G, gamma, beta)
-        self.noise_sim = NoiseSimulator(self.noise_model_qtensor)
-        exact_sim = QtreeSimulator()
+        self.noise_sim = NoiseSimulator(self.noise_model_qtensor, bucket_backend=backend)
+        exact_sim = QtreeSimulator(backend=backend)
 
         for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
             result = NoiseSimComparisonResult(self.qiskit_circ, self.qtensor_circ, self.noise_model_qiskit, 
@@ -131,12 +134,12 @@ def qtensor_qiskit_noisy_qaoa_mpi(self, num_nodes: int, num_jobs_per_node: int,
                     tools.mpi.print_stats()
                     result.print_result()
 
-    def qtensor_qiskit_noisy_qaoa_density(self, recompute_previous_ensemble: bool = False):
+    def qtensor_qiskit_noisy_qaoa_density(self, backend = NumpyBackend(), recompute_previous_ensemble: bool = False):
         self.recompute_previous_ensemble = recompute_previous_ensemble
         # Prepare circuits, simulator, and area to save results
         G, gamma, beta = get_qaoa_params(n = self.n, p = self.p, d = self.d)
         self._get_circuits(G, gamma, beta)
-        noise_sim = NoiseSimulator(self.noise_model_qtensor)
+        noise_sim = NoiseSimulator(self.noise_model_qtensor, bucket_backend=backend)
         
         # Simulate
         for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
diff --git a/qtensor/noise_simulator/NoiseSimulator.py b/qtensor/noise_simulator/NoiseSimulator.py
index 6b4e5b4a..395f8395 100644
--- a/qtensor/noise_simulator/NoiseSimulator.py
+++ b/qtensor/noise_simulator/NoiseSimulator.py
@@ -77,7 +77,7 @@ def _new_circuit(self, ideal_qc):
         
     def _apply_channel(self, gate):
         """A noisy gate has all of the proper noise channels applied to it"""
-        
+
         for i in range(len(self.noise_model.noise_gates[gate.name].channels)):
             error_name = self.noise_model.noise_gates[gate.name].channels[i].name
             if error_name == 'depolarizing':
diff --git a/qtensor/noise_simulator/comparison_mpi.py b/qtensor/noise_simulator/comparison_mpi.py
index fb3211ca..c8432141 100644
--- a/qtensor/noise_simulator/comparison_mpi.py
+++ b/qtensor/noise_simulator/comparison_mpi.py
@@ -2,8 +2,9 @@
 from NoiseModel import NoiseModel
 import qiskit.providers.aer.noise as noise
 from ComparisonSimulator import QAOAComparisonSimulator
-from datetime import datetime
 from helper_functions import save_dict_to_file
+from qtensor.contraction_backends import CuPyBackend
+from datetime import datetime
 
 ### Noise model, simulation, and results of a QAOA algorithmm  ###
 prob_1 = 0.003
@@ -32,12 +33,14 @@
 num_samples = 1
 
 num_nodes = 2
+num_cores_per_node = 1
+num_threads_per_core = 1
+num_gpus_per_node = 4
 
-"""num jobs per node is equivalent to the number of cores per node"""
-num_jobs_per_node = 1
+num_jobs_per_node = num_cores_per_node * num_threads_per_core * num_gpus_per_node
 
-# num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
-num_circs_list = [10]
+num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
+#num_circs_list = [10]
 
 outfile_name = '{}.json'.format(datetime.now().isoformat())
 results = []
@@ -66,7 +69,8 @@
             for _ in range(num_samples):
                 print("\n\nnum_circs_list:", num_circs_list)
                 comparison = QAOAComparisonSimulator(n, p, d, noise_model_qiskit, noise_model_qtensor, num_circs_list)
-                comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node)
+                comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node, backend=CuPyBackend())
+                #comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node)
                 results.extend(comparison.results)
 if results:
     save_dict_to_file(results, name = outfile_name) 

From c8e29b249e2cc6a3b01f88afbb8dbacdfca0cdad Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Mon, 23 Jan 2023 11:43:34 -0600
Subject: [PATCH 20/24] tentative cupy functionality added. Needs refactoring

---
 .../noise_simulator/ComparisonSimulator.py    | 46 ++++++++-----
 .../NoiseSimComparisonResult.py               | 64 +++++++++++++------
 qtensor/noise_simulator/NoiseSimulator.py     | 22 +++++--
 qtensor/noise_simulator/comparison_mpi.py     |  4 +-
 4 files changed, 92 insertions(+), 44 deletions(-)

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
index 93729361..840eb5fc 100644
--- a/qtensor/noise_simulator/ComparisonSimulator.py
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -11,9 +11,10 @@
 
 
 from qiskit import execute
-from qiskit.providers.aer import AerSimulator, noise
+from qiskit.providers.aer import AerSimulator, noise, AerError
 
 import numpy as np
+from qtensor.tools.lazy_import import cupy as cp
 from itertools import repeat
 import time
 
@@ -56,7 +57,7 @@ def qtensor_qiskit_noisy_qaoa(self, backend = NumpyBackend(), recompute_previous
         # Run simulation
         for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
             result = NoiseSimComparisonResult(self.qiskit_circ, self.qtensor_circ, self.noise_model_qiskit, 
-                self.noise_model_qtensor, self.n, self.p, self.d)
+                self.noise_model_qtensor, self.n, self.p, self.d, backend)
 
             if i == 0 or recompute_previous_ensemble == False:
                 self.num_circs_simulated.append(num_circs)
@@ -89,7 +90,11 @@ def qtensor_qiskit_noisy_qaoa(self, backend = NumpyBackend(), recompute_previous
 
     def _mpi_parallel_unit(self, args):
         noise_sim, qtensor_circ, num_circs, num_qubits = args
-        #print("this workers num_circs: ", num_circs)
+        # #print("this workers num_circs: ", num_circs)
+        # if noise_sim.backend == CuPyBackend():
+        #     noise_sim.simulate_batch_ensemble(cupy.array(sum(qtensor_circ, [])), num_circs, num_qubits)
+        # else: 
+        #     noise_sim.simulate_batch_ensemble(sum(qtensor_circ, []), num_circs, num_qubits)
         noise_sim.simulate_batch_ensemble(sum(qtensor_circ, []), num_circs, num_qubits)
         fraction_of_qtensor_probs = noise_sim.normalized_ensemble_probs
         # fraction_of_qtensor_probs = [0] * 2**self.n
@@ -108,7 +113,7 @@ def qtensor_qiskit_noisy_qaoa_mpi(self,  num_nodes: int, num_jobs_per_node: int,
 
         for num_circs, i in zip(self.num_circs_list, range(len(self.num_circs_list))):
             result = NoiseSimComparisonResult(self.qiskit_circ, self.qtensor_circ, self.noise_model_qiskit, 
-                self.noise_model_qtensor, self.n, self.p, self.d)
+                self.noise_model_qtensor, self.n, self.p, self.d, backend)
             self._get_args(i)
             qtensor_probs_list = tools.mpi.mpi_map(self._mpi_parallel_unit, self._arggen, pbar=pbar, total=num_nodes)
             if qtensor_probs_list:
@@ -124,7 +129,7 @@ def qtensor_qiskit_noisy_qaoa_mpi(self,  num_nodes: int, num_jobs_per_node: int,
                     self._check_correct_num_circs_simulated(i)
                     self.prev_probs = self.qtensor_probs
 
-                self.simulate_qiskit_density_matrix(self.qiskit_circ, self.noise_model_qiskit)
+                self.simulate_qiskit_density_matrix(self.qiskit_circ, self.noise_model_qiskit, backend)
                 self.exact_qtensor_amps = exact_sim.simulate_batch(sum(self.qtensor_circ, []), batch_vars=self.num_qubits)
                 # Save results 
                 result.save_result(self.qiskit_probs, self.qtensor_probs, self.exact_qtensor_amps, num_circs,
@@ -203,7 +208,7 @@ def _get_args(self, i):
             else: 
                 first_set_of_jobs = num_circs - (total_jobs * num_circs_per_job)
                 second_set_of_jobs = total_jobs - first_set_of_jobs
-                print("first_set_of_jobs: {}, secomd_set_of_jobs: {}".format(first_set_of_jobs, second_set_of_jobs))
+                print("first_set_of_jobs: {}, second_set_of_jobs: {}".format(first_set_of_jobs, second_set_of_jobs))
                 self._arggen = list(zip(repeat(self.noise_sim, first_set_of_jobs), repeat(self.qtensor_circ,  first_set_of_jobs), 
                     repeat(num_circs_per_job + 1,  first_set_of_jobs), repeat(self.n, first_set_of_jobs)))
                 self._arggen.extend(list(zip(repeat(self.noise_sim, second_set_of_jobs), repeat(self.qtensor_circ,  second_set_of_jobs), 
@@ -218,20 +223,31 @@ def _get_args(self, i):
         print("total_jobs: ", total_jobs)
 
 
-    def simulate_qiskit_density_matrix(self, circuit, noise_model_qiskit, take_trace = True):
+    def simulate_qiskit_density_matrix(self, circuit, noise_model_qiskit, backend, take_trace = True):
         start = time.time_ns() / (10 ** 9)
-        backend = AerSimulator(method='density_matrix', noise_model=noise_model_qiskit, fusion_enable=False, fusion_verbose=True)
-        result = execute(circuit, backend, shots=1).result()
-        result = backend.run(circuit).result()
+        if isinstance(backend, NumpyBackend):
+            qiskit_backend = AerSimulator(method='density_matrix', noise_model=noise_model_qiskit, fusion_enable=False, fusion_verbose=True)
+        elif isinstance(backend, CuPyBackend):
+            try: 
+                qiskit_backend = AerSimulator(method='density_matrix', noise_model=noise_model_qiskit, fusion_enable=False, fusion_verbose=True)
+                # qiskit_backend.set_options(device='GPU')
+            except AerError as e:
+                print(e)
+        result = execute(circuit, qiskit_backend, shots=1).result()
+        result = qiskit_backend.run(circuit).result()
         #result = execute(circuit, Aer.get_backend('aer_simulator_density_matrix'), shots=1, noise_model=noise_model_qiskit).result()
+        print(type(backend))
         if take_trace:
-            self.qiskit_probs = np.diagonal(result.results[0].data.density_matrix.real)
-            end = time.time_ns() / (10 ** 9)
-            self.qiskit_sim_time = end - start
+            if isinstance(backend, NumpyBackend):
+                self.qiskit_probs = np.diagonal(result.results[0].data.density_matrix.real)
+            elif isinstance(backend, CuPyBackend):
+                ## TODO: Simulate Qiskit on GPU instead. copying from CPU to GPU will have a serious performance cost
+                #self.qiskit_probs = cp.diagonal(result.results[0].data.density_matrix.real)
+                self.qiskit_probs = cp.diagonal(cp.array(result.results[0].data.density_matrix.real))
         else:
             self.qiskit_density_matrix = result.results[0].data.density_matrix.real
-            end = time.time_ns() / (10 ** 9)
-            self.qiskit_sim_time = end - start
+        end = time.time_ns() / (10 ** 9)
+        self.qiskit_sim_time = end - start
 
     def _get_circuits(self, G, gamma, beta):
         # Create Qiskit circuit 
diff --git a/qtensor/noise_simulator/NoiseSimComparisonResult.py b/qtensor/noise_simulator/NoiseSimComparisonResult.py
index db2ac9e1..24bbf0d6 100644
--- a/qtensor/noise_simulator/NoiseSimComparisonResult.py
+++ b/qtensor/noise_simulator/NoiseSimComparisonResult.py
@@ -1,5 +1,7 @@
 # from helper_functions import *
 from helper_functions import fidelity, cosine_similarity
+from qtensor.contraction_backends import CuPyBackend, NumpyBackend
+from qtensor.tools.lazy_import import cupy as cp
 
 from datetime import datetime
 import time
@@ -7,6 +9,8 @@
 import jsbeautifier
 import jsonpickle
 import numpy as np
+from numpy.linalg import norm as Norm
+
 
 from qtensor import QiskitQAOAComposer
 import qiskit.providers.aer.noise as noise
@@ -14,7 +18,7 @@
 from NoiseModel import NoiseModel
 
 class NoiseSimComparisonResult:
-    def __init__(self, qiskit_circ, qtensor_circ, qiskit_noise_model, qtensor_noise_model, n, p, d, name = ""):
+    def __init__(self, qiskit_circ, qtensor_circ, qiskit_noise_model, qtensor_noise_model, n, p, d, backend, name = ""):
         self.name = name
         self.experiment_date =  datetime.now().isoformat()
         self.experiment_start_time = time.time_ns() / (10 ** 9) # convert to floating-point seconds
@@ -34,6 +38,7 @@ def __init__(self, qiskit_circ, qtensor_circ, qiskit_noise_model, qtensor_noise_
         self.d = d
         self.qiskit_noise_model = qiskit_noise_model
         self.qtensor_noise_model = qtensor_noise_model
+        self.backend = backend
         self.data = {}
         self._get_noise_model_string()
         self.pickled_qtensor_noise_model = jsonpickle.encode(self.qtensor_noise_model)
@@ -81,6 +86,8 @@ def print_result(self):
                 f"\n{'Noiseless Fidelity:':<20}{np.round(self.noiseless_fidelity.real, 7):<10}",
                 f"\n{'Uniform Qiskit Fidelity: ':<20}{np.round(self.uniform_qiskit_fidelity.real, 7):<10}",
                 f"\n{'Uniform QTensor Fidelity: ':<20}{np.round(self.uniform_qtensor_fidelity.real, 7):<10}",
+                f"\n{'Qiskit time taken:':<20}{self.qiskit_time_taken:<10}",
+                f"\n{'QTensor time taken:':<20}{self.qtensor_time_taken:<10}",
                 f"\n{'Total time taken:':<20}{self.experiment_time_taken:<10}")
 
     def print_noise_model(self):
@@ -101,26 +108,41 @@ def _calc_similarity(self):
         self.fidelity = fidelity(self.qiskit_density_matrix.data, self.qtensor_density_matrix)
 
     def _calc_similarity_of_probs(self):
-        qiskit_probs_root = np.sqrt(self.qiskit_probs)
-        qtensor_probs_root = np.sqrt(self.qtensor_probs)
-        noiseless_qtensor_probs_root = np.sqrt(np.abs(self.exact_qtensor_amps)**2)
-        uniform_probs_root = np.sqrt(np.ones(2**self.n)/2**self.n)
-
-        # we don't need to take the conjugate, as both population density vectors are strictly 
-        # real by the time they have made it here. 
-
-        """Measures the fidelity between the qiskit density matrix noisy state and the qtensoir stochastic noisy state"""
-        self.noisy_fidelity = np.abs((np.inner(qiskit_probs_root, qtensor_probs_root)))**2
-
-        """Measures the fidelity between a noisy qtensor state and the noiseless version of the same state"""
-        self.noiseless_fidelity = np.abs((np.inner(noiseless_qtensor_probs_root, qtensor_probs_root)))**2
-
-        """Measures the fidelity between a qiskit noisy state and a uniform distribution state"""
-        self.uniform_qiskit_fidelity = np.abs((np.inner(uniform_probs_root, qiskit_probs_root)))**2
-
-        """Measures the fidelity between a qtensor noisy state and a uniform distribution state"""
-        self.uniform_qtensor_fidelity = np.abs((np.inner(uniform_probs_root, qtensor_probs_root)))**2
-        self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
+        if isinstance(self.backend, NumpyBackend):
+            qiskit_probs_root = np.sqrt(self.qiskit_probs)
+            qtensor_probs_root = np.sqrt(self.qtensor_probs)
+            noiseless_qtensor_probs_root = np.sqrt(np.abs(self.exact_qtensor_amps)**2)
+            uniform_probs_root = np.sqrt(np.ones(2**self.n)/2**self.n)
+
+            # we don't need to take the conjugate, as both population density vectors are strictly 
+            # real by the time they have made it here. 
+
+            """Measures the fidelity between the qiskit density matrix noisy state and the qtensoir stochastic noisy state"""
+            self.noisy_fidelity = np.abs((np.inner(qiskit_probs_root, qtensor_probs_root)))**2
+
+            """Measures the fidelity between a noisy qtensor state and the noiseless version of the same state"""
+            self.noiseless_fidelity = np.abs((np.inner(noiseless_qtensor_probs_root, qtensor_probs_root)))**2
+
+            """Measures the fidelity between a qiskit noisy state and a uniform distribution state"""
+            self.uniform_qiskit_fidelity = np.abs((np.inner(uniform_probs_root, qiskit_probs_root)))**2
+
+            """Measures the fidelity between a qtensor noisy state and a uniform distribution state"""
+            self.uniform_qtensor_fidelity = np.abs((np.inner(uniform_probs_root, qtensor_probs_root)))**2
+            #self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
+            self.cos_sim = np.inner(self.qiskit_probs, self.qtensor_probs) / (Norm(self.qiskit_probs)* Norm(self.qtensor_probs)) 
+
+        elif isinstance(self.backend, CuPyBackend):
+            qiskit_probs_root = cp.sqrt(self.qiskit_probs)
+            qtensor_probs_root = cp.sqrt(self.qtensor_probs)
+            noiseless_qtensor_probs_root = cp.sqrt(cp.abs(self.exact_qtensor_amps)**2)
+            uniform_probs_root = cp.sqrt(cp.ones(2**self.n)/2**self.n)
+
+            self.noisy_fidelity = cp.abs((cp.inner(qiskit_probs_root, qtensor_probs_root)))**2
+            self.noiseless_fidelity = cp.abs((cp.inner(noiseless_qtensor_probs_root, qtensor_probs_root)))**2
+            self.uniform_qiskit_fidelity = cp.abs((cp.inner(uniform_probs_root, qiskit_probs_root)))**2
+            self.uniform_qtensor_fidelity = cp.abs((cp.inner(uniform_probs_root, qtensor_probs_root)))**2
+            # self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
+            self.cos_sim = cp.inner(self.qiskit_probs, self.qtensor_probs) / (cp.linalg.norm(self.qiskit_probs)* cp.norm(self.qtensor_probs)) 
 
     def _to_dict(self):
         #self.experiment_dict['name'] = self.name
diff --git a/qtensor/noise_simulator/NoiseSimulator.py b/qtensor/noise_simulator/NoiseSimulator.py
index 395f8395..2bc99c9e 100644
--- a/qtensor/noise_simulator/NoiseSimulator.py
+++ b/qtensor/noise_simulator/NoiseSimulator.py
@@ -1,4 +1,6 @@
 from qtensor.Simulate import QtreeSimulator, NumpyBackend
+from qtensor.contraction_backends import CuPyBackend
+from qtensor.tools.lazy_import import cupy as cp
 import qtree
 from NoiseModel import NoiseModel
 
@@ -11,6 +13,7 @@ def __init__(self, noise_model, bucket_backend=NumpyBackend(), optimizer=None, m
         if not isinstance(noise_model, NoiseModel):
             raise ValueError("Error: noise_model value must be of type NoiseModel")
         self.noise_model = noise_model
+        self.backend = bucket_backend
 
 
     def simulate_batch_ensemble(self, qc, num_circs, batch_vars=0, peo=None):
@@ -21,15 +24,22 @@ def simulate_batch_ensemble(self, qc, num_circs, batch_vars=0, peo=None):
         This is because we must take the modulus squared of each statevector in the ensemble to conserve probability. Thus the user will 
         receive a probability density vector, not a probability amplitude vector.  
         """
+    
         start = time.time_ns() / (10 ** 9)
         if num_circs < 0 or not isinstance(num_circs, int):
             raise Exception("Error: The argument num_circs must be a positive integer")
-        
-        unnormalized_ensemble_probs = [0] * 2**batch_vars
-        for _ in range(num_circs):
-            noisy_state_amps = self.simulate_batch(qc, batch_vars, peo)
-            noisy_state_probs = np.square(np.absolute(noisy_state_amps))
-            unnormalized_ensemble_probs += noisy_state_probs
+        if isinstance(self.backend, CuPyBackend):
+            unnormalized_ensemble_probs = cp.zeros(shape = 2**batch_vars, dtype = cp.complex128)
+            for _ in range(num_circs):
+                noisy_state_amps = self.simulate_batch(qc, batch_vars, peo)
+                noisy_state_probs = cp.square(cp.absolute(noisy_state_amps))
+                unnormalized_ensemble_probs += noisy_state_probs
+        elif isinstance(self.backend, NumpyBackend):
+            unnormalized_ensemble_probs = [0] * 2**batch_vars
+            for _ in range(num_circs):
+                noisy_state_amps = self.simulate_batch(qc, batch_vars, peo)
+                noisy_state_probs = np.square(np.absolute(noisy_state_amps))
+                unnormalized_ensemble_probs += noisy_state_probs
         
         normalized_ensemble_probs = unnormalized_ensemble_probs / num_circs
         self.normalized_ensemble_probs = normalized_ensemble_probs
diff --git a/qtensor/noise_simulator/comparison_mpi.py b/qtensor/noise_simulator/comparison_mpi.py
index c8432141..9c469d01 100644
--- a/qtensor/noise_simulator/comparison_mpi.py
+++ b/qtensor/noise_simulator/comparison_mpi.py
@@ -69,8 +69,8 @@
             for _ in range(num_samples):
                 print("\n\nnum_circs_list:", num_circs_list)
                 comparison = QAOAComparisonSimulator(n, p, d, noise_model_qiskit, noise_model_qtensor, num_circs_list)
-                comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node, backend=CuPyBackend())
-                #comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node)
+                #comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node, backend=CuPyBackend())
+                comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node)
                 results.extend(comparison.results)
 if results:
     save_dict_to_file(results, name = outfile_name) 

From 9ca45a5ba0596d274ad52a9df57949e4f4e5bb7f Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Sun, 12 Mar 2023 23:20:02 -0500
Subject: [PATCH 21/24] added an example jupyter notebook on how to use the
 stochastic noise simulator

---
 examples/Stochastic_noise_usage.ipynb | 307 ++++++++++++++++++++++++++
 1 file changed, 307 insertions(+)
 create mode 100644 examples/Stochastic_noise_usage.ipynb

diff --git a/examples/Stochastic_noise_usage.ipynb b/examples/Stochastic_noise_usage.ipynb
new file mode 100644
index 00000000..632402fc
--- /dev/null
+++ b/examples/Stochastic_noise_usage.ipynb
@@ -0,0 +1,307 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Create noise channel(s)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qtensor.noise_simulator.NoiseChannels import DepolarizingChannel\n",
+    "from qtensor.noise_simulator.NoiseModel import NoiseModel\n",
+    "\n",
+    "# Depolarizing channel for a single qubit. The probability of the error occuring is related to 0.003\n",
+    "depol_chan_1Q = DepolarizingChannel(param = 0.003, num_qubits = 1)\n",
+    "# Depolarizing channel for two qubits. The probability of the error occuring is related to 0.03\n",
+    "depol_chan_2Q = DepolarizingChannel(param = 0.03, num_qubits = 2)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Create noise model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "noise_model = NoiseModel()\n",
+    "# If a circuit uses this noise model, all H, XPhase, and ZPhase gates will experience depolarizing noise \n",
+    "# based on the single qubit depolarizing channel created above\n",
+    "noise_model.add_channel_to_all_qubits(channel = depol_chan_1Q, gates = ['H', 'XPhase', 'ZPhase'])\n",
+    "# If a circuit uses this noise model, all cX gates will experience depolarizing noise \n",
+    "# based on the two-qubit depolarizing channel created above\n",
+    "noise_model.add_channel_to_all_qubits(channel = depol_chan_2Q, gates = ['cX'])"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Create a circuit. Here we create a QAOA circuit. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qtree.operators import from_qiskit_circuit\n",
+    "from qtensor import QiskitQAOAComposer, QtreeSimulator\n",
+    "import qtensor\n",
+    "import numpy as np\n",
+    "import networkx as nx\n",
+    "\n",
+    "# degree of graph\n",
+    "d = 3\n",
+    "num_qubits = 4\n",
+    "G = nx.random_regular_graph(d, num_qubits)\n",
+    "\n",
+    "# optimal gamma and beta parameters previously found for these values of d and n\n",
+    "gammabeta = np.array(qtensor.tools.BETHE_QAOA_VALUES[str(d)]['angles'])\n",
+    "gamma = -gammabeta[:d]\n",
+    "beta = gammabeta[d:]\n",
+    "\n",
+    "# composes a qaoa circuit of type qiskit\n",
+    "qiskit_comp = QiskitQAOAComposer(G, gamma=gamma, beta=beta)\n",
+    "qiskit_comp.ansatz_state()\n",
+    "\n",
+    "# converts the circuit to qtensor type, which is just a list \n",
+    "n, circ = from_qiskit_circuit(qiskit_comp.circuit)\n",
+    "# from_qiskiit_circuit returns each gate as a list so we remove the lists from the gates. Might want to fix that.\n",
+    "circ = sum(circ, [])"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Instatiate simulator with noise model and then simulate the circuit with noise. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([0.14060763, 0.0349968 , 0.04040751, 0.06550023, 0.05163103,\n",
+       "       0.06425988, 0.03407733, 0.06851953, 0.06851953, 0.03407733,\n",
+       "       0.06425988, 0.05163103, 0.06550023, 0.04040751, 0.0349968 ,\n",
+       "       0.14060763])"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from qtensor.noise_simulator.NoiseSimulator import NoiseSimulator\n",
+    "\n",
+    "noise_sim = NoiseSimulator(noise_model)\n",
+    "\n",
+    "num_circs = 10\n",
+    "# simulate an ensemble of 10 circuits with stochastic noise\n",
+    "noise_sim.simulate_batch_ensemble(circ, num_circs, num_qubits)\n",
+    "# get the approximate noisy probabilities \n",
+    "qtensor_probs = noise_sim.normalized_ensemble_probs\n",
+    "qtensor_probs"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This same circuit can also be simulated without noise."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([0.26962495, 0.03924902, 0.03924902, 0.02445963, 0.03924902,\n",
+       "       0.02445963, 0.02445963, 0.03924902, 0.03924902, 0.02445963,\n",
+       "       0.02445963, 0.03924902, 0.02445963, 0.03924902, 0.03924902,\n",
+       "       0.26962495])"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ideal_sim = QtreeSimulator()\n",
+    "amplitudes = ideal_sim.simulate_batch(circ, num_qubits)\n",
+    "ideal_probs = np.abs(amplitudes)**2\n",
+    "ideal_probs"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can check how well our simulation approximated the exact noisy state by running the same circuit with a density matrix simulator and then calculating the fidelity between the exact and approximate noisy states."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([0.13403107, 0.05506911, 0.05646327, 0.0514643 , 0.05139734,\n",
+       "       0.0499991 , 0.04779918, 0.05377664, 0.05377664, 0.04779918,\n",
+       "       0.0499991 , 0.05139734, 0.0514643 , 0.05646327, 0.05506911,\n",
+       "       0.13403107])"
+      ]
+     },
+     "execution_count": 26,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import qiskit.providers.aer.noise as noise\n",
+    "from qiskit.providers.aer import AerSimulator\n",
+    "\n",
+    "# Create the same noise model with the same channels and noisy gates in Qiskits framework\n",
+    "depol_chan_qiskit_1Q = noise.depolarizing_error(param = 0.003, num_qubits = 1)\n",
+    "depol_chan_qiskit_2Q = noise.depolarizing_error(param = 0.03, num_qubits = 2)\n",
+    "\n",
+    "noise_model_qiskit = noise.NoiseModel()\n",
+    "noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_1Q, ['rz', 'rx', 'h'])\n",
+    "noise_model_qiskit.add_all_qubit_quantum_error(depol_chan_qiskit_2Q, ['cx'])\n",
+    "\n",
+    "# Take the Qiskit circuit we created before, and add some final touches to it. These do not effect any measurement outcomes \n",
+    "qiskit_circ = qiskit_comp.circuit\n",
+    "\n",
+    "# Qiskit uses little endian notation by default, and Qtensor uses big endian, so we change the Qiskit circuit to match qtensor\n",
+    "qiskit_circ = qiskit_circ.reverse_bits()\n",
+    "qiskit_circ.save_density_matrix()\n",
+    "qiskit_circ.measure_all(add_bits = False)\n",
+    "\n",
+    "# Simulate the circuit\n",
+    "backend = AerSimulator(method='density_matrix', noise_model=noise_model_qiskit, fusion_enable=False, fusion_verbose=True)\n",
+    "result = backend.run(qiskit_circ).result()\n",
+    "\n",
+    "# The stochastic noisy simulation gives us a probability density vector, so we need to save the diagonal elements of the density matrix for our comparison\n",
+    "qiskit_probs = np.diagonal(np.asarray(result.results[0].data.density_matrix).real)\n",
+    "qiskit_probs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.9851105052526412"
+      ]
+     },
+     "execution_count": 24,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# The fidelity is defined for probability amplitudes, so we need to take the square root of our probability density vectors\n",
+    "A = np.sqrt(qtensor_probs)\n",
+    "B = np.sqrt(qiskit_probs)\n",
+    "fidelity = np.dot(A, B)**2\n",
+    "fidelity"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can increase our fidelity by increasing the number of circuits in the ensemble."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.9954629065353349"
+      ]
+     },
+     "execution_count": 25,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "num_circs = 100\n",
+    "noise_sim.simulate_batch_ensemble(circ, num_circs, num_qubits)\n",
+    "qtensor_probs = noise_sim.normalized_ensemble_probs\n",
+    "A = np.sqrt(qtensor_probs)\n",
+    "B = np.sqrt(qiskit_probs)\n",
+    "fidelity = np.dot(A, B)**2\n",
+    "fidelity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.2"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

From 91183d21183d8d0e58b3913ac510bd463421af0e Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Tue, 30 May 2023 10:26:09 -0500
Subject: [PATCH 22/24] state reconstruction and obtaining a classical shadow
 implemented

---
 qtensor/classical_shadows/shadow.py | 148 ++++++++++++++++++++++++++++
 1 file changed, 148 insertions(+)
 create mode 100644 qtensor/classical_shadows/shadow.py

diff --git a/qtensor/classical_shadows/shadow.py b/qtensor/classical_shadows/shadow.py
new file mode 100644
index 00000000..fc7675f4
--- /dev/null
+++ b/qtensor/classical_shadows/shadow.py
@@ -0,0 +1,148 @@
+from qtensor.Simulate import QtreeSimulator, NumpyBackend
+from qtree import operators as op
+from qtree.operators import Gate 
+import numpy as np
+from qtensor.noise_simulator.helper_functions import decimal_to_binary
+import copy
+
+class Classical_Shadow(QtreeSimulator): 
+    def __init__(self, bucket_backend=NumpyBackend(), optimizer=None, max_tw=None):
+        super().__init__(bucket_backend, optimizer, max_tw)
+
+    def get_snapshot(self, circuit, num_snapshots, num_qubits):
+        """
+        Gets the classical shadows from a given circuit
+
+        Args: 
+            circuit (list): the circuit to be simulated
+            num_snapshots (int): the number of shadows in the collection
+            num_qubits (int): the number of qubits in the circuit
+
+        Returns: 
+
+        """
+        obseravbles = [op.X, op.Y, op.Z]
+        basis_measurements = [[op.H], [op.Sdag, op.H], [op.M]]
+        observable_ids = np.random.randint(0, 3, size=(num_snapshots, num_qubits))
+        # observable_ids = np.random.randint(1, 2, size=(num_snapshots, num_qubits))
+        snapshots = np.zeros((num_snapshots, num_qubits))
+
+        for snapshot in range(num_snapshots):
+            # Generates a list of paulis applied to the i-th qubit, then appends the list to the circuit
+            snapshot_circ = copy.deepcopy(circuit)
+            # print((basis_measurements[int(observable_ids[snapshot, i])](i) forH  i in range(num_qubits)))
+            for i in range(num_qubits):
+                for clifford in basis_measurements[int(observable_ids[snapshot, i])]:
+                    snapshot_circ.append(clifford(i))
+            # snapshot_circ += (basis_measurements[int(observable_ids[snapshot, i])](i) for i in range(num_qubits))
+            print(snapshot_circ)
+            snapshots[snapshot, :] = self._get_measurement(snapshot_circ, num_qubits)
+            # print(snapshot_circ)
+        return (snapshots, observable_ids)
+
+    def _get_measurement(self, circuit, num_qubits):
+            """
+            Simulates state after applying a pauli to each qubit. Measure the state to obtain a bitstring, and then map 0 -> 1 and 1 -> -1 
+
+            Args: 
+                circuit (list): the circuit to be simulated
+                num_qubits (int): the number of qubits in the circuit
+            
+            Returns:
+                eigenvalues: the bitstring obtained probabilisitically from the statevector, with its values mapped to 1 and -1
+            """
+
+            sim = QtreeSimulator()
+            statevector = sim.simulate_batch(circuit, batch_vars = num_qubits)
+            # print("statevector:", statevector)
+            # print(statevector)
+            probs = np.square(np.absolute(statevector))
+            probs = [round(elem, 6) for elem in probs]
+            #print("statevector:", probs)
+            measurement = int(np.random.choice(np.arange(len(probs)), size = 1, p = probs))
+            #print(measurement, type(measurement))
+            length = int(np.ceil(np.log(2**num_qubits + 1)/np.log(2)) - 1)
+            measurement = str(decimal_to_binary(measurement).zfill(length))
+            #print(measurement, type(measurement))
+            measurement = np.asarray([int(measurement[i]) for i in range(len(measurement))])
+            #print(measurement, type(measurement))
+            eigenvalues = np.piecewise(measurement, [measurement == 0, measurement == 1], [1, -1])
+            # eigenvalues = []
+            # for i in range(len(measurement)):
+            #     eigenvalue = np.piecewise(measurement[i], [measurement[i] == '0', measurement[i] == '1'], [1, -1])
+            #     print(f"eiegenvalue: {eigenvalue}")
+            #     eigenvalues.append(eigenvalue)
+            #     print(f"measurement[i]: {measurement[i]}, eigevnalues[i]: {eigenvalues[i]}")
+            #eigenvalues = [np.piecewise(measurement[i], [measurement[i] == '0', measurement[i] == '1'], [1, -1]) for i in range(len(measurement))] 
+            #print(eigenvalues)
+            return eigenvalues
+    
+    def _snapshot_state(self, measurement_outcomes, obseravbles):
+        num_qubits = len(measurement_outcomes)
+
+        zero_state = np.array([[1, 0], [0, 0]])
+        one_state = np.array([[0, 0], [0, 1]])
+
+        # H = op.H(Gate).gen_tensor()
+        H = 1/np.sqrt(2)*np.array([[1,1],[1,-1]])
+        Sdag = np.array([[1,0],[0,-1j]],dtype=complex)
+        I = np.identity(2)
+        # print(H)
+        # ZPhase = op.ZPhase(Gate).gen_tensor({'alpha': np.pi/2})
+        # ZPhase = np.array([[1, 0], [0, -1j]], dtype=complex)
+        # I = op.M(Gate).gen_tensor()
+        # I = np.array([[1, 0], [0, 1]], dtype=complex)
+        #unitaries = [H, H_Sdag(Gate).gen_tensor(), I]
+        # Sdag = op.Sdag(Gate).gen_tensor()
+        unitaries = [H, H @ Sdag , I]
+        # unitaries = [op.X.gen_tensor(Gate), op.Y.gen_tensor(Gate), op.Z.gen_tensor(Gate)]
+
+
+        rho_snapshot = [1]
+        for i in range(num_qubits):
+            state = zero_state if measurement_outcomes[i] == 1 else one_state
+            U = unitaries[int(obseravbles[i])]
+
+            local_rho = 3 * U.conj().T @ state @ U - I
+            rho_snapshot = np.kron(rho_snapshot, local_rho)
+
+        return rho_snapshot
+
+
+
+    def get_approximate_state(self, shadow): 
+        """
+        Constructs and approximate state from the n snapshots in the shadow
+
+        Args: 
+            shadow (tuple): a classical shadow 
+        """
+        num_snapshots, num_qubits = shadow[0].shape
+        measurement_outcomes, observables = shadow
+
+        shadow_rho = np.zeros((2**num_qubits, 2**num_qubits), dtype=complex)
+
+        for i in range(num_snapshots):
+            shadow_rho += self._snapshot_state(measurement_outcomes[i], observables[i])
+
+        return shadow_rho / num_snapshots
+    
+
+class H_Sdag(Gate):
+    name = 'H_Sdag'
+    _changes_qubits = tuple()
+    # _changes_qubits = (0,)
+    def gen_tensor(self):
+        return op.H(Gate).gen_tensor() @ op.Sdag(Gate).gen_tensor() 
+        # return 1/np.sqrt(2) * np.array([[1.+0.j, 0.-1j],
+        #                                 [1.+0.j, 0.+1j]])
+    
+class Sdag_H(Gate):
+    name = 'Sdag_H'
+    _changes_qubits = tuple()
+    # _changes_qubits = (0,)
+    def gen_tensor(self):
+        return op.Sdag(Gate).gen_tensor() @ op.H(Gate).gen_tensor() 
+        # return 1/np.sqrt(2) * np.array([[1.+0.j, 1.+0j],
+        #                                 [0.-1.j, 0.+1j]])
+

From 4016e35a6135ac4f37f5373fcf83f7453a62a4a8 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Tue, 11 Jul 2023 20:52:32 -0500
Subject: [PATCH 23/24] fixed a few minor bugs in NoiseSimulator. Cleaned up
 some comments.

---
 .../noise_simulator/ComparisonSimulator.py    |  35 +----
 .../NoiseSimComparisonResult.py               | 122 ++++++++++++------
 qtensor/noise_simulator/NoiseSimulator.py     |  50 +++----
 qtensor/noise_simulator/comparison_mpi.py     |  15 ++-
 4 files changed, 117 insertions(+), 105 deletions(-)

diff --git a/qtensor/noise_simulator/ComparisonSimulator.py b/qtensor/noise_simulator/ComparisonSimulator.py
index 840eb5fc..ce03f7bc 100644
--- a/qtensor/noise_simulator/ComparisonSimulator.py
+++ b/qtensor/noise_simulator/ComparisonSimulator.py
@@ -2,18 +2,17 @@
 from NoiseModel import NoiseModel 
 from NoiseSimComparisonResult import NoiseSimComparisonResult
 from helper_functions import get_qaoa_params
-# from qtensor.tests.test_composers import QtreeSimulator, NumpyBackend
 from qtensor.Simulate import QtreeSimulator, NumpyBackend
 from qtensor.contraction_backends import CuPyBackend
 from qtensor import QiskitQAOAComposer
 from qtree.operators import from_qiskit_circuit
 from qtensor import tools
 
-
 from qiskit import execute
 from qiskit.providers.aer import AerSimulator, noise, AerError
 
 import numpy as np
+from sigpy import get_array_module
 from qtensor.tools.lazy_import import cupy as cp
 from itertools import repeat
 import time
@@ -32,15 +31,6 @@ def __init__(self, n: int, p: int, d: int, noise_model_qiskit: noise.NoiseModel,
         self.num_circs_list = num_circs_list
         self.num_circs_simulated = []
         self.results = []
-
-        # We can save on simulation time by not recomputing previous work. This is done by 
-        # using the results of a previous, smaller ensemble simulation: 
-        # i.e. if we ran 1,000 circuits in the previous simulation, and now we want to run 10,000 
-        # circuits using the exact same circuit and parameters, we actually only need to run
-        # 9,000 circuits, add the previous results, and renormalize. 
-        # The downside of this approach is that it requires the storage of an extra state 
-        # in memory. Therefore it should only be done for simulations that do not require 
-        # all of the available memory
         self.recompute_previous_ensemble: bool
 
         self._check_params()
@@ -62,19 +52,16 @@ def qtensor_qiskit_noisy_qaoa(self, backend = NumpyBackend(), recompute_previous
             if i == 0 or recompute_previous_ensemble == False:
                 self.num_circs_simulated.append(num_circs)
                 noise_sim.simulate_batch_ensemble(sum(self.qtensor_circ, []), num_circs, self.num_qubits)
-                #print("num_circs: {}".format(num_circs))
                 self.qtensor_probs = noise_sim.normalized_ensemble_probs
             else: 
                 actual_num_circs = num_circs - self.num_circs_list[i - 1]
                 self.num_circs_simulated.append(actual_num_circs)
-                #print("num_circs: {}, actual_num_circs: {}".format(num_circs, actual_num_circs_to_sim))
                 noise_sim.simulate_batch_ensemble(sum(self.qtensor_circ, []), actual_num_circs, self.num_qubits)
                 prev_qtensor_probs = self.prev_probs
                 curr_qtensor_probs = noise_sim.normalized_ensemble_probs
                 self.qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
             
             if recompute_previous_ensemble == True:
-                #self._check_correct_num_circs_simulated(i)
                 self.prev_probs = self.qtensor_probs
 
             qtensor_sim_time = noise_sim.time_taken
@@ -90,14 +77,8 @@ def qtensor_qiskit_noisy_qaoa(self, backend = NumpyBackend(), recompute_previous
 
     def _mpi_parallel_unit(self, args):
         noise_sim, qtensor_circ, num_circs, num_qubits = args
-        # #print("this workers num_circs: ", num_circs)
-        # if noise_sim.backend == CuPyBackend():
-        #     noise_sim.simulate_batch_ensemble(cupy.array(sum(qtensor_circ, [])), num_circs, num_qubits)
-        # else: 
-        #     noise_sim.simulate_batch_ensemble(sum(qtensor_circ, []), num_circs, num_qubits)
         noise_sim.simulate_batch_ensemble(sum(qtensor_circ, []), num_circs, num_qubits)
         fraction_of_qtensor_probs = noise_sim.normalized_ensemble_probs
-        # fraction_of_qtensor_probs = [0] * 2**self.n
         return fraction_of_qtensor_probs
 
     def qtensor_qiskit_noisy_qaoa_mpi(self,  num_nodes: int, num_jobs_per_node: int, backend = NumpyBackend(), recompute_previous_ensemble: bool = False, print_stats: bool = True, pbar: bool = True):
@@ -120,7 +101,6 @@ def qtensor_qiskit_noisy_qaoa_mpi(self,  num_nodes: int, num_jobs_per_node: int,
                 if i == 0 or recompute_previous_ensemble == False: 
                     self.qtensor_probs = sum(qtensor_probs_list) / self._total_jobs
                 else:
-                    #prev_qtensor_probs = self.results[i - 1].qtensor_probs
                     prev_qtensor_probs = self.prev_probs
                     curr_qtensor_probs = sum(qtensor_probs_list) / self._total_jobs
                     self.qtensor_probs = (curr_qtensor_probs + prev_qtensor_probs) / 2
@@ -178,10 +158,6 @@ def _get_args(self, i):
             num_circs = self.num_circs_list[i]
         else:
             num_circs = self.num_circs_list[i] - self.num_circs_list[i - 1]
-        ## We make sure that each job has a minimum number of circuits. We do this because 
-        ## if there are too few circuits for each unit of work, the overhead from 
-        ## parallelization removes any advantage gained 
-        ## TODO: determine a better minimum number of circuits. Currently 10 is chosen arbitrarily
         min_circs_per_job = min(10, num_circs)
         if self.num_nodes * self.num_jobs_per_node > num_circs / min_circs_per_job:
             num_circs_per_job = min_circs_per_job
@@ -197,7 +173,6 @@ def _get_args(self, i):
                 repeat(num_circs_per_job, total_jobs), repeat(self.n, total_jobs)))
             self._total_jobs = total_jobs
             self.num_circs_simulated.append(num_circs)
-            #print("num_circs: {}, actual num_circs simulated on this iteration: {}".format(self.num_circs_list[i], num_circs))
         else: 
             if num_circs_per_job == min_circs_per_job:
                 self._arggen = list(zip(repeat(self.noise_sim, total_jobs - 1), repeat(self.qtensor_circ, total_jobs - 1), 
@@ -208,19 +183,14 @@ def _get_args(self, i):
             else: 
                 first_set_of_jobs = num_circs - (total_jobs * num_circs_per_job)
                 second_set_of_jobs = total_jobs - first_set_of_jobs
-                print("first_set_of_jobs: {}, second_set_of_jobs: {}".format(first_set_of_jobs, second_set_of_jobs))
                 self._arggen = list(zip(repeat(self.noise_sim, first_set_of_jobs), repeat(self.qtensor_circ,  first_set_of_jobs), 
                     repeat(num_circs_per_job + 1,  first_set_of_jobs), repeat(self.n, first_set_of_jobs)))
                 self._arggen.extend(list(zip(repeat(self.noise_sim, second_set_of_jobs), repeat(self.qtensor_circ,  second_set_of_jobs), 
                     repeat(num_circs_per_job,  second_set_of_jobs), repeat(self.n, second_set_of_jobs))))  
                 actual_num_circs = first_set_of_jobs * (num_circs_per_job + 1) + second_set_of_jobs * num_circs_per_job
             self._total_jobs = total_jobs
-            print("actual_num_circs: {}, num_circs: {}, num_nodes: {}, num_jobs_per_node: {}, total_jobs: {}".format(
-                actual_num_circs, num_circs, self.num_nodes, self.num_jobs_per_node, self._total_jobs))
             assert num_circs == actual_num_circs
             self.num_circs_simulated.append(actual_num_circs)
-            #print("num_circs: {}, actual num_circs simulated on this iteration: {}, total jobs: {}".format(self.num_circs_list[i], actual_num_circs, total_jobs))
-        print("total_jobs: ", total_jobs)
 
 
     def simulate_qiskit_density_matrix(self, circuit, noise_model_qiskit, backend, take_trace = True):
@@ -235,14 +205,11 @@ def simulate_qiskit_density_matrix(self, circuit, noise_model_qiskit, backend, t
                 print(e)
         result = execute(circuit, qiskit_backend, shots=1).result()
         result = qiskit_backend.run(circuit).result()
-        #result = execute(circuit, Aer.get_backend('aer_simulator_density_matrix'), shots=1, noise_model=noise_model_qiskit).result()
-        print(type(backend))
         if take_trace:
             if isinstance(backend, NumpyBackend):
                 self.qiskit_probs = np.diagonal(result.results[0].data.density_matrix.real)
             elif isinstance(backend, CuPyBackend):
                 ## TODO: Simulate Qiskit on GPU instead. copying from CPU to GPU will have a serious performance cost
-                #self.qiskit_probs = cp.diagonal(result.results[0].data.density_matrix.real)
                 self.qiskit_probs = cp.diagonal(cp.array(result.results[0].data.density_matrix.real))
         else:
             self.qiskit_density_matrix = result.results[0].data.density_matrix.real
diff --git a/qtensor/noise_simulator/NoiseSimComparisonResult.py b/qtensor/noise_simulator/NoiseSimComparisonResult.py
index 24bbf0d6..2afc3435 100644
--- a/qtensor/noise_simulator/NoiseSimComparisonResult.py
+++ b/qtensor/noise_simulator/NoiseSimComparisonResult.py
@@ -10,6 +10,7 @@
 import jsonpickle
 import numpy as np
 from numpy.linalg import norm as Norm
+from sigpy import get_array_module
 
 
 from qtensor import QiskitQAOAComposer
@@ -101,48 +102,95 @@ def save_experiment_to_file(self, outfile):
             f.write(', ')
 
     def _calc_similarity(self):
-        qiskit_probs_root = np.sqrt(np.diagonal(self.qiskit_density_matrix))
-        qtensor_probs_root = np.sqrt(np.diagonal(self.qtensor_density_matrix))
+        xp = get_array_module(self)
+        qiskit_probs_root = xp.sqrt(xp.diagonal(self.qiskit_density_matrix))
+        qtensor_probs_root = xp.sqrt(xp.diagonal(self.qtensor_density_matrix))
+        noiseless_qtensor_probs_root = xp.sqrt(xp.abs(self.exact_qtensor_amps)**2)
+        uniform_probs_root = xp.sqrt(xp.ones(2**self.n)/2**self.n)
         
         self.cos_sim = cosine_similarity(qiskit_probs_root, qtensor_probs_root)
         self.fidelity = fidelity(self.qiskit_density_matrix.data, self.qtensor_density_matrix)
 
+        self.cos_sim = xp.inner(self.qiskit_probs, self.qtensor_probs) / (xp.linalg.norm(self.qiskit_probs)* xp.linalg.norm(self.qtensor_probs)) 
+
     def _calc_similarity_of_probs(self):
-        if isinstance(self.backend, NumpyBackend):
-            qiskit_probs_root = np.sqrt(self.qiskit_probs)
-            qtensor_probs_root = np.sqrt(self.qtensor_probs)
-            noiseless_qtensor_probs_root = np.sqrt(np.abs(self.exact_qtensor_amps)**2)
-            uniform_probs_root = np.sqrt(np.ones(2**self.n)/2**self.n)
-
-            # we don't need to take the conjugate, as both population density vectors are strictly 
-            # real by the time they have made it here. 
-
-            """Measures the fidelity between the qiskit density matrix noisy state and the qtensoir stochastic noisy state"""
-            self.noisy_fidelity = np.abs((np.inner(qiskit_probs_root, qtensor_probs_root)))**2
-
-            """Measures the fidelity between a noisy qtensor state and the noiseless version of the same state"""
-            self.noiseless_fidelity = np.abs((np.inner(noiseless_qtensor_probs_root, qtensor_probs_root)))**2
-
-            """Measures the fidelity between a qiskit noisy state and a uniform distribution state"""
-            self.uniform_qiskit_fidelity = np.abs((np.inner(uniform_probs_root, qiskit_probs_root)))**2
-
-            """Measures the fidelity between a qtensor noisy state and a uniform distribution state"""
-            self.uniform_qtensor_fidelity = np.abs((np.inner(uniform_probs_root, qtensor_probs_root)))**2
-            #self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
-            self.cos_sim = np.inner(self.qiskit_probs, self.qtensor_probs) / (Norm(self.qiskit_probs)* Norm(self.qtensor_probs)) 
-
-        elif isinstance(self.backend, CuPyBackend):
-            qiskit_probs_root = cp.sqrt(self.qiskit_probs)
-            qtensor_probs_root = cp.sqrt(self.qtensor_probs)
-            noiseless_qtensor_probs_root = cp.sqrt(cp.abs(self.exact_qtensor_amps)**2)
-            uniform_probs_root = cp.sqrt(cp.ones(2**self.n)/2**self.n)
-
-            self.noisy_fidelity = cp.abs((cp.inner(qiskit_probs_root, qtensor_probs_root)))**2
-            self.noiseless_fidelity = cp.abs((cp.inner(noiseless_qtensor_probs_root, qtensor_probs_root)))**2
-            self.uniform_qiskit_fidelity = cp.abs((cp.inner(uniform_probs_root, qiskit_probs_root)))**2
-            self.uniform_qtensor_fidelity = cp.abs((cp.inner(uniform_probs_root, qtensor_probs_root)))**2
-            # self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
-            self.cos_sim = cp.inner(self.qiskit_probs, self.qtensor_probs) / (cp.linalg.norm(self.qiskit_probs)* cp.norm(self.qtensor_probs)) 
+        xp = get_array_module(self)
+        qiskit_probs_root = xp.sqrt(self.qiskit_probs)
+        qtensor_probs_root = xp.sqrt(self.qtensor_probs)
+        noiseless_qtensor_probs_root = xp.sqrt(xp.abs(self.exact_qtensor_amps)**2)
+        uniform_probs_root = xp.sqrt(xp.ones(2**self.n)/2**self.n)
+
+        # we don't need to take the conjugate, as both population density vectors are strictly 
+        # real by the time they have made it here. 
+
+        self.cos_sim = xp.inner(self.qiskit_probs, self.qtensor_probs) / (xp.linalg.norm(self.qiskit_probs)* xp.linalg.norm(self.qtensor_probs)) 
+        print(f"{'Cosine Similarity (old):':<20}{xp.round(self.cos_sim.real, 7):<10}")
+        self.cos_sim = self._cos_sim(self.qiskit_probs, self.qtensor_probs)
+
+        """Measures the fidelity between the qiskit density matrix noisy state and the qtensor stochastic noisy state"""
+        self.noisy_fidelity = xp.abs((xp.inner(qiskit_probs_root, qtensor_probs_root)))**2
+        print(f"\n{'Noisy Fidelity (old):':<20}{xp.round(self.noisy_fidelity.real, 7):<10}")
+        self.noisy_fidelity = self._fidelity(qiskit_probs_root, qtensor_probs_root)
+
+        """Measures the fidelity between a noisy qtensor state and the noiseless version of the same state"""
+        self.noiseless_fidelity = xp.abs((xp.inner(noiseless_qtensor_probs_root, qtensor_probs_root)))**2
+        print(f"\n{'Noiseless Fidelity (old):':<20}{xp.round(self.noiseless_fidelity.real, 7):<10}")
+        self.noiseless_fidelity = self._fidelity(noiseless_qtensor_probs_root, qtensor_probs_root)
+
+        """Measures the fidelity between a qiskit noisy state and a uniform distribution state"""
+        self.uniform_qiskit_fidelity = xp.abs((xp.inner(uniform_probs_root, qiskit_probs_root)))**2
+        print(f"\n{'Uniform Qiskit Fidelity (old): ':<20}{xp.round(self.uniform_qiskit_fidelity.real, 7):<10}")
+        self.uniform_qiskit_fidelity = self._fidelity(uniform_probs_root, qiskit_probs_root)
+
+        """Measures the fidelity between a qtensor noisy state and a uniform distribution state"""
+        self.uniform_qtensor_fidelity = xp.abs((xp.inner(uniform_probs_root, qtensor_probs_root)))**2
+        print(f"\n{'Uniform QTensor Fidelity (old): ':<20}{xp.round(self.uniform_qtensor_fidelity.real, 7):<10}")
+        self.uniform_qtensor_fidelity = self._fidelity(uniform_probs_root, qtensor_probs_root)
+
+        #self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
+        # if isinstance(self.backend, NumpyBackend):
+        #     qiskit_probs_root = np.sqrt(self.qiskit_probs)
+        #     qtensor_probs_root = np.sqrt(self.qtensor_probs)
+        #     noiseless_qtensor_probs_root = np.sqrt(np.abs(self.exact_qtensor_amps)**2)
+        #     uniform_probs_root = np.sqrt(np.ones(2**self.n)/2**self.n)
+
+        #     # we don't need to take the conjugate, as both population density vectors are strictly 
+        #     # real by the time they have made it here. 
+
+        #     """Measures the fidelity between the qiskit density matrix noisy state and the qtensoir stochastic noisy state"""
+        #     self.noisy_fidelity = np.abs((np.inner(qiskit_probs_root, qtensor_probs_root)))**2
+
+        #     """Measures the fidelity between a noisy qtensor state and the noiseless version of the same state"""
+        #     self.noiseless_fidelity = np.abs((np.inner(noiseless_qtensor_probs_root, qtensor_probs_root)))**2
+
+        #     """Measures the fidelity between a qiskit noisy state and a uniform distribution state"""
+        #     self.uniform_qiskit_fidelity = np.abs((np.inner(uniform_probs_root, qiskit_probs_root)))**2
+
+        #     """Measures the fidelity between a qtensor noisy state and a uniform distribution state"""
+        #     self.uniform_qtensor_fidelity = np.abs((np.inner(uniform_probs_root, qtensor_probs_root)))**2
+        #     #self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
+        #     self.cos_sim = np.inner(self.qiskit_probs, self.qtensor_probs) / (Norm(self.qiskit_probs)* Norm(self.qtensor_probs)) 
+
+        # elif isinstance(self.backend, CuPyBackend):
+        #     qiskit_probs_root = cp.sqrt(self.qiskit_probs)
+        #     qtensor_probs_root = cp.sqrt(self.qtensor_probs)
+        #     noiseless_qtensor_probs_root = cp.sqrt(cp.abs(self.exact_qtensor_amps)**2)
+        #     uniform_probs_root = cp.sqrt(cp.ones(2**self.n)/2**self.n)
+
+        #     self.noisy_fidelity = cp.abs((cp.inner(qiskit_probs_root, qtensor_probs_root)))**2
+        #     self.noiseless_fidelity = cp.abs((cp.inner(noiseless_qtensor_probs_root, qtensor_probs_root)))**2
+        #     self.uniform_qiskit_fidelity = cp.abs((cp.inner(uniform_probs_root, qiskit_probs_root)))**2
+        #     self.uniform_qtensor_fidelity = cp.abs((cp.inner(uniform_probs_root, qtensor_probs_root)))**2
+        #     # self.cos_sim = cosine_similarity(self.qiskit_probs, self.qtensor_probs)
+        #     self.cos_sim = cp.inner(self.qiskit_probs, self.qtensor_probs) / (cp.linalg.norm(self.qiskit_probs)* cp.norm(self.qtensor_probs)) 
+
+    def _fidelity(self, A, B): 
+        xp = get_array_module(self)
+        return  xp.abs((xp.inner(A, B)))**2
+    
+    def _cos_sim(self, A, B):
+        xp = get_array_module(self)
+        return xp.inner(A, B) / (xp.linalg.norm(A) * xp.linalg.norm(B)) 
 
     def _to_dict(self):
         #self.experiment_dict['name'] = self.name
diff --git a/qtensor/noise_simulator/NoiseSimulator.py b/qtensor/noise_simulator/NoiseSimulator.py
index 2bc99c9e..f5e20b50 100644
--- a/qtensor/noise_simulator/NoiseSimulator.py
+++ b/qtensor/noise_simulator/NoiseSimulator.py
@@ -1,10 +1,11 @@
 from qtensor.Simulate import QtreeSimulator, NumpyBackend
 from qtensor.contraction_backends import CuPyBackend
-from qtensor.tools.lazy_import import cupy as cp
+# from qtensor.tools.lazy_import import cupy as cp
 import qtree
 from NoiseModel import NoiseModel
 
 import numpy as np
+from sigpy import get_array_module
 import time
 
 class NoiseSimulator(QtreeSimulator):
@@ -17,58 +18,43 @@ def __init__(self, noise_model, bucket_backend=NumpyBackend(), optimizer=None, m
 
 
     def simulate_batch_ensemble(self, qc, num_circs, batch_vars=0, peo=None):
-        """Stoachastic noise simulator using statevector approach
-        
-        This method uses significantly less memory than the density matrix approach (2^n vs 4^n). However, one should note that 
-        while this uses the statevector approach to simulate, the user never gets the amplitudes, only probabilities. 
-        This is because we must take the modulus squared of each statevector in the ensemble to conserve probability. Thus the user will 
-        receive a probability density vector, not a probability amplitude vector.  
-        """
+        """Stoachastic noise simulator using statevector approach"""
     
         start = time.time_ns() / (10 ** 9)
         if num_circs < 0 or not isinstance(num_circs, int):
             raise Exception("Error: The argument num_circs must be a positive integer")
-        if isinstance(self.backend, CuPyBackend):
-            unnormalized_ensemble_probs = cp.zeros(shape = 2**batch_vars, dtype = cp.complex128)
-            for _ in range(num_circs):
-                noisy_state_amps = self.simulate_batch(qc, batch_vars, peo)
-                noisy_state_probs = cp.square(cp.absolute(noisy_state_amps))
-                unnormalized_ensemble_probs += noisy_state_probs
-        elif isinstance(self.backend, NumpyBackend):
-            unnormalized_ensemble_probs = [0] * 2**batch_vars
-            for _ in range(num_circs):
-                noisy_state_amps = self.simulate_batch(qc, batch_vars, peo)
-                noisy_state_probs = np.square(np.absolute(noisy_state_amps))
-                unnormalized_ensemble_probs += noisy_state_probs
         
+        xp = get_array_module(self)
+
+        unnormalized_ensemble_probs = xp.zeros(shape = 2**batch_vars, dtype = complex)
+        for _ in range(num_circs):
+            noisy_state_amps = self.simulate_batch(qc, batch_vars, peo)
+            noisy_state_probs = xp.square(xp.absolute(noisy_state_amps))
+            unnormalized_ensemble_probs += noisy_state_probs
+
         normalized_ensemble_probs = unnormalized_ensemble_probs / num_circs
         self.normalized_ensemble_probs = normalized_ensemble_probs
         end = time.time_ns() / (10 ** 9)
         self.time_taken = end - start
-        #return normalized_ensemble_probs
 
     def simulate_batch_ensemble_density(self, qc, num_circs, batch_vars=0, peo=None):
-        """Stochastic noise simulator using density matrix apporach.
-        
-        Uses significantly more memory than the statevector approach, but gives more information. 
-        The user will get the relative phaseses as well as the probabilities.
-        """
+        """Stochastic noise simulator using density matrix apporach."""
         start = time.time_ns() / (10 ** 9)
         if num_circs < 0 or not isinstance(num_circs, int):
             raise Exception("Error: The argument num_circs must be a positive integer")
         
-        unnormalized_ensemble_density_matrix = np.zeros(shape=(2**batch_vars, 2**batch_vars), dtype=complex)
+        xp = get_array_module(self)
+
+        unnormalized_ensemble_density_matrix = xp.zeros(shape=(2**batch_vars, 2**batch_vars), dtype=complex)
         for _ in range(num_circs):
             noisy_state_amps = self.simulate_batch(qc, batch_vars, peo)
             conj_noisy_state_amps = noisy_state_amps.conj()
-            noisy_state_density_matrix = np.outer(conj_noisy_state_amps, noisy_state_amps)
+            noisy_state_density_matrix = xp.outer(conj_noisy_state_amps, noisy_state_amps)
             unnormalized_ensemble_density_matrix += noisy_state_density_matrix
         
-        self.normalized_ensemble_density_matrix = np.divide(unnormalized_ensemble_density_matrix, num_circs)
+        self.normalized_ensemble_density_matrix = xp.divide(unnormalized_ensemble_density_matrix, num_circs)
         end = time.time_ns() / (10 ** 9)
         self.time_taken = end - start
-        #normalized_ensemble_density_matrix = np.divide(unnormalized_ensemble_density_matrix, num_circs)
-        #return normalized_ensemble_density_matrix
     
     def simulate_ensemble(self, qc, num_circs):
         """Simulates and returns only the first amplitude of the ensemble"""
@@ -87,7 +73,7 @@ def _new_circuit(self, ideal_qc):
         
     def _apply_channel(self, gate):
         """A noisy gate has all of the proper noise channels applied to it"""
-
+        
         for i in range(len(self.noise_model.noise_gates[gate.name].channels)):
             error_name = self.noise_model.noise_gates[gate.name].channels[i].name
             if error_name == 'depolarizing':
diff --git a/qtensor/noise_simulator/comparison_mpi.py b/qtensor/noise_simulator/comparison_mpi.py
index 9c469d01..1811c9f7 100644
--- a/qtensor/noise_simulator/comparison_mpi.py
+++ b/qtensor/noise_simulator/comparison_mpi.py
@@ -39,8 +39,8 @@
 
 num_jobs_per_node = num_cores_per_node * num_threads_per_core * num_gpus_per_node
 
-num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
-#num_circs_list = [10]
+#num_circs_list = [10, 18, 32, 100, 178, 316, 1000, 1780, 3160, 10000]
+num_circs_list = [10]
 
 outfile_name = '{}.json'.format(datetime.now().isoformat())
 results = []
@@ -54,6 +54,16 @@
 min_d = 4
 max_d = 5
 
+from qtensor.tools.lazy_import import cupy as cp
+import numpy as np
+def checkType(a_list):
+    for element in a_list:
+        for i in element:
+            if isinstance(i, cp.ndarray) or isinstance(i, np.ndarray):
+                f"{i} is an ndarray type" 
+            if isinstance(element[i], cp.ndarray) or isinstance(element[i], np.ndarray):
+                f"{element[i]} is an ndarray type" 
+
 for n in range(min_n, max_n):
     for p in range(min_p, max_p):
         for d in range(min_d, max_d):
@@ -73,4 +83,5 @@
                 comparison.qtensor_qiskit_noisy_qaoa_mpi(num_nodes=num_nodes, num_jobs_per_node=num_jobs_per_node)
                 results.extend(comparison.results)
 if results:
+    #checkType(results)
     save_dict_to_file(results, name = outfile_name) 

From 23a8c5a9b31f44ca375f10c30c2068bc41796f59 Mon Sep 17 00:00:00 2001
From: William Berquist <uncreativeusername5321@gmail.com>
Date: Tue, 11 Jul 2023 20:58:33 -0500
Subject: [PATCH 24/24] removed some obsolete comments

---
 qtensor/noise_simulator/README.md | 40 -------------------------------
 1 file changed, 40 deletions(-)

diff --git a/qtensor/noise_simulator/README.md b/qtensor/noise_simulator/README.md
index 772f459b..d5874206 100644
--- a/qtensor/noise_simulator/README.md
+++ b/qtensor/noise_simulator/README.md
@@ -2,43 +2,3 @@
 
 This folder uses stochastic simulation of quantum circuits
 and benchmarks it against qiskit density matrix simulations.
-
-
-## Observations
-
-Two types of circuits are considered:
-
-1. QAOA MaxCut with CZ two-qubit gate
-2. Randomized brickwork circuit with XYZ gates and CX gate
-
-The noise model is depolarizing noise applied to every quantum gate.
-
-Two types of qiskit runs were considered: with passing noise_model to `execute` and to `backend`. These modes are named E and B below.
-Mode B is the correct one.
-
-The qiskit simulation results are then compared to qtensor simulation results and the discrepancy between the two is evaluated.
-
-### 1. QAOA circuits
-
-When running with mode E, we see that there is very small error and it reduces with number of circuits k. The same was with mode B.
-
-### 2. Generic circuits
-
-When running in mode E, we found that the error does not converge with number of circuits K. 
-In the mode B we found that the error is larger but converges better (by how much?) K, which indicates that we compare against the correct noisy output.
-
-When running in mode E, removing two-qubit gate error channels reduced the simulation error.
-
-When running in mode E, adding a layer of Hadamards in the begining of the circuit reduced the error significantly.
-
-When running in both modes E and B, increasing error rate in channel reduces the simulation error.
-
-In some cases hovewer, there was convergence of error in mode E
-
-These results are preliminary.
-
-When running in mode E with a layer of Hadamards at the beginning and end of the circuit, a depolazing error on 0.01 on single qubit gates but no error on two-qubit gates had significant error. 
-
-When running in mode E with a layer of Hadamards at the beginning and end of the circuit, a depolazing error on 0.01 on single qubit gates and 0.1 on two-qubit gates significantly reduced the error. 
-
-When running in mode E with a layer of Hadamards at the beginning and end of the circuit, a depolazing error on 0.25 on single qubit gates and 0.1 on two-qubit gates caused the error to return. 
\ No newline at end of file