Source code for ewoksxas.converters.larch

from typing import Any

import numpy as np
import numpy.typing as npt
from larch import Group




def create_groups(
    x: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
    mode: str = "energy",
    **kwargs,
) -> list[Group]:
    """Create one Larch group per spectrum from paired x and y arrays.

    Three input shapes are accepted:

    1. Both x and y are 1D with the same length: produces a single group.
    2. When x is 1D and y is 2D: one group is created per row of y, all sharing
       the same x.
    3. Both x and y are 2D with the same shape: rows are paired positionally,
       producing one group per row.

    The group attribute names are determined by ``mode``.

    Args:
        x: 1D array shared across all groups, or 2D array with one row per group.
        y: 1D array for a single group, or 2D array with one row per group.
        mode: Selects the x/y attribute names written to each group.
        kwargs: Additional attributes passed unchanged to every group.

    Raises:
        TypeError: If x or y is not a NumPy array.
        ValueError: If mode is invalid or the shapes of x and y are incompatible.
    """
    attribute_names = {
        "energy": ("energy", "mu"),
        "k": ("k", "chi"),
        "r": ("r", "chir"),
    }

    if mode not in attribute_names:
        raise ValueError(
            f"The mode should be one of {list(attribute_names)}, but got {mode}"
        )

    if not isinstance(x, np.ndarray):
        raise TypeError(f"Expected NumPy array for x, but got {type(x)}.")
    if not isinstance(y, np.ndarray):
        raise TypeError(f"Expected NumPy array for y, but got {type(y)}.")

    if mode == "energy":
        x = convert_kev_to_ev(x)

    x_name, y_name = attribute_names[mode]

    if x.ndim == 1 and y.ndim == 1 and x.shape == y.shape:
        return [Group(**{x_name: x, y_name: y}, **kwargs)]
    if x.ndim == 1 and y.ndim == 2 and x.shape[0] == y.shape[1]:
        return [Group(**{x_name: x, y_name: yi}, **kwargs) for yi in y]
    if x.ndim == 2 and y.ndim == 2 and x.shape == y.shape:
        return [
            Group(**{x_name: xi, y_name: yi}, **kwargs)
            for xi, yi in zip(x, y, strict=True)
        ]
    raise ValueError(
        f"The x values shape {x.shape} is incompatible with "
        f"the y values shape {y.shape}"
    )





def convert_kev_to_ev(
    energy: npt.NDArray[np.float64], threshold: float = 0.01
) -> npt.NDArray[np.float64]:
    """Convert energy values from keV to eV when the data appear to be in keV.

    Detection relies on the median step size between consecutive points.
    X-ray spectroscopy data in eV typically has steps of 0.1-1 eV; the same data
    in keV has steps of 0.0001-0.001 keV. A threshold of 0.01 sits cleanly
    between the two ranges for all realistic scans.

    Args:
        energy: Energy axis values (1D or 2D array).
        threshold: If the median step size is below this value the data
            are assumed to be in keV and multiplied by 1000.

    Raises:
        ValueError: If energy has more than two dimensions.
    """
    if energy.ndim not in (1, 2):
        raise ValueError(f"Expected 1D or 2D energy array, got ndim={energy.ndim}.")
    # np.diff needs at least 2 points along the last axis.
    if energy.shape[-1] < 2:
        return energy
    step = float(np.median(np.abs(np.diff(energy, axis=-1))))
    if step < threshold:
        return energy * 1000
    return energy





def get_attribute_values(
    groups: list[Group], attributes: str | list[str]
) -> dict[str, Any]:
    """Collect attribute values from a list of Larch groups into arrays.

    Args:
        groups: Larch groups to read from.
        attributes: One or more attribute names to extract.

    Returns:
        Dictionary mapping each attribute name to a NumPy array of its
        values across all groups.

    Raises:
        KeyError: If an attribute is not found in a group.
    """
    if isinstance(attributes, str):
        attributes = [attributes]

    result = {}
    for attribute in attributes:
        values = []
        for i, group in enumerate(groups):
            if not hasattr(group, attribute):
                raise KeyError(
                    f"Attribute '{attribute}' not found in group at index {i}"
                )
            values.append(getattr(group, attribute))
        result[attribute] = np.array(values)
    return result