Added named_arrays.WcsSpectralDirectionalVectorArray (#25)

* Fixed bug in `named_arrays.AbstractMatrixArray.__array_matmul__()` where multiplying a matrix by a column vector was throwing an error.

* Added `named_arrays.WcsSpectralDirectionalVectorArray` class.

named_arrays/_matrices/matrices.py

@@ -302,9 +302,12 @@ def __array_matmul__(
                     result = NotImplemented
 
             elif isinstance(x2, na.AbstractVectorArray):
-                if x1.type_vector().type_abstract == x2.type_abstract:
+                if x1.row_prototype.type_abstract == x2.type_abstract:
                     result_components = {r: components_x1[r] @ x2.cartesian_nd for r in components_x1}
-                    result = x2.type_explicit.from_cartesian_nd(na.CartesianNdVectorArray(result_components), like=x2)
+                    result = x1.type_vector.from_cartesian_nd(
+                        array=na.CartesianNdVectorArray(result_components),
+                        like=x1.column_prototype,
+                    )
 
                 else:
                     result = NotImplemented

named_arrays/_vectors/tests/test_vectors.py

@@ -658,3 +658,33 @@ class AbstractTestAbstractVectorGeometricSpace(
     named_arrays.tests.test_core.AbstractTestAbstractGeometricSpace,
 ):
     pass
+
+
+class AbstractTestAbstractWcsVector(
+    AbstractTestAbstractImplicitVectorArray,
+):
+    def test_crval(self, array: na.AbstractWcsVector):
+        result = array.crval
+        assert isinstance(result, array.type_abstract)
+
+    def test_crpix(self, array: na.AbstractWcsVector):
+        result = array.crpix
+        assert isinstance(result, na.AbstractCartesianNdVectorArray)
+
+    def test_cdelt(self, array: na.AbstractWcsVector):
+        result = array.cdelt
+        assert isinstance(result, array.type_abstract)
+
+    def test_pc(self, array: na.AbstractWcsVector):
+        result = array.pc
+        assert isinstance(result, array.type_matrix)
+        components = result.components
+        for c in components:
+            assert isinstance(components[c], na.AbstractVectorArray)
+
+    def test_shape_wcs(self, array: na.AbstractWcsVector):
+        result = array.shape_wcs
+        assert isinstance(result, dict)
+        for k in result:
+            assert isinstance(k, str)
+            assert isinstance(result[k], int)

named_arrays/_vectors/tests/test_vectors_spectral_directional.py

@@ -2,6 +2,7 @@
 import numpy as np
 import astropy.units as u
 import named_arrays as na
+from ..tests import test_vectors
 from ..cartesian.tests import test_vectors_cartesian
 
 _num_x = test_vectors_cartesian._num_x
@@ -187,3 +188,38 @@ class TestSpectralDirectionalVectorLinearSpace(
     test_vectors_cartesian.AbstractTestAbstractCartesianVectorLinearSpace,
 ):
     pass
+
+
+@pytest.mark.parametrize(
+    argnames="array",
+    argvalues=[
+        na.WcsSpectralDirectionalVectorArray(
+            crval=na.SpectralDirectionalVectorArray(
+                wavelength=500 * u.nm,
+                direction=na.Cartesian2dVectorArray(1, 1) * u.deg,
+            ),
+            crpix=na.CartesianNdVectorArray(dict(
+                wavelength=1,
+                x=2,
+                y=3,
+            )),
+            cdelt=na.SpectralDirectionalVectorArray(
+                wavelength=1 * u.nm,
+                direction=na.Cartesian2dVectorArray(1, 1) * u.arcsec,
+            ),
+            pc=na.SpectralDirectionalMatrixArray(
+                wavelength=na.CartesianNdVectorArray(dict(wavelength=1, x=0, y=0)),
+                direction=na.Cartesian2dMatrixArray(
+                    x=na.CartesianNdVectorArray(dict(wavelength=0, x=1, y=0)),
+                    y=na.CartesianNdVectorArray(dict(wavelength=0, x=0, y=1)),
+                ),
+            ),
+            shape_wcs=dict(wavelength=5, x=_num_x, y=_num_y),
+        ),
+    ],
+)
+class TestWcsSpectralDirectionalVectorArray(
+    AbstractTestAbstractImplicitSpectralDirectionalVectorArray,
+    test_vectors.AbstractTestAbstractWcsVector,
+):
+    pass

named_arrays/_vectors/vectors.py

@@ -24,6 +24,7 @@
     'AbstractVectorStratifiedRandomSpace',
     'AbstractVectorLogarithmicSpace',
     'AbstractVectorGeometricSpace',
+    'AbstractWcsVector',
 ]
 
 VectorPrototypeT = TypeVar("VectorPrototypeT", bound="AbstractVectorArray")
@@ -760,3 +761,69 @@ class AbstractVectorGeometricSpace(
     na.AbstractGeometricSpace,
 ):
     pass
+
+
+@dataclasses.dataclass(eq=False, repr=False)
+class AbstractWcsVector(
+    AbstractImplicitVectorArray,
+):
+    @property
+    @abc.abstractmethod
+    def crval(self) -> AbstractVectorArray:
+        """
+        The reference point in world coordinates
+        """
+
+    @property
+    @abc.abstractmethod
+    def crpix(self) -> na.CartesianNdVectorArray:
+        """
+        The reference point in pixel coordinates
+        """
+
+    @property
+    @abc.abstractmethod
+    def cdelt(self) -> AbstractVectorArray:
+        """
+        The plate scale at the reference point
+        """
+
+    @property
+    @abc.abstractmethod
+    def pc(self) -> na.AbstractMatrixArray:
+        """
+        The transformation matrix between pixel coordinates and
+        world coordinates
+        """
+
+    @property
+    @abc.abstractmethod
+    def shape_wcs(self) -> dict[str, int]:
+        """
+        The shape of the WCS components of the vector
+        """
+
+    @property
+    @abc.abstractmethod
+    def _components_explicit(self) -> dict[str, na.ArrayLike]:
+        """
+        The components of this vector that are not specified by the WCS parameters
+        """
+
+    @property
+    def _components_wcs(self):
+        crval = self.crval
+        r = self.crpix
+        s = self.cdelt
+        m = self.pc
+        shape_wcs = self.shape_wcs
+        p = na.CartesianNdVectorArray(na.indices(shape_wcs)) - 0.5
+        q = m @ (p - r)
+        x = s * q + crval
+        return x.components
+
+    @property
+    def explicit(self) -> na.AbstractExplicitArray:
+        components = self._components_explicit | self._components_wcs
+        return self.type_explicit.from_components(components)
+

named_arrays/_vectors/vectors_spectral_directional.py

@@ -11,6 +11,7 @@
     'AbstractParameterizedSpectralDirectionalVectorArray',
     'AbstractSpectralDirectionalVectorSpace',
     'SpectralDirectionalVectorLinearSpace',
+    'WcsSpectralDirectionalVectorArray',
 ]
 
 DirectionT = TypeVar("DirectionT", bound=na.ArrayLike)
@@ -83,3 +84,19 @@ class SpectralDirectionalVectorLinearSpace(
     na.AbstractVectorLinearSpace,
 ):
     pass
+
+
+@dataclasses.dataclass(eq=False, repr=False)
+class WcsSpectralDirectionalVectorArray(
+    AbstractImplicitSpectralDirectionalVectorArray,
+    na.AbstractWcsVector,
+):
+    crval: AbstractSpectralDirectionalVectorArray = dataclasses.MISSING
+    crpix: na.AbstractCartesianNdVectorArray = dataclasses.MISSING
+    cdelt: AbstractSpectralDirectionalVectorArray = dataclasses.MISSING
+    pc: na.AbstractSpectralDirectionalMatrixArray = dataclasses.MISSING
+    shape_wcs: dict[str, int] = dataclasses.MISSING
+
+    @property
+    def _components_explicit(self) -> dict[str, na.ArrayLike]:
+        return dict()

named_arrays/tests/test_core.py

@@ -714,6 +714,10 @@ def test_copyto(self, array: na.AbstractArray):
             ]
         )
         def test_broadcast_to(self, array: na.AbstractArray, shape: dict[str, int]):
+            if not set(array.shape).issubset(shape):
+                with pytest.raises(ValueError):
+                    np.broadcast_to(array, shape=shape)
+                return
             result = np.broadcast_to(array, shape=shape)
             assert result.shape == shape
 
@@ -730,6 +734,12 @@ def test_shape(self, array: na.AbstractArray):
         )
         def test_transpose(self, array: na.AbstractArray, axes: None | Sequence[str]):
             axes_normalized = tuple(reversed(array.axes) if axes is None else axes)
+
+            if not set(array.axes).issubset(axes_normalized):
+                with pytest.raises(ValueError):
+                    np.transpose(array, axes=axes)
+                return
+
             result = np.transpose(
                 a=array,
                 axes=axes
@@ -937,6 +947,10 @@ def test_broadcast_to(
             array: na.AbstractArray,
             shape: dict[str, int],
     ):
+        if not set(array.shape).issubset(shape):
+            with pytest.raises(ValueError):
+                array.broadcast_to(shape)
+            return
         assert np.array_equal(array.broadcast_to(shape), np.broadcast_to(array, shape))
 
     @pytest.mark.parametrize('shape', [dict(r=-1)])