Add serialization support for DoubleSliceSpan (#1989)

* Add serialization support for DoubleSliceSpan * Apply suggestions from code review * Update test/unit/test_data_serialization.py * Update qiskit_ibm_runtime/utils/json.py Co-authored-by: joshuasn <[email protected]> * Update test/unit/test_data_serialization.py --------- Co-authored-by: joshuasn <[email protected]>
Qiskit · Oct 22, 2024 · 9e3c397 · 9e3c397
1 parent 01dc4eb
commit 9e3c397
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/qiskit_ibm_runtime/utils/json.py b/qiskit_ibm_runtime/utils/json.py
@@ -75,7 +75,11 @@
 from qiskit_ibm_runtime.options.zne_options import (  # pylint: disable=ungrouped-imports
     ExtrapolatorType,
 )
-from qiskit_ibm_runtime.execution_span import SliceSpan, ExecutionSpans
+from qiskit_ibm_runtime.execution_span import (
+    DoubleSliceSpan,
+    SliceSpan,
+    ExecutionSpans,
+)
 
 from .noise_learner_result import NoiseLearnerResult
 
@@ -327,6 +331,16 @@ def default(self, obj: Any) -> Any:  # pylint: disable=arguments-differ
         if isinstance(obj, NoiseLearnerResult):
             out_val = {"data": obj.data, "metadata": obj.metadata}
             return {"__type__": "NoiseLearnerResult", "__value__": out_val}
+        if isinstance(obj, DoubleSliceSpan):
+            out_val = {
+                "start": obj.start,
+                "stop": obj.stop,
+                "data_slices": {
+                    idx: (shape, arg_sl.start, arg_sl.stop, shot_sl.start, shot_sl.stop)
+                    for idx, (shape, arg_sl, shot_sl) in obj._data_slices.items()
+                },
+            }
+            return {"__type__": "DoubleSliceSpan", "__value__": out_val}
         if isinstance(obj, SliceSpan):
             out_val = {
                 "start": obj.start,
@@ -450,6 +464,12 @@ def object_hook(self, obj: Any) -> Any:
                 return PrimitiveResult(**obj_val)
             if obj_type == "NoiseLearnerResult":
                 return NoiseLearnerResult(**obj_val)
+            if obj_type == "DoubleSliceSpan":
+                obj_val["data_slices"] = {
+                    int(idx): (tuple(shape), slice(arg0, arg1), slice(shot0, shot1))
+                    for idx, (shape, arg0, arg1, shot0, shot1) in obj_val["data_slices"].items()
+                }
+                return DoubleSliceSpan(**obj_val)
             if obj_type == "ExecutionSpan":
                 new_slices = {
                     int(idx): (tuple(shape), slice(*sl_args))

diff --git a/test/unit/test_data_serialization.py b/test/unit/test_data_serialization.py
@@ -47,7 +47,11 @@
     NoiseLearnerResult,
 )
 from qiskit_ibm_runtime.fake_provider import FakeNairobiV2
-from qiskit_ibm_runtime.execution_span import SliceSpan, ExecutionSpans
+from qiskit_ibm_runtime.execution_span import (
+    DoubleSliceSpan,
+    SliceSpan,
+    ExecutionSpans,
+)
 
 from .mock.fake_runtime_client import CustomResultRuntimeJob
 from .mock.fake_runtime_service import FakeRuntimeService
@@ -451,6 +455,19 @@ def make_test_primitive_results(self):
                         SliceSpan(
                             datetime(2024, 8, 20), datetime(2024, 8, 21), {0: ((14,), slice(2, 3))}
                         ),
+                        DoubleSliceSpan(
+                            datetime(2022, 1, 1),
+                            datetime(2023, 1, 1),
+                            {
+                                1: ((100,), slice(4, 9), slice(1, 2)),
+                                0: ((2, 5), slice(5, 7), slice(3, 4)),
+                            },
+                        ),
+                        DoubleSliceSpan(
+                            datetime(2024, 8, 20),
+                            datetime(2024, 8, 21),
+                            {0: ((14,), slice(2, 3), slice(1, 9))},
+                        ),
                     ]
                 )
             }