Source code for est.io.utils.readers.larch_reader

import logging
import functools
from typing import Optional, List, Tuple, Dict

import numpy

try:
    from larch.io import columnfile
except ImportError:
    columnfile = None

from est.units import ur
from .parse import parse_energy_mu


_logger = logging.getLogger(__name__)




def get_first_scan_title(file_path: str) -> Optional[str]:
    return None





def get_all_scan_titles(file_path: str) -> List[str]:
    return list()





def get_scan_column_names(file_path: str, scan_title: str) -> List[str]:
    return list(_read_columnfile(file_path))





@functools.lru_cache(maxsize=2)  # called twice for energy and absorption
def read_spectrum(
    column_file,
    energy_col_name=None,
    absorption_col_name=None,
    monitor_col_name=None,
    energy_unit=ur.eV,
    scan_title=None,
) -> Tuple[Optional[numpy.ndarray], Optional[numpy.ndarray]]:
    data = _read_columnfile(column_file)

    if energy_col_name is None:
        _logger.warning("Larch energy column name not provided. Try 'energy'")
        energy_col_name = "energy"
    if absorption_col_name is None:
        _logger.warning("Larch absorption column name not provided. Try 'mu'")
        absorption_col_name = "mu"
    if monitor_col_name is None:
        _logger.warning("Larch monitor column name not provided. Try 'i0'")
        monitor_col_name = "i0"

    energy = data.get(energy_col_name)
    mu = data.get(absorption_col_name)
    monitor = data.get(monitor_col_name)

    return parse_energy_mu(energy, mu, monitor, energy_unit)



def _read_columnfile(
    column_file,
) -> Dict[str, numpy.ndarray]:
    larch_group = columnfile.read_ascii(column_file)
    return dict(zip(larch_group.array_labels, larch_group.data))