Source code for ewoksfluo.units

from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
from typing import Sequence
from typing import Tuple
from typing import Union

import numpy
import pint




def unit_registry() -> pint.UnitRegistry:
    return pint.get_application_registry()





def units_as_str(units: pint.Unit) -> str:
    return f"{units:~}"





def convert_units_to_group_reference(
    values: Sequence[numpy.ndarray],
    units: Sequence[Optional[str]],
    reference_units: Optional[Dict[str, str]] = None,
    fallback: Optional[Literal["first", "largest", "smallest"]] = "smallest",
) -> Tuple[List[numpy.ndarray], List[str]]:
    """
    Convert values so that units sharing the same dimensionality
    are expressed in a common (reference) unit.

    Priority:
    1. Use explicit reference_units mapping per dimensionality
    2. Otherwise fallback to strategy ("first", "largest", "smallest")

    Example:
        units = ["um", "mm", "deg", "rad"]
        reference_units = {"[angle]": "deg"}
        fallback = "largest"

        -> ["mm", "mm", "deg", "deg"]

    :param values: sequence of numpy arrays (one per axis)
    :param units: sequence of unit strings
    :param reference_units: mapping dimensionality -> unit
    :param fallback: strategy for unspecified dimensionalities
    :returns: (converted_values, normalized_units)
    """
    ureg = unit_registry()

    normalized_units = normalize_units(units)
    groups = _group_units_by_dimensionality(normalized_units)

    new_values = list(values)
    new_units = list(normalized_units)
    if reference_units is None:
        reference_units = {}

    for dim_key, indices in groups.items():
        if dim_key in reference_units:
            ref_unit = reference_units[dim_key]
            ref_parsed = ureg.parse_units(ref_unit)
        elif fallback is not None:
            ref_idx = _select_reference_index(indices, normalized_units, ureg, fallback)
            ref_unit = normalized_units[ref_idx]
            ref_parsed = ureg.parse_units(ref_unit)
        else:
            continue

        for i in indices:
            parsed = ureg.parse_units(normalized_units[i])
            new_values[i] = (values[i] * parsed).to(ref_parsed).magnitude
            new_units[i] = ref_unit

    return new_values, new_units





def convert_values_to_units(
    values: Optional[Sequence[Union[float, Tuple[float, str]]]],
    target_units: Sequence[Optional[str]],
) -> Optional[List[float]]:
    """
    Convert values to match target units per axis.

    :param values: sequence of floats or (value, unit) tuples
    :param target_units: units per axis
    :returns: values in consistent units
    """
    if values is None:
        return None

    ureg = unit_registry()
    result: List[float] = []

    normalized_units = normalize_units(target_units)

    for value, target_unit in zip(values, normalized_units):
        if isinstance(value, tuple):
            v, u = value
            parsed = ureg.parse_units(u)
            target = ureg.parse_units(target_unit)
            v = (v * parsed).to(target).magnitude
            result.append(float(v))
        else:
            result.append(float(value))

    return result





def normalize_units(units: Sequence[Optional[str]]) -> List[str]:
    """
    Replace None with 'dimensionless'
    """
    return [u if u is not None else "dimensionless" for u in units]



def _group_units_by_dimensionality(
    units: Sequence[str],
) -> Dict[str, List[int]]:
    """
    Group indices of units by dimensionality.
    """
    ureg = unit_registry()
    groups: Dict[str, List[int]] = {}

    for i, u in enumerate(units):
        parsed = ureg.parse_units(u)
        key = str(parsed.dimensionality)
        groups.setdefault(key, []).append(i)

    return groups


def _select_reference_index(
    indices: List[int],
    units: Sequence[str],
    ureg: pint.UnitRegistry,
    strategy: Optional[Literal["first", "largest", "smallest"]],
) -> int:
    if strategy is None:
        raise ValueError(
            "No fallback strategy provided and no reference unit specified "
            "for this dimensionality"
        )

    if strategy == "first":
        return indices[0]

    if strategy == "smallest":
        return min(indices, key=lambda i: _get_scale(units[i], ureg))

    if strategy == "largest":
        return max(indices, key=lambda i: _get_scale(units[i], ureg))

    raise ValueError("fallback must be 'first', 'largest', 'smallest', or None")


def _get_scale(unit: str, ureg: pint.UnitRegistry) -> float:
    """
    Return scale of unit in base units (used for comparison).
    """
    return (1 * ureg.parse_units(unit)).to_base_units().magnitude