ScalarArray to xr.DataArray Conversions (#113)

Co-authored-by: Roy Smart <roytsmart@gmail.com>

Author: jacobdparker

docs/conf.py

@@ -122,6 +122,7 @@
     'ndfilters': ('https://ndfilters.readthedocs.io/en/stable', None),
     'colorsynth': ('https://colorsynth.readthedocs.io/en/stable', None),
     'regridding': ('https://regridding.readthedocs.io/en/stable', None),
+    'xarray': ('https://docs.xarray.dev/en/stable/', None),
 }
 
 # plt.Axes.__module__ = matplotlib.axes.__name__

named_arrays/_scalars/scalars.py

@@ -8,6 +8,7 @@
 import scipy.ndimage
 import astropy.units as u
 import named_arrays as na
+import xarray as xr
 
 __all__ = [
     'ScalarStartT',
@@ -166,6 +167,12 @@ def ndarray(self: Self) -> bool | int | float | complex | str | np.ndarray | u.Q
         This is usually an instance of :class:`numpy.ndarray` or :class:`astropy.units.Quantity`, but it can also be a
         built-in python type such as a :class:`int`, :class:`float`, or :class:`bool`
         """
+    @property
+    def to_xarray(self: Self) -> xr.DataArray:
+        """
+        Cast :class:`ScalarArray` to a :class:`xarray.DataArray`.  Useful for saving to netcdf/zarr or using Dask.
+        """
+        return xr.DataArray(self.ndarray, dims=self.axes)
 
     @property
     def value(self: Self) -> ScalarArray:
@@ -819,6 +826,13 @@ def from_scalar_array(
 
         return self
 
+    @classmethod
+    def from_xarray(cls, dataarray: xr.DataArray) -> Self:
+        """
+        Convert :class:`xarray.DataArray` into a :class:`ScalarArray`
+        """
+        return cls(dataarray.data, axes=dataarray.dims)
+
     @classmethod
     def empty(cls: Type[Self], shape: dict[str, int], dtype: Type | np.dtype = float) -> Self:
         """

named_arrays/_scalars/tests/test_scalars.py

@@ -342,6 +342,17 @@ class AbstractTestAbstractScalarArray(
     def test_ndarray(self, array: na.AbstractScalarArray):
         assert isinstance(array.ndarray, (int, float, complex, str, np.ndarray))
 
+    def test_to_xarray(self, array: na.AbstractScalarArray):
+        xarray = array.to_xarray
+        assert np.all(xarray.data == array.ndarray)
+        assert set(xarray.dims) == set(array.axes)
+
+    def test_from_xarray(self, array: na.AbstractScalarArray):
+        xarray = array.to_xarray
+        scalar_array = na.ScalarArray.from_xarray(xarray)
+        assert np.all(scalar_array.ndarray == array.ndarray)
+        assert set(scalar_array.axes) == set(array.axes)
+
     def test_axes(self, array: na.AbstractScalarArray):
         super().test_axes(array)
         assert len(array.axes) == np.ndim(array.ndarray)

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "ndfilters==0.3.0",
     "colorsynth==0.1.5",
     "regridding==0.1.1",
+    "xarray",
 ]
 dynamic = ["version"]