From ad0895091110f31fc6e92dfd6c9b15ba023bdbeb Mon Sep 17 00:00:00 2001
From: David Hoese <david.hoese@ssec.wisc.edu>
Date: Thu, 21 Dec 2023 09:32:06 -0600
Subject: [PATCH] Fix 'viirs_edr' renaming two sets of dimensions to the same
 names

---
 satpy/readers/viirs_edr.py                 | 24 ++++++++++++++--------
 satpy/tests/reader_tests/test_viirs_edr.py |  1 +
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/satpy/readers/viirs_edr.py b/satpy/readers/viirs_edr.py
index b0eaf7b7ba..eaf2f53d42 100644
--- a/satpy/readers/viirs_edr.py
+++ b/satpy/readers/viirs_edr.py
@@ -62,6 +62,7 @@
 import logging
 from typing import Iterable
 
+import dask.array as da
 import xarray as xr
 
 from satpy import DataID
@@ -93,11 +94,6 @@ def __init__(self, filename, filename_info, filetype_info):
                                       "Along_Scan_750m": -1,
                                       "Along_Track_750m": row_chunks_m,
                                   })
-        if "Columns" in self.nc.dims:
-            self.nc = self.nc.rename({"Columns": "x", "Rows": "y"})
-        elif "Along_Track_375m" in self.nc.dims:
-            self.nc = self.nc.rename({"Along_Scan_375m": "x", "Along_Track_375m": "y"})
-            self.nc = self.nc.rename({"Along_Scan_750m": "x", "Along_Track_750m": "y"})
 
         # For some reason, no 'standard_name' is defined in some netCDF files, so
         # here we manually make the definitions.
@@ -134,7 +130,8 @@ def get_dataset(self, dataset_id: DataID, info: dict) -> xr.DataArray:
             # delete the coordinates here so the base reader doesn't try to
             # make a SwathDefinition
             data_arr = data_arr.reset_coords(drop=True)
-        return data_arr
+
+        return self._rename_dims(data_arr)
 
     def _mask_invalid(self, data_arr: xr.DataArray, ds_info: dict) -> xr.DataArray:
         # xarray auto mask and scale handled any fills from the file
@@ -152,6 +149,16 @@ def _decode_flag_meanings(data_arr: xr.DataArray):
             # only handle CF-standard flag meanings
             data_arr.attrs["flag_meanings"] = [flag for flag in data_arr.attrs["flag_meanings"].split(" ")]
 
+    @staticmethod
+    def _rename_dims(data_arr: xr.DataArray) -> xr.DataArray:
+        if "Columns" in data_arr.dims:
+            data_arr = data_arr.rename({"Columns": "x", "Rows": "y"})
+        if "Along_Track_375m" in data_arr.dims:
+            data_arr = data_arr.rename({"Along_Scan_375m": "x", "Along_Track_375m": "y"})
+        if "Along_Track_750m" in data_arr.dims:
+            data_arr = data_arr.rename({"Along_Scan_750m": "x", "Along_Track_750m": "y"})
+        return data_arr
+
     @property
     def start_time(self):
         """Get first date/time when observations were recorded."""
@@ -277,7 +284,7 @@ def _mask_invalid(self, data_arr: xr.DataArray, ds_info: dict) -> xr.DataArray:
             new_data_arr = new_data_arr.where(good_mask)
         return new_data_arr
 
-    def _get_veg_index_good_mask(self) -> xr.DataArray:
+    def _get_veg_index_good_mask(self) -> da.Array:
         # each mask array should be TRUE when pixels are UNACCEPTABLE
         qf1 = self.nc["QF1 Surface Reflectance"]
         has_sun_glint = (qf1 & 0b11000000) > 0
@@ -306,8 +313,7 @@ def _get_veg_index_good_mask(self) -> xr.DataArray:
         )
         # upscale from M-band resolution to I-band resolution
         bad_mask_iband_dask = bad_mask.data.repeat(2, axis=1).repeat(2, axis=0)
-        good_mask_iband = xr.DataArray(~bad_mask_iband_dask, dims=qf1.dims)
-        return good_mask_iband
+        return ~bad_mask_iband_dask
 
 
 class VIIRSLSTHandler(VIIRSJRRFileHandler):
diff --git a/satpy/tests/reader_tests/test_viirs_edr.py b/satpy/tests/reader_tests/test_viirs_edr.py
index a6932520c0..d042576d05 100644
--- a/satpy/tests/reader_tests/test_viirs_edr.py
+++ b/satpy/tests/reader_tests/test_viirs_edr.py
@@ -152,6 +152,7 @@ def _create_surf_refl_variables() -> dict[str, xr.DataArray]:
         "750m Surface Reflectance Band M1": xr.DataArray(m_data, dims=m_dims, attrs=sr_attrs),
     }
     for data_arr in data_arrs.values():
+        data_arr.encoding["chunksizes"] = data_arr.shape
         if "scale_factor" not in data_arr.attrs:
             continue
         data_arr.encoding["dtype"] = np.int16