#93 Add function for correcting for atmospheric transmission

decode/correct.py

@@ -1 +1,71 @@
-__all__ = []
+__all__ = ["for_atmosphere"]
+
+
+# standard library
+from collections.abc import Sequence
+from datetime import datetime, timedelta
+from typing import TypeVar, Union
+
+
+# dependencies
+import xarray as xr
+from . import select
+
+
+# type hints
+T = TypeVar("T", bool, int, float, str, datetime, timedelta)
+Multiple = Union[Sequence[T], T]
+
+
+def for_atmosphere(
+    dems: xr.DataArray,
+    include_on: Multiple[T],
+    include_off: Multiple[T],
+    include_r: Multiple[T],
+    /,
+    method_off: str = "linear",
+    method_r: str = "nearest",
+    T_amb: float = 273.0,
+) -> xr.DataArray:
+    """Correct for the atmospheric transmission of DEMS.
+
+    Args:
+        dems: Input DEMS DataArray with a correct state coordinate.
+        include_on: State value(s) to be assigned to on-source samples.
+        include_off: State value(s) to be assigned to off-source samples.
+        include_r: State value(s) to be assigned to hot-load samples.
+        method_off: Interpolation method of the off-source samples
+            to the measured time of the on-source samples.
+        method_r: Interpolation method of the hot-load samples
+            to the measured time of the on-source samples.
+        T_amb: Ambient temperature assumed for correction.
+
+    Returns:
+        DEMS DataArray of the on-source samples in the Ta* scale.
+
+    """
+    Tb_on = select.by(dems, "state", include=include_on)
+    Tb_off = select.by(dems, "state", include=include_off)
+    Tb_r = select.by(dems, "state", include=include_r)
+
+    Tb_off_mean = Tb_off.groupby("scan").map(mean_in_time)
+    Tb_r_mean = Tb_r.groupby("scan").map(mean_in_time)
+
+    Tb_off_ip = Tb_off_mean.interp_like(
+        Tb_on,
+        method=method_off,  # type: ignore
+        kwargs={"fill_value": "extrapolate"},
+    ).values
+    Tb_r_ip = Tb_r_mean.interp_like(
+        Tb_on,
+        method=method_r,  # type: ignore
+        kwargs={"fill_value": "extrapolate"},
+    ).values
+
+    return T_amb * (Tb_on - Tb_off_ip) / (Tb_r_ip - Tb_off_ip)
+
+
+def mean_in_time(dems: xr.DataArray) -> xr.DataArray:
+    """Similar to DataArray.mean but keeps middle time."""
+    middle = dems[len(dems) // 2 : len(dems) // 2 + 1]
+    return xr.zeros_like(middle) + dems.mean("time")