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

imap_processing/codice/constants.py

@@ -92,6 +92,26 @@ LO_SW_SPECIES_VARIABLE_NAMES = [
     "heplus",
     "cnoplus",
 ]
+LO_SW_SOLAR_WIND_SPECIES_VARIABLE_NAMES = [
+    "hplus",
+    "heplusplus",
+    "cplus4",
+    "cplus5",
+    "cplus6",
+    "oplus5",
+    "oplus6",
+    "oplus7",
+    "oplus8",
+    "ne",
+    "mg",
+    "si",
+    "fe_loq",
+    "fe_hiq",
+]
+LO_SW_PICKUP_ION_SPECIES_VARIABLE_NAMES = [
+    "heplus",
+    "cnoplus",
+]
 LO_NSW_SPECIES_VARIABLE_NAMES = [
     "hplus",
     "heplusplus",
@@ -2259,3 +2279,44 @@ HALF_SPIN_LUT = {
     30: [116, 117, 118, 119, 120, 121],
     31: [122, 123, 124, 125, 126, 127],
 }
+
+NSW_POSITIONS = [x for x in range(3, 22)]
+SW_POSITIONS = [0]
+PUI_POSITIONS = [0, 1, 2, 22, 23]
+L2_GEOMETRIC_FACTOR = 0.013
+L2_HI_NUMBER_OF_SSD = 12.0
+
+L2_HI_SECTORED_ANGLE = np.array(
+    [
+        285.00,
+        244.11,
+        228.69,
+        225.00,
+        228.69,
+        244.11,
+        285.00,
+        325.89,
+        341.31,
+        345.00,
+        341.31,
+        325.89,
+    ]
+)
+
+HI_L2_ELEVATION_ANGLE = np.array(
+    [
+        150.0,
+        138.6,
+        115.7,
+        90.0,
+        64.3,
+        41.4,
+        30.0,
+        41.4,
+        64.3,
+        90.0,
+        115.7,
+        138.6,
+    ],
+    dtype=float,
+)

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.
 
@@ -602,7 +655,7 @@ class HiPointingSet(PointingSet):
         super().__init__(dataset, spice_reference_frame=geometry.SpiceFrame.ECLIPJ2000)
 
         # Filter out ENAs from non-selected portions of the spin.
-        if spin_phase not in ["full", "ram", "anti-ram"]:
+        if spin_phase not in ["full", "ram", "anti"]:
             raise ValueError(f"Unrecognized spin_phase value: {spin_phase}.")
         # ram only includes spin-phase interval [0, 0.5)
         # which is the first half of the spin_angle_bins
@@ -612,7 +665,7 @@ class HiPointingSet(PointingSet):
             )
         # anti-ram includes spin-phase interval [0.5, 1)
         # which is the second half of the spin_angle_bins
-        elif spin_phase == "anti-ram":
+        elif spin_phase == "anti":
             self.data = self.data.isel(
                 spin_angle_bin=slice(self.data["spin_angle_bin"].data.size // 2, None)
             )
@@ -631,16 +684,13 @@ 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 +703,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 +740,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
@@ -829,19 +876,11 @@ class AbstractSkyMap(ABC):
             )
 
         for value_key in value_keys:
-            pset_values = pointing_set.data[value_key]
-
             # 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,
+            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,10 +903,16 @@ 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
+                )
+
+                # 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,
+                    projection_indices=indices_bc.values,
                     input_valid_mask=pset_valid_mask,
                 )
                 # TODO: we may need to allow for unweighted/weighted means here by
@@ -879,7 +924,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
@@ -1120,7 +1165,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(
@@ -1267,7 +1315,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)
 
@@ -1431,7 +1479,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 +1818,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"