Merge pull request #61 from tiqi-group/pycrystal

Support Pycrystal v2.1.8, Realtime ScanParameter

Author: g-braeunlich

.gitignore

@@ -1,4 +1,5 @@
 emulator/
+output/
 poetry.lock
 .python-version
 

src/icon/server/__main__.py

@@ -59,18 +59,18 @@ def start_server() -> None:
     scheduler.start()
 
     number_of_pre_processing_workers = get_config().server.pre_processing.workers
-    pre_processing_update_queues: list[multiprocessing.Queue[UpdateQueue]] = []
+    pre_processing_update_queues: list[multiprocessing.Queue[UpdateQueue]] = [
+        multiprocessing.Queue() for _ in range(number_of_pre_processing_workers)
+    ]
 
-    for i in range(number_of_pre_processing_workers):
-        pre_processing_update_queues.append(multiprocessing.Queue())
-        pre_processing_worker = PreProcessingWorker(
+    for i, queue in enumerate(pre_processing_update_queues):
+        PreProcessingWorker(
             worker_number=i,
             hardware_processing_queue=icon.server.shared_resource_manager.hardware_processing_queue,
             pre_processing_queue=icon.server.shared_resource_manager.pre_processing_queue,
-            update_queue=pre_processing_update_queues[i],
+            update_queue=queue,
             manager=icon.server.shared_resource_manager.manager,
-        )
-        pre_processing_worker.start()
+        ).start()
 
     hardware_processing_worker = HardwareProcessingWorker(
         hardware_processing_queue=icon.server.shared_resource_manager.hardware_processing_queue,

src/icon/server/data_access/db_context/sqlite/alembic/versions/fc9af856df20_add_realtime_column_to_scanparameter_.py

@@ -0,0 +1,35 @@
+"""Add realtime column to ScanParameter model
+
+Revision ID: fc9af856df20
+Revises: 2bc7556020cf
+Create Date: 2025-09-18 13:13:16.093552
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "fc9af856df20"
+down_revision: str | None = "2bc7556020cf"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("scan_parameters", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("realtime", sa.Boolean(), nullable=False, server_default="False")
+        )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("scan_parameters", schema=None) as batch_op:
+        batch_op.drop_column("realtime")
+
+    # ### end Alembic commands ###

src/icon/server/data_access/models/sqlite/scan_parameter.py

@@ -2,6 +2,7 @@
 from typing import TYPE_CHECKING, Any
 
 import sqlalchemy
+import sqlalchemy.event
 import sqlalchemy.orm
 
 from icon.server.data_access.db_context.influxdb_v1 import DatabaseValueType
@@ -79,6 +80,9 @@ class ScanParameter(Base):
         back_populates="scan_parameters", lazy="joined"
     )
     """Relationship to the device associated with this parameter."""
+    realtime: sqlalchemy.orm.Mapped[bool] = sqlalchemy.orm.mapped_column(
+        default=False, nullable=False
+    )
 
     def __repr__(self) -> str:
         return f"<Parameter '{self.unique_id()}'>"
@@ -96,3 +100,20 @@ def unique_id(self) -> str:
             if self.device is not None
             else self.variable_id
         )
+
+
+@sqlalchemy.event.listens_for(ScanParameter, "before_insert")
+def receive_before_insert(
+    mapper: sqlalchemy.orm.Mapper[ScanParameter],
+    connection: sqlalchemy.engine.Connection,
+    target: ScanParameter,
+) -> None:
+    if target.realtime and target.variable_id != "Real Time":
+        raise ValueError(
+            f"Cannot set 'continuous' on {ScanParameter} with id {target.variable_id!r}"
+            " != 'Real Time'"
+        )
+
+
+def contains_realtime_parameter(params: list[ScanParameter]) -> bool:
+    return any(sp.realtime for sp in params)

src/icon/server/data_access/repositories/experiment_data_repository.py

@@ -11,7 +11,10 @@
 
 from icon.config.config import get_config
 from icon.server.data_access.db_context.influxdb_v1 import DatabaseValueType
-from icon.server.data_access.models.sqlite.scan_parameter import ScanParameter
+from icon.server.data_access.models.sqlite.scan_parameter import (
+    ScanParameter,
+    contains_realtime_parameter,
+)
 from icon.server.data_access.repositories.job_repository import JobRepository
 from icon.server.data_access.repositories.job_run_repository import JobRunRepository
 from icon.server.utils.h5py import get_hdf5_dtype, get_result_channels_dataset
@@ -362,6 +365,8 @@ def update_metadata_by_job_id(  # noqa: PLR0913
             h5file.attrs["experiment_id"] = job.experiment_source.experiment_id
             h5file.attrs["job_id"] = job_id
             h5file.attrs["repetitions"] = repetitions
+            h5file.attrs["realtime_scan"] = contains_realtime_parameter(parameters)
+
             if local_parameter_timestamp is not None:
                 h5file.attrs["local_parameter_timestamp"] = local_parameter_timestamp
 
@@ -582,6 +587,7 @@ def get_experiment_data_by_job_id(
             "vector_channels": {},
             "scan_parameters": {},
             "json_sequences": [],
+            "realtime_scan": False,
         }
 
         filename = get_filename_by_job_id(job_id)
@@ -596,6 +602,9 @@ def get_experiment_data_by_job_id(
             f"{ExperimentDataRepository.LOCK_EXTENSION}"
         )
         with FileLock(lock_path), h5py.File(file, "r") as h5file:
+            data["realtime_scan"] = h5file.attrs["realtime_scan"]
+            # Parse JSON strings in relevant columns back into Python objects
+
             scan_parameters: npt.NDArray = h5file["scan_parameters"][:]  # type: ignore
             data["scan_parameters"] = {
                 param: {

src/icon/server/data_access/repositories/pycrystal_library_repository.py

@@ -124,6 +124,7 @@ async def generate_json_sequence(
         exp_module_name: str,
         exp_instance_name: str,
         parameter_dict: dict[str, DatabaseValueType],
+        n_shots: int,
     ) -> str:
         """Generate a JSON sequence for an experiment.
 
@@ -140,6 +141,7 @@ async def generate_json_sequence(
             "key_val_dict": parameter_dict,
             "module_name": exp_module_name,
             "exp_instance_name": exp_instance_name,
+            "n_shots": n_shots,
         }
 
         code = PycrystalLibraryRepository._get_code(

src/icon/server/data_access/templates/generate_pycrystal_sequence.py

@@ -1,6 +1,7 @@
 import logging
 
 import pycrystal.database.local_cache
+import pycrystal.experiment
 import pycrystal.parameters
 from pycrystal.utils.helpers import get_config_from_module_name
 
@@ -9,6 +10,7 @@
 logging.getLogger("pycrystal").setLevel(log_level)
 logging.getLogger("ionpulse_sequence_generator").setLevel(log_level)
 KEY_VAL_DICT = {key_val_dict}
+N_SHOTS = {n_shots}
 
 pycrystal.parameters.Parameter.db = pycrystal.database.local_cache.LocalCache(
     key_val_dict=KEY_VAL_DICT,
@@ -32,6 +34,7 @@
 
 experiment_library.hardware_description.hardware.hardware.init()
 exp_instance._init()
+pycrystal.experiment.Experiment.shots = N_SHOTS
 exp_instance._initialize_scan(debug_level=log_level)
 sequence = exp_instance.pulse_sequence()
 sequence_json = sequence.get_json_string(exp_instance._sequence_header)

src/icon/server/hardware_processing/hardware_controller.py

@@ -40,10 +40,6 @@ def _update_zedboard_sequence(self, *, sequence: str) -> None:
         if self._zedboard is not None:
             self._zedboard.sequence_JSON_parser.Sequence_JSON = sequence  # type: ignore
 
-    def _update_number_of_shots(self, *, number_of_shots: int) -> None:
-        if self._zedboard is not None:
-            self._zedboard.sequence_JSON_parser.Shots = number_of_shots  # type: ignore
-
     def run(self, *, sequence: str, number_of_shots: int) -> ResultDict:
         if not self.connected:
             self.connect()
@@ -52,7 +48,6 @@ def run(self, *, sequence: str, number_of_shots: int) -> ResultDict:
             raise RuntimeError("Could not connect to the Zedboard")
 
         self._update_zedboard_sequence(sequence=sequence)
-        self._update_number_of_shots(number_of_shots=number_of_shots)
         self._zedboard.sequence_JSON_parser.Parse_JSON_Header()  # type: ignore
         results: tiqi_zedboard.zedboard.Result = self._zedboard.sequence_JSON_parser()  # type: ignore