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: Literal["first", "largest", "smallest"] = "first",
) -> Tuple[List[numpy.ndarray], List[str]]:
    """
    Convert values so that units sharing the same dimensionality
    are expressed in a common (reference) unit.

    The reference unit is the first occurrence in each group.

    Example:
        units = ["um", "mm", "deg", "rad"]
        -> ["um", "um", "deg", "deg"]

    :param values: sequence of numpy arrays (one per axis)
    :param units: sequence of unit strings
    :returns: (converted_values, normalized_units)
    """
    ureg = unit_registry()

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

    new_values = list(values)
    new_units = list(normalized_units)

    for indices in groups.values():
        ref_idx = indices[0]
        ref_unit = normalized_units[ref_idx]
        ref_parsed = ureg.parse_units(ref_unit)

        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 = []

    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]:
    return [u if u is not None else "dimensionless" for u in units]



def _group_units_by_dimensionality(
    units: Sequence[str],
    reference_units: Literal["first", "largest", "smallest"] = "first",
) -> Dict[str, List[int]]:
    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)

    if reference_units != "first":
        for key, indices in groups.items():
            if reference_units == "smallest":
                ref_idx = min(indices, key=lambda i: _get_scale(units[i], ureg))
            elif reference_units == "largest":
                ref_idx = max(indices, key=lambda i: _get_scale(units[i], ureg))
            else:
                raise ValueError(
                    "reference_units must be 'first', 'largest' or 'smallest'"
                )

            # move reference index to first position
            groups[key] = [ref_idx] + [i for i in indices if i != ref_idx]

    return groups


def _get_scale(unit: str, ureg) -> float:
    return (1 * ureg.parse_units(unit)).to_base_units().magnitude