Return QuTiP backend results as OQD datastores

docs/reference/qutip.md

@@ -1 +1,19 @@
 ::: oqd_trical.backend.qutip
+
+## Datastore output
+
+`QutipBackend.run(...)` returns an `oqd_dataschema.Datastore` with one
+`"emulation"` group. The group stores `tspan`, `states`, and `final_state` as
+HDF5-serializable datasets, plus run metadata in `attrs`.
+
+```python
+backend = QutipBackend()
+experiment, hilbert_space = backend.compile(circuit, fock_cutoff=4)
+datastore = backend.run(experiment, hilbert_space, timestep=1e-7)
+
+emulation = datastore["emulation"]
+tspan = emulation.tspan.data
+states = emulation.states.data
+
+datastore.model_dump_hdf5("trical_run.h5")
+```

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "scipy",
     "sympy",
     "qutip",
+    "oqd-dataschema",
     "oqd-core",
     "oqd-compiler-infrastructure",
     "dynamiqs>=0.3.1",
@@ -59,6 +60,7 @@ fixable = ["ALL"]
 [tool.uv.sources]
 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/glasses666/oqd-dataschema", branch = "codex/unitaryhack-44-trical-group" }
 
 [dependency-groups]
 dev = ["jupyter>=1.1.1", "pre-commit>=4.1.0", "ruff>=0.9.6"]

src/oqd_trical/backend/qutip/base.py

@@ -12,10 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+
+import numpy as np
 from oqd_compiler_infrastructure import Chain, Post, Pre
 from oqd_core.backend.base import BackendBase
 from oqd_core.compiler.atomic.canonicalize import canonicalize_atomic_circuit_factory
 from oqd_core.interface.atomic import AtomicCircuit
+from oqd_dataschema import Dataset, Datastore, TrICalEmulatorDataGroup
 
 from oqd_trical.backend.qutip.codegen import QutipCodeGeneration
 from oqd_trical.backend.qutip.vm import QutipVM
@@ -56,6 +60,51 @@ def __init__(
         self.solver = solver
         self.solver_options = solver_options
 
+    @staticmethod
+    def _qobj_to_array(state):
+        state_array = np.asarray(state.full(), dtype=np.complex128)
+        if state.isket:
+            return state_array.reshape(-1)
+        return state_array
+
+    @staticmethod
+    def _metadata_attrs(result, timestep, solver):
+        hilbert_space = result["hilbert_space"]
+        hilbert_space_sizes = {
+            name: int(size) for name, size in hilbert_space.size.items()
+        }
+        fock_cutoff = {
+            name: size for name, size in hilbert_space_sizes.items() if name[0] == "P"
+        }
+        frame = result["frame"]
+
+        return {
+            "backend": "qutip",
+            "solver": solver,
+            "timestep": float(timestep),
+            "hilbert_space": json.dumps(hilbert_space_sizes, sort_keys=True),
+            "fock_cutoff": json.dumps(fock_cutoff, sort_keys=True),
+            "frame": "none" if frame is None else "present",
+            "frame_type": "none" if frame is None else frame.__class__.__name__,
+        }
+
+    @classmethod
+    def _result_to_datastore(cls, result, timestep, solver):
+        states = np.stack([cls._qobj_to_array(state) for state in result["states"]])
+        final_state = cls._qobj_to_array(result["final_state"])
+        tspan = np.asarray(result["tspan"], dtype=np.float64)
+
+        return Datastore(
+            groups={
+                "emulation": TrICalEmulatorDataGroup(
+                    tspan=Dataset(data=tspan),
+                    states=Dataset(data=states),
+                    final_state=Dataset(data=final_state),
+                    attrs=cls._metadata_attrs(result, timestep, solver),
+                )
+            }
+        )
+
     def compile(self, circuit, fock_cutoff, *, relabel=True):
         """
         Compiles a AtomicCircuit or AtomicEmulatorCircuit to a [`QutipExperiment`][oqd_trical.backend.qutip.interface.QutipExperiment].
@@ -127,7 +176,7 @@ def run(self, experiment, hilbert_space, timestep, *, initial_state=None):
             timestep (float): Timestep between tracked states of the evolution.
 
         Returns:
-            result (Dict[str,Any]): Result of execution of [`QutipExperiment`][oqd_trical.backend.qutip.interface.QutipExperiment].
+            result (Datastore): Result of execution of [`QutipExperiment`][oqd_trical.backend.qutip.interface.QutipExperiment].
         """
         vm = Pre(
             QutipVM(
@@ -141,4 +190,4 @@ def run(self, experiment, hilbert_space, timestep, *, initial_state=None):
 
         vm(experiment)
 
-        return vm.children[0].result
+        return self._result_to_datastore(vm.children[0].result, timestep, self.solver)

tests/test_backends.py

@@ -14,10 +14,14 @@
 
 ########################################################################################
 import dynamiqs as dq
+import numpy as np
 import pytest
 import qutip as qt
+from oqd_dataschema import Datastore, TrICalEmulatorDataGroup
 
 from oqd_trical.backend.dynamiqs.vm import DynamiqsVM
+from oqd_trical.backend.qutip.base import QutipBackend
+from oqd_trical.backend.qutip.interface import QutipExperiment, QutipGate
 from oqd_trical.backend.qutip.vm import QutipVM
 from oqd_trical.light_matter.compiler.analysis import HilbertSpace
 
@@ -50,3 +54,34 @@ def test_dynamiqs_fail(self):
         initial_state = dq.tensor(dq.basis(2, 0), dq.basis(3, 0))
 
         DynamiqsVM(hilbert_space=hilbert_space, timestep=1, initial_state=initial_state)
+
+
+class TestQutipDatastore:
+    def test_run_returns_hdf5_roundtrippable_datastore(self, tmp_path):
+        hilbert_space = HilbertSpace(hilbert_space=dict(E0={0, 1}))
+        experiment = QutipExperiment(
+            frame=None,
+            sequence=[QutipGate(hamiltonian=None, duration=1.0)],
+        )
+
+        datastore = QutipBackend(solver_options={"progress_bar": False}).run(
+            experiment, hilbert_space, timestep=0.5
+        )
+
+        assert isinstance(datastore, Datastore)
+        emulation = datastore["emulation"]
+        assert isinstance(emulation, TrICalEmulatorDataGroup)
+        np.testing.assert_allclose(emulation.tspan.data, np.array([0.0, 0.5, 1.0]))
+        assert emulation.states.data.shape == (3, 2)
+        assert emulation.final_state.data.shape == (2,)
+        assert emulation.attrs["backend"] == "qutip"
+        assert emulation.attrs["solver"] == "SESolver"
+        assert emulation.attrs["hilbert_space"] == '{"E0": 2}'
+
+        f = tmp_path / "qutip_run.h5"
+        datastore.model_dump_hdf5(f)
+        loaded = Datastore.model_validate_hdf5(f)
+
+        np.testing.assert_allclose(
+            loaded["emulation"].states.data, emulation.states.data
+        )