add snowcover and land filters : Merge pull request #9 from ecmwf/fea…

…ture/snowmanip

add snowcover filter

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/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")

src/anemoi/transform/filters/land_parameters.py

@@ -0,0 +1,118 @@
+# (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 numpy as np
+
+from . import filter_registry
+from .base import SimpleFilter
+
+SOIL_TYPE_DIC = {
+    0: {"theta_pwp": 0, "theta_cap": 0},
+    1: {"theta_pwp": 0.059, "theta_cap": 0.244},
+    2: {"theta_pwp": 0.151, "theta_cap": 0.347},
+    3: {"theta_pwp": 0.133, "theta_cap": 0.383},
+    4: {"theta_pwp": 0.279, "theta_cap": 0.448},
+    5: {"theta_pwp": 0.335, "theta_cap": 0.541},
+    6: {"theta_pwp": 0.267, "theta_cap": 0.663},
+    7: {"theta_pwp": 0.151, "theta_cap": 0.347},
+}
+
+VEG_TYPE_DIC = {
+    0: {"veg_rsmin": 250.0, "veg_cov": 0.0, "veg_z0m": 0.013},
+    1: {"veg_rsmin": 125.0, "veg_cov": 0.9, "veg_z0m": 0.25},
+    2: {"veg_rsmin": 80.0, "veg_cov": 0.85, "veg_z0m": 0.1},
+    3: {"veg_rsmin": 395.0, "veg_cov": 0.9, "veg_z0m": 2.0},
+    4: {"veg_rsmin": 320.0, "veg_cov": 0.9, "veg_z0m": 2.0},
+    5: {"veg_rsmin": 215.0, "veg_cov": 0.9, "veg_z0m": 2.0},
+    6: {"veg_rsmin": 320.0, "veg_cov": 0.99, "veg_z0m": 2.0},
+    7: {"veg_rsmin": 100.0, "veg_cov": 0.7, "veg_z0m": 0.5},
+    8: {"veg_rsmin": 250.0, "veg_cov": 0.0, "veg_z0m": 0.013},
+    9: {"veg_rsmin": 45.0, "veg_cov": 0.5, "veg_z0m": 0.03},
+    10: {"veg_rsmin": 110.0, "veg_cov": 0.9, "veg_z0m": 0.5},
+    11: {"veg_rsmin": 45.0, "veg_cov": 0.1, "veg_z0m": 0.03},
+    12: {"veg_rsmin": 0.0, "veg_cov": 0.0, "veg_z0m": 0.0013},
+    13: {"veg_rsmin": 130.0, "veg_cov": 0.6, "veg_z0m": 0.25},
+    14: {"veg_rsmin": 0.0, "veg_cov": 0.0, "veg_z0m": 0.0001},
+    15: {"veg_rsmin": 0.0, "veg_cov": 0.0, "veg_z0m": 0.0001},
+    16: {"veg_rsmin": 230.0, "veg_cov": 0.5, "veg_z0m": 0.5},
+    17: {"veg_rsmin": 110.0, "veg_cov": 0.4, "veg_z0m": 0.1},
+    18: {"veg_rsmin": 180.0, "veg_cov": 0.9, "veg_z0m": 1.50},
+    19: {"veg_rsmin": 175.0, "veg_cov": 0.9, "veg_z0m": 1.1},
+    20: {"veg_rsmin": 150.0, "veg_cov": 0.6, "veg_z0m": 0.02},
+}
+
+
+def read_crosswalking_table(param, param_dic):
+    arrays = [np.array([param_dic[x][key] for x in param]) for key in param_dic[0].keys()]
+    return arrays
+
+
+@filter_registry.register("land_parameters")
+class LandParameters(SimpleFilter):
+    """A filter to add static parameters from table based on soil/vegetation type."""
+
+    def __init__(
+        self,
+        *,
+        # Input parameters
+        high_veg_type="tvh",
+        low_veg_type="tvl",
+        soil_type="slt",
+        # Output parameters
+        hveg_rsmin="hveg_rsmin",
+        hveg_cov="hveg_cov",
+        hveg_z0m="hveg_z0m",
+        lveg_rsmin="lveg_rsmin",
+        lveg_cov="lveg_cov",
+        lveg_z0m="lveg_z0m",
+        theta_pwp="theta_pwp",
+        theta_cap="theta_cap",
+    ):
+        self.high_veg_type = high_veg_type
+        self.low_veg_type = low_veg_type
+        self.soil_type = soil_type
+        self.hveg_rsmin = hveg_rsmin
+        self.hveg_cov = hveg_cov
+        self.hveg_z0m = hveg_z0m
+        self.lveg_rsmin = lveg_rsmin
+        self.lveg_cov = lveg_cov
+        self.lveg_z0m = lveg_z0m
+        self.theta_pwp = theta_pwp
+        self.theta_cap = theta_cap
+
+    def forward(self, data):
+        return self._transform(
+            data,
+            self.forward_transform,
+            self.high_veg_type,
+            self.low_veg_type,
+            self.soil_type,
+        )
+
+    def backward(self, data):
+        raise NotImplementedError("LandParameters is not reversible")
+
+    def forward_transform(self, tvh, tvl, sotype):
+        """Get static parameters from table based on soil/vegetation type"""
+
+        hveg_rsmin, hveg_cov, hveg_z0m = read_crosswalking_table(tvh.to_numpy(), VEG_TYPE_DIC)
+        lveg_rsmin, lveg_cov, lveg_z0m = read_crosswalking_table(tvl.to_numpy(), VEG_TYPE_DIC)
+        theta_pwp, theta_cap = read_crosswalking_table(sotype.to_numpy(), SOIL_TYPE_DIC)
+
+        yield self.new_field_from_numpy(hveg_rsmin, template=tvh, param=self.hveg_rsmin)
+        yield self.new_field_from_numpy(hveg_cov, template=tvh, param=self.hveg_cov)
+        yield self.new_field_from_numpy(hveg_z0m, template=tvh, param=self.hveg_z0m)
+        yield self.new_field_from_numpy(lveg_rsmin, template=tvl, param=self.lveg_rsmin)
+        yield self.new_field_from_numpy(lveg_cov, template=tvl, param=self.lveg_cov)
+        yield self.new_field_from_numpy(lveg_z0m, template=tvl, param=self.lveg_z0m)
+        yield self.new_field_from_numpy(theta_pwp, template=sotype, param=self.theta_pwp)
+        yield self.new_field_from_numpy(theta_cap, template=sotype, param=self.theta_cap)
+
+    def backward_transform(self, tvh, tvl, sotype):
+        raise NotImplementedError("LandParameters is not reversible")

src/anemoi/transform/filters/snow_cover.py

@@ -0,0 +1,59 @@
+# (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 numpy as np
+
+from . import filter_registry
+from .base import SimpleFilter
+
+
+def compute_snow_cover(snow_depth, snow_density):
+    """Convert snow depth to snow cover."""
+    tmp1 = (1000 * snow_depth) / snow_density
+    tmp2 = np.clip(snow_density, 100, 400)
+    snow_cover = np.clip(np.tanh((4000 * tmp1) / tmp2), 0, 1)
+    snow_cover[snow_cover > 0.99] = 1.0
+    return snow_cover
+
+
+@filter_registry.register("snow_cover")
+class SnowCover(SimpleFilter):
+    """A filter to compute snow cover from snow density and snow depth."""
+
+    def __init__(
+        self,
+        *,
+        snow_depth="sd",
+        snow_density="rsn",
+        snow_cover="snowc",
+    ):
+        self.snow_depth = snow_depth
+        self.snow_density = snow_density
+        self.snow_cover = snow_cover
+
+    def forward(self, data):
+        return self._transform(
+            data,
+            self.forward_transform,
+            self.snow_depth,
+            self.snow_density,
+        )
+
+    def backward(self, data):
+        raise NotImplementedError("SnowCover is not reversible")
+
+    def forward_transform(self, sd, rsn):
+        """Convert snow depth and snow density to snow cover"""
+
+        snow_cover = compute_snow_cover(sd.to_numpy(), rsn.to_numpy())
+
+        yield self.new_field_from_numpy(snow_cover, template=sd, param=self.snow_cover)
+
+    def backward_transform(self, sd, rsn):
+        raise NotImplementedError("SnowCover is not reversible")

src/anemoi/transform/grouping/__init__.py

@@ -45,6 +45,8 @@ def iterate(self, data, *, other=_lost):
 
         for _, group in groups.items():
             if len(group) != len(self.params):
-                raise ValueError("Missing component")
+                for p in data:
+                    print(p)
+                raise ValueError(f"Missing component. Want {sorted(self.params)}, got {sorted(group.keys())}")
 
             yield tuple(group[p] for p in self.params)