Source code for SDF.file_io.jpk

import logging
import os
from typing import Union

from SDF.config import CONFIG
from SDF.data_model import ArrayData1D, ArrayDataset1D, Instrument, Parameter, ParameterSet, SourceParameters, Workspace

if CONFIG.use_system_jpkfile:
    import jpkfile
else:
    from SDF.extern import jpkfile


logger = logging.getLogger(__name__)


def load_from_jpk(filename: str) -> Workspace:
    logger.info(f"Reading SDF workspace from file '{filename}'")
    extension = os.path.splitext(filename)[1]

    if extension.endswith("map"):
        logger.info("Treating file as map")
        source_type = "jpk-force-map"
        workspace = load_from_jpk_map(jpkfile.JPKMap(filename), label=f"JPKMap {os.path.split(filename)[1]}")
    else:
        logger.info("Treating file as single atomic force or tweezer recording")
        if "nt-force" in extension:
            source_type = "jpk-nt-force"
        elif "proc-force" in extension:
            source_type = "jpk-proc-force"
        else:
            source_type = "jpk-force"
        workspace = load_from_jpk_object(jpkfile.JPKFile(filename), label=f"JPKFile {os.path.split(filename)[1]}")

    source_par = SourceParameters(converter_name="jpk2sdf", source_file=filename, source_file_type=source_type)
    workspace.parameters.add(source_par, as_first=True)

    return workspace


def extract_original_parameters(jpk: Union[jpkfile.JPKFile, jpkfile.JPKMap, jpkfile.JPKSegment],
                                shared_params_already_stored_in_parent: bool = False) -> Instrument:
    original_params = Instrument(name="original-parameters")

    # parameters added to instrument one-by-one without a named container
    original_params.add_from_structure(jpk.parameters)

    if not isinstance(jpk, jpkfile.JPKSegment):
        # parameters from shared header
        if jpk.has_shared_header and not shared_params_already_stored_in_parent:
            shared_params = ParameterSet(name="shared-parameters")
            shared_params.add_from_structure(jpk.shared_parameters)
            original_params.add(shared_params)

    return original_params


def load_from_jpk_map(jpk: jpkfile.JPKMap, label: str) -> Workspace:
    workspace = Workspace(name=label)
    workspace.instruments.add(extract_original_parameters(jpk, shared_params_already_stored_in_parent=False))

    for i, pixel_jpk in jpk.flat_indices.items():
        workspace.workspaces.add(load_from_jpk_object(pixel_jpk, f"pixel {i} in flattened map", jpk.has_shared_header))

    return workspace


def load_from_jpk_segment(jpk: jpkfile.JPKSegment, label: str) -> Workspace:
    logger.debug(f"Loading from segment {label}")
    workspace = Workspace(name=label)
    workspace.instruments.add(extract_original_parameters(jpk, shared_params_already_stored_in_parent=False))

    for channel_name, (data, conversion_dict) in jpk.data.items():
        if channel_name == "t":
            continue

        logger.debug(f"Adding dataset {channel_name}")
        if data.ndim == 2 and 1 in data.shape:
            data = data.flatten()
        workspace.datasets.add(ArrayDataset1D(name=channel_name, data=ArrayData1D(data, try_hex_transformation=True),
                                              parameters=(Parameter("converted", "False"),)))

    return workspace


def load_from_jpk_object(jpk: jpkfile.JPKFile, label: str,
                         shared_params_already_stored_in_parent: bool = False) -> Workspace:
    workspace = Workspace(name=label)
    workspace.instruments.add(extract_original_parameters(jpk, shared_params_already_stored_in_parent))

    for i, segment_jpk in jpk.segments.items():
        workspace.workspaces.add(load_from_jpk_segment(segment_jpk, label=f"segment {i}"))

    return workspace