Source code for torx.fileio.matlab_io_m

"""Functions for reading MATLAB data '.mat' files."""
import numpy as np
from scipy.io import loadmat, savemat
from pathlib import Path
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def read_matlab_struct(filepath: Path, struct_name=None) -> dict:
    """Read a struct from a MATLAB file."""
    filepath = Path(filepath)
    assert filepath.exists(), f"File not found {filepath.absolute()}"

    mat_file = loadmat(filepath, simplify_cells=True)

    for key in ["__header__", "__version__", "__globals__"]:
        del mat_file[key]

    if struct_name is not None:
        data = mat_file[struct_name]
    else:
        data = mat_file

    return data




@autodoc_function
def write_matlab_struct(data, filepath, struct_name=None):
    """Save a data dictionary as a MATLAB struct to file."""
    filepath = Path(filepath)
    assert (
        not filepath.exists()
    ), f"File exists at {filepath.absolute()}. Unlink before writing."
    assert (
        filepath.suffix == ".mat"
    ), f"Should use a .mat suffix to indicate MATLAB files"

    if struct_name is not None:
        savemat(file_name=filepath, mdict={struct_name: data})
    else:
        savemat(file_name=filepath, mdict=data)