Allow tsm masking to be passed a dataset

pygac/reader.py

@@ -583,7 +583,18 @@ def create_counts_dataset(self):
 
     def get_calibrated_channels(self):
         """Calibrate and return the channels."""
+        calibrated_ds = self.get_calibrated_dataset()
+
+        channels = calibrated_ds["channels"].data
+        with suppress(KeyError):
+            self.meta_data["calib_coeffs_version"] = calibrated_ds.attrs["calib_coeffs_version"]
+
+        return channels
+
+    def get_calibrated_dataset(self):
+        """Create and calibrate the dataset for the pass."""
         ds = self.create_counts_dataset()
+        # calibration = {"1": "mitram", "2": "mitram", "4": {"method":"noaa", "coeff_file": "myfile.json"}}
 
         calibration_entrypoints = entry_points(group="pygac.calibration")
         calibration_function = calibration_entrypoints[self.calibration_method].load()
@@ -598,19 +609,11 @@ def get_calibrated_channels(self):
         # Apply KLM/POD specific postprocessing
         self.postproc(calibrated_ds)
 
-        channels = calibrated_ds["channels"].data
-        with suppress(KeyError):
-            self.meta_data["calib_coeffs_version"] = calibrated_ds.attrs["calib_coeffs_version"]
-
-
         # Mask pixels affected by scan motor issue
         if self.is_tsm_affected():
             LOG.info("Correcting for temporary scan motor issue")
-            self.mask_tsm_pixels(channels)
-
-
-
-        return channels
+            self.mask_tsm_pixels(calibrated_ds)
+        return calibrated_ds
 
     @abstractmethod
     def get_telemetry(self):  # pragma: no cover
@@ -1186,10 +1189,11 @@ def get_miss_lines(self):
         missing = sorted(ideal.difference(set(self.scans["scan_line_number"])))
         return np.array(missing, dtype=int)
 
-    def mask_tsm_pixels(self, channels):
+    def mask_tsm_pixels(self, ds):
         """Mask pixels affected by the scan motor issue."""
-        idx = self.get_tsm_pixels(channels)
-        channels[idx] = np.nan
+        idx = self.get_tsm_pixels(ds["channels"].values)
+        # This is because fancy/pixel indexing doesn't seem to work with xarray's `where` or `loc`
+        ds["channels"].data[idx] = np.nan
 
     @abstractmethod
     def get_tsm_pixels(self, channels):  # pragma: no cover

pygac/tests/test_reader.py

@@ -29,6 +29,7 @@
 import numpy as np
 import numpy.testing
 import pytest
+import xarray as xr
 
 from pygac.gac_pod import scanline
 from pygac.gac_reader import GACReader, ReaderError
@@ -549,15 +550,20 @@ def test_mask_tsm_pixels(self, get_tsm_pixels):
         """Test masking of pixels affected by the scan motor issue."""
         get_tsm_pixels.return_value = ([0, 1], [0, 1])
         channels = np.array([[[1., 2., 3.],
-                              [1., 2., 3.]],
-                             [[1., 2., 3.],
-                              [1., 2., 3.]]])  # (lines, pixels, channels)
+                              [4., 5., 6.]],
+                             [[7., 8., 9.],
+                              [10., 11., 12.]]])  # (lines, pixels, channels)
         masked_exp = np.array([[[np.nan, np.nan, np.nan],
-                                [1., 2., 3.]],
-                               [[1., 2., 3.],
+                                [4., 5., 6.]],
+                               [[7., 8., 9.],
                                 [np.nan, np.nan, np.nan]]])
-        self.reader.mask_tsm_pixels(channels)  # masks in-place
-        numpy.testing.assert_array_equal(channels, masked_exp)
+
+        channels = xr.DataArray(channels, dims=["scan_line_index", "columns", "channel_name"],
+                                coords=dict(channel_name=["1", "2", "3a"] ))
+        ds = xr.Dataset(dict(channels=channels))
+
+        self.reader.mask_tsm_pixels(ds)  # masks in-place
+        numpy.testing.assert_array_equal(ds["channels"].values, masked_exp)
 
     def test_correct_scan_line_numbers(self):
         """Test scanline number correction."""