diff --git a/tests/test_spatial_avg.py b/tests/test_spatial_avg.py index 32ae9419..f4c9a252 100644 --- a/tests/test_spatial_avg.py +++ b/tests/test_spatial_avg.py @@ -396,7 +396,7 @@ class TestGetWeights: def setup(self): self.ds = generate_dataset(cf_compliant=True, has_bounds=True) - def test_returns_area_weights_for_region_within_lat_and_lon(self): + def test_area_weights_for_region_within_lat_and_lon(self): result = self.ds.spatial._get_weights( axis=["lat", "lon"], lat_bounds=(-5, 5), lon_bounds=(-170, -120) ) @@ -415,7 +415,7 @@ def test_returns_area_weights_for_region_within_lat_and_lon(self): xr.testing.assert_allclose(result, expected) - def test_returns_area_weights_for_region_within_lat(self): + def test_area_weights_for_region_within_lat(self): result = self.ds.spatial._get_weights( axis=["lat", "lon"], lat_bounds=(-5, 5), lon_bounds=None ) @@ -434,7 +434,7 @@ def test_returns_area_weights_for_region_within_lat(self): xr.testing.assert_allclose(result, expected) - def test_returns_area_weights_for_region_within_lon(self): + def test_area_weights_for_region_within_lon(self): expected = xr.DataArray( data=np.array( [ @@ -459,37 +459,68 @@ class TestGetLongitudeWeights: def setup(self): self.ds = generate_dataset(cf_compliant=True, has_bounds=True) - def test_returns_area_weights_for_region_within_lon(self): - expected = xr.DataArray( - data=np.array([0.0, 0.0, 50.0, 0.0]), - coords={"lon": self.ds.lon}, - dims=["lon"], - ) + def test_returns_weights_for_region_in_lon(self): # Longitude axes orientation swaps from (-180, 180) to (0, 360). result = self.ds.spatial._get_longitude_weights( domain_bounds=self.ds.lon_bnds.copy(), region_bounds=np.array([-170.0, -120.0]), ) + expected = xr.DataArray( + data=np.array([0.0, 0.0, 50.0, 0.0]), + coords={"lon": self.ds.lon}, + dims=["lon"], + ) xr.testing.assert_allclose(result, expected) - def test_returns_area_weights_for_region_within_lon_including_prime_meridian_cell( - self, - ): + def test_weights_for_region_in_lon_domain_with_both_spanning_p_meridian(self): ds = self.ds.copy() + # Domain spans prime meridian. ds.lon_bnds.data[:] = np.array([[359, 1], [1, 90], [90, 180], [180, 359]]) + result = ds.spatial._get_longitude_weights( + domain_bounds=ds.lon_bnds, + # Region spans prime meridian. + region_bounds=np.array([359, 1]), + ) expected = xr.DataArray( - data=np.array([0.0, 0.0, 0.0, 50.0]), - coords={"lon": self.ds.lon}, + data=np.array([2.0, 0.0, 0.0, 0.0]), + coords={"lon": ds.lon}, dims=["lon"], - ) + + xr.testing.assert_allclose(result, expected) + + def test_weights_for_region_in_lon_domain_with_domain_spanning_p_meridian(self): + ds = self.ds.copy() + # Domain spans prime meridian. + ds.lon_bnds.data[:] = np.array([[359, 1], [1, 90], [90, 180], [180, 359]]) # Longitude axes orientation swaps from (-180, 180) to (0, 360). - result = self.ds.spatial._get_longitude_weights( - domain_bounds=self.ds.lon_bnds, + result = ds.spatial._get_longitude_weights( + domain_bounds=ds.lon_bnds, region_bounds=np.array([-170.0, -120.0]), ) + expected = xr.DataArray( + data=np.array([0.0, 0.0, 0.0, 50.0]), + coords={"lon": ds.lon}, + dims=["lon"], + ) + + xr.testing.assert_allclose(result, expected) + + def test_weights_for_region_in_lon_domain_with_region_spanning_p_meridian(self): + ds = self.ds.copy() + + result = ds.spatial._get_longitude_weights( + domain_bounds=ds.lon_bnds, + # Region spans prime meridian. + region_bounds=np.array([359, 1]), + ) + expected = xr.DataArray( + data=np.array([1.875, 0.0625, 0.0, 0.0625]), + coords={"lon": ds.lon}, + dims=["lon"], + ) xr.testing.assert_allclose(result, expected) @@ -528,7 +559,7 @@ class TestGetLatitudeWeights: def setup(self): self.ds = generate_dataset(cf_compliant=True, has_bounds=True) - def test_returns_area_weights_for_region_within_lat(self): + def test_weights_for_region_within_lat(self): expected = xr.DataArray( data=np.array([0.0, 0.087156, 0.087156, 0.0]), coords={"lat": self.ds.lat}, diff --git a/xcdat/spatial_avg.py b/xcdat/spatial_avg.py index 2c8f3ea8..1fa29968 100644 --- a/xcdat/spatial_avg.py +++ b/xcdat/spatial_avg.py @@ -369,11 +369,12 @@ def _get_longitude_weights( 1. Align the axis orientations of the domain and region bounds to (0, 360) to ensure compatibility in the proceeding steps. - 2. Handle grid cells that cross the prime meridian (e.g., [-1, 1]) - by recreating the axis with two additional grid cells ([0, 1] and - [359, 360]) to ensure alignment with the (0, 360) axis orientation. - The prime meridian grid cell is returned as a variable to handle - the length of weights in a proceeding step. + 2. Handle grid cells that cross the prime meridian (e.g., [-1, 1]) by + breaking such grid cells into two (e.g., [0, 1] and [359, 360]) to + ensure alignment with the (0, 360) axis orientation. This results in + a bounds axis of length(nlon)+1. The index of the grid cell that + crosses the prime meridian is returned in order to reduce the length + of weights to nlon. 3. Scale the domain down to a region (if selected). 4. Calculate weights using the domain bounds. 5. If the prime meridian grid cell exists, use this cell's index to