feat: add CosmoDownscaling diagnostic model (ERA5 -> COSMO-REA6/REA2)

CHANGELOG.md

@@ -30,6 +30,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   reflectivity, rain rate, and 1-hour accumulation (`OPERA`)
 - Added support for cumulative variables in ARCO data source
 - Added DLESyM-v0-ISCCP-ERA5 climate model
+- Added COSMO-REA downscaling diagnostic model (`CorrDiffCosmoEra5`) and
+  `CosmoLexicon` for regression (mean) and diffusion downscaling of ERA5 to
+  COSMO-REA6 (6 km) and COSMO-REA2 (2.2 km), with sub-domain support via
+  `set_domain`
 
 ### Changed
 

docs/modules/models_dx.rst

@@ -40,6 +40,7 @@ etc.
       CBottleSR
       CBottleTCGuidance
       CorrDiffCMIP6
+      CorrDiffCosmoEra5
       CorrDiffTaiwan
       ClimateNet
       DerivedRH

docs/userguide/about/install.md

@@ -586,6 +586,30 @@ pip install earth2studio[corrdiff]
 uv add earth2studio --extra corrdiff
 ```
 
+:::
+::::
+:::::
+:::::{tab-item} CorrDiff COSMO-ERA5
+Notes: Additional dependencies for the `CorrDiffCosmoEra5` model. This model needs
+the RoPE / NATTEN attention backend from `nvidia-physicsnemo`, which is not on PyPI
+yet, so physicsnemo must be installed from a pinned git commit. The `uv` path picks
+this up automatically from `[tool.uv.sources]`; the `pip` path installs it explicitly.
+
+::::{tab-set}
+:::{tab-item} pip
+
+```bash
+pip install "nvidia-physicsnemo @ git+https://github.com/NVIDIA/physicsnemo.git@ced75d93d014f70bb691372788eee2d201171c12"
+pip install earth2studio[cosmo]
+```
+
+:::
+:::{tab-item} uv
+
+```bash
+uv add earth2studio --extra cosmo
+```
+
 :::
 ::::
 :::::

earth2studio/lexicon/__init__.py

@@ -21,6 +21,7 @@
 from .cds import CDSLexicon
 from .cfs import CFSFluxLexicon, CFSLexicon
 from .cmip6 import CMIP6Lexicon
+from .cosmo import CosmoLexicon
 from .ecmwf import AIFSLexicon, IFSLexicon
 from .gdas import GDASObsConvLexicon
 from .gefs import GEFSLexicon, GEFSLexiconSel

earth2studio/lexicon/cosmo.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+from collections.abc import Callable
+
+import numpy as np
+
+from .base import LexiconType
+
+
+class CosmoLexicon(metaclass=LexiconType):
+    """COSMO-REA Lexicon
+
+    Maps between Earth2Studio vocabulary and the COSMO-REA output variable names
+    used by :class:`~earth2studio.models.dx.CorrDiffCosmoEra5` (COSMO-REA6 and
+    COSMO-REA2). The vocabulary is the union over both resolutions: where a field
+    has a canonical Earth2Studio name and matching units it maps to that name
+    (e.g. ``t2m``, ``sp``); COSMO-specific fields with no canonical equivalent
+    (model-level winds/temperatures/humidity, turbulent fluxes, TKE, PBL height,
+    ...) keep a descriptive lowercase name.
+
+    Each entry's value is ``(cosmo_name(s), scale)``: ``cosmo_name`` is the
+    interior COSMO output name (a tuple when the two resolutions spell the same
+    field differently, e.g. ``U_10M`` in REA6 vs ``10U`` in REA2), and ``scale``
+    converts the COSMO value to the Earth2Studio name's units (e.g. ``CLCT`` is a
+    percentage and ``tcc`` is a 0-1 fraction, so ``scale = 0.01``).
+
+    The COSMO-REA regional reanalyses are DWD products:
+
+    * COSMO-REA6 (~6 km): https://reanalysis.meteo.uni-bonn.de/?COSMO-REA6
+    * COSMO-REA2 (~2.2 km): https://reanalysis.meteo.uni-bonn.de/?COSMO-REA2
+    """
+
+    @staticmethod
+    def build_vocab() -> dict[str, tuple[tuple[str, ...] | str, float]]:
+        """Create the COSMO-REA vocab dictionary."""
+        vocab: dict[str, tuple[tuple[str, ...] | str, float]] = {
+            # Earth2Studio canonical name : (COSMO output name(s), unit scale)
+            "u10m": (("U_10M", "10U"), 1.0),  # REA6 / REA2 spellings
+            "v10m": (("V_10M", "10V"), 1.0),
+            "t2m": (("T_2M", "2MT"), 1.0),
+            "d2m": ("TD_2M", 1.0),  # 2 m dewpoint (K) -- same units
+            "sp": ("PS", 1.0),  # surface pressure (Pa)
+            "tcc": ("CLCT", 0.01),  # cloud cover: COSMO % -> E2S 0-1 fraction
+            # COSMO-specific fields (no canonical Earth2Studio name): the name is
+            # the lowercased COSMO output name, units unchanged.
+            "tot_precip": ("TOT_PRECIP", 1.0),
+            "aswdifd_s": ("ASWDIFD_S", 1.0),
+            "aswdir_s": ("ASWDIR_S", 1.0),
+            "alwu_s": ("ALWU_S", 1.0),
+            "athd_s": ("ATHD_S", 1.0),
+            "qv_2m": ("QV_2M", 1.0),
+            "vmax_10m": ("VMAX_10M", 1.0),
+            "lhfl_s": ("LHFL_S", 1.0),
+            "shfl_s": ("SHFL_S", 1.0),
+            "h_pbl": ("H_PBL", 1.0),
+        }
+        # Model-level fields. REA6 ships levels 35-40 as ``<P>_L<lvl>``; REA2 ships
+        # levels 45-50 as ``<P>3D_L<lvl>``. Distinct names, no overlap.
+        for lvl in range(35, 41):
+            for p in ("U", "V", "T", "TKE", "Q"):
+                vocab[f"{p.lower()}_l{lvl}"] = (f"{p}_L{lvl}", 1.0)
+        for lvl in range(45, 51):
+            for p in ("U3D", "V3D"):
+                vocab[f"{p.lower()}_l{lvl}"] = (f"{p}_L{lvl}", 1.0)
+        return vocab
+
+    VOCAB = build_vocab()
+
+    @classmethod
+    def get_item(cls, val: str) -> tuple[tuple[str, ...] | str, Callable]:
+        """Return the COSMO name(s) and unit modifier for an Earth2Studio name.
+
+        Parameters
+        ----------
+        val : str
+            Name in Earth2Studio terminology.
+
+        Returns
+        -------
+        tuple[str | tuple[str, ...], Callable]
+            COSMO output name(s) and a modifier converting the COSMO value to the
+            Earth2Studio name's units.
+        """
+        cosmo_name, scale = cls.VOCAB[val]
+
+        def mod(x: np.ndarray) -> np.ndarray:
+            """Convert COSMO units to the Earth2Studio name's units."""
+            return x * scale
+
+        return cosmo_name, mod
+
+    @classmethod
+    def to_e2studio(cls, cosmo_name: str) -> tuple[str, float]:
+        """Map a COSMO output variable name to ``(earth2studio_name, scale)``.
+
+        The reverse of :attr:`VOCAB`, accepting either resolution's spelling.
+        Unknown names fall back to their lowercased form with unit scale 1.0, so a
+        new COSMO field still gets a sensible name rather than raising.
+        """
+        reverse = getattr(cls, "_REVERSE", None)
+        if reverse is None:
+            reverse = {}
+            for e2s_name, (cosmo, scale) in cls.VOCAB.items():
+                names = cosmo if isinstance(cosmo, tuple) else (cosmo,)
+                for name in names:
+                    reverse[name] = (e2s_name, scale)
+            cls._REVERSE = reverse
+        return reverse.get(cosmo_name, (cosmo_name.lower(), 1.0))

earth2studio/models/dx/__init__.py

@@ -24,6 +24,7 @@
     CorrDiffTaiwan,
 )
 from earth2studio.models.dx.corrdiff_cmip6 import CorrDiffCMIP6
+from earth2studio.models.dx.corrdiff_cosmo_era5 import CorrDiffCosmoEra5
 from earth2studio.models.dx.derived import (
     DerivedRH,
     DerivedRHDewpoint,
@@ -53,6 +54,7 @@
     "ClimateNet",
     "CorrDiff",
     "CorrDiffCMIP6",
+    "CorrDiffCosmoEra5",
     "CorrDiffTaiwan",
     "DLESyMv0_ISCCP_ERA5Precip",
     "PrecipitationAFNO",