From 476bbcb5a3cc811299a300366376ad9c2750a553 Mon Sep 17 00:00:00 2001 From: David Hoese Date: Tue, 29 Sep 2026 21:33:29 -0500 Subject: [PATCH 1/5] Fix VII interpolator producing Fortran-ordered arrays --- geotiepoints/tests/test_viiinterpolator.py | 6 ++++++ geotiepoints/viiinterpolator.py | 7 +++++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/geotiepoints/tests/test_viiinterpolator.py b/geotiepoints/tests/test_viiinterpolator.py index 37eae0c..fbfa4c7 100644 --- a/geotiepoints/tests/test_viiinterpolator.py +++ b/geotiepoints/tests/test_viiinterpolator.py @@ -175,6 +175,8 @@ def test_tie_points_interpolation(self): # Along track delta_axis_1 = [0., 2., 4., 6., 12., 14., 16., 18] self.assertTrue(np.allclose(result_valid[:, 0], delta_axis_1)) + # Consumers like pyresample's EWA resampling require C-contiguous arrays + self.assertTrue(result_valid.values.flags.c_contiguous) # Test the interpolation routine with invalid input pytest.raises(ValueError, tie_points_interpolation, @@ -193,6 +195,8 @@ def test_tie_points_geo_interpolation(self): ) self.assertTrue(np.allclose(lon, TEST_LON_1)) self.assertTrue(np.allclose(lat, TEST_LAT_1)) + self.assertTrue(lon.values.flags.c_contiguous) + self.assertTrue(lat.values.flags.c_contiguous) lon, lat = tie_points_geo_interpolation( self.longitude_over360, @@ -211,6 +215,8 @@ def test_tie_points_geo_interpolation(self): ) self.assertTrue(np.allclose(lon, TEST_LON_3)) self.assertTrue(np.allclose(lat, TEST_LAT_3)) + self.assertTrue(lon.values.flags.c_contiguous) + self.assertTrue(lat.values.flags.c_contiguous) # Test the interpolation routine with invalid input (different dimensions of the two arrays) with self.assertRaises(ValueError): diff --git a/geotiepoints/viiinterpolator.py b/geotiepoints/viiinterpolator.py index 6b4add8..78ec9ed 100644 --- a/geotiepoints/viiinterpolator.py +++ b/geotiepoints/viiinterpolator.py @@ -70,9 +70,12 @@ def tie_points_interpolation(data_on_tie_points, scan_alt_tie_points, tie_points # Interpolate using the xarray interp function twice: first across, then along the scan # (much faster than interpolating directly in the two dimensions) + # The order also matters for the memory layout: interpolating along the last dimension + # produces Fortran-ordered arrays, so doing it first gives C-contiguous final arrays + # as required by consumers like pyresample's EWA resampling. data = data.assign_coords({dim_alt: tie_grid_alt, dim_act: tie_grid_act}) - data_pixel = data.interp({dim_alt: pixel_grid_alt}, assume_sorted=True) \ - .interp({dim_act: pixel_grid_act}, assume_sorted=True).drop_vars([dim_alt, dim_act]) + data_pixel = data.interp({dim_act: pixel_grid_act}, assume_sorted=True) \ + .interp({dim_alt: pixel_grid_alt}, assume_sorted=True).drop_vars([dim_alt, dim_act]) data_on_pixel_points.append(data_pixel) From f6f238250a0dd7eb109c297af2e8b8fcb4e5bfdd Mon Sep 17 00:00:00 2001 From: David Hoese Date: Tue, 29 Sep 2026 21:39:21 -0500 Subject: [PATCH 2/5] Convert VII interpolator tests to pytest --- geotiepoints/tests/test_viiinterpolator.py | 214 +++++++-------------- 1 file changed, 73 insertions(+), 141 deletions(-) diff --git a/geotiepoints/tests/test_viiinterpolator.py b/geotiepoints/tests/test_viiinterpolator.py index fbfa4c7..e9c44f8 100644 --- a/geotiepoints/tests/test_viiinterpolator.py +++ b/geotiepoints/tests/test_viiinterpolator.py @@ -6,10 +6,10 @@ """ -import unittest import numpy as np -import xarray as xr import pytest +import xarray as xr + from geotiepoints.viiinterpolator import tie_points_interpolation, tie_points_geo_interpolation @@ -42,7 +42,7 @@ [18., 18.5, 19., 19.5, 20., 20.5]] ) -# Results of latitude/longitude interpolation on cartesian coordinates (latitude above 60 degrees) +# Results of latitude/longitude interpolation on cartesian coordinates (longitude with a 360 degrees step) TEST_LON_2 = np.array( [[-12., -11.50003808, -11., -10.50011426, -10., -9.50019052], [-10.00243991, -9.5032411, -9.00366173, -8.50454031, -8.00488578, -7.5058423], @@ -65,7 +65,7 @@ [17.98865968, 18.48759253, 18.98798235, 19.48686863, 19.98729684, 20.48613573]] ) -# Results of latitude/longitude interpolation on cartesian coordinates (longitude with a 360 degrees step) +# Results of latitude/longitude interpolation on cartesian coordinates (latitude above 60 degrees) TEST_LON_3 = np.array( [[-12., -11.50444038, -11., -10.50459822, -10., -9.50476197], [-10.07492627, -9.58101155, -9.07759836, -8.5839056, - 8.0803761, -7.58691614], @@ -89,140 +89,72 @@ ) -class TestViiInterpolator(unittest.TestCase): - """Test the vii_utils module.""" - - def setUp(self): - """Set up the test.""" - # Create the arrays for the interpolation test - # The first has a valid number of n_tie_alt points (multiple of SCAN_ALT_TIE_POINTS) - self.valid_data_for_interpolation = xr.DataArray( - np.arange( - TEST_VALID_ALT_TIE_POINTS * TEST_ACT_TIE_POINTS, - dtype=np.float64, - ).reshape(TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS), - dims=('num_tie_points_alt', 'num_tie_points_act'), - ) - # The second has an invalid number of n_tie_alt points (not multiple of SCAN_ALT_TIE_POINTS) - self.invalid_data_for_interpolation = xr.DataArray( - np.arange( - TEST_INVALID_ALT_TIE_POINTS * TEST_ACT_TIE_POINTS, - dtype=np.float64, - ).reshape(TEST_INVALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS), - dims=('num_tie_points_alt', 'num_tie_points_act'), - ) - # Then two arrays containing valid longitude and latitude data - self.longitude = xr.DataArray( - np.linspace( - -12, - 11, - num=TEST_VALID_ALT_TIE_POINTS * TEST_ACT_TIE_POINTS, - dtype=np.float64, - ).reshape(TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS), - dims=('num_tie_points_alt', 'num_tie_points_act'), - ) - self.latitude = xr.DataArray( - np.linspace( - 0, - 23, - num=TEST_VALID_ALT_TIE_POINTS * TEST_ACT_TIE_POINTS, - dtype=np.float64, - ).reshape(TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS), - dims=('num_tie_points_alt', 'num_tie_points_act'), - ) - # Then one containing latitude data above 60 degrees - self.latitude_over60 = xr.DataArray( - np.linspace( - 45, - 68, - num=TEST_VALID_ALT_TIE_POINTS * TEST_ACT_TIE_POINTS, - dtype=np.float64, - ).reshape(TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS), - dims=('num_tie_points_alt', 'num_tie_points_act'), - ) - # Then one containing longitude data with a 360 degrees step - self.longitude_over360 = xr.DataArray( - np.linspace( - -12, - 11, - num=TEST_VALID_ALT_TIE_POINTS * TEST_ACT_TIE_POINTS, - dtype=np.float64, - ).reshape(TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS) % 360., - dims=('num_tie_points_alt', 'num_tie_points_act'), - ) - - def tearDown(self): - """Tear down the test.""" - # Nothing to do - pass - - def test_tie_points_interpolation(self): - """# Test the interpolation routine with valid and invalid input.""" - # Test the interpolation routine with valid input - result_valid = tie_points_interpolation( - [self.valid_data_for_interpolation], - TEST_SCAN_ALT_TIE_POINTS, - TEST_TIE_POINTS_FACTOR - )[0] - - act_points_interp = (TEST_ACT_TIE_POINTS - 1) * TEST_TIE_POINTS_FACTOR - num_scans = TEST_VALID_ALT_TIE_POINTS // TEST_SCAN_ALT_TIE_POINTS - scan_alt_points_interp = (TEST_SCAN_ALT_TIE_POINTS - 1) * TEST_TIE_POINTS_FACTOR - - # Across the track - delta_axis_0 = [0., 0.5, 1., 1.5, 2., 2.5] - self.assertTrue(np.allclose(result_valid[0, :], delta_axis_0)) - # Along track - delta_axis_1 = [0., 2., 4., 6., 12., 14., 16., 18] - self.assertTrue(np.allclose(result_valid[:, 0], delta_axis_1)) - # Consumers like pyresample's EWA resampling require C-contiguous arrays - self.assertTrue(result_valid.values.flags.c_contiguous) - - # Test the interpolation routine with invalid input - pytest.raises(ValueError, tie_points_interpolation, - [self.invalid_data_for_interpolation], - TEST_SCAN_ALT_TIE_POINTS, - TEST_TIE_POINTS_FACTOR) - - def test_tie_points_geo_interpolation(self): - """# Test the coordinates interpolation routine with valid and invalid input.""" - # Test the interpolation routine with valid input - lon, lat = tie_points_geo_interpolation( - self.longitude, - self.latitude, - TEST_SCAN_ALT_TIE_POINTS, - TEST_TIE_POINTS_FACTOR - ) - self.assertTrue(np.allclose(lon, TEST_LON_1)) - self.assertTrue(np.allclose(lat, TEST_LAT_1)) - self.assertTrue(lon.values.flags.c_contiguous) - self.assertTrue(lat.values.flags.c_contiguous) - - lon, lat = tie_points_geo_interpolation( - self.longitude_over360, - self.latitude, - TEST_SCAN_ALT_TIE_POINTS, - TEST_TIE_POINTS_FACTOR - ) - self.assertTrue(np.allclose(lon, TEST_LON_2)) - self.assertTrue(np.allclose(lat, TEST_LAT_2)) - - lon, lat = tie_points_geo_interpolation( - self.longitude, - self.latitude_over60, - TEST_SCAN_ALT_TIE_POINTS, - TEST_TIE_POINTS_FACTOR - ) - self.assertTrue(np.allclose(lon, TEST_LON_3)) - self.assertTrue(np.allclose(lat, TEST_LAT_3)) - self.assertTrue(lon.values.flags.c_contiguous) - self.assertTrue(lat.values.flags.c_contiguous) - - # Test the interpolation routine with invalid input (different dimensions of the two arrays) - with self.assertRaises(ValueError): - tie_points_geo_interpolation( - self.longitude, - self.invalid_data_for_interpolation, - TEST_SCAN_ALT_TIE_POINTS, - TEST_TIE_POINTS_FACTOR - ) +def _tie_points_data_array(data): + """Wrap tie point values in a DataArray with the VII tie point dimensions.""" + return xr.DataArray(data, dims=('num_tie_points_alt', 'num_tie_points_act')) + + +def _arange_tie_points(n_tie_alt): + """Create tie points counting up from 0 with ``n_tie_alt`` points along the track.""" + data = np.arange(n_tie_alt * TEST_ACT_TIE_POINTS, dtype=np.float64).reshape(n_tie_alt, TEST_ACT_TIE_POINTS) + return _tie_points_data_array(data) + + +def _linspace_tie_points(start, stop): + """Create evenly spaced tie point values from ``start`` to ``stop`` for a valid number of scans.""" + data = np.linspace(start, stop, num=TEST_VALID_ALT_TIE_POINTS * TEST_ACT_TIE_POINTS, dtype=np.float64) + return data.reshape(TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS) + + +def test_tie_points_interpolation(): + """Test the interpolation routine with valid input.""" + data = _arange_tie_points(TEST_VALID_ALT_TIE_POINTS) + result = tie_points_interpolation([data], TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR)[0] + + # Across the track + np.testing.assert_allclose(result[0, :], [0., 0.5, 1., 1.5, 2., 2.5]) + # Along the track + np.testing.assert_allclose(result[:, 0], [0., 2., 4., 6., 12., 14., 16., 18.]) + # Consumers like pyresample's EWA resampling require C-contiguous arrays + assert result.values.flags.c_contiguous + + +def test_tie_points_interpolation_invalid_alt_tie_points(): + """Test that the number of tie points along the track must be a multiple of the tie points per scan.""" + data = _arange_tie_points(TEST_INVALID_ALT_TIE_POINTS) + with pytest.raises(ValueError, match="must be a multiple"): + tie_points_interpolation([data], TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) + + +@pytest.mark.parametrize( + ("longitude", "latitude", "exp_lon", "exp_lat"), + [ + pytest.param(_linspace_tie_points(-12, 11), _linspace_tie_points(0, 23), TEST_LON_1, TEST_LAT_1, + id="lonlat"), + pytest.param(_linspace_tie_points(-12, 11) % 360., _linspace_tie_points(0, 23), TEST_LON_2, TEST_LAT_2, + id="cartesian_lon_360_step"), + pytest.param(_linspace_tie_points(-12, 11), _linspace_tie_points(45, 68), TEST_LON_3, TEST_LAT_3, + id="cartesian_lat_over_60"), + ], +) +def test_tie_points_geo_interpolation(longitude, latitude, exp_lon, exp_lat): + """Test the coordinates interpolation routine in geodetic and cartesian coordinates.""" + lon, lat = tie_points_geo_interpolation( + _tie_points_data_array(longitude), + _tie_points_data_array(latitude), + TEST_SCAN_ALT_TIE_POINTS, + TEST_TIE_POINTS_FACTOR + ) + np.testing.assert_allclose(lon, exp_lon) + np.testing.assert_allclose(lat, exp_lat) + # Consumers like pyresample's EWA resampling require C-contiguous arrays + assert lon.values.flags.c_contiguous + assert lat.values.flags.c_contiguous + + +def test_tie_points_geo_interpolation_mismatched_shapes(): + """Test that longitude and latitude must have the same shape.""" + longitude = _tie_points_data_array(_linspace_tie_points(-12, 11)) + latitude = _arange_tie_points(TEST_INVALID_ALT_TIE_POINTS) + with pytest.raises(ValueError, match="don't match"): + tie_points_geo_interpolation(longitude, latitude, TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) From 5c2346a8dde6cd1d84733fcd0eaef5f46d2f0704 Mon Sep 17 00:00:00 2001 From: David Hoese Date: Wed, 30 Sep 2026 10:30:19 -0500 Subject: [PATCH 3/5] Add dask arrays to tests and revert interp dim swap --- geotiepoints/tests/test_viiinterpolator.py | 87 ++++++++++++++++------ geotiepoints/viiinterpolator.py | 32 ++++++-- 2 files changed, 88 insertions(+), 31 deletions(-) diff --git a/geotiepoints/tests/test_viiinterpolator.py b/geotiepoints/tests/test_viiinterpolator.py index e9c44f8..aa12c91 100644 --- a/geotiepoints/tests/test_viiinterpolator.py +++ b/geotiepoints/tests/test_viiinterpolator.py @@ -6,12 +6,16 @@ """ +import dask +import dask.array as da import numpy as np import pytest import xarray as xr from geotiepoints.viiinterpolator import tie_points_interpolation, tie_points_geo_interpolation +from .utils import CustomScheduler + TEST_N_SCANS = 2 TEST_TIE_POINTS_FACTOR = 2 @@ -19,6 +23,8 @@ TEST_VALID_ALT_TIE_POINTS = TEST_SCAN_ALT_TIE_POINTS * TEST_N_SCANS TEST_INVALID_ALT_TIE_POINTS = TEST_SCAN_ALT_TIE_POINTS * TEST_N_SCANS + 1 TEST_ACT_TIE_POINTS = 4 +TEST_SCAN_ALT_PIXELS = (TEST_SCAN_ALT_TIE_POINTS - 1) * TEST_TIE_POINTS_FACTOR +TEST_ACT_PIXELS = (TEST_ACT_TIE_POINTS - 1) * TEST_TIE_POINTS_FACTOR # Results of latitude/longitude interpolation with simple interpolation on coordinates TEST_LON_1 = np.array( @@ -89,15 +95,27 @@ ) -def _tie_points_data_array(data): - """Wrap tie point values in a DataArray with the VII tie point dimensions.""" +@pytest.fixture(params=[False, True], ids=["numpy", "dask"]) +def use_dask(request): + """Run a test with numpy-backed and with dask-backed tie points.""" + return request.param + + +def _tie_points_data_array(data, use_dask): + """Wrap tie point values in a DataArray with the VII tie point dimensions. + + Dask arrays are chunked along the track in whole scans like Satpy's METimage readers do. + + """ + if use_dask: + data = da.from_array(data, chunks=(TEST_SCAN_ALT_TIE_POINTS, -1)) return xr.DataArray(data, dims=('num_tie_points_alt', 'num_tie_points_act')) -def _arange_tie_points(n_tie_alt): +def _arange_tie_points(n_tie_alt, use_dask): """Create tie points counting up from 0 with ``n_tie_alt`` points along the track.""" data = np.arange(n_tie_alt * TEST_ACT_TIE_POINTS, dtype=np.float64).reshape(n_tie_alt, TEST_ACT_TIE_POINTS) - return _tie_points_data_array(data) + return _tie_points_data_array(data, use_dask) def _linspace_tie_points(start, stop): @@ -106,22 +124,41 @@ def _linspace_tie_points(start, stop): return data.reshape(TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS) -def test_tie_points_interpolation(): +def _assert_pixel_array(data_arr, use_dask): + """Check the array type, chunks, and memory layout of interpolated data. + + Satpy's METimage readers expect one chunk of pixel rows per chunk of tie point rows, each + spanning the whole width of the swath. Consumers like pyresample's EWA resampling require + C-contiguous arrays. Dask blocks are checked individually as computing the whole array would + concatenate them into a new C-contiguous array and hide the problem. + + """ + if use_dask: + assert isinstance(data_arr.data, da.Array) + assert data_arr.chunks == ((TEST_SCAN_ALT_PIXELS,) * TEST_N_SCANS, (TEST_ACT_PIXELS,)) + blocks = dask.compute(*data_arr.data.to_delayed().ravel()) + else: + assert isinstance(data_arr.data, np.ndarray) + blocks = [data_arr.data] + assert all(block.flags.c_contiguous for block in blocks) + + +def test_tie_points_interpolation(use_dask): """Test the interpolation routine with valid input.""" - data = _arange_tie_points(TEST_VALID_ALT_TIE_POINTS) - result = tie_points_interpolation([data], TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR)[0] + data = _arange_tie_points(TEST_VALID_ALT_TIE_POINTS, use_dask) + with dask.config.set(scheduler=CustomScheduler(max_computes=0)): + result = tie_points_interpolation([data], TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR)[0] + _assert_pixel_array(result, use_dask) # Across the track np.testing.assert_allclose(result[0, :], [0., 0.5, 1., 1.5, 2., 2.5]) # Along the track np.testing.assert_allclose(result[:, 0], [0., 2., 4., 6., 12., 14., 16., 18.]) - # Consumers like pyresample's EWA resampling require C-contiguous arrays - assert result.values.flags.c_contiguous -def test_tie_points_interpolation_invalid_alt_tie_points(): +def test_tie_points_interpolation_invalid_alt_tie_points(use_dask): """Test that the number of tie points along the track must be a multiple of the tie points per scan.""" - data = _arange_tie_points(TEST_INVALID_ALT_TIE_POINTS) + data = _arange_tie_points(TEST_INVALID_ALT_TIE_POINTS, use_dask) with pytest.raises(ValueError, match="must be a multiple"): tie_points_interpolation([data], TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) @@ -137,24 +174,26 @@ def test_tie_points_interpolation_invalid_alt_tie_points(): id="cartesian_lat_over_60"), ], ) -def test_tie_points_geo_interpolation(longitude, latitude, exp_lon, exp_lat): +def test_tie_points_geo_interpolation(longitude, latitude, exp_lon, exp_lat, use_dask): """Test the coordinates interpolation routine in geodetic and cartesian coordinates.""" - lon, lat = tie_points_geo_interpolation( - _tie_points_data_array(longitude), - _tie_points_data_array(latitude), - TEST_SCAN_ALT_TIE_POINTS, - TEST_TIE_POINTS_FACTOR - ) + # Choosing between geodetic and cartesian interpolation computes the latitude and longitude ranges + with dask.config.set(scheduler=CustomScheduler(max_computes=2)): + lon, lat = tie_points_geo_interpolation( + _tie_points_data_array(longitude, use_dask), + _tie_points_data_array(latitude, use_dask), + TEST_SCAN_ALT_TIE_POINTS, + TEST_TIE_POINTS_FACTOR + ) + + _assert_pixel_array(lon, use_dask) + _assert_pixel_array(lat, use_dask) np.testing.assert_allclose(lon, exp_lon) np.testing.assert_allclose(lat, exp_lat) - # Consumers like pyresample's EWA resampling require C-contiguous arrays - assert lon.values.flags.c_contiguous - assert lat.values.flags.c_contiguous -def test_tie_points_geo_interpolation_mismatched_shapes(): +def test_tie_points_geo_interpolation_mismatched_shapes(use_dask): """Test that longitude and latitude must have the same shape.""" - longitude = _tie_points_data_array(_linspace_tie_points(-12, 11)) - latitude = _arange_tie_points(TEST_INVALID_ALT_TIE_POINTS) + longitude = _tie_points_data_array(_linspace_tie_points(-12, 11), use_dask) + latitude = _arange_tie_points(TEST_INVALID_ALT_TIE_POINTS, use_dask) with pytest.raises(ValueError, match="don't match"): tie_points_geo_interpolation(longitude, latitude, TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) diff --git a/geotiepoints/viiinterpolator.py b/geotiepoints/viiinterpolator.py index 78ec9ed..57267d8 100644 --- a/geotiepoints/viiinterpolator.py +++ b/geotiepoints/viiinterpolator.py @@ -68,21 +68,21 @@ def tie_points_interpolation(data_on_tie_points, scan_alt_tie_points, tie_points if data.shape != (n_tie_alt, n_tie_act) or data.dims != (dim_alt, dim_act): raise ValueError("The dimensions of the arrays are not consistent") - # Interpolate using the xarray interp function twice: first across, then along the scan + # Interpolate using the xarray interp function twice: first along, then across the scan # (much faster than interpolating directly in the two dimensions) - # The order also matters for the memory layout: interpolating along the last dimension - # produces Fortran-ordered arrays, so doing it first gives C-contiguous final arrays - # as required by consumers like pyresample's EWA resampling. + # Interpolating across the scan last keeps dask chunks of whole pixel rows, one per chunk of + # tie point rows, but it produces Fortran-ordered arrays. Make them C-contiguous as required + # by consumers like pyresample's EWA resampling. data = data.assign_coords({dim_alt: tie_grid_alt, dim_act: tie_grid_act}) - data_pixel = data.interp({dim_act: pixel_grid_act}, assume_sorted=True) \ - .interp({dim_alt: pixel_grid_alt}, assume_sorted=True).drop_vars([dim_alt, dim_act]) + data_pixel = data.interp({dim_alt: pixel_grid_alt}, assume_sorted=True) \ + .interp({dim_act: pixel_grid_act}, assume_sorted=True).drop_vars([dim_alt, dim_act]) + data_pixel = _as_c_contiguous(data_pixel) data_on_pixel_points.append(data_pixel) return data_on_pixel_points - def tie_points_geo_interpolation(longitude, latitude, scan_alt_tie_points, tie_points_factor, lat_threshold_use_cartesian=60., @@ -179,3 +179,21 @@ def _xyz2lonlat(x_coords, y_coords, z_coords, z_threshold_use_xy=0.8): np.sign(z_coords) * (90. - np.rad2deg(np.arcsin(r / MEAN_EARTH_RADIUS))) ) return lons, lats + + +def _as_c_contiguous(data_arr): + """Get a copy of the DataArray with a C-contiguous numpy array or C-contiguous dask blocks. + + Args: + data_arr: xarray DataArray backed by a numpy or dask array. + + Returns: + xarray DataArray with the same dimensions, coordinates and attributes. + + """ + data = data_arr.data + if isinstance(data, da.Array): + data = data.map_blocks(np.ascontiguousarray, dtype=data.dtype, meta=np.array((), dtype=data.dtype)) + else: + data = np.ascontiguousarray(data) + return data_arr.copy(data=data) From 25ac2b7cf069d0a6e4e381141f05afa87f174e9b Mon Sep 17 00:00:00 2001 From: David Hoese Date: Wed, 30 Sep 2026 13:20:40 -0500 Subject: [PATCH 4/5] Rewrite VII interpolation to use map_blocks --- AGENTS.md | 12 +- geotiepoints/tests/test_viiinterpolator.py | 138 +++++++++- geotiepoints/viiinterpolator.py | 300 ++++++++++++++------- 3 files changed, 339 insertions(+), 111 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index fd5d9b0..7d32441 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -17,7 +17,7 @@ These share almost no code and are **not** interchangeable. |---|---|---| | Generic spline/grid | `interpolator.py`, `geointerpolator.py` | Two generations in one file (below) | | MODIS Cython | `modisinterpolator.py` + `_modis_interpolator.pyx`; `simple_modis_interpolator.py` + `_simple_modis_interpolator.pyx`; `_modis_utils.pyx` | Modern MODIS path | -| EPS-SG VII | `viiinterpolator.py` | **Not VIIRS.** Pure xarray, no Cython | +| EPS-SG VII | `viiinterpolator.py` | **Not VIIRS.** Pure numpy per scan (dask `map_blocks`), no Cython | | Multilinear | `multilinear.py` + `multilinear_cython.pyx` | Regular cartesian grid; unrelated to geolocation | Non-obvious details: @@ -92,8 +92,14 @@ be exact. That is why the test fixture is an ocean scene and why tests assert *g 2. **`AbstractSingleInterpolator.interpolate_dask`** (`interpolator.py`) hand-builds a dask graph (`normalize_chunks` + `tokenize` + a task dict), deliberately *not* `map_blocks`. Enabled by passing `chunks=` to `interpolate()` / `interpolate_to_shape()`; `chunks="auto"` works. -3. **`viiinterpolator`** relies entirely on `xarray.DataArray.interp()` and *requires* xarray input. It - interpolates across-track then along-track as two 1-D passes because that is much faster than 2-D. +3. **`viiinterpolator`** *requires* `DataArray` input (numpy- or dask-backed) and returns `DataArray`s + that keep the name, attrs, and non-dimension coords. It is its own `map_blocks` wrapper, not + `scanline_mapblocks`: it rechunks to whole scans (`scan_alt_tie_points` tie point rows) and the + full width, and stacks `(lon, lat)` along a new first axis the same way. Each scan is interpolated + on its own with numpy (along then across track, tie-aligned pixels copied rather than weighted), so + results are identical for any chunking; tests enforce that, plus C-contiguous blocks and one + pixel-row chunk per tie-point-row chunk (Satpy's METimage reader relies on both). + `tie_points_geo_interpolation` computes the input once, to choose geodetic vs cartesian globally. Everything else (the `__init__.py` legacy helpers, Gen-1 `Interpolator`/`GeoInterpolator`, `multilinear`) is numpy-only; `__init__.py` parallelizes with `multiprocessing.Pool` instead. diff --git a/geotiepoints/tests/test_viiinterpolator.py b/geotiepoints/tests/test_viiinterpolator.py index aa12c91..caf778d 100644 --- a/geotiepoints/tests/test_viiinterpolator.py +++ b/geotiepoints/tests/test_viiinterpolator.py @@ -118,10 +118,10 @@ def _arange_tie_points(n_tie_alt, use_dask): return _tie_points_data_array(data, use_dask) -def _linspace_tie_points(start, stop): - """Create evenly spaced tie point values from ``start`` to ``stop`` for a valid number of scans.""" - data = np.linspace(start, stop, num=TEST_VALID_ALT_TIE_POINTS * TEST_ACT_TIE_POINTS, dtype=np.float64) - return data.reshape(TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS) +def _linspace_tie_points(start, stop, n_tie_alt=TEST_VALID_ALT_TIE_POINTS): + """Create evenly spaced tie point values from ``start`` to ``stop`` with ``n_tie_alt`` points along the track.""" + data = np.linspace(start, stop, num=n_tie_alt * TEST_ACT_TIE_POINTS, dtype=np.float64) + return data.reshape(n_tie_alt, TEST_ACT_TIE_POINTS) def _assert_pixel_array(data_arr, use_dask): @@ -177,7 +177,7 @@ def test_tie_points_interpolation_invalid_alt_tie_points(use_dask): def test_tie_points_geo_interpolation(longitude, latitude, exp_lon, exp_lat, use_dask): """Test the coordinates interpolation routine in geodetic and cartesian coordinates.""" # Choosing between geodetic and cartesian interpolation computes the latitude and longitude ranges - with dask.config.set(scheduler=CustomScheduler(max_computes=2)): + with dask.config.set(scheduler=CustomScheduler(max_computes=1)): lon, lat = tie_points_geo_interpolation( _tie_points_data_array(longitude, use_dask), _tie_points_data_array(latitude, use_dask), @@ -197,3 +197,131 @@ def test_tie_points_geo_interpolation_mismatched_shapes(use_dask): latitude = _arange_tie_points(TEST_INVALID_ALT_TIE_POINTS, use_dask) with pytest.raises(ValueError, match="don't match"): tie_points_geo_interpolation(longitude, latitude, TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) + + +@pytest.mark.parametrize( + ("nan_tie_rows", "nan_tie_cols", "exp_nan_rows", "exp_nan_cols"), + [ + pytest.param([0, 1, 2], [], [0, 1, 2, 3], [], id="missing_first_scan"), + pytest.param([2], [], [3], [], id="edge_tie_row_of_first_scan"), + pytest.param([1], [], [1, 2, 3], [], id="middle_tie_row_of_first_scan"), + pytest.param([3], [], [4, 5], [], id="first_tie_row_of_second_scan"), + pytest.param([], [1], [], [1, 2, 3], id="tie_column"), + ], +) +def test_tie_points_interpolation_invalid_tie_points(nan_tie_rows, nan_tie_cols, exp_nan_rows, exp_nan_cols, + use_dask): + """Test that invalid (NaN) tie points only invalidate the pixels interpolated from them. + + Pixels of other scans and pixels coinciding with a valid tie point stay valid. + + """ + data = np.arange(TEST_VALID_ALT_TIE_POINTS * TEST_ACT_TIE_POINTS, dtype=np.float64) + data = data.reshape(TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS) + data[nan_tie_rows, :] = np.nan + data[:, nan_tie_cols] = np.nan + result = tie_points_interpolation([_tie_points_data_array(data, use_dask)], + TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR)[0] + + exp_nan = np.zeros(result.shape, dtype=bool) + exp_nan[exp_nan_rows, :] = True + exp_nan[:, exp_nan_cols] = True + np.testing.assert_array_equal(np.isnan(result.values), exp_nan) + + +def test_tie_points_interpolation_keeps_tie_point_values(use_dask): + """Test that pixels coinciding with a tie point get exactly its value.""" + data = np.random.default_rng(42).uniform(-180, 180, (TEST_VALID_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS)) + result = tie_points_interpolation([_tie_points_data_array(data, use_dask)], + TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR)[0] + + # Every tie point but the edge ones, at the end of each scan and of the swath width, coincides with a pixel + tie_rows = [scan * TEST_SCAN_ALT_TIE_POINTS + row + for scan in range(TEST_N_SCANS) for row in range(TEST_SCAN_ALT_TIE_POINTS - 1)] + pixel_rows = [scan * TEST_SCAN_ALT_PIXELS + row * TEST_TIE_POINTS_FACTOR + for scan in range(TEST_N_SCANS) for row in range(TEST_SCAN_ALT_TIE_POINTS - 1)] + np.testing.assert_array_equal(result.values[pixel_rows, ::TEST_TIE_POINTS_FACTOR], data[tie_rows, :-1]) + + +@pytest.mark.parametrize( + ("dtype", "exp_dtype"), + [(np.float64, np.float64), (np.float32, np.float32), (np.int16, np.float64)], +) +def test_tie_points_interpolation_dtype_and_metadata(dtype, exp_dtype, use_dask): + """Test that floating point tie points keep their type, integers become float64, and metadata is kept.""" + data = _arange_tie_points(TEST_VALID_ALT_TIE_POINTS, use_dask).astype(dtype) + data = data.rename("tie_data").assign_attrs(units="1").assign_coords(scalar=1.0) + result = tie_points_interpolation([data], TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR)[0] + + assert result.dtype == exp_dtype + assert result.values.dtype == exp_dtype + assert result.name == "tie_data" + assert result.attrs == {"units": "1"} + assert list(result.coords) == ["scalar"] + assert result.dims == data.dims + + +@pytest.mark.parametrize("lat_range", [(0, 23), (45, 68)], ids=["lonlat", "cartesian"]) +def test_tie_points_geo_interpolation_float32(lat_range, use_dask): + """Test that 32-bit floating point longitudes and latitudes stay 32-bit.""" + longitude = _tie_points_data_array(_linspace_tie_points(-12, 11).astype(np.float32), use_dask) + latitude = _tie_points_data_array(_linspace_tie_points(*lat_range).astype(np.float32), use_dask) + lon, lat = tie_points_geo_interpolation(longitude, latitude, TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) + + for data_arr in (lon, lat): + assert data_arr.dtype == np.float32 + assert data_arr.values.dtype == np.float32 + + +TEST_CHUNKS = [ + pytest.param((TEST_SCAN_ALT_TIE_POINTS, -1), id="one_scan"), + pytest.param((2 * TEST_SCAN_ALT_TIE_POINTS, -1), id="two_scans"), + pytest.param((-1, -1), id="single_chunk"), + pytest.param((3 * TEST_SCAN_ALT_TIE_POINTS, -1), id="uneven_scans"), + pytest.param((TEST_SCAN_ALT_TIE_POINTS - 1, -1), id="partial_scans"), + pytest.param((2 * TEST_SCAN_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS // 2), id="partial_width"), +] +TEST_N_SCANS_CHUNKED = 4 + + +def _assert_same_as_numpy(result, expected): + """Check that dask results are made of whole scan chunks across the swath width and equal the numpy ones.""" + assert isinstance(result.data, da.Array) + assert all(rows % TEST_SCAN_ALT_PIXELS == 0 for rows in result.chunks[0]) + assert result.chunks[1] == (result.shape[1],) + np.testing.assert_array_equal(result.values, expected.values) + + +@pytest.mark.parametrize("chunks", TEST_CHUNKS) +def test_tie_points_interpolation_independent_of_chunks(chunks): + """Test that dask results are the same as numpy ones, whatever the chunks of the tie points.""" + data = np.random.default_rng(42).uniform( + -180, 180, (TEST_N_SCANS_CHUNKED * TEST_SCAN_ALT_TIE_POINTS, TEST_ACT_TIE_POINTS)) + expected = tie_points_interpolation([_tie_points_data_array(data, False)], + TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR)[0] + dask_data = xr.DataArray(da.from_array(data, chunks=chunks), dims=('num_tie_points_alt', 'num_tie_points_act')) + with dask.config.set(scheduler=CustomScheduler(max_computes=0)): + result = tie_points_interpolation([dask_data], TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR)[0] + + _assert_same_as_numpy(result, expected) + + +@pytest.mark.parametrize("chunks", TEST_CHUNKS) +@pytest.mark.parametrize("lat_range", [(0, 23), (45, 68)], ids=["lonlat", "cartesian"]) +def test_tie_points_geo_interpolation_independent_of_chunks(lat_range, chunks): + """Test that dask results are the same as numpy ones, whatever the chunks of the tie points.""" + n_tie_alt = TEST_N_SCANS_CHUNKED * TEST_SCAN_ALT_TIE_POINTS + longitude = _linspace_tie_points(-12, 11, n_tie_alt) + latitude = _linspace_tie_points(*lat_range, n_tie_alt) + exp_lon, exp_lat = tie_points_geo_interpolation( + _tie_points_data_array(longitude, False), _tie_points_data_array(latitude, False), + TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) + dims = ('num_tie_points_alt', 'num_tie_points_act') + with dask.config.set(scheduler=CustomScheduler(max_computes=1)): + lon, lat = tie_points_geo_interpolation( + xr.DataArray(da.from_array(longitude, chunks=chunks), dims=dims), + xr.DataArray(da.from_array(latitude, chunks=chunks), dims=dims), + TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) + + _assert_same_as_numpy(lon, exp_lon) + _assert_same_as_numpy(lat, exp_lat) diff --git a/geotiepoints/viiinterpolator.py b/geotiepoints/viiinterpolator.py index 57267d8..6982f9a 100644 --- a/geotiepoints/viiinterpolator.py +++ b/geotiepoints/viiinterpolator.py @@ -4,17 +4,20 @@ Tiepoints are typically subsampled by a factor 8 with respect to the pixels, along and across the satellite track. Because of the bowtie effect, tiepoints at the scan edges are not continuous between neighbouring scans, therefore each scan has its own edge tiepoints in the along-track direction. -However, for computational efficiency, the edge tie points that lie outside the original data point raster -are excluded from the interpolation grid which is then carried out per granule, rather than per scan -at a (very) small geolocation accuracy cost at the swath edge (investigation to quantify this ongoing). -The interpolation functions are implemented for xarray.DataArrays as input. +The interpolation is carried out per scan, so the tie points of one scan never affect the pixels of another. +Pixels that coincide with a tie point take its value, the others are linearly interpolated between +the two neighbouring tie points, first along then across the track. +The interpolation functions are implemented for xarray.DataArrays as input, backed by numpy or dask arrays. +Dask arrays are interpolated chunk by chunk, each chunk holding whole scans of tie points across the whole +width of the swath; they are rechunked that way when needed. The results do not depend on the chunk size. This version works with vii test data V2 to be released Jan 2022 which has the data stored in alt, act (row,col) format instead of act,alt (col,row) """ -import xarray as xr +import dask import dask.array as da import numpy as np +import xarray as xr # MEAN EARTH RADIUS AS DEFINED BY IUGG MEAN_EARTH_RADIUS = 6371008.7714 # [m] @@ -22,64 +25,30 @@ def tie_points_interpolation(data_on_tie_points, scan_alt_tie_points, tie_points_factor): """Interpolate the data from the tie points to the pixel points. + The data are provided as a list of xarray DataArray objects, allowing to interpolate on several arrays at the same time; however the individual arrays must have exactly the same dimensions. + Args: data_on_tie_points: list of xarray DataArray objects containing the values defined on the tie points. scan_alt_tie_points: number of tie points along the satellite track for each scan tie_points_factor: sub-sampling factor of tie points wrt pixel points + Returns: list of xarray DataArray objects containing the interpolated values on the pixel points. + They keep the dimension names, name and attributes of the tie point arrays. + """ - # Extract the dimensions of the tie points array across and along track n_tie_alt, n_tie_act = data_on_tie_points[0].shape - dim_alt, dim_act = data_on_tie_points[0].dims - - # Check that the number of tie points along track is multiple of the number of tie points per scan - if n_tie_alt % scan_alt_tie_points != 0: - raise ValueError("The number of tie points in the along-route dimension must be a multiple of %d", - scan_alt_tie_points) - - # Compute the number of scans - n_scans = n_tie_alt // scan_alt_tie_points - - # Compute the dimensions of the pixel points array across and along track - n_pixel_act = (n_tie_act - 1) * tie_points_factor - n_pixel_alt = (n_tie_alt - 1) * tie_points_factor - - # Create the grids used for interpolation across the track - tie_grid_act = np.arange(0, n_pixel_act + 1, tie_points_factor) - pixel_grid_act = np.arange(0, n_pixel_act) + dims = data_on_tie_points[0].dims + _check_whole_scans(n_tie_alt, scan_alt_tie_points) - # Create the grids used for the interpolation along the track (must not include the spurious points between scans) - tie_grid_alt = np.arange(0, n_pixel_alt + 1, tie_points_factor) - n_pixel_alt_per_scan = (scan_alt_tie_points - 1) * tie_points_factor - pixel_grid_alt = [] - - for j_scan in range(n_scans): - start_index_scan = j_scan * scan_alt_tie_points * tie_points_factor - pixel_grid_alt.append(np.arange(start_index_scan, start_index_scan + n_pixel_alt_per_scan)) - pixel_grid_alt = np.concatenate(pixel_grid_alt) - - # Loop on all arrays data_on_pixel_points = [] for data in data_on_tie_points: - - if data.shape != (n_tie_alt, n_tie_act) or data.dims != (dim_alt, dim_act): + if data.shape != (n_tie_alt, n_tie_act) or data.dims != dims: raise ValueError("The dimensions of the arrays are not consistent") - - # Interpolate using the xarray interp function twice: first along, then across the scan - # (much faster than interpolating directly in the two dimensions) - # Interpolating across the scan last keeps dask chunks of whole pixel rows, one per chunk of - # tie point rows, but it produces Fortran-ordered arrays. Make them C-contiguous as required - # by consumers like pyresample's EWA resampling. - data = data.assign_coords({dim_alt: tie_grid_alt, dim_act: tie_grid_act}) - data_pixel = data.interp({dim_alt: pixel_grid_alt}, assume_sorted=True) \ - .interp({dim_act: pixel_grid_act}, assume_sorted=True).drop_vars([dim_alt, dim_act]) - data_pixel = _as_c_contiguous(data_pixel) - - data_on_pixel_points.append(data_pixel) - + pixels = _map_scan_blocks(_interpolate_scans, [data.data], scan_alt_tie_points, tie_points_factor) + data_on_pixel_points.append(_pixels_to_data_array(pixels, data)) return data_on_pixel_points @@ -89,15 +58,9 @@ def tie_points_geo_interpolation(longitude, latitude, z_threshold_use_xy=0.8): """Interpolate the geographical position from the tie points to the pixel points. - The longitude and latitude values are provided as xarray DataArray objects. - - Args: - data_on_tie_points: list of xarray DataArray objects containing the values defined on the tie points. - scan_alt_tie_points: number of tie points along the satellite track for each scan - tie_points_factor: sub-sampling factor of tie points wrt pixel points - - Returns: - list of xarray DataArray objects containing the interpolated values on the pixel points. + The interpolation is done on the longitude and latitude values, or on cartesian coordinates if the + tie points reach latitudes above ``lat_threshold_use_cartesian`` or span more than 180 degrees of + longitude. The choice is made once for all the tie points, which for dask arrays requires computing them. Args: longitude: xarray DataArray containing the longitude values defined on the tie points (degrees). @@ -109,32 +72,176 @@ def tie_points_geo_interpolation(longitude, latitude, Returns: two xarray DataArray objects containing the interpolated longitude and latitude values on the pixel points. + They keep the dimension names, name and attributes of the tie point arrays. """ # Check that the two arrays have the same dimensions if longitude.shape != latitude.shape: raise ValueError("The dimensions of longitude and latitude don't match") + _check_whole_scans(longitude.shape[0], scan_alt_tie_points) - # Determine if the interpolation should be done in cartesian or geodetic coordinates - to_cart = np.max(np.fabs(latitude)) > lat_threshold_use_cartesian or (np.max(longitude) - np.min(longitude)) > 180. + to_cart = _use_cartesian(longitude, latitude, lat_threshold_use_cartesian) + interp_lonlat = _map_scan_blocks(_geo_interpolate_scans, [longitude.data, latitude.data], + scan_alt_tie_points, tie_points_factor, n_outputs=2, + to_cart=to_cart, z_threshold_use_xy=z_threshold_use_xy) + return _pixels_to_data_array(interp_lonlat[0], longitude), _pixels_to_data_array(interp_lonlat[1], latitude) + + +def _check_whole_scans(n_tie_alt, scan_alt_tie_points): + """Check that the number of tie points along track is a multiple of the number of tie points per scan.""" + if n_tie_alt % scan_alt_tie_points != 0: + raise ValueError("The number of tie points in the along-route dimension must be a multiple of " + f"{scan_alt_tie_points}") - if to_cart: - x, y, z = _lonlat2xyz(longitude, latitude) +def _use_cartesian(longitude, latitude, lat_threshold_use_cartesian): + """Check if the geographical interpolation must be done on cartesian coordinates. - interp_x, interp_y, interp_z = tie_points_interpolation([x, y, z], - scan_alt_tie_points, - tie_points_factor) + That is the case at high latitudes, or when the longitudes cross the antimeridian. + Both ranges are computed together to compute dask arrays only once. + + """ + max_abs_lat, lon_range = dask.compute(abs(latitude).max(), longitude.max() - longitude.min()) + return bool(max_abs_lat > lat_threshold_use_cartesian or lon_range > 180.) + + +def _map_scan_blocks(func, tie_arrays, scan_alt_tie_points, tie_points_factor, n_outputs=None, **kwargs): + """Call ``func`` on whole scans of tie points, chunk by chunk for dask arrays. + + Args: + func: function taking the numpy tie point arrays, ``scan_alt_tie_points`` and ``tie_points_factor`` + and returning the pixel array, or ``n_outputs`` pixel arrays stacked along a new first dimension. + tie_arrays: numpy or dask arrays of the tie points, all with the same shape. + scan_alt_tie_points: number of tie points along the satellite track for each scan. + tie_points_factor: sub-sampling factor of tie points wrt pixel points. + n_outputs: number of stacked pixel arrays returned by ``func``, or None if it returns only one. + kwargs: other keyword arguments passed to ``func``. - interp_longitude, interp_latitude = _xyz2lonlat(interp_x, interp_y, interp_z, z_threshold_use_xy) + Returns: + numpy or dask array of the pixel values, stacked along the first dimension if ``n_outputs`` is given. + """ + if not any(isinstance(arr, da.Array) for arr in tie_arrays): + return func(*tie_arrays, scan_alt_tie_points, tie_points_factor, **kwargs) + + tie_arrays = _rechunk_to_whole_scans([da.asarray(arr) for arr in tie_arrays], scan_alt_tie_points) + row_chunks, col_chunks = tie_arrays[0].chunks + pixel_chunks = ( + tuple(rows // scan_alt_tie_points * (scan_alt_tie_points - 1) * tie_points_factor for rows in row_chunks), + ((col_chunks[0] - 1) * tie_points_factor,), + ) + dtype = _pixel_dtype(*tie_arrays) + if n_outputs is None: + map_kwargs = {"chunks": pixel_chunks} else: + map_kwargs = {"chunks": ((n_outputs,),) + pixel_chunks, "new_axis": 0} + return da.map_blocks(func, *tie_arrays, scan_alt_tie_points, tie_points_factor, **kwargs, + dtype=dtype, meta=np.array((), dtype=dtype), **map_kwargs) + + +def _rechunk_to_whole_scans(tie_arrays, scan_alt_tie_points): + """Rechunk dask tie point arrays to whole scans along the track and the whole swath width across it. + + Arrays already chunked this way, like the ones of Satpy's METimage readers, are returned unchanged. + + """ + row_chunks, col_chunks = tie_arrays[0].chunks + whole_scans = all(rows % scan_alt_tie_points == 0 for rows in row_chunks) and len(col_chunks) == 1 + if whole_scans and all(arr.chunks == tie_arrays[0].chunks for arr in tie_arrays[1:]): + return tie_arrays + rows = max(row_chunks[0] // scan_alt_tie_points, 1) * scan_alt_tie_points + return [arr.rechunk((rows, -1)) for arr in tie_arrays] + + +def _pixel_dtype(*tie_arrays): + """Get the data type of the pixels: floating point tie points keep theirs, integers become float64.""" + return np.result_type(*(arr.dtype for arr in tie_arrays), 1.0) - interp_longitude, interp_latitude = tie_points_interpolation([longitude, latitude], - scan_alt_tie_points, - tie_points_factor) - return interp_longitude, interp_latitude +def _pixels_to_data_array(pixels, tie_points): + """Wrap pixel values in a DataArray like the tie points, keeping only the coordinates not on their dimensions.""" + coords = {name: coord for name, coord in tie_points.coords.items() if not set(coord.dims) & set(tie_points.dims)} + return xr.DataArray(pixels, dims=tie_points.dims, coords=coords, name=tie_points.name, attrs=tie_points.attrs) + + +def _interpolate_scans(tie_points, scan_alt_tie_points, tie_points_factor, out=None): + """Interpolate whole scans of tie points to pixel points, first along then across the track. + + Pixels that coincide with a tie point get its value. The others are linearly interpolated between + the two neighbouring tie points of the same scan. Taking the tie point value, instead of giving a + zero weight to a neighbouring tie point, keeps an invalid (NaN) neighbour from invalidating the pixel. + + Args: + tie_points: numpy array of whole scans (rows) of tie points over the whole swath width (columns). + scan_alt_tie_points: number of tie points along the satellite track for each scan. + tie_points_factor: sub-sampling factor of tie points wrt pixel points. + out: C-contiguous numpy array to write the pixel values to, or None to create a new one. + + Returns: + C-contiguous numpy array of the pixel values. + + """ + dtype = _pixel_dtype(tie_points) + n_tie_alt, n_tie_act = tie_points.shape + n_scans = n_tie_alt // scan_alt_tie_points + # Position of the pixels between two tie points, excluding the one on the first tie point + weights = np.arange(1, tie_points_factor, dtype=dtype) / tie_points_factor + + # Along the track, only between the tie points of the same scan: + # (scans, tie point intervals per scan, pixels per interval, tie point columns) + scans = tie_points.astype(dtype, copy=False).reshape(n_scans, scan_alt_tie_points, 1, n_tie_act) + along = np.empty((n_scans, scan_alt_tie_points - 1, tie_points_factor, n_tie_act), dtype=dtype) + _interpolate_intervals(scans[:, :-1], scans[:, 1:], weights[:, np.newaxis], along) + along = along.reshape(-1, n_tie_act) + + # Across the track: (pixel rows, tie point intervals, pixels per interval) + if out is None: + out = np.empty((along.shape[0], (n_tie_act - 1) * tie_points_factor), dtype=dtype) + across = out.reshape(along.shape[0], n_tie_act - 1, tie_points_factor) + _interpolate_intervals(along[:, :-1, np.newaxis], along[:, 1:, np.newaxis], weights, across) + return out + + +def _interpolate_intervals(start, end, weights, out): + """Linearly interpolate tie point intervals into ``out``, whose third dimension holds the pixels of an interval. + + The first pixel of each interval gets the value of the start tie point, the others are interpolated + using ``weights``. The values are written in place to avoid temporary arrays as large as ``out``. + + """ + out[:, :, :1] = start + interpolated = out[:, :, 1:] + np.multiply(end - start, weights, out=interpolated) + interpolated += start + + +def _geo_interpolate_scans(longitude, latitude, scan_alt_tie_points, tie_points_factor, to_cart, z_threshold_use_xy): + """Interpolate whole scans of longitude and latitude tie points to pixel points. + + Args: + longitude: numpy array of longitude tie points (degrees). + latitude: numpy array of latitude tie points (degrees). + scan_alt_tie_points: number of tie points along the satellite track for each scan. + tie_points_factor: sub-sampling factor of tie points wrt pixel points. + to_cart: interpolate on cartesian coordinates instead of longitudes and latitudes. + z_threshold_use_xy: z threshold to compute latitude from x and y in cartesian coordinates. + + Returns: + numpy array of the pixel longitudes and latitudes, stacked along the first dimension. + + """ + n_tie_alt, n_tie_act = longitude.shape + shape = (2, n_tie_alt // scan_alt_tie_points * (scan_alt_tie_points - 1) * tie_points_factor, + (n_tie_act - 1) * tie_points_factor) + lonlat = np.empty(shape, dtype=_pixel_dtype(longitude, latitude)) + if to_cart: + x_coords, y_coords, z_coords = (_interpolate_scans(coords, scan_alt_tie_points, tie_points_factor) + for coords in _lonlat2xyz(longitude, latitude)) + _xyz2lonlat(x_coords, y_coords, z_coords, z_threshold_use_xy, out=lonlat) + else: + _interpolate_scans(longitude, scan_alt_tie_points, tie_points_factor, out=lonlat[0]) + _interpolate_scans(latitude, scan_alt_tie_points, tie_points_factor, out=lonlat[1]) + return lonlat def _lonlat2xyz(lons, lats): @@ -156,44 +263,31 @@ def _lonlat2xyz(lons, lats): return x_coords, y_coords, z_coords -def _xyz2lonlat(x_coords, y_coords, z_coords, z_threshold_use_xy=0.8): +def _xyz2lonlat(x_coords, y_coords, z_coords, z_threshold_use_xy, out): """Get longitudes and latitudes from cartesian coordinates. + The computations are done in place as much as possible to limit the number of temporary pixel arrays. + Args: - x_coords: array containing the x values in meters. - y_coords: array containing the y values in meters. - z_coords: array containing the z values in meters. + x_coords: numpy array containing the x values in meters. + y_coords: numpy array containing the y values in meters. + z_coords: numpy array containing the z values in meters. z_threshold_use_xy: z threshold to compute latitude from x and y in cartesian coordinates. - - Returns: - tuple of arrays containing the longitude and latitude values in degrees. + out: numpy array to write the longitude and latitude values in degrees to, stacked along the first dimension. """ + lons, lats = out r = np.sqrt(x_coords ** 2 + y_coords ** 2) - thr_z = z_threshold_use_xy * MEAN_EARTH_RADIUS - lons = np.rad2deg(np.arccos(x_coords / r)) * np.sign(y_coords) + np.divide(x_coords, r, out=lons) + np.arccos(lons, out=lons) + np.rad2deg(lons, out=lons) + lons *= np.sign(y_coords) # Compute latitude from z at low z and from x and y at high z - lats = xr.where( - np.logical_and(np.less(z_coords, thr_z), np.greater(z_coords, -thr_z)), - 90. - np.rad2deg(np.arccos(z_coords / MEAN_EARTH_RADIUS)), - np.sign(z_coords) * (90. - np.rad2deg(np.arcsin(r / MEAN_EARTH_RADIUS))) - ) - return lons, lats - - -def _as_c_contiguous(data_arr): - """Get a copy of the DataArray with a C-contiguous numpy array or C-contiguous dask blocks. - - Args: - data_arr: xarray DataArray backed by a numpy or dask array. - - Returns: - xarray DataArray with the same dimensions, coordinates and attributes. - - """ - data = data_arr.data - if isinstance(data, da.Array): - data = data.map_blocks(np.ascontiguousarray, dtype=data.dtype, meta=np.array((), dtype=data.dtype)) - else: - data = np.ascontiguousarray(data) - return data_arr.copy(data=data) + np.divide(z_coords, MEAN_EARTH_RADIUS, out=lats) + np.arccos(lats, out=lats) + np.rad2deg(lats, out=lats) + np.subtract(90., lats, out=lats) + thr_z = z_threshold_use_xy * MEAN_EARTH_RADIUS + use_xy = ~np.logical_and(np.less(z_coords, thr_z), np.greater(z_coords, -thr_z)) + if use_xy.any(): + lats[use_xy] = np.sign(z_coords[use_xy]) * (90. - np.rad2deg(np.arcsin(r[use_xy] / MEAN_EARTH_RADIUS))) From 523e6ed9ccfd6231927fe5664ee6a549722c70b2 Mon Sep 17 00:00:00 2001 From: David Hoese Date: Wed, 30 Sep 2026 13:38:18 -0500 Subject: [PATCH 5/5] Move to_cart condition into dask --- AGENTS.md | 4 +- geotiepoints/tests/test_viiinterpolator.py | 12 ++++-- geotiepoints/viiinterpolator.py | 46 ++++++++++++++-------- 3 files changed, 41 insertions(+), 21 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 7d32441..4d8e9e2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -99,7 +99,9 @@ be exact. That is why the test fixture is an ocean scene and why tests assert *g on its own with numpy (along then across track, tie-aligned pixels copied rather than weighted), so results are identical for any chunking; tests enforce that, plus C-contiguous blocks and one pixel-row chunk per tie-point-row chunk (Satpy's METimage reader relies on both). - `tie_points_geo_interpolation` computes the input once, to choose geodetic vs cartesian globally. + `tie_points_geo_interpolation` never computes: the reductions that choose geodetic vs cartesian + (max |lat|, longitude range) are passed to `map_blocks` as 0-d dask arrays, so the choice is made + once per compute, over all the tie points, and is the same for every chunk. Never decide per chunk. Everything else (the `__init__.py` legacy helpers, Gen-1 `Interpolator`/`GeoInterpolator`, `multilinear`) is numpy-only; `__init__.py` parallelizes with `multiprocessing.Pool` instead. diff --git a/geotiepoints/tests/test_viiinterpolator.py b/geotiepoints/tests/test_viiinterpolator.py index caf778d..bd679ab 100644 --- a/geotiepoints/tests/test_viiinterpolator.py +++ b/geotiepoints/tests/test_viiinterpolator.py @@ -176,8 +176,7 @@ def test_tie_points_interpolation_invalid_alt_tie_points(use_dask): ) def test_tie_points_geo_interpolation(longitude, latitude, exp_lon, exp_lat, use_dask): """Test the coordinates interpolation routine in geodetic and cartesian coordinates.""" - # Choosing between geodetic and cartesian interpolation computes the latitude and longitude ranges - with dask.config.set(scheduler=CustomScheduler(max_computes=1)): + with dask.config.set(scheduler=CustomScheduler(max_computes=0)): lon, lat = tie_points_geo_interpolation( _tie_points_data_array(longitude, use_dask), _tie_points_data_array(latitude, use_dask), @@ -309,7 +308,12 @@ def test_tie_points_interpolation_independent_of_chunks(chunks): @pytest.mark.parametrize("chunks", TEST_CHUNKS) @pytest.mark.parametrize("lat_range", [(0, 23), (45, 68)], ids=["lonlat", "cartesian"]) def test_tie_points_geo_interpolation_independent_of_chunks(lat_range, chunks): - """Test that dask results are the same as numpy ones, whatever the chunks of the tie points.""" + """Test that dask results are the same as numpy ones, whatever the chunks of the tie points. + + The choice of cartesian coordinates must be the same for every chunk too: in the cartesian case, + only the chunks of the last scans reach latitudes above 60 degrees on their own. + + """ n_tie_alt = TEST_N_SCANS_CHUNKED * TEST_SCAN_ALT_TIE_POINTS longitude = _linspace_tie_points(-12, 11, n_tie_alt) latitude = _linspace_tie_points(*lat_range, n_tie_alt) @@ -317,7 +321,7 @@ def test_tie_points_geo_interpolation_independent_of_chunks(lat_range, chunks): _tie_points_data_array(longitude, False), _tie_points_data_array(latitude, False), TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) dims = ('num_tie_points_alt', 'num_tie_points_act') - with dask.config.set(scheduler=CustomScheduler(max_computes=1)): + with dask.config.set(scheduler=CustomScheduler(max_computes=0)): lon, lat = tie_points_geo_interpolation( xr.DataArray(da.from_array(longitude, chunks=chunks), dims=dims), xr.DataArray(da.from_array(latitude, chunks=chunks), dims=dims), diff --git a/geotiepoints/viiinterpolator.py b/geotiepoints/viiinterpolator.py index 6982f9a..5fad3fb 100644 --- a/geotiepoints/viiinterpolator.py +++ b/geotiepoints/viiinterpolator.py @@ -14,7 +14,6 @@ in alt, act (row,col) format instead of act,alt (col,row) """ -import dask import dask.array as da import numpy as np import xarray as xr @@ -60,7 +59,8 @@ def tie_points_geo_interpolation(longitude, latitude, The interpolation is done on the longitude and latitude values, or on cartesian coordinates if the tie points reach latitudes above ``lat_threshold_use_cartesian`` or span more than 180 degrees of - longitude. The choice is made once for all the tie points, which for dask arrays requires computing them. + longitude. The choice is made from all the tie points, so it is the same for every dask chunk. + For dask arrays it is made when the result is computed, so nothing is computed by this function. Args: longitude: xarray DataArray containing the longitude values defined on the tie points (degrees). @@ -80,10 +80,14 @@ def tie_points_geo_interpolation(longitude, latitude, raise ValueError("The dimensions of longitude and latitude don't match") _check_whole_scans(longitude.shape[0], scan_alt_tie_points) - to_cart = _use_cartesian(longitude, latitude, lat_threshold_use_cartesian) + # Lazy for dask arrays: computed once with the result and given to every chunk + max_abs_lat = abs(latitude).max().data + lon_range = (longitude.max() - longitude.min()).data interp_lonlat = _map_scan_blocks(_geo_interpolate_scans, [longitude.data, latitude.data], - scan_alt_tie_points, tie_points_factor, n_outputs=2, - to_cart=to_cart, z_threshold_use_xy=z_threshold_use_xy) + scan_alt_tie_points, tie_points_factor, + reductions=(max_abs_lat, lon_range), n_outputs=2, + lat_threshold_use_cartesian=lat_threshold_use_cartesian, + z_threshold_use_xy=z_threshold_use_xy) return _pixels_to_data_array(interp_lonlat[0], longitude), _pixels_to_data_array(interp_lonlat[1], latitude) @@ -94,26 +98,33 @@ def _check_whole_scans(n_tie_alt, scan_alt_tie_points): f"{scan_alt_tie_points}") -def _use_cartesian(longitude, latitude, lat_threshold_use_cartesian): +def _use_cartesian(max_abs_lat, lon_range, lat_threshold_use_cartesian): """Check if the geographical interpolation must be done on cartesian coordinates. That is the case at high latitudes, or when the longitudes cross the antimeridian. - Both ranges are computed together to compute dask arrays only once. + + Args: + max_abs_lat: largest absolute latitude of all the tie points (degrees), NaN if there are none. + lon_range: difference between the largest and smallest longitudes of all the tie points (degrees). + lat_threshold_use_cartesian: latitude threshold to use cartesian coordinates. """ - max_abs_lat, lon_range = dask.compute(abs(latitude).max(), longitude.max() - longitude.min()) return bool(max_abs_lat > lat_threshold_use_cartesian or lon_range > 180.) -def _map_scan_blocks(func, tie_arrays, scan_alt_tie_points, tie_points_factor, n_outputs=None, **kwargs): +def _map_scan_blocks(func, tie_arrays, scan_alt_tie_points, tie_points_factor, reductions=(), n_outputs=None, + **kwargs): """Call ``func`` on whole scans of tie points, chunk by chunk for dask arrays. Args: - func: function taking the numpy tie point arrays, ``scan_alt_tie_points`` and ``tie_points_factor`` - and returning the pixel array, or ``n_outputs`` pixel arrays stacked along a new first dimension. + func: function taking the numpy tie point arrays, the ``reductions``, ``scan_alt_tie_points`` and + ``tie_points_factor``, and returning the pixel array, or ``n_outputs`` pixel arrays stacked along a + new first dimension. tie_arrays: numpy or dask arrays of the tie points, all with the same shape. scan_alt_tie_points: number of tie points along the satellite track for each scan. tie_points_factor: sub-sampling factor of tie points wrt pixel points. + reductions: numpy or dask 0-d arrays computed from all the tie points. Dask ones are computed once and + the same values are given to every chunk. n_outputs: number of stacked pixel arrays returned by ``func``, or None if it returns only one. kwargs: other keyword arguments passed to ``func``. @@ -122,7 +133,7 @@ def _map_scan_blocks(func, tie_arrays, scan_alt_tie_points, tie_points_factor, n """ if not any(isinstance(arr, da.Array) for arr in tie_arrays): - return func(*tie_arrays, scan_alt_tie_points, tie_points_factor, **kwargs) + return func(*tie_arrays, *reductions, scan_alt_tie_points, tie_points_factor, **kwargs) tie_arrays = _rechunk_to_whole_scans([da.asarray(arr) for arr in tie_arrays], scan_alt_tie_points) row_chunks, col_chunks = tie_arrays[0].chunks @@ -135,7 +146,7 @@ def _map_scan_blocks(func, tie_arrays, scan_alt_tie_points, tie_points_factor, n map_kwargs = {"chunks": pixel_chunks} else: map_kwargs = {"chunks": ((n_outputs,),) + pixel_chunks, "new_axis": 0} - return da.map_blocks(func, *tie_arrays, scan_alt_tie_points, tie_points_factor, **kwargs, + return da.map_blocks(func, *tie_arrays, *reductions, scan_alt_tie_points, tie_points_factor, **kwargs, dtype=dtype, meta=np.array((), dtype=dtype), **map_kwargs) @@ -215,15 +226,18 @@ def _interpolate_intervals(start, end, weights, out): interpolated += start -def _geo_interpolate_scans(longitude, latitude, scan_alt_tie_points, tie_points_factor, to_cart, z_threshold_use_xy): +def _geo_interpolate_scans(longitude, latitude, max_abs_lat, lon_range, scan_alt_tie_points, tie_points_factor, + lat_threshold_use_cartesian, z_threshold_use_xy): """Interpolate whole scans of longitude and latitude tie points to pixel points. Args: longitude: numpy array of longitude tie points (degrees). latitude: numpy array of latitude tie points (degrees). + max_abs_lat: largest absolute latitude of all the tie points, not only these ones (degrees). + lon_range: longitude range of all the tie points, not only these ones (degrees). scan_alt_tie_points: number of tie points along the satellite track for each scan. tie_points_factor: sub-sampling factor of tie points wrt pixel points. - to_cart: interpolate on cartesian coordinates instead of longitudes and latitudes. + lat_threshold_use_cartesian: latitude threshold to use cartesian coordinates. z_threshold_use_xy: z threshold to compute latitude from x and y in cartesian coordinates. Returns: @@ -234,7 +248,7 @@ def _geo_interpolate_scans(longitude, latitude, scan_alt_tie_points, tie_points_ shape = (2, n_tie_alt // scan_alt_tie_points * (scan_alt_tie_points - 1) * tie_points_factor, (n_tie_act - 1) * tie_points_factor) lonlat = np.empty(shape, dtype=_pixel_dtype(longitude, latitude)) - if to_cart: + if _use_cartesian(max_abs_lat, lon_range, lat_threshold_use_cartesian): x_coords, y_coords, z_coords = (_interpolate_scans(coords, scan_alt_tie_points, tie_points_factor) for coords in _lonlat2xyz(longitude, latitude)) _xyz2lonlat(x_coords, y_coords, z_coords, z_threshold_use_xy, out=lonlat)