Implement RemoteEmuMPSEmulator (#45)

Author: Yoric

examples/tutorial 1 - Using a Quantum Device to Extract Machine-Learning Features.ipynb

@@ -156,7 +156,7 @@
     "\n",
     "if HAVE_PASQAL_ACCOUNT:\n",
     "    # Use the QPU Extractor.\n",
-    "    extractor = qek_extractors.QPUExtractor(\n",
+    "    extractor = qek_extractors.RemoteQPUExtractor(\n",
     "        # Once computing is complete, data will be saved in this file.\n",
     "        path=\"saved_data.json\",\n",
     "        compiler = compiler,\n",
@@ -191,10 +191,11 @@
     "\n",
     "There are two main ways to deal with this:\n",
     "\n",
-    "1. `QPUExtractor` can be attached to an ongoing job from batch ids, so that you can resume your work\n",
+    "1. `RemoteQPUExtractor` can be attached to an ongoing job from batch ids, so that you can resume your work\n",
     "    e.g. after turning off your computer.\n",
     "2. Pasqal CLOUD offers access to high-performance hardware-based emulators, with dramatically\n",
-    "    shorter waiting lines.\n",
+    "    shorter waiting lines. For instance, in the snippet above, you may replace `RemoteQPUExtractor`\n",
+    "    with `RemoteEmuMPSExtractor` to use the emu-mps emulator.\n",
     "\n",
     "See [the documentation](https://pqs.pages.pasqal.com/quantum-evolution-kernel/) for more details."
    ]

examples/tutorial 1a - Using a Quantum Device to Extract Machine-Learning Features - low-level.ipynb

@@ -37,7 +37,7 @@
    "source": [
     "# Load the original PTC-FM dataset\n",
     "import torch_geometric.datasets as pyg_dataset\n",
-    "og_ptcfm = [data for data in pyg_dataset.TUDataset(root=\"dataset\", name=\"PTC_FM\")]\n",
+    "og_ptcfm = pyg_dataset.TUDataset(root=\"dataset\", name=\"PTC_FM\")\n",
     "\n",
     "display(\"Loaded %s samples\" % (len(og_ptcfm), ))"
    ]

qek/data/extractors.py

@@ -12,7 +12,10 @@
 from numpy.typing import NDArray
 from pasqal_cloud import SDK
 from pasqal_cloud.batch import Batch
+from pasqal_cloud.device import BaseConfig, EmuTNConfig, EmulatorType
+from pasqal_cloud.job import Job
 from pasqal_cloud.utils.filters import BatchFilters
+from pathlib import Path
 import numpy as np
 import pulser as pl
 from pulser.devices import Device
@@ -134,14 +137,14 @@ def features(self, size_max: int | None) -> list[Feature]:
 
         return [Feature(dataset.dist_excitation(state, size_max)) for state in self.states]
 
-    def save_dataset(self, file_path: str) -> None:
+    def save_dataset(self, file_path: Path) -> None:
         """Saves the processed dataset to a JSON file.
 
         Note: This does NOT attempt to save the graphs.
 
         Args:
             dataset: The dataset to be saved.
-            file_path (str): The path where the dataset will be saved as a JSON
+            file_path: The path where the dataset will be saved as a JSON
                 file.
 
         Note:
@@ -249,7 +252,7 @@ class BaseExtractor(abc.ABC, Generic[GraphType]):
     """
 
     def __init__(
-        self, device: Device, compiler: BaseGraphCompiler[GraphType], path: str | None = None
+        self, device: Device, compiler: BaseGraphCompiler[GraphType], path: Path | None = None
     ) -> None:
         self.path = path
 
@@ -380,7 +383,7 @@ def __init__(
         self,
         compiler: BaseGraphCompiler[GraphType],
         device: Device = pl.devices.AnalogDevice,
-        path: str | None = None,
+        path: Path | None = None,
     ):
         super().__init__(path=path, device=device, compiler=compiler)
         self.graphs: list[BaseGraph]
@@ -463,7 +466,7 @@ def __init__(
         self,
         compiler: BaseGraphCompiler[GraphType],
         device: Device = pl.devices.AnalogDevice,
-        path: str | None = None,
+        path: Path | None = None,
     ):
         super().__init__(device=device, compiler=compiler, path=path)
         self.graphs: list[BaseGraph]
@@ -534,7 +537,7 @@ def run(self, max_qubits: int = 10, dt: int = 10) -> BaseExtracted:
 SLEEP_DELAY_S = 2
 
 
-class CloudExtracted(BaseExtracted):
+class PasqalCloudExtracted(BaseExtracted):
     """
     Data extracted from the cloud API, i.e. we need wait for a remote server.
 
@@ -548,13 +551,29 @@ class CloudExtracted(BaseExtracted):
     """
 
     def __init__(
-        self, compiled: list[Compiled], batch_ids: list[str], sdk: SDK, path: str | None = None
+        self,
+        compiled: list[Compiled],
+        batch_ids: list[str],
+        sdk: SDK,
+        state_extractor: Callable[[Job, pl.Sequence], dict[str, int] | None],
+        path: Path | None = None,
     ):
+        """
+        Prepare for reception of data.
+
+        Arguments:
+            compiled: The result of compiling a set of graphs.
+            batch_ids: The ids of the batches on the cloud API, in the same order as `compiled`.
+            state_extractor: A callback used to extract the counter from a job.
+                Used as various cloud back-ends return different formats.
+            path: If provided, a path at which to save the results once they're available.
+        """
         self._compiled = compiled
         self._batch_ids = batch_ids
         self._results: SyncExtracted | None = None
         self._path = path
         self._sdk = sdk
+        self._state_extractor = state_extractor
 
     def _wait(self) -> None:
         """
@@ -643,10 +662,16 @@ def _ingest(self, batches: dict[str, Batch]) -> None:
             compiled = self._compiled[i]
             # Note: There's only one job per batch.
             assert len(batch.jobs) == 1
-            for _, job in batch.jobs.items():
+            for job in batch.jobs.values():
                 if job.status == "DONE":
-                    state_dict = job.result
-                    assert state_dict is not None
+                    state_dict = self._state_extractor(job, compiled.sequence)
+                    if state_dict is None:
+                        logger.warning(
+                            "Batch %s (graph %s) did not return a usable state, skipping",
+                            i,
+                            compiled.graph.id,
+                        )
+                        continue
                     raw_data.append(compiled.graph)
                     if compiled.graph.target is not None:
                         targets.append(compiled.graph.target)
@@ -698,14 +723,18 @@ def states(self) -> list[dict[str, int]]:
         return self._results.states
 
 
-class QPUExtractor(BaseExtractor[GraphType]):
+class BaseRemoteExtractor(BaseExtractor[GraphType], Generic[GraphType]):
     """
-    A Extractor that uses a distant physical QPU to run sequences
-    compiled from graphs.
+    An Extractor that uses a remote Quantum Device published
+    on Pasqal Cloud, to run sequences compiled from graphs.
 
-    Performance note: as of this writing, the waiting lines for a QPU
-    may be very long. You may use this Extractor to resume your workflow
-    with a computation that has been previously started.
+    Performance note (servers and interactive applications only):
+        If your code is meant to be executed as part of an interactive application or
+        a server, you should consider calling `await extracted` before your first call
+        to any of the methods of `extracted`. Otherwise, you will block the main thread.
+
+        If you are running this as part of an experiment, a Jupyter notebook, etc. you
+        may ignore this performance note.
 
     Args:
         path: Path to store the result of the run, for future uses.
@@ -726,10 +755,10 @@ def __init__(
         compiler: BaseGraphCompiler[GraphType],
         project_id: str,
         username: str,
+        device_name: str,
         password: str | None = None,
-        device_name: str = "FRESNEL",
         batch_ids: list[str] | None = None,
-        path: str | None = None,
+        path: Path | None = None,
     ):
         sdk = SDK(username=username, project_id=project_id, password=password)
 
@@ -745,13 +774,33 @@ def __init__(
     def batch_ids(self) -> list[str] | None:
         return self._batch_ids
 
-    def run(self) -> CloudExtracted:
+    @abc.abstractmethod
+    def run(
+        self,
+    ) -> PasqalCloudExtracted:
+        """
+        Launch the extraction.
+        """
+        raise Exception("Not implemented")
+
+    def _run(
+        self,
+        state_extractor: Callable[[Job, pl.Sequence], dict[str, int] | None],
+        emulator: EmulatorType | None,
+        config: BaseConfig | None,
+    ) -> PasqalCloudExtracted:
         if len(self.sequences) == 0:
             logger.warning("No sequences to run, did you forget to call compile()?")
-            return CloudExtracted(compiled=[], batch_ids=[], sdk=self._sdk, path=self.path)
+            return PasqalCloudExtracted(
+                compiled=[],
+                batch_ids=[],
+                sdk=self._sdk,
+                path=self.path,
+                state_extractor=state_extractor,
+            )
 
         device: pl.devices.Device = self.sequences[0].sequence.device
-        # As of this writing, the API doesn't support run longer than 500 jobs.
+        # As of this writing, the API doesn't support runs longer than 500 jobs.
         # If we want to add more runs, we'll need to split them across several jobs.
         max_runs = device.max_runs if isinstance(device.max_runs, int) else 500
 
@@ -764,6 +813,8 @@ def run(self) -> CloudExtracted:
                     compiled.sequence.to_abstract_repr(),
                     jobs=[{"runs": max_runs}],
                     wait=False,
+                    emulator=emulator,
+                    configuration=config,
                 )
                 logger.info(
                     "Remote execution of compiled graph #%s starting, batched with id %s",
@@ -778,6 +829,134 @@ def run(self) -> CloudExtracted:
             )
         assert len(self._batch_ids) == len(self.sequences)
 
-        return CloudExtracted(
-            compiled=self.sequences, batch_ids=self._batch_ids, sdk=self._sdk, path=self.path
+        return PasqalCloudExtracted(
+            compiled=self.sequences,
+            batch_ids=self._batch_ids,
+            sdk=self._sdk,
+            path=self.path,
+            state_extractor=state_extractor,
+        )
+
+
+class RemoteQPUExtractor(BaseRemoteExtractor[GraphType]):
+    """
+    An Extractor that uses a remote QPU published
+    on Pasqal Cloud, to run sequences compiled from graphs.
+
+    Performance note:
+        as of this writing, the waiting lines for a QPU
+        may be very long. You may use this Extractor to resume your workflow
+        with a computation that has been previously started.
+
+    Performance note (servers and interactive applications only):
+        If your code is meant to be executed as part of an interactive application or
+        a server, you should consider calling `await extracted` before your first call
+        to any of the methods of `extracted`. Otherwise, you will block the main thread.
+
+        If you are running this as part of an experiment, a Jupyter notebook, etc. you
+        may ignore this performance note.
+
+    Args:
+        path: Path to store the result of the run, for future uses.
+            To reload the result of a previous run, use `LoadExtractor`.
+        project_id: The ID of the project on the Pasqal Cloud API.
+        username: Your username on the Pasqal Cloud API.
+        password: Your password on the Pasqal Cloud API. If you leave
+            this to None, you will need to enter your password manually.
+        device_name: The name of the device to use. As of this writing,
+            the default value of "FRESNEL" represents the latest QPU
+            available through the Pasqal Cloud API.
+        batch_id: Use this to resume a workflow e.g. after turning off
+            your computer while the QPU was executing your sequences.
+    """
+
+    def __init__(
+        self,
+        compiler: BaseGraphCompiler[GraphType],
+        project_id: str,
+        username: str,
+        device_name: str = "FRESNEL",
+        password: str | None = None,
+        batch_ids: list[str] | None = None,
+        path: Path | None = None,
+    ):
+        super().__init__(
+            compiler=compiler,
+            project_id=project_id,
+            username=username,
+            device_name=device_name,
+            password=password,
+            batch_ids=batch_ids,
+            path=path,
+        )
+
+    def run(self) -> PasqalCloudExtracted:
+        return self._run(emulator=None, config=None, state_extractor=lambda job, _seq: job.result)
+
+
+class RemoteEmuMPSExtractor(BaseRemoteExtractor[GraphType]):
+    """
+    An Extractor that uses a remote high-performance emulator (EmuMPS)
+    published on Pasqal Cloud, to run sequences compiled from graphs.
+
+    Performance note (servers and interactive applications only):
+        If your code is meant to be executed as part of an interactive application or
+        a server, you should consider calling `await extracted` before your first call
+        to any of the methods of `extracted`. Otherwise, you will block the main thread.
+
+        If you are running this as part of an experiment, a Jupyter notebook, etc. you
+        may ignore this performance note.
+
+    Args:
+        path: Path to store the result of the run, for future uses.
+            To reload the result of a previous run, use `LoadExtractor`.
+        project_id: The ID of the project on the Pasqal Cloud API.
+        username: Your username on the Pasqal Cloud API.
+        password: Your password on the Pasqal Cloud API. If you leave
+            this to None, you will need to enter your password manually.
+        device_name: The name of the device to use. As of this writing,
+            the default value of "FRESNEL" represents the latest QPU
+            available through the Pasqal Cloud API.
+        batch_id: Use this to resume a workflow e.g. after turning off
+            your computer while the QPU was executing your sequences.
+    """
+
+    def __init__(
+        self,
+        compiler: BaseGraphCompiler[GraphType],
+        project_id: str,
+        username: str,
+        device_name: str = "FRESNEL",
+        password: str | None = None,
+        batch_ids: list[str] | None = None,
+        path: Path | None = None,
+    ):
+        super().__init__(
+            compiler=compiler,
+            project_id=project_id,
+            username=username,
+            device_name=device_name,
+            password=password,
+            batch_ids=batch_ids,
+            path=path,
+        )
+
+    def run(self, dt: int = 10) -> PasqalCloudExtracted:
+        def extractor(job: Job, sequence: pl.Sequence) -> dict[str, int] | None:
+            cutoff_duration = int(ceil(sequence.get_duration() / dt) * dt)
+            full_result = job.full_result
+            if full_result is None:
+                return None
+            result = full_result["bitstring"][cutoff_duration]
+            if result is None:
+                return None
+            assert isinstance(result, dict)
+            return result
+
+        return self._run(
+            emulator=EmulatorType.EMU_TN,
+            config=EmuTNConfig(
+                dt=dt,
+            ),
+            state_extractor=extractor,
         )