feat: adding-general-purpose-regridding

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 2035 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/commands/make-regrid-matrix.py

@@ -0,0 +1,76 @@
+# (C) Copyright 2025 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.
+
+
+import os
+
+from . import Command
+
+
+class MakeRegridMatrix(Command):
+    """Extract the grid from a pair GRIB or NetCDF files extract the MIR interpolation matrix to be used
+    by earthkit-regrid."""
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument("--source1", default="file", help="EKD Source type (default: file).")
+        command_parser.add_argument("--source2", default="file", help="EKD Source type (default: file).")
+
+        command_parser.add_argument("--mir", default=os.environ.get("MIR_COMMAND", "mir"), help="MIR command")
+
+        command_parser.add_argument("input1", help="Input file (GRIB or NetCDF).")
+        command_parser.add_argument("input2", help="Input file (GRIB or NetCDF).")
+
+        command_parser.add_argument("output", help="Output NPZ file.")
+
+        command_parser.add_argument("kwargs", nargs="*", help="MIR arguments.")
+
+    def run(self, args):
+        import numpy as np
+        from earthkit.data import from_source
+        from earthkit.regrid.utils.mir import mir_make_matrix
+
+        _, ext1 = os.path.splitext(args.input1)
+        if ext1 in (".npz", ".npy"):
+            ds1 = np.load(args.input1)
+            lat1 = ds1["latitudes"]
+            lon1 = ds1["longitudes"]
+        else:
+            ds1 = from_source(args.source1, args.input1)
+            lat1, lon1 = ds1[0].grid_points()
+
+        _, ext2 = os.path.splitext(args.input2)
+        if ext2 in (".npz", ".npy"):
+            ds2 = np.load(args.input2)
+            lat2 = ds2["latitudes"]
+            lon2 = ds2["longitudes"]
+        else:
+            ds2 = from_source(args.source2, args.input2)
+            lat2, lon2 = ds2[0].grid_points()
+
+        kwargs = {}
+        for arg in args.kwargs:
+            key, value = arg.split("=")
+            kwargs[key] = value
+
+        sparse_array = mir_make_matrix(lat1, lon1, lat2, lon2, output=None, mir=args.mir, **kwargs)
+
+        np.savez(
+            args.output,
+            matrix_data=sparse_array.data,
+            matrix_indices=sparse_array.indices,
+            matrix_indptr=sparse_array.indptr,
+            matrix_shape=sparse_array.shape,
+            in_latitudes=lat1,
+            in_longitudes=lon1,
+            out_latitudes=lat2,
+            out_longitudes=lon2,
+        )
+
+
+command = MakeRegridMatrix

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)