32 create output image for global and local results

bris-adapt/README.md

@@ -56,3 +56,77 @@ uv run bris-adapt checkpoint move-domain --orography-file ghana.tiff --grid 0.05
 ```
 
 Note that the downloaded grid data must be _larger_ than the target area for the checkpoint.
+
+#### Creating image with globale and local area data
+
+We need two configuration files to output both global and local data. anemoi-inference must be run two times, once for each configuration.
+
+Global forcast configuration: create a configuration file with the following content.
+
+```yaml
+ ...
+post_processors: 
+  - extract_from_state: 
+      cutout_source: 'global'
+
+output:
+  tee:
+    outputs:
+      - printer
+      - netcdf: 
+          path: global.nc
+          variables:
+            - '2t'
+            - msl
+            - '10u'
+            - '10v'
+            - 'tp'
+
+...
+```
+
+Local forecast configuration: create a configuration file with the following content.
+
+```yaml
+...
+post_processors: 
+  - extract_from_state: 
+      cutout_source: 'lam_0'
+
+output:
+  tee:
+    outputs:
+      - printer
+      - netcdf: 
+          path: local.nc
+
+```
+
+**Note** At the momment, with the post_processors.extract_from_state, the cutout_source global and lam_0 are exlusive so we need to run two inferences.
+
+The netcdf output from the inference is not grided. The first step is to convert local.nc and global.nc
+to a regular grid.
+
+Convert **local.nc** to a regular grid.
+
+```shell
+uv run bris-adapt process make-grid --config bris-adapt/etc/mkgrid.json local.nc local-gridded.nc
+```
+
+The **global.nc** netcdf is scattered points over the globe. We need to interpolate this on a regular grid. 
+
+```shell
+uv run bris-adapt process mkglobal-grid global.nc global_0_25deg.nc
+```
+
+We can now create an image with both global and local forecast.
+
+```shell
+uv run  bris-adapt process create-image global_0_25deg.nc local-gridded.nc --map-type temperature --map-area africa --timestep 1 --output-dir images
+```
+
+To create an animated gif with data from all timesteps.
+
+```shell
+uv run  bris-adapt process create-image global_0_25deg.nc local-gridded.nc --map-type temperature --map-area africa --timestep -1 --create-animated-gif --output-dir images
+```

bris-adapt/bris_adapt/process/areas.py

@@ -0,0 +1,59 @@
+from typing import Dict
+import json
+from pydantic import BaseModel
+
+
+class Area(BaseModel):
+    north: float
+    west: float
+    south: float
+    east: float
+    name: str | None = None
+
+    def as_tuple(self) -> tuple[float, float, float, float]:
+        return (self.north, self.west, self.south, self.east)
+
+    def as_list(self) -> list[float]:
+        return [self.north, self.west, self.south, self.east]
+
+    def set_name(self, name: str):
+        self.name = name
+
+    def get_name(self) -> str:
+        return self.name if self.name is not None else "<unnamed>"
+
+    def __str__(self) -> str:
+        name = f"({self.get_name()})" if self.name != '<unnamed>' else ""
+        return f"{self.north}/{self.west}/{self.south}/{self.east} {name}"
+
+
+class AreasConfig(BaseModel):
+    areas: Dict[str, Area]
+
+    def get_area(self, name: str) -> Area:
+        if name not in self.areas:
+            raise ValueError(f"Area '{name}' not found in configuration.")
+        return self.areas[name]
+
+    def list_area_names(self) -> list[str]:
+        return list(self.areas)
+
+
+def load_areas(config: str) -> AreasConfig:
+    '''Load areas from a JSON configuration file.'''
+    with open(config) as f:
+        config_json = json.load(f)
+        areas = AreasConfig.model_validate(config_json)
+        for name, area in areas.areas.items():
+            area.set_name(name)
+        return areas
+
+
+def parse_area_from_str(area: str) -> Area:
+    '''Parse an area string in the format "north/west/south/east" into an Area object.'''
+    parts = area.split('/')
+    if len(parts) != 4:
+        raise ValueError(
+            "Area string must be in the format 'north/west/south/east'.")
+    north, west, south, east = map(float, parts)
+    return Area(north=north, west=west, south=south, east=east)

bris-adapt/bris_adapt/process/config.py

@@ -0,0 +1,35 @@
+from typing import Dict
+import pydantic
+import json
+
+
+class VariableConfig(pydantic.BaseModel):
+    variable_name: str
+    attributes: Dict[str, object]
+    # if set, convert from this unit to the unit in attributes
+    assumed_input_units: str | None = None
+
+
+class SurfaceVariablesConfig(pydantic.BaseModel):
+    variables: Dict[str, VariableConfig]
+
+
+class PressureLevelVariablesConfig(pydantic.BaseModel):
+    levels: list[int]
+    variables: Dict[str, VariableConfig]
+
+
+class VariablesConfig(pydantic.BaseModel):
+    sfc: SurfaceVariablesConfig
+    pl: PressureLevelVariablesConfig
+
+
+class MkGridConfig(pydantic.BaseModel):
+    variables: VariablesConfig
+
+
+def open_config(config: str) -> MkGridConfig:
+    '''Load netCDF variable mappings from a JSON configuration file.'''
+    with open(config) as f:
+        config_json = json.load(f)
+        return MkGridConfig.model_validate(config_json)

bris-adapt/bris_adapt/process/configutil.py

@@ -0,0 +1,14 @@
+from typing import Final
+import os
+
+
+def find_config_file(filename: str) -> str:
+    '''Find configuration file in current or parent directories.'''
+    search_paths: Final = ['.', 'etc',
+                           'bris-adapt/etc', '/etc', '/usr/local/etc']
+    for path in search_paths:
+        full_path = f"{path}/{filename}"
+
+        if os.path.isfile(full_path):
+            return full_path
+    raise FileNotFoundError(f"Configuration file '{filename}' not found.")

bris-adapt/bris_adapt/process/interpolate.py

@@ -0,0 +1,162 @@
+import rasterio
+import numpy as np
+import xarray as xr
+from bris_adapt.process.ncutil import get_variable_by_standard_name
+from sklearn.neighbors import BallTree
+
+MEAN_EARTH_RADIUS_KM: float = 6371.0  # Mean Earth radius in km
+
+
+def create_target_grid_from_geotiff(topofile: str) -> tuple[np.ndarray, np.ndarray]:
+    '''
+    Create a target grid from a GeoTIFF file with elevation data.
+    The returned grid has the same latitudes and longitudes as the GeoTIFF.
+
+    Arguments:
+        topofile: path to the GeoTIFF file
+    Returns: (latitudes, longitudes) as 2D numpy arrays
+    '''
+    with rasterio.open_rasterio(topofile) as topo:
+        lat, lon = np.meshgrid(
+            topo['y'].values,
+            topo['x'].values,
+            indexing='ij'
+        )
+        return lat, lon
+
+
+def create_target_grid_from_area(area: tuple[float, float, float, float], resolution: float) -> tuple[np.ndarray, np.ndarray]:
+    '''
+    Create a target grid for the specified area and resolution.
+
+    If area is None, a global grid is created.
+
+    Arguments:
+        area: tuple of (north, west, south, east) in degrees
+        resolution: grid resolution in degrees
+
+    Returns: (latitudes, longitudes) as 2D numpy arrays
+    '''
+    north, west, south, east = area
+    lat = np.arange(north, south - resolution, -resolution)
+    lon = np.arange(west, east + resolution, resolution)
+    lat2d, lon2d = np.meshgrid(lat, lon, indexing='ij')  # (nlat, nlon)
+    return lat2d, lon2d
+
+
+class Interpolator:
+    def __init__(self, ds: xr.Dataset, target_grid2d: tuple[np.ndarray, np.ndarray], method: str = 'idw',
+                 k: int = 4, power: float = 2.0, radius_km: float | None = None, earth_km: float = MEAN_EARTH_RADIUS_KM):
+        '''Initialize the interpolator with source data and parameters. Precompute neighbor indices and distances. 
+        If target_grid_lat and target_grid_lon are provided, they will be used as the target grid instead of generating a new one.
+        If target_grid_lat and target_grid_lon are None, a global grid will be created based on the specified resolution.'''
+        self.ds = ds
+        self.method = method
+        self.k = k
+        self.power = power
+        self.radius_km = radius_km
+        self.earth_km = earth_km
+        self.tree = None
+        self.within = None
+        self.idx = None
+        self.dist = None
+        self.target_grid2d = target_grid2d
+
+        self.ntime = self.ds.sizes["time"]
+        lat_src = get_variable_by_standard_name(ds, "latitude")  # (values,)
+        lon_src = get_variable_by_standard_name(ds, "longitude")  # (values,)
+
+        # Wrap longitudes to [-180, 180)
+        lon_src = (lon_src + 180.0) % 360.0 - 180.0
+        self._setup(lat_src, lon_src)
+
+    def _setup(self, lat_src: np.ndarray, lon_src: np.ndarray):
+        lat2d, lon2d = self.target_grid2d
+
+        # ----------------- BUILD NEIGHBOR MAP (ONCE) -----------------
+        src_rad = np.deg2rad(np.c_[lat_src, lon_src])  # (values, 2)
+        target_rad = np.deg2rad(
+            np.c_[lat2d.ravel(), lon2d.ravel()])      # (ncell,  2)
+        print("target grid shape:", target_rad.shape)
+        self.tree = BallTree(src_rad, metric="haversine")
+
+        # k for nearest-only can be 1; for IDW use >= 2
+        if self.method == "nearest" and self.radius_km is None:
+            self.kq = 1
+        else:
+            self.kq = max(1, self.k)
+
+        # dist in radians, sorted by distance
+        self.dist, self.idx = self.tree.query(
+            target_rad, k=self.kq, return_distance=True)  # (ncell, kq)
+
+        if self.radius_km is not None:
+            radius_rad = self.radius_km / self.earth_km
+            self.within = self.dist <= radius_rad
+        else:
+            self.within = np.ones_like(self.dist, dtype=bool)
+
+    def latitude(self) -> np.ndarray:
+        return self.target_grid2d[0][:, 0]
+
+    def longitude(self) -> np.ndarray:
+        return self.target_grid2d[1][0, :]
+
+    def interpolate_var(self, var_da: xr.DataArray) -> np.ndarray:  # type: ignore
+        """Map (time, values) -> (time, nlat, nlon) using precomputed idx/dist."""
+        nlat, nlon = self.target_grid2d[0].shape
+        ncell = nlat * nlon
+
+        ntime = self.ntime
+        out = np.full((ntime, nlat, nlon), np.nan, dtype=np.float32)
+
+        for ti in range(ntime):
+            src = var_da.isel(time=ti).values  # (values,)
+            vals = src[self.idx]               # (ncell, kq)
+            valid_vals = np.isfinite(vals)
+            valid = valid_vals & self.within
+
+            if self.method == "nearest":
+                # choose nearest *within radius* (or NaN if none)
+                # first True along axis (neighbors are distance-sorted)
+                choose = np.argmax(valid, axis=1)
+                has_any = valid.any(axis=1)
+                res_flat = np.full(ncell, np.nan, dtype=np.float32)
+                if self.kq == 1:
+                    res_flat = np.where(
+                        has_any, vals[:, 0], np.nan).astype(np.float32)
+                else:
+                    res_flat[has_any] = vals[np.arange(
+                        ncell)[has_any], choose[has_any]].astype(np.float32)
+            elif self.method == "idw":
+                # IDW over neighbors (common, smooth)
+                eps = 1e-12
+                # (ncell, kq)
+                w = 1.0 / (self.dist + eps)**self.power
+                w *= valid
+                wsum = w.sum(axis=1)
+                # Replace NaNs with 0 for multiplication, they'll be masked by weights
+                num = np.nansum(w * np.nan_to_num(vals), axis=1)
+                res_flat = np.where(wsum > 0, num / wsum,
+                                    np.nan).astype(np.float32)
+            else:
+                raise ValueError(
+                    f"Unknown interpolation method: {self.method}")
+
+            out[ti] = res_flat.reshape(nlat, nlon)
+
+        return out
+
+    def interpolate_var_name(self, name: str) -> np.ndarray:
+        da = self.get_var(name)
+        return self.interpolate_var(da)
+
+    def shape(self) -> tuple[int, int]:
+        return self.target_grid2d[0].shape
+
+    def get_var(self, name: str) -> xr.DataArray:
+        if name in self.ds:
+            return self.ds[name]
+        if ("\\" + name) in self.ds:
+            return self.ds["\\" + name]  # handles ncdump-escaped names
+        raise KeyError(f"Variable '{name}' not found")

bris-adapt/bris_adapt/process/ncutil.py

@@ -0,0 +1,21 @@
+import xarray as xr
+import numpy as np
+
+
+def get_variable_name_by_standard_name(ds: xr.Dataset, standard_name: str) -> str:
+    for var in ds.variables:
+        if "standard_name" in ds[var].attrs and ds[var].attrs["standard_name"] == standard_name:
+            return var  # type: ignore
+
+    # Fallback to long_name if standard_name not found
+    for var in ds.variables:
+        if "long_name" in ds[var].attrs and ds[var].attrs["long_name"] == standard_name:
+            return var  # type: ignore
+
+    raise ValueError(f"Variable with standard_name {standard_name} not found")
+
+
+def get_variable_by_standard_name(ds: xr.Dataset, standard_name: str, dtype: np.dtype = np.float64) -> np.ndarray:  # type: ignore
+    '''Get a variable from an xarray Dataset by its standard_name attribute and convert it to the specified dtype.'''
+    var = get_variable_name_by_standard_name(ds, standard_name)
+    return ds[var].values.astype(dtype)  # type: ignore

bris-adapt/bris_adapt/scripts/process/__init__.py

@@ -1,9 +1,15 @@
 import click
 from .make_grid import make_grid
+from .mkglobal_grid import mkglobal_grid
+from .create_image import create_image
+
 
 @click.group()
 def process():
     '''Manipulate FIAB output files.'''
     pass
 
-process.add_command(make_grid)
+
+process.add_command(make_grid)
+process.add_command(mkglobal_grid)
+process.add_command(create_image)