Merge branch 'main' into develop

.gitignore

@@ -177,3 +177,4 @@ bar
 dev/
 *.out
 _version.py
+*.tar

CITATION.cff

@@ -0,0 +1,33 @@
+# This CITATION.cff file was generated with cffinit.
+# Visit https://bit.ly/cffinit to generate yours today!
+
+cff-version: 1.2.0
+title: ai-models
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file.
+type: software
+authors:
+  - given-names: Baudouin
+    family-names: Raoult
+    affiliation: ECMWF
+  - given-names: Florian
+    family-names: Pinault
+  - given-names: Gert
+    family-names: Mertes
+    affiliation: ECMWF
+  - given-names: Jesper Sören
+    family-names: Dramsch
+    affiliation: ECMWF
+    orcid: 'https://orcid.org/0000-0001-8273-905X'
+  - given-names: Matthew
+    family-names: Chantry
+    affiliation: ECMWF
+repository-code: 'https://github.com/ecmwf-lab/ai-models'
+abstract: >-
+  ai-models is used to run AI-based weather forecasting
+  models. These models need to be installed independently.
+keywords:
+  - ai
+  - weather forecasting
+license: Apache-2.0

README.md

@@ -1,5 +1,11 @@
 # ai-models
 
+**DISCLAIMER**
+This project is **BETA** and will be **Experimental** for the foreseeable future.
+Interfaces and functionality are likely to change, and the project itself may be scrapped.
+**DO NOT** use this software in any project/software that is operational.
+
+
 The `ai-models` command is used to run AI-based weather forecasting models. These models need to be installed independently.
 
 ## Usage
@@ -195,3 +201,25 @@ It has the following options:
 - `--expver EXPVER`: The experiment version of the model output.
 - `--class CLASS`: The 'class' metadata of the model output.
 - `--metadata KEY=VALUE`: Additional metadata metadata in the model output
+
+## License
+
+```
+Copyright 2022, European Centre for Medium Range Weather Forecasts.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+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.
+```

pyproject.toml

@@ -40,22 +40,26 @@ classifiers = [
 ]
 
 dependencies = [
-    "entrypoints",
-    "requests",
-    "earthkit-data>=0.10.1",
+    "cdsapi",
+    "earthkit-data>=0.10.3",
+    "earthkit-meteo",
+    "earthkit-regrid",
     "eccodes>=2.37",
-    "multiurl",
+    "ecmwf-api-client",
+    "ecmwf-opendata",
+    "entrypoints",
     "gputil",
-    "earthkit-meteo",
+    "multiurl",
     "pyyaml",
+    "requests",
     "tqdm",
 ]
 
 
 [project.urls]
-Homepage = "https://github.com/ecmwf/ai-models/"
-Repository = "https://github.com/ecmwf/ai-models/"
-Issues = "https://github.com/ecmwf/ai-models/issues"
+Homepage = "https://github.com/ecmwf-lab/ai-models/"
+Repository = "https://github.com/ecmwf-lab/ai-models/"
+Issues = "https://github.com/ecmwf-lab/ai-models/issues"
 
 [project.scripts]
 ai-models = "ai_models.__main__:main"
@@ -64,10 +68,11 @@ ai-models = "ai_models.__main__:main"
 version_file = "src/ai_models/_version.py"
 
 [project.entry-points."ai_models.input"]
-file = "ai_models.inputs:FileInput"
-mars = "ai_models.inputs:MarsInput"
-cds = "ai_models.inputs:CdsInput"
-opendata = "ai_models.inputs:OpenDataInput"
+file = "ai_models.inputs.file:FileInput"
+mars = "ai_models.inputs.mars:MarsInput"
+cds = "ai_models.inputs.cds:CdsInput"
+ecmwf-open-data = "ai_models.inputs.opendata:OpenDataInput"
+opendata = "ai_models.inputs.opendata:OpenDataInput"
 
 [project.entry-points."ai_models.output"]
 file = "ai_models.outputs:FileOutput"

src/ai_models/inputs/__init__.py

@@ -9,198 +9,13 @@
 from functools import cached_property
 
 import earthkit.data as ekd
+import earthkit.regrid as ekr
 import entrypoints
+from earthkit.data.indexing.fieldlist import FieldArray
 
 LOG = logging.getLogger(__name__)
 
 
-class RequestBasedInput:
-    def __init__(self, owner, **kwargs):
-        self.owner = owner
-
-    def _patch(self, **kargs):
-        r = dict(**kargs)
-        self.owner.patch_retrieve_request(r)
-        return r
-
-    @cached_property
-    def fields_sfc(self):
-        param = self.owner.param_sfc
-        if not param:
-            return ekd.from_source("empty")
-
-        LOG.info(f"Loading surface fields from {self.WHERE}")
-
-        return ekd.from_source(
-            "multi",
-            [
-                self.sfc_load_source(
-                    **self._patch(
-                        date=date,
-                        time=time,
-                        param=param,
-                        grid=self.owner.grid,
-                        area=self.owner.area,
-                        **self.owner.retrieve,
-                    )
-                )
-                for date, time in self.owner.datetimes()
-            ],
-        )
-
-    @cached_property
-    def fields_pl(self):
-        param, level = self.owner.param_level_pl
-        if not (param and level):
-            return ekd.from_source("empty")
-
-        LOG.info(f"Loading pressure fields from {self.WHERE}")
-        return ekd.from_source(
-            "multi",
-            [
-                self.pl_load_source(
-                    **self._patch(
-                        date=date,
-                        time=time,
-                        param=param,
-                        level=level,
-                        grid=self.owner.grid,
-                        area=self.owner.area,
-                    )
-                )
-                for date, time in self.owner.datetimes()
-            ],
-        )
-
-    @cached_property
-    def fields_ml(self):
-        param, level = self.owner.param_level_ml
-        if not (param and level):
-            return ekd.from_source("empty")
-
-        LOG.info(f"Loading model fields from {self.WHERE}")
-        return ekd.from_source(
-            "multi",
-            [
-                self.ml_load_source(
-                    **self._patch(
-                        date=date,
-                        time=time,
-                        param=param,
-                        level=level,
-                        grid=self.owner.grid,
-                        area=self.owner.area,
-                    )
-                )
-                for date, time in self.owner.datetimes()
-            ],
-        )
-
-    @cached_property
-    def all_fields(self):
-        return self.fields_sfc + self.fields_pl + self.fields_ml
-
-
-class MarsInput(RequestBasedInput):
-    WHERE = "MARS"
-
-    def __init__(self, owner, **kwargs):
-        self.owner = owner
-
-    def pl_load_source(self, **kwargs):
-        kwargs["levtype"] = "pl"
-        logging.debug("load source mars %s", kwargs)
-        return ekd.from_source("mars", kwargs)
-
-    def sfc_load_source(self, **kwargs):
-        kwargs["levtype"] = "sfc"
-        logging.debug("load source mars %s", kwargs)
-        return ekd.from_source("mars", kwargs)
-
-    def ml_load_source(self, **kwargs):
-        kwargs["levtype"] = "ml"
-        logging.debug("load source mars %s", kwargs)
-        return ekd.from_source("mars", kwargs)
-
-
-class CdsInput(RequestBasedInput):
-    WHERE = "CDS"
-
-    def pl_load_source(self, **kwargs):
-        kwargs["product_type"] = "reanalysis"
-        return ekd.from_source("cds", "reanalysis-era5-pressure-levels", kwargs)
-
-    def sfc_load_source(self, **kwargs):
-        kwargs["product_type"] = "reanalysis"
-        return ekd.from_source("cds", "reanalysis-era5-single-levels", kwargs)
-
-    def ml_load_source(self, **kwargs):
-        raise NotImplementedError("CDS does not support model levels")
-
-
-class OpenDataInput(RequestBasedInput):
-    WHERE = "OPENDATA"
-
-    RESOLS = {(0.25, 0.25): "0p25"}
-
-    def __init__(self, owner, **kwargs):
-        self.owner = owner
-
-    def _adjust(self, kwargs):
-        if "level" in kwargs:
-            # OpenData uses levelist instead of level
-            kwargs["levelist"] = kwargs.pop("level")
-
-        grid = kwargs.pop("grid")
-        if isinstance(grid, list):
-            grid = tuple(grid)
-
-        kwargs["resol"] = self.RESOLS[grid]
-        r = dict(**kwargs)
-        r.update(self.owner.retrieve)
-        return r
-
-    def pl_load_source(self, **kwargs):
-        self._adjust(kwargs)
-        kwargs["levtype"] = "pl"
-        logging.debug("load source ecmwf-open-data %s", kwargs)
-        return ekd.from_source("ecmwf-open-data", **kwargs)
-
-    def sfc_load_source(self, **kwargs):
-        self._adjust(kwargs)
-        kwargs["levtype"] = "sfc"
-        logging.debug("load source ecmwf-open-data %s", kwargs)
-        return ekd.from_source("ecmwf-open-data", **kwargs)
-
-    def ml_load_source(self, **kwargs):
-        self._adjust(kwargs)
-        kwargs["levtype"] = "ml"
-        logging.debug("load source ecmwf-open-data %s", kwargs)
-        return ekd.from_source("ecmwf-open-data", **kwargs)
-
-
-class FileInput:
-    def __init__(self, owner, file, **kwargs):
-        self.file = file
-        self.owner = owner
-
-    @cached_property
-    def fields_sfc(self):
-        return self.all_fields.sel(levtype="sfc")
-
-    @cached_property
-    def fields_pl(self):
-        return self.all_fields.sel(levtype="pl")
-
-    @cached_property
-    def fields_ml(self):
-        return self.all_fields.sel(levtype="ml")
-
-    @cached_property
-    def all_fields(self):
-        return ekd.from_source("file", self.file)
-
-
 def get_input(name, *args, **kwargs):
     return available_inputs()[name].load()(*args, **kwargs)