refactor: Use metkit paramdb

docs/modules/settings.rst

@@ -106,8 +106,12 @@ List fields are set as JSON-encoded arrays:
 
 .. code-block:: bash
 
-   # Change the default GRIB parameter origin
-   export ANEMOI_SETTINGS_PARAMDB__DEFAULT_ORIGIN="ecmf"
+   # Change the default filters
+   export ANEMOI_SETTINGS_PARAMDB__DEFAULT_FILTERS='{"origin":98}'
+
+   # Add context
+   export ANEMOI_SETTINGS_PARAMDB__DEFAULT_FILTERS='{"context": {"class":"ai","levtype":"sfc","type":"fc","stream":"oper"}}'
+
 
    # Extend the parameter cache lifetime to 90 days
    export ANEMOI_SETTINGS_PARAMDB__CACHE_LENGTH="90"

pyproject.toml

@@ -43,6 +43,7 @@ dependencies = [
   "peekle",
   "pydantic>=2.9",
   "pydantic-settings>=2.14.1",
+  "pymetkit",
   "python-dateutil",
   "pyyaml",
   "rich",

src/anemoi/utils/grib.py

@@ -10,223 +10,59 @@
 
 """Utilities for working with GRIB parameters.
 
+Aliases Pymetkit's ParamDB to provide a consistent interface for GRIB parameter lookups across Anemoi, with configuration controlled by Anemoi settings.
+Use `PARAMDB` for direct access.
+
 See https://codes.ecmwf.int/grib/param-db/ for more information.
 """
 
-import json
 import logging
-import os
-import re
-import warnings
+from datetime import timedelta
 from functools import cache
+from typing import Any
 
-import requests
-
-from .caching import cached
-from .settings import AnemoiSettings
+from pymetkit import ParamDB
 
-LOG = logging.getLogger(__name__)
+from .settings import SETTINGS
 
-SETTINGS = AnemoiSettings()
 """Anemoi settings, loaded on module import."""
 
-
-@cached(collection="grib", expires=SETTINGS.paramdb.cache_length * 24 * 3600)
-def _units() -> dict[str, str]:
-    """Fetch and cache GRIB parameter units.
-
-    Returns
-    -------
-    dict
-        A dictionary mapping unit ids to their names.
-    """
-    r = requests.get("https://codes.ecmwf.int/parameter-database/api/v1/unit/")
-    r.raise_for_status()
-    units = r.json()
-    return {str(u["id"]): u["name"] for u in units}
-
-
-@cache
-def _get_local_db(local_db: str) -> list[dict[str, str | int | list[str]]]:
-    """Open the local GRIB parameter database.
-
-    Parameters
-    ----------
-    local_db : str
-        Path to the local cache file.
-
-    Returns
-    -------
-    list[dict[str, str | int | list[str]]]
-        A list of dictionaries containing parameter details.
-
-    Raises
-    ------
-    FileNotFoundError
-        If the local db file is not found.
-    """
-
-    if not os.path.exists(local_db):
-        raise FileNotFoundError(f"Local cache file {local_db} not found.")
-
-    return json.load(open(local_db, "r"))
+PARAMDB_SETTINGS = SETTINGS.paramdb
+LOG = logging.getLogger(__name__)
 
 
 @cache
-def _local_search_param(name: str) -> list[dict[str, str | int | list[str]]]:
-    """Search for a GRIB parameter by name in the local cache.
-
-    This is used to avoid making API calls when the local cache is available.
-
-    Parameters
-    ----------
-    name : str
-        Parameter name to search for.
+def get_paramdb() -> ParamDB:
+    """Return the global ParamDB instance.
 
     Returns
     -------
-    list
-        A list of dictionaries containing parameter details.
-
-    Raises
-    ------
-    KeyError
-        If no parameter is found.
+    ParamDB
+        The global ParamDB instance.
     """
-    local_cache = SETTINGS.paramdb.local_cache
-    assert local_cache is not None, "Local cache is not configured."
-
-    local_param_db = _get_local_db(local_cache)
 
-    matched_params = [param for param in local_param_db if param["shortname"] == name]
-    if matched_params:
-        return matched_params
-
-    raise KeyError(f"{name} not found in local cache.")
-
-
-@cached(collection="grib", expires=SETTINGS.paramdb.cache_length * 24 * 3600)
-def _online_search_param(name: str, **filters) -> list[dict[str, str | int | list[str]]]:
-    """Search for a GRIB parameter by name using the online API.
-
-    Parameters
-    ----------
-    name : str
-        Parameter name to search for.
-    filters : Any
-        Additional filters to disambiguate parameters with the same shortname (e.g. origin, encoding, table, discipline, category).
-
-    Returns
-    -------
-    list
-        A list of dictionaries containing parameter details.
-    """
-    r = requests.get(
-        f"https://codes.ecmwf.int/parameter-database/api/v1/param/?search=^{name}$&regex=true{''.join(f'&{k}={v}' for k, v in filters.items())}"
+    return ParamDB(
+        mode=PARAMDB_SETTINGS.mode,
+        cache_path=PARAMDB_SETTINGS.cache_path,
+        cache_ttl=timedelta(days=PARAMDB_SETTINGS.cache_length),
+        yaml_path=PARAMDB_SETTINGS.local_data,
     )
-    r.raise_for_status()
-    return r.json()
-
-
-def _search_param(name: str, **filters) -> dict[str, str | int | list[str]]:
-    """Search for a GRIB parameter by name.
-
-    Parameters
-    ----------
-    name : str
-        Parameter name to search for.
-    filters : Any
-        Additional filters to disambiguate parameters with the same shortname (e.g. origin, encoding, table, discipline, category).
-
-    Returns
-    -------
-    dict
-        A dictionary containing parameter details.
-
-    Raises
-    ------
-    KeyError
-        If no parameter is found.
-    """
-    if "origin" in filters and isinstance(filters["origin"], str):
-        filters["origin"] = origin(filters["origin"])["id"]
-
-    name = re.escape(name)
-
-    if SETTINGS.paramdb.local_cache is not None:
-        if filters:
-            warnings.warn("Filters are ignored when using local cache.")
-        results = _local_search_param(name)
-    else:
-        results = _online_search_param(name, **filters)
-
-    if len(results) == 0:
-        raise KeyError(f"{name} not found in parameter database.")
-
-    if len(results) > 1:
-        names = [f"{r.get('id')} ({r.get('name')})" for r in results]
-        dissemination = [r for r in results if "dissemination" in r.get("access_ids", [])]  # type: ignore[reportOperatorIssue]
-        if len(dissemination) == 1:
-            return dissemination[0]
-
-        warnings.warn(f"{name} is ambiguous: {', '.join(names)}.")
-        if "origin" not in filters and SETTINGS.paramdb.local_cache is None:
-            warnings.warn(f"Applying origin='{SETTINGS.paramdb.default_origin}' in an attempt to disambiguate {name}.")
-            try:
-                filtered_param = _search_param(name, **{**filters, "origin": SETTINGS.paramdb.default_origin})
-                warnings.warn(
-                    f"Disambiguated {name} to id: {filtered_param['id']} ({filtered_param.get('name', 'unknown')})."
-                )
-                return filtered_param
-            except KeyError:
-                pass
 
-        warnings.warn(f"Failed to disambiguate {name}'. Returning the first match: {names[0]}.")
-        results = sorted(results, key=lambda x: x["id"])
 
-    return results[0]
-
-
-@cached(collection="grib", expires=SETTINGS.paramdb.cache_length * 24 * 3600)
-def origin(name: str) -> dict[str, str | int]:
-    """Search for an id of an origin by name.
-
-    Parameters
-    ----------
-    name : str
-        Origin name to search for.
-
-    Returns
-    -------
-    dict
-        A dictionary containing origin details.
-
-    Raises
-    ------
-    KeyError
-        If no origin is found.
-    """
-    name = re.escape(name)
-    r = requests.get("https://codes.ecmwf.int/parameter-database/api/v1/origin/")
-    r.raise_for_status()
-    results = r.json()
-
-    for result in results:
-        if result["abbreviation"] == name:
-            return result
-
-    raise KeyError(f"{name} not found in origin database.")
+def shortname_to_paramid(shortname: str, **filters: Any) -> int:
+    """Return the GRIB parameter id given its shortname.
 
+    SETTINGS.paramdb.default_filters can be used to provide default filters for disambiguation.
 
-def shortname_to_paramid(shortname: str, **filters) -> int:
-    """Return the GRIB parameter id given its shortname.
+    If a collision is detected (i.e. multiple parameters with the same shortname), a warning will be logged and the first parameter id will be returned.
+    Additional filters can be provided to disambiguate.
 
     Parameters
     ----------
     shortname : str
         Parameter shortname.
     filters : Any
-        Additional filters to disambiguate parameters with the same shortname (e.g. origin, encoding, table, discipline, category).
+        Additional filters to disambiguate parameters with the same shortname (e.g. origin, access, table).
 
     Returns
     -------
@@ -236,18 +72,31 @@ def shortname_to_paramid(shortname: str, **filters) -> int:
     >>> shortname_to_paramid("2t")
     167
     """
-    return _search_param(shortname, **filters)["id"]  # type: ignore[reportReturnType]
-
-
-def paramid_to_shortname(paramid: int, **filters) -> str:
+    filters = filters or PARAMDB_SETTINGS.default_filters or {}
+
+    if get_paramdb().shortname_has_collisions(shortname) and not filters:
+        candidates = get_paramdb().shortname_to_param_id_candidates(shortname)
+        log_message = (
+            f"Shortname '{shortname}' has collisions. Candidates: {[c.param_id for c in candidates]}. "
+            "Consider providing additional filters to disambiguate, will use empty context by default."
+            f"\nTo select one of the ids use: \n"
+            + "\n".join(
+                f"  - {c.param_id}: {f'context = {c.mars_request_context}' if c.mars_request_context is not None else c.hard_filter_selector}"
+                for c in candidates
+            )
+        )
+        LOG.warning(log_message)
+        filters = {"context": {}}
+    return get_paramdb().shortname_to_param_id(shortname, **filters)
+
+
+def paramid_to_shortname(paramid: int) -> str:
     """Return the shortname of a GRIB parameter given its id.
 
     Parameters
     ----------
     paramid : int
         Parameter id.
-    filters : Any
-        Additional filters to disambiguate parameters with the same shortname (e.g. origin, encoding, table, discipline, category).
 
     Returns
     -------
@@ -257,7 +106,7 @@ def paramid_to_shortname(paramid: int, **filters) -> str:
     >>> paramid_to_shortname(167)
     '2t'
     """
-    return _search_param(str(paramid), **filters)["shortname"]  # type: ignore[reportReturnType]
+    return get_paramdb().param_id_to_shortname(paramid)
 
 
 def units(param: int | str) -> str:
@@ -273,12 +122,10 @@ def units(param: int | str) -> str:
     str
         Parameter unit.
 
-    >>> unit(167)
+    >>> units(167)
     'K'
     """
-
-    unit_id = str(_search_param(str(param))["unit_id"])
-    return _units()[unit_id]
+    return get_paramdb().get_units(param)
 
 
 def must_be_positive(param: int | str) -> bool:

src/anemoi/utils/settings_schema/paramdb.py

@@ -7,8 +7,11 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+from typing import Any
+from typing import Literal
 from typing import Optional
 
+from pydantic import Field
 from pydantic import FilePath
 
 from .base import AnemoiBaseSettingsSchema
@@ -22,11 +25,17 @@ class ParamDBConfig(AnemoiBaseSettingsSchema):
     and cache lifetime).
     """
 
-    default_origin: str = "ecmf"
-    """Default origin to use when disambiguating parameters with the same shortname."""
+    mode: Literal["online", "offline"] = "offline"
+    """Mode for GRIB parameter lookups. 'online' uses the ECMWF API, 'offline' uses a local cache file."""
+
+    cache_path: Optional[FilePath] = None
+    """Directory in which to store the cache file. Defaults to the OS-appropriate user cache directory (requires ``platformdirs``). Only relevant when ``mode="online"``.."""
 
     cache_length: int = 30
     """Cache length in days for GRIB parameter lookups."""
 
-    local_cache: Optional[FilePath] = None
-    """Path to a local JSON cache file for GRIB parameters. If set, used instead of the online API."""
+    local_data: Optional[FilePath] = None
+    """Path to a local YAML cache file for GRIB parameters. If set, used instead of the online API."""
+
+    default_filters: dict[str, Any] | None = Field(default={"access": "dissemination"})
+    """Default filters to disambiguate GRIB parameters with the same shortname (e.g. origin, access, table, or context with class, type, stream, etc.)."""