From 852770eb27f26a3ae6becda6365ce34239e8dcf5 Mon Sep 17 00:00:00 2001
From: "Christopher J. Wood" <cjwood@us.ibm.com>
Date: Tue, 24 Jan 2023 12:12:16 -0500
Subject: [PATCH] Add mit tomo analysis option

Add mitigated tomography analysis option to optionally compute the unmitigated tomography fit in addition to the mitigated tomography fit.
---
 .../tomography/mit_tomography_analysis.py     | 39 ++++++++++++++-----
 1 file changed, 29 insertions(+), 10 deletions(-)

diff --git a/qiskit_experiments/library/tomography/mit_tomography_analysis.py b/qiskit_experiments/library/tomography/mit_tomography_analysis.py
index 435d2bc814..e23e746d65 100644
--- a/qiskit_experiments/library/tomography/mit_tomography_analysis.py
+++ b/qiskit_experiments/library/tomography/mit_tomography_analysis.py
@@ -62,6 +62,9 @@ def _default_options(cls):
             preparation_qubits (Sequence[int]): Optional, the physical qubits with tomographic
                 preparations. If not specified will be set to ``[0, ..., N-1]`` for N-qubit
                 tomographic preparations.
+            unmitigated_fit (bool): If True also run tomography fit without readout error
+                mitigation and include both mitigated and unmitigated analysis results. If
+                False only compute mitigated results (Default: False)
             target (Any): Optional, target object for fidelity comparison of the fit
                 (Default: None).
         """
@@ -73,12 +76,18 @@ def _default_options(cls):
         options.rescale_trace = True
         options.measurement_qubits = None
         options.preparation_qubits = None
+        options.unmitigated_fit = False
         options.target = None
         return options
 
     def set_options(self, **fields):
-        super().set_options(**fields)
-        self._analyses[1].set_options(**fields)
+        # filter fields
+        self_fields = {key: val for key, val in fields.items() if hasattr(self.options, key)}
+        super().set_options(**self_fields)
+        tomo_fields = {
+            key: val for key, val in fields.items() if hasattr(self._analyses[1].options, key)
+        }
+        self._analyses[1].set_options(**tomo_fields)
 
     def _run_analysis(self, experiment_data):
         # Return list of experiment data containers for each component experiment
@@ -88,21 +97,31 @@ def _run_analysis(self, experiment_data):
 
         # Run readout error analysis
         roerror_analysis.run(roerror_data, replace_results=True).block_for_results()
-        
+
         # Construct noisy measurement basis
         mitigator = roerror_data.analysis_results(0).value
-        
+
         # Construct noisy measurement basis
         measurement_basis = PauliMeasurementBasis(mitigator=mitigator)
         tomo_analysis.set_options(measurement_basis=measurement_basis)
-        
-        # Run tomography analysis
+
+        # Run mitigated tomography analysis
         tomo_analysis.run(tomo_data, replace_results=True).block_for_results()
+        for res in tomo_data.analysis_results(block=False):
+            res.extra["mitigated"] = True
+
+        # Combine results so that tomography results are ordered first
+        combined_data = [tomo_data, roerror_data]
+
+        # Run unmitigated tomography analysis
+        if self.options.unmitigated_fit:
+            tomo_analysis.set_options(measurement_basis=PauliMeasurementBasis())
+            nomit_data = tomo_analysis.run(tomo_data, replace_results=False).block_for_results()
+            for res in nomit_data.analysis_results(block=False):
+                res.extra["mitigated"] = False
+            combined_data.append(nomit_data)
 
-        # Optionally flatten results from all component experiments
-        # for adding to the main experiment data container
         if self._flatten_results:
-            # Combine results so that tomography results are ordered first
-            return self._combine_results([tomo_data, roerror_data])
+            return self._combine_results(combined_data)
 
         return [], []