progress

qiskit_ibm_runtime/utils/estimator_pub_result.py

@@ -23,7 +23,8 @@ def plot_zne(self, tol: float = 2e-1) -> go.Figure:
         """Plot the zero noise extrapolation data contained in this estimator pub result.
 
         Args:
-            tol: The tolerance. If ``stds_extrapolated`` is greater than this value for a expectation value and extrapolator, the fit is omitted from the plot.
+            tol: The tolerance. If ``stds_extrapolated`` is greater than this value for a expectation value
+            and extrapolator, the fit is omitted from the plot.
 
         Returns:
             A plotly figure.

qiskit_ibm_runtime/visualization/__init__.py

@@ -27,6 +27,8 @@
     :toctree: ../stubs/
 
     draw_layer_error_map
+    plot_zne
 """
 
 from .draw_layer_error_map import draw_layer_error_map
+from .plot_zne import plot_zne

qiskit_ibm_runtime/visualization/plot_zne.py

@@ -10,57 +10,80 @@
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
 
-from typing import TYPE_CHECKING
-
-from plotly.colors import DEFAULT_PLOTLY_COLORS
+from itertools import product
+from plotly.colors import sample_colorscale, diverging
 from plotly.subplots import make_subplots
+import numpy as np
 import plotly.graph_objects as go
 
-from qiskit.primitives.containers import PubResult
+from ..utils.estimator_pub_result import EstimatorPubResult
+
+
+def plot_zne(result: EstimatorPubResult, tol: float = 2.0e-1) -> go.Figure:
+    """Plot an :class:`~.EstimatorPubResult` with zero noise extrapolation data.
 
-from ..options.zne_options import ExtrapolatorType
+    Args:
+        result: An :class:`~.EstimatorPubResult`.
+        tol: The tolerance. If ``stds_extrapolated`` is greater than this value for a expectation value
+        and extrapolator, the fit is omitted from the plot.
 
+    Returns:
+        A plotly figure.
 
-def plot_zne(
-    result: PubResult,
-    extrapolated_noise_factors: list[int],
-    models: list[ExtrapolatorType],
-) -> go.Figure:
-    if not result.metadata["resilience"].get("zne"):
+    Raises:
+        ValueError: If ``result`` does not contain zero noise extrapolation data.
+    """
+    if not (zne_metadata := result.metadata["resilience"].get("zne")):
         raise ValueError("Result does not contain ZNE data.")
 
-    nfs_fill = extrapolated_noise_factors + extrapolated_noise_factors[::-1]
+    noise_factors = zne_metadata["noise_factors"]
+    extrapolated_noise_factors = zne_metadata["extrapolated_noise_factors"]
+    extrapolators = zne_metadata["extrapolators"]
 
-    shape = (-1, len(extrapolated_noise_factors), len(models))
-    evs_reshaped = result.data.evs_extrapolated.reshape(shape)
-    evs_std_reshaped = result.data.stds_extrapolated.reshape(shape)
+    fig = make_subplots(cols=len(extrapolators), subplot_titles=extrapolators)
 
-    fig = make_subplots(cols=len(models), subplot_titles=models)
+    # generate the x coordinates for the fill around the fit
+    e_nfs_fill = extrapolated_noise_factors + extrapolated_noise_factors[::-1]
 
-    for idx in range(result.data.size):
+    colors = sample_colorscale(diverging.Spectral, np.linspace(0, 1, result.data.size))
+
+    for idx, obs in enumerate(product(*(range(s) for s in result.data.shape))):
         show_legend = True
-        for idx_m in range(len(models)):
-            evs = evs_reshaped[(idx, idx_m)]
-            evs_std = evs_std_reshaped[(idx, idx_m)]
-            evs_fill_up = (evs + evs_std).tolist()
-            evs_fill_down = (evs - evs_std).tolist()
+        color = colors[idx]
+        for idx_m in range(len(extrapolators)):
+            evs = result.data.evs_extrapolated[obs + (idx_m,)]
+            evs_std = result.data.stds_extrapolated[obs + (idx_m,)]
+            if any(evs_std > tol):
+                continue
+
+            # generate the y coordinates for the fill around the fit
+            evs_fill = (evs + evs_std).tolist() + (evs - evs_std).tolist()[::-1]
 
             fig.add_traces(
                 [
+                    go.Scatter(
+                        x=noise_factors,
+                        y=result.data.evs_noise_factors[obs],
+                        name=f"evs_{obs if len(obs) != 1 else obs[0]}",
+                        mode="markers",
+                        marker=dict(color=color),
+                        error_y=dict(array=result.data.stds_noise_factors[obs]),
+                        legendgroup=idx,
+                        showlegend=show_legend,
+                    ),
                     go.Scatter(
                         x=extrapolated_noise_factors,
                         y=evs,
-                        name=f"observable_{idx}",
-                        marker=dict(color=DEFAULT_PLOTLY_COLORS[idx]),
-                        error_y=dict(array=evs_std),
+                        mode="lines",
+                        line=dict(color=color),
                         legendgroup=idx,
-                        showlegend=show_legend,
+                        showlegend=False,
                     ),
                     go.Scatter(
-                        x=nfs_fill,
-                        y=evs_fill_up + evs_fill_down[::-1],  # upper, then lower reversed
+                        x=e_nfs_fill,
+                        y=evs_fill,
                         fill="toself",
-                        fillcolor=DEFAULT_PLOTLY_COLORS[idx],
+                        fillcolor=color,
                         line=dict(color="rgba(255,255,255,0)"),
                         opacity=0.2,
                         legendgroup=idx,
@@ -73,4 +96,11 @@ def plot_zne(
             )
             show_legend = False
 
+    fig.update_yaxes(
+        range=[
+            np.min(result.data.evs_noise_factors) - 0.5,
+            np.max(result.data.evs_noise_factors) + 0.5,
+        ],
+    )
+
     return fig