From 7969bedf5de1432dbf2ab6f9111a1de55378ab29 Mon Sep 17 00:00:00 2001 From: Junna Sugiyama Date: Thu, 23 May 2024 02:33:41 +0000 Subject: [PATCH 01/23] added masking and binning function --- sotodlib/tod_ops/fft_ops.py | 268 +++++++++++++++++++++++++++++++++++- 1 file changed, 264 insertions(+), 4 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index e31fe51de..3db1d125c 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -269,12 +269,13 @@ def calc_psd( freqs, Pxx = welch(signal[:, start:stop], fs, **kwargs) if merge: - aman.merge( core.AxisManager(core.OffsetAxis("nusamps", len(freqs)))) if overwrite: + if "freqs" in aman._fields: aman.move("freqs", None) if "Pxx" in aman._fields: aman.move("Pxx", None) + aman.merge( core.AxisManager(core.OffsetAxis("nusamps", len(freqs)))) aman.wrap("freqs", freqs, [(0,"nusamps")]) aman.wrap("Pxx", Pxx, [(0,"dets"),(1,"nusamps")]) return freqs, Pxx @@ -338,6 +339,147 @@ def neglnlike(params, x, y): return 1.0e30 return output +def calc_masked_psd( + aman, + f=None, + pxx=None, + hwpss=False, + hwp_freq=None, + f_max=100, + width_for_1f=(-0.4, +0.6), + width_for_Nf=(-0.2, +0.2), + peak=False, + peak_freq=None, + peak_width=(-0.002, +0.002), + merge=False, + overwrite=True, +): + """ + Function that masks hwpss or single peak in PSD. + + Arguments + --------- + aman : AxisManager + Axis manager which has samps axis aligned with signal. + f : nparray + Frequency of PSD of signal. + pxx : nparray + PSD of signal. + hwpss : bool + If True, hwpss are masked. + hwp_freq : float + HWP frequency. + f_max : float + Maximum frequency to include in the hwpss masking. + width_for_1f : tuple + Range to mask 1f hwpss. The default masked range will be + (hwp_freq - 0.4) < f < (hwp_freq + 0.6). + Usually 1f hwpss distrubites wider than other hwpss. + width_for_Nf : tuple + Range to mask Nf hwpss. The default masked range will be + (hwp_freq*N - 0.2) < f < (hwp_freq*N + 0.2). + peak : bool + If True, single peak is masked. + peak_freq : float + Center frequency of the mask. + peak_width : tuple + Range to mask signal. The default masked range will be + (peak_freq - 0.002) < f < (peak_freq + 0.002). + merge : bool + if True merge results into axismanager. + overwrite: bool + if true will overwrite f, pxx axes. You cannot overwrite with data + longer than before. + + Returns + ------- + p_mask: MaskedArray of the input PSD. p_mask.data returns the PSD before masking, + while p_mask.mask returns the bool array of the mask. + """ + if f is None or pxx is None: + f = aman.freqs + pxx = aman.Pxx + if hwpss: + pxx_masked = [] + if hwp_freq is None: + hwp_freq = np.median(aman['hwp_solution']['raw_approx_hwp_freq_1'][1]) + for i in range(aman.dets.count): + masked_signal = mask_hwpss(f, pxx[i], hwp_freq, f_max=f_max, width_for_1f=width_for_1f, width_for_Nf=width_for_Nf) + pxx_masked.append(masked_signal.compressed()) + pxx = np.array(pxx_masked) + f = np.ma.masked_array(f, mask=masked_signal.mask).compressed() + if peak: + pxx_masked = [] + for i in range(aman.dets.count): + masked_signal = mask_peak(f, pxx[i], peak_freq, peak_width=peak_width) + pxx_masked.append(masked_signal.compressed()) + pxx = np.array(pxx_masked) + f = np.ma.masked_array(f, mask=masked_signal.mask).compressed() + if merge: + aman.merge( core.AxisManager(core.OffsetAxis("nusamps", len(f)))) + if overwrite: + if "freqs" in aman._fields: + aman.move("freqs", None) + if "Pxx" in aman._fields: + aman.move("Pxx", None) + aman.wrap("freqs", f, [(0,"nusamps")]) + aman.wrap("Pxx", pxx, [(0,"dets"),(1,"nusamps")]) + return f, pxx + +def calc_binned_psd( + aman, + f=None, + pxx=None, + unbinned_mode=10, + base=2, + merge=False, + overwrite=True, +): + """ + Function that masks hwpss in PSD. + + Arguments + --------- + aman : AxisManager + Axis manager which has samps axis aligned with signal. + f : nparray + Frequency of PSD of signal. + pxx : nparray + PSD of signal. + hwpss : bool + If True, hwpss are masked. + unbinned_mode : int + First Fourier modes up to this number are left un-binned. + base : float (> 1) + Base of the logspace bins. + merge : bool + if True merge results into axismanager. + overwrite: bool + if true will overwrite f, pxx axes. + + Returns + ------- + p_mask: MaskedArray of the input PSD. p_mask.data returns the PSD before masking, + while p_mask.mask returns the bool array of the mask. + """ + if f is None or pxx is None: + f = aman.freqs + pxx = aman.Pxx + f_binned = binning_psd(f, unbinned_mode=unbinned_mode, base=base) + pxx_binned = [] + for i in range(aman.dets.count): + pxx_binned.append(binning_psd(pxx[i], unbinned_mode=unbinned_mode, base=base)) + pxx_binned = np.array(pxx_binned) + if merge: + aman.merge( core.AxisManager(core.OffsetAxis("nusamps_bin", len(f_binned)))) + if overwrite: + if "freqs_bin" in aman._fields: + aman.move("freqs_bin", None) + if "Pxx_bin" in aman._fields: + aman.move("Pxx_bin", None) + aman.wrap("freqs_bin", f_binned, [(0,"nusamps_bin")]) + aman.wrap("Pxx_bin", pxx_binned, [(0,"dets"),(1,"nusamps_bin")]) + return f_binned, pxx_binned def fit_noise_model( aman, @@ -348,6 +490,7 @@ def fit_noise_model( fwhite=(10, 100), lowf=1, merge_fit=False, + f_min=None, f_max=100, merge_name="noise_fit_stats", merge_psd=True, @@ -382,6 +525,9 @@ def fit_noise_model( for initial guess passed to ``scipy.signal.curve_fit``. merge_fit : bool Merges fit and fit statistics into input axis manager. + f_min : float + Minimum frequency to include in the fitting. + Default is None which selects f_min as the second index of f. f_max : float Maximum frequency to include in the fitting. This is particularly important for lowpass filtered data such as that post demodulation @@ -414,9 +560,12 @@ def fit_noise_model( **psdargs, ) eix = np.argmin(np.abs(f - f_max)) - f = f[1:eix] - pxx = pxx[:, 1:eix] - + if f_min is None: + six = 1 + else: + six = np.argmin(np.abs(f - f_min)) + f = f[six:eix] + pxx = pxx[:, six: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)) @@ -438,6 +587,11 @@ def fit_noise_model( f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping." ) covout[i] = np.full((3, 3), np.nan) + except IndexError: + print( + f"Cannot fit 1/f curve for detector {aman.dets.vals[i]} skipping." + ) + covout[i] = np.full((3, 3), np.nan) fitout[i] = res.x noise_model_coeffs = ["fknee", "white_noise", "alpha"] @@ -457,3 +611,109 @@ def fit_noise_model( if merge_fit: aman.wrap(merge_name, noise_fit_stats) return noise_fit_stats + +def mask_hwpss(f, p, hwp_freq, f_max=100, width_for_1f=(-0.4, +0.6), width_for_Nf=(-0.2, +0.2)): + """ + Function that masks hwpss in PSD. + + Arguments + --------- + f : nparray + Frequency of PSD of signal. + + p : nparray + PSD of signal. + + hwp_freq : float + HWP frequency. + + f_max : float + Maximum frequency to include in the fitting. + + width_for_1f : tuple + Range to mask 1f hwpss. The default masked range will be + (hwp_freq - 0.4) < f < (hwp_freq + 0.6). + Usually 1f hwpss distrubites wider than other hwpss. + width_for_Nf : tuple + Range to mask Nf hwpss. The default masked range will be + (hwp_freq*N - 0.2) < f < (hwp_freq*N + 0.2). + + Returns + ------- + p_mask: MaskedArray of the input PSD. p_mask.data returns the PSD before masking, + while p_mask.mask returns the bool array of the mask. + """ + mask_arrays = [((f > hwp_freq + width_for_1f[0]) & (f < hwp_freq + width_for_1f[1]))] + for n in range(int(f_max//hwp_freq-1)): + mask_arrays.append(((f > hwp_freq*(n+2) + width_for_Nf[0]) & (f < hwp_freq*(n+2) + width_for_Nf[1]))) + mask = np.any(np.array(mask_arrays), axis=0) + p_mask = np.ma.masked_where(mask, p) + return p_mask + +def mask_peak(f, p, peak_freq, peak_width=(-0.002, +0.002)): + """ + Function that masks single peak (e.g. scan synchronous signal) in PSD. + + Arguments + --------- + f : nparray + Frequency of PSD of signal. + + p : nparray + PSD of signal. + + peak_freq : float + Center frequency of the mask. + + peak_width : tuple + Range to mask signal. The default masked range will be + (peak_freq - 0.002) < f < (peak_freq + 0.002). + + Returns + ------- + p_mask: MaskedArray of the input PSD. p_mask.data returns the PSD before masking, + while p_mask.mask returns the bool array of the mask. + """ + mask = (f > peak_freq + peak_width[0]) & (f < peak_freq + peak_width[1]) + p_mask = np.ma.masked_where(mask, p) + return p_mask + +def binning_psd(psd, unbinned_mode=10, base=2, drop_nan=False): + """ + Function to bin PSD. + First several Fourier modes are left un-binned. + Fourier modes higher than that are averaged into logspace bins. + + Arguments + --------- + psd : nparray + PSD of signal. + + unbinned_mode : int + First Fourier modes up to this number are left un-binned. + + base : float (> 1) + Base of the logspace bins. + + drop_nan : bool + If True, drop the index where p is nan. + + Returns + ------- + binned_psd: nparray + Binned PSD. + """ + binned_psd = [] + for i in np.linspace(0, unbinned_mode, unbinned_mode+1, dtype = int): + binned_psd.append(psd[i]) + N = int(np.emath.logn(base, len(psd)-unbinned_mode)) + binning_idx = np.logspace(base, N, N, base=base, dtype = int)+unbinned_mode-1 + for i in range(N-1): + if binning_idx[i] == binning_idx[i+1]: + continue + else: + binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) + binned_psd = np.array(binned_psd) + if drop_nan: + binned_psd = binned_psd[~np.isnan(binned_psd)] + return binned_psd \ No newline at end of file From 5e88252c5e5139a58c430d58f3eed5e8b64f53d6 Mon Sep 17 00:00:00 2001 From: Junna Sugiyama Date: Thu, 23 May 2024 02:38:30 +0000 Subject: [PATCH 02/23] hwp_angle can be declared --- sotodlib/hwp/hwp.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/sotodlib/hwp/hwp.py b/sotodlib/hwp/hwp.py index facfba500..ea5a851ce 100644 --- a/sotodlib/hwp/hwp.py +++ b/sotodlib/hwp/hwp.py @@ -489,7 +489,7 @@ def subtract_hwpss(aman, signal=None, hwpss_template=None, signal, hwpss_template), [(0, 'dets'), (1, 'samps')]) -def demod_tod(aman, signal_name='signal', demod_mode=4, +def demod_tod(aman, signal_name='signal', hwp_angle=None, demod_mode=4, bpf_cfg=None, lpf_cfg=None): """ Demodulate TOD based on HWP angle @@ -524,8 +524,10 @@ def demod_tod(aman, signal_name='signal', demod_mode=4, the demodulated signal imaginary component filtered with `lpf` and multiplied by 2. """ + if hwp_angle is None: + hwp_angle = aman.hwp_angle # HWP speed in Hz - speed = (np.sum(np.abs(np.diff(np.unwrap(aman.hwp_angle)))) / + speed = (np.sum(np.abs(np.diff(np.unwrap(hwp_angle)))) / (aman.timestamps[-1] - aman.timestamps[0])) / (2 * np.pi) if bpf_cfg is None: @@ -544,7 +546,7 @@ def demod_tod(aman, signal_name='signal', demod_mode=4, 'trans_width': 0.1} lpf = filters.get_lpf(lpf_cfg) - phasor = np.exp(demod_mode * 1.j * aman.hwp_angle) + phasor = np.exp(demod_mode * 1.j * hwp_angle) demod = tod_ops.fourier_filter(aman, bpf, detrend=None, signal_name=signal_name) * phasor From 3fc19aec5a3a1b5f95d5ee64f63242c7887bc49c Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Mon, 3 Jun 2024 11:06:33 +0000 Subject: [PATCH 03/23] Modified not to use Numpy.MaskedArray --- sotodlib/tod_ops/fft_ops.py | 71 ++++++++++++++++--------------------- 1 file changed, 31 insertions(+), 40 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index 3db1d125c..a2d03d80e 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -346,8 +346,8 @@ def calc_masked_psd( hwpss=False, hwp_freq=None, f_max=100, - width_for_1f=(-0.4, +0.6), - width_for_Nf=(-0.2, +0.2), + width_for_1f=(-0.5, +0.7), + width_for_Nf=(-0.3, +0.3), peak=False, peak_freq=None, peak_width=(-0.002, +0.002), @@ -403,18 +403,13 @@ def calc_masked_psd( pxx_masked = [] if hwp_freq is None: hwp_freq = np.median(aman['hwp_solution']['raw_approx_hwp_freq_1'][1]) - for i in range(aman.dets.count): - masked_signal = mask_hwpss(f, pxx[i], hwp_freq, f_max=f_max, width_for_1f=width_for_1f, width_for_Nf=width_for_Nf) - pxx_masked.append(masked_signal.compressed()) - pxx = np.array(pxx_masked) - f = np.ma.masked_array(f, mask=masked_signal.mask).compressed() + mask_idx = hwpss_mask(f, hwp_freq, f_max=f_max, width_for_1f=width_for_1f, width_for_Nf=width_for_Nf) + f = f[mask_idx] + pxx = pxx[:, mask_idx] if peak: - pxx_masked = [] - for i in range(aman.dets.count): - masked_signal = mask_peak(f, pxx[i], peak_freq, peak_width=peak_width) - pxx_masked.append(masked_signal.compressed()) - pxx = np.array(pxx_masked) - f = np.ma.masked_array(f, mask=masked_signal.mask).compressed() + mask_idx = peak_mask(f, peak_freq, peak_width=peak_width) + f = f[mask_idx] + pxx = pxx[:, mask_idx] if merge: aman.merge( core.AxisManager(core.OffsetAxis("nusamps", len(f)))) if overwrite: @@ -612,18 +607,15 @@ def fit_noise_model( aman.wrap(merge_name, noise_fit_stats) return noise_fit_stats -def mask_hwpss(f, p, hwp_freq, f_max=100, width_for_1f=(-0.4, +0.6), width_for_Nf=(-0.2, +0.2)): +def hwpss_mask(f, hwp_freq, f_max=100, width_for_1f=(-0.4, +0.6), width_for_Nf=(-0.2, +0.2)): """ - Function that masks hwpss in PSD. + Function that returns boolean array to mask hwpss in PSD. Arguments --------- f : nparray Frequency of PSD of signal. - p : nparray - PSD of signal. - hwp_freq : float HWP frequency. @@ -640,28 +632,24 @@ def mask_hwpss(f, p, hwp_freq, f_max=100, width_for_1f=(-0.4, +0.6), width_for_N Returns ------- - p_mask: MaskedArray of the input PSD. p_mask.data returns the PSD before masking, - while p_mask.mask returns the bool array of the mask. + mask: Boolean array to mask frequency and power of the given PSD. + False in this array stands for the index of hwpss to mask. """ - mask_arrays = [((f > hwp_freq + width_for_1f[0]) & (f < hwp_freq + width_for_1f[1]))] + mask_arrays = [((f < hwp_freq + width_for_1f[0])|(f > hwp_freq + width_for_1f[1]))] for n in range(int(f_max//hwp_freq-1)): - mask_arrays.append(((f > hwp_freq*(n+2) + width_for_Nf[0]) & (f < hwp_freq*(n+2) + width_for_Nf[1]))) - mask = np.any(np.array(mask_arrays), axis=0) - p_mask = np.ma.masked_where(mask, p) - return p_mask + mask_arrays.append(((f < hwp_freq*(n+2) + width_for_Nf[0])|(f > hwp_freq*(n+2) + width_for_Nf[1]))) + mask = np.all(np.array(mask_arrays), axis=0) + return mask -def mask_peak(f, p, peak_freq, peak_width=(-0.002, +0.002)): +def peak_mask(f, peak_freq, peak_width=(-0.002, +0.002)): """ - Function that masks single peak (e.g. scan synchronous signal) in PSD. + Function that returns boolean array to masks single peak (e.g. scan synchronous signal) in PSD. Arguments --------- f : nparray Frequency of PSD of signal. - p : nparray - PSD of signal. - peak_freq : float Center frequency of the mask. @@ -671,12 +659,11 @@ def mask_peak(f, p, peak_freq, peak_width=(-0.002, +0.002)): Returns ------- - p_mask: MaskedArray of the input PSD. p_mask.data returns the PSD before masking, - while p_mask.mask returns the bool array of the mask. + mask: Boolean array to mask the given PSD. + False in this array stands for the index of the single peak to mask. """ - mask = (f > peak_freq + peak_width[0]) & (f < peak_freq + peak_width[1]) - p_mask = np.ma.masked_where(mask, p) - return p_mask + mask = (f < peak_freq + peak_width[0])|(f > peak_freq + peak_width[1]) + return mask def binning_psd(psd, unbinned_mode=10, base=2, drop_nan=False): """ @@ -704,15 +691,19 @@ def binning_psd(psd, unbinned_mode=10, base=2, drop_nan=False): Binned PSD. """ binned_psd = [] - for i in np.linspace(0, unbinned_mode, unbinned_mode+1, dtype = int): + for i in range(0, unbinned_mode+1): binned_psd.append(psd[i]) N = int(np.emath.logn(base, len(psd)-unbinned_mode)) binning_idx = np.logspace(base, N, N, base=base, dtype = int)+unbinned_mode-1 - for i in range(N-1): - if binning_idx[i] == binning_idx[i+1]: - continue - else: + if base >= 2: + for i in range(N-1): binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) + else: + for i in range(N-1): + if binning_idx[i] == binning_idx[i+1]: + continue + else: + binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) binned_psd = np.array(binned_psd) if drop_nan: binned_psd = binned_psd[~np.isnan(binned_psd)] From 307ab1610e22e5fd705afc064da10028585717e8 Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Mon, 8 Jul 2024 03:53:16 +0000 Subject: [PATCH 04/23] modified binning and fitting function --- sotodlib/tod_ops/fft_ops.py | 479 ++++++++++++++++++++++++++++++------ 1 file changed, 404 insertions(+), 75 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index a2d03d80e..fee889c95 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -4,6 +4,7 @@ import numpy as np import pyfftw import so3g +from so3g.proj import Ranges, RangesMatrix from scipy.optimize import minimize from scipy.signal import welch from sotodlib import core @@ -324,22 +325,61 @@ def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10): return wn -def noise_model(f, p): +def noise_model(f, params, **fixed_param): """ Noise model for power spectrum with white noise, and 1/f noise. + If any fixed param is handed, that parameter is fixed in the fit. + 'alpha' or 'wn' can be fixed. """ - fknee, w, alpha = p[0], p[1], p[2] - return w * (1 + (fknee / f) ** alpha) + if 'alpha' in fixed_param.keys(): + if len(params)==2: + alpha = fixed_param['alpha'] + fknee, wn = params[0], params[1] + else: + raise ValueError('The number of fit parameters are invalid.') + return + elif 'wn' in fixed_param.keys(): + if len(params)==2: + wn = fixed_param['wn'] + fknee, alpha = params[0], params[1] + else: + raise ValueError('The number of fit parameters are invalid.') + return + elif len(fixed_param)==0: + if len(params)==3: + fknee, wn, alpha = params[0], params[1], params[2] + else: + raise ValueError('The number of fit parameters are invalid.') + return + else: + raise ValueError('"alpha" or "wn" can be a fixed parameter.') + return + return wn * (1 + (fknee / f) ** alpha) -def neglnlike(params, x, y): - model = noise_model(x, params) +def neglnlike(params, x, y, **fixed_param): + """ + For unbinned PSD fitting + """ + model = noise_model(x, params, **fixed_param) output = np.sum(np.log(model) + y / model) if not np.isfinite(output): return 1.0e30 return output -def calc_masked_psd( +def leastsq(params, x, y, y_err, **fixed_param): + """ + For binned PSD fitting + """ + model = noise_model(x, params, **fixed_param) + output = np.sum((model-y)**2/y_err**2) + if not np.isfinite(output): + return 1.0e30 + return output + + + +def calc_psd_mask( aman, f=None, pxx=None, @@ -351,11 +391,12 @@ def calc_masked_psd( peak=False, peak_freq=None, peak_width=(-0.002, +0.002), - merge=False, - overwrite=True, + merge=True, + mode="add", + overwrite=True ): """ - Function that masks hwpss or single peak in PSD. + Function that calculates masks for hwpss or single peak in PSD. Arguments --------- @@ -387,46 +428,49 @@ def calc_masked_psd( (peak_freq - 0.002) < f < (peak_freq + 0.002). merge : bool if True merge results into axismanager. + mode: str + if "replace", existing PSD mask is replaced to new mask. + If "add", new mask range is added to the existing one. overwrite: bool - if true will overwrite f, pxx axes. You cannot overwrite with data - longer than before. + if true will overwrite aman.PSD_mask. Returns ------- - p_mask: MaskedArray of the input PSD. p_mask.data returns the PSD before masking, - while p_mask.mask returns the bool array of the mask. + PSD_mask (nusamps): Ranges array. If merge == True, "PSD_mask" is added to the aman. """ + PSD_mask = np.zeros(aman.nusamps.count, dtype=bool) if f is None or pxx is None: f = aman.freqs pxx = aman.Pxx if hwpss: - pxx_masked = [] if hwp_freq is None: hwp_freq = np.median(aman['hwp_solution']['raw_approx_hwp_freq_1'][1]) - mask_idx = hwpss_mask(f, hwp_freq, f_max=f_max, width_for_1f=width_for_1f, width_for_Nf=width_for_Nf) - f = f[mask_idx] - pxx = pxx[:, mask_idx] + PSD_mask = PSD_mask | hwpss_mask(f, hwp_freq, f_max=f_max, width_for_1f=width_for_1f, width_for_Nf=width_for_Nf) if peak: - mask_idx = peak_mask(f, peak_freq, peak_width=peak_width) - f = f[mask_idx] - pxx = pxx[:, mask_idx] + PSD_mask = PSD_mask | peak_mask(f, peak_freq, peak_width=peak_width) + if mode == "add": + if "PSD_mask" in aman: + PSD_mask = PSD_mask | aman.PSD_mask.mask() + elif mode == "replace": + pass + else: + print('Select mode from "add" or "replace".') + PSD_mask = Ranges.from_bitmask(PSD_mask) if merge: - aman.merge( core.AxisManager(core.OffsetAxis("nusamps", len(f)))) if overwrite: - if "freqs" in aman._fields: - aman.move("freqs", None) - if "Pxx" in aman._fields: - aman.move("Pxx", None) - aman.wrap("freqs", f, [(0,"nusamps")]) - aman.wrap("Pxx", pxx, [(0,"dets"),(1,"nusamps")]) - return f, pxx + if "PSD_mask" in aman: + aman.move("PSD_mask", None) + aman.wrap("PSD_mask", PSD_mask, [(0,"nusamps")]) + return PSD_mask def calc_binned_psd( aman, f=None, pxx=None, - unbinned_mode=10, - base=2, + mask=False, + unbinned_mode=3, + base=1.2, + limit_N=5, merge=False, overwrite=True, ): @@ -441,12 +485,16 @@ def calc_binned_psd( Frequency of PSD of signal. pxx : nparray PSD of signal. - hwpss : bool - If True, hwpss are masked. + mask : bool + if True calculate binned psd with mask. unbinned_mode : int First Fourier modes up to this number are left un-binned. base : float (> 1) Base of the logspace bins. + limit_N : int + If data points in a bin is less than limit_N, that bin is handled with + chi2 distribution and its error is . Otherwise the central limit theorem + is applied to the bin and the error of the bin is std(psd)/sqrt(len(psd)). merge : bool if True merge results into axismanager. overwrite: bool @@ -454,27 +502,41 @@ def calc_binned_psd( Returns ------- - p_mask: MaskedArray of the input PSD. p_mask.data returns the PSD before masking, - while p_mask.mask returns the bool array of the mask. + f_binned, pxx_binned, pxx_err: binned frequency and PSD. """ if f is None or pxx is None: f = aman.freqs pxx = aman.Pxx - f_binned = binning_psd(f, unbinned_mode=unbinned_mode, base=base) - pxx_binned = [] + if mask: + if 'PSD_mask' in aman: + mask = ~aman.PSD_mask.mask() + f = f[mask] + pxx = pxx[:, mask] + else: + print('"PSD_mask" is not in aman. Masking is skipped.') + + f_bin = binning_psd(f, unbinned_mode=unbinned_mode, base=base, limit_N=limit_N, err=False) + pxx_bin = [] + pxx_err = [] for i in range(aman.dets.count): - pxx_binned.append(binning_psd(pxx[i], unbinned_mode=unbinned_mode, base=base)) - pxx_binned = np.array(pxx_binned) + binned, err = binning_psd(pxx[i], unbinned_mode=unbinned_mode, base=base, limit_N=limit_N, err=True) + pxx_bin.append(binned) + pxx_err.append(err) + pxx_bin = np.array(pxx_bin) + pxx_err = np.array(pxx_err) if merge: - aman.merge( core.AxisManager(core.OffsetAxis("nusamps_bin", len(f_binned)))) + aman.merge(core.AxisManager(core.OffsetAxis("nusamps_bin", len(f_bin)))) if overwrite: if "freqs_bin" in aman._fields: aman.move("freqs_bin", None) if "Pxx_bin" in aman._fields: aman.move("Pxx_bin", None) - aman.wrap("freqs_bin", f_binned, [(0,"nusamps_bin")]) - aman.wrap("Pxx_bin", pxx_binned, [(0,"dets"),(1,"nusamps_bin")]) - return f_binned, pxx_binned + if "Pxx_bin_err" in aman._fields: + aman.move("Pxx_bin_err", None) + aman.wrap("freqs_bin", f_bin, [(0,"nusamps_bin")]) + aman.wrap("Pxx_bin", pxx_bin, [(0,"dets"),(1,"nusamps_bin")]) + aman.wrap("Pxx_bin_err", pxx_err, [(0,"dets"),(1,"nusamps_bin")]) + return f_bin, pxx_bin, pxx_err def fit_noise_model( aman, @@ -489,6 +551,8 @@ def fit_noise_model( f_max=100, merge_name="noise_fit_stats", merge_psd=True, + mask=False, + fixed_parameter=None, ): """ Fits noise model with white and 1/f noise to the PSD of signal. @@ -530,7 +594,200 @@ def fit_noise_model( merge_name : bool 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. + If ``merge_psd`` is True then adds fres and Pxx to the axis manager. + mask : bool + If ``mask`` is True then PSD is masked with ``aman.PSD_mask``. + fixed_parameter : str + This accepts None or 'wn' or 'alpha'. If 'wn' ('alpha') is given, + white noise level (alpha) is fixed to the initially estimated value. + Returns + ------- + noise_fit_stats : AxisManager + If merge_fit is False then axis manager with fit and fit statistics + is returned otherwise nothing is returned and axis manager is wrapped + into input aman. + """ + if signal is None: + 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, + ) + if mask: + if 'PSD_mask' in aman: + mask = ~aman.PSD_mask.mask() + f = f[mask] + pxx = pxx[:, mask] + else: + print('"PSD_mask" is not in aman. Masking is skipped.') + + eix = np.argmin(np.abs(f - f_max)) + if f_min is None: + six = 1 + else: + six = np.argmin(np.abs(f - f_min)) + f = f[six:eix] + pxx = pxx[:, six: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)) + fixed_param = {} + for i in range(aman.dets.count): + p = pxx[i] + wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))])*1.5 #conversion factor from median to mean + if fixed_parameter == None: + 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]] + elif fixed_parameter == 'wn': + 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], -pfit[0]] + fixed_param['wn'] = wnest + elif fixed_parameter == 'alpha': + pfit, vfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1, cov=True) + fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) + p0 = [f[fidx], wnest] + fixed_param['alpha'] = -pfit[0] + else: + print('"fixed_parameter" is invalid. Parameter fix is skipped.') + p0 = [f[fidx], wnest, -pfit[0]] + res = minimize(lambda params: neglnlike(params, f, p, **fixed_param), + p0, method="Nelder-Mead") + try: + #Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True) + Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p, **fixed_param), 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." + ) + covout_i = np.full((len(p0), len(p0)), np.nan) + except IndexError: + print( + f"Cannot fit 1/f curve for detector {aman.dets.vals[i]} skipping." + ) + covout_i = np.full((len(p0), len(p0)), np.nan) + fitout_i = res.x + if fixed_parameter == 'alpha': + alpha_err = np.sqrt(vfit[0][0]) + covout_i = np.insert(covout_i, 2, 0, axis=0) + covout_i = np.insert(covout_i, 2, 0, axis=1) + covout_i[2][2] = alpha_err + fitout_i = np.insert(fitout_i, 2, -pfit[0]) + elif fixed_parameter == 'wn': + wnest_err = np.std(p[((f > fwhite[0]) & (f < fwhite[1]))]) + covout_i = np.insert(covout_i, 1, 0, axis=0) + covout_i = np.insert(covout_i, 1, 0, axis=1) + covout_i[1][1] = wnest_err + fitout_i = np.insert(fitout_i, 1, wnest) + covout[i] = covout_i + fitout[i] = fitout_i + + noise_model_coeffs = ["fknee", "white_noise", "alpha"] + noise_fit_stats = core.AxisManager( + aman.dets, + core.LabelAxis( + name="noise_model_coeffs", vals=np.array(noise_model_coeffs, dtype=" 1) + Base of the logspace bins. + limit_N : int + If the number of data points in a bin is less than ``limit_N``, that bin is handled with + chi2 distribution and its error is . Otherwise the central limit theorem + is applied to the bin and the error of the bin is std(psd)/sqrt(len(psd)). + Returns ------- noise_fit_stats : AxisManager @@ -554,6 +811,14 @@ def fit_noise_model( merge=merge_psd, **psdargs, ) + if mask: + if 'PSD_mask' in aman: + mask = ~aman.PSD_mask.mask() + f = f[mask] + pxx = pxx[:, mask] + else: + print('"PSD_mask" is not in aman. Masking is skipped.') + eix = np.argmin(np.abs(f - f_max)) if f_min is None: six = 1 @@ -561,33 +826,68 @@ def fit_noise_model( six = np.argmin(np.abs(f - f_min)) f = f[six:eix] pxx = pxx[:, six:eix] + # binning + f, pxx, pxx_err = calc_binned_psd(aman, f=f, pxx=pxx, unbinned_mode=unbinned_mode, + base=base, limit_N=limit_N, 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)) + fixed_param = {} 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]] - res = minimize(neglnlike, p0, args=(f, p), method="Nelder-Mead") + p_err = pxx_err[i] + wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))])*1.5 #conversion factor from median to mean + if fixed_parameter == None: + 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]] + elif fixed_parameter == 'wn': + 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], -pfit[0]] + fixed_param['wn'] = wnest + elif fixed_parameter == 'alpha': + pfit, vfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1, cov=True) + fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) + p0 = [f[fidx], wnest] + fixed_param['alpha'] = -pfit[0] + else: + print('"fixed_parameter" is invalid. Parameter fix is skipped.') + p0 = [f[fidx], wnest, -pfit[0]] + res = minimize(lambda params: leastsq(params, f, p, p_err, **fixed_param), + p0, method="Nelder-Mead") try: - Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True) + #Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True) + Hfun = ndt.Hessian(lambda params: leastsq(params, f, p, p_err, **fixed_param), 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) + covout_i = np.linalg.inv(hessian_ndt) except np.linalg.LinAlgError: print( f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping." ) - covout[i] = np.full((3, 3), np.nan) + covout_i = np.full((len(p0), len(p0)), np.nan) except IndexError: print( f"Cannot fit 1/f curve for detector {aman.dets.vals[i]} skipping." ) - covout[i] = np.full((3, 3), np.nan) - fitout[i] = res.x + covout_i = np.full((len(p0), len(p0)), np.nan) + fitout_i = res.x + if fixed_parameter == 'alpha': + alpha_err = np.sqrt(vfit[0][0]) + covout_i = np.insert(covout_i, 2, 0, axis=0) + covout_i = np.insert(covout_i, 2, 0, axis=1) + covout_i[2][2] = alpha_err + fitout_i = np.insert(fitout_i, 2, -pfit[0]) + elif fixed_parameter == 'wn': + wnest_err = np.std(p[((f > fwhite[0]) & (f < fwhite[1]))]) + covout_i = np.insert(covout_i, 1, 0, axis=0) + covout_i = np.insert(covout_i, 1, 0, axis=1) + covout_i[1][1] = wnest_err + fitout_i = np.insert(fitout_i, 1, wnest) + covout[i] = covout_i + fitout[i] = fitout_i noise_model_coeffs = ["fknee", "white_noise", "alpha"] noise_fit_stats = core.AxisManager( @@ -633,12 +933,12 @@ def hwpss_mask(f, hwp_freq, f_max=100, width_for_1f=(-0.4, +0.6), width_for_Nf=( Returns ------- mask: Boolean array to mask frequency and power of the given PSD. - False in this array stands for the index of hwpss to mask. + True in this array stands for the index of hwpss to mask. """ - mask_arrays = [((f < hwp_freq + width_for_1f[0])|(f > hwp_freq + width_for_1f[1]))] + mask_arrays = [((f > hwp_freq + width_for_1f[0])&(f < hwp_freq + width_for_1f[1]))] for n in range(int(f_max//hwp_freq-1)): - mask_arrays.append(((f < hwp_freq*(n+2) + width_for_Nf[0])|(f > hwp_freq*(n+2) + width_for_Nf[1]))) - mask = np.all(np.array(mask_arrays), axis=0) + mask_arrays.append(((f > hwp_freq*(n+2) + width_for_Nf[0])&(f < hwp_freq*(n+2) + width_for_Nf[1]))) + mask = np.any(np.array(mask_arrays), axis=0) return mask def peak_mask(f, peak_freq, peak_width=(-0.002, +0.002)): @@ -660,12 +960,12 @@ def peak_mask(f, peak_freq, peak_width=(-0.002, +0.002)): Returns ------- mask: Boolean array to mask the given PSD. - False in this array stands for the index of the single peak to mask. + True in this array stands for the index of the single peak to mask. """ - mask = (f < peak_freq + peak_width[0])|(f > peak_freq + peak_width[1]) + mask = (f > peak_freq + peak_width[0])&(f < peak_freq + peak_width[1]) return mask -def binning_psd(psd, unbinned_mode=10, base=2, drop_nan=False): +def binning_psd(psd, unbinned_mode=3, base=1.2, limit_N=5, err=True, drop_nan=False): """ Function to bin PSD. First several Fourier modes are left un-binned. @@ -675,12 +975,20 @@ def binning_psd(psd, unbinned_mode=10, base=2, drop_nan=False): --------- psd : nparray PSD of signal. - + unbinned_mode : int First Fourier modes up to this number are left un-binned. base : float (> 1) - Base of the logspace bins. + Base of the logspace bins. + + limit_N : int + If data points in a bin is less than limit_N, that bin is handled with + chi2 distribution and its error is . Otherwise the central limit theorem + is applied to the bin and the error of the bin is std(psd)/sqrt(len(psd)). + + err : bool + If True, standard deviations are calculated for bins. drop_nan : bool If True, drop the index where p is nan. @@ -688,23 +996,44 @@ def binning_psd(psd, unbinned_mode=10, base=2, drop_nan=False): Returns ------- binned_psd: nparray - Binned PSD. + binned_psd_err: nparray """ binned_psd = [] - for i in range(0, unbinned_mode+1): - binned_psd.append(psd[i]) - N = int(np.emath.logn(base, len(psd)-unbinned_mode)) - binning_idx = np.logspace(base, N, N, base=base, dtype = int)+unbinned_mode-1 - if base >= 2: + if err == True: + binned_psd_err = [] + for i in range(0, unbinned_mode+1): + binned_psd.append(psd[i]) + binned_psd_err.append(psd[i]) + N = int(np.emath.logn(base, len(psd)-unbinned_mode)) + binning_idx = np.logspace(base, N, N, base=base, dtype = int)+unbinned_mode-1 for i in range(N-1): - binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) + # chi2 distributed bins + if limit_N > binning_idx[i+1] - binning_idx[i]: + binned_psd.append(psd[binning_idx[i]]) + binned_psd_err.append(psd[binning_idx[i]]) + else: + binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) + binned_psd_err.append(np.std(psd[binning_idx[i]:binning_idx[i+1]])/np.sqrt(binning_idx[i+1]-binning_idx[i]+1)) + + binned_psd = np.array(binned_psd) + binned_psd_err = np.array(binned_psd_err) + if drop_nan: + binned_psd = binned_psd[~np.isnan(binned_psd)] + binned_psd_err = binned_psd_err[~np.isnan(binned_psd_err)] + return binned_psd, binned_psd_err + else: + for i in range(0, unbinned_mode+1): + binned_psd.append(psd[i]) + N = int(np.emath.logn(base, len(psd)-unbinned_mode)) + binning_idx = np.logspace(base, N, N, base=base, dtype = int)+unbinned_mode-1 for i in range(N-1): - if binning_idx[i] == binning_idx[i+1]: - continue + if limit_N > binning_idx[i+1] - binning_idx[i]: + binned_psd.append(psd[binning_idx[i]]) else: binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) - binned_psd = np.array(binned_psd) - if drop_nan: - binned_psd = binned_psd[~np.isnan(binned_psd)] - return binned_psd \ No newline at end of file + + binned_psd = np.array(binned_psd) + if drop_nan: + binned_psd = binned_psd[~np.isnan(binned_psd)] + return binned_psd \ No newline at end of file From 3c96965d55357e753c4fcb3f9445217a46fca1fa Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Mon, 8 Jul 2024 09:35:06 +0000 Subject: [PATCH 05/23] Reflected Tomoki's comments --- sotodlib/tod_ops/fft_ops.py | 96 +++++++++++++++++++++++++------------ 1 file changed, 65 insertions(+), 31 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index fee889c95..dc6e58e6a 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -377,7 +377,10 @@ def leastsq(params, x, y, y_err, **fixed_param): return 1.0e30 return output - +def get_hwp_freq(timestamps=None, hwp_angle=None): + hwp_freq = (np.sum(np.abs(np.diff(np.unwrap(hwp_angle)))) / + (timestamps[-1] - timestamps[0])) / (2 * np.pi) + return hwp_freq def calc_psd_mask( aman, @@ -385,9 +388,8 @@ def calc_psd_mask( pxx=None, hwpss=False, hwp_freq=None, - f_max=100, - width_for_1f=(-0.5, +0.7), - width_for_Nf=(-0.3, +0.3), + max_hwpss_mode=10, + hwpss_width=((-0.4, 0.6),(-0.2, 0.2)), peak=False, peak_freq=None, peak_width=(-0.002, +0.002), @@ -410,15 +412,18 @@ def calc_psd_mask( If True, hwpss are masked. hwp_freq : float HWP frequency. - f_max : float - Maximum frequency to include in the hwpss masking. - width_for_1f : tuple - Range to mask 1f hwpss. The default masked range will be - (hwp_freq - 0.4) < f < (hwp_freq + 0.6). + max_hwpss_mode : int + Maximum hwpss mode to subtract. + hwpss_width : float/tuple/list/array + If given in float, hwpss will be masked like + (hwp_freq - width/2) < f < (hwp_freq + width/2). + If given in array like [1.0, 0.4], 1f hwpss is masked like + (hwp_freq - 0.5) < f < (hwp_freq + 0.5) and Nf hwpss are masked like + (hwp_freq - 0.2) < f < (hwp_freq + 0.2). + If given in array like [[-0.4, 0.6], [-0.2, 0.3]], + 1f hwpss is masked like (hwp_freq - 0.4) < f < (hwp_freq + 0.6) + and Nf are masked like (hwp_freq - 0.2) < f < (hwp_freq + 0.3). Usually 1f hwpss distrubites wider than other hwpss. - width_for_Nf : tuple - Range to mask Nf hwpss. The default masked range will be - (hwp_freq*N - 0.2) < f < (hwp_freq*N + 0.2). peak : bool If True, single peak is masked. peak_freq : float @@ -444,10 +449,10 @@ def calc_psd_mask( pxx = aman.Pxx if hwpss: if hwp_freq is None: - hwp_freq = np.median(aman['hwp_solution']['raw_approx_hwp_freq_1'][1]) - PSD_mask = PSD_mask | hwpss_mask(f, hwp_freq, f_max=f_max, width_for_1f=width_for_1f, width_for_Nf=width_for_Nf) + hwp_freq = get_hwp_freq(aman.timestamps, aman.hwp_solution.hwp_angle) + PSD_mask = PSD_mask | get_mask_for_hwpss(f, hwp_freq, max_mode=max_hwpss_mode, width=hwpss_width) if peak: - PSD_mask = PSD_mask | peak_mask(f, peak_freq, peak_width=peak_width) + PSD_mask = PSD_mask | get_mask_for_single_peak(f, peak_freq, peak_width=peak_width) if mode == "add": if "PSD_mask" in aman: PSD_mask = PSD_mask | aman.PSD_mask.mask() @@ -907,41 +912,70 @@ def fit_binned_noise_model( aman.wrap(merge_name, noise_fit_stats) return noise_fit_stats -def hwpss_mask(f, hwp_freq, f_max=100, width_for_1f=(-0.4, +0.6), width_for_Nf=(-0.2, +0.2)): +def get_mask_for_hwpss(freq, hwp_freq, max_mode=10, width=((-0.4, 0.6), (-0.2, 0.2))): """ Function that returns boolean array to mask hwpss in PSD. Arguments --------- - f : nparray + freq : nparray Frequency of PSD of signal. hwp_freq : float HWP frequency. - f_max : float - Maximum frequency to include in the fitting. - - width_for_1f : tuple - Range to mask 1f hwpss. The default masked range will be - (hwp_freq - 0.4) < f < (hwp_freq + 0.6). + max_mode : int + Maximum hwpss mode to subtract. + + width : float/tuple/list/array + If given in float, hwpss will be masked like + (hwp_freq - width/2) < f < (hwp_freq + width/2). + If given in array like [1.0, 0.4], 1f hwpss is masked like + (hwp_freq - 0.5) < f < (hwp_freq + 0.5) and Nf hwpss are masked like + (hwp_freq - 0.2) < f < (hwp_freq + 0.2). + If given in array like [[-0.4, 0.6], [-0.2, 0.3]], + 1f hwpss is masked like (hwp_freq - 0.4) < f < (hwp_freq + 0.6) + and Nf are masked like (hwp_freq - 0.2) < f < (hwp_freq + 0.3). Usually 1f hwpss distrubites wider than other hwpss. - width_for_Nf : tuple - Range to mask Nf hwpss. The default masked range will be - (hwp_freq*N - 0.2) < f < (hwp_freq*N + 0.2). Returns ------- mask: Boolean array to mask frequency and power of the given PSD. True in this array stands for the index of hwpss to mask. """ - mask_arrays = [((f > hwp_freq + width_for_1f[0])&(f < hwp_freq + width_for_1f[1]))] - for n in range(int(f_max//hwp_freq-1)): - mask_arrays.append(((f > hwp_freq*(n+2) + width_for_Nf[0])&(f < hwp_freq*(n+2) + width_for_Nf[1]))) + if isinstance(width, (float, int)): + width_minus = -width/2 + width_plus = width/2 + mask_arrays = [] + for n in range(max_mode): + mask_arrays.append(get_mask_for_single_peak(freq, hwp_freq*(n+1), peak_width=(width_minus, width_plus))) + elif isinstance(width, (np.ndarray, list, tuple)): + width = np.array(width) + if len(width.shape) == 1: + # mask for 1f + width_minus = -width[0]/2 + width_plus = width[0]/2 + mask_arrays = [get_mask_for_single_peak(freq, hwp_freq, peak_width=(width_minus, width_plus))] + # masks for Nf + width_minus = -width[1]/2 + width_plus = width[1]/2 + for n in range(max_mode-1): + mask_arrays.append(get_mask_for_single_peak(freq, hwp_freq*(n+2), peak_width=(width_minus, width_plus))) + elif len(width.shape) == 2: + # mask for 1f + width_minus = width[0][0] + width_plus = width[0][1] + mask_arrays = [get_mask_for_single_peak(freq, hwp_freq, peak_width=(width_minus, width_plus))] + # masks for Nf + width_minus = width[1][0] + width_plus = width[1][1] + for n in range(max_mode-1): + mask_arrays.append(get_mask_for_single_peak(freq, hwp_freq*(n+2), peak_width=(width_minus, width_plus))) mask = np.any(np.array(mask_arrays), axis=0) return mask - -def peak_mask(f, peak_freq, peak_width=(-0.002, +0.002)): + + +def get_mask_for_single_peak(f, peak_freq, peak_width=(-0.002, +0.002)): """ Function that returns boolean array to masks single peak (e.g. scan synchronous signal) in PSD. From 9575bf99651e2781c99e4ea5aaa8ca393bfafa39 Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Tue, 16 Jul 2024 11:22:25 +0000 Subject: [PATCH 06/23] Modified fit_binned_noise_model --- sotodlib/tod_ops/fft_ops.py | 99 +++++++++++++------------------------ 1 file changed, 33 insertions(+), 66 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index dc6e58e6a..5c94307fc 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -357,25 +357,12 @@ def noise_model(f, params, **fixed_param): return wn * (1 + (fknee / f) ** alpha) -def neglnlike(params, x, y, **fixed_param): - """ - For unbinned PSD fitting - """ +def neglnlike(params, x, y, bin_size=1, **fixed_param): model = noise_model(x, params, **fixed_param) - output = np.sum(np.log(model) + y / model) + output = np.sum((np.log(model) + y / model)*bin_size) if not np.isfinite(output): return 1.0e30 - return output - -def leastsq(params, x, y, y_err, **fixed_param): - """ - For binned PSD fitting - """ - model = noise_model(x, params, **fixed_param) - output = np.sum((model-y)**2/y_err**2) - if not np.isfinite(output): - return 1.0e30 - return output + return output def get_hwp_freq(timestamps=None, hwp_angle=None): hwp_freq = (np.sum(np.abs(np.diff(np.unwrap(hwp_angle)))) / @@ -475,7 +462,6 @@ def calc_binned_psd( mask=False, unbinned_mode=3, base=1.2, - limit_N=5, merge=False, overwrite=True, ): @@ -496,10 +482,6 @@ def calc_binned_psd( First Fourier modes up to this number are left un-binned. base : float (> 1) Base of the logspace bins. - limit_N : int - If data points in a bin is less than limit_N, that bin is handled with - chi2 distribution and its error is . Otherwise the central limit theorem - is applied to the bin and the error of the bin is std(psd)/sqrt(len(psd)). merge : bool if True merge results into axismanager. overwrite: bool @@ -507,7 +489,7 @@ def calc_binned_psd( Returns ------- - f_binned, pxx_binned, pxx_err: binned frequency and PSD. + f_binned, pxx_binned, bin_size: binned frequency and PSD. """ if f is None or pxx is None: f = aman.freqs @@ -520,15 +502,12 @@ def calc_binned_psd( else: print('"PSD_mask" is not in aman. Masking is skipped.') - f_bin = binning_psd(f, unbinned_mode=unbinned_mode, base=base, limit_N=limit_N, err=False) + f_bin, bin_size = binning_psd(f, unbinned_mode=unbinned_mode, base=base, return_bin_size=True) pxx_bin = [] - pxx_err = [] for i in range(aman.dets.count): - binned, err = binning_psd(pxx[i], unbinned_mode=unbinned_mode, base=base, limit_N=limit_N, err=True) + binned = binning_psd(pxx[i], unbinned_mode=unbinned_mode, base=base) pxx_bin.append(binned) - pxx_err.append(err) pxx_bin = np.array(pxx_bin) - pxx_err = np.array(pxx_err) if merge: aman.merge(core.AxisManager(core.OffsetAxis("nusamps_bin", len(f_bin)))) if overwrite: @@ -536,12 +515,12 @@ def calc_binned_psd( aman.move("freqs_bin", None) if "Pxx_bin" in aman._fields: aman.move("Pxx_bin", None) - if "Pxx_bin_err" in aman._fields: - aman.move("Pxx_bin_err", None) + if "bin_size" in aman._fields: + aman.move("bin_size", None) aman.wrap("freqs_bin", f_bin, [(0,"nusamps_bin")]) aman.wrap("Pxx_bin", pxx_bin, [(0,"dets"),(1,"nusamps_bin")]) - aman.wrap("Pxx_bin_err", pxx_err, [(0,"dets"),(1,"nusamps_bin")]) - return f_bin, pxx_bin, pxx_err + aman.wrap("bin_size", bin_size, [(0,"dets"),(1,"nusamps_bin")]) + return f_bin, pxx_bin, bin_size def fit_noise_model( aman, @@ -738,7 +717,6 @@ def fit_binned_noise_model( fixed_parameter=None, unbinned_mode=3, base=1.2, - limit_N=5, ): """ Fits noise model with white and 1/f noise to the PSD of binned signal. @@ -832,15 +810,14 @@ def fit_binned_noise_model( f = f[six:eix] pxx = pxx[:, six:eix] # binning - f, pxx, pxx_err = calc_binned_psd(aman, f=f, pxx=pxx, unbinned_mode=unbinned_mode, - base=base, limit_N=limit_N, merge=False) + f, pxx, bin_size = calc_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)) fixed_param = {} - for i in range(aman.dets.count): + for i in range(len(pxx)): p = pxx[i] - p_err = pxx_err[i] wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))])*1.5 #conversion factor from median to mean if fixed_parameter == None: pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) @@ -859,11 +836,11 @@ def fit_binned_noise_model( else: print('"fixed_parameter" is invalid. Parameter fix is skipped.') p0 = [f[fidx], wnest, -pfit[0]] - res = minimize(lambda params: leastsq(params, f, p, p_err, **fixed_param), + res = minimize(lambda params: neglnlike(params, f, p, bin_size=bin_size, **fixed_param), p0, method="Nelder-Mead") try: #Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True) - Hfun = ndt.Hessian(lambda params: leastsq(params, f, p, p_err, **fixed_param), full_output=True) + Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p, bin_size=bin_size, **fixed_param), 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. @@ -875,7 +852,7 @@ def fit_binned_noise_model( covout_i = np.full((len(p0), len(p0)), np.nan) except IndexError: print( - f"Cannot fit 1/f curve for detector {aman.dets.vals[i]} skipping." + f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping." ) covout_i = np.full((len(p0), len(p0)), np.nan) fitout_i = res.x @@ -999,7 +976,7 @@ def get_mask_for_single_peak(f, peak_freq, peak_width=(-0.002, +0.002)): mask = (f > peak_freq + peak_width[0])&(f < peak_freq + peak_width[1]) return mask -def binning_psd(psd, unbinned_mode=3, base=1.2, limit_N=5, err=True, drop_nan=False): +def binning_psd(psd, unbinned_mode=3, base=1.2, return_bin_size=False, drop_nan=False): """ Function to bin PSD. First several Fourier modes are left un-binned. @@ -1015,14 +992,9 @@ def binning_psd(psd, unbinned_mode=3, base=1.2, limit_N=5, err=True, drop_nan=Fa base : float (> 1) Base of the logspace bins. - - limit_N : int - If data points in a bin is less than limit_N, that bin is handled with - chi2 distribution and its error is . Otherwise the central limit theorem - is applied to the bin and the error of the bin is std(psd)/sqrt(len(psd)). - - err : bool - If True, standard deviations are calculated for bins. + + return_bin_size : bool + If True, the numbers of data points in the bins are returned. drop_nan : bool If True, drop the index where p is nan. @@ -1030,43 +1002,38 @@ def binning_psd(psd, unbinned_mode=3, base=1.2, limit_N=5, err=True, drop_nan=Fa Returns ------- binned_psd: nparray - binned_psd_err: nparray + bin_size: int """ binned_psd = [] - if err == True: - binned_psd_err = [] + if return_bin_size == True: + bin_size = [] for i in range(0, unbinned_mode+1): binned_psd.append(psd[i]) - binned_psd_err.append(psd[i]) + bin_size.append(1) N = int(np.emath.logn(base, len(psd)-unbinned_mode)) - binning_idx = np.logspace(base, N, N, base=base, dtype = int)+unbinned_mode-1 + binning_idx = np.logspace(base, N, N, base=base, dtype=int)+unbinned_mode-1 for i in range(N-1): - # chi2 distributed bins - if limit_N > binning_idx[i+1] - binning_idx[i]: - binned_psd.append(psd[binning_idx[i]]) - binned_psd_err.append(psd[binning_idx[i]]) + if binning_idx[i+1] == binning_idx[i]: + continue else: binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) - binned_psd_err.append(np.std(psd[binning_idx[i]:binning_idx[i+1]])/np.sqrt(binning_idx[i+1]-binning_idx[i]+1)) - + bin_size.append(binning_idx[i+1] - binning_idx[i]) binned_psd = np.array(binned_psd) - binned_psd_err = np.array(binned_psd_err) + bin_size = np.array(bin_size) if drop_nan: binned_psd = binned_psd[~np.isnan(binned_psd)] - binned_psd_err = binned_psd_err[~np.isnan(binned_psd_err)] - return binned_psd, binned_psd_err - + bin_size = bin_size[~np.isnan(bin_size)] + return binned_psd, bin_size else: for i in range(0, unbinned_mode+1): binned_psd.append(psd[i]) N = int(np.emath.logn(base, len(psd)-unbinned_mode)) binning_idx = np.logspace(base, N, N, base=base, dtype = int)+unbinned_mode-1 for i in range(N-1): - if limit_N > binning_idx[i+1] - binning_idx[i]: - binned_psd.append(psd[binning_idx[i]]) + if binning_idx[i+1] == binning_idx[i]: + continue else: binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) - binned_psd = np.array(binned_psd) if drop_nan: binned_psd = binned_psd[~np.isnan(binned_psd)] From 81fbde8f792f0a1d185c3c66732299ced3fb0d79 Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Wed, 17 Jul 2024 06:39:28 +0000 Subject: [PATCH 07/23] Minor debug --- sotodlib/tod_ops/fft_ops.py | 78 ++++++++++++++++++++++++++++++------- 1 file changed, 64 insertions(+), 14 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index 5c94307fc..fa86895ce 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -618,6 +618,8 @@ def fit_noise_model( eix = np.argmin(np.abs(f - f_max)) if f_min is None: six = 1 + elif f_min < f[0]: + six = 1 else: six = np.argmin(np.abs(f - f_min)) f = f[six:eix] @@ -630,22 +632,46 @@ def fit_noise_model( p = pxx[i] wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))])*1.5 #conversion factor from median to mean if fixed_parameter == None: - pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) + try: + pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) + except np.linalg.LinAlgError: + print( + f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." + ) + covout[i] = np.full((3, 3), np.nan) + fitout[i] = np.full(3, np.nan) + continue fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) p0 = [f[fidx], wnest, -pfit[0]] elif fixed_parameter == 'wn': - pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) + try: + pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) + except np.linalg.LinAlgError: + print( + f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." + ) + covout[i] = np.full((3, 3), np.nan) + fitout[i] = np.full(3, np.nan) + continue fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) p0 = [f[fidx], -pfit[0]] fixed_param['wn'] = wnest elif fixed_parameter == 'alpha': - pfit, vfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1, cov=True) + try: + pfit, vfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1, cov=True) + except np.linalg.LinAlgError: + print( + f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." + ) + covout[i] = np.full((3, 3), np.nan) + fitout[i] = np.full(3, np.nan) + continue fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) p0 = [f[fidx], wnest] fixed_param['alpha'] = -pfit[0] else: - print('"fixed_parameter" is invalid. Parameter fix is skipped.') - p0 = [f[fidx], wnest, -pfit[0]] + print('"fixed_parameter" is invalid.') + return res = minimize(lambda params: neglnlike(params, f, p, **fixed_param), p0, method="Nelder-Mead") try: @@ -657,12 +683,12 @@ def fit_noise_model( covout_i = np.linalg.inv(hessian_ndt) except np.linalg.LinAlgError: print( - f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping." + f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping. (LinAlgError)" ) covout_i = np.full((len(p0), len(p0)), np.nan) except IndexError: print( - f"Cannot fit 1/f curve for detector {aman.dets.vals[i]} skipping." + f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping. (IndexError)" ) covout_i = np.full((len(p0), len(p0)), np.nan) fitout_i = res.x @@ -820,22 +846,46 @@ def fit_binned_noise_model( p = pxx[i] wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))])*1.5 #conversion factor from median to mean if fixed_parameter == None: - pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) + try: + pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) + except np.linalg.LinAlgError: + print( + f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." + ) + covout[i] = np.full((3, 3), np.nan) + fitout[i] = np.full(3, np.nan) + continue fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) p0 = [f[fidx], wnest, -pfit[0]] elif fixed_parameter == 'wn': - pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) + try: + pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) + except np.linalg.LinAlgError: + print( + f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." + ) + covout[i] = np.full((3, 3), np.nan) + fitout[i] = np.full(3, np.nan) + continue fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) p0 = [f[fidx], -pfit[0]] fixed_param['wn'] = wnest elif fixed_parameter == 'alpha': - pfit, vfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1, cov=True) + try: + pfit, vfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1, cov=True) + except np.linalg.LinAlgError: + print( + f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." + ) + covout[i] = np.full((3, 3), np.nan) + fitout[i] = np.full(3, np.nan) + continue fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) p0 = [f[fidx], wnest] fixed_param['alpha'] = -pfit[0] else: - print('"fixed_parameter" is invalid. Parameter fix is skipped.') - p0 = [f[fidx], wnest, -pfit[0]] + print('"fixed_parameter" is invalid.') + return res = minimize(lambda params: neglnlike(params, f, p, bin_size=bin_size, **fixed_param), p0, method="Nelder-Mead") try: @@ -847,12 +897,12 @@ def fit_binned_noise_model( covout_i = np.linalg.inv(hessian_ndt) except np.linalg.LinAlgError: print( - f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping." + f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping. (LinAlgError)" ) covout_i = np.full((len(p0), len(p0)), np.nan) except IndexError: print( - f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping." + f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping. (IndexError)" ) covout_i = np.full((len(p0), len(p0)), np.nan) fitout_i = res.x From 8116be2b07d2349c6c8f376062e714e08807c519 Mon Sep 17 00:00:00 2001 From: tterasaki Date: Thu, 1 Aug 2024 02:33:39 +0000 Subject: [PATCH 08/23] moved get_hwp_freq to sotodlib.hwp.hwp --- sotodlib/hwp/hwp.py | 17 +++++++++++++++-- sotodlib/tod_ops/fft_ops.py | 11 +++-------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/sotodlib/hwp/hwp.py b/sotodlib/hwp/hwp.py index 596e22620..8b27672fc 100644 --- a/sotodlib/hwp/hwp.py +++ b/sotodlib/hwp/hwp.py @@ -491,6 +491,20 @@ def subtract_hwpss(aman, signal=None, hwpss_template=None, aman.wrap(subtract_name, np.subtract( signal, hwpss_template), [(0, 'dets'), (1, 'samps')]) +def get_hwp_freq(timestamps, hwp_angle): + """ + Calculate the frequency of HWP rotation. + + Parameters: + timestamps (array-like): An array of timestamps. + hwp_angle (array-like): An array of HWP angles in radian + + Returns: + float: The frequency of the HWP rotation in Hz. + """ + hwp_freq = (np.sum(np.abs(np.diff(np.unwrap(hwp_angle)))) / + (timestamps[-1] - timestamps[0])) / (2 * np.pi) + return hwp_freq def demod_tod(aman, signal_name='signal', hwp_angle=None, demod_mode=4, bpf_cfg=None, lpf_cfg=None): @@ -530,8 +544,7 @@ def demod_tod(aman, signal_name='signal', hwp_angle=None, demod_mode=4, if hwp_angle is None: hwp_angle = aman.hwp_angle # HWP speed in Hz - speed = (np.sum(np.abs(np.diff(np.unwrap(hwp_angle)))) / - (aman.timestamps[-1] - aman.timestamps[0])) / (2 * np.pi) + speed = get_hwp_freq(timestamps=aman.timestamps, hwp_angle=hwp_angle): if bpf_cfg is None: bpf_center = demod_mode * speed diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index fa86895ce..7f8bcbeaa 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -8,6 +8,7 @@ from scipy.optimize import minimize from scipy.signal import welch from sotodlib import core +from sotodlib.hwp.hwp import get_hwp_freq from . import detrend_tod @@ -362,12 +363,7 @@ def neglnlike(params, x, y, bin_size=1, **fixed_param): output = np.sum((np.log(model) + y / model)*bin_size) if not np.isfinite(output): return 1.0e30 - return output - -def get_hwp_freq(timestamps=None, hwp_angle=None): - hwp_freq = (np.sum(np.abs(np.diff(np.unwrap(hwp_angle)))) / - (timestamps[-1] - timestamps[0])) / (2 * np.pi) - return hwp_freq + return output def calc_psd_mask( aman, @@ -435,8 +431,7 @@ def calc_psd_mask( f = aman.freqs pxx = aman.Pxx if hwpss: - if hwp_freq is None: - hwp_freq = get_hwp_freq(aman.timestamps, aman.hwp_solution.hwp_angle) + hwp_freq = get_hwp_freq(aman.timestamps, aman.hwp_solution.hwp_angle) PSD_mask = PSD_mask | get_mask_for_hwpss(f, hwp_freq, max_mode=max_hwpss_mode, width=hwpss_width) if peak: PSD_mask = PSD_mask | get_mask_for_single_peak(f, peak_freq, peak_width=peak_width) From f5673c83f8091eb7173361882d29e53aa054c586 Mon Sep 17 00:00:00 2001 From: tterasaki Date: Thu, 1 Aug 2024 02:53:28 +0000 Subject: [PATCH 09/23] small bug fix --- sotodlib/hwp/__init__.py | 2 +- sotodlib/hwp/hwp.py | 6 +++--- sotodlib/tod_ops/fft_ops.py | 8 +++----- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/sotodlib/hwp/__init__.py b/sotodlib/hwp/__init__.py index d78991ee9..7128b11f0 100644 --- a/sotodlib/hwp/__init__.py +++ b/sotodlib/hwp/__init__.py @@ -6,7 +6,7 @@ """ from .g3thwp import G3tHWP -from .hwp import (get_hwpss, subtract_hwpss, demod_tod) +from .hwp import (get_hwpss, subtract_hwpss, demod_tod, get_hwp_freq) from .sim_hwp import I_to_P_param from .sim_hwp import sim_hwpss from .sim_hwp import sim_hwpss_2f4f diff --git a/sotodlib/hwp/hwp.py b/sotodlib/hwp/hwp.py index 8b27672fc..77c12f1df 100644 --- a/sotodlib/hwp/hwp.py +++ b/sotodlib/hwp/hwp.py @@ -1,7 +1,7 @@ import numpy as np from scipy.optimize import curve_fit from sotodlib import core, tod_ops -from sotodlib.tod_ops import bin_signal, filters +from sotodlib.tod_ops import filters import logging logger = logging.getLogger(__name__) @@ -210,7 +210,7 @@ def get_binned_hwpss(aman, signal=None, hwp_angle=None, if hwp_angle is None: hwp_angle = aman['hwp_angle'] - binning_dict = bin_signal(aman, bin_by=hwp_angle, range=[0, 2*np.pi], + binning_dict = tod_ops.bin_signal(aman, bin_by=hwp_angle, range=[0, 2*np.pi], bins=bins, signal=signal, flags=flags) bin_centers = binning_dict['bin_centers'] @@ -544,7 +544,7 @@ def demod_tod(aman, signal_name='signal', hwp_angle=None, demod_mode=4, if hwp_angle is None: hwp_angle = aman.hwp_angle # HWP speed in Hz - speed = get_hwp_freq(timestamps=aman.timestamps, hwp_angle=hwp_angle): + speed = get_hwp_freq(timestamps=aman.timestamps, hwp_angle=hwp_angle) if bpf_cfg is None: bpf_center = demod_mode * speed diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index 7f8bcbeaa..fc33af882 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -7,10 +7,8 @@ from so3g.proj import Ranges, RangesMatrix from scipy.optimize import minimize from scipy.signal import welch -from sotodlib import core -from sotodlib.hwp.hwp import get_hwp_freq - -from . import detrend_tod +from sotodlib import core, hwp +from sotodlib.tod_ops import detrend_tod def _get_num_threads(): @@ -431,7 +429,7 @@ def calc_psd_mask( f = aman.freqs pxx = aman.Pxx if hwpss: - hwp_freq = get_hwp_freq(aman.timestamps, aman.hwp_solution.hwp_angle) + hwp_freq = hwp.get_hwp_freq(aman.timestamps, aman.hwp_solution.hwp_angle) PSD_mask = PSD_mask | get_mask_for_hwpss(f, hwp_freq, max_mode=max_hwpss_mode, width=hwpss_width) if peak: PSD_mask = PSD_mask | get_mask_for_single_peak(f, peak_freq, peak_width=peak_width) From 912ac2e4216657028b14d8c71af54f4090bb9429 Mon Sep 17 00:00:00 2001 From: tterasaki Date: Thu, 1 Aug 2024 03:14:20 +0000 Subject: [PATCH 10/23] improved calc_psd merging style --- sotodlib/tod_ops/fft_ops.py | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index fc33af882..7ffcf56b8 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -211,6 +211,7 @@ def calc_psd( prefer='center', freq_spacing=None, merge=False, + merge_suffix=None, overwrite=True, **kwargs ): @@ -231,7 +232,8 @@ def calc_psd( If None the default nperseg of ~1/50th the signal length is used. 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. + merge_suffix (str, optional): Suffix to append to the Pxx field name in aman. Defaults to None (merged as Pxx). + overwrite (bool): if true will overwrite f, Pxx axes. **kwargs: keyword args to be passed to signal.welch(). Returns: @@ -269,15 +271,23 @@ def calc_psd( freqs, Pxx = welch(signal[:, start:stop], fs, **kwargs) if merge: + if 'nusamps' not in list(aman._axes.keys()): + aman.merge(core.AxisManager(core.OffsetAxis("nusamps", len(freqs)))) + aman.wrap("freqs", freqs, [(0,"nusamps")]) + else: + if len(freqs) != aman.nusamps.count: + raise ValueError('New freqs does not match the shape of nusamps\ + To avoid this, use the same value for nperseg') + + if merge_suffix is None: + Pxx_name = 'Pxx' + else: + Pxx_name = f'Pxx_{merge_suffix}' + if overwrite: - - if "freqs" in aman._fields: - aman.move("freqs", None) - if "Pxx" in aman._fields: - aman.move("Pxx", None) - aman.merge( core.AxisManager(core.OffsetAxis("nusamps", len(freqs)))) - aman.wrap("freqs", freqs, [(0,"nusamps")]) - aman.wrap("Pxx", Pxx, [(0,"dets"),(1,"nusamps")]) + if Pxx_name in aman._fields: + aman.move(Pxx_name, None) + aman.wrap(Pxx_name, Pxx, [(0,"dets"), (1,"nusamps")]) return freqs, Pxx From 97d212731ed18af71b2e1fd24837379cf98ab235 Mon Sep 17 00:00:00 2001 From: tterasaki Date: Thu, 1 Aug 2024 04:53:00 +0000 Subject: [PATCH 11/23] from capital PSD_mask to psd_mask --- sotodlib/tod_ops/fft_ops.py | 111 ++++++++++++++++-------------------- 1 file changed, 50 insertions(+), 61 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index 7ffcf56b8..c0d0e8f17 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -373,20 +373,10 @@ def neglnlike(params, x, y, bin_size=1, **fixed_param): return 1.0e30 return output -def calc_psd_mask( - aman, - f=None, - pxx=None, - hwpss=False, - hwp_freq=None, - max_hwpss_mode=10, - hwpss_width=((-0.4, 0.6),(-0.2, 0.2)), - peak=False, - peak_freq=None, - peak_width=(-0.002, +0.002), - merge=True, - mode="add", - overwrite=True +def calc_psd_mask(aman, psd_mask=None, f=None, pxx=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), + merge=True, overwrite=True ): """ Function that calculates masks for hwpss or single peak in PSD. @@ -394,28 +384,29 @@ def calc_psd_mask( Arguments --------- aman : AxisManager - Axis manager which has samps axis aligned with signal. + Axis manager which has 'nusamps' axis. + psd_mask : numpy.ndarray or Ranges + Existing psd_mask to be updated. If None, a new mask is created. f : nparray - Frequency of PSD of signal. + Frequency of PSD of signal. If None, aman.freqs are used. pxx : nparray - PSD of signal. - hwpss : bool - If True, hwpss are masked. + PSD of signal. If None, aman.Pxx are used. + mask_hwpss : bool + If True, hwpss are masked. Defaults to True. hwp_freq : float - HWP frequency. + HWP frequency. If None, calculated based on aman.hwp_angle max_hwpss_mode : int Maximum hwpss mode to subtract. - hwpss_width : float/tuple/list/array - If given in float, hwpss will be masked like - (hwp_freq - width/2) < f < (hwp_freq + width/2). - If given in array like [1.0, 0.4], 1f hwpss is masked like - (hwp_freq - 0.5) < f < (hwp_freq + 0.5) and Nf hwpss are masked like - (hwp_freq - 0.2) < f < (hwp_freq + 0.2). + hwpss_width : array-like + If given in float, + nf-hwpss will be masked like (n * hwp_freq - width/2) < f < (n * hwp_freq + width/2). + If given in array like [1.0, 0.4], + 1f hwpss is masked like (hwp_freq - 0.5) < f < (hwp_freq + 0.5) and + nf hwpss are masked like (n * hwp_freq - 0.2) < f < (n * hwp_freq + 0.2). If given in array like [[-0.4, 0.6], [-0.2, 0.3]], - 1f hwpss is masked like (hwp_freq - 0.4) < f < (hwp_freq + 0.6) - and Nf are masked like (hwp_freq - 0.2) < f < (hwp_freq + 0.3). - Usually 1f hwpss distrubites wider than other hwpss. - peak : bool + 1f hwpss is masked like (hwp_freq - 0.4) < f < (hwp_freq + 0.6) and + nf are masked like (n * hwp_freq - 0.2) < f < (n * hwp_freq + 0.3). + mask_peak : bool If True, single peak is masked. peak_freq : float Center frequency of the mask. @@ -428,35 +419,33 @@ def calc_psd_mask( if "replace", existing PSD mask is replaced to new mask. If "add", new mask range is added to the existing one. overwrite: bool - if true will overwrite aman.PSD_mask. + if true will overwrite aman.psd_mask. Returns ------- - PSD_mask (nusamps): Ranges array. If merge == True, "PSD_mask" is added to the aman. + psd_mask (nusamps): Ranges array. If merge == True, "psd_mask" is added to the aman. """ - PSD_mask = np.zeros(aman.nusamps.count, dtype=bool) - if f is None or pxx is None: + if psd_mask is None: + psd_mask = np.zeros(aman.nusamps.count, dtype=bool) + elif isinstance(psd_mask, so3g.RangesInt32): + psd_mask = psd_mask.mask() + if f is None: f = aman.freqs + if pxx is None: pxx = aman.Pxx - if hwpss: + if mask_hwpss: hwp_freq = hwp.get_hwp_freq(aman.timestamps, aman.hwp_solution.hwp_angle) - PSD_mask = PSD_mask | get_mask_for_hwpss(f, hwp_freq, max_mode=max_hwpss_mode, width=hwpss_width) - if peak: - PSD_mask = PSD_mask | get_mask_for_single_peak(f, peak_freq, peak_width=peak_width) - if mode == "add": - if "PSD_mask" in aman: - PSD_mask = PSD_mask | aman.PSD_mask.mask() - elif mode == "replace": - pass - else: - print('Select mode from "add" or "replace".') - PSD_mask = Ranges.from_bitmask(PSD_mask) + psd_mask = psd_mask | get_mask_for_hwpss(f, hwp_freq, max_mode=max_hwpss_mode, width=hwpss_width) + if mask_peak: + psd_mask = psd_mask | get_mask_for_single_peak(f, peak_freq, peak_width=peak_width) + + psd_mask = Ranges.from_bitmask(psd_mask) if merge: if overwrite: - if "PSD_mask" in aman: - aman.move("PSD_mask", None) - aman.wrap("PSD_mask", PSD_mask, [(0,"nusamps")]) - return PSD_mask + if "psd_mask" in aman: + aman.move("psd_mask", None) + aman.wrap("psd_mask", psd_mask, [(0,"nusamps")]) + return psd_mask def calc_binned_psd( aman, @@ -498,12 +487,12 @@ def calc_binned_psd( f = aman.freqs pxx = aman.Pxx if mask: - if 'PSD_mask' in aman: - mask = ~aman.PSD_mask.mask() + if 'psd_mask' in aman: + mask = ~aman.psd_mask.mask() f = f[mask] pxx = pxx[:, mask] else: - print('"PSD_mask" is not in aman. Masking is skipped.') + print('"psd_mask" is not in aman. Masking is skipped.') f_bin, bin_size = binning_psd(f, unbinned_mode=unbinned_mode, base=base, return_bin_size=True) pxx_bin = [] @@ -583,7 +572,7 @@ def fit_noise_model( merge_psd : bool If ``merge_psd`` is True then adds fres and Pxx to the axis manager. mask : bool - If ``mask`` is True then PSD is masked with ``aman.PSD_mask``. + If ``mask`` is True then PSD is masked with ``aman.psd_mask``. fixed_parameter : str This accepts None or 'wn' or 'alpha'. If 'wn' ('alpha') is given, white noise level (alpha) is fixed to the initially estimated value. @@ -611,12 +600,12 @@ def fit_noise_model( **psdargs, ) if mask: - if 'PSD_mask' in aman: - mask = ~aman.PSD_mask.mask() + if 'psd_mask' in aman: + mask = ~aman.psd_mask.mask() f = f[mask] pxx = pxx[:, mask] else: - print('"PSD_mask" is not in aman. Masking is skipped.') + print('"psd_mask" is not in aman. Masking is skipped.') eix = np.argmin(np.abs(f - f_max)) if f_min is None: @@ -787,7 +776,7 @@ def fit_binned_noise_model( merge_psd : bool If ``merge_psd`` is True then adds fres and Pxx to the axis manager. mask : bool - If ``mask`` is True then PSD is masked with ``aman.PSD_mask``. + If ``mask`` is True then PSD is masked with ``aman.psd_mask``. fixed_parameter : str This accepts None or 'wn' or 'alpha'. If 'wn' ('alpha') is given, white noise level (alpha) is fixed to the initially estimated value. @@ -824,12 +813,12 @@ def fit_binned_noise_model( **psdargs, ) if mask: - if 'PSD_mask' in aman: - mask = ~aman.PSD_mask.mask() + if 'psd_mask' in aman: + mask = ~aman.psd_mask.mask() f = f[mask] pxx = pxx[:, mask] else: - print('"PSD_mask" is not in aman. Masking is skipped.') + print('"psd_mask" is not in aman. Masking is skipped.') eix = np.argmin(np.abs(f - f_max)) if f_min is None: From 63b9608d0e44612359db9a95a3d6b14aa6309d31 Mon Sep 17 00:00:00 2001 From: tterasaki Date: Thu, 1 Aug 2024 05:30:04 +0000 Subject: [PATCH 12/23] binning_psd can accept 2d array --- sotodlib/tod_ops/fft_ops.py | 119 ++++++++++++++++++------------------ 1 file changed, 59 insertions(+), 60 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index c0d0e8f17..ce334fe44 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -451,7 +451,6 @@ def calc_binned_psd( aman, f=None, pxx=None, - mask=False, unbinned_mode=3, base=1.2, merge=False, @@ -483,16 +482,10 @@ def calc_binned_psd( ------- f_binned, pxx_binned, bin_size: binned frequency and PSD. """ - if f is None or pxx is None: + if f is None: f = aman.freqs + if pxx is None: pxx = aman.Pxx - if mask: - if 'psd_mask' in aman: - mask = ~aman.psd_mask.mask() - f = f[mask] - pxx = pxx[:, mask] - else: - print('"psd_mask" is not in aman. Masking is skipped.') f_bin, bin_size = binning_psd(f, unbinned_mode=unbinned_mode, base=base, return_bin_size=True) pxx_bin = [] @@ -1024,59 +1017,65 @@ def binning_psd(psd, unbinned_mode=3, base=1.2, return_bin_size=False, drop_nan= First several Fourier modes are left un-binned. Fourier modes higher than that are averaged into logspace bins. - Arguments - --------- - psd : nparray - PSD of signal. - - unbinned_mode : int - First Fourier modes up to this number are left un-binned. - - base : float (> 1) - Base of the logspace bins. - - return_bin_size : bool - If True, the numbers of data points in the bins are returned. - - drop_nan : bool - If True, drop the index where p is nan. + Parameters + ---------- + psd : numpy.ndarray + PSD of the signal. Can be a 1D or 2D array. + unbinned_mode : int, optional + First Fourier modes up to this number are left un-binned. Defaults to 3. + base : float, optional + Base of the logspace bins. Must be greater than 1. Defaults to 1.2. + return_bin_size : bool, optional + If True, the number of data points in the bins are returned. Defaults to False. + drop_nan : bool, optional + If True, drop the indices where psd is NaN. Defaults to False. Returns ------- - binned_psd: nparray - bin_size: int + binned_psd : numpy.ndarray + The binned PSD. If the input is 2D, the output will also be 2D with the same number of rows. + bin_size : numpy.ndarray, optional + The number of data points in each bin, only returned if return_bin_size is True. """ - binned_psd = [] - if return_bin_size == True: - bin_size = [] - for i in range(0, unbinned_mode+1): - binned_psd.append(psd[i]) - bin_size.append(1) - N = int(np.emath.logn(base, len(psd)-unbinned_mode)) - binning_idx = np.logspace(base, N, N, base=base, dtype=int)+unbinned_mode-1 - for i in range(N-1): - if binning_idx[i+1] == binning_idx[i]: - continue - else: - binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) - bin_size.append(binning_idx[i+1] - binning_idx[i]) - binned_psd = np.array(binned_psd) - bin_size = np.array(bin_size) - if drop_nan: - binned_psd = binned_psd[~np.isnan(binned_psd)] - bin_size = bin_size[~np.isnan(bin_size)] + if base <= 1: + raise ValueError("base must be greater than 1") + + # Ensure psd is at least 2D for consistent processing + psd = np.atleast_2d(psd) + num_signals, num_samples = psd.shape + + # Initialize the binned PSD and optionally the bin sizes + binned_psd = np.zeros((num_signals, unbinned_mode + 1)) + binned_psd[:, :unbinned_mode + 1] = psd[:, :unbinned_mode + 1] + bin_size = np.ones((num_signals, unbinned_mode + 1)) if return_bin_size else None + + # Determine the number of bins and their indices + N = int(np.ceil(np.emath.logn(base, num_samples - unbinned_mode))) + binning_idx = np.unique(np.logspace(base, N, N, base=base, dtype=int) + unbinned_mode - 1) + + # Bin the PSD values for each signal + new_binned_psd = [] + new_bin_size = [] + for start, end in zip(binning_idx[:-1], binning_idx[1:]): + if start < end: # Ensure the slice is not empty + bin_mean = np.nanmean(psd[:, start:end], axis=1) + new_binned_psd.append(bin_mean) + if return_bin_size: + new_bin_size.append(end - start) + + # Convert lists to numpy arrays and concatenate with initial values + new_binned_psd = np.array(new_binned_psd).T # Transpose to match dimensions + binned_psd = np.hstack([binned_psd, new_binned_psd]) + if return_bin_size: + new_bin_size = np.array(new_bin_size) + bin_size = np.hstack([bin_size, np.tile(new_bin_size, (num_signals, 1))]) + + if drop_nan: + valid_indices = ~np.isnan(binned_psd).any(axis=0) + binned_psd = binned_psd[:, valid_indices] + if return_bin_size: + bin_size = bin_size[:, valid_indices] + + if return_bin_size: return binned_psd, bin_size - else: - for i in range(0, unbinned_mode+1): - binned_psd.append(psd[i]) - N = int(np.emath.logn(base, len(psd)-unbinned_mode)) - binning_idx = np.logspace(base, N, N, base=base, dtype = int)+unbinned_mode-1 - for i in range(N-1): - if binning_idx[i+1] == binning_idx[i]: - continue - else: - binned_psd.append(np.mean(psd[binning_idx[i]:binning_idx[i+1]])) - binned_psd = np.array(binned_psd) - if drop_nan: - binned_psd = binned_psd[~np.isnan(binned_psd)] - return binned_psd \ No newline at end of file + return binned_psd From c94607190dcb961f268227fe47b23b37eec3b3db Mon Sep 17 00:00:00 2001 From: tterasaki Date: Thu, 1 Aug 2024 06:28:23 +0000 Subject: [PATCH 13/23] binning_psd can accept 2d array --- sotodlib/tod_ops/fft_ops.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index ce334fe44..ef102001c 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -488,11 +488,8 @@ def calc_binned_psd( pxx = aman.Pxx f_bin, bin_size = binning_psd(f, unbinned_mode=unbinned_mode, base=base, return_bin_size=True) - pxx_bin = [] - for i in range(aman.dets.count): - binned = binning_psd(pxx[i], unbinned_mode=unbinned_mode, base=base) - pxx_bin.append(binned) - pxx_bin = np.array(pxx_bin) + pxx_bin = binning_psd(pxx, unbinned_mode=unbinned_mode, base=base, return_bin_size=False) + if merge: aman.merge(core.AxisManager(core.OffsetAxis("nusamps_bin", len(f_bin)))) if overwrite: From 65f4125df087899bdcd7c35dce21a2490260283c Mon Sep 17 00:00:00 2001 From: tterasaki Date: Thu, 1 Aug 2024 10:45:19 +0000 Subject: [PATCH 14/23] log_binning modification --- sotodlib/tod_ops/fft_ops.py | 45 +++++++++++++++++++++++++------------ 1 file changed, 31 insertions(+), 14 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index ef102001c..36c044c71 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -452,7 +452,7 @@ def calc_binned_psd( f=None, pxx=None, unbinned_mode=3, - base=1.2, + base=1.05, merge=False, overwrite=True, ): @@ -487,8 +487,8 @@ def calc_binned_psd( if pxx is None: pxx = aman.Pxx - f_bin, bin_size = binning_psd(f, unbinned_mode=unbinned_mode, base=base, return_bin_size=True) - pxx_bin = binning_psd(pxx, unbinned_mode=unbinned_mode, base=base, return_bin_size=False) + f_bin, bin_size = log_binning(f, unbinned_mode=unbinned_mode, base=base, return_bin_size=True) + pxx_bin = log_binning(pxx, unbinned_mode=unbinned_mode, base=base, return_bin_size=False) if merge: aman.merge(core.AxisManager(core.OffsetAxis("nusamps_bin", len(f_bin)))) @@ -724,7 +724,7 @@ def fit_binned_noise_model( mask=False, fixed_parameter=None, unbinned_mode=3, - base=1.2, + base=1.05, ): """ Fits noise model with white and 1/f noise to the PSD of binned signal. @@ -1008,20 +1008,23 @@ def get_mask_for_single_peak(f, peak_freq, peak_width=(-0.002, +0.002)): mask = (f > peak_freq + peak_width[0])&(f < peak_freq + peak_width[1]) return mask -def binning_psd(psd, unbinned_mode=3, base=1.2, return_bin_size=False, drop_nan=False): +def log_binning(psd, unbinned_mode=3, base=1.05, mask=None, + return_bin_size=False, drop_nan=False): """ - Function to bin PSD. - First several Fourier modes are left un-binned. + Function to bin PSD or frequency. First several Fourier modes are left un-binned. Fourier modes higher than that are averaged into logspace bins. Parameters ---------- psd : numpy.ndarray - PSD of the signal. Can be a 1D or 2D array. + PSD (or frequency) to be binned. Can be a 1D or 2D array. unbinned_mode : int, optional First Fourier modes up to this number are left un-binned. Defaults to 3. base : float, optional - Base of the logspace bins. Must be greater than 1. Defaults to 1.2. + Base of the logspace bins. Must be greater than 1. Defaults to 1.05. + mask : numpy.ndarray, optional + Mask for psd. If all values in a bin are masked, the value becomes np.nan. + Should be a 1D array. return_bin_size : bool, optional If True, the number of data points in the bins are returned. Defaults to False. drop_nan : bool, optional @@ -1037,10 +1040,19 @@ def binning_psd(psd, unbinned_mode=3, base=1.2, return_bin_size=False, drop_nan= if base <= 1: raise ValueError("base must be greater than 1") + is_1d = psd.ndim == 1 + # Ensure psd is at least 2D for consistent processing psd = np.atleast_2d(psd) num_signals, num_samples = psd.shape + if mask is not None: + # Ensure mask is at least 2D and has the same shape as psd + mask = np.atleast_2d(mask) + if mask.shape[1] != num_samples: + raise ValueError("Mask must have the same number of columns as psd") + psd = np.ma.masked_array(psd, mask=np.tile(mask, (num_signals, 1))) + # Initialize the binned PSD and optionally the bin sizes binned_psd = np.zeros((num_signals, unbinned_mode + 1)) binned_psd[:, :unbinned_mode + 1] = psd[:, :unbinned_mode + 1] @@ -1054,11 +1066,10 @@ def binning_psd(psd, unbinned_mode=3, base=1.2, return_bin_size=False, drop_nan= new_binned_psd = [] new_bin_size = [] for start, end in zip(binning_idx[:-1], binning_idx[1:]): - if start < end: # Ensure the slice is not empty - bin_mean = np.nanmean(psd[:, start:end], axis=1) - new_binned_psd.append(bin_mean) - if return_bin_size: - new_bin_size.append(end - start) + bin_mean = np.nanmean(psd[:, start:end], axis=1) + new_binned_psd.append(bin_mean) + if return_bin_size: + new_bin_size.append(end - start) # Convert lists to numpy arrays and concatenate with initial values new_binned_psd = np.array(new_binned_psd).T # Transpose to match dimensions @@ -1073,6 +1084,12 @@ def binning_psd(psd, unbinned_mode=3, base=1.2, return_bin_size=False, drop_nan= if return_bin_size: bin_size = bin_size[:, valid_indices] + if is_1d: + binned_psd = binned_psd.flatten() + if return_bin_size: + bin_size = bin_size.flatten() + if return_bin_size: return binned_psd, bin_size return binned_psd + From f138b700ac192ba78b983e23d5157e0edea9ee7e Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Thu, 8 Aug 2024 08:18:49 +0000 Subject: [PATCH 15/23] integrated fit_noise_model and fit_binned_noise_model --- sotodlib/tod_ops/fft_ops.py | 218 ++---------------------------------- 1 file changed, 7 insertions(+), 211 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index 36c044c71..ac09d6c32 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -504,211 +504,8 @@ def calc_binned_psd( aman.wrap("bin_size", bin_size, [(0,"dets"),(1,"nusamps_bin")]) return f_bin, pxx_bin, bin_size -def fit_noise_model( - aman, - signal=None, - f=None, - pxx=None, - psdargs=None, - fwhite=(10, 100), - lowf=1, - merge_fit=False, - f_min=None, - f_max=100, - merge_name="noise_fit_stats", - merge_psd=True, - mask=False, - fixed_parameter=None, -): - """ - Fits noise model with white and 1/f noise to the PSD of signal. - This uses a MLE method that minimizes a log likelihood. This is - better for chi^2 distributed data like the PSD. - - Reference: http://keatonb.github.io/archivers/powerspectrumfits - Args - ---- - aman : AxisManager - Axis manager which has samps axis aligned with signal. - signal : nparray - Signal sized ndets x nsamps to fit noise model to. - Default is None which corresponds to aman.signal. - f : nparray - Frequency of PSD of signal. - Default is None which calculates f, pxx from signal. - pxx : nparray - PSD sized ndets x len(f) which is fit to with model. - Default is None which calculates f, pxx from signal. - psdargs : dict - Dictionary of optional argument for ``scipy.signal.welch`` - fwhite : tuple - Low and high frequency used to estimate white noise for initial - guess passed to ``scipy.signal.curve_fit``. - lowf : tuple - Frequency below which estimate of 1/f noise index and knee are estimated - for initial guess passed to ``scipy.signal.curve_fit``. - merge_fit : bool - Merges fit and fit statistics into input axis manager. - f_min : float - Minimum frequency to include in the fitting. - Default is None which selects f_min as the second index of f. - f_max : float - Maximum frequency to include in the fitting. This is particularly - important for lowpass filtered data such as that post demodulation - if the data is not downsampled after lowpass filtering. - merge_name : bool - If ``merge_fit`` is True then addes into axis manager with merge_name. - merge_psd : bool - If ``merge_psd`` is True then adds fres and Pxx to the axis manager. - mask : bool - If ``mask`` is True then PSD is masked with ``aman.psd_mask``. - fixed_parameter : str - This accepts None or 'wn' or 'alpha'. If 'wn' ('alpha') is given, - white noise level (alpha) is fixed to the initially estimated value. - Returns - ------- - noise_fit_stats : AxisManager - If merge_fit is False then axis manager with fit and fit statistics - is returned otherwise nothing is returned and axis manager is wrapped - into input aman. - """ - if signal is None: - 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, - ) - if mask: - if 'psd_mask' in aman: - mask = ~aman.psd_mask.mask() - f = f[mask] - pxx = pxx[:, mask] - else: - print('"psd_mask" is not in aman. Masking is skipped.') - - eix = np.argmin(np.abs(f - f_max)) - if f_min is None: - six = 1 - elif f_min < f[0]: - six = 1 - else: - six = np.argmin(np.abs(f - f_min)) - f = f[six:eix] - pxx = pxx[:, six: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)) - fixed_param = {} - for i in range(aman.dets.count): - p = pxx[i] - wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))])*1.5 #conversion factor from median to mean - if fixed_parameter == None: - try: - pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) - except np.linalg.LinAlgError: - print( - f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." - ) - covout[i] = np.full((3, 3), np.nan) - fitout[i] = np.full(3, np.nan) - continue - fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) - p0 = [f[fidx], wnest, -pfit[0]] - elif fixed_parameter == 'wn': - try: - pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) - except np.linalg.LinAlgError: - print( - f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." - ) - covout[i] = np.full((3, 3), np.nan) - fitout[i] = np.full(3, np.nan) - continue - fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) - p0 = [f[fidx], -pfit[0]] - fixed_param['wn'] = wnest - elif fixed_parameter == 'alpha': - try: - pfit, vfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1, cov=True) - except np.linalg.LinAlgError: - print( - f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." - ) - covout[i] = np.full((3, 3), np.nan) - fitout[i] = np.full(3, np.nan) - continue - fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) - p0 = [f[fidx], wnest] - fixed_param['alpha'] = -pfit[0] - else: - print('"fixed_parameter" is invalid.') - return - res = minimize(lambda params: neglnlike(params, f, p, **fixed_param), - p0, method="Nelder-Mead") - try: - #Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True) - Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p, **fixed_param), 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: - print( - f"Cannot calculate Hessian for detector {aman.dets.vals[i]} skipping. (IndexError)" - ) - covout_i = np.full((len(p0), len(p0)), np.nan) - fitout_i = res.x - if fixed_parameter == 'alpha': - alpha_err = np.sqrt(vfit[0][0]) - covout_i = np.insert(covout_i, 2, 0, axis=0) - covout_i = np.insert(covout_i, 2, 0, axis=1) - covout_i[2][2] = alpha_err - fitout_i = np.insert(fitout_i, 2, -pfit[0]) - elif fixed_parameter == 'wn': - wnest_err = np.std(p[((f > fwhite[0]) & (f < fwhite[1]))]) - covout_i = np.insert(covout_i, 1, 0, axis=0) - covout_i = np.insert(covout_i, 1, 0, axis=1) - covout_i[1][1] = wnest_err - fitout_i = np.insert(fitout_i, 1, wnest) - covout[i] = covout_i - fitout[i] = fitout_i - - noise_model_coeffs = ["fknee", "white_noise", "alpha"] - noise_fit_stats = core.AxisManager( - aman.dets, - core.LabelAxis( - name="noise_model_coeffs", vals=np.array(noise_model_coeffs, dtype=" 1) Base of the logspace bins. - limit_N : int - If the number of data points in a bin is less than ``limit_N``, that bin is handled with - chi2 distribution and its error is . Otherwise the central limit theorem - is applied to the bin and the error of the bin is std(psd)/sqrt(len(psd)). - + Returns ------- noise_fit_stats : AxisManager @@ -817,9 +611,11 @@ def fit_binned_noise_model( six = np.argmin(np.abs(f - f_min)) f = f[six:eix] pxx = pxx[:, six:eix] + bin_size = 1 # binning - f, pxx, bin_size = calc_binned_psd(aman, f=f, pxx=pxx, unbinned_mode=unbinned_mode, - base=base, merge=False) + if binning == True: + f, pxx, bin_size = calc_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)) From 15a83a92ea4600937700d28f40902eac63a1ea8b Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Thu, 8 Aug 2024 09:51:47 +0000 Subject: [PATCH 16/23] fit_noise_model can have a constant fixed parameter's value --- sotodlib/tod_ops/fft_ops.py | 54 +++++++++++++++---------------------- 1 file changed, 21 insertions(+), 33 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index ac09d6c32..cce9c6c35 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -519,7 +519,7 @@ def fit_noise_model( merge_name="noise_fit_stats", merge_psd=True, mask=False, - fixed_parameter=None, + fixed_parameter={}, binning=False, unbinned_mode=3, base=1.05, @@ -565,9 +565,9 @@ def fit_noise_model( If ``merge_psd`` is True then adds fres and Pxx to the axis manager. mask : bool If ``mask`` is True then PSD is masked with ``aman.psd_mask``. - fixed_parameter : str - This accepts None or 'wn' or 'alpha'. If 'wn' ('alpha') is given, - white noise level (alpha) is fixed to the initially estimated value. + fixed_parameter : dict + This accepts the key in 'wn' or 'alpha'. If 'wn' ('alpha') is given, + white noise level (alpha) is fixed to the given value. unbinned_mode : int First Fourier modes up to this number are left un-binned. base : float (> 1) @@ -622,8 +622,8 @@ def fit_noise_model( fixed_param = {} for i in range(len(pxx)): p = pxx[i] - wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))])*1.5 #conversion factor from median to mean - if fixed_parameter == None: + if fixed_parameter == {}: + wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))]) try: pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) except np.linalg.LinAlgError: @@ -635,7 +635,8 @@ def fit_noise_model( continue fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) p0 = [f[fidx], wnest, -pfit[0]] - elif fixed_parameter == 'wn': + elif 'wn' in fixed_parameter.keys(): + wnest = fixed_parameter["wn"] try: pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) except np.linalg.LinAlgError: @@ -647,28 +648,17 @@ def fit_noise_model( continue fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) p0 = [f[fidx], -pfit[0]] - fixed_param['wn'] = wnest - elif fixed_parameter == 'alpha': - try: - pfit, vfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1, cov=True) - except np.linalg.LinAlgError: - print( - f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." - ) - covout[i] = np.full((3, 3), np.nan) - fitout[i] = np.full(3, np.nan) - continue - fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) - p0 = [f[fidx], wnest] - fixed_param['alpha'] = -pfit[0] + elif 'alpha' in fixed_parameter.keys(): + alpha = fixed_parameter['alpha'] + wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))]) + p0 = [lowf, wnest] else: print('"fixed_parameter" is invalid.') return - res = minimize(lambda params: neglnlike(params, f, p, bin_size=bin_size, **fixed_param), + res = minimize(lambda params: neglnlike(params, f, p, bin_size=bin_size, **fixed_parameter), p0, method="Nelder-Mead") try: - #Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True) - Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p, bin_size=bin_size, **fixed_param), full_output=True) + Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p, bin_size=bin_size, **fixed_parameter), 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. @@ -684,18 +674,16 @@ def fit_noise_model( ) covout_i = np.full((len(p0), len(p0)), np.nan) fitout_i = res.x - if fixed_parameter == 'alpha': - alpha_err = np.sqrt(vfit[0][0]) - covout_i = np.insert(covout_i, 2, 0, axis=0) - covout_i = np.insert(covout_i, 2, 0, axis=1) - covout_i[2][2] = alpha_err - fitout_i = np.insert(fitout_i, 2, -pfit[0]) - elif fixed_parameter == 'wn': - wnest_err = np.std(p[((f > fwhite[0]) & (f < fwhite[1]))]) + if 'wn' in fixed_parameter.keys(): covout_i = np.insert(covout_i, 1, 0, axis=0) covout_i = np.insert(covout_i, 1, 0, axis=1) - covout_i[1][1] = wnest_err + covout_i[1][1] = wnest fitout_i = np.insert(fitout_i, 1, wnest) + elif 'alpha' in fixed_parameter.keys(): + covout_i = np.insert(covout_i, 2, 0, axis=0) + covout_i = np.insert(covout_i, 2, 0, axis=1) + covout_i[2][2] = alpha + fitout_i = np.insert(fitout_i, 2, alpha) covout[i] = covout_i fitout[i] = fitout_i From d3eaabd9a9abdca7447748558d383794a649db4c Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Thu, 8 Aug 2024 10:06:52 +0000 Subject: [PATCH 17/23] put nan in covout where the parameter is fixed --- sotodlib/tod_ops/fft_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index cce9c6c35..b17c343fd 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -677,12 +677,12 @@ def fit_noise_model( if 'wn' in fixed_parameter.keys(): covout_i = np.insert(covout_i, 1, 0, axis=0) covout_i = np.insert(covout_i, 1, 0, axis=1) - covout_i[1][1] = wnest + covout_i[1][1] = np.nan fitout_i = np.insert(fitout_i, 1, wnest) elif 'alpha' in fixed_parameter.keys(): covout_i = np.insert(covout_i, 2, 0, axis=0) covout_i = np.insert(covout_i, 2, 0, axis=1) - covout_i[2][2] = alpha + covout_i[2][2] = np.nan fitout_i = np.insert(fitout_i, 2, alpha) covout[i] = covout_i fitout[i] = fitout_i From b57326e990b352a0bc615982e2738dbfd1858fff Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Thu, 8 Aug 2024 10:42:52 +0000 Subject: [PATCH 18/23] changed some function's name --- sotodlib/tod_ops/fft_ops.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index b17c343fd..4c6df4f76 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -373,13 +373,13 @@ def neglnlike(params, x, y, bin_size=1, **fixed_param): return 1.0e30 return output -def calc_psd_mask(aman, psd_mask=None, f=None, pxx=None, +def get_psd_mask(aman, psd_mask=None, f=None, pxx=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), merge=True, overwrite=True ): """ - Function that calculates masks for hwpss or single peak in PSD. + Function to get masks for hwpss or single peak in PSD. Arguments --------- @@ -447,7 +447,7 @@ def calc_psd_mask(aman, psd_mask=None, f=None, pxx=None, aman.wrap("psd_mask", psd_mask, [(0,"nusamps")]) return psd_mask -def calc_binned_psd( +def get_binned_psd( aman, f=None, pxx=None, @@ -614,7 +614,7 @@ def fit_noise_model( bin_size = 1 # binning if binning == True: - f, pxx, bin_size = calc_binned_psd(aman, f=f, pxx=pxx, unbinned_mode=unbinned_mode, + 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 From 57f5c244dad385599a4b2467bdf1caeb786f905c Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Thu, 8 Aug 2024 11:00:07 +0000 Subject: [PATCH 19/23] minor bug fix --- sotodlib/tod_ops/fft_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index 4c6df4f76..75cc5ae86 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -691,7 +691,7 @@ def fit_noise_model( noise_fit_stats = core.AxisManager( aman.dets, core.LabelAxis( - name="noise_model_coeffs", vals=np.array(noise_model_coeffs, dtype=" Date: Thu, 8 Aug 2024 12:47:22 +0000 Subject: [PATCH 20/23] give initial estimates of parameters for fit_noise_model --- sotodlib/tod_ops/fft_ops.py | 104 ++++++++++++++++++------------------ 1 file changed, 51 insertions(+), 53 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index 75cc5ae86..d81aab127 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -511,15 +511,16 @@ def fit_noise_model( f=None, pxx=None, psdargs=None, - fwhite=(10, 100), - lowf=1, + fknee_est=1, + wn_est=2E-9, + alpha_est=3.4, merge_fit=False, f_min=None, f_max=100, merge_name="noise_fit_stats", merge_psd=True, mask=False, - fixed_parameter={}, + fixed_param=None, binning=False, unbinned_mode=3, base=1.05, @@ -544,12 +545,15 @@ def fit_noise_model( Default is None which calculates f, pxx from signal. psdargs : dict Dictionary of optional argument for ``scipy.signal.welch`` - fwhite : tuple - Low and high frequency used to estimate white noise for initial - guess passed to ``scipy.signal.curve_fit``. - lowf : tuple - Frequency below which estimate of 1/f noise index and knee are estimated - for initial guess passed to ``scipy.signal.curve_fit``. + fknee_est : float or nparray + Initial estimation value of knee frequency. If it is given in array, + initial values are applied for each detector. + wn_est : float or nparray + Initial estimation value of white noise. If it is given in array, + initial values are applied for each detector. + alpha_est : float or nparray + Initial estimation value of alpha. If it is given in array, + initial values are applied for each detector. merge_fit : bool Merges fit and fit statistics into input axis manager. f_min : float @@ -565,9 +569,9 @@ def fit_noise_model( If ``merge_psd`` is True then adds fres and Pxx to the axis manager. mask : bool If ``mask`` is True then PSD is masked with ``aman.psd_mask``. - fixed_parameter : dict - This accepts the key in 'wn' or 'alpha'. If 'wn' ('alpha') is given, - white noise level (alpha) is fixed to the given value. + fixed_param : str + This accepts 'wn' or 'alpha' or None. If 'wn' ('alpha') is given, + white noise level (alpha) is fixed to the wn_est (alpha_est). unbinned_mode : int First Fourier modes up to this number are left un-binned. base : float (> 1) @@ -615,50 +619,44 @@ def fit_noise_model( # binning if binning == True: f, pxx, bin_size = get_binned_psd(aman, f=f, pxx=pxx, unbinned_mode=unbinned_mode, - base=base, merge=False) + 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)) - fixed_param = {} + if isinstance(fknee_est, (int, float)): + fknee_est = np.zeros(aman.dets.count)+fknee_est + elif len(fknee_est)!=aman.dets.count: + print('Size of fknee_est must be equal to aman.dets.count or a single value.') + return + if isinstance(wn_est, (int, float)): + wn_est = np.zeros(aman.dets.count)+wn_est + elif len(wn_est)!=aman.dets.count: + print('Size of wn_est must be equal to aman.dets.count or a single value.') + return + if isinstance(alpha_est, (int, float)): + alpha_est = np.zeros(aman.dets.count)+alpha_est + elif len(alpha_est)!=aman.dets.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([fknee_est, wn_est, alpha_est]) + if fixed_param == "wn": + initial_params = np.array([fknee_est, alpha_est]) + fixed = wn_est + if fixed_param == "alpha": + initial_params = np.array([fknee_est, wn_est]) + fixed = alpha_est + for i in range(len(pxx)): p = pxx[i] - if fixed_parameter == {}: - wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))]) - try: - pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) - except np.linalg.LinAlgError: - print( - f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." - ) - covout[i] = np.full((3, 3), np.nan) - fitout[i] = np.full(3, np.nan) - continue - fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) - p0 = [f[fidx], wnest, -pfit[0]] - elif 'wn' in fixed_parameter.keys(): - wnest = fixed_parameter["wn"] - try: - pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1) - except np.linalg.LinAlgError: - print( - f"Cannot fit 1/f for detector {aman.dets.vals[i]} skipping." - ) - covout[i] = np.full((3, 3), np.nan) - fitout[i] = np.full(3, np.nan) - continue - fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest)) - p0 = [f[fidx], -pfit[0]] - elif 'alpha' in fixed_parameter.keys(): - alpha = fixed_parameter['alpha'] - wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))]) - p0 = [lowf, wnest] - else: - print('"fixed_parameter" is invalid.') - return - res = minimize(lambda params: neglnlike(params, f, p, bin_size=bin_size, **fixed_parameter), + p0 = initial_params.T[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, method="Nelder-Mead") try: - Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p, bin_size=bin_size, **fixed_parameter), full_output=True) + 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. @@ -674,16 +672,16 @@ def fit_noise_model( ) covout_i = np.full((len(p0), len(p0)), np.nan) fitout_i = res.x - if 'wn' in fixed_parameter.keys(): + if fixed_param == "wn": covout_i = np.insert(covout_i, 1, 0, axis=0) covout_i = np.insert(covout_i, 1, 0, axis=1) covout_i[1][1] = np.nan - fitout_i = np.insert(fitout_i, 1, wnest) - elif 'alpha' in fixed_parameter.keys(): + fitout_i = np.insert(fitout_i, 1, wn_est[i]) + elif fixed_param == "alpha": covout_i = np.insert(covout_i, 2, 0, axis=0) covout_i = np.insert(covout_i, 2, 0, axis=1) covout_i[2][2] = np.nan - fitout_i = np.insert(fitout_i, 2, alpha) + fitout_i = np.insert(fitout_i, 2, alpha_est[i]) covout[i] = covout_i fitout[i] = fitout_i From 91c43812089bbf198602da03eb754594ef6bd40e Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Thu, 12 Sep 2024 10:52:05 +0000 Subject: [PATCH 21/23] reflected Tomoki's comments --- sotodlib/tod_ops/fft_ops.py | 47 +++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 23 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index d81aab127..351b80d75 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -339,31 +339,32 @@ def noise_model(f, params, **fixed_param): Noise model for power spectrum with white noise, and 1/f noise. If any fixed param is handed, that parameter is fixed in the fit. 'alpha' or 'wn' can be fixed. + params = [wn, fknee, alpha] """ - if 'alpha' in fixed_param.keys(): + if 'wn' in fixed_param.keys(): if len(params)==2: - alpha = fixed_param['alpha'] - fknee, wn = params[0], params[1] + wn = fixed_param['wn'] + fknee, alpha = params[0], params[1] else: raise ValueError('The number of fit parameters are invalid.') return - elif 'wn' in fixed_param.keys(): + elif 'alpha' in fixed_param.keys(): if len(params)==2: - wn = fixed_param['wn'] - fknee, alpha = params[0], params[1] + alpha = fixed_param['alpha'] + wn, fknee = params[0], params[1] else: raise ValueError('The number of fit parameters are invalid.') return elif len(fixed_param)==0: if len(params)==3: - fknee, wn, alpha = params[0], params[1], params[2] + wn, fknee, alpha = params[0], params[1], params[2] else: raise ValueError('The number of fit parameters are invalid.') return else: raise ValueError('"alpha" or "wn" can be a fixed parameter.') return - return wn * (1 + (fknee / f) ** alpha) + return wn**2 * (1 + (fknee / f) ** alpha) def neglnlike(params, x, y, bin_size=1, **fixed_param): @@ -512,7 +513,7 @@ def fit_noise_model( pxx=None, psdargs=None, fknee_est=1, - wn_est=2E-9, + wn_est=4E-5, alpha_est=3.4, merge_fit=False, f_min=None, @@ -623,28 +624,28 @@ def fit_noise_model( 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)) - if isinstance(fknee_est, (int, float)): - fknee_est = np.zeros(aman.dets.count)+fknee_est - elif len(fknee_est)!=aman.dets.count: - print('Size of fknee_est must be equal to aman.dets.count or a single value.') - return if isinstance(wn_est, (int, float)): - wn_est = np.zeros(aman.dets.count)+wn_est + wn_est = np.full(aman.dets.count, wn_est) elif len(wn_est)!=aman.dets.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: + 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.zeros(aman.dets.count)+alpha_est + alpha_est = np.full(aman.dets.count, alpha_est) elif len(alpha_est)!=aman.dets.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([fknee_est, wn_est, alpha_est]) + initial_params = np.array([wn_est, fknee_est, alpha_est]) if fixed_param == "wn": initial_params = np.array([fknee_est, alpha_est]) fixed = wn_est if fixed_param == "alpha": - initial_params = np.array([fknee_est, wn_est]) + initial_params = np.array([wn_est, fknee_est]) fixed = alpha_est for i in range(len(pxx)): @@ -673,10 +674,10 @@ def fit_noise_model( covout_i = np.full((len(p0), len(p0)), np.nan) fitout_i = res.x if fixed_param == "wn": - covout_i = np.insert(covout_i, 1, 0, axis=0) - covout_i = np.insert(covout_i, 1, 0, axis=1) - covout_i[1][1] = np.nan - fitout_i = np.insert(fitout_i, 1, wn_est[i]) + covout_i = np.insert(covout_i, 0, 0, axis=0) + covout_i = np.insert(covout_i, 0, 0, axis=1) + covout_i[0][0] = np.nan + fitout_i = np.insert(fitout_i, 0, wn_est[i]) elif fixed_param == "alpha": covout_i = np.insert(covout_i, 2, 0, axis=0) covout_i = np.insert(covout_i, 2, 0, axis=1) @@ -685,7 +686,7 @@ def fit_noise_model( covout[i] = covout_i fitout[i] = fitout_i - noise_model_coeffs = ["fknee", "white_noise", "alpha"] + noise_model_coeffs = ["white_noise", "fknee", "alpha"] noise_fit_stats = core.AxisManager( aman.dets, core.LabelAxis( From f10e0f0083b36950299490ce3d8e6c745d8e3b5c Mon Sep 17 00:00:00 2001 From: 17-sugiyama Date: Tue, 24 Sep 2024 12:04:30 +0000 Subject: [PATCH 22/23] add unittest --- tests/test_tod_ops.py | 50 +++++++++++++++++++++++++++++++++++++------ 1 file changed, 44 insertions(+), 6 deletions(-) diff --git a/tests/test_tod_ops.py b/tests/test_tod_ops.py index c2bd93754..af8870b9c 100644 --- a/tests/test_tod_ops.py +++ b/tests/test_tod_ops.py @@ -9,7 +9,9 @@ import numpy as np import pylab as pl import scipy.signal - +from numpy.fft import rfftfreq,irfft +from numpy.fft import fftfreq,rfft +from scipy import interpolate from numpy.testing import assert_array_equal, assert_allclose from sotodlib import core, tod_ops, sim_flags @@ -298,11 +300,47 @@ def test_jumpfinder(self): class FFTTest(unittest.TestCase): def test_psd(self): - tod = get_tod("white") - f, Pxx = tod_ops.fft_ops.calc_psd(tod, nperseg=256) - self.assertEqual(len(f), 129) # nperseg/2 + 1 - f, Pxx = tod_ops.fft_ops.calc_psd(tod, freq_spacing=.1) - self.assertEqual(np.round(np.median(np.diff(f)), 1), .1) + fs = 200. + dets = core.LabelAxis('dets', [f'det{di:003}' for di in range(20)]) + nsamps = 200*3600 + + aman = core.AxisManager(dets) + ndets = aman.dets.count + + white_noise_amp_array_input = 50 + np.random.randn(ndets) #W/sqrt{Hz} + fknee_array_input = 1 + 0.1*np.random.randn(ndets) + alpha_array_input = 3 + 0.2*np.random.randn(ndets) + + freqs = rfftfreq(nsamps, d=1/fs) + params = [white_noise_amp_array_input[:, np.newaxis], + fknee_array_input[:, np.newaxis], + alpha_array_input[:, np.newaxis]] + pxx_input = tod_ops.fft_ops.noise_model(freqs, params) + + pxx_input[:, 0] = 0 + nusamps_input = core.OffsetAxis('nusamps_input', len(freqs)) + + T = nsamps/fs + ft_amps = np.sqrt(pxx_input * T * fs**2 / 2) + + ft_phases = np.random.uniform(0, 2*np.pi, size=ft_amps.shape) + ft_coefs = ft_amps * np.exp(1.0j*ft_phases) + realized_noise = irfft(ft_coefs) + timestamps = 1700000000 + np.arange(0, realized_noise.shape[1])/fs + aman.wrap('timestamps', timestamps, [(0, core.OffsetAxis('samps', len(timestamps)))]) + aman.wrap('signal', realized_noise, [(0, 'dets'), (1, 'samps')]) + + tod_ops.detrend_tod(aman) + freqs_output, Pxx_output = tod_ops.fft_ops.calc_psd(aman, nperseg=200*100) + fit_result = tod_ops.fft_ops.fit_noise_model(aman, wn_est=50, fknee_est=1.0, alpha_est=3.3, f_min=0.05, f_max=5, + binning=True, psdargs={'nperseg':200*1000}) + wnl_fit = fit_result.fit[:, 0] + fk_fit = fit_result.fit[:, 1] + alpha_fit = fit_result.fit[:, 2] + + self.assertTrue(np.abs(np.median(white_noise_amp_array_input - wnl_fit)) < 1) + self.assertTrue(np.abs(np.median(fknee_array_input - fk_fit)) < 0.1) + self.assertTrue(np.abs(np.median(alpha_array_input - alpha_fit)) < 0.1) if __name__ == '__main__': unittest.main() From fe6d5beef711f1e4a1f8d1306c5b58fcc0090d86 Mon Sep 17 00:00:00 2001 From: Erik Rosenberg Date: Mon, 13 Jan 2025 10:37:19 -0800 Subject: [PATCH 23/23] fft_ops cleanups and subscan merge fixes --- sotodlib/tod_ops/fft_ops.py | 81 ++++++++++++++++++++++++------------- 1 file changed, 52 insertions(+), 29 deletions(-) diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py index fde63ed22..390824378 100644 --- a/sotodlib/tod_ops/fft_ops.py +++ b/sotodlib/tod_ops/fft_ops.py @@ -376,7 +376,6 @@ def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10): wn = np.sqrt(wn2) return wn - def noise_model(f, params, **fixed_param): """ Noise model for power spectrum with white noise, and 1/f noise. @@ -390,26 +389,21 @@ def noise_model(f, params, **fixed_param): fknee, alpha = params[0], params[1] else: raise ValueError('The number of fit parameters are invalid.') - return elif 'alpha' in fixed_param.keys(): if len(params)==2: alpha = fixed_param['alpha'] wn, fknee = params[0], params[1] else: raise ValueError('The number of fit parameters are invalid.') - return elif len(fixed_param)==0: if len(params)==3: wn, fknee, alpha = params[0], params[1], params[2] else: raise ValueError('The number of fit parameters are invalid.') - return else: raise ValueError('"alpha" or "wn" can be a fixed parameter.') - return return wn**2 * (1 + (fknee / f) ** alpha) - def neglnlike(params, x, y, bin_size=1, **fixed_param): model = noise_model(x, params, **fixed_param) output = np.sum((np.log(model) + y / model)*bin_size) @@ -417,7 +411,7 @@ def neglnlike(params, x, y, bin_size=1, **fixed_param): return 1.0e30 return output -def get_psd_mask(aman, psd_mask=None, f=None, pxx=None, +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), merge=True, overwrite=True @@ -433,8 +427,6 @@ def get_psd_mask(aman, psd_mask=None, f=None, pxx=None, Existing psd_mask to be updated. If None, a new mask is created. f : nparray Frequency of PSD of signal. If None, aman.freqs are used. - pxx : nparray - PSD of signal. If None, aman.Pxx are used. mask_hwpss : bool If True, hwpss are masked. Defaults to True. hwp_freq : float @@ -469,14 +461,13 @@ def get_psd_mask(aman, psd_mask=None, f=None, pxx=None, ------- psd_mask (nusamps): Ranges array. If merge == True, "psd_mask" is added to the aman. """ + if f is None: + f = aman.freqs if psd_mask is None: - psd_mask = np.zeros(aman.nusamps.count, dtype=bool) + psd_mask = np.zeros(f.shape, dtype=bool) elif isinstance(psd_mask, so3g.RangesInt32): psd_mask = psd_mask.mask() - if f is None: - f = aman.freqs - if pxx is None: - pxx = aman.Pxx + if mask_hwpss: hwp_freq = hwp.get_hwp_freq(aman.timestamps, aman.hwp_solution.hwp_angle) psd_mask = psd_mask | get_mask_for_hwpss(f, hwp_freq, max_mode=max_hwpss_mode, width=hwpss_width) @@ -488,6 +479,8 @@ def get_psd_mask(aman, psd_mask=None, f=None, pxx=None, if overwrite: if "psd_mask" in aman: aman.move("psd_mask", None) + if 'nusamps' not in list(aman._axes.keys()): + aman.merge(core.AxisManager(core.OffsetAxis("nusamps", len(f)))) aman.wrap("psd_mask", psd_mask, [(0,"nusamps")]) return psd_mask @@ -544,8 +537,12 @@ def get_binned_psd( if "bin_size" in aman._fields: aman.move("bin_size", None) aman.wrap("freqs_bin", f_bin, [(0,"nusamps_bin")]) - aman.wrap("Pxx_bin", pxx_bin, [(0,"dets"),(1,"nusamps_bin")]) - aman.wrap("bin_size", bin_size, [(0,"dets"),(1,"nusamps_bin")]) + aman.wrap("bin_size", bin_size, [(0,"nusamps_bin")]) + if pxx_bin.ndim > 2 and 'subscans' in aman and pxx_bin.shape[-1] == aman.subscans.count: + aman.wrap("Pxx_bin", pxx_bin, [(0, "dets"), (1, "nusamps_bin"), (2, "subscans")]) + else: + aman.wrap("Pxx_bin", pxx_bin, [(0,"dets"),(1,"nusamps_bin")]) + return f_bin, pxx_bin, bin_size @@ -614,10 +611,16 @@ def fit_noise_model( merge_psd : bool If ``merge_psd`` is True then adds fres and Pxx to the axis manager. mask : bool - If ``mask`` is True then PSD is masked with ``aman.psd_mask``. + If True, (nusamps,) mask is taken from aman.psd_mask, or calculated on the fly. + Can also be a 1d (nusamps,) bool array True for good samples to keep. + Can also be a Ranges in which case it will be inverted before application. + fixed_param : str This accepts 'wn' or 'alpha' or None. If 'wn' ('alpha') is given, white noise level (alpha) is fixed to the wn_est (alpha_est). + binning : bool + True to bin the psd before fitting. + The binning is determined by the 'unbinned_mode' and 'base' params. unbinned_mode : int First Fourier modes up to this number are left un-binned. base : float (> 1) @@ -646,17 +649,24 @@ def fit_noise_model( subscan=subscan, **psdargs, ) - if mask: - if 'psd_mask' in aman: - mask = ~aman.psd_mask.mask() - f = f[mask] - pxx = pxx[:, mask] + if np.any(mask): + if isinstance(mask, np.ndarray): + pass + elif isinstance(mask, Ranges): + mask = ~mask.mask() + elif mask == True: + if 'psd_mask' in aman: + mask = ~aman.psd_mask.mask() + else: # Calculate on the fly + mask = ~(get_psd_mask(aman, f=f, merge=False).mask()) else: - print('"psd_mask" is not in aman. Masking is skipped.') + raise ValueError("mask should be an ndarray or True") + f = f[mask] + pxx = pxx[:, mask] if subscan: fit_noise_model_kwargs = {"fknee_est": fknee_est, "wn_est": wn_est, "alpha_est": alpha_est, - "f_min": f_min, "f_max": f_max, "mask": mask, "fixed_param": fixed_param, + "f_min": f_min, "f_max": f_max, "fixed_param": fixed_param, "binning": binning, "unbinned_mode": unbinned_mode, "base": base, "freq_spacing": freq_spacing} fitout, covout = _fit_noise_model_subscan(aman, signal, f, pxx, fit_noise_model_kwargs) @@ -880,17 +890,19 @@ def log_binning(psd, unbinned_mode=3, base=1.05, mask=None, # Ensure psd is at least 2D for consistent processing psd = np.atleast_2d(psd) - num_signals, num_samples = psd.shape + num_signals, num_samples = psd.shape[:2] if mask is not None: # Ensure mask is at least 2D and has the same shape as psd mask = np.atleast_2d(mask) if mask.shape[1] != num_samples: raise ValueError("Mask must have the same number of columns as psd") - psd = np.ma.masked_array(psd, mask=np.tile(mask, (num_signals, 1))) + mask = np.tile(mask, (num_signals,) + psd.shape[2:] + (1,)) + mask = np.moveaxis(mask, -1, 1) + psd = np.ma.masked_array(psd, mask=mask) # Initialize the binned PSD and optionally the bin sizes - binned_psd = np.zeros((num_signals, unbinned_mode + 1)) + binned_psd = np.zeros((num_signals, unbinned_mode + 1) + psd.shape[2:]) binned_psd[:, :unbinned_mode + 1] = psd[:, :unbinned_mode + 1] bin_size = np.ones((num_signals, unbinned_mode + 1)) if return_bin_size else None @@ -908,7 +920,8 @@ def log_binning(psd, unbinned_mode=3, base=1.05, mask=None, new_bin_size.append(end - start) # Convert lists to numpy arrays and concatenate with initial values - new_binned_psd = np.array(new_binned_psd).T # Transpose to match dimensions + new_binned_psd = np.array(new_binned_psd) + new_binned_psd = np.moveaxis(new_binned_psd, 0, 1)# Transpose to match dimensions binned_psd = np.hstack([binned_psd, new_binned_psd]) if return_bin_size: new_bin_size = np.array(new_bin_size) @@ -944,12 +957,22 @@ def _fit_noise_model_subscan( fitout = np.empty((aman.dets.count, 3, aman.subscan_info.subscans.count)) covout = np.empty((aman.dets.count, 3, 3, aman.subscan_info.subscans.count)) + per_subscan = {} + for entry in ["fknee_est", "wn_est", "alpha_est"]: + if (entry in fit_noise_model_kwargs): + val = fit_noise_model_kwargs[entry] + if isinstance(val, np.ndarray) and val.ndim > 1 and val.shape[-1] == aman.subscan_info.subscans.count: + per_subscan[entry] = fit_noise_model_kwargs.pop(entry) + + kwargs = fit_noise_model_kwargs.copy() if len(per_subscan) > 0 else fit_noise_model_kwargs for isub in range(aman.subscan_info.subscans.count): if np.all(np.isnan(pxx[...,isub])): # Subscan has been fully cut fitout[..., isub] = np.full((aman.dets.count, 3), np.nan) covout[..., isub] = np.full((aman.dets.count, 3, 3), np.nan) else: - noise_model = fit_noise_model(aman, f=f, pxx=pxx[...,isub], merge_fit=False, merge_psd=False, subscan=False, **fit_noise_model_kwargs) + for entry in list(per_subscan.keys()): + kwargs[entry] = per_subscan[entry][..., isub] + noise_model = fit_noise_model(aman, f=f, pxx=pxx[...,isub], merge_fit=False, merge_psd=False, subscan=False, **kwargs) fitout[..., isub] = noise_model.fit covout[..., isub] = noise_model.cov