Skip to main content
`Ctrl`+`K`

# Source code for astropy.units.structured

    # Licensed under a 3-clause BSD style license - see LICENSE.rst
    """
    This module defines structured units and quantities.
    """

    # Standard library
    import operator
    from collections.abc import Collection
    from functools import cached_property
    from typing import Self

    import numpy as np

    from .core import UNITY, Unit, UnitBase

    __all__ = ["StructuredUnit"]

    DTYPE_OBJECT = np.dtype("O")

    def _names_from_dtype(dtype):
        """Recursively extract field names from a dtype."""
        names = []
        for name in dtype.names:
            subdtype = dtype.fields[name][0].base
            if subdtype.names:
                names.append([name, _names_from_dtype(subdtype)])
            else:
                names.append(name)
        return tuple(names)

    def _normalize_names(names):
        """Recursively normalize, inferring upper level names for unadorned tuples.

        Generally, we want the field names to be organized like dtypes, as in
        ``(['pv', ('p', 'v')], 't')``.  But we automatically infer upper
        field names if the list is absent from items like ``(('p', 'v'), 't')``,
        by concatenating the names inside the tuple.
        """
        result = []
        for name in names:
            if isinstance(name, str) and len(name) > 0:
                result.append(name)
            elif (
                isinstance(name, list)
                and len(name) == 2
                and isinstance(name[0], str)
                and len(name[0]) > 0
                and isinstance(name[1], tuple)
                and len(name[1]) > 0
            ):
                result.append([name[0], _normalize_names(name[1])])
            elif isinstance(name, tuple) and len(name) > 0:
                new_tuple = _normalize_names(name)
                name = "".join([(i[0] if isinstance(i, list) else i) for i in new_tuple])
                result.append([name, new_tuple])
            else:
                raise ValueError(
                    f"invalid entry {name!r}. Should be a name, "
                    "tuple of names, or 2-element list of the "
                    "form [name, tuple of names]."
                )

        return tuple(result)

    class StructuredUnit:
        """Container for units for a structured Quantity.

        Parameters
        ----------
        units : unit-like, tuple of unit-like, or `~astropy.units.StructuredUnit`
            Tuples can be nested.  If a `~astropy.units.StructuredUnit` is passed
            in, it will be returned unchanged unless different names are requested.
        names : tuple of str, tuple or list; `~numpy.dtype`; or `~astropy.units.StructuredUnit`, optional
            Field names for the units, possibly nested. Can be inferred from a
            structured `~numpy.dtype` or another `~astropy.units.StructuredUnit`.
            For nested tuples, by default the name of the upper entry will be the
            concatenation of the names of the lower levels.  One can pass in a
            list with the upper-level name and a tuple of lower-level names to
            avoid this.  For tuples, not all levels have to be given; for any level
            not passed in, default field names of 'f0', 'f1', etc., will be used.

        Notes
        -----
        It is recommended to initialize the class indirectly, using
        `~astropy.units.Unit`.  E.g., ``u.Unit('AU,AU/day')``.

        When combined with a structured array to produce a structured
        `~astropy.units.Quantity`, array field names will take precedence.
        Generally, passing in ``names`` is needed only if the unit is used
        unattached to a `~astropy.units.Quantity` and one needs to access its
        fields.

        Examples
        --------
        Various ways to initialize a `~astropy.units.StructuredUnit`::

            >>> import astropy.units as u
            >>> su = u.Unit('(AU,AU/day),yr')
            >>> su
            Unit("((AU, AU / d), yr)")
            >>> su.field_names
            (['f0', ('f0', 'f1')], 'f1')
            >>> su['f1']
            Unit("yr")
            >>> su2 = u.StructuredUnit(((u.AU, u.AU/u.day), u.yr), names=(('p', 'v'), 't'))
            >>> su2 == su
            True
            >>> su2.field_names
            (['pv', ('p', 'v')], 't')
            >>> su3 = u.StructuredUnit((su2['pv'], u.day), names=(['p_v', ('p', 'v')], 't'))
            >>> su3.field_names
            (['p_v', ('p', 'v')], 't')
            >>> su3.keys()
            ('p_v', 't')
            >>> su3.values()
            (Unit("(AU, AU / d)"), Unit("d"))

        Structured units share most methods with regular units::

            >>> su.physical_type
            astropy.units.structured.Structure((astropy.units.structured.Structure((PhysicalType('length'), PhysicalType({'speed', 'velocity'})), dtype=[('f0', 'O'), ('f1', 'O')]), PhysicalType('time')), dtype=[('f0', 'O'), ('f1', 'O')])
            >>> su.si
            Unit("((1.49598e+11 m, 1.73146e+06 m / s), 3.15576e+07 s)")

        """

        def __new__(cls, units, names=None):
            dtype = None
            if names is not None:
                if isinstance(names, StructuredUnit):
                    dtype = names._units.dtype
                    names = names.field_names
                elif isinstance(names, np.dtype):
                    if not names.fields:
                        raise ValueError("dtype should be structured, with fields.")
                    dtype = np.dtype([(name, DTYPE_OBJECT) for name in names.names])
                    names = _names_from_dtype(names)
                else:
                    if not isinstance(names, tuple):
                        names = (names,)
                    names = _normalize_names(names)

            if not isinstance(units, tuple):
                units = Unit(units)
                if isinstance(units, StructuredUnit):
                    # Avoid constructing a new StructuredUnit if no field names
                    # are given, or if all field names are the same already anyway.
                    if names is None or units.field_names == names:
                        return units

                    # Otherwise, turn (the upper level) into a tuple, for renaming.
                    units = units.values()
                else:
                    # Single regular unit: make a tuple for iteration below.
                    units = (units,)

            if names is None:
                names = tuple(f"f{i}" for i in range(len(units)))

            elif len(units) != len(names):
                raise ValueError("lengths of units and field names must match.")

            converted = []
            for unit, name in zip(units, names):
                if isinstance(name, list):
                    # For list, the first item is the name of our level,
                    # and the second another tuple of names, i.e., we recurse.
                    unit = cls(unit, name[1])
                    name = name[0]
                else:
                    # We are at the lowest level.  Check unit.
                    unit = Unit(unit)
                    if dtype is not None and isinstance(unit, StructuredUnit):
                        raise ValueError(
                            "units do not match in depth with field "
                            "names from dtype or structured unit."
                        )

                converted.append(unit)

            self = super().__new__(cls)
            if dtype is None:
                dtype = np.dtype(
                    [
                        ((name[0] if isinstance(name, list) else name), DTYPE_OBJECT)
                        for name in names
                    ]
                )
            # Decay array to void so we can access by field name and number.
            self._units = np.array(tuple(converted), dtype)[()]
            return self

        def __getnewargs__(self):
            """When de-serializing, e.g. pickle, start with a blank structure."""
            return (), None

        @property
        def field_names(self):
            """Possibly nested tuple of the field names of the parts."""
            return tuple(
                ([name, unit.field_names] if isinstance(unit, StructuredUnit) else name)
                for name, unit in self.items()
            )

        # Allow StructuredUnit to be treated as an (ordered) mapping.
        def __len__(self):
            return len(self._units.dtype.names)

        def __getitem__(self, item):
            # Since we are based on np.void, indexing by field number works too.
            return self._units[item]

        def values(self):
            return self._units.item()

        def keys(self):
            return self._units.dtype.names

        def items(self):
            return tuple(zip(self._units.dtype.names, self._units.item()))

        def __iter__(self):
            yield from self._units.dtype.names

        # Helpers for methods below.
        def _recursively_apply(self, func, cls=None):
            """Apply func recursively.

            Parameters
            ----------
            func : callable
                Function to apply to all parts of the structured unit,
                recursing as needed.
            cls : type, optional
                If given, should be a subclass of `~numpy.void`. By default,
                will return a new `~astropy.units.StructuredUnit` instance.
            """
            results = np.void(tuple(map(func, self.values())), self._units.dtype)
            if cls is not None:
                return results.view((cls, results.dtype))

            # Short-cut; no need to interpret field names, etc.
            result = super().__new__(self.__class__)
            result._units = results
            return result

        def _recursively_get_dtype(self, value, enter_lists=True):
            """Get structured dtype according to value, using our field names.

            This is useful since ``np.array(value)`` would treat tuples as lower
            levels of the array, rather than as elements of a structured array.
            The routine does presume that the type of the first tuple is
            representative of the rest.  Used in ``get_converter``.

            For the special value of ``UNITY``, all fields are assumed to be 1.0,
            and hence this will return an all-float dtype.

            """
            if enter_lists:
                while isinstance(value, list):
                    value = value[0]
            if value is UNITY:
                value = (UNITY,) * len(self)
            elif not isinstance(value, tuple) or len(self) != len(value):
                raise ValueError(f"cannot interpret value {value} for unit {self}.")
            descr = []
            for (name, unit), part in zip(self.items(), value):
                if isinstance(unit, StructuredUnit):
                    descr.append(
                        (name, unit._recursively_get_dtype(part, enter_lists=False))
                    )
                else:
                    # Got a part associated with a regular unit. Gets its dtype.
                    # Like for Quantity, we cast integers to float.
                    part = np.array(part)
                    part_dtype = part.dtype
                    if part_dtype.kind in "iu":
                        part_dtype = np.dtype(float)
                    descr.append((name, part_dtype, part.shape))
            return np.dtype(descr)

        @property
        def si(self):
            """The `StructuredUnit` instance in SI units."""
            return self._recursively_apply(operator.attrgetter("si"))

        @property
        def cgs(self):
            """The `StructuredUnit` instance in cgs units."""
            return self._recursively_apply(operator.attrgetter("cgs"))

        # Needed to pass through Unit initializer, so might as well use it.
        @cached_property
        def _physical_type_id(self):
            return self._recursively_apply(
                operator.attrgetter("_physical_type_id"), cls=Structure
            )

        @property
        def physical_type(self):
            """Physical types of all the fields."""
            return self._recursively_apply(
                operator.attrgetter("physical_type"), cls=Structure
            )

        def decompose(self, bases: Collection[UnitBase] = ()) -> Self:
            """The `StructuredUnit` composed of only irreducible units.

            Parameters
            ----------
            bases : sequence of `~astropy.units.UnitBase`, optional
                The bases to decompose into.  When not provided,
                decomposes down to any irreducible units.  When provided,
                the decomposed result will only contain the given units.
                This will raises a `UnitsError` if it's not possible
                to do so.

            Returns
            -------
            `~astropy.units.StructuredUnit`
                With the unit for each field containing only irreducible units.
            """
            return self._recursively_apply(operator.methodcaller("decompose", bases=bases))

        def is_equivalent(self, other, equivalencies=[]):
            """`True` if all fields are equivalent to the other's fields.

            Parameters
            ----------
            other : `~astropy.units.StructuredUnit`
                The structured unit to compare with, or what can initialize one.
            equivalencies : list of tuple, optional
                A list of equivalence pairs to try if the units are not
                directly convertible.  See :ref:`unit_equivalencies`.
                The list will be applied to all fields.

            Returns
            -------
            bool
            """
            try:
                other = StructuredUnit(other)
            except Exception:
                return False

            if len(self) != len(other):
                return False

            for self_part, other_part in zip(self.values(), other.values()):
                if not self_part.is_equivalent(other_part, equivalencies=equivalencies):
                    return False

            return True

        def get_converter(self, other, equivalencies=[]):
            if not isinstance(other, type(self)):
                other = self.__class__(other, names=self)

            converters = [
                self_part.get_converter(other_part, equivalencies=equivalencies)
                for (self_part, other_part) in zip(self.values(), other.values())
            ]

            def converter(value):
                if not hasattr(value, "dtype"):
                    value = np.array(value, self._recursively_get_dtype(value))
                result = np.empty_like(value)
                for name, converter_ in zip(result.dtype.names, converters):
                    result[name] = converter_(value[name])
                # Index with empty tuple to decay array scalars to numpy void.
                return result if result.shape else result[()]

            return converter

        get_converter.__doc__ = UnitBase.get_converter.__doc__

        def to(self, other, value=np._NoValue, equivalencies=[]):
            """Return values converted to the specified unit.

            Parameters
            ----------
            other : `~astropy.units.StructuredUnit`
                The unit to convert to.  If necessary, will be converted to
                a `~astropy.units.StructuredUnit` using the dtype of ``value``.
            value : array-like, optional
                Value(s) in the current unit to be converted to the
                specified unit.  If a sequence, the first element must have
                entries of the correct type to represent all elements (i.e.,
                not have, e.g., a ``float`` where other elements have ``complex``).
                If not given, assumed to have 1. in all fields.
            equivalencies : list of tuple, optional
                A list of equivalence pairs to try if the units are not
                directly convertible.  See :ref:`unit_equivalencies`.
                This list is in addition to possible global defaults set by, e.g.,
                `set_enabled_equivalencies`.
                Use `None` to turn off all equivalencies.

            Returns
            -------
            values : scalar or array
                Converted value(s).

            Raises
            ------
            UnitsError
                If units are inconsistent
            """
            if value is np._NoValue:
                # We do not have UNITY as a default, since then the docstring
                # would list 1.0 as default, yet one could not pass that in.
                value = UNITY
            return self.get_converter(other, equivalencies=equivalencies)(value)

        def to_string(self, format="generic"):
            """Output the unit in the given format as a string.

            Units are separated by commas.

            Parameters
            ----------
            format : `astropy.units.format.Base` subclass or str
                The name of a format or a formatter class.  If not
                provided, defaults to the generic format.

            Notes
            -----
            Structured units can be written to all formats, but can be
            re-read only with 'generic'.

            """
            parts = [part.to_string(format) for part in self.values()]
            out_fmt = "({})" if len(self) > 1 else "({},)"
            if format.startswith("latex"):
                # Strip $ from parts and add them on the outside.
                parts = [part[1:-1] for part in parts]
                out_fmt = "$" + out_fmt + "$"
            return out_fmt.format(", ".join(parts))

        def _repr_latex_(self):
            return self.to_string("latex")

        __array_ufunc__ = None

        def __mul__(self, other):
            if isinstance(other, str):
                try:
                    other = Unit(other, parse_strict="silent")
                except Exception:
                    return NotImplemented
            if isinstance(other, UnitBase):
                new_units = tuple(part * other for part in self.values())
                return self.__class__(new_units, names=self)
            if isinstance(other, StructuredUnit):
                return NotImplemented

            # Anything not like a unit, try initialising as a structured quantity.
            try:
                from .quantity import Quantity

                return Quantity(other, unit=self)
            except Exception:
                return NotImplemented

        def __rmul__(self, other):
            return self.__mul__(other)

        def __truediv__(self, other):
            if isinstance(other, str):
                try:
                    other = Unit(other, parse_strict="silent")
                except Exception:
                    return NotImplemented

            if isinstance(other, UnitBase):
                new_units = tuple(part / other for part in self.values())
                return self.__class__(new_units, names=self)
            return NotImplemented

        def __rlshift__(self, m):
            try:
                from .quantity import Quantity

                return Quantity(m, self, copy=False, subok=True)
            except Exception:
                return NotImplemented

        def __str__(self):
            return self.to_string()

        def __repr__(self):
            return f'Unit("{self.to_string()}")'

        def __hash__(self):
            return hash(self.values())

        def __eq__(self, other):
            try:
                other = StructuredUnit(other)
            except Exception:
                return NotImplemented

            return self.values() == other.values()

        def __ne__(self, other):
            if not isinstance(other, type(self)):
                try:
                    other = StructuredUnit(other)
                except Exception:
                    return NotImplemented

            return self.values() != other.values()

    class Structure(np.void):
        """Single element structure for physical type IDs, etc.

        Behaves like a `~numpy.void` and thus mostly like a tuple which can also
        be indexed with field names, but overrides ``__eq__`` and ``__ne__`` to
        compare only the contents, not the field names.  Furthermore, this way no
        `FutureWarning` about comparisons is given.

        """

        # Note that it is important for physical type IDs to not be stored in a
        # tuple, since then the physical types would be treated as alternatives in
        # :meth:`~astropy.units.UnitBase.is_equivalent`.  (Of course, in that
        # case, they could also not be indexed by name.)

        def __eq__(self, other):
            if isinstance(other, np.void):
                other = other.item()

            return self.item() == other

        def __ne__(self, other):
            if isinstance(other, np.void):
                other = other.item()

            return self.item() != other

    def _structured_unit_like_dtype(
        unit: UnitBase | StructuredUnit, dtype: np.dtype
    ) -> StructuredUnit:
        """Make a `StructuredUnit` of one unit, with the structure of a `numpy.dtype`.

        Parameters
        ----------
        unit : UnitBase
            The unit that will be filled into the structure.
        dtype : `numpy.dtype`
            The structure for the StructuredUnit.

        Returns
        -------
        StructuredUnit
        """
        if isinstance(unit, StructuredUnit):
            # If unit is structured, it should match the dtype. This function is
            # only used in Quantity, which performs this check, so it's fine to
            # return as is.
            return unit

        # Make a structured unit
        units = []
        for name in dtype.names:
            subdtype = dtype.fields[name][0]
            if subdtype.names is not None:
                units.append(_structured_unit_like_dtype(unit, subdtype))
            else:
                units.append(unit)
        return StructuredUnit(tuple(units), names=dtype.names)