diff --git a/AGENTS.md b/AGENTS.md index fd5d9b0..4d8e9e2 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,16 @@ 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` 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 37eae0c..bd679ab 100644 --- a/geotiepoints/tests/test_viiinterpolator.py +++ b/geotiepoints/tests/test_viiinterpolator.py @@ -6,12 +6,16 @@ """ -import unittest +import dask +import dask.array as da 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 +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( @@ -42,7 +48,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 +71,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,134 +95,237 @@ ) -class TestViiInterpolator(unittest.TestCase): - """Test the vii_utils module.""" +@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 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 _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, 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, use_dask) + + +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): + """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, 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.]) + + +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, use_dask) + with pytest.raises(ValueError, match="must be a multiple"): + tie_points_interpolation([data], TEST_SCAN_ALT_TIE_POINTS, TEST_TIE_POINTS_FACTOR) - 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)) - - # 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)) +@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, use_dask): + """Test the coordinates interpolation routine in geodetic and cartesian coordinates.""" + with dask.config.set(scheduler=CustomScheduler(max_computes=0)): lon, lat = tie_points_geo_interpolation( - self.longitude_over360, - self.latitude, + _tie_points_data_array(longitude, use_dask), + _tie_points_data_array(latitude, use_dask), 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)) + _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) + + +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), 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. + + 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) + 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=0)): 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)) - - # 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 - ) + 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 6b4add8..5fad3fb 100644 --- a/geotiepoints/viiinterpolator.py +++ b/geotiepoints/viiinterpolator.py @@ -4,17 +4,19 @@ 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.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,79 +24,43 @@ 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 across, then along the scan - # (much faster than interpolating directly in the two dimensions) - 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_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 - def tie_points_geo_interpolation(longitude, latitude, scan_alt_tie_points, tie_points_factor, lat_threshold_use_cartesian=60., 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 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). @@ -106,32 +72,190 @@ 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) + + # 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, + 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) + + +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}") + + +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. - # 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. + 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. - if to_cart: + """ + return bool(max_abs_lat > lat_threshold_use_cartesian or lon_range > 180.) - x, y, z = _lonlat2xyz(longitude, latitude) - interp_x, interp_y, interp_z = tie_points_interpolation([x, y, z], - scan_alt_tie_points, - tie_points_factor) +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. - interp_longitude, interp_latitude = _xyz2lonlat(interp_x, interp_y, interp_z, z_threshold_use_xy) + Args: + 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``. + 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, *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 + 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, *reductions, scan_alt_tie_points, tie_points_factor, **kwargs, + dtype=dtype, meta=np.array((), dtype=dtype), **map_kwargs) - interp_longitude, interp_latitude = tie_points_interpolation([longitude, latitude], - scan_alt_tie_points, - tie_points_factor) - return interp_longitude, interp_latitude +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) + + +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, 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. + 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: + 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 _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) + 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): @@ -153,26 +277,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 + 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)))