diff --git a/nion/data/AnnotatedArray.py b/nion/data/AnnotatedArray.py new file mode 100644 index 0000000..dca12fa --- /dev/null +++ b/nion/data/AnnotatedArray.py @@ -0,0 +1,631 @@ +"""Annotated n-dimensional arrays with calibrated, labeled axes. + +This module pairs a numpy array with a structural :class:`ArrayDescriptor` and +contextual :class:`ArrayMetadata`. An :class:`ArrayHeader` packages those +objects with the storage data type when the buffer is passed separately. Axes +are grouped into :class:`Axis`/:class:`AxisGroup` objects with per-axis and +intensity :class:`Calibration` mappings. +""" + +from __future__ import annotations + +import dataclasses +import datetime +import typing +import types +import numpy +import tzlocal +import zoneinfo + +from nion.data import Calibration as LegacyCalibration +from nion.data import DataAndMetadata + + +class ValueType: + """Supported value types for array data.""" + SCALAR = "scalar" + COMPLEX = "complex" + RGB = "rgb" + RGBA = "rgba" + VECTOR = "vector" + + +def infer_value_type(dtype: numpy.dtype) -> str: + """Infer the value type from a numpy dtype. + + Args: + dtype: The numpy dtype to infer from. + + Returns: + A ValueType string (scalar, complex, rgb, rgba, or vector). + + Raises: + ValueError: If the dtype cannot be mapped to a supported value type. + """ + dtype = numpy.dtype(typing.cast(typing.Any, dtype)) + + # Complex types + if dtype.kind == 'c': + return ValueType.COMPLEX + + # Structured/record types for RGB and RGBA + if dtype.names is not None: + if len(dtype.names) == 3 and set(dtype.names) >= {'r', 'g', 'b'}: + return ValueType.RGB + elif len(dtype.names) == 4 and set(dtype.names) >= {'r', 'g', 'b', 'a'}: + return ValueType.RGBA + # Other structured types are treated as vectors + return ValueType.VECTOR + + # Numeric types (int, uint, float, bool) + if dtype.kind in ('b', 'i', 'u', 'f'): + return ValueType.SCALAR + + raise ValueError(f"Unsupported dtype {dtype} for value type inference") + + + + +class Calibration(typing.Protocol): + """Protocol converting between array indices and physical coordinates.""" + + @property + def unit(self) -> str: ... + + def to_coordinate(self, index: float) -> float: ... + def to_index(self, coordinate: float) -> float: ... + + +@dataclasses.dataclass(frozen=True) +class AffineCalibration(Calibration): + """Linear calibration of the form ``coordinate = offset + index * scale``.""" + + scale: float = 1.0 + offset: float = 0.0 + unit: str = "" + + def __post_init__(self) -> None: + if self.scale == 0.0: + raise ValueError("AffineCalibration scale must be non-zero") + + def to_coordinate(self, index: float) -> float: + return self.offset + index * self.scale + + def to_index(self, coordinate: float) -> float: + return (coordinate - self.offset) / self.scale + + +@dataclasses.dataclass(frozen=True) +class CalibrationSet: + """Keyed calibrations with an optional primary calibration key.""" + + calibrations: typing.Mapping[str, Calibration] = dataclasses.field(default_factory=dict) + primary_key: str | None = None + + def __post_init__(self) -> None: + calibrations = dict(self.calibrations) + if self.primary_key is not None and self.primary_key not in calibrations: + raise ValueError(f"primary_key {self.primary_key!r} is not present in calibrations") + if not calibrations and self.primary_key is not None: + raise ValueError("primary_key must be None when calibrations is empty") + object.__setattr__(self, "calibrations", types.MappingProxyType(calibrations)) + + @property + def calibration_keys(self) -> tuple[str, ...]: + return tuple(self.calibrations.keys()) + + @property + def primary_calibration(self) -> Calibration: + if self.primary_key is None: + raise KeyError("No primary calibration is designated") + return self.calibrations[self.primary_key] + + def has_calibration(self, key: str) -> bool: + return key in self.calibrations + + def get_calibration(self, key: str | None = None) -> Calibration: + target_key = self.primary_key if key is None else key + if target_key is None: + raise KeyError("No primary calibration is designated") + calibration = self.calibrations.get(target_key) + if calibration is None: + raise KeyError(f"Unknown calibration {target_key!r}") + return calibration + + @staticmethod + def from_calibration(calibration: Calibration, key: str) -> CalibrationSet: + return CalibrationSet(calibrations={key: calibration}, primary_key=key) + + def with_primary_calibration(self, key: str) -> CalibrationSet: + if key not in self.calibrations: + raise KeyError(f"Unknown calibration {key!r}") + return dataclasses.replace(self, primary_key=key) + + def with_calibration(self, key: str, calibration: Calibration, *, make_primary: bool = False) -> CalibrationSet: + calibrations = dict(self.calibrations) + calibrations[key] = calibration + primary_key = key if make_primary else self.primary_key + return CalibrationSet(calibrations=calibrations, primary_key=primary_key) + + +@dataclasses.dataclass(frozen=True) +class Axis: + """A dimension label, size, and coordinate metadata.""" + + label: str = "" + size: int = 1 + + def __post_init__(self) -> None: + if self.size <= 0: + raise ValueError(f"Axis size must be positive, got {self.size}") + + +@dataclasses.dataclass(frozen=True) +class CoordinateCalibration: + """A labeled per-axis coordinate calibration bundle for an axis group.""" + + calibrations: tuple[Calibration, ...] = dataclasses.field(default_factory=tuple) + label: str | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "calibrations", tuple(self.calibrations)) + + +@dataclasses.dataclass(frozen=True) +class AxisGroup: + """An ordered group of sized axes with optional coordinate calibrations.""" + + axes: tuple[Axis, ...] = dataclasses.field(default_factory=tuple) + coordinate_system_id: str | None = None + coordinate_calibrations: typing.Mapping[str, CoordinateCalibration] = dataclasses.field(default_factory=dict) + primary_calibration_key: str | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "axes", tuple(self.axes)) + coordinate_calibrations = dict(self.coordinate_calibrations) + normalized_coordinate_calibrations: dict[str, CoordinateCalibration] = dict() + for key, coordinate_calibration in coordinate_calibrations.items(): + if len(coordinate_calibration.calibrations) != self.rank: + raise ValueError(f"coordinate calibration {key!r} rank {len(coordinate_calibration.calibrations)} does not match axis group rank {self.rank}") + normalized_coordinate_calibrations[key] = coordinate_calibration + object.__setattr__(self, "coordinate_calibrations", types.MappingProxyType(normalized_coordinate_calibrations)) + + if self.primary_calibration_key is not None and self.primary_calibration_key not in normalized_coordinate_calibrations: + raise ValueError(f"primary_calibration_key {self.primary_calibration_key!r} is not present in coordinate_calibrations") + if not normalized_coordinate_calibrations and self.primary_calibration_key is not None: + raise ValueError("primary_calibration_key must be None when coordinate_calibrations is empty") + + @staticmethod + def from_1d_size( + size: int, + *, + label: str = "x", + coordinate_calibrations: typing.Mapping[str, CoordinateCalibration] | None = None, + primary_calibration_key: str | None = None, + ) -> AxisGroup: + """Create a 1D axis group with optional coordinate calibrations.""" + calibrations: typing.Mapping[str, CoordinateCalibration] = coordinate_calibrations or {} + + return AxisGroup( + axes=(Axis(label, size),), + coordinate_calibrations=calibrations, + primary_calibration_key=primary_calibration_key, + ) + + @staticmethod + def from_2d_size( + size: tuple[int, int], + *, + labels: tuple[str, str] = ("x", "y"), + coordinate_calibrations: typing.Mapping[str, CoordinateCalibration] | None = None, + primary_calibration_key: str | None = None, + ) -> AxisGroup: + """Create a 2D axis group with optional coordinate calibrations.""" + calibrations: typing.Mapping[str, CoordinateCalibration] = coordinate_calibrations or {} + + size_x, size_y = size + x_label, y_label = labels + return AxisGroup( + axes=( + Axis(x_label, size_x), + Axis(y_label, size_y), + ), + coordinate_calibrations=calibrations, + primary_calibration_key=primary_calibration_key, + ) + + @property + def shape(self) -> tuple[int, ...]: + return tuple(axis.size for axis in self.axes) + + @property + def rank(self) -> int: + return len(self.axes) + + @property + def units(self) -> list[str]: + if self.primary_calibration_key is None: + return [] + return [calibration.unit for calibration in self.get_coordinate_calibration().calibrations] + + @property + def calibration_keys(self) -> tuple[str, ...]: + return tuple(self.coordinate_calibrations.keys()) + + def with_coordinate_calibration(self, key: str, coordinate_calibration: CoordinateCalibration) -> AxisGroup: + """Replace all coordinate calibrations with one calibration and make it primary.""" + return dataclasses.replace(self, coordinate_calibrations={key: coordinate_calibration}, primary_calibration_key=key) + + def with_added_coordinate_calibration(self, key: str, coordinate_calibration: CoordinateCalibration, *, make_primary: bool = False) -> AxisGroup: + """Add or replace one keyed coordinate calibration.""" + coordinate_calibrations = dict(self.coordinate_calibrations) + coordinate_calibrations[key] = coordinate_calibration + primary_calibration_key = key if make_primary else self.primary_calibration_key + return dataclasses.replace(self, coordinate_calibrations=coordinate_calibrations, primary_calibration_key=primary_calibration_key) + + def get_coordinate_calibration(self, key: str | None = None) -> CoordinateCalibration: + target_key = self.primary_calibration_key if key is None else key + if target_key is None: + raise KeyError("No primary coordinate calibration is designated") + coordinate_calibration = self.coordinate_calibrations.get(target_key) + if coordinate_calibration is None: + raise KeyError(f"Unknown coordinate calibration {target_key!r}") + return coordinate_calibration + + def get_calibration(self, axis: int, key: str | None = None) -> Calibration: + return self.get_coordinate_calibration(key).calibrations[axis] + + +@dataclasses.dataclass(frozen=True) +class ArrayDescriptor: + """Intrinsic structure and calibrations used to interpret an array.""" + + axis_groups: tuple[AxisGroup, ...] + intensity_calibrations: CalibrationSet = dataclasses.field(default_factory=CalibrationSet) + value_type: str = ValueType.SCALAR + + def __post_init__(self) -> None: + object.__setattr__(self, "axis_groups", tuple(self.axis_groups)) + if not self.axis_groups: + raise ValueError("ArrayDescriptor requires at least one axis group") + if any(axis_group.rank == 0 for axis_group in self.axis_groups[:-1]): + raise ValueError("Only the final axis group may have rank 0") + valid_types = {ValueType.SCALAR, ValueType.COMPLEX, ValueType.RGB, ValueType.RGBA, ValueType.VECTOR} + if self.value_type not in valid_types: + raise ValueError(f"value_type must be one of {valid_types}, got {self.value_type!r}") + + @property + def shape(self) -> tuple[int, ...]: + return tuple(dimension for axis_group in self.axis_groups for dimension in axis_group.shape) + + @property + def ndim(self) -> int: + return sum(axis_group.rank for axis_group in self.axis_groups) + + def get_intensity_calibration(self, key: str | None = None) -> Calibration: + if key is None and self.intensity_calibrations.primary_key is None: + return AffineCalibration() + return self.intensity_calibrations.get_calibration(key) + + def with_primary_intensity_calibration(self, key: str) -> ArrayDescriptor: + return dataclasses.replace(self, intensity_calibrations=self.intensity_calibrations.with_primary_calibration(key)) + + def with_intensity_calibration(self, key: str, calibration: Calibration, *, make_primary: bool = False) -> ArrayDescriptor: + return dataclasses.replace( + self, + intensity_calibrations=self.intensity_calibrations.with_calibration(key, calibration, make_primary=make_primary), + ) + + @property + def intensity_calibration_keys(self) -> tuple[str, ...]: + return self.intensity_calibrations.calibration_keys + + @property + def primary_intensity_calibration_key(self) -> str | None: + return self.intensity_calibrations.primary_key + + def with_value_type(self, value_type: str) -> ArrayDescriptor: + return dataclasses.replace(self, value_type=value_type) + + +@dataclasses.dataclass(frozen=True) +class ExtensionRecord: + """A versioned, encoded payload whose semantics are defined outside this module.""" + + extension_type_id: str + schema_version: int + payload: str + + def __post_init__(self) -> None: + if not self.extension_type_id: + raise ValueError("ExtensionRecord extension_type_id must not be empty") + if self.schema_version < 1: + raise ValueError("ExtensionRecord schema_version must be positive") + if not isinstance(self.payload, str): + raise TypeError("ExtensionRecord payload must be str") + + +@dataclasses.dataclass(frozen=True, init=False) +class ArrayMetadata: + """Context carried with an array but not required to interpret its data.""" + + created: datetime.datetime + attributes: typing.Mapping[str, typing.Any] + __extensions: typing.Mapping[str, ExtensionRecord] = dataclasses.field(repr=False) + + def __init__( + self, + *, + created: datetime.datetime | None = None, + attributes: typing.Mapping[str, typing.Any] | None = None, + extensions: typing.Iterable[ExtensionRecord] = (), + ) -> None: + created = created or datetime.datetime.now(tz=tzlocal.get_localzone()) + if created.tzinfo is None or created.utcoffset() is None: + raise ValueError("created must be timezone-aware") + extension_map: dict[str, ExtensionRecord] = dict() + for extension in extensions: + if extension.extension_type_id in extension_map: + raise ValueError(f"Duplicate extension type {extension.extension_type_id!r}") + extension_map[extension.extension_type_id] = extension + object.__setattr__(self, "created", created) + object.__setattr__(self, "attributes", types.MappingProxyType(dict(attributes or {}))) + object.__setattr__(self, "_ArrayMetadata__extensions", types.MappingProxyType(extension_map)) + + @property + def extension_type_ids(self) -> tuple[str, ...]: + return tuple(self.__extensions.keys()) + + def has_extension(self, extension_type_id: str) -> bool: + return extension_type_id in self.__extensions + + def get_extension(self, extension_type_id: str) -> ExtensionRecord: + extension = self.__extensions.get(extension_type_id) + if extension is None: + raise KeyError(f"Unknown extension {extension_type_id!r}") + return extension + + def with_extension(self, extension: ExtensionRecord) -> ArrayMetadata: + extensions = dict(self.__extensions) + extensions[extension.extension_type_id] = extension + return ArrayMetadata(created=self.created, attributes=self.attributes, extensions=extensions.values()) + + def without_extension(self, extension_type_id: str) -> ArrayMetadata: + extensions = dict(self.__extensions) + if extension_type_id not in extensions: + raise KeyError(f"Unknown extension {extension_type_id!r}") + del extensions[extension_type_id] + return ArrayMetadata(created=self.created, attributes=self.attributes, extensions=extensions.values()) + + def with_attributes(self, attributes: typing.Mapping[str, typing.Any]) -> ArrayMetadata: + return ArrayMetadata(created=self.created, attributes=attributes, extensions=self.__extensions.values()) + + + +@dataclasses.dataclass(frozen=True) +class ArrayHeader: + """Everything needed to describe an array without its underlying buffer.""" + + descriptor: ArrayDescriptor + dtype: numpy.typing.DTypeLike + metadata: ArrayMetadata = dataclasses.field(default_factory=ArrayMetadata) + + def __post_init__(self) -> None: + object.__setattr__(self, "dtype", numpy.dtype(typing.cast(typing.Any, self.dtype))) + + @property + def shape(self) -> tuple[int, ...]: + return self.descriptor.shape + + +class ArrayProtocol(typing.Protocol): + """Duck-array interface required by AnnotatedArray (satisfied by numpy.ndarray, h5py.Dataset, zarr.Array, etc.).""" + + @property + def shape(self) -> tuple[int, ...]: ... + + @property + def dtype(self) -> numpy.dtype[typing.Any]: ... + + def __array__(self, dtype: typing.Any = None, /, *, copy: bool | None = None) -> numpy.typing.NDArray[typing.Any]: ... + + +@dataclasses.dataclass +class AnnotatedArray: + """A numpy array paired with its descriptor and contextual metadata.""" + + data: ArrayProtocol + descriptor: ArrayDescriptor + metadata: ArrayMetadata = dataclasses.field(default_factory=ArrayMetadata) + + def __post_init__(self) -> None: + if self.descriptor.shape != self.data.shape: + raise ValueError(f"Descriptor shape {self.descriptor.shape} does not match array shape {self.data.shape}") + inferred_value_type = infer_value_type(self.data.dtype) + if self.descriptor.value_type != inferred_value_type: + raise ValueError(f"Descriptor value_type {self.descriptor.value_type!r} does not match inferred value_type {inferred_value_type!r} from array dtype {self.data.dtype}") + + def __array__(self, dtype: typing.Any = None, /, *, copy: bool | None = None) -> numpy.typing.NDArray[typing.Any]: + """Allow AnnotatedArray to be passed directly to numpy functions.""" + return numpy.asarray(self.data, dtype=dtype, copy=copy) + + @property + def header(self) -> ArrayHeader: + return ArrayHeader(self.descriptor, self.data.dtype, self.metadata) + + @classmethod + def from_header(cls, data: ArrayProtocol, header: ArrayHeader) -> AnnotatedArray: + if header.dtype != data.dtype: + raise ValueError(f"Header dtype {header.dtype} does not match array dtype {data.dtype}") + if header.descriptor.shape != data.shape: + raise ValueError(f"Header shape {header.descriptor.shape} does not match array shape {data.shape}") + inferred_value_type = infer_value_type(data.dtype) + if header.descriptor.value_type != inferred_value_type: + raise ValueError(f"Header value_type {header.descriptor.value_type!r} does not match inferred value_type {inferred_value_type!r} from array dtype {data.dtype}") + return cls(data, header.descriptor, header.metadata) + + def get_intensity_calibration(self, key: str | None = None) -> Calibration: + return self.descriptor.get_intensity_calibration(key) + + + +def zeros_annotated_array(axis_groups: typing.Sequence[AxisGroup], dtype: numpy.typing.DTypeLike = numpy.float64) -> AnnotatedArray: + """Create an AnnotatedArray filled with zeros. + + The value_type is inferred from the provided dtype. + + Args: + axis_groups: Sequence of AxisGroup objects describing the array structure. + dtype: The numpy dtype for the array. Defaults to float64 (scalar). + + Returns: + An AnnotatedArray with the specified structure, filled with zeros. + """ + dtype_obj = numpy.dtype(typing.cast(typing.Any, dtype)) + value_type = infer_value_type(dtype_obj) + descriptor = ArrayDescriptor(axis_groups=tuple(axis_groups), value_type=value_type) + return AnnotatedArray(data=numpy.zeros(descriptor.shape, dtype=dtype_obj), descriptor=descriptor) + + +def _make_axis_group(shape: tuple[int, ...], dimensional_calibrations: tuple[typing.Any, ...]) -> AxisGroup: + axes = list[Axis]() + coordinate_calibrations = list[Calibration]() + for i, size in enumerate(shape): + calibration = dimensional_calibrations[i] if i < len(dimensional_calibrations) else None + if calibration is None: + affine_calibration = AffineCalibration() + else: + affine_calibration = AffineCalibration(calibration.scale, calibration.offset, calibration.units) + axes.append(Axis(str(i), size)) + coordinate_calibrations.append(affine_calibration) + return AxisGroup( + axes=tuple(axes), + coordinate_calibrations={"calibrated": CoordinateCalibration(calibrations=tuple(coordinate_calibrations))} if coordinate_calibrations else {}, + primary_calibration_key="calibrated" if coordinate_calibrations else None, + ) + + +def _parse_timezone_offset(timezone_offset: str) -> datetime.timezone | None: + if len(timezone_offset) != 5 or timezone_offset[0] not in ("+", "-"): + return None + if not timezone_offset[1:].isdigit(): + return None + hours = int(timezone_offset[1:3]) + minutes = int(timezone_offset[3:5]) + delta = datetime.timedelta(hours=hours, minutes=minutes) + if timezone_offset[0] == "-": + delta = -delta + return datetime.timezone(delta) + + +def _legacy_timestamp_to_created( + timestamp: datetime.datetime, + timezone_name: str | None, + timezone_offset: str | None, +) -> datetime.datetime: + utc_timestamp = timestamp if timestamp.tzinfo else timestamp.replace(tzinfo=datetime.timezone.utc) + + if timezone_name: + try: + return utc_timestamp.astimezone(zoneinfo.ZoneInfo(timezone_name)) + except Exception: + pass + + if timezone_offset: + timezone = _parse_timezone_offset(timezone_offset) + if timezone is not None: + return utc_timestamp.astimezone(timezone) + + return utc_timestamp.astimezone(datetime.timezone.utc) + + +def from_data_and_metadata(xdata: DataAndMetadata.DataAndMetadata) -> AnnotatedArray: + """Convert a legacy DataAndMetadata instance into an AnnotatedArray.""" + data = xdata.data + if data is None: + raise ValueError("DataAndMetadata input data is missing.") + + data_descriptor = xdata.data_descriptor + dimensional_calibrations = tuple(xdata.dimensional_calibrations) + shape = xdata.data_shape + + axis_groups = list[AxisGroup]() + calibration_index = 0 + shape_index = 0 + if data_descriptor.is_sequence: + axis_groups.append(_make_axis_group(shape[shape_index:shape_index + 1], dimensional_calibrations[calibration_index:calibration_index + 1])) + calibration_index += 1 + shape_index += 1 + if data_descriptor.collection_dimension_count: + count = data_descriptor.collection_dimension_count + axis_groups.append(_make_axis_group(shape[shape_index:shape_index + count], dimensional_calibrations[calibration_index:calibration_index + count])) + calibration_index += count + shape_index += count + datum_count = data_descriptor.datum_dimension_count + axis_groups.append(_make_axis_group(shape[shape_index:shape_index + datum_count], dimensional_calibrations[calibration_index:calibration_index + datum_count])) + + intensity_calibration = xdata.intensity_calibration + descriptor = ArrayDescriptor( + axis_groups=tuple(axis_groups), + intensity_calibrations=CalibrationSet.from_calibration( + AffineCalibration(intensity_calibration.scale, intensity_calibration.offset, intensity_calibration.units), + "calibrated", + ), + value_type=infer_value_type(data.dtype), + ) + + created = _legacy_timestamp_to_created(xdata.timestamp, xdata.timezone, xdata.timezone_offset) + metadata = ArrayMetadata(created=created, attributes=dict(xdata.metadata)) + return AnnotatedArray(data=data, descriptor=descriptor, metadata=metadata) + + +def to_data_and_metadata(annotated_array: AnnotatedArray) -> DataAndMetadata.DataAndMetadata: + """Convert an AnnotatedArray into a legacy DataAndMetadata instance.""" + axis_groups = annotated_array.descriptor.axis_groups + if len(axis_groups) == 1: + is_sequence = False + collection_dimension_count = 0 + datum_dimension_count = axis_groups[0].rank + elif len(axis_groups) == 2: + is_sequence = False + collection_dimension_count = axis_groups[0].rank + datum_dimension_count = axis_groups[1].rank + elif len(axis_groups) == 3 and axis_groups[0].rank == 1: + is_sequence = True + collection_dimension_count = axis_groups[1].rank + datum_dimension_count = axis_groups[2].rank + else: + raise ValueError("AnnotatedArray axis groups are incompatible with DataAndMetadata conversion.") + + dimensional_calibrations = list[LegacyCalibration.Calibration]() + for axis_group in axis_groups: + for axis_index in range(axis_group.rank): + calibration = axis_group.get_calibration(axis_index) if axis_group.primary_calibration_key else AffineCalibration() + if not isinstance(calibration, AffineCalibration): + raise ValueError("Only affine coordinate calibrations are supported for DataAndMetadata conversion.") + dimensional_calibrations.append(LegacyCalibration.Calibration(calibration.offset, calibration.scale, calibration.unit)) + + intensity_calibration = LegacyCalibration.Calibration() + if annotated_array.descriptor.primary_intensity_calibration_key: + calibration = annotated_array.get_intensity_calibration() + if not isinstance(calibration, AffineCalibration): + raise ValueError("Only affine intensity calibrations are supported for DataAndMetadata conversion.") + intensity_calibration = LegacyCalibration.Calibration(calibration.offset, calibration.scale, calibration.unit) + + created = annotated_array.metadata.created + timezone_name = typing.cast(str | None, getattr(created.tzinfo, "key", None)) or created.tzname() + timestamp = created.astimezone(datetime.timezone.utc).replace(tzinfo=None) + timezone_offset = created.strftime("%z") + + return DataAndMetadata.new_data_and_metadata( + data=numpy.asarray(annotated_array.data), + intensity_calibration=intensity_calibration, + dimensional_calibrations=tuple(dimensional_calibrations), + metadata=dict(annotated_array.metadata.attributes), + data_descriptor=DataAndMetadata.DataDescriptor(is_sequence, collection_dimension_count, datum_dimension_count), + timestamp=timestamp, + timezone=timezone_name, + timezone_offset=timezone_offset, + ) + + diff --git a/nion/data/test/AnnotatedArray_test.py b/nion/data/test/AnnotatedArray_test.py new file mode 100644 index 0000000..d79d188 --- /dev/null +++ b/nion/data/test/AnnotatedArray_test.py @@ -0,0 +1,371 @@ +import datetime +import io +import typing +import unittest +import zoneinfo + +import h5py +import numpy + +from nion.data import AnnotatedArray +from nion.data import Calibration +from nion.data import DataAndMetadata + + +class TestAnnotatedArray(unittest.TestCase): + + _descriptor_variants: tuple[tuple[bool, int, int], ...] = ( + (False, 0, 1), + (False, 0, 2), + (False, 1, 1), + (False, 1, 2), + (False, 2, 1), + (False, 2, 2), + (True, 0, 1), + (True, 0, 2), + (True, 1, 1), + (True, 1, 2), + (True, 2, 1), + (True, 2, 2), + ) + + def test_calibration_set_uses_explicit_calibration_accessors(self) -> None: + primary = AnnotatedArray.AffineCalibration(unit="nm") + alternate = AnnotatedArray.AffineCalibration(unit="rad") + calibrations = AnnotatedArray.CalibrationSet.from_calibration(primary, key="primary").with_calibration("alternate", alternate) + + self.assertTrue(calibrations.has_calibration("alternate")) + self.assertIs(primary, calibrations.primary_calibration) + self.assertIs(alternate, calibrations.get_calibration("alternate")) + self.assertIs(alternate, calibrations.with_primary_calibration("alternate").primary_calibration) + + def test_array_descriptor_describes_shape_and_rank(self) -> None: + collection = AnnotatedArray.AxisGroup.from_1d_size(3) + signal = AnnotatedArray.AxisGroup.from_2d_size((4, 5)) + descriptor = AnnotatedArray.ArrayDescriptor((collection, signal)) + + self.assertEqual((3, 4, 5), descriptor.shape) + self.assertEqual(3, descriptor.ndim) + + def test_array_descriptor_uses_identity_intensity_calibration_when_no_primary_is_designated(self) -> None: + descriptor = AnnotatedArray.ArrayDescriptor((AnnotatedArray.AxisGroup.from_1d_size(3),)) + + calibration = descriptor.get_intensity_calibration() + self.assertIsInstance(calibration, AnnotatedArray.AffineCalibration) + affine_calibration = typing.cast(AnnotatedArray.AffineCalibration, calibration) + self.assertEqual(1.0, affine_calibration.scale) + self.assertEqual(0.0, affine_calibration.offset) + self.assertEqual("", affine_calibration.unit) + + with self.assertRaises(KeyError): + descriptor.get_intensity_calibration("missing") + + def test_axis_group_size_factories_accept_optional_coordinate_calibrations(self) -> None: + single_calibration = AnnotatedArray.CoordinateCalibration(calibrations=(AnnotatedArray.AffineCalibration(unit="nm"),)) + group_1d = AnnotatedArray.AxisGroup.from_1d_size( + 3, + coordinate_calibrations={"spatial": single_calibration}, + primary_calibration_key="spatial", + ) + self.assertEqual(("spatial",), group_1d.calibration_keys) + self.assertEqual("spatial", group_1d.primary_calibration_key) + + map_calibration = AnnotatedArray.CoordinateCalibration( + calibrations=(AnnotatedArray.AffineCalibration(unit="nm"), AnnotatedArray.AffineCalibration(unit="nm")) + ) + group_2d = AnnotatedArray.AxisGroup.from_2d_size( + (2, 3), + coordinate_calibrations={"camera": map_calibration}, + primary_calibration_key="camera", + ) + self.assertEqual(("camera",), group_2d.calibration_keys) + self.assertEqual("camera", group_2d.primary_calibration_key) + + def test_array_descriptor_requires_valid_axis_group_layout(self) -> None: + with self.assertRaisesRegex(ValueError, "at least one"): + AnnotatedArray.ArrayDescriptor(()) + + scalar_group = AnnotatedArray.AxisGroup() + vector_group = AnnotatedArray.AxisGroup.from_1d_size(3) + with self.assertRaisesRegex(ValueError, "Only the final"): + AnnotatedArray.ArrayDescriptor((scalar_group, vector_group)) + + def test_array_metadata_controls_extension_access(self) -> None: + extension = AnnotatedArray.ExtensionRecord("org.nion.test", 1, "value=42") + metadata = AnnotatedArray.ArrayMetadata(extensions=(extension,)) + + self.assertEqual("org.nion.test", extension.extension_type_id) + self.assertEqual(("org.nion.test",), metadata.extension_type_ids) + self.assertTrue(metadata.has_extension("org.nion.test")) + self.assertIs(extension, metadata.get_extension("org.nion.test")) + with self.assertRaises(KeyError): + metadata.get_extension("org.nion.missing") + + replacement = AnnotatedArray.ExtensionRecord("org.nion.test", 2, "value=43") + replaced_metadata = metadata.with_extension(replacement) + self.assertEqual(2, replaced_metadata.get_extension("org.nion.test").schema_version) + self.assertEqual(1, metadata.get_extension("org.nion.test").schema_version) + self.assertFalse(replaced_metadata.without_extension("org.nion.test").has_extension("org.nion.test")) + + def test_array_metadata_validates_extensions(self) -> None: + with self.assertRaisesRegex(ValueError, "must not be empty"): + AnnotatedArray.ExtensionRecord("", 1, "") + with self.assertRaisesRegex(ValueError, "must be positive"): + AnnotatedArray.ExtensionRecord("org.nion.test", 0, "") + with self.assertRaisesRegex(TypeError, "must be str"): + AnnotatedArray.ExtensionRecord("org.nion.test", 1, bytearray()) # type: ignore[arg-type] + + extension = AnnotatedArray.ExtensionRecord("org.nion.test", 1, "") + with self.assertRaisesRegex(ValueError, "Duplicate"): + AnnotatedArray.ArrayMetadata(extensions=(extension, extension)) + + def test_array_metadata_snapshots_attributes_and_requires_timezone(self) -> None: + source = {"note": "original"} + created = datetime.datetime(2026, 7, 16, tzinfo=datetime.timezone.utc) + metadata = AnnotatedArray.ArrayMetadata(created=created, attributes=source) + source["note"] = "changed" + + self.assertEqual("original", metadata.attributes["note"]) + self.assertEqual("replacement", metadata.with_attributes({"note": "replacement"}).attributes["note"]) + self.assertEqual(datetime.timedelta(0), metadata.created.utcoffset()) + with self.assertRaises(TypeError): + metadata.attributes["note"] = "changed" # type: ignore[index] + with self.assertRaisesRegex(ValueError, "timezone-aware"): + AnnotatedArray.ArrayMetadata(created=datetime.datetime(2026, 7, 16)) + + def test_array_metadata_created_retains_iana_zone_and_offset(self) -> None: + created = datetime.datetime(2026, 7, 16, 12, tzinfo=zoneinfo.ZoneInfo("America/Los_Angeles")) + metadata = AnnotatedArray.ArrayMetadata(created=created) + + self.assertEqual("America/Los_Angeles", getattr(metadata.created.tzinfo, "key", None)) + self.assertEqual(datetime.timedelta(hours=-7), metadata.created.utcoffset()) + + def test_array_metadata_default_created_has_iana_timezone(self) -> None: + # tzlocal.get_localzone() returns a zoneinfo.ZoneInfo with a proper IANA key + # (e.g. "America/New_York"). A plain datetime.timezone fixed-offset object + # would not satisfy the isinstance check and would have no .key attribute. + metadata = AnnotatedArray.ArrayMetadata() + + self.assertIsInstance(metadata.created.tzinfo, zoneinfo.ZoneInfo) + self.assertIsNotNone(metadata.created.tzinfo.key) # type: ignore[union-attr] + + def test_array_header_can_be_passed_without_data(self) -> None: + descriptor = AnnotatedArray.ArrayDescriptor((AnnotatedArray.AxisGroup.from_2d_size((4, 5)),)) + metadata = AnnotatedArray.ArrayMetadata(attributes={"note": "test"}) + header = AnnotatedArray.ArrayHeader(descriptor, "float32", metadata) + + self.assertEqual((4, 5), header.shape) + self.assertEqual(numpy.dtype(numpy.float32), header.dtype) + + first = AnnotatedArray.AnnotatedArray.from_header(numpy.zeros(header.shape, dtype=header.dtype), header) + second = AnnotatedArray.AnnotatedArray.from_header(numpy.ones(header.shape, dtype=header.dtype), header) + self.assertEqual(header, first.header) + self.assertIsNot(header, first.header) + self.assertEqual(first.header, second.header) + self.assertIs(descriptor, first.descriptor) + self.assertIs(metadata, first.metadata) + + def test_annotated_array_validates_data_against_header(self) -> None: + descriptor = AnnotatedArray.ArrayDescriptor((AnnotatedArray.AxisGroup.from_1d_size(3),)) + header = AnnotatedArray.ArrayHeader(descriptor, numpy.float32) + + with self.assertRaisesRegex(ValueError, "shape"): + AnnotatedArray.AnnotatedArray(numpy.zeros((4,), dtype=numpy.float32), descriptor) + with self.assertRaisesRegex(ValueError, "dtype"): + AnnotatedArray.AnnotatedArray.from_header(numpy.zeros((3,), dtype=numpy.float64), header) + + def test_zeros_annotated_array_constructs_matching_shape_and_dtype(self) -> None: + group = AnnotatedArray.AxisGroup.from_2d_size((2, 3)) + array = AnnotatedArray.zeros_annotated_array((group,), dtype=numpy.float32) + + self.assertEqual((2, 3), array.data.shape) + self.assertEqual(numpy.dtype(numpy.float32), array.header.dtype) + + def test_annotated_array_is_numpy_passable(self) -> None: + group = AnnotatedArray.AxisGroup.from_1d_size(4) + array = AnnotatedArray.zeros_annotated_array((group,), dtype=numpy.float64) + + # AnnotatedArray itself is directly usable with numpy functions via __array__ + self.assertEqual(0.0, float(numpy.sum(array))) + as_array = numpy.asarray(array) + self.assertIsInstance(as_array, numpy.ndarray) + self.assertEqual((4,), as_array.shape) + self.assertEqual(numpy.dtype(numpy.float64), as_array.dtype) + + def test_annotated_array_data_is_numpy_passable(self) -> None: + group = AnnotatedArray.AxisGroup.from_2d_size((2, 3)) + array = AnnotatedArray.zeros_annotated_array((group,), dtype=numpy.float32) + + # data satisfies ArrayProtocol including __array__, so it is directly usable with numpy + self.assertEqual(0.0, float(numpy.sum(array.data))) + result = numpy.asarray(array.data) + self.assertIsInstance(result, numpy.ndarray) + self.assertEqual((2, 3), result.shape) + self.assertEqual(numpy.dtype(numpy.float32), result.dtype) + + def test_annotated_array_accepts_h5py_dataset_as_data(self) -> None: + # h5py.Dataset satisfies ArrayProtocol: it exposes .shape, .dtype, and __array__. + # AnnotatedArray must accept it without forcing an eager copy into memory. + buf = io.BytesIO() + with h5py.File(buf, "w") as f: + ds = f.create_dataset("data", data=numpy.arange(6, dtype=numpy.float32).reshape(2, 3)) + group = AnnotatedArray.AxisGroup.from_2d_size((2, 3)) + annotated = AnnotatedArray.AnnotatedArray(ds, AnnotatedArray.ArrayDescriptor((group,))) + + # shape and dtype are read directly from the dataset without materialising it + self.assertEqual((2, 3), annotated.data.shape) + self.assertEqual(numpy.dtype(numpy.float32), annotated.data.dtype) + + # numpy functions work via __array__ (materialises on demand) + result = numpy.asarray(annotated) + self.assertIsInstance(result, numpy.ndarray) + self.assertEqual((2, 3), result.shape) + numpy.testing.assert_array_equal(result, numpy.arange(6, dtype=numpy.float32).reshape(2, 3)) + + def test_from_data_and_metadata_covers_sequence_collection_and_datum_variations(self) -> None: + def flatten_affine_calibrations(annotated_array: AnnotatedArray.AnnotatedArray) -> tuple[AnnotatedArray.AffineCalibration, ...]: + calibrations = list[AnnotatedArray.AffineCalibration]() + for axis_group in annotated_array.descriptor.axis_groups: + for axis_index in range(axis_group.rank): + calibration = axis_group.get_calibration(axis_index) if axis_group.primary_calibration_key else AnnotatedArray.AffineCalibration() + calibrations.append(typing.cast(AnnotatedArray.AffineCalibration, calibration)) + return tuple(calibrations) + + for is_sequence, collection_rank, datum_rank in self._descriptor_variants: + with self.subTest(is_sequence=is_sequence, collection_rank=collection_rank, datum_rank=datum_rank): + total_rank = (1 if is_sequence else 0) + collection_rank + datum_rank + shape = tuple(range(2, 2 + total_rank)) + data = numpy.arange(int(numpy.prod(shape)), dtype=numpy.float32).reshape(shape) + dimensional_calibrations = tuple(Calibration.Calibration(offset=i + 0.25, scale=i + 1.5, units=f"u{i}") for i in range(total_rank)) + xdata = DataAndMetadata.new_data_and_metadata( + data=data, + intensity_calibration=Calibration.Calibration(offset=7.0, scale=3.0, units="counts"), + dimensional_calibrations=dimensional_calibrations, + metadata={"descriptor_variant": f"legacy-{int(is_sequence)}-{collection_rank}-{datum_rank}"}, + data_descriptor=DataAndMetadata.DataDescriptor(is_sequence, collection_rank, datum_rank), + timestamp=datetime.datetime(2026, 7, 16, 19, 0, 0), + timezone="America/Los_Angeles", + timezone_offset="-0700", + ) + + annotated = AnnotatedArray.from_data_and_metadata(xdata) + expected_axis_group_ranks = (1, datum_rank) if is_sequence and collection_rank == 0 else (1, collection_rank, datum_rank) if is_sequence else (datum_rank,) if collection_rank == 0 else (collection_rank, datum_rank) + self.assertEqual(expected_axis_group_ranks, tuple(axis_group.rank for axis_group in annotated.descriptor.axis_groups)) + self.assertEqual(xdata.data_shape, annotated.descriptor.shape) + self.assertEqual("counts", typing.cast(AnnotatedArray.AffineCalibration, annotated.get_intensity_calibration()).unit) + + flattened_annotated_calibrations = flatten_affine_calibrations(annotated) + self.assertEqual(len(xdata.dimensional_calibrations), len(flattened_annotated_calibrations)) + for i, (legacy_calibration, annotated_calibration) in enumerate(zip(xdata.dimensional_calibrations, flattened_annotated_calibrations)): + with self.subTest(calibration_index=i): + self.assertEqual(legacy_calibration.offset, annotated_calibration.offset) + self.assertEqual(legacy_calibration.scale, annotated_calibration.scale) + self.assertEqual(legacy_calibration.units, annotated_calibration.unit) + + self.assertEqual(xdata.metadata["descriptor_variant"], annotated.metadata.attributes["descriptor_variant"]) + + def test_to_data_and_metadata_covers_sequence_collection_and_datum_variations(self) -> None: + def flatten_affine_calibrations(annotated_array: AnnotatedArray.AnnotatedArray) -> tuple[AnnotatedArray.AffineCalibration, ...]: + calibrations = list[AnnotatedArray.AffineCalibration]() + for axis_group in annotated_array.descriptor.axis_groups: + for axis_index in range(axis_group.rank): + calibration = axis_group.get_calibration(axis_index) if axis_group.primary_calibration_key else AnnotatedArray.AffineCalibration() + calibrations.append(typing.cast(AnnotatedArray.AffineCalibration, calibration)) + return tuple(calibrations) + + for is_sequence, collection_rank, datum_rank in self._descriptor_variants: + with self.subTest(is_sequence=is_sequence, collection_rank=collection_rank, datum_rank=datum_rank): + if is_sequence and collection_rank == 0: + # (1, datum_rank) is ambiguous in legacy form and round-trips as non-sequence collection rank 1. + continue + + expected_axis_group_ranks = (1, collection_rank, datum_rank) if is_sequence else (datum_rank,) if collection_rank == 0 else (collection_rank, datum_rank) + dim_sizes = tuple(range(2, 2 + sum(expected_axis_group_ranks))) + + dim_index = 0 + axis_groups = list[AnnotatedArray.AxisGroup]() + for axis_group_rank in expected_axis_group_ranks: + axes = list[AnnotatedArray.Axis]() + calibrations = list[AnnotatedArray.AffineCalibration]() + for _ in range(axis_group_rank): + size = dim_sizes[dim_index] + calibration_index = dim_index + axes.append(AnnotatedArray.Axis(label=f"a{calibration_index}", size=size)) + calibrations.append(AnnotatedArray.AffineCalibration(offset=calibration_index + 0.5, scale=calibration_index + 1.25, unit=f"ua{calibration_index}")) + dim_index += 1 + axis_groups.append( + AnnotatedArray.AxisGroup( + axes=tuple(axes), + coordinate_calibrations={"calibrated": AnnotatedArray.CoordinateCalibration(calibrations=tuple(calibrations))}, + primary_calibration_key="calibrated", + ) + ) + + descriptor = AnnotatedArray.ArrayDescriptor( + axis_groups=tuple(axis_groups), + intensity_calibrations=AnnotatedArray.CalibrationSet.from_calibration( + AnnotatedArray.AffineCalibration(offset=4.0, scale=2.0, unit="counts"), + "calibrated", + ), + ) + metadata = AnnotatedArray.ArrayMetadata( + created=datetime.datetime(2026, 7, 16, 12, tzinfo=zoneinfo.ZoneInfo("America/Los_Angeles")), + attributes={"descriptor_variant": f"annotated-{int(is_sequence)}-{collection_rank}-{datum_rank}"}, + ) + annotated = AnnotatedArray.AnnotatedArray( + data=numpy.arange(int(numpy.prod(dim_sizes)), dtype=numpy.float32).reshape(dim_sizes), + descriptor=descriptor, + metadata=metadata, + ) + + xdata = AnnotatedArray.to_data_and_metadata(annotated) + self.assertEqual(is_sequence, xdata.is_sequence) + self.assertEqual(collection_rank, xdata.collection_dimension_count) + self.assertEqual(datum_rank, xdata.datum_dimension_count) + self.assertEqual(annotated.descriptor.shape, xdata.data_shape) + self.assertEqual("counts", xdata.intensity_calibration.units) + + flattened_annotated_calibrations = flatten_affine_calibrations(annotated) + self.assertEqual(len(flattened_annotated_calibrations), len(xdata.dimensional_calibrations)) + for i, (annotated_calibration, legacy_calibration) in enumerate(zip(flattened_annotated_calibrations, xdata.dimensional_calibrations)): + with self.subTest(calibration_index=i): + self.assertEqual(annotated_calibration.offset, legacy_calibration.offset) + self.assertEqual(annotated_calibration.scale, legacy_calibration.scale) + self.assertEqual(annotated_calibration.unit, legacy_calibration.units) + + round_tripped = AnnotatedArray.from_data_and_metadata(xdata) + self.assertEqual(tuple(axis_group.rank for axis_group in annotated.descriptor.axis_groups), tuple(axis_group.rank for axis_group in round_tripped.descriptor.axis_groups)) + self.assertEqual(annotated.descriptor.shape, round_tripped.descriptor.shape) + self.assertEqual(annotated.metadata.attributes["descriptor_variant"], round_tripped.metadata.attributes["descriptor_variant"]) + numpy.testing.assert_array_equal(numpy.asarray(annotated.data), numpy.asarray(round_tripped.data)) + + def test_data_and_metadata_timezone_round_trip_through_annotated_array(self) -> None: + xdata = DataAndMetadata.new_data_and_metadata( + data=numpy.arange(6, dtype=numpy.float32).reshape(2, 3), + intensity_calibration=Calibration.Calibration(offset=7.0, scale=3.0, units="counts"), + dimensional_calibrations=( + Calibration.Calibration(offset=0.25, scale=1.5, units="u0"), + Calibration.Calibration(offset=1.25, scale=2.5, units="u1"), + ), + metadata={"descriptor_variant": "legacy-0-1-1"}, + data_descriptor=DataAndMetadata.DataDescriptor(False, 1, 1), + timestamp=datetime.datetime(2026, 7, 16, 19, 0, 0), + timezone="America/Los_Angeles", + timezone_offset="-0700", + ) + + annotated = AnnotatedArray.from_data_and_metadata(xdata) + self.assertEqual("America/Los_Angeles", getattr(annotated.metadata.created.tzinfo, "key", None)) + self.assertEqual(datetime.timedelta(hours=-7), annotated.metadata.created.utcoffset()) + + round_tripped = AnnotatedArray.to_data_and_metadata(annotated) + self.assertEqual(xdata.timestamp, round_tripped.timestamp) + self.assertEqual("America/Los_Angeles", round_tripped.timezone) + self.assertEqual("-0700", round_tripped.timezone_offset) + + + +if __name__ == "__main__": + unittest.main() + + diff --git a/pyproject.toml b/pyproject.toml index 8e78a8d..4ea2571 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,9 +21,10 @@ classifiers = [ ] urls = { "Homepage" = "https://github.com/nion-software/niondata" } dependencies = [ - "nionutils >=4.14.0,<5.0.0", - "numpy >=2.0,<3.0", - "scipy" + "nionutils >=4.14.0,<5.0.0", + "numpy >=2.0,<3.0", + "scipy", + "tzlocal" ] [tool.setuptools.packages.find] diff --git a/test-requirements.txt b/test-requirements.txt index 2b670f8..ef8963d 100755 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -8,3 +8,4 @@ git+https://github.com/nion-software/nionutils.git#egg=nionutils numpy scipy h5py +tzlocal