diff --git a/decode/qlook.py b/decode/qlook.py index 1da71fa..328549d 100644 --- a/decode/qlook.py +++ b/decode/qlook.py @@ -32,6 +32,7 @@ # constants +ABBA_PHASES = {0, 1, 2, 3} DATA_FORMATS = "csv", "nc", "zarr", "zarr.zip" DEFAULT_DATA_TYPE = "auto" DEFAULT_DEBUG = False @@ -374,10 +375,23 @@ def pswsc( da_despiked = despike(da) + # calculate ABBA cycles and phases + da_onoff = select.by(da_despiked, "state", ["ON", "OFF"]) + scan_onoff = utils.phaseof(da_onoff.state) + chop_per_scan = da_onoff.beam.groupby(scan_onoff).apply(utils.phaseof) + is_second_half = chop_per_scan.groupby(scan_onoff).apply( + lambda group: (group >= group.mean()) + ) + abba_cycle = (scan_onoff * 2 + is_second_half - 1) // 4 + abba_phase = (scan_onoff * 2 + is_second_half - 1) % 4 + # make spectrum - da_scan = select.by(da_despiked, "state", ["ON", "OFF"]) - da_sub = da_scan.groupby("scan").map(subtract_per_scan) - spec = da_sub.mean("scan") + spec = ( + da_onoff.assign_coords(abba_cycle=abba_cycle, abba_phase=abba_phase) + .groupby("abba_cycle") + .map(subtract_per_abba_cycle) + .mean("abba_cycle") + ) # save result suffixes = f".{suffix}.{format}" @@ -1133,31 +1147,52 @@ def mean_in_time(dems: xr.DataArray) -> xr.DataArray: return xr.zeros_like(middle) + dems.mean("time") -def subtract_per_scan(dems: xr.DataArray) -> xr.DataArray: - """Apply source-sky subtraction to a single-scan DEMS.""" +def subtract_per_abba_cycle(dems: xr.DataArray, /) -> xr.DataArray: + """Subtract sky from source with atmospheric correction for each ABBA cycle. + + Args: + dems: 2D DataArray (time x chan) of DEMS per ABBA cycle. + + Returns: + 1D DataArray (chan) of the mean spectrum after subtraction and correction. + If ABBA phases per cycle are incomplete, i.e., some phases are missing, + a spectrum filled with NaN will be returned instead. + + """ + if not set(np.unique(dems.abba_phase)) == ABBA_PHASES: + return dems.mean("time") * np.nan + + return dems.groupby("abba_phase").map(subtract_per_abba_phase).mean("abba_phase") + + +def subtract_per_abba_phase(dems: xr.DataArray, /) -> xr.DataArray: + """Subtract sky from source with atmospheric correction for each ABBA phase. + + Args: + dems: 2D DataArray (time x chan) of DEMS per ABBA phase. + + Returns: + 1D DataArray (chan) of the mean spectrum after subtraction and correction. + + Raises: + ValueError: Raised if ``dems.state`` is not ON-only nor OFF-only. + + """ t_amb = 273.15 + if len(states := np.unique(dems.state)) != 1: raise ValueError("State must be unique.") - if (state := states[0]) == "ON": - src = select.by(dems, "beam", include="A") - sky = select.by(dems, "beam", include="B") - return ( - t_amb - * (src.mean("time") - sky.mean("time").data) - / ((t_amb - sky.mean("time"))) - ) - - if state == "OFF": - src = select.by(dems, "beam", include="B") - sky = select.by(dems, "beam", include="A") - return ( - t_amb - * (src.mean("time") - sky.mean("time").data) - / ((t_amb - sky.mean("time"))) - ) + if states[0] == "ON": + src = select.by(dems, "beam", include="A").mean("time") + sky = select.by(dems, "beam", include="B").mean("time") + elif states[0] == "OFF": + src = select.by(dems, "beam", include="B").mean("time") + sky = select.by(dems, "beam", include="A").mean("time") + else: + raise ValueError("State must be either ON or OFF.") - raise ValueError("State must be either ON or OFF.") + return t_amb * (src - sky.data) / (t_amb - sky) def get_chan_weight(