Return Datastore from QutipBackend.run

README.md

@@ -77,9 +77,15 @@ args = TaskArgsAnalog(
 task = Task(program=circuit, args=args)
 
 backend = QutipBackend()
-results = backend.run(task=task)
+datastore = backend.run(task=task)
+datastore.model_dump_hdf5("rabi_run.h5")
 
-plt.plot(results.times, results.metrics["Z"], label=f"$\\langle Z \\rangle$")
+sim = datastore.groups["emulation"]
+times = sim.times.data
+z_idx = sim.metrics.attrs["metric_labels"].split(",").index("Z")
+z = sim.metrics.data[:, z_idx]
+
+plt.plot(times, z, label=f"$\\langle Z \\rangle$")
 ```
 
 ### Where in the stack

docs/explanation/qutip-simulation.md

@@ -17,5 +17,5 @@ After compilation, the time dynamical evolution is emulated using the Qutip impl
     User_Input --> OpenQSIM: The user inputs are put in the OpenQSIM AST's Task object
     OpenQSIM --> OpenQSIM(canonicalized): RewriteRule (Canonicalization of Operators)
     OpenQSIM(canonicalized) --> QutipExperiment: compile the program of Task to QutipExperiment. And also convert the analog
-    QutipExperiment --> TaskResultAnalog: Run the experiment using the QutipExperimentVM Virtual Machine.
+    QutipExperiment --> Datastore: Run the experiment using the QutipExperimentVM Virtual Machine.
 ```

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
   "setuptools",
   "oqd-compiler-infrastructure@git+https://github.com/openquantumdesign/oqd-compiler-infrastructure",
   "oqd-core@git+https://github.com/openquantumdesign/oqd-core",
+  "oqd-dataschema@git+https://github.com/OpenQuantumDesign/oqd-dataschema",
 ]
 
 [project.optional-dependencies]

src/oqd_analog_emulator/conversion.py

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
 import itertools
 import time
+
+import numpy as np
 import qutip as qt
 
 from oqd_compiler_infrastructure import ConversionRule, RewriteRule, Chain
-from oqd_core.backend.task import TaskResultAnalog
 from oqd_core.compiler.math.passes import (
     evaluate_math_expr,
     simplify_math_expr,
@@ -44,6 +44,12 @@
     "QutipExperimentVM",
 ]
 
+
+def _qobj_to_state_vector(state: qt.Qobj) -> np.ndarray:
+    """Extract a 1D complex128 state vector from a QuTiP Qobj."""
+    return np.asarray(state.full(), dtype=np.complex128).squeeze()
+
+
 ########################################################################################
 
 
@@ -95,63 +101,53 @@ def map_EntanglementEntropyVN(self, model, operands):
 
 class QutipExperimentVM(RewriteRule):
     """
-    This is a Virtual Machine which takes in a QutipExperiment object, simulates the experiment and then produces the results
-
-    Args:
-        model (QutipExperiment): This is the compiled  [`QutipExperiment`][oqd_analog_emulator.qutip_backend.QutipExperiment] object
-
-    Returns:
-        task (TaskResultAnalog):
-
-    Note:
-        n_qreg and n_qmode are given as compiler parameters
+    Virtual machine that simulates a [`QutipExperiment`][oqd_analog_emulator.interface.QutipExperiment]
+    and collects run data as plain numpy-friendly attributes for datastore export.
+
+    Attributes populated during execution:
+        times: simulation time points
+        metric_labels: ordered metric names
+        metrics: expectation time-series per metric
+        state_trajectory: state vectors at each time step
+        measurements: sampled qubit outcomes after a measure instruction
+        runtime: total solver wall time in seconds
     """
 
     def __init__(self, qt_metrics, n_shots, fock_cutoff, dt):
         super().__init__()
-        self.results = TaskResultAnalog(runtime=0)
         self._qt_metrics = qt_metrics
         self._n_shots = n_shots
         self._fock_cutoff = fock_cutoff
         self._dt = dt
 
+        self.metric_labels: list[str] = list(qt_metrics.keys())
+        self.times: list[float] = []
+        self.metrics: dict[str, list[float]] = {key: [] for key in self.metric_labels}
+        self.state_trajectory: list[np.ndarray] = []
+        self.measurements: np.ndarray | None = None
+        self.runtime: float = 0.0
+
     def map_QutipExperiment(self, model):
         dims = model.n_qreg * [2] + model.n_qmode * [self._fock_cutoff]
         self.n_qreg = model.n_qreg
         self.n_qmode = model.n_qmode
         self.current_state = qt.tensor([qt.basis(d, 0) for d in dims])
 
-        self.results.times.append(0.0)
-        self.results.state = list(
-            self.current_state.full().squeeze(),
-        )
-        self.results.metrics.update(
-            {
-                key: [self._qt_metrics[key](0.0, self.current_state)]
-                for key in self._qt_metrics.keys()
-            }
-        )
+        self.times.append(0.0)
+        self.state_trajectory.append(_qobj_to_state_vector(self.current_state))
+        for key in self.metric_labels:
+            self.metrics[key].append(self._qt_metrics[key](0.0, self.current_state))
 
     def map_QutipMeasurement(self, model):
-        if self._n_shots is None:
-            self.results.counts = {}
-        else:
+        if self._n_shots is not None:
             probs = np.power(np.abs(self.current_state.full()), 2).squeeze()
-            n_shots = self._n_shots
-            inds = np.random.choice(len(probs), size=n_shots, p=probs)
+            inds = np.random.choice(len(probs), size=self._n_shots, p=probs)
             opts = self.n_qreg * [[0, 1]] + self.n_qmode * [
                 list(range(self._fock_cutoff))
             ]
             bases = list(itertools.product(*opts))
             shots = np.array([bases[ind] for ind in inds])
-            bitstrings = ["".join(map(str, shot)) for shot in shots]
-            self.results.counts = {
-                bitstring: bitstrings.count(bitstring) for bitstring in bitstrings
-            }
-
-        self.results.state = list(
-            self.current_state.full().squeeze(),
-        )
+            self.measurements = shots[:, : self.n_qreg].astype(np.int64)
 
     def map_QutipOperation(self, model):
         duration = model.duration
@@ -171,18 +167,18 @@ def map_QutipOperation(self, model):
             e_ops=self._qt_metrics,
             options={"store_states": True},
         )
-        self.results.runtime = time.time() - start_runtime + self.results.runtime
+        self.runtime += time.time() - start_runtime
 
-        self.results.times.extend([t + self.results.times[-1] for t in tspan][1:])
+        self.times.extend([t + self.times[-1] for t in tspan][1:])
 
-        for idx, key in enumerate(self.results.metrics.keys()):
-            self.results.metrics[key].extend(result_qobj.expect[idx].tolist()[1:])
+        for idx, key in enumerate(self.metric_labels):
+            self.metrics[key].extend(result_qobj.expect[idx].tolist()[1:])
 
-        self.current_state = result_qobj.final_state
+        if result_qobj.states is not None:
+            for st in result_qobj.states[1:]:
+                self.state_trajectory.append(_qobj_to_state_vector(st))
 
-        self.results.state = list(
-            result_qobj.final_state.full().squeeze(),
-        )
+        self.current_state = result_qobj.final_state
 
 
 class QutipBackendCompiler(ConversionRule):

src/oqd_analog_emulator/datastore.py

@@ -0,0 +1,102 @@
+# Copyright 2024-2025 Open Quantum Design
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from importlib.metadata import version
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from oqd_dataschema import AnalogEmulatorDataGroup, Datastore, Dataset
+
+if TYPE_CHECKING:
+    from oqd_analog_emulator.conversion import QutipExperimentVM
+    from oqd_analog_emulator.interface import TaskArgsQutip
+
+########################################################################################
+
+__all__ = [
+    "EMULATION_GROUP_KEY",
+    "METRIC_LABELS_ATTR",
+    "vm_to_datastore",
+]
+
+EMULATION_GROUP_KEY = "emulation"
+METRIC_LABELS_ATTR = "metric_labels"
+
+########################################################################################
+
+
+def vm_to_datastore(vm: QutipExperimentVM, args: TaskArgsQutip) -> Datastore:
+    """
+    Build a [`Datastore`][oqd_dataschema.datastore.Datastore] from VM run data.
+
+    Expects ``vm`` to hold plain numpy-friendly collections populated during
+    simulation (no ``TaskResultAnalog`` intermediate).
+    """
+    times = np.asarray(vm.times, dtype=np.float64)
+
+    metrics_dataset = None
+    if vm.metric_labels:
+        metrics_data = np.column_stack(
+            [
+                np.asarray(vm.metrics[label], dtype=np.float64)
+                for label in vm.metric_labels
+            ]
+        )
+        metrics_dataset = Dataset(
+            data=metrics_data,
+            attrs={METRIC_LABELS_ATTR: ",".join(vm.metric_labels)},
+        )
+
+    state_dataset = None
+    if vm.state_trajectory:
+        state_dataset = Dataset(data=np.vstack(vm.state_trajectory))
+
+    measurements_dataset = None
+    if vm.measurements is not None and vm.measurements.size > 0:
+        measurements_dataset = Dataset(
+            data=vm.measurements,
+            attrs={
+                "axis_0": "shots",
+                "axis_1": "qubits",
+            },
+        )
+
+    try:
+        pkg_version = version("oqd-analog-emulator")
+    except Exception:
+        pkg_version = "unknown"
+
+    group_attrs = {
+        "backend": "qutip",
+        "version": pkg_version,
+        "dt": args.dt,
+        "fock_cutoff": args.fock_cutoff,
+        "layer": args.layer,
+        "runtime": vm.runtime,
+    }
+    if args.n_shots is not None:
+        group_attrs["n_shots"] = args.n_shots
+
+    emulation = AnalogEmulatorDataGroup(
+        attrs=group_attrs,
+        times=Dataset(data=times),
+        metrics=metrics_dataset,
+        state=state_dataset,
+        measurements=measurements_dataset,
+    )
+
+    return Datastore(groups={EMULATION_GROUP_KEY: emulation})

src/oqd_analog_emulator/passes.py

@@ -17,6 +17,7 @@
     QutipExperimentVM,
     QutipMetricConversion,
 )
+from oqd_analog_emulator.datastore import vm_to_datastore
 
 from oqd_compiler_infrastructure import Post, Pre
 
@@ -63,27 +64,26 @@ def compiler_analog_args_to_qutipIR(model):
 
 def run_qutip_experiment(model: QutipExperimentVM, args):
     """
-    This takes in a [`QutipExperiment`][oqd_analog_emulator.interface.QutipExperiment] and produces a TaskResultAnalog object
+    Run a [`QutipExperiment`][oqd_analog_emulator.interface.QutipExperiment] and return a dataschema [`Datastore`][oqd_dataschema.datastore.Datastore].
 
     Args:
         model (QutipExperiment):
-        args: (Qutip
+        args: Compiled QuTiP task arguments.
 
     Returns:
-        (TaskResultAnalog): Contains results of the simulation
+        Datastore containing an [`AnalogEmulatorDataGroup`][oqd_dataschema.groups.analog_emulator.AnalogEmulatorDataGroup] under the ``emulation`` key.
 
     """
     n_qreg = model.n_qreg
     n_qmode = model.n_qmode
     metrics = Post(QutipMetricConversion(n_qreg=n_qreg, n_qmode=n_qmode))(args.metrics)
-    interpreter = Pre(
-        QutipExperimentVM(
-            qt_metrics=metrics,
-            n_shots=args.n_shots,
-            fock_cutoff=args.fock_cutoff,
-            dt=args.dt,
-        )
+    vm = QutipExperimentVM(
+        qt_metrics=metrics,
+        n_shots=args.n_shots,
+        fock_cutoff=args.fock_cutoff,
+        dt=args.dt,
     )
+    interpreter = Pre(vm)
     interpreter(model=model)
 
-    return interpreter.children[0].results
+    return vm_to_datastore(vm, args)

src/oqd_analog_emulator/qutip_backend.py

@@ -100,7 +100,8 @@ def run(
             task (Optional[Task]): Run experiment from a [`Task`][oqd_core.backend.task.Task] object
 
         Returns:
-            TaskResultAnalog object containing the simulation results.
+            [`Datastore`][oqd_dataschema.datastore.Datastore] with simulation results
+            under the ``emulation`` group.
 
         Note:
             only one of task or experiment must be provided.