imap-processing 1.0.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

imap_processing/ena_maps/ena_maps.py

@@ -73,7 +73,7 @@ def match_coords_to_indices(
     input_object: PointingSet | AbstractSkyMap,
     output_object: PointingSet | AbstractSkyMap,
     event_et: float | None = None,
-) -> NDArray:
+) -> NDArray | xr.DataArray:
     """
     Find the output indices corresponding to each input coord between 2 spatial objects.
 
@@ -87,7 +87,7 @@ def match_coords_to_indices(
     This function always "pushes" the pixels of the input object to corresponding pixels
     in the output object's unwrapped rectangular grid or healpix tessellation;
     however, by swapping the input and output objects, one can apply the "pull" method
-    of index  matching.
+    of index matching.
 
     At present, the allowable inputs are either:
     - A PointingSet object and a SkyMap object, in either order of input/output.
@@ -97,7 +97,7 @@ def match_coords_to_indices(
     Parameters
     ----------
     input_object : PointingSet | AbstractSkyMap
-        An object containing 1D spatial pixel centers in azimuth and elevation,
+        An object containing spatial pixel centers in azimuth and elevation,
         which will be matched to 1D indices of spatial pixels in the output frame.
         Must contain the Spice frame in which the pixel centers are defined.
     output_object : PointingSet | AbstractSkyMap
@@ -114,11 +114,13 @@ def match_coords_to_indices(
 
     Returns
     -------
-    flat_indices_input_grid_output_frame : NDArray
-        1D array of pixel indices of the output object corresponding to each pixel in
-        the input object. The length of the array is equal to the number of pixels in
-        the input object, and may contain 0, 1, or multiple occurrences of the same
-        output index.
+    flat_indices_input_grid_output_frame : xr.DataArray
+        Array of pixel indices mapping each input object pixel center to a pixel
+        in the output object. The output xr.DataArray will have the same leading
+        dimension labels preserved. The shape of the output array is (..., n)
+        where ... matches the non-spatial dimensions of the input object, and n
+        is the number of spatial pixels in the input object. Output indices may
+        contain 0, 1, or multiple occurrences of the same output index.
 
     Raises
     ------
@@ -166,14 +168,14 @@ def match_coords_to_indices(
         # use ravel_multi_index to get the 1D indices of the pixels in the output frame.
         az_indices = (
             np.digitize(
-                input_obj_az_el_output_frame[:, 0],
+                input_obj_az_el_output_frame[..., 0],
                 output_object.sky_grid.az_bin_edges,
             )
             - 1
         )
         el_indices = (
             np.digitize(
-                input_obj_az_el_output_frame[:, 1],
+                input_obj_az_el_output_frame[..., 1],
                 output_object.sky_grid.el_bin_edges,
             )
             - 1
@@ -191,8 +193,8 @@ def match_coords_to_indices(
         # which directly returns the index on the output frame's Healpix tessellation.
         flat_indices_input_grid_output_frame = hp.ang2pix(
             nside=output_object.nside,
-            theta=input_obj_az_el_output_frame[:, 0],  # Lon in degrees
-            phi=input_obj_az_el_output_frame[:, 1],  # Lat in degrees
+            theta=input_obj_az_el_output_frame[..., 0],  # Lon in degrees
+            phi=input_obj_az_el_output_frame[..., 1],  # Lat in degrees
             nest=output_object.nested,
             lonlat=True,
         )
@@ -202,6 +204,14 @@ def match_coords_to_indices(
             f"Received: {output_object.tiling_type}"
         )
 
+    # Wrap the output indices in a DataArray with the same leading dimensions as
+    # the input object az_el_points to preserve broadcasting information
+    input_dims = input_obj_az_el_input_frame.dims[:-1]
+    flat_indices_input_grid_output_frame = xr.DataArray(
+        flat_indices_input_grid_output_frame,
+        dims=input_dims,
+    )
+
     return flat_indices_input_grid_output_frame
 
 
@@ -236,9 +246,10 @@ class PointingSet(ABC):
     spice_reference_frame: geometry.SpiceFrame
 
     # ======== Attributes required to be set in a subclass ========
-    # Azimuth and elevation coordinates of each spatial pixel. The ndarray should
-    # have the shape (n, 2) where n is the number of spatial pixels
-    az_el_points: np.ndarray
+    # Azimuth and elevation coordinates of each spatial pixel. Must be an
+    # xr.DataArray with dimensions (..., spatial_dim, az_el_coord) to preserve
+    # dimension labels
+    az_el_points: xr.DataArray
     # Tuple containing the names of each spatial coordinate of the xarray.Dataset
     # stored in the data attribute
     spatial_coords: tuple[str, ...]
@@ -274,7 +285,9 @@ class PointingSet(ABC):
         num_points: int
             The number of spatial pixels in the pointing set.
         """
-        return self.az_el_points.shape[0]
+        # Last dimension is az/el vector, the second to last dimension is
+        # the number of pixels.
+        return self.az_el_points.shape[-2]
 
     @property
     def epoch(self) -> int:
@@ -430,11 +443,12 @@ class RectangularPointingSet(PointingSet):
         # into shape (number of points in tiling of the sky, 2) where
         # column 0 (az_el_points[:, 0]) is the azimuth of that point and
         # column 1 (az_el_points[:, 1]) is the elevation of that point.
-        self.az_el_points = np.column_stack(
-            (
-                self.sky_grid.az_grid.ravel(),
-                self.sky_grid.el_grid.ravel(),
-            )
+        self.az_el_points = xr.DataArray(
+            np.stack(
+                (self.sky_grid.az_grid.ravel(), self.sky_grid.el_grid.ravel()),
+                axis=-1,
+            ),
+            dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value],
         )
 
 
@@ -540,8 +554,9 @@ class UltraPointingSet(HealpixPointingSet):
         # The coordinates of the healpix pixel centers are stored as a 2D array
         # of shape (num_points, 2) where column 0 is the lon/az
         # and column 1 is the lat/el.
-        self.az_el_points = np.column_stack(
-            (azimuth_pixel_center, elevation_pixel_center)
+        self.az_el_points = xr.DataArray(
+            np.stack((azimuth_pixel_center, elevation_pixel_center), axis=-1),
+            dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value],
         )
 
     @property
@@ -586,7 +601,45 @@ class UltraPointingSet(HealpixPointingSet):
         )
 
 
-class HiPointingSet(PointingSet):
+class LoHiBasePointingSet(PointingSet):
+    """
+    Base class for Lo and Hi pointing sets with HAE coordinate data.
+
+    This class provides common functionality for pointing sets that contain
+    hae_longitude and hae_latitude coordinates in the dataset.
+    """
+
+    tiling_type: SkyTilingType = SkyTilingType.RECTANGULAR
+
+    def update_az_el_points(self) -> None:
+        """
+        Update the az_el_points instance variable with new az/el coordinates.
+
+        The values store in the "hae_longitude" and "hae_latitude" variables
+        are used to construct the azimuth and elevation coordinates.
+        """
+        # Get lon/lat coordinates, squeeze the epoch dimension and stack along
+        # the spatial dimensions. xarray.stack() takes possibly multiple spatial
+        # dimensions and reshapes those into a single dimension.
+        az_stacked = (
+            self.data["hae_longitude"]
+            .squeeze("epoch")
+            .stack({CoordNames.GENERIC_PIXEL.value: self.spatial_coords})
+        )
+        el_stacked = (
+            self.data["hae_latitude"]
+            .squeeze("epoch")
+            .stack({CoordNames.GENERIC_PIXEL.value: self.spatial_coords})
+        )
+
+        # Stack lon/lat along last axis to create shape (..., 2)
+        self.az_el_points = xr.DataArray(
+            np.stack([az_stacked.values, el_stacked.values], axis=-1),
+            dims=[*az_stacked.dims, CoordNames.AZ_EL_VECTOR.value],
+        )
+
+
+class HiPointingSet(LoHiBasePointingSet):
     """
     PointingSet object specific to Hi L1C PSet data.
 
@@ -594,28 +647,21 @@ class HiPointingSet(PointingSet):
     ----------
     dataset : xarray.Dataset | str | Path
         Hi L1C pointing set data loaded in a xarray.DataArray.
-    spin_phase : str
-        Include ENAs from "full", "ram" or "anti-ram" phases of the spin.
     """
 
-    def __init__(self, dataset: xr.Dataset | str | Path, spin_phase: str):
-        super().__init__(dataset, spice_reference_frame=geometry.SpiceFrame.ECLIPJ2000)
+    def __init__(self, dataset: xr.Dataset | str | Path):
+        super().__init__(dataset, spice_reference_frame=geometry.SpiceFrame.IMAP_HAE)
+
+        self.spatial_coords = ("spin_angle_bin",)
 
-        # Filter out ENAs from non-selected portions of the spin.
-        if spin_phase not in ["full", "ram", "anti-ram"]:
-            raise ValueError(f"Unrecognized spin_phase value: {spin_phase}.")
+        # Naively generate the ram_mask variable assuming spacecraft frame
+        # binning. The ram_mask variable gets updated in the CG correction
+        # code if the CG correction is applied.
+        ram_mask = xr.zeros_like(self.data["spin_angle_bin"], dtype=bool)
         # ram only includes spin-phase interval [0, 0.5)
         # which is the first half of the spin_angle_bins
-        elif spin_phase == "ram":
-            self.data = self.data.isel(
-                spin_angle_bin=slice(0, self.data["spin_angle_bin"].data.size // 2)
-            )
-        # anti-ram includes spin-phase interval [0.5, 1)
-        # which is the second half of the spin_angle_bins
-        elif spin_phase == "anti-ram":
-            self.data = self.data.isel(
-                spin_angle_bin=slice(self.data["spin_angle_bin"].data.size // 2, None)
-            )
+        ram_mask[slice(0, self.data["spin_angle_bin"].data.size // 2)] = True
+        self.data["ram_mask"] = ram_mask
 
         # Rename some PSET vars to match L2 variables
         self.data = self.data.rename(
@@ -631,16 +677,11 @@ class HiPointingSet(PointingSet):
             self.data["exposure_factor"], self.data["epoch"].values[0]
         )
 
-        self.az_el_points = np.column_stack(
-            (
-                np.squeeze(self.data["hae_longitude"]),
-                np.squeeze(self.data["hae_latitude"]),
-            )
-        )
-        self.spatial_coords = ("spin_angle_bin",)
+        # Update az_el_points using the base class method
+        self.update_az_el_points()
 
 
-class LoPointingSet(PointingSet):
+class LoPointingSet(LoHiBasePointingSet):
     """
     PointingSet object specific to Lo L1C PSet data.
 
@@ -653,15 +694,11 @@ class LoPointingSet(PointingSet):
     def __init__(self, dataset: xr.Dataset):
         super().__init__(dataset, spice_reference_frame=geometry.SpiceFrame.IMAP_HAE)
 
-        # The HAE centers are stored in the pset as (1, 3600, 40) arrays
-        self.az_el_points = np.column_stack(
-            (
-                np.squeeze(self.data["hae_longitude"]).values.ravel(),
-                np.squeeze(self.data["hae_latitude"]).values.ravel(),
-            )
-        )
         self.spatial_coords = ("spin_angle", "off_angle")
 
+        # Update az_el_points using the base class method
+        self.update_az_el_points()
+
 
 # Define the Map classes
 class AbstractSkyMap(ABC):
@@ -694,9 +731,10 @@ class AbstractSkyMap(ABC):
     max_epoch: int
 
     # ======== Attributes required to be set in a subclass ========
-    # Azimuth and elevation coordinates of each spatial pixel. The ndarray should
-    # have the shape (n, 2) where n is the number of spatial pixels
-    az_el_points: np.ndarray
+    # Azimuth and elevation coordinates of each spatial pixel. The xarray.DataArray
+    # should have the shape (n, 2) where n is the number of spatial pixels.
+    # Always a simple numpy array for maps (no need for multi-dimensional coords).
+    az_el_points: xr.DataArray
     # Type of sky tiling
     tiling_type: SkyTilingType
     # Dictionary of xr.DataArray objects for each non-spatial coordinate in the SkyMap
@@ -763,12 +801,12 @@ class AbstractSkyMap(ABC):
         """
         return self.az_el_points.shape[0]
 
-    def project_pset_values_to_map(
+    def project_pset_values_to_map(  # noqa: PLR0912
         self,
         pointing_set: PointingSet,
         value_keys: list[str] | None = None,
         index_match_method: IndexMatchMethod = IndexMatchMethod.PUSH,
-        pset_valid_mask: NDArray | None = None,
+        pset_valid_mask: NDArray | xr.DataArray | None = None,
     ) -> None:
         """
         Project a pointing set's values to the map grid.
@@ -790,7 +828,7 @@ class AbstractSkyMap(ABC):
         index_match_method : IndexMatchMethod, optional
             The method of index matching to use for all values.
             Default is IndexMatchMethod.PUSH.
-        pset_valid_mask : NDArray, optional
+        pset_valid_mask : xarray.DataArray or NDArray, optional
             A boolean mask of shape (number of pointing set pixels,) indicating
             which pixels in the pointing set should be considered valid for projection.
             If None, all pixels are considered valid. Default is None.
@@ -802,9 +840,9 @@ class AbstractSkyMap(ABC):
         """
         if value_keys is None:
             value_keys = list(pointing_set.data.data_vars.keys())
-        for value_key in value_keys:
-            if value_key not in pointing_set.data.data_vars:
-                raise ValueError(f"Value key {value_key} not found in pointing set.")
+
+        if missing_keys := set(value_keys) - set(pointing_set.data.data_vars):
+            raise KeyError(f"Value keys not found in pointing set: {missing_keys}")
 
         if pset_valid_mask is None:
             pset_valid_mask = np.ones(pointing_set.num_points, dtype=bool)
@@ -829,19 +867,14 @@ class AbstractSkyMap(ABC):
             )
 
         for value_key in value_keys:
-            pset_values = pointing_set.data[value_key]
+            if value_key not in pointing_set.data.data_vars:
+                raise ValueError(f"Value key {value_key} not found in pointing set.")
 
             # If multiple spatial axes present
             # (i.e (az, el) for rectangular coordinate PSET),
-            # flatten them in the values array to match the raveled indices
-            non_spatial_axes_shape = tuple(
-                size
-                for key, size in pset_values.sizes.items()
-                if key not in pointing_set.spatial_coords
-            )
-            raveled_pset_data = pset_values.data.reshape(
-                *non_spatial_axes_shape,
-                pointing_set.num_points,
+            # stack them into a single coordinate to match the raveled indices
+            raveled_pset_data = pointing_set.data[value_key].stack(
+                {CoordNames.GENERIC_PIXEL.value: pointing_set.spatial_coords}
             )
 
             if value_key not in self.data_1d.data_vars:
@@ -864,11 +897,26 @@ class AbstractSkyMap(ABC):
             if index_match_method is IndexMatchMethod.PUSH:
                 # Bin the values at the matched indices. There may be multiple
                 # pointing set pixels that correspond to the same sky map pixel.
+                # Broadcast all arrays together using xarray dimension alignment
+                data_bc, indices_bc = xr.broadcast(
+                    raveled_pset_data, matched_indices_push
+                )
+                # If the valid mask is a xr.DataArray, broadcast it to the same shape
+                if isinstance(pset_valid_mask, xr.DataArray):
+                    stacked_valid_mask = pset_valid_mask.stack(
+                        {CoordNames.GENERIC_PIXEL.value: pointing_set.spatial_coords}
+                    )
+                    pset_valid_mask_bc, _ = xr.broadcast(data_bc, stacked_valid_mask)
+                    pset_valid_mask_values = pset_valid_mask_bc.values
+                else:
+                    pset_valid_mask_values = pset_valid_mask
+
+                # Extract numpy arrays for bincount operation
                 pointing_projected_values = map_utils.bin_single_array_at_indices(
-                    value_array=raveled_pset_data,
+                    value_array=data_bc.values,
                     projection_grid_shape=self.binning_grid_shape,
-                    projection_indices=matched_indices_push,
-                    input_valid_mask=pset_valid_mask,
+                    projection_indices=indices_bc.values,
+                    input_valid_mask=pset_valid_mask_values,
                 )
                 # TODO: we may need to allow for unweighted/weighted means here by
                 # dividing pointing_projected_values by some binned weights.
@@ -879,7 +927,7 @@ class AbstractSkyMap(ABC):
                 valid_map_mask = pset_valid_mask[matched_indices_pull]
                 # We know that there will only be one value per sky map pixel,
                 # so we can use the matched indices directly
-                pointing_projected_values = raveled_pset_data[
+                pointing_projected_values = raveled_pset_data.values[
                     ..., matched_indices_pull[valid_map_mask]
                 ]
                 # TODO: we may need to allow for unweighted/weighted means here by
@@ -889,10 +937,6 @@ class AbstractSkyMap(ABC):
                 self.data_1d[value_key].values[..., valid_map_mask] += (
                     pointing_projected_values
                 )
-            else:
-                raise NotImplementedError(
-                    "Only PUSH and PULL index matching methods are supported."
-                )
 
         # TODO: The max epoch needs to include the pset duration. Right now it
         #     is just capturing the start epoch. See issue #1747
@@ -1120,7 +1164,10 @@ class RectangularSkyMap(AbstractSkyMap):
         el_points = self.sky_grid.el_grid.ravel()
 
         # Stack so axis 0 is different pixels, and axis 1 is (az, el) of the pixel
-        self.az_el_points = np.column_stack((az_points, el_points))
+        self.az_el_points = xr.DataArray(
+            np.column_stack((az_points, el_points)),
+            dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value],
+        )
 
         # Calculate solid angles of each pixel in the map grid in units of steradians
         self.solid_angle_grid = spatial_utils.build_solid_angle_map(
@@ -1218,12 +1265,13 @@ class RectangularSkyMap(AbstractSkyMap):
             coords={**self.non_spatial_coords, **self.spatial_coords},
         )
 
-    def build_cdf_dataset(
+    def build_cdf_dataset(  # noqa: PLR0912
         self,
         instrument: str,
         level: str,
         descriptor: str,
         sensor: str | None = None,
+        drop_vars_with_no_attributes: bool = True,
     ) -> xr.Dataset:
         """
         Format the data into a xarray.Dataset and add required CDF variables.
@@ -1238,6 +1286,12 @@ class RectangularSkyMap(AbstractSkyMap):
             Descriptor for filename.
         sensor : str, optional
             Sensor number "45" or "90".
+        drop_vars_with_no_attributes : bool, optional
+            Default behavior is to drop any dataset variables that don't have
+            attributes defined in the CDF attribute manager. This ensures that
+            the output CDF doesn't have any of the intermedeiate variables left
+            over from computations. Sometimes, it is useful to output the
+            intermedeiate variables. To do so, set this to False.
 
         Returns
         -------
@@ -1267,7 +1321,7 @@ class RectangularSkyMap(AbstractSkyMap):
             if ("L2" in name)
         ]
         l2_coords.append(CoordNames.TIME.value)
-        for map_coord in cdf_ds.dims.keys():
+        for map_coord in cdf_ds.dims:
             if map_coord not in l2_coords:
                 cdf_ds = cdf_ds.drop_dims(map_coord)
 
@@ -1341,13 +1395,18 @@ class RectangularSkyMap(AbstractSkyMap):
                     variable_name=name,
                     check_schema=check_schema,
                 )
-            except KeyError as e:
-                raise KeyError(
-                    f"Attributes for variable {name} not found in "
-                    f"loaded variable attributes."
-                ) from e
-
-            cdf_ds[name].attrs.update(var_attrs)
+                cdf_ds[name].attrs.update(var_attrs)
+            except KeyError:
+                if drop_vars_with_no_attributes:
+                    logger.debug(
+                        f"Dropping variable '{name}' that has no attributes defined."
+                    )
+                    cdf_ds = cdf_ds.drop_vars(name)
+                else:
+                    logger.debug(
+                        f"Variable '{name}' has no attributes defined. It will "
+                        f"be included in the output dataset with no attributes."
+                    )
 
         # Manually adjust epoch attributes
         cdf_ds["epoch"].attrs.update(
@@ -1431,7 +1490,10 @@ class HealpixSkyMap(AbstractSkyMap):
             nside=nside, ipix=np.arange(hp.nside2npix(nside)), nest=nested, lonlat=True
         )
         # Stack so axis 0 is different pixels, and axis 1 is (az, el) of the pixel
-        self.az_el_points = np.column_stack((pixel_az, pixel_el))
+        self.az_el_points = xr.DataArray(
+            np.column_stack((pixel_az, pixel_el)),
+            dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value],
+        )
 
         self.spatial_coords = {
             CoordNames.HEALPIX_INDEX.value: xr.DataArray(
@@ -1767,7 +1829,7 @@ class HealpixSkyMap(AbstractSkyMap):
                     value_array=healpix_values_array,
                     max_subdivision_depth=max_subdivision_depth,
                 )
-                for lon_lat in rect_map.az_el_points
+                for lon_lat in rect_map.az_el_points.values
             ]
 
             # Separate the best value and the recursion depth for each pixel

imap_processing/ena_maps/utils/coordinates.py

@@ -18,3 +18,8 @@ class CoordNames(Enum):
     ELEVATION_L1C = "latitude"
     AZIMUTH_L2 = "longitude"
     ELEVATION_L2 = "latitude"
+
+    # Common name for dimension along azimuth/elevation vector
+    AZ_EL_VECTOR = "az_el"
+    # Commoon name for dimension along Cartesian vector
+    CARTESIAN_VECTOR = "x_y_z"