New functions for fft_ops #869

17-sugiyama · 2024-05-23T10:43:23Z

This PR suggests adding calc_masked_psd and calc_binned_psd to fft_ops.
calc_masked_psd masks hwpss peaks and other peaks in the given PSD. This enables applying 1/f curve fitting to data with hwpss.
calc_binned_psd puts given PSD into bins. This makes fit_noise_model work much faster.
Here is a demonstration for new functions.
https://github.com/17-sugiyama/scratch/blob/main/20240513_fft_ops_test.ipynb

kmharrington

I'm noticing this PR brings in masked numpy arrays, I think for the first time in sotodlib. I'm worried that will break our IO setup going into and out of hdf5 files.

Is there a reason we can't use Ranges / RangesMatrices to pass around masks here?

I would also like the calc_masked_psd to accept more general masks. The function would be much more useful if it was generalized beyond HWPSS lines. Ex: what if I also want to mask PTC lines?

17-sugiyama · 2024-06-04T01:32:39Z

I was using MaskedArray for no reason, so I modified the functions not to use it.
The change also made the calculation faster. Thanks for your comments!
It's not apparent but the function calc_masked_psd has an option to mask a single peak like PTC lines.
fft_ops.calc_masked_psd(aman, peak=True, peak_freq=peak_freq, peak_width=peak_width) will do that.
Currently peak_freq accepts single frequency, so we need to run the code several times to mask multiple peaks.

tterasaki · 2024-06-07T06:55:49Z

sotodlib/hwp/hwp.py

    # 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)



Can you implement below and use here and elsewhere?

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

tterasaki · 2024-06-07T06:56:24Z

sotodlib/tod_ops/fft_ops.py

+    if hwpss:
+        pxx_masked = []
+        if hwp_freq is None:
+            hwp_freq = np.median(aman['hwp_solution']['raw_approx_hwp_freq_1'][1])


Let's use the code I suggested above.

tterasaki · 2024-06-07T07:15:41Z

sotodlib/tod_ops/fft_ops.py

@@ -457,3 +606,105 @@ def fit_noise_model(
    if merge_fit:
        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)):


I think we can implement it more straight forward.
Like

def get_mask_for_hwpss(freq, hwp_freq, modes=[1,2,3,4,5,6], width=0.4): if isinstance(width, float): "apply (-width/2, width/2) for all nf" elif isinstance(width, [np.array, list, tuple]): width = np.array(width) if len(width.shape) == 1: "apply symmetric mask for each nf" elif len(width.shape) == 2: "apply asymmetric mask for each nf"

tterasaki · 2024-06-07T07:18:51Z

sotodlib/tod_ops/fft_ops.py

+    """
+    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])))


Let's use get_mask_for_single_peak defined below.

tterasaki · 2024-06-07T07:19:36Z

sotodlib/tod_ops/fft_ops.py

+    mask = np.all(np.array(mask_arrays), axis=0)
+    return mask
+
+def peak_mask(f, peak_freq, peak_width=(-0.002, +0.002)):


I think get_mask_for_single_peak is more straightforward naming.

17-sugiyama · 2024-07-08T10:05:04Z

Thank you so much for your comments!
Besides applying your comments, I modified several things.

calc_pds_mask has a new option, mode. If mode==replace, it will calculate new mask array. If mode==add, newly calculated mask will be stacked on the existing one.
fit_noise_model has a new option, fixed_parameter. You can declare a parameter to fix during the PSD fitting.
fit_binned_noise_model is a new function to make PSD fittings much faster. besides knee frequency and white noise (fit parameters) are unbiased, the calculation becomes more than 20 times faster.

The demonstration of new features locates here.
https://github.com/17-sugiyama/scratch/blob/main/20240705_fft_ops_PR.ipynb

17-sugiyama · 2024-07-16T11:51:43Z

Extended neglnlike function to apply to binned PSD.
This made fit results of binned and unbinned PSDs almost identical.

sotodlib/sotodlib/tod_ops/fft_ops.py

Lines 360 to 365 in 9575bf9

    
           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) 
        
               if not np.isfinite(output): 
        
                   return 1.0e30 
        
               return output

tterasaki

Ok, you are almost there to get approved with minor my comments.
Another thing I want to ask you to add is a unit test for the fitting function.

I have already written basic component for the unit test like below.
Can you add the unit test using those at https://github.com/simonsobs/sotodlib/tree/fft_ops_hwpss_mask/tests ?
https://github.com/simonsobs/sotodlib/blob/fft_ops_hwpss_mask/tests/test_subpolyf.py (or some other tests) would be a good reference.

import numpy as np
from sotodlib import core
from numpy.fft import rfftfreq,irfft
from numpy.fft import fftfreq,rfft
from scipy import interpolate
import matplotlib.pyplot as plt
from sotodlib.tod_ops import fft_ops

def model_func(x, sigma, fk, alpha):
    return sigma**2 * (1 + (fk/x)**alpha)

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}
white_noise_power_array_input = white_noise_amp_array_input**2 #W^2/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)
pxx_input = model_func(freqs, 
           sigma=white_noise_amp_array_input[:, np.newaxis], 
           fk=fknee_array_input[:, np.newaxis],
           alpha=alpha_array_input[:, np.newaxis])

pxx_input[:, 0] = 0
nusamps_input = core.OffsetAxis('nusamps_input', len(freqs))

aman.wrap('freqs_input', freqs, [(0, nusamps_input)])
aman.wrap('wnl_input', white_noise_amp_array_input, [(0, 'dets')])
aman.wrap('fk_input', fknee_array_input, [(0, 'dets')])
aman.wrap('alpha_input', alpha_array_input, [(0, 'dets')])
aman.wrap('Pxx_input', pxx_input, [(0, 'dets'), (1, 'nusamps_input')])

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')])

tterasaki · 2024-06-07T07:45:43Z

sotodlib/tod_ops/fft_ops.py

+    if peak:
+        mask_idx = peak_mask(f, peak_freq, peak_width=peak_width)
+        f = f[mask_idx]
+        pxx = pxx[:, mask_idx]


It is not good to discard noise spectrum.
I think what Katie meant is to use flags manager for masking.

tterasaki · 2024-07-30T08:29:45Z

sotodlib/tod_ops/fft_ops.py

-    return output
+    return output   
+
+def get_hwp_freq(timestamps=None, hwp_angle=None):


Can you move this function to sotodlib.hwp.hwp?

tterasaki · 2024-08-26T09:58:46Z

sotodlib/tod_ops/fft_ops.py

+    else:
+        raise ValueError('"alpha" or "wn" can be a fixed parameter.')
+        return
+    return wn * (1 + (fknee / f) ** alpha)


Can you modify the definition of white noise level (wn)?
The unit of wn here (=W^2/Hz) is inconsistent with the white noise level in other place(=W/sqrt{Hz}).

tterasaki · 2024-08-26T10:01:51Z

sotodlib/tod_ops/fft_ops.py

+            return
+    elif len(fixed_param)==0:
+        if len(params)==3:
+            fknee, wn, alpha = params[0], params[1], params[2]


Can you modify the order of parameters?
I think the order of [wn, fknee, alpha] is more natural and less confusing.

tterasaki · 2024-08-26T10:02:57Z

sotodlib/tod_ops/fft_ops.py

    fitout = np.zeros((aman.dets.count, 3))
    # This is equal to np.sqrt(np.diag(cov)) when doing curve_fit
    covout = np.zeros((aman.dets.count, 3, 3))
-    for i in range(aman.dets.count):
+    if isinstance(fknee_est, (int, float)):
+        fknee_est = np.zeros(aman.dets.count)+fknee_est


np.full is a bit smarter.

17-sugiyama · 2024-09-24T15:37:27Z

I've reflected Tomoki's latest comments and added the unittest for fft_ops.
Investigating test failures

To check.

Junna Sugiyama added 2 commits May 23, 2024 02:33

added masking and binning function

7969bed

hwp_angle can be declared

5e88252

17-sugiyama requested review from msilvafe, RemingtonGerras and tterasaki May 23, 2024 10:43

kmharrington reviewed May 28, 2024

View reviewed changes

Modified not to use Numpy.MaskedArray

3fc19ae

tterasaki reviewed Jun 7, 2024

View reviewed changes

17-sugiyama added 2 commits July 8, 2024 03:53

modified binning and fitting function

307ab16

Reflected Tomoki's comments

3c96965

Modified fit_binned_noise_model

9575bf9

17-sugiyama and others added 12 commits July 17, 2024 06:39

Minor debug

81fbde8

Merge branch 'master' into fft_ops_hwpss_mask

5bbde3d

moved get_hwp_freq to sotodlib.hwp.hwp

8116be2

small bug fix

f5673c8

improved calc_psd merging style

912ac2e

from capital PSD_mask to psd_mask

97d2127

binning_psd can accept 2d array

63b9608

binning_psd can accept 2d array

c946071

log_binning modification

65f4125

integrated fit_noise_model and fit_binned_noise_model

f138b70

fit_noise_model can have a constant fixed parameter's value

15a83a9

put nan in covout where the parameter is fixed

d3eaabd

changed some function's name

b57326e

tterasaki self-requested a review August 8, 2024 10:57

17-sugiyama added 2 commits August 8, 2024 11:00

minor bug fix

57f5c24

give initial estimates of parameters for fit_noise_model

3d02df5

tterasaki requested changes Aug 26, 2024

View reviewed changes

17-sugiyama added 2 commits September 12, 2024 10:52

reflected Tomoki's comments

91c4381

add unittest

f10e0f0

tterasaki force-pushed the fft_ops_hwpss_mask branch from a1bf7bd to f10e0f0 Compare September 26, 2024 05:20

tterasaki previously approved these changes Sep 26, 2024

View reviewed changes

tterasaki self-requested a review September 26, 2024 05:23

Merge branch 'master' into fft_ops_hwpss_mask

db40980

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

New functions for fft_ops #869

New functions for fft_ops #869

17-sugiyama commented May 23, 2024

kmharrington left a comment

17-sugiyama commented Jun 4, 2024

tterasaki Jun 7, 2024

tterasaki Jun 7, 2024

tterasaki Jun 7, 2024 •

edited

Loading

tterasaki Jun 7, 2024

tterasaki Jun 7, 2024

17-sugiyama commented Jul 8, 2024

17-sugiyama commented Jul 16, 2024

tterasaki left a comment

tterasaki Jun 7, 2024

tterasaki Jul 30, 2024

tterasaki Aug 26, 2024

tterasaki Aug 26, 2024

tterasaki Aug 26, 2024

17-sugiyama commented Sep 24, 2024 •

edited

Loading

New functions for fft_ops #869

Are you sure you want to change the base?

New functions for fft_ops #869

Conversation

17-sugiyama commented May 23, 2024

kmharrington left a comment

Choose a reason for hiding this comment

17-sugiyama commented Jun 4, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tterasaki Jun 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

17-sugiyama commented Jul 8, 2024

17-sugiyama commented Jul 16, 2024

tterasaki left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

17-sugiyama commented Sep 24, 2024 • edited Loading

tterasaki Jun 7, 2024 •

edited

Loading

17-sugiyama commented Sep 24, 2024 •

edited

Loading