Add variable

src/anemoi/transform/variables/__init__.py

@@ -0,0 +1,49 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 abc import ABC
+from abc import abstractmethod
+
+
+class Variable(ABC):
+    """Variable is a class that represents a variable during
+    training and inference.
+    """
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        from .variables import VariableFromDict
+
+        return VariableFromDict(data)
+
+    @classmethod
+    def from_earthkit(cls, field):
+        from .variables import VariableFromEarthkit
+
+        return VariableFromEarthkit(field)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.name})"
+
+    def __hash__(self) -> int:
+        return hash(self.name)
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, Variable):
+            return False
+        return self.name == other.name
+
+    @abstractmethod
+    def is_pressure_level(self):
+        pass
+
+    @abstractmethod
+    def level(self):
+        pass

src/anemoi/transform/variables/variables.py

@@ -0,0 +1,47 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 Variable
+
+
+class VariableFromMarsVocabulary(Variable):
+    """A variable that is defined by the Mars vocabulary."""
+
+    def __init__(self, name, data: dict) -> None:
+        super().__init__(name)
+        self.data = data
+
+    def is_pressure_level(self):
+        return self.data.get("levtype", None) == "pl"
+
+    def level(self):
+        return self.data.get("levelist", None)
+
+
+class VariableFromDict(VariableFromMarsVocabulary):
+    """A variable that is defined by a user provided dictionary."""
+
+    def __init__(self, name, data: dict) -> None:
+
+        if "mars" in data:
+            data = data["mars"]
+
+        super().__init__(name, data)
+
+
+class VariableFromEarthkit(VariableFromMarsVocabulary):
+    """A variable that is defined by an EarthKit field."""
+
+    def __init__(self, name, field, namespace="mars") -> None:
+        super().__init__(name, field.metadata(namespace=namespace))
+        self.field = field
+
+    def is_pressure_level(self):
+        return self.field.is_pressure_level()
+
+    def level(self):
+        return self.field.level()