Separate lab environment logic from icon

[g-braeunlich]

Currently, icon has hidden dependencies to pycrystal, experiment_library and ionpulse_sequence_generator via the templates in src/icon/server/data_access/templates/.

The goal is to introduce a clean architectural boundary between icon and lab environments.

Proposal

• Introduce a lab environment API on the icon side, which has to implemented on the lab environment side
• The implementation could either be integrated into experiment_library or could be placed in an own python package.
• The implementation would have to be installed in the same virtual environment as icon and has to be configured in icons config file

API

class LabEnvironment:
    """Abstraction API for lab environments.

    This API must be implemented by a lab environment plugin and
    will be loaded by the lab environment repository.
    """

    experiment_metadata: ExperimentDict
    """Dictionary mapping the unique experiment identifier to its metadata."""
    parameter_metadata: ParameterMetadataDict
    """Dictionary of parameter metadata."""

    def generate_json_sequence(
        self,
        *,
        exp_module_name: str,
        exp_instance_name: str,
        parameter_dict: dict[str, DatabaseValueType],
        n_shots: int,
    ) -> str:
        """Generate a JSON sequence for an experiment.

        Args:
            exp_module_name: Module name of the experiment.
            exp_instance_name: Name of the experiment instance.
            parameter_dict: Mapping of parameter IDs to values.

        Returns:
            JSON string containing the generated sequence.
        """
        raise NotImplementedError("Must be implemented by a subclass")

    def get_experiment_readout_metadata(
        self,
        *,
        exp_module_name: str,
        exp_instance_name: str,
        parameter_dict: dict[str, DatabaseValueType],
    ) -> ReadoutMetadata:
        """Fetch readout metadata for an experiment.

        Args:
            exp_module_name: Module name of the experiment.
            exp_instance_name: Name of the experiment instance.
            parameter_dict: Mapping of parameter IDs to values.

        Returns:
            Dictionary containing readout metadata for the experiment.
        """
        raise NotImplementedError("Must be implemented by a subclass")

Integration in icon

The API could be consumed by icon like this and would replace PycrystalLibraryRepository:

class LabEnvironmentRepository:
    def __init__(self, config: ExperimentLibraryConfigV1) -> None:
        # TODO: add "module" and "lab_environment_class" config options in `config`.
        module = importlib.import_module(config.module)
        lab_environment_class = getattr(module, config.lab_environment_class)
        self.lab_environment: LabEnvironment = lab_environment_class()

    def get_experiment_and_parameter_metadata(self) -> ParameterAndExperimentMetadata:
        """Fetch the experiment and parameter metadata.

        Returns:
            Dictionary with experiment metadata and parameter metadata.
        """
        return {
            "experiment_metadata": self.lab_environment.experiment_metadata,
            "parameter_metadata": self.lab_environment.parameter_metadata,
        }

    async def generate_json_sequence(
        self,
        *,
        exp_module_name: str,
        exp_instance_name: str,
        parameter_dict: dict[str, DatabaseValueType],
        n_shots: int,
    ) -> str:
        """Generate a JSON sequence for an experiment.

        Args:
            exp_module_name: Module name of the experiment.
            exp_instance_name: Name of the experiment instance.
            parameter_dict: Mapping of parameter IDs to values.

        Returns:
            JSON string containing the generated sequence.
        """
        return await asyncio.to_thread(
            self.lab_environment.generate_json_sequence,
            exp_module_name=exp_module_name,
            exp_instance_name=exp_instance_name,
            parameter_dict=parameter_dict,
            n_shots=n_shots,
        )

    async def get_experiment_readout_metadata(
        self,
        *,
        exp_module_name: str,
        exp_instance_name: str,
        parameter_dict: dict[str, DatabaseValueType],
    ) -> ReadoutMetadata:
        """Fetch readout metadata for an experiment.

        Args:
            exp_module_name: Module name of the experiment.
            exp_instance_name: Name of the experiment instance.
            parameter_dict: Mapping of parameter IDs to values.

        Returns:
            Dictionary containing readout metadata for the experiment.
        """
        return await asyncio.to_thread(
            self.lab_environment.get_experiment_readout_metadata,
            exp_module_name=exp_module_name,
            exp_instance_name=exp_instance_name,
            parameter_dict=parameter_dict,
        )

Implementation

On the lab environment side, an implementation could look like (basically the logic of the templates):

import logging
import pkgutil

import experiment_library.experiments
import experiment_library.hardware_description.hardware
import pycrystal.database.local_cache
import pycrystal.parameters
from ionpulse_sequence_generator import System
from pycrystal.parameters import Parameter
from pycrystal.utils.helpers import (
    ExperimentDict,
    get_config_from_module_name,
    get_experiment_metadata,
)

from icon.server.data_access.db_context.influxdb_v1 import DatabaseValueType
from icon.server.data_access.repositories.experiment_data_repository import (
    ReadoutMetadata,
)
from icon.server.data_access.repositories.lab_environment import LabEnvironment

log_level = logging.ERROR
logging.basicConfig(level=log_level)
logging.getLogger("pycrystal").setLevel(log_level)
logging.getLogger("ionpulse_sequence_generator").setLevel(log_level)


class PyCrystalLabEnvironment(LabEnvironment):
    def __init__(self) -> None:
        experiments: ExperimentDict = {}
        for mod_info in pkgutil.iter_modules(experiment_library.experiments.__path__):
            experiment_module = (
                experiment_library.experiments.__name__ + "." + mod_info.name
            )
            experiments.update(
                get_experiment_metadata(experiment_module=experiment_module)
            )
        self.experiment_metadata = experiments

        self.parameter_metadata = {
            "all parameters": Parameter.registry.all_parameters,
            "display groups": {
                f"{namespace} ({display_group})": parameter_dict
                for namespace, display_groups in Parameter.registry.namespace_registry.items()
                for display_group, parameter_dict in display_groups.items()
            },
        }

    def generate_json_sequence(
        self,
        *,
        exp_module_name: str,
        exp_instance_name: str,
        parameter_dict: dict[str, DatabaseValueType],
        n_shots: int,
    ) -> str:
        """Generate a JSON sequence for an experiment.

        Args:
            exp_module_name: Module name of the experiment.
            exp_instance_name: Name of the experiment instance.
            parameter_dict: Mapping of parameter IDs to values.

        Returns:
            JSON string containing the generated sequence.
        """
        pycrystal.parameters.Parameter.db = pycrystal.database.local_cache.LocalCache(
            key_val_dict=parameter_dict,
        )

        config = get_config_from_module_name(exp_module_name)
        exp_config = next(
            instance
            for instance in config["experiment_instances"]
            if instance[1]["name"] == exp_instance_name
        )
        exp_class = exp_config[0]
        exp_kwargs = exp_config[1]
        exp_instance = exp_class(**exp_kwargs)

        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()
        return sequence.get_json_string(exp_instance._sequence_header)

    def get_experiment_readout_metadata(
        self,
        *,
        exp_module_name: str,
        exp_instance_name: str,
        parameter_dict: dict[str, DatabaseValueType],
    ) -> ReadoutMetadata:
        """Fetch readout metadata for an experiment.

        Args:
            exp_module_name: Module name of the experiment.
            exp_instance_name: Name of the experiment instance.
            parameter_dict: Mapping of parameter IDs to values.

        Returns:
            Dictionary containing readout metadata for the experiment.
        """
        pycrystal.parameters.Parameter.db = pycrystal.database.local_cache.LocalCache(
            key_val_dict=parameter_dict,
        )

        config = get_config_from_module_name(exp_module_name)
        exp_config = next(
            instance
            for instance in config["experiment_instances"]
            if instance[1]["name"] == exp_instance_name
        )
        exp_class = exp_config[0]
        exp_kwargs = exp_config[1]
        exp_instance = exp_class(**exp_kwargs)

        experiment_library.hardware_description.hardware.hardware.init()
        exp_instance._init()
        exp_instance._initialize_scan(debug_level=log_level)

        return System().readout.to_dict()

[g-braeunlich]

@aferk , @antti-ethz This would be my proposal
Might be also interesting to @mosmuell @mastadle