Source code for ewoksfluo.tasks.positioner_utils

import logging
import numbers
from typing import Dict
from typing import Optional

import h5py
import numpy

from ..io import hdf5

_ENERGY_UNITS = {"kev": 1, "ev": 0.001}
_ENERGY_TEMPLATE = "instrument/positioners_start/{}"

_logger = logging.getLogger(__name__)




def get_energy_suburi(bliss_scan_uri: str) -> Optional[str]:
    return _get_unit_position_suburi(bliss_scan_uri, _ENERGY_UNITS)





def get_energy(
    bliss_scan_uri: str,
    energy_name: Optional[str] = None,
    energy_uri_template: Optional[str] = None,
) -> Optional[float]:
    if energy_name:
        if not energy_uri_template:
            energy_uri_template = _ENERGY_TEMPLATE
        try:
            return _get_template_position_value(
                bliss_scan_uri, energy_name, energy_uri_template, _ENERGY_UNITS
            )
        except KeyError:
            _logger.warning(
                "'%s' does not exist. Do not modify the primary beam energy.",
                energy_name,
            )
            return None
    return _get_unit_position_value(bliss_scan_uri, _ENERGY_UNITS)



def _get_unit_position_suburi(
    bliss_scan_uri: str, units: Dict[str, float]
) -> Optional[str]:
    """Get scan sub-URI for a positioner with specific units"""
    scan_filename, scan_h5path = hdf5.split_h5uri(bliss_scan_uri)
    with hdf5.FileReadAccess(scan_filename) as nxroot:
        positioners = nxroot[f"{scan_h5path}/instrument/positioners_start"]
        name = _get_positioner_name(positioners, units)
        if name is not None:
            return f"instrument/positioners_start/{name}"


def _get_template_position_value(
    bliss_scan_uri: str,
    position_name: str,
    position_uri_template: str,
    units: Dict[str, float],
) -> float:
    """Get position value from scan"""
    scan_filename, scan_h5path = hdf5.split_h5uri(bliss_scan_uri)
    suburi = position_uri_template.format(position_name)
    with hdf5.FileReadAccess(scan_filename) as nxroot:
        data_path = f"{scan_h5path}/{suburi}"
        dset = nxroot[data_path]
        punits = dset.attrs.get("units", "").lower()
        return _asscalar(dset[()]) * units.get(punits, 1)


def _get_unit_position_value(
    bliss_scan_uri: str, units: Dict[str, float]
) -> Optional[float]:
    """Get position value from scan"""
    scan_filename, scan_h5path = hdf5.split_h5uri(bliss_scan_uri)
    with hdf5.FileReadAccess(scan_filename) as nxroot:
        positioners = nxroot[f"{scan_h5path}/instrument/positioners_start"]
        return _get_positioner_value(positioners, units)


def _get_positioner_value(
    positioners: hdf5.GroupType, units: Dict[str, float]
) -> Optional[h5py.Dataset]:
    for name in positioners:
        positioner = positioners[name]
        punits = positioner.attrs.get("units", "").lower()
        if punits in units:
            return _asscalar(positioner[()]) * units[punits]


def _asscalar(value: numpy.ndarray, rtol=1.0e-5, atol=1.0e-8) -> numbers.Number:
    if numpy.isscalar(value):
        return value

    if not isinstance(value, numpy.ndarray):
        raise TypeError(f"Unsupported type: {type(value)}")

    if value.size == 0:
        raise ValueError("Cannot convert empty array to scalar")
    if value.size == 1:
        return value.item()
    ref_value = numpy.nanmedian(value)
    if numpy.allclose(value, ref_value, atol=atol, rtol=rtol):
        return ref_value
    raise ValueError(f"Array values differ beyond tolerance (shape {value.shape})")


def _get_positioner_name(
    positioners: hdf5.GroupType, units: Dict[str, float]
) -> Optional[str]:
    for name in positioners:
        positioner = positioners[name]
        punits = positioner.attrs.get("units", "").lower()
        if punits in units:
            return name