Merge branch 'main' of https://github.com/Kankelborg-Group/named_arrays

named_arrays/_core.py

@@ -17,8 +17,6 @@
     'named_array_like',
     'get_dtype',
     'value',
-    'unit',
-    'unit_normalized',
     'type_array',
     'broadcast_shapes',
     'shape_broadcasted',
@@ -148,31 +146,6 @@ def value(a: float | u.Quantity | AbstractArray):
     else:
         return a
 
-def unit(
-        value: float | complex | np.ndarray | u.UnitBase | u.Quantity | AbstractArray
-) -> None | u.UnitBase:
-    if isinstance(value, u.UnitBase):
-        return value
-    elif isinstance(value, u.Quantity):
-        return value.unit
-    elif isinstance(value, AbstractArray):
-        if isinstance(value, na.AbstractScalar):
-            return value.unit
-        else:
-            raise ValueError("non-scalar instances of `na.AbstractArray` may not be represented by a single unit")
-    else:
-        return None
-
-
-def unit_normalized(
-        value: float | complex | np.ndarray | u.UnitBase | u.Quantity | AbstractArray
-) -> u.UnitBase | dict[str, u.UnitBase]:
-    result = unit(value)
-    if result is None:
-        return u.dimensionless_unscaled
-    else:
-        return result
-
 
 def type_array(
         *values: bool | int | float | complex | str | np.ndarray | u.Quantity | AbstractArray,

named_arrays/_functions/function_named_array_functions.py

@@ -9,6 +9,14 @@
 ]
 
 ASARRAY_LIKE_FUNCTIONS = named_arrays._scalars.scalar_named_array_functions.ASARRAY_LIKE_FUNCTIONS
+HANDLED_FUNCTIONS = dict()
+
+def _implements(function: Callable):
+    """Register a __named_array_function__ implementation for AbstractScalarArray objects."""
+    def decorator(func):
+        HANDLED_FUNCTIONS[function] = func
+        return func
+    return decorator
 
 
 def asarray_like(
@@ -64,3 +72,13 @@ def asarray_like(
             like=like_outputs,
         ),
     )
+
+
+@_implements(na.unit)
+def unit(a: na.AbstractFunctionArray) -> None | u.UnitBase | na.AbstractArray:
+    return na.unit(a.outputs)
+
+
+@_implements(na.unit_normalized)
+def unit_normalized(a: na.AbstractFunctionArray) -> u.UnitBase | na.AbstractArray:
+    return na.unit_normalized(a.outputs)

named_arrays/_functions/functions.py

@@ -409,6 +409,9 @@ def __named_array_function__(self, func, *args, **kwargs):
         if func in function_named_array_functions.ASARRAY_LIKE_FUNCTIONS:
             return function_named_array_functions.asarray_like(func=func, *args, **kwargs)
 
+        if func in function_named_array_functions.HANDLED_FUNCTIONS:
+            return function_named_array_functions.HANDLED_FUNCTIONS[func](*args, **kwargs)
+
         return NotImplemented
 
     def pcolormesh(

named_arrays/_named_array_functions.py

@@ -15,6 +15,8 @@
     'geomspace',
     'ndim',
     'shape',
+    'unit',
+    'unit_normalized',
     'broadcast_to',
     'stack',
     'concatenate',
@@ -439,6 +441,65 @@ def shape(a: na.ArrayLike) -> dict[str, int]:
     return np.shape(a)
 
 
+def unit(a: Any) -> None | u.UnitBase | na.AbstractArray:
+    """
+    Isolate the physical units associated with the given object.
+
+    If the array has no physical units, this function returns :obj:`None`.
+
+    Parameters
+    ----------
+    a
+        object to isolate the units of
+
+    See Also
+    --------
+    :func:`unit_normalized` : version of this function that returns :obj:`astropy.units.dimensionless_unscaled`
+        instead of :obj:`None` if there is no unit associated with the given object
+    """
+    if isinstance(a, u.UnitBase):
+        return a
+    elif isinstance(a, u.Quantity):
+        return a.unit
+    elif isinstance(a, na.AbstractArray):
+        return na._named_array_function(
+            func=unit,
+            a=a,
+        )
+    else:
+        return None
+
+
+def unit_normalized(a: Any) -> u.UnitBase | na.AbstractArray:
+    """
+    Isolate the physical units associated with a given object, normalizing to dimensionless units
+    if the object does not have associated units.
+
+
+    Parameters
+    ----------
+    a
+        object to isolate the units of
+
+    See Also
+    --------
+    :func:`unit` : version of this function that returns :obj:`None` instead of
+        :obj:`astropy.units.dimensionless_unscaled` if there is no unit associated with the given object.
+
+    """
+    if isinstance(a, u.UnitBase):
+        return a
+    elif isinstance(a, u.Quantity):
+        return a.unit
+    elif isinstance(a, na.AbstractArray):
+        return na._named_array_function(
+            func=unit_normalized,
+            a=a,
+        )
+    else:
+        return u.dimensionless_unscaled
+
+
 @overload
 def broadcast_to(array: float | complex | np.ndarray | u.Quantity, shape: dict[str, int]) -> na.ScalarArray:
     ...

named_arrays/_scalars/scalar_named_array_functions.py

@@ -113,6 +113,19 @@ def arange(
     )
 
 
+@_implements(na.unit)
+def unit(a: na.AbstractScalarArray) -> None | u.UnitBase:
+    return na.unit(a.ndarray)
+
+
+@_implements(na.unit_normalized)
+def unit_normalized(a: na.AbstractScalarArray) -> u.UnitBase:
+    result = na.unit(a)
+    if result is None:
+        result = u.dimensionless_unscaled
+    return result
+
+
 def random(
         func: Callable,
         *args: float | u.Quantity | na.AbstractScalarArray,

named_arrays/_scalars/scalars.py

@@ -80,17 +80,17 @@ def dtype(self: Self) -> np.dtype:
         """
 
     @property
-    @abc.abstractmethod
-    def unit(self: Self) -> None | u.UnitBase | dict[str, None | u.UnitBase]:
+    def unit(self: Self) -> None | u.UnitBase:
         """
         Unit associated with the array.
 
         If :attr:`ndarray` is an instance of :class:`astropy.units.Quantity`, return :attr:`astropy.units.Quantity.unit`,
         otherwise return :class:`None`.
         """
+        return na.unit(self)
 
     @property
-    def unit_normalized(self: Self) -> u.UnitBase | dict[str, u.UnitBase]:
+    def unit_normalized(self: Self) -> u.UnitBase:
         """
         Similar to :attr:`unit` but returns :attr:`astropy.units.dimensionless_unscaled` if :attr:`ndarray` is not an
         instance of :class:`astropy.units.Quantity`.
@@ -865,10 +865,6 @@ def size(self: Self) -> int:
     def dtype(self: Self) -> np.dtype:
         return na.get_dtype(self.ndarray)
 
-    @property
-    def unit(self: Self) -> None | u.UnitBase:
-        return na.unit(self.ndarray)
-
     @property
     def explicit(self: Self) -> Self:
         return self
@@ -973,10 +969,6 @@ def ndarray(self: Self) -> na.QuantityLike:
     def dtype(self: Self) -> np.dtype:
         return self.explicit.dtype
 
-    @property
-    def unit(self: Self) -> None | u.Unit:
-        return self.explicit.unit
-
     def _attr_normalized(self, name: str) -> ScalarArray:
 
         attr = getattr(self, name)

named_arrays/_scalars/uncertainties/uncertainties.py

@@ -120,10 +120,6 @@ def value(self) -> UncertainScalarArray:
             distribution=na.value(self.distribution),
         )
 
-    @property
-    def unit(self: Self) -> None | u.Unit:
-        return na.unit(self.nominal)
-
     def astype(
             self,
             dtype: str | np.dtype | Type,

named_arrays/_scalars/uncertainties/uncertainties_named_array_functions.py

@@ -81,6 +81,16 @@ def asarray_like(
     )
 
 
+@_implements(na.unit)
+def unit(a: na.AbstractUncertainScalarArray) -> None | u.UnitBase:
+    return na.unit(a.nominal)
+
+
+@_implements(na.unit_normalized)
+def unit_normalized(a: na.AbstractUncertainScalarArray) -> u.UnitBase:
+    return na.unit_normalized(a.nominal)
+
+
 def random(
         func: Callable,
         *args: float | u.Quantity | na.AbstractScalar,

named_arrays/_vectors/vector_named_array_functions.py

@@ -113,6 +113,31 @@ def arange(
     return prototype.type_explicit.from_components(components)
 
 
+@_implements(na.unit)
+def unit(a: na.AbstractVectorArray) -> None | u.UnitBase | na.AbstractVectorArray:
+    components = a.components
+    components = {c: na.unit(components[c]) for c in components}
+    iter_components = iter(components)
+    component_0 = components[next(iter_components)]
+    if all(component_0 == components[c] for c in components):
+        return component_0
+    else:
+        components = {c: 1 if components[c] is None else 1 * components[c] for c in components}
+        return a.type_explicit.from_components(components,)
+
+
+@_implements(na.unit_normalized)
+def unit_normalized(a: na.AbstractVectorArray) -> u.UnitBase | na.AbstractVectorArray:
+    components = a.components
+    components = {c: na.unit_normalized(components[c]) for c in components}
+    iter_components = iter(components)
+    component_0 = components[next(iter_components)]
+    if all(component_0 == components[c] for c in components):
+        return component_0
+    else:
+        return a.type_explicit.from_components(components)
+
+
 def random(
         func: Callable,
         *args: float | u.Quantity |na.AbstractScalar | na.AbstractVectorArray,

named_arrays/tests/test_core.py

@@ -1012,6 +1012,15 @@ def test_transpose(
 
     class TestNamedArrayFunctions(abc.ABC):
 
+        def test_unit(self, array: na.AbstractArray):
+            result = na.unit(array)
+            if result is not None:
+                assert isinstance(result, (u.UnitBase, na.AbstractArray))
+
+        def test_unit_normalized(self, array: na.AbstractArray):
+            result = na.unit_normalized(array)
+            assert isinstance(result, (u.UnitBase, na.AbstractArray))
+
         @pytest.mark.parametrize(
             argnames="func",
             argvalues=[