Modify build_cdf_dataset to automatically drop variables that don't have attributes defined. (IMAP-Science-Operations-Center#2326)

Author: tmplummer

imap_processing/ena_maps/ena_maps.py

@@ -1265,12 +1265,13 @@ def to_dataset(self) -> xr.Dataset:
             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.
@@ -1285,6 +1286,12 @@ def build_cdf_dataset(
             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
         -------
@@ -1388,13 +1395,18 @@ def build_cdf_dataset(
                     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(

imap_processing/tests/ena_maps/test_ena_maps.py

@@ -885,12 +885,20 @@ def mock_data_for_build_cdf_dataset(self):
             name="ena_intesity",
             dims=[k for k in coord_sizes.keys()][:-1],
         )
-        # Add one variable that is expected to get removed
-        mock_dataset["foo_var"] = xr.DataArray(
+        # Add one variable that is expected to get removed because it has a
+        # dimension that is not in the list of `coord_names`
+        mock_dataset["extra_dimension_var"] = xr.DataArray(
             np.ones(tuple(s for s in coord_sizes.values())),
-            name="foo_var",
+            name="extra_dimension_var",
             dims=[k for k in coord_sizes.keys()],
         )
+        # Add a variable that is expected to get removed because it has no
+        # attributes defined
+        mock_dataset["no_attr_var"] = xr.DataArray(
+            np.ones(tuple(s for s in coord_sizes.values())[:-1]),
+            name="no_attr_var",
+            dims=[k for k in coord_sizes.keys()][:-1],
+        )
         # Add required energy delta variables
         for side in ["minus", "plus"]:
             mock_dataset[f"{CoordNames.ENERGY_L2.value}_delta_{side}"] = xr.DataArray(
@@ -910,11 +918,12 @@ def test_build_cdf_dataset(self, mock_to_dataset, mock_data_for_build_cdf_datase
         skymap.min_epoch = 10
         skymap.max_epoch = 15
         cdf_dataset = skymap.build_cdf_dataset(
-            "hi", "l2", "foo_descriptor", sensor="45"
+            "hi", "l2", "foo_descriptor", sensor="45", drop_vars_with_no_attributes=True
         )
 
-        # Check that expected var gets removed
-        assert "foo_var" not in cdf_dataset
+        # Check that expected vars gets removed
+        assert "extra_dimension_var" not in cdf_dataset
+        assert "no_attr_var" not in cdf_dataset
         # Check the epoch values
         assert CoordNames.TIME.value in cdf_dataset
         assert cdf_dataset[CoordNames.TIME.value].values[0] == skymap.min_epoch
@@ -961,30 +970,15 @@ def test_build_cdf_dataset_key_error(
     ):
         """Test build_cdf_dataset raising a KeyError."""
         mock_dataset = mock_data_for_build_cdf_dataset
-        # Add ena intensity variable
-        mock_dataset["no_attrs_var"] = xr.DataArray(
-            np.ones(
-                tuple(s for s in mock_data_for_build_cdf_dataset.coords.sizes.values())[
-                    :-1
-                ]
-            ),
-            name="no_attrs_var",
-            dims=[k for k in mock_data_for_build_cdf_dataset.coords.sizes.keys()][:-1],
-        )
         mock_to_dataset.return_value = mock_dataset
 
         skymap = ena_maps.RectangularSkyMap(6, geometry.SpiceFrame.ECLIPJ2000)
         skymap.min_epoch = 10
         skymap.max_epoch = 15
-        # Test that variables with no attributes defined raise KeyError
-        with pytest.raises(
-            KeyError, match="Attributes for variable no_attrs_var not found"
-        ):
-            _ = skymap.build_cdf_dataset("hi", "l2", "foo_descriptor", sensor="45")
 
         # Test that missing energy delta variable raise KeyError
         # Test for missing energy_delta_plus
-        mock_dataset = mock_dataset.drop(["no_attrs_var", "energy_delta_plus"])
+        mock_dataset = mock_dataset.drop(["energy_delta_plus"])
         mock_to_dataset.return_value = mock_dataset
         with pytest.raises(
             KeyError,
@@ -1000,6 +994,28 @@ def test_build_cdf_dataset_key_error(
         ):
             _ = skymap.build_cdf_dataset("hi", "l2", "foo_descriptor", sensor="45")
 
+    @mock.patch("imap_processing.ena_maps.ena_maps.RectangularSkyMap.to_dataset")
+    def test_keep_vars_with_no_attributes(
+        self, mock_to_dataset, mock_data_for_build_cdf_dataset
+    ):
+        """Test that variables with no attributes are kept when desired."""
+        # Set up the mock
+        mock_to_dataset.return_value = mock_data_for_build_cdf_dataset
+
+        skymap = ena_maps.RectangularSkyMap(6, geometry.SpiceFrame.ECLIPJ2000)
+        skymap.min_epoch = 10
+        skymap.max_epoch = 15
+        cdf_dataset = skymap.build_cdf_dataset(
+            "hi",
+            "l2",
+            "foo_descriptor",
+            sensor="45",
+            drop_vars_with_no_attributes=False,
+        )
+
+        # Check that expected var was not removed
+        assert "no_attr_var" in cdf_dataset
+
 
 class TestHealpixSkyMap:
     @pytest.fixture(autouse=True)