#121 Merge pull request from deshima-dev/astropenguin/issue120

Use cube class of DEMS in making cube

.devcontainer/devcontainer.json

@@ -11,19 +11,20 @@
             "extensions": [
                 "github.vscode-pull-request-github",
                 "mhutchie.git-graph",
+                "ms-python.black-formatter",
                 "ms-python.python",
                 "streetsidesoftware.code-spell-checker",
                 "tamasfe.even-better-toml"
             ],
             "settings": {
                 "files.insertFinalNewline": true,
                 "files.trimTrailingWhitespace": true,
-                "python.formatting.provider": "black",
-                "python.languageServer": "Pylance",
                 "[python]": {
+                    "editor.defaultFormatter": "ms-python.black-formatter",
                     "editor.formatOnSave": true,
                     "editor.insertSpaces": true,
-                    "editor.tabSize": 4
+                    "editor.tabSize": 4,
+                    "python.languageServer": "Pylance"
                 }
             }
         }

decode/convert.py

@@ -1,4 +1,4 @@
-__all__ = ["frame", "units"]
+__all__ = ["coord_units", "frame", "units"]
 
 
 # standard library
@@ -15,6 +15,30 @@
 UnitLike = Union[Unit, str]
 
 
+def coord_units(
+    da: xr.DataArray,
+    coord_name: str,
+    new_units: UnitLike,
+    /,
+    *,
+    equivalencies: Optional[Equivalency] = None,
+) -> xr.DataArray:
+    """Convert units of a coordinate of a DataArray.
+
+    Args:
+        da: Input DataArray.
+        coord_name: Name of the coordinate to be converted.
+        new_units: Units to be converted from the current ones.
+        equivalencies: Optional Astropy equivalencies.
+
+    Returns:
+        DataArray with the units of the coordinate converted.
+
+    """
+    new_coord = units(da[coord_name], new_units, equivalencies=equivalencies)
+    return da.assign_coords({coord_name: new_coord})
+
+
 def frame(
     dems: xr.DataArray,
     new_frame: str,
@@ -57,40 +81,29 @@ def frame(
     lon_origin[:] = 0.0
     lat_origin[:] = 0.0
 
-    return dems
+    return dems.assign_coords(frame=dems.frame.copy(False, new_frame))
 
 
 def units(
-    dems: xr.DataArray,
-    coord_name: str,
+    da: xr.DataArray,
     new_units: UnitLike,
     /,
     *,
     equivalencies: Optional[Equivalency] = None,
-    inplace: bool = False,
 ) -> xr.DataArray:
-    """Convert units of a coordinate of DEMS.
+    """Convert units of a DataArray.
 
     Args:
-        dems: Target DEMS DataArray.
-        coord_name: Name of the coordinate for the conversion.
+        da: Input DataArray.
         new_units: Units to be converted from the current ones.
         equivalencies: Optional Astropy equivalencies.
-        inplace: Whether the units are converted in-place.
 
     Returns:
-        DEMS DataArray with the coordinate units converted.
+        DataArray with the units converted.
 
     """
-    if not inplace:
-        # deepcopy except for data
-        dems = dems.copy(data=dems.data)
-
-    coord = dems[coord_name]
-
-    if (units := coord.attrs.get("units")) is None:
-        return dems
+    if (units := da.attrs.get("units")) is None:
+        raise ValueError("Units must exist in DataArray attrs.")
 
-    coord.values = Quantity(coord, units).to(new_units, equivalencies)
-    coord.attrs["units"] = new_units
-    return dems
+    new_data = Quantity(da, units).to(new_units, equivalencies)
+    return da.copy(False, new_data).assign_attrs(units=str(new_units))

decode/make.py

@@ -2,90 +2,20 @@
 
 
 # standard library
-from dataclasses import dataclass
-from typing import Any, Literal, Tuple, Union
+from typing import Union
 
 
 # dependencies
 import numpy as np
 import xarray as xr
 from astropy.units import Quantity, Unit
-from xarray_dataclasses import AsDataArray, Attr, Coordof, Data
+from dems.d2 import Cube
 from . import convert
 
 
 # type hints
 QuantityLike = Union[Quantity, str]
 UnitLike = Union[Unit, str]
-Ln = Literal["lon"]
-Lt = Literal["lat"]
-Ch = Literal["chan"]
-
-
-@dataclass
-class Weight:
-    data: Data[Tuple[Ln, Lt, Ch], float]
-    long_name: Attr[str] = "Data weights"
-
-
-@dataclass
-class Lon:
-    data: Data[Ln, float]
-    long_name: Attr[str] = "Sky longitude"
-    units: Attr[str] = "deg"
-
-
-@dataclass
-class Lat:
-    data: Data[Lt, float]
-    long_name: Attr[str] = "Sky latitude"
-    units: Attr[str] = "deg"
-
-
-@dataclass
-class Chan:
-    data: Data[Ch, int]
-    long_name: Attr[str] = "Channel ID"
-
-
-@dataclass
-class Frame:
-    data: Data[Tuple[()], str]
-    long_name: Attr[str] = "Sky coordinate frame"
-
-
-@dataclass
-class D2MkidID:
-    data: Data[Ch, int]
-    long_name: Attr[str] = "[DESHIMA 2.0] MKID ID"
-
-
-@dataclass
-class D2MkidType:
-    data: Data[Ch, str]
-    long_name: Attr[str] = "[DESHIMA 2.0] MKID type"
-
-
-@dataclass
-class D2MkidFrequency:
-    data: Data[Ch, float]
-    long_name: Attr[str] = "[DESHIMA 2.0] MKID center frequency"
-    units: Attr[str] = "Hz"
-
-
-@dataclass
-class Cube(AsDataArray):
-    """Cube of DESHIMA 2.0."""
-
-    data: Data[Tuple[Ln, Lt, Ch], Any]
-    weight: Coordof[Weight] = 1.0
-    lon: Coordof[Lon] = 0.0
-    lat: Coordof[Lat] = 0.0
-    chan: Coordof[Chan] = 0
-    frame: Coordof[Frame] = "altaz"
-    d2_mkid_frequency: Coordof[D2MkidFrequency] = 0.0
-    d2_mkid_id: Coordof[D2MkidID] = 0.0
-    d2_mkid_type: Coordof[D2MkidType] = ""
 
 
 def cube(
@@ -106,8 +36,11 @@ def cube(
         Cube DataArray.
 
     """
-    dlon = Quantity(skycoord_grid).to(dems.lon.attrs["units"]).value
-    dlat = Quantity(skycoord_grid).to(dems.lat.attrs["units"]).value
+    dems = convert.coord_units(dems, "lon", "deg")
+    dems = convert.coord_units(dems, "lat", "deg")
+    dlon = Quantity(skycoord_grid).to("deg").value
+    dlat = Quantity(skycoord_grid).to("deg").value
+
     lon_min = np.floor(dems.lon.min() / dlon) * dlon
     lon_max = np.ceil(dems.lon.max() / dlon) * dlon
     lat_min = np.floor(dems.lat.min() / dlat) * dlat
@@ -123,21 +56,22 @@ def cube(
 
     dems = dems.copy(data=dems.data)
     dems.coords.update({"index": index})
-    grid = dems.groupby("index").mean("time")
+    gridded = dems.groupby("index").mean("time")
 
-    temp = np.full([n_lat * n_lon, n_chan], np.nan)
-    temp[grid.index.values] = grid.values
-    data = temp.reshape((n_lat, n_lon, n_chan)).swapaxes(0, 1)
+    data = np.full([n_lat * n_lon, n_chan], np.nan)
+    data[gridded.index.values] = gridded.values
+    data = data.reshape(n_lat, n_lon, n_chan).transpose(2, 0, 1)
 
     cube = Cube.new(
         data=data,
-        lon=lon.values,
-        lat=lat.values,
-        frame=dems.frame.values,
-        d2_mkid_id=dems.d2_mkid_id.values,
-        d2_mkid_frequency=dems.d2_mkid_frequency.values,
-        d2_mkid_type=dems.d2_mkid_type.values,
+        lat=lat,
+        lon=lon,
+        chan=dems.chan,
+        frame=dems.frame,
+        d2_mkid_id=dems.d2_mkid_id,
+        d2_mkid_frequency=dems.d2_mkid_frequency,
+        d2_mkid_type=dems.d2_mkid_type,
     )
-    cube = convert.units(cube, "lon", skycoord_units)
-    cube = convert.units(cube, "lat", skycoord_units)
+    cube = convert.coord_units(cube, "lon", skycoord_units)
+    cube = convert.coord_units(cube, "lat", skycoord_units)
     return cube

decode/plot.py

@@ -7,13 +7,10 @@
 
 # dependencies
 import xarray as xr
+from dems.d2 import MS_DIMS as DEMS_DIMS
 from matplotlib.artist import Artist
 
 
-# constants
-DEMS_DIMS = "time", "chan"
-
-
 def data(
     data: xr.DataArray,
     /,