From 3da7c1a8dadcbae43f84ffd802c548b30a6ad467 Mon Sep 17 00:00:00 2001
From: Akio Taniguchi <taniguchi@a.phys.nagoya-u.ac.jp>
Date: Fri, 9 Aug 2024 23:04:59 +0000
Subject: [PATCH] #189 Add df/f-to-brightness conversion to load.dems

---
 decode/convert.py | 31 ++++++++++++++++++++++++++++++-
 decode/load.py    | 33 ++++++++++++++++++++++++++-------
 2 files changed, 56 insertions(+), 8 deletions(-)

diff --git a/decode/convert.py b/decode/convert.py
index 37f399e..f4c9b46 100644
--- a/decode/convert.py
+++ b/decode/convert.py
@@ -1,4 +1,4 @@
-__all__ = ["coord_units", "frame", "units"]
+__all__ = ["coord_units", "dfof_to_brightness", "frame", "units"]
 
 
 # standard library
@@ -17,6 +17,11 @@
 UnitLike = Union[xr.DataArray, Unit, str]
 
 
+# constants
+DEFAULT_T_ATM = 273.0  # K
+DEFAULT_T_ROOM = 293.0  # K
+
+
 def coord_units(
     da: xr.DataArray,
     coord_names: Multiple[str],
@@ -51,6 +56,30 @@ def coord_units(
     return da
 
 
+def dfof_to_brightness(da: xr.DataArray, /) -> xr.DataArray:
+    """Convert a DataArray of df/f to that of brightness."""
+    if np.isnan(T_atm := da.temperature.mean().data):
+        T_atm = DEFAULT_T_ATM
+
+    if np.isnan(T_room := da.aste_cabin_temperature.mean().data):
+        T_room = DEFAULT_T_ROOM
+
+    fwd = da.d2_resp_fwd.data
+    p0 = da.d2_resp_p0.data
+    T0 = da.d2_resp_t0.data
+
+    return (
+        da.copy(
+            deep=True,
+            data=(da.data + p0 * np.sqrt(T_room + T0)) ** 2 / (p0**2 * fwd)
+            - T0 / fwd
+            - (1 - fwd) / fwd * T_atm,
+        )
+        .astype(da.dtype)
+        .assign_attrs(long_name="Brightness", units="K")
+    )
+
+
 def frame(da: xr.DataArray, new_frame: str, /) -> xr.DataArray:
     """Convert the skycoord frame of a DataArray.
 
diff --git a/decode/load.py b/decode/load.py
index 6e1102a..70b10fa 100644
--- a/decode/load.py
+++ b/decode/load.py
@@ -11,6 +11,7 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
+from . import convert
 
 
 # constants
@@ -66,20 +67,24 @@ def atm(*, type: Literal["eta", "tau"] = "tau") -> xr.DataArray:
         raise ValueError("Type must be either eta or tau.")
 
 
-def dems(dems: Union[Path, str], /, **options: Any) -> xr.DataArray:
+def dems(
+    dems: Union[Path, str],
+    /,
+    measure: Literal["brightness", "df/f"] = "brightness",
+    **options: Any,
+) -> xr.DataArray:
     """Load a DEMS file as a DataArray.
 
     Args:
         dems: Path of the DEMS file.
+        measure: Measure of the DataArray (either brightness or df/f).
+        **options: Options to be passed to ``xarray.open_dataarray``.
 
-    Keyword Args:
-        options: Arguments to be passed to ``xarray.open_dataarray``.
-
-    Return:
+    Returns:
         Loaded DEMS DataArray.
 
     Raises:
-        ValueError: Raised if the file type is not supported.
+        ValueError: Raised if file type or measure is not supported.
 
     """
     suffixes = Path(dems).suffixes
@@ -101,4 +106,18 @@ def dems(dems: Union[Path, str], /, **options: Any) -> xr.DataArray:
             "Use netCDF (.nc) or Zarr (.zarr, .zarr.zip)."
         )
 
-    return xr.open_dataarray(dems, **options)
+    da = xr.open_dataarray(dems, **options)
+
+    if da.long_name == "Brightness" and measure == "brightness":
+        return da
+
+    if da.long_name == "df/f" and measure == "df/f":
+        return da
+
+    if da.long_name == "df/f" and measure == "brightness":
+        return convert.dfof_to_brightness(da)
+
+    if da.long_name == "Brightness" and measure == "df/f":
+        raise ValueError("Brightness-to-df/f conversion is not supported.")
+
+    raise ValueError("Measure must be either brightness or df/f.")