Source code for SDF.file_io.force_sdf

import logging

from scipy.io import savemat
from typing import BinaryIO, Union, TextIO

from SDF.data_model import ArrayDataset1D, SDFObject, SourceParameters, Workspace
from SDF.force_sdf import ForceSDF, jpk_sdf_to_force_sdf, mfp_sdf_to_force_sdf, is_mfp_force_curve
from SDF.file_io.sdf import load_from_sdf

logger = logging.getLogger(__name__)


def is_convertible_to_force_sdf(sdf: SDFObject) -> bool:
    if not isinstance(sdf, Workspace):
        return False
    source_type = sdf.parameters["source-of-sdf-file"]["source-file-type"].value
    if source_type == "jpk-force":
        return True
    elif source_type == "mfp-file" and is_mfp_force_curve(sdf):
        return True
    return False


def is_force_sdf(sdf: SDFObject) -> bool:
    if not isinstance(sdf, Workspace):
        return False
    return sdf.name.startswith("ForceSDF")


def load_from_force_sdf(filename: Union[str, TextIO]) -> ForceSDF:
    sdf = load_from_sdf(filename)
    if not isinstance(sdf, Workspace):
        raise TypeError(f"Expected a workspace, got {type(sdf)}")
    if not is_force_sdf(sdf):
        raise ValueError("Given workspace is no ForceSDF")
    return ForceSDF(sdf)


def sdf_to_force_sdf(sdf: SDFObject) -> ForceSDF:
    if not isinstance(sdf, Workspace):
        raise TypeError(f"Expected a workspace, got {type(sdf)}")
    if is_force_sdf(sdf):
        return ForceSDF(sdf)
    if not is_convertible_to_force_sdf(sdf):
        raise ValueError("Given workspace cannot be converted to ForceSDF")

    original_source_type = sdf.parameters["source-of-sdf-file"]["source-file-type"].value
    if original_source_type == "jpk-force":
        logger.info("Determined file to be of type 'jpk-force'")
        curve = jpk_sdf_to_force_sdf(sdf)
        source_type = "converted-jpk-force"
    elif original_source_type == "mfp-file":
        logger.info("Determined file to be of type 'mfp-file'")
        curve = mfp_sdf_to_force_sdf(sdf)
        source_type = "converted-mfp-file"
    else:
        raise ValueError(f"Unknown source type: '{original_source_type}'")

    source_file = sdf.parameters["source-of-sdf-file"]["source-file"].value
    source_par = SourceParameters(converter_name="sdf2forcesdf", source_file=source_file, source_file_type=source_type)
    curve.parameters.add(source_par, as_first=True)

    return curve


def force_sdf_to_mat(sdf: ForceSDF, mat_file: Union[str, BinaryIO]) -> None:
    if not isinstance(sdf, ForceSDF):
        raise TypeError(f"Can only save ForceSDF objects to .mat, got {type(sdf)}")
    spring_constant = sdf.instruments["parameters"]["spring_constant"]
    sensitivity = sdf.instruments["parameters"]["sensitivity"]
    mat = dict(
        spring_constant=spring_constant.value,
        spring_constant_unit=str(spring_constant.unit),
        sensitivity=sensitivity.value,
        sensitivity_unit=str(sensitivity.unit),
    )

    for segment in sdf.workspaces:
        if not segment.name.startswith("segment"):  # skip non-segment workspaces
            continue
        mat_segment = dict()
        for dataset in segment.datasets:
            if not isinstance(dataset, ArrayDataset1D):
                continue
            mat_segment[dataset.name] = dataset.data
            mat_segment[f"{dataset.name}_unit"] = str(dataset.unit)
        for parameter in segment.instruments["segment-parameters"]:
            mat_segment[parameter.name] = parameter.value
            if parameter.unit is not None:
                mat_segment[parameter.name + "_unit"] = parameter.unit
        mat[segment.name.replace(" ", "_")] = mat_segment

    savemat(mat_file, mat)