diff --git a/sotodlib/obs_ops/splits.py b/sotodlib/obs_ops/splits.py index ffebebc97..5d789c398 100644 --- a/sotodlib/obs_ops/splits.py +++ b/sotodlib/obs_ops/splits.py @@ -150,11 +150,19 @@ def get_split_flags(aman, proc_aman=None, split_cfg=None): if 't2p' in proc_aman: # T2P Leakage split - if not 'high_leakage' in split_cfg: + if 'high_leakage' not in split_cfg: split_cfg['high_leakage'] = 1e-3 - fm.wrap_dets('high_leakage', np.sqrt(proc_aman.t2p.lamQ**2 + proc_aman.t2p.lamU**2) > split_cfg['high_leakage']) - fm.wrap_dets('low_leakage', np.sqrt(proc_aman.t2p.lamQ**2 + proc_aman.t2p.lamU**2) <= split_cfg['high_leakage']) - split_aman.wrap('leakage_avg', np.nanmean(np.sqrt(proc_aman.t2p.lamQ**2 + proc_aman.t2p.lamU**2)), + if 'AQ' in proc_aman.t2p: + leak_fldQ = 'lamQ' + leak_fldU = 'lamU' + elif 'coeffsQ' in proc_aman.t2p: + leak_fldQ = 'coeffsQ' + leak_fldU = 'coeffsU' + else: + raise ValueError('no leakage coefficients found in axis manager') + fm.wrap_dets('high_leakage', np.sqrt(proc_aman.t2p[leak_fldQ]**2 + proc_aman.t2p[leak_fldU]**2) > split_cfg['high_leakage']) + fm.wrap_dets('low_leakage', np.sqrt(proc_aman.t2p[leak_fldQ]**2 + proc_aman.t2p[leak_fldU]**2) <= split_cfg['high_leakage']) + split_aman.wrap('leakage_avg', np.nanmean(np.sqrt(proc_aman.t2p[leak_fldQ]**2 + proc_aman.t2p[leak_fldU]**2)), [(0, 'dets')]) if 'hwpss_stats' in proc_aman: # High 2f amplitude split diff --git a/sotodlib/preprocess/preprocess_util.py b/sotodlib/preprocess/preprocess_util.py index fd88b8394..3d8b2f96f 100644 --- a/sotodlib/preprocess/preprocess_util.py +++ b/sotodlib/preprocess/preprocess_util.py @@ -562,11 +562,19 @@ def multilayer_load_and_preprocess_sim(obs_id, configs_init, configs_proc, pipe_proc = Pipeline(configs_proc["process_pipe"], logger=logger) logger.info("Restricting detectors on all proc pipeline processes") - keep_all = np.ones(meta_proc.dets.count, dtype=bool) - for process in pipe_proc[:]: - keep = process.select(meta_proc, in_place=False) - if isinstance(keep, np.ndarray): - keep_all &= keep + + if ( + 'valid_data' in meta_proc.preprocess and + isinstance(meta_proc.preprocess.valid_data, core.AxisManager) + ): + keep_all = has_any_cuts(meta_proc.preprocess.valid_data.valid_data) + else: + keep_all = np.ones(meta_proc.dets.count, dtype=bool) + for process in pipe_proc[:]: + keep = process.select(meta_proc, in_place=False) + if isinstance(keep, np.ndarray): + keep_all &= keep + meta_proc.restrict("dets", meta_proc.dets.vals[keep_all]) meta_init.restrict('dets', meta_proc.dets.vals) aman = context_init.get_obs(meta_proc, no_signal=True) diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py index 61a4d77d0..2e67ed696 100644 --- a/sotodlib/preprocess/processes.py +++ b/sotodlib/preprocess/processes.py @@ -376,7 +376,6 @@ 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. @@ -593,6 +592,47 @@ def plot(self, aman, proc_aman, filename): plot_signal(aman, signal_name=self.signal, x_name="timestamps", filename=filename, **self.plot_cfgs) +class CutBadDistribution(_Preprocess): + """ + Detector cuts to keep a statistic within some bounds of a gaussian distribution. + + Example config:: + + - name: "cut_bad_dist" + select: + param_name: wn_signal + outlier_range: [0.5, 2.0] + kurtosis_threshold: 2.0 + blame_max: False + blame_min: False + + For parameter options see: :func:`sotodlib.tod_ops.flags.get_good_distribution_flags` + """ + name = "cut_bad_dist" + + def select(self, meta, proc_aman=None, in_place=True): + if self.select_cfgs is None: + return meta + + if proc_aman is None: + proc_aman = meta.preprocess + + stat_name = self.select_cfgs.get('param_name', 'wn_signal') + if stat_name in meta: + keep = tod_ops.flags.get_good_distribution_flags(meta, **self.select_cfgs) + elif stat_name in proc_aman: + keep = tod_ops.flags.get_good_distribution_flags(proc_aman, **self.select_cfgs) + else: + warnings.warn(f'{stat_name} not found in aman or proc_aman not applying bad dist cut.') + return meta + + if in_place: + meta.restrict("dets", meta.dets.vals[keep]) + return meta + else: + return keep + + class Noise(_Preprocess): """ Estimate the white noise levels in the data. Assumes the PSD has been @@ -621,9 +661,16 @@ class Noise(_Preprocess): wn_est: noise fixed_param: 'wn' binning: True + fit_method: log_curve_fit #or likelihood + curve_fit_kwargs: + maxfev: 20000 save: True select: max_noise: 2000 + require_finite_fit: True + + Set ``select.require_finite_fit`` to ``True`` to drop detectors whose fit + parameters contain NaNs (indicating a failed noise fit). Example config block for calculating white noise only:: @@ -641,7 +688,7 @@ class Noise(_Preprocess): If ``fit: True`` this operation will run :func:`sotodlib.tod_ops.fft_ops.fit_noise_model`, else it will run :func:`sotodlib.tod_ops.fft_ops.calc_wn`. - + """ name = "noise" _influx_field = "median_white_noise" @@ -681,6 +728,12 @@ def check_frequency_cutoff(fmin, fmax): if self.fit: fcfgs = copy.deepcopy(self.calc_cfgs) + fit_method = fcfgs.pop('fit_method', None) + curve_fit_kwargs = fcfgs.pop('curve_fit_kwargs', None) + if curve_fit_kwargs is None and fit_method == 'log_curve_fit': + curve_fit_kwargs = {} + if curve_fit_kwargs is not None and not isinstance(curve_fit_kwargs, dict): + raise TypeError("calc.curve_fit_kwargs must be a mapping when provided") fixed_param = fcfgs.get('fixed_param', []) wn_est = fcfgs.get('wn_est', None) @@ -703,11 +756,16 @@ def check_frequency_cutoff(fmin, fmax): fcfgs['subscan'] = self.subscan fcfgs.pop('fwhite', None) f_max = check_frequency_cutoff(fcfgs.get("lowf", None), fcfgs.pop("f_max", 100)) + fit_kwargs = dict(fcfgs) + if fit_method is not None: + fit_kwargs['fit_method'] = fit_method + if curve_fit_kwargs is not None: + fit_kwargs['curve_fit_kwargs'] = curve_fit_kwargs calc_aman = tod_ops.fft_ops.fit_noise_model(aman, pxx=pxx, f=psd.freqs, merge_fit=True, f_max=f_max, - **fcfgs) + **fit_kwargs) if calc_wn or wn_est is None: if not self.subscan: calc_aman.wrap("white_noise", fcfgs['wn_est'], [(0,"dets")]) @@ -783,6 +841,10 @@ def select(self, meta, proc_aman=None, in_place=True): keep &= (wn >= np.float64(self.select_cfgs["min_noise"])) if fk is not None and "max_fknee" in self.select_cfgs.keys(): keep &= (fk <= np.float64(self.select_cfgs["max_fknee"])) + if self.fit and self.select_cfgs.get("require_finite_fit", False): + fit_vals = np.asarray(noise_aman.fit) + fit_flat = fit_vals.reshape(fit_vals.shape[0], -1) + keep &= np.all(np.isfinite(fit_flat), axis=1) if in_place: meta.restrict("dets", meta.dets.vals[keep]) return meta @@ -815,6 +877,8 @@ class Calibrate(_Preprocess): process: kind: "array" cal_array: "cal.array" + select: + cut_array: "cal.missing_cal" # should be 0 where cal is good 1 where missing. """ name = "calibrate" @@ -846,6 +910,16 @@ def process(self, aman, proc_aman, sim=False): raise ValueError(f"Entry '{self.process_cfgs['kind']}'" " not understood") return aman, proc_aman + + def select(self, meta, in_place=True): + if self.select_cfgs is None: + return meta + keep = meta[self.select_cfgs['cut_array']] == 0 + if in_place: + meta.restrict("dets", meta.dets.vals[keep]) + return meta + else: + return keep class EstimateHWPSS(_Preprocess): """ @@ -1236,6 +1310,21 @@ def process(self, aman, proc_aman, sim=False): else: tod_ops.azss.get_azss(aman, **self.calc_cfgs) return aman, proc_aman + + def select(self, meta, in_place=True): + if self.select_cfgs is None: + return meta + if 'bad_dets' in meta[self.azss_stats_name]: + keep = ~meta[self.azss_stats_name]['bad_dets'] + else: + print('No bad_dets field in azss_stats, skipping det selection.') + return meta + + if in_place: + meta.restrict("dets", meta.dets.vals[keep]) + return meta + else: + return keep class SubtractAzSSTemplate(_Preprocess): @@ -1882,11 +1971,12 @@ def process(self, aman, proc_aman, sim=False): filters = [] for spec in filt_list: fname = spec.get("name") - params = tod_ops.fft_ops.build_hpf_params_dict( + params = tod_ops.fft_ops.build_filt_params_dict( fname, noise_fit=noise_fit, filter_params=spec.get("filter_params") ) + ffun = getattr(tod_ops.filters, fname) filters.append(ffun(**params)) @@ -2073,11 +2163,15 @@ def plot(self, aman, proc_aman, filename): class PCAFilter(_Preprocess): """ - Applies a pca filter to the data. + Applies a pca filter to the data. `model_signal` is used to calculate + the PCA modes, which are then subtracted from `signal`. If `model_signal` + is not provided, `signal` is used for both. An example use case is to + use a low-pass filtered version of the signal to calculate the PCA modes. example config file entry:: - name: "pca_filter" + model_signal: "lpf_signal" # optional, if not provided, use signal signal: "signal" process: n_modes: 10 @@ -2088,16 +2182,21 @@ class PCAFilter(_Preprocess): def __init__(self, step_cfgs): self.signal = step_cfgs.get('signal', 'signal') + self.model_signal = step_cfgs.get('model_signal', None) super().__init__(step_cfgs) def process(self, aman, proc_aman, sim=False): n_modes = self.process_cfgs.get('n_modes') signal = aman.get(self.signal) + if self.model_signal: + model_signal = aman.get(self.model_signal) + else: + model_signal = signal if aman.dets.count < n_modes: raise ValueError(f'The number of pca modes {n_modes} is ' f'larger than the number of detectors {aman.dets.count}.') - model = tod_ops.pca.get_pca_model(aman, signal=signal, n_modes=n_modes) + model = tod_ops.pca.get_pca_model(aman, signal=model_signal, n_modes=n_modes) _ = tod_ops.pca.add_model(aman, model, signal=signal, scale=-1) return aman, proc_aman @@ -2108,20 +2207,37 @@ class GetCommonMode(_Preprocess): example config file entry:: - name: "get_common_mode" + noise_fit: True + f_max: 2.0 + wrap_name: "common_demodQ" calc: signal: "signal" method: "median" - wrap: "signal_commonmode" + save: True + If ``noise_fit`` is True, the 1/f noise fit parameters of the common mode + is wrapped together. .. autofunction:: sotodlib.tod_ops.pca.get_common_mode """ name = 'get_common_mode' + def __init__(self, step_cfgs): + self.noise_fit = step_cfgs.get('noise_fit', False) + self.f_max = step_cfgs.get('f_max', None) + self.wrap_name = step_cfgs.get('wrap_name', 'common_mode') + + super().__init__(step_cfgs) def calc_and_save(self, aman, proc_aman): common_mode = tod_ops.pca.get_common_mode(aman, **self.calc_cfgs) - common_aman = core.AxisManager(aman.samps) - common_aman.wrap(self.calc_cfgs['wrap'], common_mode, [(0, 'samps')]) + if self.noise_fit: + common_aman = tod_ops.fft_ops.get_common_noise_params(aman, signal=common_mode, + f_max=self.f_max) + samps = core.OffsetAxis('samps', aman.samps.count) + common_aman.wrap(self.wrap_name, common_mode, [(0, samps)]) + else: + common_aman = core.AxisManager(aman.samps) + common_aman.wrap(self.wrap_name, common_mode, [(0, 'samps')]) self.save(proc_aman, common_aman) return aman, proc_aman @@ -2129,7 +2245,7 @@ def save(self, proc_aman, common_aman): if self.save_cfgs is None: return if self.save_cfgs: - proc_aman.wrap(self.calc_cfgs['wrap'], common_aman) + proc_aman.wrap(self.wrap_name, common_aman) class FilterForSources(_Preprocess): """ @@ -2268,6 +2384,7 @@ class EstimateT2P(_Preprocess): Example config block:: - name : "estimate_t2p" + fit_in_freq : False calc: T_sig_name: 'dsT' Q_sig_name: 'demodQ' @@ -2285,8 +2402,16 @@ class EstimateT2P(_Preprocess): """ name = "estimate_t2p" + def __init__(self, step_cfgs): + self.fit_in_freq = step_cfgs.get('fit_in_freq', False) + + super().__init__(step_cfgs) + def calc_and_save(self, aman, proc_aman): - t2p_aman = tod_ops.t2pleakage.get_t2p_coeffs(aman, **self.calc_cfgs) + if self.fit_in_freq: + t2p_aman = tod_ops.t2pleakage.get_t2p_coeffs_in_freq(aman, **self.calc_cfgs) + else: + t2p_aman = tod_ops.t2pleakage.get_t2p_coeffs(aman, **self.calc_cfgs) self.save(proc_aman, t2p_aman) return aman, proc_aman @@ -2297,6 +2422,18 @@ def save(self, proc_aman, t2p_aman): if self.save_cfgs: proc_aman.wrap("t2p", t2p_aman) + def select(self, meta, proc_aman=None, in_place=True): + if self.select_cfgs is None: + return meta + if proc_aman is None: + proc_aman = meta.preprocess + keep = tod_ops.t2pleakage.get_t2p_cuts(meta, t2p_aman=proc_aman.t2p, in_place=False, **self.select_cfgs) + if in_place: + meta.restrict("dets", meta.dets.vals[keep]) + return meta + else: + return keep + class SubtractT2P(_Preprocess): """Subtract T to P leakage. @@ -2312,8 +2449,7 @@ class SubtractT2P(_Preprocess): name = "subtract_t2p" def process(self, aman, proc_aman, sim=False): - tod_ops.t2pleakage.subtract_t2p(aman, proc_aman['t2p'], - **self.process_cfgs) + tod_ops.t2pleakage.subtract_t2p(aman, proc_aman['t2p']) return aman, proc_aman class SplitFlags(_Preprocess): @@ -2523,6 +2659,35 @@ def process(self, aman, proc_aman, sim=False): self.signal_name_U, merge=True) return aman, proc_aman +class ScanFreqCut(_Preprocess): + """Apply high-pass cut at the scan frequency. + + Example config block:: + + - name : 'scan_freq_cut' + signal_name_T: 'dsT' + signal_name_Q: 'demodQ' + signal_name_U: 'demodU' + process: True + + """ + name = "scan_freq_cut" + + def __init__(self, step_cfgs): + self.signal_name_T = step_cfgs.get('signal_name_T', 'dsT') + self.signal_name_Q = step_cfgs.get('signal_name_Q', 'demodQ') + self.signal_name_U = step_cfgs.get('signal_name_U', 'demodU') + super().__init__(step_cfgs) + + def process(self, aman, proc_aman, sim=False): + scan_freq = tod_ops.utils.get_scan_freq(aman) + hpf_cfg = {'type': 'sine2', 'cutoff': scan_freq, 'trans_width': scan_freq/10} + filt = tod_ops.get_hpf(hpf_cfg) + aman[self.signal_name_T] = tod_ops.fourier_filter(aman, filt, signal_name=self.signal_name_T, detrend='linear') + aman[self.signal_name_Q] = tod_ops.fourier_filter(aman, filt, signal_name=self.signal_name_Q, detrend='linear') + aman[self.signal_name_U] = tod_ops.fourier_filter(aman, filt, signal_name=self.signal_name_U, detrend='linear') + return aman, proc_aman + class FocalplaneNanFlags(_Preprocess): """Find additional detectors which have nans in their focal plane coordinates. @@ -2786,7 +2951,6 @@ def process(self, aman, proc_aman, sim=False): aman.move(**self.process_cfgs) return aman, proc_aman - _Preprocess.register(SplitFlags) _Preprocess.register(SubtractT2P) _Preprocess.register(EstimateT2P) @@ -2828,7 +2992,8 @@ def process(self, aman, proc_aman, sim=False): _Preprocess.register(CombineFlags) _Preprocess.register(RotateFocalPlane) _Preprocess.register(RotateQU) -_Preprocess.register(SubtractQUCommonMode) +_Preprocess.register(SubtractQUCommonMode) +_Preprocess.register(ScanFreqCut) _Preprocess.register(FocalplaneNanFlags) _Preprocess.register(PointingModel) _Preprocess.register(BadSubscanFlags) @@ -2838,4 +3003,4 @@ def process(self, aman, proc_aman, sim=False): _Preprocess.register(SmurfGapsFlags) _Preprocess.register(GetTauHWP) _Preprocess.register(Move) - +_Preprocess.register(CutBadDistribution) diff --git a/sotodlib/tod_ops/__init__.py b/sotodlib/tod_ops/__init__.py index 57a82133f..7e5700e12 100644 --- a/sotodlib/tod_ops/__init__.py +++ b/sotodlib/tod_ops/__init__.py @@ -1,6 +1,6 @@ from .detrend import detrend_tod from .fft_ops import rfft -from .filters import fourier_filter, fft_trim +from .filters import fourier_filter, fft_trim, get_hpf from .gapfill import \ get_gap_fill, get_gap_fill_single, \ get_gap_model, get_gap_model_single @@ -12,3 +12,4 @@ from .azss import get_azss from .t2pleakage import get_t2p_coeffs, subtract_t2p from . import deproject +from . import utils \ No newline at end of file diff --git a/sotodlib/tod_ops/azss.py b/sotodlib/tod_ops/azss.py index 66c772cb3..a3a9baecb 100644 --- a/sotodlib/tod_ops/azss.py +++ b/sotodlib/tod_ops/azss.py @@ -8,6 +8,40 @@ logger = logging.getLogger(__name__) +def _check_azcoverage(aman, flags, az=None, coverage_threshold=0.95, + exclude_turnarounds=True): + """ + Check if the az coverage is sufficient for fitting. + + Parameters + ---------- + az: array-like + A 1D numpy array representing the azimuth angles. + flags: RangesMatrix + Flags indicating invalid samples. + coverage_threshold: float, optional + Threshold for azimuth coverage to consider it sufficient. Default is 0.95. + exclude_turnarounds: bool, optional + Whether to exclude turnaround regions in the azimuth coverage check. Default is True. + + Returns + ------- + bad_dets: boolean array + A boolean array indicating which detectors have insufficient azimuth coverage. + coverages: float array + An array of azimuth coverage fractions for each detector. + """ + if az is None: + az = aman.boresight.az + if exclude_turnarounds: + tot_range = np.ptp(az[~aman.flags.turnarounds.mask()]) + else: + tot_range = np.ptp(az) + if isinstance(flags, str): + flags = aman.flags.get(flags) + coverages = np.array([np.ptp(az[fl])/tot_range if sum(fl) != 0 else 0 for fl in ~flags.mask()]) + return coverages < coverage_threshold, coverages + def bin_by_az(aman, signal=None, az=None, azrange=None, bins=100, flags=None, apodize_edges=True, apodize_edges_samps=1600, @@ -86,7 +120,7 @@ def bin_by_az(aman, signal=None, az=None, azrange=None, bins=100, flags=None, range=azrange, bins=bins, flags=flags, weight_for_signal=weight_for_signal) return binning_dict -def fit_azss(az, azss_stats, max_mode, fit_range=None): +def fit_azss(az, azss_stats, max_mode, fit_range=None, overwrite=False): """ Function for fitting Legendre polynomials to signal binned in azimuth. @@ -103,6 +137,10 @@ def fit_azss(az, azss_stats, max_mode, fit_range=None): Azimuth range used to renormalized to the [-1,1] range spanned by the Legendre polynomials for fitting. Default is the max-min span in the ``binned_az`` array passed in via ``azss_stats``. + overwrite: bool + If overwrite is true will refit the data even if the fit parameters are + already stored in azss_stats. If False will just use the stored + parameters in azss_stats to compute and return the model. Returns ------- @@ -123,6 +161,8 @@ def fit_azss(az, azss_stats, max_mode, fit_range=None): x_legendre = (2 * az - (az_min+az_max)) / (az_max - az_min) x_legendre_bin_centers = (2 * azss_stats.binned_az - (az_min+az_max)) / (az_max - az_min) x_legendre_bin_centers = np.where(~m, np.nan, x_legendre_bin_centers) + if ('coeffs' in azss_stats) and not overwrite: + return L.legval(x_legendre, azss_stats.coeffs.T) coeffs = L.legfit(x_legendre_bin_centers[m], azss_stats.binned_signal[:, m].T, max_mode) coeffs = coeffs.T @@ -147,7 +187,8 @@ def get_azss(aman, signal='signal', az=None, azrange=None, bins=100, flags=None, apply_prefilt=True, prefilt_cfg=None, prefilt_detrend='linear', method='interpolate', max_mode=None, subtract_in_place=False, merge_stats=True, azss_stats_name='azss_stats', - merge_model=True, azss_model_name='azss_model'): + merge_model=True, azss_model_name='azss_model', coverage_threshold=0.95, + exclude_turnarounds=True, return_det_mask=False): """ Derive azss (Azimuth Synchronous Signal) statistics and model from the given axismanager data. **NOTE:** This function does not modify the ``signal`` unless ``subtract_in_place = True``. @@ -207,6 +248,12 @@ def get_azss(aman, signal='signal', az=None, azrange=None, bins=100, flags=None, Boolean flag indicating whether to merge the azss model with the aman. Defaults to True. azss_model_name: string, optional The name to assign to the merged azss model. Defaults to 'azss_model'. + coverage_threshold: float, optional + Minimum azimuth coverage fraction required. Default 0.8 + exclude_turnarounds: bool, optional + Exclude turnarounds when checking coverage. Default True + return_det_mask: bool, optional + If True, return detector mask along with model. Default False Returns ------- @@ -265,7 +312,14 @@ def get_azss(aman, signal='signal', az=None, azrange=None, bins=100, flags=None, azss_stats.wrap('binned_signal', binned_signal, [(0, 'dets'), (1, 'bin_az_samps')]) azss_stats.wrap('binned_signal_sigma', binned_signal_sigma, [(0, 'dets'), (1, 'bin_az_samps')]) azss_stats.wrap('uniform_binned_signal_sigma', uniform_binned_signal_sigma, [(0, 'dets')]) - + if return_det_mask: + bad_dets, coverages = _check_azcoverage(aman, flags=flags, az=az, + coverage_threshold=coverage_threshold, + exclude_turnarounds=exclude_turnarounds) + azss_stats.wrap('bad_dets', bad_dets, [(0, 'dets')]) + azss_stats.wrap('az_coverage', coverages, [(0, 'dets')]) + else: + bad_dets = None model_sig_tod = get_azss_model(aman, azss_stats, az, method, max_mode, azrange) if merge_stats: @@ -277,9 +331,32 @@ def get_azss(aman, signal='signal', az=None, azrange=None, bins=100, flags=None, return azss_stats, model_sig_tod -def get_azss_model(aman, azss_stats, az=None, method='interpolate', max_mode=None, azrange=None): +def get_azss_model(aman, azss_stats, az=None, method='interpolate', + max_mode=None, azrange=None): """ Function to return the azss template for subtraction given the azss_stats AxisManager + + Parameters + ---------- + aman: AxisManager + The axis manager containing the signal + azss_stats: AxisManager + Contains binned azss statistics + az: array-like, optional + Azimuth array. If None, uses aman.boresight.az + method: str + Method for modeling: 'interpolate' or 'fit' + max_mode: int, optional + Maximum Legendre mode for 'fit' method + azrange: list, optional + Azimuth range for fitting + + Returns + ------- + model: array-like + AZSS model for each detector + good_dets_mask: array-like (optional) + Boolean mask of detectors with sufficient coverage (if return_det_mask=True) """ if az is None: az = aman.boresight.az @@ -298,6 +375,9 @@ def get_azss_model(aman, azss_stats, az=None, method='interpolate', max_mode=Non if np.any(~np.isfinite(model)): logger.warning('azss model has nan. set zero to nan but this may make glitch') model[~np.isfinite(model)] = 0 + + if 'bad_dets' in azss_stats: + model[azss_stats['bad_dets'], :] = 0 return model @@ -407,7 +487,12 @@ def subtract_azss_template( if isinstance(signal, str): signal = aman.get(signal) if isinstance(azss, str): - azss = aman.get(azss) + if azss in aman: + azss = aman.get(azss) + elif azss in aman.preprocess: + azss = aman.preprocess.get(azss) + else: + raise ValueError(f'{azss} field not found in aman or aman.preprocess.') if scan_flags is None: scan_flags = np.ones(aman.samps.count, dtype=bool) elif isinstance(scan_flags, str): diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index 19f6ce9a2..c78408ad6 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -11,7 +11,7 @@ import so3g import sys from so3g.proj import Ranges -from scipy.optimize import minimize +from scipy.optimize import curve_fit, minimize from scipy.signal import welch from scipy.stats import chi2 from sotodlib import core, hwp @@ -276,6 +276,7 @@ def calc_psd( overwrite=True, subscan=False, full_output=False, + label_axis='dets', **kwargs ): """Calculates the power spectrum density of an input signal using signal.welch(). @@ -300,6 +301,8 @@ def calc_psd( subscan (bool): if True, compute psd on subscans. full_output: if True this also outputs nseg, the number of segments used for welch, for correcting bias of median white noise estimation by calc_wn. + label_axis (str): The name of LabelAxis in the input aman. + Default is ``dets``. **kwargs: keyword args to be passed to signal.welch(). Returns: @@ -328,7 +331,7 @@ def calc_psd( freqs, Pxx = _calc_psd_subscan(aman, signal=signal, freq_spacing=freq_spacing, **kwargs) - axis_map_pxx = [(0, "dets"), (1, "nusamps"), (2, "subscans")] + axis_map_pxx = [(0, label_axis), (1, "nusamps"), (2, "subscans")] axis_map_nseg = [(0, "subscans")] else: if timestamps is None: @@ -364,7 +367,7 @@ def calc_psd( nseg = int(max_samples / kwargs["nperseg"]) freqs, Pxx = welch(signal[:, start:stop], fs, **kwargs) - axis_map_pxx = [(0, aman.dets), (1, "nusamps")] + axis_map_pxx = [(0, aman[label_axis]), (1, "nusamps")] axis_map_nseg = None if merge: @@ -438,7 +441,7 @@ def _calc_psd_subscan(aman, signal=None, freq_spacing=None, full_output=False, * else: return freqs, Pxx -def calc_wn(aman, pxx=None, freqs=None, nseg=None, low_f=5, high_f=10): +def calc_wn(aman, pxx=None, freqs=None, nseg=None, low_f=5, high_f=10, method='median'): """ Function that calculates the white noise level as a median PSD value between two frequencies. Defaults to calculation of white noise between 5 and 10Hz. @@ -467,6 +470,9 @@ def calc_wn(aman, pxx=None, freqs=None, nseg=None, low_f=5, high_f=10): high_f (float): high frequency cutoff to calculate median psd value. Defaults to 10Hz + + method (str): + median or mean to compute white noise. Default is median. Returns ------- @@ -488,13 +494,23 @@ def calc_wn(aman, pxx=None, freqs=None, nseg=None, low_f=5, high_f=10): '`noverlap=0, full_output=True`') debias = None else: - debias = 2 * nseg / chi2.ppf(0.5, 2 * nseg) + if method == 'median': + debias = 2 * nseg / chi2.ppf(0.5, 2 * nseg) + else: + debias = 1 + + if method == 'median': + _f = np.median + elif method == 'mean': + _f = np.mean + else: + raise ValueError('Only median and mean allowed') fmsk = np.all([freqs >= low_f, freqs <= high_f], axis=0) if pxx.ndim == 1: - wn2 = np.median(pxx[fmsk]) + wn2 = _f(pxx[fmsk]) else: - wn2 = np.median(pxx[:, fmsk], axis=1) + wn2 = _f(pxx[:, fmsk], axis=1) if debias is not None: if pxx.ndim == 3: wn2 *= debias[None, :] @@ -585,6 +601,15 @@ def neglnlike(params, x, y, bin_size=1, **fixed_param): return 1.0e30 return output + +def _curve_fit_bounds_array(bounds): + lower = [] + upper = [] + for low, high in bounds: + lower.append(-np.inf if low is None else low) + upper.append(np.inf if high is None else high) + return np.array(lower, dtype=float), np.array(upper, dtype=float) + def get_psd_mask(aman, psd_mask=None, f=None, mask_hwpss=True, hwp_freq=None, max_hwpss_mode=10, hwpss_width=((-0.4, 0.6), (-0.2, 0.2)), mask_peak=False, peak_freq=None, peak_width=(-0.002, +0.002), @@ -740,7 +765,11 @@ def fit_noise_model( unbinned_mode=3, base=1.05, freq_spacing=None, - subscan=False + subscan=False, + label_axis='dets', + bounds=None, + fit_method="likelihood", + curve_fit_kwargs=None, ): """ Fits noise model with white and 1/f noise to the PSD of binned signal. @@ -803,6 +832,19 @@ def fit_noise_model( The approximate desired frequency spacing of the PSD. Passed to calc_psd. subscan : bool If True, fit noise on subscans. + label_axis : str + The name of LabelAxis in the input aman. + Default is ``dets``. + bounds : list of tuple + List length 2 with number of non-fixed params each list entry is a tuple + first is lower bounds second is upper bounds to constrain fit. + fit_method : str + Selects optimizer. "likelihood" (default) uses the Nelder-Mead + negative log-likelihood fit. "log_curve_fit" performs a + log-space non-linear least square fit via scipy.optimize.curve_fit. + curve_fit_kwargs : dict + Extra options forwarded to curve_fit when fit_method is + "log_curve_fit". Ignored otherwise. Returns ------- noise_fit_stats : AxisManager @@ -824,6 +866,7 @@ def fit_noise_model( merge=merge_psd, subscan=subscan, full_output=True, + label_axis=label_axis, **psdargs, ) if np.any(mask): @@ -842,19 +885,32 @@ def fit_noise_model( pxx = pxx[:, mask] if subscan: - fit_noise_model_kwargs = {"fknee_est": fknee_est, "wn_est": wn_est, "alpha_est": alpha_est, - "lowf": lowf, "f_max": f_max, "fixed_param": fixed_param, - "binning": binning, "unbinned_mode": unbinned_mode, "base": base, - "freq_spacing": freq_spacing} + fit_noise_model_kwargs = { + "fknee_est": fknee_est, + "wn_est": wn_est, + "alpha_est": alpha_est, + "lowf": lowf, + "f_max": f_max, + "fixed_param": fixed_param, + "binning": binning, + "unbinned_mode": unbinned_mode, + "base": base, + "freq_spacing": freq_spacing, + "bounds": bounds, + "fit_method": fit_method, + "curve_fit_kwargs": curve_fit_kwargs, + } fitout, covout = _fit_noise_model_subscan(aman, signal, f, pxx, fit_noise_model_kwargs) - 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)] + axis_map_fit = [(0, label_axis), (1, "noise_model_coeffs"), (2, aman.subscans)] + axis_map_cov = [(0, label_axis), (1, "noise_model_coeffs"), (2, "noise_model_coeffs"), (3, aman.subscans)] else: eix = np.argmin(np.abs(f - f_max)) if lowf is None: six = 1 else: six = np.argmin(np.abs(f - lowf)) + if six == 0: + six = 1 f = f[six:eix] pxx = pxx[:, six:eix] bin_size = 1 @@ -862,66 +918,137 @@ def fit_noise_model( if binning == True: f, pxx, bin_size = get_binned_psd(aman, f=f, pxx=pxx, unbinned_mode=unbinned_mode, base=base, merge=False) - 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)) + fitout = np.zeros((aman[label_axis].count, 3)) + covout = np.zeros((aman[label_axis].count, 3, 3)) if isinstance(wn_est, (int, float)): - wn_est = np.full(aman.dets.count, wn_est) - elif len(wn_est)!=aman.dets.count: + wn_est = np.full(aman[label_axis].count, wn_est) + elif len(wn_est)!=aman[label_axis].count: print('Size of wn_est must be equal to aman.dets.count or a single value.') return if isinstance(fknee_est, (int, float)): - fknee_est = np.full(aman.dets.count, fknee_est) - elif len(fknee_est)!=aman.dets.count: + fknee_est = np.full(aman[label_axis].count, fknee_est) + elif len(fknee_est)!=aman[label_axis].count: print('Size of fknee_est must be equal to aman.dets.count or a single value.') return if isinstance(alpha_est, (int, float)): - alpha_est = np.full(aman.dets.count, alpha_est) - elif len(alpha_est)!=aman.dets.count: + alpha_est = np.full(aman[label_axis].count, alpha_est) + elif len(alpha_est)!=aman[label_axis].count: print('Size of alpha_est must be equal to aman.dets.count or a single value.') return - if fixed_param == None: - initial_params = np.array([wn_est, fknee_est, alpha_est]) - bounds= ((sys.float_info.min, None), (sys.float_info.min, None), (0, 10)) - if fixed_param == "wn": - initial_params = np.array([fknee_est, alpha_est]) - fixed = wn_est - bounds= ((sys.float_info.min, None), (0, 10)) - if fixed_param == "alpha": - initial_params = np.array([wn_est, fknee_est]) - fixed = alpha_est - bounds= ((sys.float_info.min, None), (sys.float_info.min, None)) - - for i in range(len(pxx)): - p = pxx[i] - p0 = initial_params.T[i] + + wn_est = np.asarray(wn_est, dtype=float) + fknee_est = np.asarray(fknee_est, dtype=float) + alpha_est = np.asarray(alpha_est, dtype=float) + + fixed_vals = None + if fixed_param is None: + initial_params = np.vstack((wn_est, fknee_est, alpha_est)) + if bounds is None: + bounds = ( + (sys.float_info.min, None), + (sys.float_info.min, None), + (0, 10), + ) + elif fixed_param == "wn": + initial_params = np.vstack((fknee_est, alpha_est)) + fixed_vals = np.asarray(wn_est, dtype=float) + if bounds is None: + bounds = ( + (sys.float_info.min, None), + (0, 10), + ) + elif fixed_param == "alpha": + initial_params = np.vstack((wn_est, fknee_est)) + fixed_vals = np.asarray(alpha_est, dtype=float) + if bounds is None: + bounds = ( + (sys.float_info.min, None), + (sys.float_info.min, None), + ) + else: + raise ValueError("fixed_param must be 'wn', 'alpha', or None.") + + n_free_params = initial_params.shape[0] + if isinstance(bounds, list): + bounds = [tuple(b) for b in bounds] + if isinstance(bounds, tuple): + pass + else: + bounds = tuple(bounds) + + if len(bounds) != n_free_params: + raise ValueError("Number of bounds entries must match free parameters.") + + curve_bounds = None + if fit_method == "log_curve_fit": + curve_bounds = _curve_fit_bounds_array(bounds) + + for i, p in enumerate(pxx): + p0 = initial_params[:, i] _fixed = {} - if fixed_param != None: - _fixed = {fixed_param: fixed[i]} - res = minimize(lambda params: neglnlike(params, f, p, bin_size=bin_size, **_fixed), - p0, bounds=bounds, method="Nelder-Mead") - with warnings.catch_warnings(): - warnings.filterwarnings("error") - try: - Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p, bin_size=bin_size, **_fixed), 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: - print( - f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping. (LinAlgError)" - ) - covout_i = np.full((len(p0), len(p0)), np.nan) - except IndexError: + if fixed_param is not None: + _fixed = {fixed_param: fixed_vals[i]} + + if fit_method == "likelihood": + res = minimize( + lambda params: neglnlike(params, f, p, bin_size=bin_size, **_fixed), + p0, + bounds=bounds, + method="Nelder-Mead", + ) + with warnings.catch_warnings(): + warnings.filterwarnings("error") + try: + Hfun = ndt.Hessian( + lambda params: neglnlike(params, f, p, bin_size=bin_size, **_fixed), + full_output=True, + ) + hessian_ndt, _ = Hfun(res["x"]) + covout_i = np.linalg.inv(hessian_ndt) + except np.linalg.LinAlgError: + print( + f"Cannot calculate Hessian for detector {aman[label_axis].vals[i]} skipping. (LinAlgError)" + ) + covout_i = np.full((len(p0), len(p0)), np.nan) + except IndexError: + print( + f"Cannot calculate Hessian for detector {aman[label_axis].vals[i]} skipping. (IndexError)" + ) + covout_i = np.full((len(p0), len(p0)), np.nan) + except RuntimeWarning as e: + covout_i = np.full((len(p0), len(p0)), np.nan) + print( + f"RuntimeWarning: {e}\n Hessian failed because results are: {res['x']}, for det: {aman[label_axis].vals[i]}" + ) + fitout_i = res.x + else: + valid = np.isfinite(p) & (p > 0) & np.isfinite(f) + if np.count_nonzero(valid) < n_free_params: print( - f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping. (IndexError)" + f"Insufficient valid samples for detector {aman[label_axis].vals[i]} in log_curve_fit." ) - covout_i = np.full((len(p0), len(p0)), np.nan) - except RuntimeWarning as e: - covout_i = np.full((len(p0), len(p0)), np.nan) - print(f'RuntimeWarning: {e}\n Hessian failed because results are: {res["x"]}, for det: {aman.dets.vals[i]}') - fitout_i = res.x + fitout_i = np.full(n_free_params, np.nan) + covout_i = np.full((n_free_params, n_free_params), np.nan) + else: + try: + popt, pcov = curve_fit( + lambda freq, *params: np.log(noise_model(freq, params, **_fixed)), + f[valid], + np.log(p[valid]), + p0=p0, + bounds=curve_bounds, + **curve_fit_kwargs, + ) + except (RuntimeError, ValueError) as err: + print( + f"curve_fit failed for detector {aman[label_axis].vals[i]}: {err}" + ) + fitout_i = np.full(n_free_params, np.nan) + covout_i = np.full((n_free_params, n_free_params), np.nan) + else: + fitout_i = popt + covout_i = pcov + if fixed_param == "wn": covout_i = np.insert(covout_i, 0, 0, axis=0) covout_i = np.insert(covout_i, 0, 0, axis=1) @@ -932,15 +1059,16 @@ def fit_noise_model( covout_i = np.insert(covout_i, 2, 0, axis=1) covout_i[2][2] = np.nan fitout_i = np.insert(fitout_i, 2, alpha_est[i]) + covout[i] = covout_i fitout[i] = fitout_i - axis_map_fit = [(0, "dets"), (1, "noise_model_coeffs")] - axis_map_cov = [(0, "dets"), (1, "noise_model_coeffs"), (2, "noise_model_coeffs")] + axis_map_fit = [(0, label_axis), (1, "noise_model_coeffs")] + axis_map_cov = [(0, label_axis), (1, "noise_model_coeffs"), (2, "noise_model_coeffs")] noise_model_coeffs = ["white_noise", "fknee", "alpha"] noise_fit_stats = core.AxisManager( - aman.dets, + aman[label_axis], core.LabelAxis( name="noise_model_coeffs", vals=np.array(noise_model_coeffs, dtype="= rn_range[0]) ranges.append(detcal.r_n <= rn_range[1]) - if si_nan: - ranges.append(np.isnan(detcal.s_i) == False) + if si_range is not None: + ranges.append(detcal.s_i >= si_range[0]) + ranges.append(detcal.s_i <= si_range[1]) if phase_to_pW is not None: ranges.append(detcal.phase_to_pW >= phase_to_pW[0]) ranges.append(detcal.phase_to_pW <= phase_to_pW[1]) @@ -129,6 +130,11 @@ def get_det_bias_flags(aman, detcal=None, rfrac_range=(0.1, 0.7), msk_names.extend(['r_n_gt', 'r_n_lt']) ranges.append(detcal.r_n >= rn_range[0]) ranges.append(detcal.r_n <= rn_range[1]) + + if si_range is not None: + msk_names.extend(['s_i_gt', 's_i_lt']) + ranges.append(detcal.s_i >= si_range[0]) + ranges.append(detcal.s_i <= si_range[1]) if phase_to_pW is not None: msk_names.extend(['phase_to_pW_gt', 'phase_to_pW_lt']) @@ -1170,3 +1176,59 @@ def expand_smurfgaps_flags(aman, buffer=200, name='smurfgaps', merge=True): if merge: aman.flags.wrap('smurfgaps', smurfgaps, [(0, 'dets'), (1, 'samps')]) return smurfgaps + +def get_good_distribution_flags(aman, param_name='wn_signal', + outlier_range=(0.5, 2.), kurtosis_threshold=2., + blame_max=False, blame_min=False): + """ + Get detector cuts based on the detectors which fit within a gaussian distribution. + Outlier and kurtosis thresholds set the cuts and param name determines which + statistic's distribution is evaluated to generate the cuts. + + Parameters + ---------- + aman : AxisManager + Input axis manager. + param_name : str + Field in aman where to evaluate distribution of. + outlier_range : tuple + Min and max of distribution allowed + kurtosis_threshold : float + Maximum allowed kurtosis of final distribution + blame_max : bool + ADD DESCRIPTION + blame_min : bool + ADD DESCRIPTION + + Returns + ------- + det_mask : list(bool) + Boolean mask for cuts True for keep and False for cut. + """ + det_mask = np.full(aman.dets.count, True, dtype=bool) + ratio = aman[param_name]/np.nanmedian(aman[param_name]) + outlier_mask = (ratio 0: + distributions = aman[param_name][det_mask] + else: + break + kurtosis_wn = stats.kurtosis(distributions) + if np.abs(kurtosis_wn) < kurtosis_threshold: + break + else: + assert (blame_max==False) or (blame_min==False) + if blame_max: + det_mask[aman[param_name] >= np.nanmax(distributions)] = False + elif blame_min: + det_mask[aman[param_name] <= np.nanmin(distributions)] = False + else: + max_is_bad_factor = np.nanmax(distributions)/np.nanmedian(distributions) + min_is_bad_factor = np.nanmedian(distributions)/np.nanmin(distributions) + if max_is_bad_factor > min_is_bad_factor: + det_mask[aman[param_name] >= np.nanmax(distributions)] = False + else: + det_mask[aman[param_name] <= np.nanmin(distributions)] = False + return det_mask diff --git a/sotodlib/tod_ops/pca.py b/sotodlib/tod_ops/pca.py index b8d4ef663..8f0bafa83 100644 --- a/sotodlib/tod_ops/pca.py +++ b/sotodlib/tod_ops/pca.py @@ -129,7 +129,8 @@ def get_pca(tod=None, cov=None, signal=None, wrap=None, mask=None): E, R = np.linalg.eigh(cov) E[np.isnan(E)] = 0. - E, R = E.real, R.real + if np.isrealobj(signal): + E, R = E.real, R.real idx = np.argsort(-E) diff --git a/sotodlib/tod_ops/t2pleakage.py b/sotodlib/tod_ops/t2pleakage.py index 2f43cc5f6..6a3804aa8 100644 --- a/sotodlib/tod_ops/t2pleakage.py +++ b/sotodlib/tod_ops/t2pleakage.py @@ -6,6 +6,7 @@ from sotodlib.tod_ops.fft_ops import calc_psd, calc_wn from scipy.odr import ODR, Model, RealData from lmfit import Model as LmfitModel +from scipy.fft import rfft, rfftfreq def t2p_fit(aman, T_sig_name='dsT', Q_sig_name='demodQ', U_sig_name='demodU', @@ -300,6 +301,126 @@ def get_t2p_coeffs(aman, T_sig_name='dsT', Q_sig_name='demodQ', U_sig_name='demo return out_aman +def get_t2p_coeffs_in_freq(aman, T_sig_name='dsT', Q_sig_name='demodQ', U_sig_name='demodU', + fs=None, fit_freq_range=(0.01, 0.1), wn_freq_range=(0.2, 1.9), + subtract_sig=False, merge_stats=True, t2p_stats_name='t2p_stats', + ): + """ + Compute the leakage coefficients from temperature (T) to polarization (Q and U) in Fourier + domain. Return an axismanager of the coefficients with their statistical uncertainties and + reduced chi-squared values for the fit. + + Parameters + ---------- + aman : AxisManager + AxisManager object containing the TOD data. + T_sig_name : str + Name of the temperature signal in `aman`. Default is 'dsT'. + Q_sig_name : str + Name of the Q polarization signal in `aman`. Default is 'demodQ'. + U_sig_name : str + Name of the U polarization signal in `aman`. Default is 'demodU'. + fs: float + The sampling frequency. If it is None, it will be calculated. Default is None. + fit_range_freq: tuple + The start/end frequencies of the t2p fit. Default is (0.01, 0.1). + wn_freq_range: tuple + The start/end frequencies to calculate the white noise level of demod signal. + Default is (0.2, 1.9). + subtract_sig : bool + Whether to subtract the calculated leakage from the polarization signals. Default is False. + merge_stats : bool + Whether to merge the calculated statistics back into `aman`. Default is True. + t2p_stats_name : str + Name under which to wrap the output AxisManager containing statistics. Default is 't2p_stats'. + + Returns + ------- + out_aman : AxisManager + An AxisManager containing leakage coefficients, their errors, and reduced + chi-squared statistics. + """ + if fs is None: + fs = np.median(1/np.diff(aman.timestamps)) + N = aman.samps.count + freqs = rfftfreq(N, d=1/fs) + I_fs = rfft(aman[T_sig_name], axis=1) + Q_fs = rfft(aman[Q_sig_name], axis=1) + U_fs = rfft(aman[U_sig_name], axis=1) + + coeffsQ = np.zeros(aman.dets.count) + errorsQ = np.zeros(aman.dets.count) + redchi2sQ = np.zeros(aman.dets.count) + coeffsU = np.zeros(aman.dets.count) + errorsU = np.zeros(aman.dets.count) + redchi2sU = np.zeros(aman.dets.count) + + fit_mask = (fit_freq_range[0] < freqs) & (freqs < fit_freq_range[1]) + wn_mask = (wn_freq_range[0] < freqs) & (freqs < wn_freq_range[1]) + + def leakage_model(B, x): + return B[0] * x + + model = Model(leakage_model) + + for i in range(aman.dets.count): + # fit Q + Q_wnl = np.nanmean(np.abs(Q_fs[i][wn_mask])) + x = np.real(I_fs[i])[fit_mask] + y = np.real(Q_fs[i])[fit_mask] + sx = Q_wnl / np.sqrt(2) * np.ones_like(x) + sy = Q_wnl * np.ones_like(y) + try: + data = RealData(x=x, + y=y, + sx=sx, + sy=sy) + odr = ODR(data, model, beta0=[1e-3]) + output = odr.run() + coeffsQ[i] = output.beta[0] + errorsQ[i] = output.sd_beta[0] + redchi2sQ[i] = output.sum_square / (len(x) - 2) + except: + coeffsQ[i] = np.nan + errorsQ[i] = np.nan + redchi2sQ[i] = np.nan + + #fit U + U_wnl = np.nanmean(np.abs(U_fs[i][wn_mask])) + x = np.real(I_fs[i])[fit_mask] + y = np.real(U_fs[i])[fit_mask] + sx = U_wnl / np.sqrt(2) * np.ones_like(x) + sy = U_wnl * np.ones_like(y) + try: + data = RealData(x=x, + y=y, + sx=sx, + sy=sy) + odr = ODR(data, model, beta0=[1e-3]) + output = odr.run() + coeffsU[i] = output.beta[0] + errorsU[i] = output.sd_beta[0] + redchi2sU[i] = output.sum_square / (len(x) - 2) + except: + coeffsU[i] = np.nan + errorsU[i] = np.nan + redchi2sU[i] = np.nan + + out_aman = core.AxisManager(aman.dets, aman.samps) + out_aman.wrap('coeffsQ', coeffsQ, [(0, 'dets')]) + out_aman.wrap('errorsQ', errorsQ, [(0, 'dets')]) + out_aman.wrap('redchi2sQ', redchi2sQ, [(0, 'dets')]) + out_aman.wrap('coeffsU', coeffsU, [(0, 'dets')]) + out_aman.wrap('errorsU', errorsU, [(0, 'dets')]) + out_aman.wrap('redchi2sU', redchi2sU, [(0, 'dets')]) + + if subtract_sig: + subtract_t2p(aman, out_aman) + if merge_stats: + aman.wrap(t2p_stats_name, out_aman) + + return out_aman + def subtract_t2p(aman, t2p_aman, T_signal=None): """ Subtract T to P leakage. @@ -329,3 +450,63 @@ def subtract_t2p(aman, t2p_aman, T_signal=None): aman.demodU -= np.multiply(T_signal.T, t2p_aman.coeffsU).T else: raise ValueError('no leakage coefficients found in axis manager') + +def get_t2p_cuts(aman, t2p_aman=None, redchi2s=True, error=True, lam=False, + redchi2s_lims=(0.1, 3), error_lims=(0, 0.03), lam_lims=(0, 0.01), + in_place=False): + """ + Restrict detectors based on the t2p fit stats or t2p coefficient. + + Parameters + ---------- + aman : AxisManager + The tod. + t2p_aman : AxisManager + Axis manager with Q and U leakage coeffients. + If joint fitting was used in get_t2p_coeffs, Q coeffs are in + fields ``lamQ`` and ``AQ`` and U coeffs are in ``lamU`` and + ``AU``. Otherwise Q coeff is in field ``coeffsQ`` and U coeff + in ``coeffsU``. + redchi2s : bool + If True, restrict detectors based on redchi2s values. + error : bool + If True, restrict detectors based on fit errors of lamQ and lamU. + lam : bool + If True, restrict detectors based on the amplitude of lamQ and lamU. + redchi2s_lims : tuple + The lower and upper limit of acceptable redchi2s. + error_lims : tuple + The lower and upper limit of acceptable errors. + lam_lims : tuple + The lower and upper limit of acceptable leakage coefficient. + """ + if t2p_aman is None: + if 't2p_stats' not in aman: + raise ValueError('t2p_aman must be provided or already in aman') + t2p_aman = aman.t2p_stats + mask = np.ones(aman.dets.count, dtype=bool) + if redchi2s: + if 'redchi2s' in t2p_aman: + redchi2s_t2p = t2p_aman.redchi2s + elif 'redchi2sQ' in t2p_aman: + redchi2s_t2p = t2p_aman.redchi2sQ + t2p_aman.redchi2sU + mask_redchi2s = (redchi2s_lims[0] < redchi2s_t2p) & (redchi2s_t2p < redchi2s_lims[1]) + mask = np.logical_and(mask, mask_redchi2s) + if error: + if 'lamQ_error' in t2p_aman: + error_t2p = np.sqrt(t2p_aman.lamQ_error**2+t2p_aman.lamU_error**2) + elif 'errorsQ' in t2p_aman: + error_t2p = np.sqrt(t2p_aman.errorsQ**2+t2p_aman.errorsU**2) + mask_error = (error_lims[0] < error_t2p) & (error_t2p < error_lims[1]) + mask = np.logical_and(mask, mask_error) + if lam: + if 'lamQ' in t2p_aman: + lam_t2p = np.sqrt(t2p_aman.lamQ**2 + t2p_aman.lamU**2) + elif 'coeffsQ' in t2p_aman: + lam_t2p = np.sqrt(t2p_aman.coeffsQ**2 + t2p_aman.coeffsU**2) + mask_lam = (lam_lims[0] < lam_t2p) & (lam_t2p < lam_lims[1]) + mask = np.logical_and(mask, mask_lam) + if in_place: + aman.restrict('dets', aman.dets.vals[mask]) + else: + return mask diff --git a/sotodlib/tod_ops/utils.py b/sotodlib/tod_ops/utils.py index 25b16a91c..0e198f158 100644 --- a/sotodlib/tod_ops/utils.py +++ b/sotodlib/tod_ops/utils.py @@ -71,3 +71,51 @@ def get_block_moment( block_moment64(tod, output, block_size, moment, central, shift) return output.reshape(orig_shape) + +def get_scan_speed(aman, in_deg=False, wrap=False, wrap_name='scanspeed'): + """ + Compute the azimuth scan speed of the telescope in [rad/s]. + + Arguments: + + aman (axismanager): Observation TOD with boresight data. + + in_deg (bool): If True, this returns scan speed in [deg/s]. + + wrap (bool): If True, scan speed is wrapped in TOD. + + wrap_name (str): The name to wrap the scan speed. + Returns: + + scanspeed (float): The azimuth scan speed. + """ + scanspeed = np.median(np.abs(np.diff(aman.boresight.az))/np.diff(aman.timestamps)) + if in_deg: + scanspeed = 180/np.pi * scanspeed + if wrap: + aman.wrap(wrap_name, scanspeed) + return scanspeed + +def get_scan_freq(aman, wrap=False, wrap_name='scanfreq'): + """ + Compute the azimuth scan frequency of the telescope in [1/s]. + Here the scan frequency is defined as the inverse of the period + of one set of left-going and right-going scan. + + Arguments: + + aman (axismanager): Observation TOD with boresight data. + + wrap (bool): If True, scan frequency is wrapped in TOD. + + wrap_name (str): The name to wrap the scan frequency. + Returns: + + scanfreq (float): The azimuth scan frequency. + """ + scanspeed = get_scan_speed(aman, in_deg=True) + scanfreq = 1/(aman.obs_info.az_throw*4/scanspeed) + if wrap: + aman.wrap(wrap_name, scanfreq) + return scanfreq +