Add joint demodulated Q/U Fourier Nmat filter - #1709
Conversation
Adds tod_ops.nmat_filter.apply_joint_qu_nmat_filter, which whitens and jointly analyzes the demodulated Q and U detector streams. Modes inconsistent with independent noise, per a Marchenko-Pastur plus Tracy-Widom threshold, are inverse-covariance weighted in Fourier space. Exposed as the JointQUNmatFilter preprocess step, which supports use_data_aman so the noise operator can be estimated from a real-data snapshot rather than from a signal-only simulation.
ykyohei
left a comment
There was a problem hiding this comment.
I haven't read everything yet, but I have one suggestion.
I believe we can save the relatively small set of statistics required to reproduce the NMat filter and use them for the pre-processing run using the archive. This is similar to how we apply the 1/f counter filter using the saved values (wn, fknee, alpha). Then, we do not need to handle model_aman for this.
msilvafe
left a comment
There was a problem hiding this comment.
This is an awesome addition @susannaaz thanks for implementing this so quickly and the initial results are super promising. I have a few comments but we should talk over this review to decide what you want to implement now vs later (or at all) as I know this is something you want to get to running on more data ASAP.
The main structural thing I want to flag is as written, apply_joint_qu_nmat_filter always derives the operator (whitening, covariance, mode detection,
I've broken out 4 inline comments walking through one way to address this: 1) splitting apply_joint_qu_nmat_filter into a fit step (everything derived from real data) and an apply step (everything derived from the target TOD), 2) adding calc_and_save/save so a sim run can reload rather than refit it, 3) a small utility for reconstructing what's actually being subtracted from a given detector for debugging, and 4) some stats on how many modes get selected vs. rejected and why (MP threshold, singleness veto, or the n_modes_max cap).
There are a few more inline comments which I think I'd like to get a response on but take-it-or-leave-it in terms of actually implementing anything in code.
Addresses review: the operator is no longer derived and applied in the same call, so it can be fit on real data and reapplied to signal-only sims without redoing the real-data analysis. - fit_joint_qu_nmat_operator: everything derived from the model TOD (whitening sigma, mode selection, D(f), E(f)), returned as an AxisManager on the dets/nmat_modes/nmat_profile/nmat_bins axes. - apply_joint_qu_nmat_operator: applies a stored operator to a target. apply_joint_qu_nmat_filter is now a thin wrapper over the two. - New joint_qu_nmat_model preprocess step (calc/save) writes the operator to proc_aman; joint_qu_nmat_filter reads it back, mirroring the noise -> fourier_filter pattern. Drops use_data_aman. - Record per-bin selection stats: n_above_lambda_plus, n_above_threshold, n_failed_singleness, n_capped_by_max, n_selected, plus per-mode singleness. get_marchenko_pastur_threshold gained return_edge. - get_nmat_subtraction returns what the filter removes, for comparing against other filters in TOD space. - Drop the mode_shrink operator. - Reuse mapmaking.utils.makebins for the profile bins, with rfun=floor and clamped endpoints; the default ceil drops the lowest Fourier mode. - irfft normalize=True rather than "phys", and trim input validation.
The operator is stored on the dets axis, so reloading it from an archive alongside an observation with a tighter detector cut silently restricts the mode vectors, which are then no longer orthonormal over the retained subset. Record the detector count the modes were fit over and warn if it does not match at apply time.
Restricting D and V yields the marginal covariance of the subset under the fitted model, which Woodbury inverts correctly; the earlier text claimed the operator was only approximate because the mode vectors are no longer orthonormal, which is not the reason. The warning still matters because a marginalised operator differs from one fit on the subset.
With use_data_aman: True the step fits the operator directly from the supplied real-data AxisManager rather than reloading a stored one, taking the fit parameters from process.fit. The stored-operator path remains the default.
cdb98ae to
d5ee70a
Compare
Thanks @msilvafe , I've implemented all four structural points, and the rest is answered inline. Summary:
I also dropped |
msilvafe
left a comment
There was a problem hiding this comment.
Ok one more round of inline comments to reply to. Much simpler this time hit-rereview for me when ready and I'll do next round approve fast. Will try to checkout this branch and start running some data with it now. Thank you!!
| ndet, nsamp = q.shape | ||
| nfreq = nsamp // 2 + 1 | ||
| real_dtype = np.result_type(q.dtype, u.dtype) | ||
| complex_dtype = np.complex64 if real_dtype.itemsize <= 4 else np.complex128 |
There was a problem hiding this comment.
Why have these different complex datatypes?
| cfgs = dict(self.process_cfgs) | ||
| model_name = cfgs.pop("nmat_model", "nmat_qu") | ||
| fit_cfgs = cfgs.pop("fit", None) | ||
|
|
There was a problem hiding this comment.
self.process_cfgs should already be a dict, also in_place should also be popped as you hardcode that in L2462
| op.wrap("fmin", float(freqs[band_lo])) | ||
| op.wrap("fmax", float(freqs[band_hi - 1])) | ||
| op.wrap("profile_diagonal_floor", float(profile_diagonal_floor)) | ||
| op.wrap("nsamp", int(nsamp)) |
There was a problem hiding this comment.
fmin, fmax, nsamp wrapped but not used again when op is loaded by other functions. Use them in the loaded functions instead of recomputing or what was your thought on their intended purpose?
| if self.use_data_aman: | ||
| model_aman = data_aman if data_aman is not None else aman | ||
| logger.info( | ||
| "Fitting the Nmat operator from the supplied AxisManager " | ||
| "(%d detectors)", model_aman.dets.count | ||
| ) | ||
| operator = tod_ops.nmat_filter.fit_joint_qu_nmat_operator( | ||
| model_aman, | ||
| signal_Q=self.signal_Q, | ||
| signal_U=self.signal_U, | ||
| **(fit_cfgs or {}), | ||
| ) |
There was a problem hiding this comment.
We should probably remove this block now, or when do you think we should go down this branch?
| if n0 < 0 or n1 <= n0: | ||
| raise ValueError("noise_band must satisfy 0 <= low < high") |
There was a problem hiding this comment.
This should be checked against against the demodulation lowpass filter and not use any frequencies above the lowpass. That was the intention of adding proc_aman["frequency_cutoffs"]. But the appropriate place for this is probably in processes.py though so you can use info in proc_aman. There's an existing example in Noise.calc_and_save in processes.py Lines 745-755. Here you can at least check against the nyquist frequency i.e. n1 must be < freqs[-1]
| def _split(arr_over_channels, fill=0.0): | ||
| """Scatter a valid-channel array back to (dets,) Q and U halves.""" | ||
| full = np.full(2 * ndet, fill, dtype=float) | ||
| full[valid_idx] = arr_over_channels | ||
| return full[:ndet], full[ndet:] |
There was a problem hiding this comment.
This function seems to never be used.
Adds tod_ops.nmat_filter.apply_joint_qu_nmat_filter, which whitens and jointly analyzes the demodulated Q and U detector streams. Modes inconsistent with independent noise, per a Marchenko-Pastur plus Tracy-Widom threshold, are inverse-covariance weighted in Fourier space.
Included as the
JointQUNmatFilterpreprocess step, which supportsuse_data_amanso the noise operator can be estimated from a real-data snapshot rather than from a signal-only simulation for transfer function purposes.