Point to Grid Translator for MetNet3 (OMO weather stations) (#85)

* initial commit for point to grid translator

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: trudramukerji14

metnet/data/__init__.py

@@ -1,5 +1 @@
 """Data loading and preprocessing utilities for MetNet."""
-
-from .merge_single_band import merge_086um_band, merge_two_arrays
-
-__all__ = ["merge_two_arrays", "merge_086um_band"]

metnet/data/point_to_grid.py

@@ -0,0 +1,160 @@
+"""Point-source to grid translator for MetNet-3 weather station data.
+
+Translates sparse weather station observations (lat/lon point measurements)
+to a dense 2D grid, as described in the MetNet-3 paper:
+
+'For surface variables, we take the OMO station point measurements and map
+them to a 4km by 4km pixel in which the station lies. If there are multiple
+stations in a given region, we take the average of their measurements.'
+
+The 14 variables used are based on the OMO variable list at:
+https://madis.ncep.noaa.gov/sfc_OMO_variable_list.shtml
+"""
+
+from typing import Dict, List
+
+import numpy as np
+import torch
+from pyproj import Transformer
+
+# Best approximation of the 14 OMO input variables using MADIS ASOS
+OMO_VARIABLES = [
+    "T",  # air temperature (K)
+    "TD",  # dewpoint temperature (K)
+    "RH",  # relative humidity (%)
+    "P",  # station pressure (Pa)
+    "DD",  # wind direction (deg)
+    "FF",  # wind speed (m/s)
+    "U",  # u wind component (m/s)
+    "V",  # v wind component (m/s)
+    "FFGUST",  # wind gust (m/s)
+    "VIS",  # visibility (m)
+    "PCPRATE",  # precipitation rate (kg/m²/s)
+    "PCP1H",  # accumulated precip 1h (m)
+    "Q",  # specific humidity (kg/kg)
+    "ALTSE",  # altimeter pressure (Pa)
+]
+
+
+class PointToGridTranslator:
+    """
+    Translates sparse weather station observations to a dense grid.
+
+    Supports Lambert Conformal Conic projection, matches CONUS setup
+    """
+
+    def __init__(
+        self,
+        grid_height: int,
+        grid_width: int,
+        lat_min: float,
+        lat_max: float,
+        lon_min: float,
+        lon_max: float,
+        resolution_km: float = 4.0,
+        variables: List[str] = OMO_VARIABLES,
+        fill_value: float = 0.0,
+    ):
+        """
+        Initialize the translator.
+
+        Args:
+            grid_height: Number of grid points in y direction
+            grid_width: Number of grid points in x direction
+            lat_min: Minimum latitude of grid domain
+            lat_max: Maximum latitude of grid domain
+            lon_min: Minimum longitude of grid domain
+            lon_max: Maximum longitude of grid domain
+            resolution_km: Grid resolution in km, 4.0 as per MetNet-3
+            variables: List of variable names to extract from station data
+            fill_value: Value for grid cells with no station, default 0.0
+        """
+        self.grid_height = grid_height
+        self.grid_width = grid_width
+        self.lat_min = lat_min
+        self.lat_max = lat_max
+        self.lon_min = lon_min
+        self.lon_max = lon_max
+        self.resolution_km = resolution_km
+        self.variables = variables
+        self.fill_value = fill_value
+        self.num_variables = len(variables)
+
+        self.transformer = Transformer.from_crs(
+            "EPSG:4326",  # WGS84 lat/lon
+            "+proj=lcc +lat_1=25 +lat_2=60 +lat_0=40 +lon_0=-96",  # CONUS LCC
+            always_xy=True,
+        )
+        self.x_min, self.y_min = self.transformer.transform(lon_min, lat_min)
+
+    def _latlon_to_pixel(self, lat: float, lon: float):
+        """Convert lat/lon to grid pixel coordinates."""
+        x, y = self.transformer.transform(lon, lat)
+        col = int((x - self.x_min) / (self.resolution_km * 1000))
+        row = int((y - self.y_min) / (self.resolution_km * 1000))
+        return row, col
+
+    def translate(
+        self,
+        stations: List[Dict],
+    ) -> torch.Tensor:
+        """
+        Translate station observations to a dense grid.
+
+        Args:
+            stations: List of dicts, each with keys:
+                - 'lat': station latitude
+                - 'lon': station longitude
+                - variable names matching self.variables with float values
+
+        Returns:
+            Dense grid tensor of shape [num_variables, grid_height, grid_width]
+            with fill_value where no station is present
+        """
+        # Initialize grid and count grids for averaging
+        grid = np.full(
+            (self.num_variables, self.grid_height, self.grid_width),
+            self.fill_value,
+            dtype=np.float32,
+        )
+        counts = np.zeros(
+            (self.grid_height, self.grid_width),
+            dtype=np.int32,
+        )
+
+        for station in stations:
+            lat = station["lat"]
+            lon = station["lon"]
+
+            # Skip stations outside grid bounds
+            if not (self.lat_min <= lat <= self.lat_max):
+                continue
+            if not (self.lon_min <= lon <= self.lon_max):
+                continue
+
+            row, col = self._latlon_to_pixel(lat, lon)
+
+            # Skip if outside grid
+            if not (0 <= row < self.grid_height and 0 <= col < self.grid_width):
+                continue
+
+            # Add station values
+            for i, var in enumerate(self.variables):
+                if var in station and station[var] is not None:
+                    value = station[var]
+                    if np.isnan(value):
+                        continue
+                    if counts[row, col] == 0:
+                        grid[i, row, col] = value
+                    else:
+                        # Running sum for averaging
+                        grid[i, row, col] += value
+
+            counts[row, col] += 1
+
+        # Average where multiple stations fall in same pixel
+        multiple = counts > 1
+        if np.any(multiple):
+            grid[:, multiple] /= counts[multiple]
+
+        return torch.from_numpy(grid)

requirements.txt

@@ -5,3 +5,4 @@ antialiased_cnns
 axial_attention
 pytorch_msssim
 huggingface_hub
+pyproj>=3.0.0

tests/test_data.py

@@ -0,0 +1,78 @@
+import torch
+from metnet.data.point_to_grid import PointToGridTranslator
+
+
+def make_translator(
+    grid_height=10,
+    grid_width=10,
+    # Domain sized to match a 10x10 grid at the default 4km resolution
+    # (40km x 40km), with SAMPLE_STATION (40.0, -95.0) landing near the
+    # center. A domain much larger than grid_height/width * resolution_km
+    # will place stations outside the array, so translate() silently
+    # drops them instead of raising.
+    lat_min=39.8,
+    lat_max=40.2,
+    lon_min=-95.25,
+    lon_max=-94.5,
+    fill_value=0.0,
+):
+    return PointToGridTranslator(
+        grid_height=grid_height,
+        grid_width=grid_width,
+        lat_min=lat_min,
+        lat_max=lat_max,
+        lon_min=lon_min,
+        lon_max=lon_max,
+        fill_value=fill_value,
+    )
+
+
+SAMPLE_STATION = {
+    "lat": 40.0,
+    "lon": -95.0,
+    "T": 295.0,
+    "TD": 288.0,
+    "RH": 70.0,
+    "P": 101325.0,
+    "DD": 180.0,
+    "FF": 5.0,
+    "U": 0.0,
+    "V": -5.0,
+    "FFGUST": 8.0,
+    "VIS": 10000.0,
+    "PCPRATE": 0.0,
+    "PCP1H": 0.0,
+    "Q": 0.01,
+    "ALTSE": 101325.0,
+}
+
+
+def test_output_shape():
+    output = make_translator().translate([SAMPLE_STATION])
+    assert output.shape == (14, 10, 10)
+
+
+def test_fill_value():
+    output = make_translator().translate([])
+    assert (output == 0.0).all()
+
+
+def test_nan_handling():
+    station = {**SAMPLE_STATION, "T": float("nan")}
+    output = make_translator().translate([station])
+    assert output[0].max() == 0.0
+
+
+def test_outside_bounds():
+    station = {**SAMPLE_STATION, "lat": 60.0}
+    output = make_translator().translate([station])
+    assert (output == 0.0).all()
+
+
+def test_averaging():
+    translator = make_translator(fill_value=-999.0)  # distinct fill value
+    station1 = {**SAMPLE_STATION, "T": 290.0}
+    station2 = {**SAMPLE_STATION, "T": 300.0}  # exact same lat/lon
+    output = translator.translate([station1, station2])
+    non_fill = output[0][output[0] != -999.0]
+    assert torch.isclose(non_fill, torch.tensor(295.0)).any()