Acquisition notebook manager.

AcquisitionAnalysisManager.

Init

aqm = AcquisitionAnalysisManager("tmp_data/", use_magic=False, save_files=False)
aqm.set_config_file("configuration.py")

acquisition_cell:

aqm.acquisition_cell('simple_sine')
...
aqm.save_acquisition(x=x, y=y)

analysis_cell:

aqm.analysis_cell()
...
plt.plot(aqm.d.x, aqm.d.y)
aqm.save_fig()

Source code in labmate/acquisition_notebook/acquisition_analysis_manager.py

class AcquisitionAnalysisManager(AcquisitionManager):
    """AcquisitionAnalysisManager.

    # Init
    ```
    aqm = AcquisitionAnalysisManager("tmp_data/", use_magic=False, save_files=False)
    aqm.set_config_file("configuration.py")
    ```
    # acquisition_cell:
    ```
    aqm.acquisition_cell('simple_sine')
    ...
    aqm.save_acquisition(x=x, y=y)
    ```

    # analysis_cell:
    ```
    aqm.analysis_cell()
    ...
    plt.plot(aqm.d.x, aqm.d.y)
    aqm.save_fig()
    ```
    """

    _analysis_data: Optional[AnalysisData] = None
    _analysis_cell_str = None
    _is_old_data = False
    _last_fig_name = None
    _default_config_files: Tuple[str, ...] = ()
    _acquisition_started = 0
    _linting_external_vars = None
    _connected_widgets: Optional[List["display_widget.WidgetProtocol"]] = None

    def __init__(
        self,
        data_directory: Optional[Union[str, Any]] = None,
        *,
        config_files: Optional[List[str]] = None,
        save_files: bool = False,
        use_magic: bool = False,
        save_on_edit: bool = True,
        save_on_edit_analysis: Optional[bool] = None,
        save_fig_inside_h5: bool = False,
        shell: Any = True,
        hooks: Optional[LifecycleHooks] = None,
    ):
        """
        AcquisitionAnalysisManager.

        Args:
            data_directory:
                Path to data_directory. Should be explicitly set here or as environ parameter.
            config_files:
                List of paths to config files. Defaults to empty.
            save_files:
                True to additionally save config files and the cells to files. Defaults to False.
                So all information is saved inside the h5 file.
            use_magic:
                True to register the magic cells. Defaults to False.
            save_on_edit:
                True to save data for every change.
            save_on_edit_analysis:
                save_on_edit parameter for AnalysisManager i.e. data inside analysis_cell
            shell:
                could be provided or explicitly set to None. Defaults to get_ipython().
            hooks:
                Optional shared :class:`~labmate.acquisition.hooks.LifecycleHooks`.
                If omitted, a new instance is created on the base manager.
        """
        if shell is False or shell is True:  # behavior by default shell
            try:
                from IPython.core.getipython import get_ipython

                self.shell = get_ipython()
            except ImportError:
                self.shell = None
        else:  # if any shell is provided. even None
            self.shell = shell

        if use_magic:
            from .acquisition_magic_class import (  # noqa: I001
                load_ipython_extension,
            )

            load_ipython_extension(aqm=self, shell=self.shell)

        if save_on_edit_analysis is None:
            save_on_edit_analysis = save_on_edit

        self._save_on_edit_analysis = save_on_edit_analysis
        self._save_fig_inside_h5 = save_fig_inside_h5

        self._logger = logger
        super().__init__(
            data_directory=str(data_directory),
            config_files=config_files,
            save_files=save_files,
            save_on_edit=save_on_edit,
            hooks=hooks,
        )

    @property
    def logger(self):
        return self._logger

    @property
    def current_acquisition(self):
        """Return current acquisition if it's not an old data analyses."""
        if self._is_old_data:
            return None
        return super().current_acquisition

    @property
    def current_analysis(self):
        """Return the current analysis. Class where you cannot change existed keys."""
        return self._analysis_data

    @property
    def data(self):
        """Same as current_analysis."""
        if self._analysis_data is None:
            raise ValueError("No data set")
        return self._analysis_data

    @property
    def d(self):  # pylint: disable=invalid-name
        """Shorter alias for data."""
        return self.data

    def save_fig_only(
        self,
        fig: Optional["FigureProtocol"] = None,
        name: Optional[Union[str, int]] = None,
        extensions: Optional[str] = None,
        **kwds,
    ) -> "AcquisitionAnalysisManager":
        """Save the figure as a file."""
        self.data.save_fig(fig=fig, name=name, extensions=extensions, **kwds)
        return self

    def save_analysis_cell(
        self,
        name: Optional[Union[str, int]] = None,
        cell: Optional[Union[str, Literal["none"]]] = None,
    ) -> "AcquisitionAnalysisManager":
        if name is None:
            name = self.data.figure_last_name

        if name is not None:
            name = str(name)

        cell = cell or self._analysis_cell_str
        self.data.save_analysis_cell(code=cell, code_name=name)
        return self

    def save_fig(
        self,
        fig_or_name: Optional[Union["FigureProtocol", str, int]] = None,
        /,
        *,
        fig: Optional["FigureProtocol"] = None,
        name: Optional[Union[str, int]] = None,
        cell: Optional[str] = None,
        **kwds,
    ) -> "AcquisitionAnalysisManager":
        if fig_or_name is not None:
            if isinstance(fig_or_name, (str, int)):
                name = name or fig_or_name
            else:
                fig = fig or fig_or_name

        self.save_fig_only(fig=fig, name=name, **kwds)
        self.save_analysis_cell(name=name, cell=cell)

        if self.current_acquisition is not None:
            self.hooks.dispatch_figure_saved(self.current_acquisition)

        if self._connected_widgets:
            display_widget.display_widgets(
                self._connected_widgets,
                aqm=self,
                fig=fig,
            )
        return self

    def __setitem__(self, __key: str, __value: Any) -> None:
        if self._analysis_data is not None:
            raise ValueError(
                "This is the way to save acquisition data. But analysis data was loaded. "
                "So you possibly run it outside of acquisition_cell"
            )
        acq_data = self.current_acquisition
        if acq_data is None:
            raise ValueError(
                "Cannot save data to acquisition as current acquisition is None."
                "Possibly because you have never run `acquisition_cell(..)` or it's an old data"
            )
        acq_data[__key] = __value  # pylint: disable=E1137

    def save_acquisition(
        self,
        update_: bool = True,
        /,
        file_suffix: Optional[str] = None,
        **kwds,
    ) -> "AcquisitionAnalysisManager":
        acquisition_finished = time.time()
        if not self._once_saved:
            additional_info: Dict[str, Any] = {
                "acquisition_duration": acquisition_finished - self._acquisition_started,
                "logs": self.logger.getvalue(),
                "prints": self.logger.get_stdout(),
            }
            if self._default_config_files:
                additional_info.update({"default_config_files": self._default_config_files})
            kwds.update({"info": additional_info})

        super().save_acquisition(update_, file_suffix=file_suffix, **kwds)
        self._load_analysis_data()
        return self

    def _load_analysis_data(self, filepath: Optional[str] = None):
        filepath = filepath or str(self.current_filepath)

        self._analysis_data = self.load_file(filepath)

        if self._save_on_edit_analysis is False:
            self._analysis_data.save()

        return self._analysis_data

    def load_file(self, filename) -> "AnalysisData":
        filename = self._get_full_filename(filename)
        if not os.path.exists(filename if filename.endswith(".h5") else filename + ".h5"):  # noqa: PTH110
            raise ValueError(f"File {filename} cannot be found")

        data = AnalysisData(
            filepath=filename,
            save_files=self._save_files,
            save_on_edit=self._save_on_edit_analysis,
            save_fig_inside_h5=self._save_fig_inside_h5,
            open_on_init=False,
        )

        if not data.get("useful", True):
            data.unlock_data("useful").update(**{"useful": True}).lock_data("useful")

        if self._default_config_files:
            data.set_default_config_files(self._default_config_files)

        return data

    def acquisition_cell(
        self,
        name: str,
        cell: Optional[str] = None,
        prerun: Optional[Union[_CallableWithNoArgs, List[_CallableWithNoArgs]]] = None,
        save_on_edit: Optional[bool] = None,
        step: int = 1,
    ) -> "AcquisitionAnalysisManager":
        self._analysis_cell_str = None
        self._analysis_data = None
        self._is_old_data = False
        self._acquisition_started = time.time()

        cell = cell or get_current_cell(self.shell)
        if step == 1:
            self.logger.reset()
            self.new_acquisition(name=name, cell=cell, save_on_edit=save_on_edit)
        elif self._current_acquisition is None:
            raise ValueError("Acquisition should start from step 1")
        elif self._current_acquisition.current_step == step:
            raise ValueError(
                "This step was already run. Please run the next step or restart from step 1"
            )
        else:
            if self._current_acquisition.experiment_name != name:
                raise ValueError(
                    f"Current acquisition ('{self.current_experiment_name}') "
                    f"isn't the one expected ('{name}') for this acquisition. "
                    f"Possible solutions: run acquisition '{name}' with step 1; "
                    f"or change current acquisition name to '{self.current_experiment_name}'"
                )
            self._current_acquisition.current_step = step
            self._current_acquisition.set_cell(cell, step=step)
            self._current_acquisition.save_cell(cell, suffix=str(step))
            configs_modified = self._get_configs_last_modified()
            if configs_modified != self._configs_last_modified:
                raise ValueError(
                    "Config files were modified since the previous acquisition step. "
                    "Please rerun the acquisition from the first step."
                )
            self.logger.stdout_flush()

        self.logger.info(f"{step}:{self.current_filepath.basename}")  # pylint: disable=W1203

        if step == 1:
            self.hooks.dispatch_acquisition_cell_ready()

        utils.run_functions(prerun)
        return self

    def analysis_cell(
        self,
        filename: Optional[Union[str, "Path"]] = None,
        *,
        acquisition_name=None,
        cell: Optional[str] = None,
        filepath: Optional[Union[str, "Path"]] = None,
        prerun: Optional[Union[_CallableWithNoArgs, List[_CallableWithNoArgs]]] = None,
    ) -> "AcquisitionAnalysisManager":
        self._analysis_cell_str = cell or get_current_cell(self.shell)
        if filename or filepath:  # getting old data
            self._is_old_data = True
            if self.shell is not None:
                from labmate.display.html_output import display_warning

                display_warning("Old data analysis")

            filename = str(filepath or self._get_full_filename(filename))  # type: ignore[arg-type]
            filename = (filename.rsplit(".h5", 1)[0]) if filename.endswith(".h5") else filename
        else:
            self._is_old_data = False
            if acquisition_name is not None:
                import re

                if (
                    len(acquisition_name) == 0
                    or (
                        acquisition_name[0] != r"^"
                        and acquisition_name != self.current_experiment_name
                    )
                    or (
                        acquisition_name[0] == r"^"
                        and re.match(acquisition_name, self.current_experiment_name) is None
                    )
                ):
                    raise ValueError(
                        f"Current acquisition ('{self.current_experiment_name}') "
                        f"isn't the one expected ('{acquisition_name}') for this analysis"
                    )
            filename = str(self.current_filepath)  # without h5

        self.logger.info(os.path.basename(filename))  # noqa: PTH119

        if (
            (not self._is_old_data)
            and (self.shell is not None)
            and (
                "acquisition_cell(" in self.shell.last_execution_result.info.raw_cell  # type: ignore
                and not self.shell.last_execution_result.success  # type: ignore
            )
        ):
            raise ChildProcessError(
                "Last executed cell was probably an `acquisition_cell` and failed to run. "
                "Check if everything is ok and executive again"
            )

        full_h5 = filename + ".h5"

        self.hooks.dispatch_analysis_data_loading(full_h5)

        if os.path.exists(full_h5):  # noqa: PTH110
            self._load_analysis_data(filename)
        else:
            if self._is_old_data:
                raise ValueError(f"Cannot load data from {filename}")
            self._analysis_data = None

        if cell is not None:
            self.save_analysis_cell(cell=cell)

        if (self._analysis_cell_str is not None) and (self._linting_external_vars is not None):
            from ..acquisition import custom_lint
            from ..utils import lint

            lint_result = lint.find_variables_from_code(
                self._analysis_cell_str,
                self._linting_external_vars,
                run_on_call=custom_lint.on_call_functions,
            )
            for var in lint_result.external_vars:
                self.logger.warning("External variable used inside the analysis code: %s", var)
            for error in lint_result.errors:
                self.logger.warning(error)

        self.hooks.dispatch_analysis_cell_ready()
        utils.run_functions(prerun)
        return self

    def get_analysis_code(self, look_inside: bool = True) -> str:
        code = self.data.get_analysis_code(update_code=look_inside)

        if self.shell is not None:
            self.shell.set_next_input(code)  # type: ignore
        return code

    def _get_full_filename(self, filename: Union[str, "Path"]) -> str:
        if filename is None:
            raise ValueError("Filename cannot be None")

        filepath = utils.get_path_from_filename(filename)
        if isinstance(filepath, tuple):
            return os.path.join(self.data_directory, *filepath)  # noqa: PTH118
        return filepath

    def parse_config_file(self, config_file_name: str, /) -> "ConfigFile":
        return self.data.parse_config_file(config_file_name)

    def parse_config(self, config_files: Optional[Tuple[str, ...]] = None) -> "ConfigFile":
        return self.data.parse_config(config_files=config_files)

    @property
    def cfg(self) -> "ConfigFile":
        return self.data.cfg

    def parse_config_str(self, values: List[str], /, max_length: Optional[int] = None) -> str:
        return self.data.parse_config_str(values, max_length=max_length)

    def linting(
        self,
        allowed_variables: Optional[Iterable[str]] = None,
        init_file: Optional[str] = None,
    ):
        from ..utils import lint

        allowed = set() if allowed_variables is None else set(allowed_variables)
        if init_file is not None:
            allowed.update(lint.find_variables_from_file(init_file)[0])
        self._linting_external_vars = allowed

    def set_default_config_files(self, config_files: Union[str, Tuple[str, ...], List[str]], /):
        self._default_config_files = (
            (config_files,) if isinstance(config_files, str) else tuple(config_files)
        )
        if self._analysis_data:
            self._analysis_data.set_default_config_files(self._default_config_files)

    def set_analysis_cell_prerun_hook(
        self,
        hook: Union[
            _CallableWithNoArgs,
            List[_CallableWithNoArgs],
            Tuple[_CallableWithNoArgs, ...],
        ],
    ):
        self.hooks.set_analysis_cell_ready(hook)

    def set_acquisition_cell_prerun_hook(
        self,
        hook: Union[
            _CallableWithNoArgs,
            List[_CallableWithNoArgs],
            Tuple[_CallableWithNoArgs, ...],
        ],
    ):
        self.hooks.set_acquisition_cell_ready(hook)

    def find_param_in_config(self, param: str) -> Optional[Tuple[str, int]]:
        for file in self._default_config_files:
            for line_no, line in enumerate(self.d["configs", file].split("\n")):
                if line.startswith(param):
                    return file, line_no + 1
        return None

    def display_param_link(
        self,
        params: Union[str, List[str], List[Tuple[str, str]]],
        after_text: Optional[str] = None,
        title: Optional[str] = None,
    ):
        if after_text is not None:
            if not isinstance(params, str):
                raise ValueError(
                    "Cannot use after_text with multiple params. "
                    "Use params=[(param, after_text), ...] instead."
                )
            return self.display_param_link(params=[(params, after_text)], title=title)

        if isinstance(params, str):
            params = [params]

        links = "" if not title else title + "<br/>"
        for param in params:
            if not isinstance(param, str):
                param_text, after_text = param
            else:
                param_text, after_text = param, None

            res = self.find_param_in_config(param_text)
            if res is None:
                self.logger.warning(
                    "Parameter '%s' cannot be found in default config files.", param
                )
                continue

            file, line_no = res
            file = self._config_files_names_to_path.get(file, file)
            link = display.links.create_link(param_text, file, line_no, after_text)
            links += link + "<br/>"
        return display.display_html(links)

    def display_cfg_link(self, parameters: Dict[str, Any], update_button: bool = False):
        from labmate.display import html_output

        links = []
        for param, value in parameters.items():
            param_eq = f"{param.strip()} = "
            res = self.find_param_in_config(param_eq)
            if res is None:
                self.logger.warning(
                    "Parameter '%s' cannot be found in default config files.", param
                )
                continue

            file, line_no = res
            file = self._config_files_names_to_path.get(file, file)

            def update_value(param, value):
                self.update_config_params_on_disk({param: value})

            buttons = (
                [display.buttons.create_button(update_value, param, value, name="Update")]
                if update_button
                else None
            )

            link = html_output.create_link_row(
                link_text=f"{param} = ",
                link_url=f"{file}:{line_no}",
                text=str(value),
                buttons=buttons,  # type: ignore
            )
            links.append(link)
        return display.display_widgets_vertically(links, class_="labmate-params")

    def update_config_params_on_disk(self, params: Dict[str, Any]):
        for file in self.config_files:
            file = self._config_files_names_to_path.get(file, file)
            utils.file_read.update_file_variable(file, params)
        return self

    def connect_default_widget(
        self,
        objs: Union["display_widget.WidgetProtocol", List["display_widget.WidgetProtocol"]],
    ):
        if not isinstance(objs, (list, tuple)):
            objs = [objs]
        if self._connected_widgets is None:
            self._connected_widgets = []
        self._connected_widgets.extend(objs)

acquisition_tmp_data property writable

acquisition_tmp_data

Return information about the current acquisition.

Returns class attribute or read it from temp.json file if first is not set. Returns: AcquisitionTmpData(NamedTuple): experiment_name: current experiment name time_stamp: time stamp of the current acquisition configs: dict of configurations files to save directory: directory where the data is stored

current_acquisition property

current_acquisition

Return current acquisition if it's not an old data analyses.

current_analysis property

current_analysis

Return the current analysis. Class where you cannot change existed keys.

d property

d

Shorter alias for data.

data property

data

Same as current_analysis.

data_directory property writable

data_directory

Return the path to the directory where the data is stored.

Returns:

Name	Type	Description
Path	Path	The path to the directory where the data is stored.

__init__

__init__(data_directory=None, *, config_files=None, save_files=False, use_magic=False, save_on_edit=True, save_on_edit_analysis=None, save_fig_inside_h5=False, shell=True, hooks=None)

AcquisitionAnalysisManager.

Parameters:

Name	Type	Description	Default
data_directory	Optional[Union[str, Any]]	Path to data_directory. Should be explicitly set here or as environ parameter.	None
config_files	Optional[List[str]]	List of paths to config files. Defaults to empty.	None
save_files	bool	True to additionally save config files and the cells to files. Defaults to False. So all information is saved inside the h5 file.	False
use_magic	bool	True to register the magic cells. Defaults to False.	False
save_on_edit	bool	True to save data for every change.	True
save_on_edit_analysis	Optional[bool]	save_on_edit parameter for AnalysisManager i.e. data inside analysis_cell	None
shell	Any	could be provided or explicitly set to None. Defaults to get_ipython().	True
hooks	Optional[LifecycleHooks]	Optional shared :class:~labmate.acquisition.hooks.LifecycleHooks. If omitted, a new instance is created on the base manager.	None

Source code in labmate/acquisition_notebook/acquisition_analysis_manager.py

def __init__(
    self,
    data_directory: Optional[Union[str, Any]] = None,
    *,
    config_files: Optional[List[str]] = None,
    save_files: bool = False,
    use_magic: bool = False,
    save_on_edit: bool = True,
    save_on_edit_analysis: Optional[bool] = None,
    save_fig_inside_h5: bool = False,
    shell: Any = True,
    hooks: Optional[LifecycleHooks] = None,
):
    """
    AcquisitionAnalysisManager.

    Args:
        data_directory:
            Path to data_directory. Should be explicitly set here or as environ parameter.
        config_files:
            List of paths to config files. Defaults to empty.
        save_files:
            True to additionally save config files and the cells to files. Defaults to False.
            So all information is saved inside the h5 file.
        use_magic:
            True to register the magic cells. Defaults to False.
        save_on_edit:
            True to save data for every change.
        save_on_edit_analysis:
            save_on_edit parameter for AnalysisManager i.e. data inside analysis_cell
        shell:
            could be provided or explicitly set to None. Defaults to get_ipython().
        hooks:
            Optional shared :class:`~labmate.acquisition.hooks.LifecycleHooks`.
            If omitted, a new instance is created on the base manager.
    """
    if shell is False or shell is True:  # behavior by default shell
        try:
            from IPython.core.getipython import get_ipython

            self.shell = get_ipython()
        except ImportError:
            self.shell = None
    else:  # if any shell is provided. even None
        self.shell = shell

    if use_magic:
        from .acquisition_magic_class import (  # noqa: I001
            load_ipython_extension,
        )

        load_ipython_extension(aqm=self, shell=self.shell)

    if save_on_edit_analysis is None:
        save_on_edit_analysis = save_on_edit

    self._save_on_edit_analysis = save_on_edit_analysis
    self._save_fig_inside_h5 = save_fig_inside_h5

    self._logger = logger
    super().__init__(
        data_directory=str(data_directory),
        config_files=config_files,
        save_files=save_files,
        save_on_edit=save_on_edit,
        hooks=hooks,
    )

create_acquisition

create_acquisition(name=None, cell=None, save_on_edit=None)

Create a new acquisition with the given experiment name.

Source code in labmate/acquisition/acquisition_manager.py

def create_acquisition(
    self,
    name: Optional[str] = None,
    cell: Optional[str] = None,
    save_on_edit: Optional[bool] = None,
) -> NotebookAcquisitionData:
    """Create a new acquisition with the given experiment name."""
    configs = read_files(self.config_files)

    if self.config_files_eval:
        configs = append_values_from_modules_to_files(configs, self.config_files_eval)

    if name is None:
        name = self.current_experiment_name + "_item"

    dic = AcquisitionTmpData(
        experiment_name=name,
        time_stamp=get_timestamp(),
        configs=configs,
        directory=self.data_directory,
    )

    filepath = self.create_path_from_tmp_data(dic)
    configs = configs if configs else None
    save_on_edit = save_on_edit if save_on_edit is not None else self._save_on_edit

    acquisition = NotebookAcquisitionData(
        filepath=str(filepath),
        configs=configs,
        cell=cell or self.cell,
        overwrite=False,
        save_on_edit=save_on_edit,
        save_files=self._save_files,
    )
    self.hooks.dispatch_acquisition_data_loaded(acquisition)

    return acquisition

new_acquisition

new_acquisition(name, cell=None, save_on_edit=None)

Create a new acquisition with the given experiment name.

Source code in labmate/acquisition/acquisition_manager.py

def new_acquisition(
    self, name: str, cell: Optional[str] = None, save_on_edit: Optional[bool] = None
) -> NotebookAcquisitionData:
    """Create a new acquisition with the given experiment name."""
    self._current_acquisition = None
    self._once_saved = False
    self.cell = cell
    configs = read_files(self.config_files)
    self._configs_last_modified = self._get_configs_last_modified()

    if self.config_files_eval:
        configs = append_values_from_modules_to_files(configs, self.config_files_eval)

    dic = AcquisitionTmpData(
        experiment_name=name,
        time_stamp=get_timestamp(),
        configs=configs,
        directory=self.data_directory,
    )

    self.acquisition_tmp_data = dic

    self._current_acquisition = self.get_acquisition(replace=True, save_on_edit=save_on_edit)

    return self.current_acquisition

save_fig_only

save_fig_only(fig=None, name=None, extensions=None, **kwds)

Save the figure as a file.

Source code in labmate/acquisition_notebook/acquisition_analysis_manager.py

def save_fig_only(
    self,
    fig: Optional["FigureProtocol"] = None,
    name: Optional[Union[str, int]] = None,
    extensions: Optional[str] = None,
    **kwds,
) -> "AcquisitionAnalysisManager":
    """Save the figure as a file."""
    self.data.save_fig(fig=fig, name=name, extensions=extensions, **kwds)
    return self

set_config_file

set_config_file(filename)

Set self.config_file to filename. Verify if exists. Only set config file for future acquisition and will not change current acquisition

Source code in labmate/acquisition/acquisition_manager.py

def set_config_file(
    self, filename: Union[str, List[str], Tuple[str, ...]]
) -> "AcquisitionManager":
    """Set self.config_file to filename. Verify if exists.
    Only set config file for future acquisition and will not change current acquisition
    """
    if not isinstance(filename, (list, set, tuple)):
        filename = [str(filename)]

    self.config_files = list(filename)
    self._config_files_names_to_path = {Path(file).name: file for file in self.config_files}

    for config_file in self.config_files:
        if not Path(config_file).exists():
            raise ValueError(f"Configuration file at {config_file} does not exist")

    return self