prepare 0.1.0: Merge pull request #13 from ecmwf/develop

prepare 0.1.0

.pre-commit-config.yaml

@@ -59,11 +59,6 @@ repos:
   hooks:
   - id: rstfmt
     exclude: 'cli/.*' # Because we use argparse
-- repo: https://github.com/b8raoult/pre-commit-docconvert
-  rev: "0.1.5"
-  hooks:
-  - id: docconvert
-    args: ["numpy"]
 - repo: https://github.com/tox-dev/pyproject-fmt
   rev: "2.2.3"
   hooks:

pyproject.toml

@@ -1,14 +1,12 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (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.
 
-# https://packaging.python.org/en/latest/guides/writing-pyproject-toml/
-
 [build-system]
 requires = [ "setuptools>=60", "setuptools-scm>=8" ]
 

src/anemoi/transform/__init__.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.

src/anemoi/transform/__main__.py

@@ -1,12 +1,11 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (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 anemoi.utils.cli import cli_main
 from anemoi.utils.cli import make_parser

src/anemoi/transform/commands/__init__.py

@@ -1,12 +1,11 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (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.
-#
 
 import os
 

src/anemoi/transform/data/__init__.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.

src/anemoi/transform/fields.py

@@ -0,0 +1,106 @@
+# (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.
+
+import logging
+
+from earthkit.data.indexing.fieldlist import FieldArray
+
+LOG = logging.getLogger(__name__)
+
+
+def new_fieldlist_from_list(fields):
+    return FieldArray(fields)
+
+
+def new_empty_fieldlist():
+    return FieldArray([])
+
+
+class WrappedField:
+    """A wrapper around a earthkit-data field object."""
+
+    def __init__(self, field):
+        self._field = field
+
+    def __getattr__(self, name):
+        if name not in (
+            "mars_area",
+            "mars_grid",
+            "to_numpy",
+            "metadata",
+        ):
+            LOG.warning(f"NewField: forwarding `{name}`")
+        return getattr(self._field, name)
+
+    def __repr__(self) -> str:
+        return repr(self._field)
+
+
+class NewDataField(WrappedField):
+    """Change the data of a field."""
+
+    def __init__(self, field, data):
+        super().__init__(field)
+        self._data = data
+        self.shape = data.shape
+
+    def to_numpy(self, flatten=False, dtype=None, index=None):
+        data = self._data
+        if dtype is not None:
+            data = data.astype(dtype)
+        if flatten:
+            data = data.flatten()
+        if index is not None:
+            data = data[index]
+        return data
+
+
+class NewMetadataField(WrappedField):
+    """Change the metadata of a field."""
+
+    def __init__(self, field, **kwargs):
+        super().__init__(field)
+        self._metadata = kwargs
+
+    def metadata(self, *args, **kwargs):
+
+        if kwargs.get("namespace"):
+            assert kwargs.get("namespace") == "mars", kwargs
+            assert len(args) == 0, (args, kwargs)
+            mars = self._field.metadata(**kwargs).copy()
+            for k in list(mars.keys()):
+                if k in self._metadata:
+                    mars[k] = self._metadata[k]
+            return mars
+
+        if len(args) == 1 and args[0] in self._metadata:
+            return self._metadata[args[0]]
+
+        return self._field.metadata(*args, **kwargs)
+
+
+class NewValidDateTimeField(NewMetadataField):
+    """Change the valid_datetime of a field."""
+
+    def __init__(self, field, valid_datetime):
+        date = int(valid_datetime.date().strftime("%Y%m%d"))
+        assert valid_datetime.time().minute == 0, valid_datetime.time()
+        time = valid_datetime.time().hour
+
+        self.valid_datetime = valid_datetime
+
+        super().__init__(field, date=date, time=time, step=0, valid_datetime=valid_datetime.isoformat())
+
+
+def new_field_from_numpy(array, *, template, **metadata):
+    return NewMetadataField(NewDataField(template, array), **metadata)
+
+
+def new_field_with_valid_datetime(template, date):
+    return NewValidDateTimeField(template, date)

src/anemoi/transform/filters/__init__.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.

src/anemoi/transform/filters/base.py

@@ -8,13 +8,16 @@
 # nor does it submit to any jurisdiction.
 
 
+import logging
 from abc import abstractmethod
 
-import earthkit.data as ekd
-
+from ..fields import new_field_from_numpy
+from ..fields import new_fieldlist_from_list
 from ..filter import Filter
 from ..grouping import GroupByMarsParam
 
+LOG = logging.getLogger(__name__)
+
 
 class SimpleFilter(Filter):
     """A filter to convert only some fields.
@@ -35,14 +38,10 @@ def _transform(self, data, transform, *group_by):
 
     def new_field_from_numpy(self, array, *, template, param):
         """Create a new field from a numpy array."""
-        md = template.metadata().override(shortName=param)
-        # return ekd.ArrayField(array, md)
-        return ekd.FieldList.from_array(array, md)[0]
+        return new_field_from_numpy(array, template=template, param=param)
 
     def new_fieldlist_from_list(self, fields):
-        from earthkit.data.indexing.fieldlist import FieldArray
-
-        return FieldArray(fields)
+        return new_fieldlist_from_list(fields)
 
     @abstractmethod
     def forward_transform(self, *fields):

src/anemoi/transform/filters/glacier_mask.py

@@ -0,0 +1,55 @@
+# (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.
+
+import earthkit.data as ekd
+import numpy as np
+
+from . import filter_registry
+from .base import SimpleFilter
+
+
+def mask_glaciers(snow_depth, glacier_mask):
+    snow_depth[glacier_mask] = np.nan
+    return snow_depth
+
+
+@filter_registry.register("glacier_mask")
+class SnowDepthMasked(SimpleFilter):
+    """A filter to mask about glacier in snow depth."""
+
+    def __init__(
+        self,
+        *,
+        glacier_mask,
+        snow_depth="sd",
+        snow_depth_masked="sd_masked",
+    ):
+        self.glacier_mask = ekd.from_source("file", glacier_mask)[0].to_numpy().astype(bool)
+        self.snow_depth = snow_depth
+        self.snow_depth_masked = snow_depth_masked
+
+    def forward(self, data):
+        return self._transform(
+            data,
+            self.forward_transform,
+            self.snow_depth,
+        )
+
+    def backward(self, data):
+        raise NotImplementedError("SnowDepthMasked is not reversible")
+
+    def forward_transform(self, sd):
+        """Mask out glaciers in snow depth"""
+
+        snow_depth_masked = mask_glaciers(sd.to_numpy(), self.glacier_mask)
+
+        yield self.new_field_from_numpy(snow_depth_masked, template=sd, param=self.snow_depth_masked)
+
+    def backward_transform(self, sd):
+        raise NotImplementedError("SnowDepthMasked is not reversible")