Source code for quantify_core.data.experiment

# Repository: https://gitlab.com/quantify-os/quantify-core
# Licensed according to the LICENCE file on the main branch
"""Utilities for managing experiment data."""

import os
from pathlib import Path
from typing import Dict, Any, Optional

import xarray as xr
from quantify_core.data.handling import (
    locate_experiment_container,
    load_dataset,
    DATASET_NAME,
    write_dataset,
    create_exp_folder,
)
from quantify_core.data.handling import snapshot as create_snapshot
from quantify_core.data.types import TUID
from quantify_core.utilities.general import save_json, load_json


[docs]class QuantifyExperiment: """ Class which represents all data related to an experiment. This allows the user to run experiments and store data without the `quantify_core.measurement.control.MeasurementControl`. The class serves as an initial interface for other data storage backends. """
[docs] def __init__(self, tuid: Optional[str], dataset=None): """ Creates an instance of the QuantifyExperiment. Parameters ---------- tuid TUID to use dataset If the TUID is None, use the TUID from this dataset """ if tuid is None: self.tuid = dataset.tuid self.dataset = dataset else: self.tuid = tuid self.dataset = None self.tuid = TUID(tuid)
def __repr__(self) -> str: classname = ".".join([self.__module__, self.__class__.__qualname__]) idx = "%x" % id(self) return f"<{classname} at %x{idx}>: TUID {self.tuid}" @property def experiment_directory(self) -> Path: """ Returns a path to the experiment directory containing the TUID set within the class. Returns ------- : """ experiment_directory = locate_experiment_container(tuid=self.tuid) return Path(experiment_directory) def _get_or_create_experiment_directory(self) -> Path: """ Create the experiment directory containing the TUID set within the class, if it does not exist already. To be used by methods that write/save. The experiment directory will be created on the first write/save, not before. Methods that load should not create an experiment directory. Returns ------- : The path to the experiment directory. """ try: experiment_directory = self.experiment_directory except FileNotFoundError: experiment_directory = create_exp_folder(tuid=self.tuid) return Path(experiment_directory)
[docs] def load_dataset(self) -> xr.Dataset: """ Loads the quantify dataset associated with the TUID set within the class. Returns ------- : """ self.dataset = load_dataset(self.tuid) return self.dataset
[docs] def write_dataset(self, dataset: xr.Dataset): """ Writes the quantify dataset to the directory specified by `~.experiment_directory`. Parameters ---------- dataset The dataset to be written to the directory """ path = self._get_or_create_experiment_directory() / DATASET_NAME write_dataset(path, dataset)
[docs] def load_snapshot(self) -> Dict[str, Any]: """ Loads the snapshot from the directory specified by `~.experiment_directory`. Returns ------- : The loaded snapshot from disk """ return load_json(full_path=self.experiment_directory / "snapshot.json")
[docs] def save_snapshot(self, snapshot: Optional[Dict[str, Any]] = None): """ Writes the snapshot to disk as specified by `~.experiment_directory`. Parameters ---------- snapshot The snapshot to be written to the directory """ if snapshot is None: snapshot = create_snapshot() save_json( directory=self._get_or_create_experiment_directory(), filename="snapshot.json", data=snapshot, )
[docs] def load_text(self, rel_path: str) -> str: """ Loads a string from a text file from the path specified by `~.experiment_directory` / rel_path. Parameters ---------- rel_path path relative to the base directory of the experiment, e.g. "data.json" or "my_folder/data.txt" Returns ------- : The loaded text from disk """ file_path = self.experiment_directory / rel_path text = file_path.read_text(encoding="utf-8") return text
[docs] def save_text(self, text: str, rel_path: str) -> None: """ Saves a string to a text file in the path specified by `~.experiment_directory` / rel_path. Parameters ---------- text text to be saved rel_path path relative to the base directory of the experiment, e.g. "data.json" or "my_folder/data.txt" """ directory = (self._get_or_create_experiment_directory() / rel_path).parent os.makedirs(directory, exist_ok=True) file_path = self.experiment_directory / rel_path file_path.write_text(text, encoding="utf-8")