Merge pull request #2584 from djhoese/feature-ahi-abi-res-chunking

satpy/modifiers/angles.py

@@ -399,6 +399,7 @@ def get_cos_sza(data_arr: xr.DataArray) -> xr.DataArray:
 @cache_to_zarr_if("cache_lonlats", sanitize_args_func=_sanitize_args_with_chunks)
 def _get_valid_lonlats(area: PRGeometry, chunks: Union[int, str, tuple] = "auto") -> tuple[da.Array, da.Array]:
     with ignore_invalid_float_warnings():
+        # NOTE: This defaults to 64-bit floats due to needed precision for X/Y coordinates
         lons, lats = area.get_lonlats(chunks=chunks)
         lons = da.where(lons >= 1e30, np.nan, lons)
         lats = da.where(lats >= 1e30, np.nan, lats)
@@ -526,7 +527,7 @@ def _sunzen_corr_cos_ndarray(data: np.ndarray,
     max_sza_rad = np.deg2rad(max_sza) if max_sza is not None else max_sza
 
     # Cosine correction
-    corr = 1. / cos_zen
+    corr = (1. / cos_zen).astype(data.dtype, copy=False)
     if max_sza is not None:
         # gradually fall off for larger zenith angle
         grad_factor = (np.arccos(cos_zen) - limit_rad) / (max_sza_rad - limit_rad)
@@ -538,7 +539,11 @@ def _sunzen_corr_cos_ndarray(data: np.ndarray,
     else:
         # Use constant value (the limit) for larger zenith angles
         grad_factor = 1.
-    corr = np.where(cos_zen > limit_cos, corr, grad_factor / limit_cos)
+    corr = np.where(
+        cos_zen > limit_cos,
+        corr,
+        (grad_factor / limit_cos).astype(data.dtype, copy=False)
+    )
     # Force "night" pixels to 0 (where SZA is invalid)
     corr[np.isnan(cos_zen)] = 0
     return data * corr

satpy/readers/ahi_hsd.py

@@ -78,7 +78,7 @@
     np2str,
     unzip_file,
 )
-from satpy.utils import get_chunk_size_limit
+from satpy.utils import normalize_low_res_chunks
 
 AHI_CHANNEL_NAMES = ("1", "2", "3", "4", "5",
                      "6", "7", "8", "9", "10",
@@ -615,15 +615,18 @@ def _read_header(self, fp_):
 
         return header
 
-    def _read_data(self, fp_, header):
+    def _read_data(self, fp_, header, resolution):
         """Read data block."""
         nlines = int(header["block2"]['number_of_lines'][0])
         ncols = int(header["block2"]['number_of_columns'][0])
-        chunks = da.core.normalize_chunks("auto",
-                                          shape=(nlines, ncols),
-                                          limit=get_chunk_size_limit(),
-                                          dtype='f8',
-                                          previous_chunks=(550, 550))
+        chunks = normalize_low_res_chunks(
+            ("auto", "auto"),
+            (nlines, ncols),
+            # 1100 minimum chunk size for 500m, 550 for 1km, 225 for 2km
+            (1100, 1100),
+            (int(resolution / 500), int(resolution / 500)),
+            np.float32,
+        )
         return da.from_array(np.memmap(self.filename, offset=fp_.tell(),
                                        dtype='<u2', shape=(nlines, ncols), mode='r'),
                              chunks=chunks)
@@ -642,7 +645,7 @@ def read_band(self, key, ds_info):
         """Read the data."""
         with open(self.filename, "rb") as fp_:
             self._header = self._read_header(fp_)
-            res = self._read_data(fp_, self._header)
+            res = self._read_data(fp_, self._header, key["resolution"])
         res = self._mask_invalid(data=res, header=self._header)
         res = self.calibrate(res, key['calibration'])
 
@@ -671,7 +674,7 @@ def _get_metadata(self, key, ds_info):
             units=ds_info['units'],
             standard_name=ds_info['standard_name'],
             wavelength=ds_info['wavelength'],
-            resolution='resolution',
+            resolution=ds_info['resolution'],
             id=key,
             name=key['name'],
             platform_name=self.platform_name,
@@ -732,12 +735,12 @@ def convert_to_radiance(self, data):
             dn_gain, dn_offset = get_user_calibration_factors(self.band_name,
                                                               self.user_calibration)
 
-        data = (data * dn_gain + dn_offset)
+        data = (data * np.float32(dn_gain) + np.float32(dn_offset))
         # If using radiance correction factors from GSICS or similar, apply here
         if correction_type == 'RAD':
             user_slope, user_offset = get_user_calibration_factors(self.band_name,
                                                                    self.user_calibration)
-            data = apply_rad_correction(data, user_slope, user_offset)
+            data = apply_rad_correction(data, np.float32(user_slope), np.float32(user_offset))
         return data
 
     def _get_user_calibration_correction_type(self):
@@ -750,7 +753,7 @@ def _get_user_calibration_correction_type(self):
     def _vis_calibrate(self, data):
         """Visible channel calibration only."""
         coeff = self._header["calibration"]["coeff_rad2albedo_conversion"]
-        return (data * coeff * 100).clip(0)
+        return (data * np.float32(coeff) * 100).clip(0)
 
     def _ir_calibrate(self, data):
         """IR calibration."""

satpy/tests/reader_tests/test_ahi_hsd.py

@@ -293,11 +293,12 @@ def test_bad_calibration(self):
     def test_actual_satellite_position(self, round_actual_position, expected_result):
         """Test that rounding of the actual satellite position can be controlled."""
         with _fake_hsd_handler(fh_kwargs={"round_actual_position": round_actual_position}) as fh:
-            ds_id = make_dataid(name="B01")
+            ds_id = make_dataid(name="B01", resolution=1000)
             ds_info = {
                 "units": "%",
                 "standard_name": "some_name",
                 "wavelength": (0.1, 0.2, 0.3),
+                "resolution": 1000,
             }
             metadata = fh._get_metadata(ds_id, ds_info)
             orb_params = metadata["orbital_parameters"]
@@ -365,11 +366,20 @@ def test_read_band_from_actual_file(self, hsd_file_jp01):
         filename_info = {"segment": 1, "total_segments": 1}
         filetype_info = {"file_type": "blahB01"}
         fh = AHIHSDFileHandler(hsd_file_jp01, filename_info, filetype_info)
-        key = {"name": "B01", "calibration": "counts"}
+        key = {"name": "B01", "calibration": "counts", "resolution": 1000}
         import dask
-        with dask.config.set({"array.chunk-size": "16MiB"}):
-            data = fh.read_band(key, {"units": "%", "standard_name": "toa_bidirectional_reflectance", "wavelength": 2})
+        with dask.config.set({"array.chunk-size": "32MiB"}):
+            data = fh.read_band(
+                key,
+                {
+                    "units": "%",
+                    "standard_name": "toa_bidirectional_reflectance",
+                    "wavelength": 2,
+                    "resolution": 1000,
+                })
             assert data.chunks == ((1100,) * 10, (1100,) * 10)
+            assert data.dtype == data.compute().dtype
+            assert data.dtype == np.float32
 
     @mock.patch('satpy.readers.ahi_hsd.AHIHSDFileHandler._read_data')
     @mock.patch('satpy.readers.ahi_hsd.AHIHSDFileHandler._mask_invalid')
@@ -495,8 +505,8 @@ def setUp(self, *mocks):
                             'cali_offset_count2rad_conversion': [self.upd_cali[1]]},
         }
 
-        self.counts = da.array(np.array([[0., 1000.],
-                                         [2000., 5000.]]))
+        self.counts = da.array(np.array([[0, 1000],
+                                         [2000, 5000]], dtype=np.uint16))
         self.fh = fh
 
     def test_default_calibrate(self, *mocks):
@@ -564,7 +574,10 @@ def test_user_calibration(self):
         self.fh.user_calibration = {'B13': {'slope': 0.95,
                                             'offset': -0.1}}
         self.fh.band_name = 'B13'
-        rad = self.fh.calibrate(data=self.counts, calibration='radiance').compute()
+        rad = self.fh.calibrate(data=self.counts, calibration='radiance')
+        rad_np = rad.compute()
+        assert rad.dtype == rad_np.dtype
+        assert rad.dtype == np.float32
         rad_exp = np.array([[16.10526316, 12.21052632],
                             [8.31578947, -3.36842105]])
         self.assertTrue(np.allclose(rad, rad_exp))
@@ -574,7 +587,10 @@ def test_user_calibration(self):
                                             'offset': 15.20},
                                     'type': 'DN'}
         self.fh.band_name = 'B13'
-        rad = self.fh.calibrate(data=self.counts, calibration='radiance').compute()
+        rad = self.fh.calibrate(data=self.counts, calibration='radiance')
+        rad_np = rad.compute()
+        assert rad.dtype == rad_np.dtype
+        assert rad.dtype == np.float32
         rad_exp = np.array([[15.2, 12.],
                             [8.8, -0.8]])
         self.assertTrue(np.allclose(rad, rad_exp))

satpy/tests/test_modifiers.py

@@ -110,17 +110,23 @@ def sunz_sza():
 class TestSunZenithCorrector:
     """Test case for the zenith corrector."""
 
-    def test_basic_default_not_provided(self, sunz_ds1):
+    @pytest.mark.parametrize("as_32bit", [False, True])
+    def test_basic_default_not_provided(self, sunz_ds1, as_32bit):
         """Test default limits when SZA isn't provided."""
         from satpy.modifiers.geometry import SunZenithCorrector
+
+        if as_32bit:
+            sunz_ds1 = sunz_ds1.astype(np.float32)
         comp = SunZenithCorrector(name='sza_test', modifiers=tuple())
         res = comp((sunz_ds1,), test_attr='test')
         np.testing.assert_allclose(res.values, np.array([[22.401667, 22.31777], [22.437503, 22.353533]]))
         assert 'y' in res.coords
         assert 'x' in res.coords
         ds1 = sunz_ds1.copy().drop_vars(('y', 'x'))
         res = comp((ds1,), test_attr='test')
-        np.testing.assert_allclose(res.values, np.array([[22.401667, 22.31777], [22.437503, 22.353533]]))
+        res_np = res.compute()
+        np.testing.assert_allclose(res_np.values, np.array([[22.401667, 22.31777], [22.437503, 22.353533]]))
+        assert res.dtype == res_np.dtype
         assert 'y' not in res.coords
         assert 'x' not in res.coords