From 0ad0a3c0b7ad8fb36aaa0ba2ee3c69ec27317189 Mon Sep 17 00:00:00 2001
From: Nina Burgdorfer <83002751+ninaburg@users.noreply.github.com>
Date: Tue, 16 Jul 2024 17:50:39 +0200
Subject: [PATCH] Update metadata BRN operator (#31)

---
 src/meteodatalab/operators/brn.py   | 43 +++++++++++++++++++++++++++--
 tests/test_meteodatalab/test_brn.py |  8 ++++++
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/src/meteodatalab/operators/brn.py b/src/meteodatalab/operators/brn.py
index 1dd7838..a0a8446 100644
--- a/src/meteodatalab/operators/brn.py
+++ b/src/meteodatalab/operators/brn.py
@@ -5,13 +5,47 @@
 import xarray as xr
 
 # Local
+from .. import metadata
 from .. import physical_constants as pc
 from .destagger import destagger
 from .thetav import fthetav
 
 
-def fbrn(p, t, qv, u, v, hhl, hsurf):
-    """Bulk Richardson Number (BRN)."""
+def fbrn(
+    p: xr.DataArray,
+    t: xr.DataArray,
+    qv: xr.DataArray,
+    u: xr.DataArray,
+    v: xr.DataArray,
+    hhl: xr.DataArray,
+    hsurf: xr.DataArray,
+) -> xr.DataArray:
+    """Bulk Richardson Number (BRN).
+
+    Parameters
+    ----------
+    p : xarray.DataArray
+        pressure in Pa
+    t : xarray.DataArray
+        air temperature in K
+    qv : xarray.DataArray
+        specific humidity (dimensionless)
+    u : xarray.DataArray
+        the x component of the wind velocity in m/s
+    v : xarray.DataArray
+        the y component of the wind velocity in m/s
+    hhl : xarray.DataArray
+        Heights of the interfaces between vertical layers in m amsl
+    hsurf : xarray.DataArray
+        earth surface height in m amsl
+
+
+    Returns
+    -------
+    xarray.DataArray
+        Bulk Richardson Number (dimensionless)
+
+    """
     nlevels = p.sizes["z"]
 
     thetav = fthetav(p, t, qv)
@@ -34,4 +68,7 @@ def fbrn(p, t, qv, u, v, hhl, hsurf):
         / (thetav_sum * (u_**2 + v_**2))
     )
 
-    return brn
+    return xr.DataArray(
+        data=brn,
+        attrs=metadata.override(p.message, shortName="BRN"),
+    )
diff --git a/tests/test_meteodatalab/test_brn.py b/tests/test_meteodatalab/test_brn.py
index 193b052..8e1f79d 100644
--- a/tests/test_meteodatalab/test_brn.py
+++ b/tests/test_meteodatalab/test_brn.py
@@ -20,6 +20,14 @@ def test_brn(data_dir, fieldextra):
         ds["P"], ds["T"], ds["QV"], ds["U"], ds["V"], ds["HHL"], ds["HSURF"]
     )
 
+    assert brn.parameter == {
+        "centre": "lssw",
+        "paramId": 503154,
+        "shortName": "BRN",
+        "units": "Numeric",
+        "name": "Bulk Richardson number",
+    }
+
     fs_ds = fieldextra("BRN")
 
     assert_allclose(fs_ds["BRN"], brn, rtol=5e-3, atol=5e-2)