diff --git a/docs/preprocess.rst b/docs/preprocess.rst index 6ce35a95a..1845208f5 100644 --- a/docs/preprocess.rst +++ b/docs/preprocess.rst @@ -252,6 +252,7 @@ Flagging and Products .. autoclass:: sotodlib.preprocess.processes.FlagTurnarounds .. autoclass:: sotodlib.preprocess.processes.DarkDets .. autoclass:: sotodlib.preprocess.processes.SourceFlags +.. autoclass:: sotodlib.preprocess.processes.GetStats HWP Related ::::::::::: diff --git a/sotodlib/preprocess/pcore.py b/sotodlib/preprocess/pcore.py index e6bf5a42e..0f99f9037 100644 --- a/sotodlib/preprocess/pcore.py +++ b/sotodlib/preprocess/pcore.py @@ -5,6 +5,7 @@ from .. import core from so3g.proj import Ranges, RangesMatrix from scipy.sparse import csr_array +from matplotlib import pyplot as plt class _Preprocess(object): """The base class for Preprocessing modules which defines the required @@ -270,7 +271,7 @@ def _expand(new, full, wrap_valid=True): continue out.wrap_new( k, new._assignments[k], cls=_zeros_cls(v)) oidx=[]; nidx=[] - for a in new._assignments[k]: + for ii, a in enumerate(new._assignments[k]): if a == 'dets': oidx.append(fs_dets) nidx.append(ns_dets) @@ -278,8 +279,19 @@ def _expand(new, full, wrap_valid=True): oidx.append(fs_samps) nidx.append(ns_samps) else: - oidx.append(slice(None)) - nidx.append(slice(None)) + if (ii == 0) and isinstance(out[k], RangesMatrix): # Treat like dets + # _ranges_matrix_match expects oidx[0] and nidx[0] to be list(inds), not slice. + # Unknown axes treated as dets if first entry, else like samps. Added to support (subscans, samps) RangesMatrix. + if a in full._axes: + _, fs, ns = full[a].intersection(new[a], return_slices=True) + else: + fs = range(new[a].count) + ns = range(new[a].count) + oidx.append(fs) + nidx.append(ns) + else: # Treat like samps + oidx.append(slice(None)) + nidx.append(slice(None)) oidx = tuple(oidx) nidx = tuple(nidx) if isinstance(out[k], RangesMatrix): @@ -456,6 +468,7 @@ def run(self, aman, proc_aman=None, select=True, sim=False, update_plot=False): update_full_aman( proc_aman, full, self.wrap_valid) if update_plot: process.plot(aman, proc_aman, filename=os.path.join(self.plot_dir, '{ctime}/{obsid}', f'{step+1}_{{name}}.png')) + plt.close() if select: process.select(aman, proc_aman) proc_aman.restrict('dets', aman.dets.vals) diff --git a/sotodlib/preprocess/preprocess_plot.py b/sotodlib/preprocess/preprocess_plot.py index 83f3325d0..3985cfc81 100644 --- a/sotodlib/preprocess/preprocess_plot.py +++ b/sotodlib/preprocess/preprocess_plot.py @@ -390,6 +390,44 @@ def plot_trending_flags(aman, trend_aman, filename='./trending_flags.png'): os.makedirs(head_tail[0], exist_ok=True) plt.savefig(filename) +def plot_signal(aman, signal=None, xx=None, signal_name="signal", x_name="timestamps", plot_ds_factor=50, plot_ds_factor_dets=None, xlim=None, alpha=0.2, yscale='linear', y_unit=None, filename="./signal.png"): + from operator import attrgetter + if plot_ds_factor_dets is None: + plot_ds_factor_dets = plot_ds_factor + if signal is None: + signal = attrgetter(signal_name)(aman) + if xx is None: + xx = attrgetter(x_name)(aman) + yy = signal[::plot_ds_factor_dets, 1::plot_ds_factor].copy() # (dets, samps); (dets, nusamps); (dets, nusamps, subscans) + xx = xx[1::plot_ds_factor].copy() # (samps); (nusamps) + if x_name == "timestamps": + xx -= xx[0] + if yy.ndim > 2: # Flatten subscan axis into dets + yy = yy.swapaxes(1,2).reshape(-1, yy.shape[1]) + + if xlim is not None: + xinds = np.logical_and(xx >= xlim[0], xx <= xlim[1]) + xx = xx[xinds] + yy = yy[:,xinds] + + fig, ax = plt.subplots(1, 1, figsize=(6.4, 4.8)) + ax.plot(xx, yy.T, color='k', alpha=0.2) + ax.set_yscale(yscale) + if "freqs" in x_name: + ax.set_xlabel("freq [Hz]") + else: + ax.set_xlabel(f"{x_name} [s]") + y_unit = "" if y_unit is None else f" [{y_unit}]" + ax.set_ylabel(f"{signal_name.replace('.Pxx', '')}{y_unit}") + plt.suptitle(f"{aman.obs_info.obs_id}, dT = {np.ptp(aman.timestamps)/60:.1f} min") + plt.tight_layout() + head_tail = os.path.split(filename) + os.makedirs(head_tail[0], exist_ok=True) + plt.savefig(filename) + +def plot_psd(aman, signal=None, xx=None, signal_name="psd.Pxx", x_name="psd.freqs", plot_ds_factor=4, plot_ds_factor_dets=20, xlim=None, alpha=0.2, yscale='log', y_unit=None, filename="./psd.png"): + return plot_signal(aman, signal, xx, signal_name, x_name, plot_ds_factor, plot_ds_factor_dets, xlim, alpha, yscale, y_unit, filename) + def plot_signal_diff(aman, flag_aman, flag_type="glitches", flag_threshold=10, plot_ds_factor=50, filename="./glitch_signal_diff.png"): """ Function for plotting the difference in signal before and after cuts from either glitches or jumps. diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py index 909fe3ecf..a142bf6a2 100644 --- a/sotodlib/preprocess/processes.py +++ b/sotodlib/preprocess/processes.py @@ -172,6 +172,7 @@ class GlitchDetection(_FracFlaggedMixIn, _Preprocess): buffer: 10 hp_fc: 1 n_sig: 10 + subscan: False save: True plot: plot_ds_factor: 50 @@ -340,6 +341,7 @@ def plot(self, aman, proc_aman, filename): plot_ds_factor=self.plot_cfgs.get("plot_ds_factor", 50), filename=filename.replace('{name}', f'{ufm}_jump_signal_diff')) plot_flag_stats(aman, proc_aman[name], flag_type='jumps', filename=filename.replace('{name}', f'{ufm}_jumps_stats')) + class PSDCalc(_Preprocess): """ Calculate the PSD of the data and add it to the Preprocessing AxisManager under the "psd" field. @@ -353,6 +355,7 @@ class PSDCalc(_Preprocess): "psd_cfgs": # optional, kwargs to scipy.welch "nperseg": 1024 "wrap_name": "psd" # optional + "subscan": False "save": True .. autofunction:: sotodlib.tod_ops.fft_ops.calc_psd @@ -368,21 +371,105 @@ def __init__(self, step_cfgs): def calc_and_save(self, aman, proc_aman): freqs, Pxx = tod_ops.fft_ops.calc_psd(aman, signal=aman[self.signal], **self.calc_cfgs) - fft_aman = core.AxisManager( - aman.dets, - core.OffsetAxis("nusamps",len(freqs)) - ) + + fft_aman = core.AxisManager(aman.dets, + core.OffsetAxis("nusamps", len(freqs))) + pxx_axis_map = [(0, "dets"), (1, "nusamps")] + if self.calc_cfgs.get('subscan', False): + fft_aman.wrap("Pxx_ss", Pxx, pxx_axis_map+[(2, aman.subscans)]) + Pxx = np.nanmean(Pxx, axis=-1) # Mean of subscans + fft_aman.wrap("freqs", freqs, [(0,"nusamps")]) - fft_aman.wrap("Pxx", Pxx, [(0,"dets"), (1,"nusamps")]) + fft_aman.wrap("Pxx", Pxx, pxx_axis_map) + self.save(proc_aman, fft_aman) def save(self, proc_aman, fft_aman): if not(self.save_cfgs is None): proc_aman.wrap(self.wrap, fft_aman) + def plot(self, aman, proc_aman, filename): + if self.plot_cfgs is None: + return + if self.plot_cfgs: + from .preprocess_plot import plot_psd + + filename = filename.replace('{ctime}', f'{str(aman.timestamps[0])[:5]}') + filename = filename.replace('{obsid}', aman.obs_info.obs_id) + det = aman.dets.vals[0] + ufm = det.split('_')[2] + filename = filename.replace('{name}', f'{ufm}_{self.wrap}') + + plot_psd(aman, signal=attrgetter(f"{self.wrap}.Pxx")(proc_aman), + xx=attrgetter(f"{self.wrap}.freqs")(proc_aman), filename=filename, **self.plot_cfgs) + + +class GetStats(_Preprocess): + """ Get basic statistics from a TOD or its power spectrum. + + Example config block: + + - name : "tod_stats" + signal: "signal" # optional + wrap: "tod_stats" # optional + calc: + stat_names: ["median", "std"] + split_subscans: False # optional + psd_mask: # optional, for cutting a power spectrum in frequency + freqs: "psd.freqs" + low_f: 1 + high_f: 10 + save: True + + """ + name = "tod_stats" + def __init__(self, step_cfgs): + self.signal = step_cfgs.get('signal', 'signal') + self.wrap = step_cfgs.get('wrap', 'tod_stats') + + super().__init__(step_cfgs) + + def calc_and_save(self, aman, proc_aman): + if self.calc_cfgs.get('psd_mask') is not None: + mask_dict = self.calc_cfgs.get('psd_mask') + _f = attrgetter(mask_dict['freqs']) + try: + freqs = _f(aman) + except KeyError: + freqs = _f(proc_aman) + low_f, high_f = mask_dict['low_f'], mask_dict['high_f'] + fmask = np.all([freqs >= low_f, freqs <= high_f], axis=0) + self.calc_cfgs['mask'] = fmask + del self.calc_cfgs['psd_mask'] + + _f = attrgetter(self.signal) + try: + signal = _f(aman) + except KeyError: + signal = _f(proc_aman) + stats_aman = tod_ops.flags.get_stats(aman, signal, **self.calc_cfgs) + self.save(proc_aman, stats_aman) + + def save(self, proc_aman, stats_aman): + if not(self.save_cfgs is None): + proc_aman.wrap(self.wrap, stats_aman) + + def plot(self, aman, proc_aman, filename): + if self.plot_cfgs is None: + return + if self.plot_cfgs: + from .preprocess_plot import plot_signal + + filename = filename.replace('{ctime}', f'{str(aman.timestamps[0])[:5]}') + filename = filename.replace('{obsid}', aman.obs_info.obs_id) + det = aman.dets.vals[0] + ufm = det.split('_')[2] + filename = filename.replace('{name}', f'{ufm}_{self.signal}') + + plot_signal(aman, signal_name=self.signal, x_name="timestamps", filename=filename, **self.plot_cfgs) class Noise(_Preprocess): """Estimate the white noise levels in the data. Assumes the PSD has been - wrapped into the preprocessing AxisManager. All calculation configs goes to `calc_wn`. + wrapped into the preprocessing AxisManager. All calculation configs goes to `calc_wn`. Saves the results into the "noise" field of proc_aman. @@ -391,6 +478,8 @@ class Noise(_Preprocess): Example config block:: - name: "noise" + fit: False + subscan: False calc: low_f: 5 high_f: 10 @@ -408,6 +497,7 @@ class Noise(_Preprocess): def __init__(self, step_cfgs): self.psd = step_cfgs.get('psd', 'psd') self.fit = step_cfgs.get('fit', False) + self.subscan = step_cfgs.get('subscan', False) super().__init__(step_cfgs) @@ -415,21 +505,28 @@ def calc_and_save(self, aman, proc_aman): if self.psd not in proc_aman: raise ValueError("PSD is not saved in Preprocessing AxisManager") psd = proc_aman[self.psd] - + pxx = psd.Pxx_ss if self.subscan else psd.Pxx + if self.calc_cfgs is None: self.calc_cfgs = {} - + if self.fit: - calc_aman = tod_ops.fft_ops.fit_noise_model(aman, pxx=psd.Pxx, + if self.calc_cfgs.get('subscan') is None: + self.calc_cfgs['subscan'] = self.subscan + calc_aman = tod_ops.fft_ops.fit_noise_model(aman, pxx=pxx, f=psd.freqs, merge_fit=True, **self.calc_cfgs) else: - wn = tod_ops.fft_ops.calc_wn(aman, pxx=psd.Pxx, + wn = tod_ops.fft_ops.calc_wn(aman, pxx=pxx, freqs=psd.freqs, **self.calc_cfgs) - calc_aman = core.AxisManager(aman.dets) - calc_aman.wrap("white_noise", wn, [(0,"dets")]) + if not self.subscan: + calc_aman = core.AxisManager(aman.dets) + calc_aman.wrap("white_noise", wn, [(0,"dets")]) + else: + calc_aman = core.AxisManager(aman.dets, aman.subscan_info.subscans) + calc_aman.wrap("white_noise", wn, [(0,"dets"), (1,"subscans")]) self.save(proc_aman, calc_aman) @@ -457,10 +554,12 @@ def select(self, meta, proc_aman=None): self.select_cfgs['name'] = self.select_cfgs.get('name','noise') if self.fit: - keep = proc_aman[self.select_cfgs['name']].fit[:,1] <= self.select_cfgs["max_noise"] + wn = proc_aman[self.select_cfgs['name']].fit[:,1] else: - keep = proc_aman[self.select_cfgs['name']].white_noise <= self.select_cfgs["max_noise"] - + wn = proc_aman[self.select_cfgs['name']].white_noise + if self.subscan: + wn = np.nanmean(wn, axis=-1) # Mean over subscans + keep = wn <= np.float64(self.select_cfgs["max_noise"]) meta.restrict("dets", meta.dets.vals[keep]) return meta @@ -786,6 +885,9 @@ def calc_and_save(self, aman, proc_aman): calc_aman = core.AxisManager(aman.dets, aman.samps) calc_aman.wrap('turnarounds', ta, [(0, 'dets'), (1, 'samps')]) + if ('merge_subscans' not in self.calc_cfgs) or (self.calc_cfgs['merge_subscans']): + calc_aman.wrap('subscan_info', aman.subscan_info) + self.save(proc_aman, calc_aman) def save(self, proc_aman, turn_aman): @@ -1083,9 +1185,9 @@ class PCARelCal(_Preprocess): yfac: 1.5 calc_good_medianw: True lpf: - type: "low_pass_sine2" + type: "sine2" cutoff: 1 - width: 0.1 + trans_width: 0.1 trim_samps: 2000 save: True plot: @@ -1102,6 +1204,7 @@ def __init__(self, step_cfgs): super().__init__(step_cfgs) def calc_and_save(self, aman, proc_aman): + self.plot_signal = self.signal if self.calc_cfgs.get("lpf") is not None: filt = tod_ops.filters.get_lpf(self.calc_cfgs.get("lpf")) filt_tod = tod_ops.fourier_filter(aman, filt, signal_name='signal') @@ -1117,6 +1220,8 @@ def calc_and_save(self, aman, proc_aman): proc_aman.samps.offset + proc_aman.samps.count - trim)) filt_aman.restrict('samps', (filt_aman.samps.offset + trim, filt_aman.samps.offset + filt_aman.samps.count - trim)) + if self.plot_cfgs: + self.plot_signal = filt_aman[self.signal] bands = np.unique(aman.det_info.wafer.bandpass) bands = bands[bands != 'NC'] @@ -1184,7 +1289,7 @@ def plot(self, aman, proc_aman, filename): for band in bands: pca_aman = aman.restrict('dets', aman.dets.vals[proc_aman[self.run_name][f'{band}_idx']], in_place=False) band_aman = proc_aman[self.run_name].restrict('dets', aman.dets.vals[proc_aman[self.run_name][f'{band}_idx']], in_place=False) - plot_pcabounds(pca_aman, band_aman, filename=filename.replace('{name}', f'{ufm}_{band}_pca'), signal=self.signal, band=band, plot_ds_factor=self.plot_cfgs.get('plot_ds_factor', 20)) + plot_pcabounds(pca_aman, band_aman, filename=filename.replace('{name}', f'{ufm}_{band}_pca'), signal=self.plot_signal, band=band, plot_ds_factor=self.plot_cfgs.get('plot_ds_factor', 20)) class PTPFlags(_Preprocess): @@ -1384,3 +1489,4 @@ def save(self, proc_aman, split_flg_aman): _Preprocess.register(DarkDets) _Preprocess.register(SourceFlags) _Preprocess.register(HWPAngleModel) +_Preprocess.register(GetStats) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index 4f173e8e6..9b64f43fd 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -213,7 +213,8 @@ def calc_psd( prefer='center', freq_spacing=None, merge=False, - overwrite=True, + overwrite=True, + subscan=False, **kwargs ): """Calculates the power spectrum density of an input signal using signal.welch(). @@ -234,6 +235,7 @@ def calc_psd( If an nperseg is explicitly passed then that will be used. merge (bool): if True merge results into axismanager. overwrite (bool): if true will overwrite f, pxx axes. + subscan (bool): if True, compute psd on subscans. **kwargs: keyword args to be passed to signal.welch(). Returns: @@ -242,34 +244,40 @@ def calc_psd( """ if signal is None: signal = aman.signal - if timestamps is None: - timestamps = aman.timestamps - - n_samps = signal.shape[-1] - if n_samps <= max_samples: - start = 0 - stop = n_samps + if subscan: + freqs, Pxx = _calc_psd_subscan(aman, signal=signal, freq_spacing=freq_spacing, **kwargs) + axis_map_pxx = [(0, "dets"), (1, "nusamps"), (2, "subscans")] else: - offset = n_samps - max_samples - if prefer == "left": - offset = 0 - elif prefer == "center": - offset //= 2 - elif prefer == "right": - pass - else: - raise ValueError(f"Invalid choise prefer='{prefer}'") - start = offset - stop = offset + max_samples - fs = 1 / np.nanmedian(np.diff(timestamps[start:stop])) - if "nperseg" not in kwargs: - if freq_spacing is not None: - nperseg = int(2 ** (np.around(np.log2(fs / freq_spacing)))) + if timestamps is None: + timestamps = aman.timestamps + + n_samps = signal.shape[-1] + if n_samps <= max_samples: + start = 0 + stop = n_samps else: - nperseg = int(2 ** (np.around(np.log2((stop - start) / 50.0)))) - kwargs["nperseg"] = nperseg + offset = n_samps - max_samples + if prefer == "left": + offset = 0 + elif prefer == "center": + offset //= 2 + elif prefer == "right": + pass + else: + raise ValueError(f"Invalid choice prefer='{prefer}'") + start = offset + stop = offset + max_samples + fs = 1 / np.nanmedian(np.diff(timestamps[start:stop])) + if "nperseg" not in kwargs: + if freq_spacing is not None: + nperseg = int(2 ** (np.around(np.log2(fs / freq_spacing)))) + else: + nperseg = int(2 ** (np.around(np.log2((stop - start) / 50.0)))) + kwargs["nperseg"] = nperseg + + freqs, Pxx = welch(signal[:, start:stop], fs, **kwargs) + axis_map_pxx = [(0, aman.dets), (1, "nusamps")] - freqs, Pxx = welch(signal[:, start:stop], fs, **kwargs) if merge: aman.merge( core.AxisManager(core.OffsetAxis("nusamps", len(freqs)))) if overwrite: @@ -278,9 +286,42 @@ def calc_psd( if "Pxx" in aman._fields: aman.move("Pxx", None) aman.wrap("freqs", freqs, [(0,"nusamps")]) - aman.wrap("Pxx", Pxx, [(0,"dets"),(1,"nusamps")]) + aman.wrap("Pxx", Pxx, axis_map_pxx) return freqs, Pxx +def _calc_psd_subscan(aman, signal=None, freq_spacing=None, **kwargs): + """ + Calculate the power spectrum density of subscans using signal.welch(). + Data defaults to aman.signal. aman.timestamps is used for times. + aman.subscan_info is used to identify subscans. + See calc_psd for arguments. + """ + from .flags import get_subscan_signal + if signal is None: + signal = aman.signal + + fs = 1 / np.nanmedian(np.diff(aman.timestamps)) + if "nperseg" not in kwargs: + if freq_spacing is not None: + nperseg = int(2 ** (np.around(np.log2(fs / freq_spacing)))) + else: + duration_samps = np.asarray([np.ptp(x.ranges()) if x.ranges().size > 0 else 0 for x in aman.subscan_info.subscan_flags]) + duration_samps = duration_samps[duration_samps > 0] + nperseg = int(2 ** (np.around(np.log2(np.median(duration_samps) / 4)))) + kwargs["nperseg"] = nperseg + + Pxx = [] + for iss in range(aman.subscan_info.subscans.count): + signal_ss = get_subscan_signal(aman, signal, iss) + axis = -1 if "axis" not in kwargs else kwargs["axis"] + if signal_ss.shape[axis] >= kwargs["nperseg"]: + freqs, pxx_sub = welch(signal_ss, fs, **kwargs) + Pxx.append(pxx_sub) + else: + Pxx.append(np.full((signal.shape[0], kwargs["nperseg"]//2+1), np.nan)) # Add nans if subscan is too short + Pxx = np.array(Pxx) + Pxx = Pxx.transpose(1, 2, 0) # Dets, nusamps, subscans + return freqs, Pxx def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10): """ @@ -346,13 +387,15 @@ def fit_noise_model( signal=None, f=None, pxx=None, - psdargs=None, + psdargs={}, fwhite=(10, 100), lowf=1, merge_fit=False, f_max=100, merge_name="noise_fit_stats", merge_psd=True, + freq_spacing=None, + subscan=False ): """ Fits noise model with white and 1/f noise to the PSD of signal. @@ -392,6 +435,10 @@ def fit_noise_model( If ``merge_fit`` is True then addes into axis manager with merge_name. merge_psd : bool If ``merg_psd`` is True then adds fres and Pxx to the axis manager. + freq_spacing : float + The approximate desired frequency spacing of the PSD. Passed to calc_psd. + subscan : bool + If True, fit noise on subscans. Returns ------- noise_fit_stats : AxisManager @@ -403,42 +450,48 @@ def fit_noise_model( signal = aman.signal if f is None or pxx is None: - if psdargs is None: - f, pxx = calc_psd( - aman, signal=signal, timestamps=aman.timestamps, merge=merge_psd - ) - else: - f, pxx = calc_psd( - aman, - signal=signal, - timestamps=aman.timestamps, - merge=merge_psd, - **psdargs, - ) - eix = np.argmin(np.abs(f - f_max)) - f = f[1:eix] - pxx = pxx[:, 1:eix] - - fitout = np.zeros((aman.dets.count, 3)) - # This is equal to np.sqrt(np.diag(cov)) when doing curve_fit - covout = np.zeros((aman.dets.count, 3, 3)) - for i in range(aman.dets.count): - p = pxx[i] - wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))]) - pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) - fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) - p0 = [f[fidx], wnest, -pfit[0]] - bounds = [(0, None), (sys.float_info.min, None), (None, None)] - res = minimize(neglnlike, p0, args=(f, p), bounds=bounds, method="Nelder-Mead") - try: - Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True) - hessian_ndt, _ = Hfun(res["x"]) - # Inverse of the hessian is an estimator of the covariance matrix - # sqrt of the diagonals gives you the standard errors. - covout[i] = np.linalg.inv(hessian_ndt) - except np.linalg.LinAlgError: - covout[i] = np.full((3, 3), np.nan) - fitout[i] = res.x + f, pxx = calc_psd( + aman, + signal=signal, + timestamps=aman.timestamps, + freq_spacing=freq_spacing, + merge=merge_psd, + subscan=subscan, + **psdargs, + ) + if subscan: + fitout, covout = _fit_noise_model_subscan(aman, signal, f, pxx, psdargs=psdargs, + fwhite=fwhite, lowf=lowf, f_max=f_max, + freq_spacing=freq_spacing) + axis_map_fit = [(0, "dets"), (1, "noise_model_coeffs"), (2, aman.subscans)] + axis_map_cov = [(0, "dets"), (1, "noise_model_coeffs"), (2, "noise_model_coeffs"), (3, aman.subscans)] + else: + eix = np.argmin(np.abs(f - f_max)) + f = f[1:eix] + pxx = pxx[:, 1:eix] + + fitout = np.zeros((aman.dets.count, 3)) + # This is equal to np.sqrt(np.diag(cov)) when doing curve_fit + covout = np.zeros((aman.dets.count, 3, 3)) + for i in range(aman.dets.count): + p = pxx[i] + wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))]) + pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) + fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) + p0 = [f[fidx], wnest, -pfit[0]] + bounds = [(0, None), (sys.float_info.min, None), (None, None)] + res = minimize(neglnlike, p0, args=(f, p), bounds=bounds, method="Nelder-Mead") + try: + Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True) + hessian_ndt, _ = Hfun(res["x"]) + # Inverse of the hessian is an estimator of the covariance matrix + # sqrt of the diagonals gives you the standard errors. + covout[i] = np.linalg.inv(hessian_ndt) + except np.linalg.LinAlgError: + covout[i] = np.full((3, 3), np.nan) + fitout[i] = res.x + axis_map_fit = [(0, "dets"), (1, "noise_model_coeffs")] + axis_map_cov = [(0, "dets"), (1, "noise_model_coeffs"), (2, "noise_model_coeffs")] noise_model_coeffs = ["fknee", "white_noise", "alpha"] noise_fit_stats = core.AxisManager( @@ -447,18 +500,46 @@ def fit_noise_model( name="noise_model_coeffs", vals=np.array(noise_model_coeffs, dtype=" iqr_range[:, None] * n_sig + + if subscan: + # We include turnarounds + subscan_indices = np.concatenate([aman.flags.left_scan.ranges(), (~aman.flags.left_scan).ranges()]) + else: + subscan_indices = np.array([[0, fvec.shape[1]]]) + + msk = np.zeros_like(fvec, dtype='bool') + for ss in subscan_indices: + iqr_range = 0.741 * stats.iqr(fvec[:,ss[0]:ss[1]:ds], axis=1) + # get flags + msk[:,ss[0]:ss[1]] = fvec[:,ss[0]:ss[1]] > iqr_range[:, None] * n_sig msk[:,:edge_guard] = False msk[:,-edge_guard:] = False flag = RangesMatrix([Ranges.from_bitmask(m) for m in msk]) @@ -672,3 +692,208 @@ def get_inv_var_flags(aman, signal_name='signal', nsigma=5, aman.flags.wrap(inv_var_flag_name, mskinvar, [(0, 'dets'), (1, 'samps')]) return mskinvar + +def get_subscans(aman, merge=True, include_turnarounds=False, overwrite=True): + """ + Returns an axis manager with information about subscans. + This includes direction and a ranges matrix (subscans samps) + True inside each subscan. + + Parameters + ---------- + aman : AxisManager + Input AxisManager. + merge : bool + Merge into aman as 'subscan_info' + include_turnarounds : bool + Include turnarounds in the subscan ranges + overwrite : bool + If true, write over subscan_info. + + Returns + ------- + subscan_aman : AxisManager + AxisManager containing information about the subscans. + "direction" is a (subscans,) array of strings 'left' or 'right' + "subscan_flags" is a (subscans, samps) RangesMatrix; True inside the subscan. + """ + if not include_turnarounds: + ss_ind = (~aman.flags.turnarounds).ranges() # sliceable indices (first inclusive, last exclusive) for subscans + else: + left = aman.flags.left_scan.ranges() + right = aman.flags.right_scan.ranges() + start_left = 0 if (left[0,0] < right[0,0]) else 1 + ss_ind = np.empty((left.shape[0] + right.shape[0], 2), dtype=left.dtype) + ss_ind[start_left::2] = left + ss_ind[(start_left-1)%2::2] = right + + start_inds, end_inds = ss_ind.T + n_subscan = ss_ind.shape[0] + tt = aman.timestamps + subscan_aman = core.AxisManager(aman.samps, core.IndexAxis("subscans", n_subscan)) + + is_left = aman.flags.left_scan.mask()[start_inds] + subscan_aman.wrap('direction', np.array(['left' if is_left[ii] else 'right' for ii in range(n_subscan)]), [(0, 'subscans')]) + + rm = RangesMatrix([Ranges.from_array(np.atleast_2d(ss), tt.size) for ss in ss_ind]) + subscan_aman.wrap('subscan_flags', rm, [(0, 'subscans'), (1, 'samps')]) # True in the subscan + if merge: + name = 'subscan_info' + if overwrite and name in aman: + aman.move(name, None) + aman.wrap(name, subscan_aman) + return subscan_aman + +def get_subscan_signal(aman, arr, isub=None, trim=False): + """ + Split an array into subscans. + + Parameters + ---------- + aman : AxisManager + Input AxisManager. + arr : Array + Input array. + isub : int + Index of the desired subscan. May also be a list of indices. + If None, all are used. + trim : bool + Do not include size-zero arrays from empty subscans in the output. + + Returns + ------- + out : list + If isub is a scalar, return an Array of arr cut on the samps axis to the given subscan. + If isub is a list or None, return a list of such Arrays. + """ + if isinstance(arr, str): + arr = aman[arr] + if np.isscalar(isub): + out = apply_rng(arr, aman.subscan_info.subscan_flags[isub]) + if trim and out.size == 0: + out = None + else: + if isub is None: + isub = range(len(aman.subscan_info.subscan_flags)) + out = [apply_rng(arr, aman.subscan_info.subscan_flags[ii]) for ii in isub] + if trim: + out = [x for x in out if x.size > 0] + + return out + + +def apply_rng(arr, rng): + """ + Apply a Ranges object to an array. rng should be True on the samples you want to keep. + + Parameters + ---------- + arr : Array + Array containing the signal. Should have one axis of len (samps). + rng : Ranges + Ranges object of len (samps) selecting the desired range. + """ + if rng.ranges().size == 0: + slices = [slice(0,0)] # Return an empty array if rng is empty + else: + slices = [slice(*irng) for irng in rng.ranges()] + + # Identify the samps axis + isamps = np.where(np.array(arr.shape) == rng.count)[0] + if isamps.size != 1: + # Check for axis mismatch between arr and rng, or multiple axes with the same size + raise RuntimeError("Could not identify axis matching Ranges") + # Apply ranges + out = [] + for slc in slices: + ndslice = tuple((slice(None) if ii != isamps[0] else slc for ii in range(arr.ndim))) + out.append(arr[ndslice]) + return np.concatenate(out, axis=isamps[0]) + +def wrap_stats(aman, info_aman_name, info, info_names, merge=True): + """ + Wrap multiple stats into a new aman, checking for subscan information. Stats can be (dets,) or (dets, subscans). + + Parameters + ---------- + aman : AxisManager + Input AxisManager. + info_aman_name : str + Name for info_aman when wrapped into input. + info : Array + (stats, dets,) or (stats, dets, subscans) containing the information you want to wrap. + info_names : list + List of str names for each entry in the new aman. + merge : bool + If True merge info_aman into aman. + + Returns + ------- + info_aman : AxisManager + (dets,) or (dets, subscans) aman with a field for each item in info_names. + """ + info_names = np.atleast_1d(info_names) + info = np.atleast_2d(info) + if info.shape == (len(info_names), aman.dets.count): # (stats, dets) + if len(info_names) == aman.dets.count and aman.dets.count == aman.subscan_info.subscans.count: + raise RuntimeError("Cannot infer axis mapping") # Catch corner case + info_aman = core.AxisManager(aman.dets) + axmap = [(0, 'dets')] + + else: + info = np.atleast_3d(info) # (stats, dets, subscans) + info_aman = core.AxisManager(aman.dets, aman.subscan_info.subscans) + axmap = [(0, 'dets'), (1, 'subscans')] + + for ii in range(len(info_names)): + info_aman.wrap(info_names[ii], info[ii], axmap) + if merge: + if info_aman_name in aman.keys(): + aman[info_aman_name].merge(info_aman) + else: + aman.wrap(info_aman_name, info_aman) + return info_aman + +def get_stats(aman, signal, stat_names, split_subscans=False, mask=None, name="stats", merge=False): + """ + Calculate basic statistics on a TOD or power spectrum. + + Parameters + ---------- + aman : AxisManager + Input AxisManager. + signal : Array + Input signal. Statistics will be computed over *axis 1*. + stat_names : list + List of strings identifying which statistics to run. + split_subscans : bool + If True statistics will be computed on subscans. Assumes aman.subscan_info exists already. + mask : Array + Mask to apply before computation. 1d array for advanced indexing (keep True), or a slice object. + name : str + Name of axis manager to add to aman if merge is True. + """ + stat_names = np.atleast_1d(stat_names) + fn_dict = {'mean': np.mean, 'median': np.median, 'ptp': np.ptp, 'std': np.std, + 'kurtosis': stats.kurtosis, 'skew': stats.skew} + + if isinstance(signal, str): + signal = aman[signal] + if split_subscans: + if mask is not None: + raise ValueError("Cannot mask samples and split subscans") + stats_arr = [] + for iss in range(aman.subscan_info.subscans.count): + data = get_subscan_signal(aman, signal, iss) + if data.size > 0: + stats_arr.append([fn_dict[name](data, axis=1) for name in stat_names]) # Samps axis assumed to be 1 + else: + stats_arr.append(np.full((len(stat_names), signal.shape[0]), np.nan)) # Add nans if subscan has been entirely cut + stats_arr = np.array(stats_arr).transpose(1, 2, 0) # stat, dets, subscan + else: + if mask is None: + mask = slice(None) + stats_arr = np.array([fn_dict[name](signal[:, mask], axis=1) for name in stat_names]) # Samps axis assumed to be 1 + + info_aman = wrap_stats(aman, name, stats_arr, stat_names, merge) + return info_aman