Support weather station evaluation

One highly requested feature is to evaluate against weather stations. We are currently rewriting the evaluation pipeline to support a variety of observation types. However, since this will likely take a few months, I want to share the outline of an initial attempt to support station evaluation in the current framework. This is not officially supported but rather is intended as inspiration for anyone that wants to build on it.

You can find the draft in the `station_support` branch: https://github.com/google-research/weatherbench2/tree/station_support

---

First, the pipeline now has a DataReader class that allows arbitrary data sources to be read for forecast and ground truth.
```
@dataclasses.dataclass()
class DataReader:
path: str
variables: Optional[Sequence[str]] = None
rename_variables: Optional[str] = None

def get_chunk(self, time_chunk: xr.Dataset) -> xr.Dataset:
"""Return chunk of data for given time_chunk.

Args:
time_chunk: xr.Dataset with coordinates: init_time, lead_time and their
combination valid_time

Returns:
chunk: xr.Dataset with corresponding data chunk
"""
raise NotImplementedError()
```

This allows us to then implement classes like “GriddedForecastFromZarr” and “SparseGroundTruthFromParquet”. This could eventually also be extended as an online model API.

Time alignment of sparse observations: There is the option to use an exact time or allow for some tolerance around the valid_time in question. If for the given time range several observations are present for a single station, the closest time is chosen.

Before the metrics computation, an additional grid-to-sparse step is required. Currently I am using “nearest”. Linear probably gives slightly smaller errors.

```
def interpolate_grid2sparse(
fc: xr.Dataset, gt: xr.Dataset, method: str
) -> tuple[xr.Dataset, xr.Dataset]:
fc_like_gt = fc.interp(
latitude=gt.latitude, longitude=gt.longitude, method=method
) # pytype: disable=wrong-arg-types
return fc_like_gt, gt
```

Next, I removed the hard-coded lat/lon dependencies in the metrics and separated the point-wise metric/statistic computation and the spatial aggregation.

```
@dataclasses.dataclass
class Metric:
"""Base class for metrics."""

def compute_statistic(
self,
forecast: xr.Dataset,
truth: xr.Dataset,
) -> xr.Dataset:
"""Compute point-wise statistic.

Args:
forecast: dataset of forecasts to evaluate.
truth: dataset of ground truth. Should have the same variables as
forecast.

Returns:
Dataset with point-wise statistic.
"""
raise NotImplementedError

def compute_statistic_and_average_in_time(
self,
forecast: xr.Dataset,
truth: xr.Dataset,
) -> xr.Dataset:
"""Evaluate this metric on datasets with full temporal coverages."""
if "time" in forecast.dims:
avg_dim = "time"
elif "init_time" in forecast.dims:
avg_dim = "init_time"
else:
raise ValueError(
f"Forecast has neither valid_time or init_time dimension {forecast}"
)
return self.compute_statistic(forecast, truth).mean(avg_dim)

class MSE(Metric):

def compute_statistic(
self, forecast: xr.Dataset, truth: xr.Dataset
) -> xr.Dataset:
statistic = (forecast - truth) ** 2
return statistic
```

The aggregations can then be defined separately (replacing the previous regions).
```
@dataclasses.dataclass
class Aggregation:
skipna: bool = False

def aggregate_in_space(self, statistic):
raise NotImplementedError

@dataclasses.dataclass
class NoAggregation(Aggregation):

def aggregate_in_space(self, statistic):
return statistic

@dataclasses.dataclass
class LatLonAverage(Aggregation):

def aggregate_in_space(self, statistic):
weights = metrics.get_lat_weights(statistic)
return statistic.weighted(weights).mean(
('latitude', 'longitude'), skipna=self.skipna
)

@dataclasses.dataclass
class UnweightedAverage(Aggregation):
dims: list[str] = dataclasses.field(default_factory=list)

def aggregate_in_space(self, statistic):
return statistic.mean(self.dims, skipna=self.skipna)

@dataclasses.dataclass
class WeightedStationAverage(Aggregation):
station_dim: str = 'stationName'
weights: Optional[xr.DataArray] = None
alpha_0: float = 0.75
min_weight: float = 1
max_weight: float = 10

def aggregate_in_space(self, statistic):
if self.weights is None:
self.weights = compute_station_weights(
statistic,
self.station_dim,
self.alpha_0,
self.min_weight,
self.max_weight,
)
return statistic.weighted(self.weights).mean(
self.station_dim, skipna=self.skipna
)
```

A sketch of what the pipeline along with the configs could look like is in `scripts/evaluate_stations.py`

Start with the station_support branch and scripts/evaluate_stations.py, then read the proposed DataReader, interpolate_grid2sparse, Metric, and Aggregation interfaces in the issue. The issue is an architectural sketch rather than a scoped task, so a contributor would first need to define the supported station workflow, validation, and acceptance criteria before implementation.

Written by the indexing model from the issue text.