Adds a regular grid Latitude and longitude binning method.

PiperOrigin-RevId: 809092370

Author: WeatherBenchX authors

weatherbenchX/binning.py

@@ -49,30 +49,43 @@ def create_bin_mask(
     """
 
 
-def _region_to_mask(
-    lat: xr.DataArray,
-    lon: xr.DataArray,
-    lat_lims: Tuple[int, int],
-    lon_lims: Tuple[int, int],
+def _create_lat_mask(
+    lat: xr.DataArray, lat_lims: Tuple[int, int]
 ) -> xr.DataArray:
-  """Computes a boolean mask for a lat/lon limits region."""
+  """Computes a boolean mask for a latitude limits region."""
   if lat_lims[0] >= lat_lims[1]:
     raise ValueError(
         f'`lat_lims[0]` must be smaller than `lat_lims[1]`, got {lat_lims}`'
     )
-  lat_mask = np.logical_and(lat >= lat_lims[0], lat <= lat_lims[1])
+  return np.logical_and(lat >= lat_lims[0], lat <= lat_lims[1])
 
+
+def _create_lon_mask(
+    lon: xr.DataArray, lon_lims: Tuple[int, int]
+) -> xr.DataArray:
+  """Computes a boolean mask for a longitude limits region."""
   # Make sure we are in the [0, 360] interval.
   lon = np.mod(lon, 360)
   lon_lims = np.mod(lon_lims[0], 360), np.mod(lon_lims[1], 360)
-
   if lon_lims[1] > lon_lims[0]:
     # Same as the latitude.
     lon_mask = np.logical_and(lon >= lon_lims[0], lon <= lon_lims[1])
   else:
     # In this case it means we need to wrap longitude around the other side of
     # the globe.
     lon_mask = np.logical_or(lon <= lon_lims[1], lon >= lon_lims[0])
+  return lon_mask
+
+
+def _region_to_mask(
+    lat: xr.DataArray,
+    lon: xr.DataArray,
+    lat_lims: Tuple[int, int],
+    lon_lims: Tuple[int, int],
+) -> xr.DataArray:
+  """Computes a boolean mask for a lat/lon limits region."""
+  lat_mask = _create_lat_mask(lat, lat_lims)
+  lon_mask = _create_lon_mask(lon, lon_lims)
   return np.logical_and(lat_mask, lon_mask)
 
 
@@ -185,6 +198,93 @@ def create_bin_mask(
     return masks
 
 
+class LatitudeBins(Binning):
+  """Class for binning by latitude bands."""
+
+  def __init__(
+      self,
+      degrees: float,
+      lat_range: Tuple[int, int] = (-90, 90),
+      bin_dim_name: str = 'latitude_bins',
+  ):
+    """Init.
+
+    Args:
+      degrees: Grid spacing in degrees.
+      lat_range: Tuple of (min_lat, max_lat).
+      bin_dim_name: Name of binning dimension.
+    """
+    super().__init__(bin_dim_name)
+    self._degrees = degrees
+    self._lat_bins = np.arange(
+        lat_range[0], lat_range[1] + self._degrees, self._degrees
+    )
+
+  def create_bin_mask(
+      self,
+      statistic: xr.DataArray,
+  ) -> xr.DataArray:
+    """Creates a bin mask for a statistic."""
+    masks = []
+    for lat_start in self._lat_bins[:-1]:
+      lat_end = lat_start + self._degrees
+      mask = _create_lat_mask(
+          statistic.latitude,
+          (lat_start, lat_end),
+      )
+      # Broadcast the mask to the shape of statistic
+      mask = mask.broadcast_like(statistic)
+      mask = mask.expand_dims(dim=self.bin_dim_name, axis=0)
+      mask.coords[self.bin_dim_name] = np.array([lat_start])
+      masks.append(mask)
+    return xr.concat(masks, dim=self.bin_dim_name)
+
+
+class LongitudeBins(Binning):
+  """Class for binning by longitude bands."""
+
+  def __init__(
+      self,
+      degrees: float,
+      lon_range: Tuple[int, int] = (0, 360),
+      bin_dim_name: str = 'longitude_bins',
+  ):
+    """Init.
+
+    Args:
+      degrees: Grid spacing in degrees.
+      lon_range: Tuple of (min_lon, max_lon).
+      bin_dim_name: Name of binning dimension.
+    """
+    super().__init__(bin_dim_name)
+    self._degrees = degrees
+    lon_end = lon_range[1]
+    if lon_range[0] >= lon_range[1]:
+      lon_end += 360
+    self._lon_bins = np.arange(
+        lon_range[0], lon_end + self._degrees, self._degrees
+    )
+
+  def create_bin_mask(
+      self,
+      statistic: xr.DataArray,
+  ) -> xr.DataArray:
+    """Creates a bin mask for a statistic."""
+    masks = []
+    for lon_start in self._lon_bins[:-1]:
+      lon_end = lon_start + self._degrees
+      mask = _create_lon_mask(
+          statistic.longitude,
+          (lon_start, lon_end),
+      )
+      # Broadcast the mask to the shape of statistic
+      mask = mask.broadcast_like(statistic)
+      mask = mask.expand_dims(dim=self.bin_dim_name, axis=0)
+      mask.coords[self.bin_dim_name] = np.array([np.mod(lon_start, 360)])
+      masks.append(mask)
+    return xr.concat(masks, dim=self.bin_dim_name)
+
+
 def vectorized_coord_mask(
     coord: xr.DataArray,
     coord_name: str,

weatherbenchX/binning_test.py

@@ -151,6 +151,37 @@ def test_by_time_unit_from_seconds_binning(self):
         mask['prediction_timedelta_sec_hour'], np.arange(0, 6)
     )
 
+  @parameterized.parameters(
+      ('second', None, np.arange(0, 60)),
+      ('second', [0, 15, 30, 45], [0, 15, 30, 45]),
+      ('minute', None, np.arange(0, 60)),
+      ('minute', [0, 30], [0, 30]),
+      ('hour', None, np.arange(0, 24)),
+      ('hour', [0, 6, 12, 18], [0, 6, 12, 18]),
+  )
+  def test_by_time_unit_from_seconds_binning_with_units(
+      self, unit, bins, expected_bins
+  ):
+    statistic_values = test_utils.mock_prediction_data(
+        time_start='2020-01-01T00',
+        time_stop='2020-01-01T01',
+        time_resolution='1 hr',
+        lead_resolution='1 second',
+        lead_stop='24 hour',
+    )['2m_temperature']
+    statistic_values = statistic_values.assign_coords({
+        'prediction_timedelta_sec': (
+            statistic_values.prediction_timedelta.dt.total_seconds()
+        )
+    })
+    binning_obj = binning.ByTimeUnitFromSeconds(
+        unit, 'prediction_timedelta_sec', bins=bins
+    )
+    mask = binning_obj.create_bin_mask(statistic_values)
+    np.testing.assert_array_equal(
+        mask[f'prediction_timedelta_sec_{unit}'].values, expected_bins
+    )
+
   def test_by_coord_bins(self):
     target_path = resources.files('weatherbenchX').joinpath(
         'test_data/metar-timeNominal-by-month'
@@ -222,6 +253,67 @@ def test_by_sets(self):
     self.assertEqual(mask.sum('index').sel(station_subset='wrong_set'), 0)
     self.assertLen(statistic, mask.sum('index').sel(station_subset='global'))
 
+  @parameterized.parameters(
+      (10, (-90, 90), 18),
+      (30, (-90, 90), 6),
+      (20, (0, 60), 3),
+  )
+  def test_latitude_bins(self, degrees, lat_range, expected_bins):
+    statistic_values = test_utils.mock_prediction_data(
+        time_start='2020-01-01T00', time_stop='2020-01-01T01'
+    )['2m_temperature']
+    binning_obj = binning.LatitudeBins(degrees, lat_range)
+    mask = binning_obj.create_bin_mask(statistic_values)
+    self.assertEqual(mask.latitude_bins.shape[0], expected_bins)
+    self.assertTrue(np.all(mask.latitude_bins.values >= lat_range[0]))
+    self.assertTrue(np.all(mask.latitude_bins.values < lat_range[1]))
+    self.assertEqual(mask.shape, (expected_bins,) + statistic_values.shape)
+    # Check that a point is in the correct bin
+    # Find the latitude closest to 25
+    lat_val = 25
+    if not (lat_range[0] <= lat_val < lat_range[1]):
+      # If 25 is not in range, pick a value that is.
+      lat_val = (lat_range[0] + lat_range[1]) / 2
+    lat_idx = np.argmin(np.abs(statistic_values.latitude.values - lat_val))
+    lon_idx = np.argmin(np.abs(statistic_values.longitude.values - 0))
+
+    # Find the bin that contains the selected latitude
+    expected_bin_idx = (statistic_values.latitude.values[lat_idx] - lat_range[0]) // degrees
+    self.assertTrue(
+        mask.isel(latitude_bins=int(expected_bin_idx), latitude=lat_idx, longitude=lon_idx).values.all()
+    )
+
+  @parameterized.parameters(
+      (10, (0, 360), 36, 10),
+      (30, (0, 360), 12, 150),
+      (60, (-180, 180), 6, 0),
+      (90, (270, 360), 1, 300),
+  )
+  def test_longitude_bins(self, degrees, lon_range, expected_bins, test_lon):
+    statistic_values = test_utils.mock_prediction_data(
+        time_start='2020-01-01T00', time_stop='2020-01-01T01'
+    )['2m_temperature']
+    binning_obj = binning.LongitudeBins(degrees, lon_range)
+    mask = binning_obj.create_bin_mask(statistic_values)
+    self.assertEqual(mask.longitude_bins.shape[0], expected_bins)
+    self.assertEqual(mask.shape, (expected_bins,) + statistic_values.shape)
+    # Check wrapping
+    if lon_range == (-180, 180):
+      self.assertTrue(0 in mask.longitude_bins.values)
+
+    # Find the longitude closest to test_lon
+    lon_idx = np.argmin(np.abs(statistic_values.longitude.values - test_lon))
+    lat_idx = np.argmin(np.abs(statistic_values.latitude.values - 0))
+    lon_val = statistic_values.longitude.values[lon_idx]
+
+    # Calculate expected bin index: (lon_val - lon_range[0]) // degrees
+    # This works even with wrapping ranges because lon_bins is constructed correctly.
+    expected_bin_idx = (lon_val - lon_range[0]) // degrees
+
+    self.assertTrue(
+        mask.isel(longitude_bins=int(expected_bin_idx), latitude=lat_idx, longitude=lon_idx).values.all()
+    )
+
 
 if __name__ == '__main__':
   absltest.main()