add get-grid

CHANGELOG.md

@@ -14,6 +14,7 @@ Keep it human-readable, your future self will thank you!
 
 - Add regrid filter
 - Added repeat-member #18
+- Add `get-grid` command
 
 ## [0.1.0](https://github.com/ecmwf/anemoi-utils/transform/0.0.5...HEAD/compare/0.0.8...0.1.0) - 2024-11-18
 

src/anemoi/transform/commands/get-grid.py

@@ -0,0 +1,31 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+
+from . import Command
+
+
+class GetGrid(Command):
+    """Extract the grid from a GRIB or NetCDF file and save it as a NPZ file."""
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument("--source", default="file", help="EKD Source type (default: file).")
+        command_parser.add_argument("input", help="Input file (GRIB or NetCDF).")
+        command_parser.add_argument("output", help="Output NPZ file.")
+
+    def run(self, args):
+        import numpy as np
+        from earthkit.data import from_source
+
+        ds = from_source(args.source, args.input)
+        lat, lon = ds[0].grid_points()
+        np.savez(args.output, latitudes=lat, longitudes=lon)
+
+
+command = GetGrid

src/anemoi/transform/fields.py

@@ -29,13 +29,21 @@ def __init__(self, field):
         self._field = field
 
     def __getattr__(self, name):
+        if name in (
+            "clone",
+            "copy",
+        ):
+            raise AttributeError(f"NewField: forwarding of `{name}` is not supported")
+
         if name not in (
             "mars_area",
             "mars_grid",
             "to_numpy",
             "metadata",
+            "shape",
         ):
             LOG.warning(f"NewField: forwarding `{name}`")
+
         return getattr(self._field, name)
 
     def __repr__(self) -> str:
@@ -61,6 +69,18 @@ def to_numpy(self, flatten=False, dtype=None, index=None):
         return data
 
 
+class NewGridField(WrappedField):
+    """Change the grid of a field."""
+
+    def __init__(self, field, latitudes, longitudes):
+        super().__init__(field)
+        self._latitudes = latitudes
+        self._longitudes = longitudes
+
+    def grid_points(self):
+        return self._latitudes, self._longitudes
+
+
 class NewMetadataField(WrappedField):
     """Change the metadata of a field."""
 
@@ -89,7 +109,7 @@ class NewValidDateTimeField(NewMetadataField):
 
     def __init__(self, field, valid_datetime):
         date = int(valid_datetime.date().strftime("%Y%m%d"))
-        assert valid_datetime.time().minute == 0, valid_datetime.time()
+        assert valid_datetime.time().minute == 0, valid_datetime
         time = valid_datetime.time().hour
 
         self.valid_datetime = valid_datetime
@@ -107,3 +127,7 @@ def new_field_with_valid_datetime(template, date):
 
 def new_field_with_metadata(template, **metadata):
     return NewMetadataField(template, **metadata)
+
+
+def new_field_from_latitudes_longitudes(template, latitudes, longitudes):
+    return NewGridField(template, latitudes, longitudes)

src/anemoi/transform/filters/regrid.py

@@ -13,6 +13,7 @@
 
 from earthkit.data.core.fieldlist import Field
 
+from ..fields import new_field_from_latitudes_longitudes
 from ..fields import new_field_from_numpy
 from ..fields import new_fieldlist_from_list
 from ..filter import Filter
@@ -24,21 +25,29 @@
 def as_gridspec(grid):
     if grid is None:
         return None
+
     if isinstance(grid, str):
         return {"grid": grid}
+
     return grid
 
 
 def as_griddata(grid):
+    if grid is None:
+        return None
+
     if isinstance(grid, Field):
-        lat, lon = grid.gridpoints()
+        lat, lon = grid.grid_points()
         return dict(latitudes=lat, longitudes=lon)
+
     if isinstance(grid, dict) and "latitudes" in grid and "longitudes" in grid:
         return grid
+
     if isinstance(grid, str):
         from anemoi.utils.grids import grids
 
         return grids(grid)
+
     raise ValueError(f"Invalid grid: {grid}")
 
 
@@ -124,7 +133,7 @@ def __call__(self, field):
         assert data.shape == self.ingrid["longitudes"].shape, (data.shape, self.ingrid["longitudes"].shape)
 
         data = data[..., self.nearest_grid_points]
-        return new_field_from_numpy(data, template=field)
+        return new_field_from_latitudes_longitudes(new_field_from_numpy(data, template=field), **self.outgrid)
 
 
 def interpolator(in_grid, out_grid, method):