Add FT-beam and window function CLI; fix k_parallel scaling of exact WFs by vincentmackay · Pull Request #514 · HERA-Team/hera_pspec

vincentmackay · 2026-07-07T19:44:45Z

Addresses #451 (will be fully closed when I also address hera_pipelines accordingly). Moves the FT-beam and window function production scripts from hera_pipelines into hera_pspec as pspec compute-ft-beam / pspec compute-window-functions.

Note that working on this exposed a bug that affected the IDR3.1.1a theory release window functions: UVWindow.from_uvpspec assumed the FT beam and the data share a channel width, but I was accidentally computing the FT beam on a 97.66 kHz grid vs the data's 122.07 kHz, scaling the WFs' k_parallel axis by 1.25x. To prevent this error in the future, I forced the FT beam to be computed on the data's exact channels (FTBeam.select_freqs; interpolates with a warning on mismatch, raises an exception if not covered).

commit 1: the bug fix + regression tests
commit 2: vectorized cylindrical binning (~20x faster; new test pins it to the old loop implementation at 1e-12)
commit 3: FTBeam.from_beam (port of @adeliegorce's routine, az_za + healpix), write_hdf5, the two CLI commands

Parts of this were done with the help of Claude (Fable 5), e.g., for finding the vectorization, and documenting/cleaning up.

codecov · 2026-07-07T19:49:46Z

Codecov Report

❌ Patch coverage is 60.64257% with 98 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.58%. Comparing base (9e886f3) to head (43e0a3f).

Files with missing lines	Patch %	Lines
src/hera_pspec/cli.py	15.09%	90 Missing ⚠️
src/hera_pspec/uvwindow.py	94.40%	8 Missing ⚠️

Additional details and impacted files

@@              Coverage Diff               @@
##           release/v1     #514      +/-   ##
==============================================
- Coverage       95.90%   94.58%   -1.32%     
==============================================
  Files              19       19              
  Lines            6562     6776     +214     
==============================================
+ Hits             6293     6409     +116     
- Misses            269      367      +98

Flag	Coverage Δ
unittests	`94.58% <60.64%> (-1.32%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

steven-murray · 2026-07-08T07:51:50Z

@vincentmackay I'm very sorry I should have remembered that we're currently merging bigger feature changes to the release/v1 branch. There are a lot of nice updates on that branch concerning how the CLI is specified (we use cyclopts instead of typer -- they're similar but cyclopts has more features). If we merge this to main there are going to be a lot of annoying merge conflicts etc that we need to deal with.

Could you get Fable to reapply essentially the same logic but branching off of release/v1? Sorry about that!

…data frequency grids differ UVWindow.from_uvpspec extracted len(freq_array) FTBeam channels starting at the data's first frequency, silently assuming the FTBeam and the data had identical channel widths. Because the window functions' k_parallel axis is built from the FTBeam's freq_array (via delta_nu in _get_wf_for_tau), any mismatch scaled that axis by delta_nu_data / delta_nu_ftbeam: 122.07/97.66 = 1.25x for the H6C IDR3 window functions (FT beam computed as 2048 channels over 50-250 MHz vs the data's native 122.07 kHz grid). Delta^2 and its errors were unaffected (binned on the data's own grid in spherical_average). - New FTBeam.select_freqs(freq_array): restrict the FTBeam to the exact frequency channels of the data; channels are extracted when they coincide (within atol) and the FT of the beam is interpolated in frequency (with a warning) otherwise. from_uvpspec now uses it, and raises if the data frequencies are not covered by the FTBeam. - FTBeam.from_file accepts freq_array and reads only the relevant channel range from disk (large memory saving for wide-band FT beams). - Fix pol-comparison bug in from_uvpspec: 'pol[ip] == pol[0]' compared characters of the pol string, so equal-pol pairs (e.g. pI,pI) were never detected (and line 474 tested a stale loop variable). - ftbeam=None in from_uvpspec now raises a clear error instead of calling the unimplemented FTBeam.from_beam with a dummy beamfile. - Tests: select_freqs (extract/interpolate/out-of-range), partial-read from_file, and an end-to-end regression test checking that a mismatched-grid FTBeam yields the same window functions as one built directly on the data channels. test_from_uvpspec's uvp_nocosmo now uses spw (175, 195): the test FT-beam file only covers data channels 170-199, which the old code silently ignored. Co-authored-by: Claude <noreply@anthropic.com>

The sky-plane average in get_cylindrical_wf rebuilt a boolean mask over the whole (kperp_x, kperp_y) grid for every (frequency, kperp bin) pair -- although the mask does not depend on frequency -- and _get_wf_for_tau ran a nested Python loop over (kperp bin, kpara bin) for every delay. Both are replaced by bin-index computations (np.digitize) and weighted bincounts / a matrix product, with identical bin-edge semantics (edges[m] <= k < edges[m+1], empty bins left at zero). ~20x faster at H6C problem sizes. Results are identical to the loop implementation: a new test compares against a verbatim copy of the old loops on real FT-beam data (agreement at 1e-12, i.e. float round-off). Co-authored-by: Claude <noreply@anthropic.com>

…unctions CLI commands Moves the FT-beam and window-function production scripts from hera_pipelines into hera_pspec (HERA-Team/hera_pipelines#35 review), following the cyclopts CLI conventions of this branch. - FTBeam.from_beam: compute the FT of the beam from a UVBeam simulation file (port of A. Gorce's FT_beam_routine, using UVBeam.interp at the Cartesian-grid sky coordinates so both az_za and healpix beams work). Frequencies slightly outside the simulation coverage (e.g. H6C band 1 starts at 49.12 MHz vs the Vivaldi simulation's 50 MHz floor) get the beam evaluated at the nearest covered frequency, with a warning. - FTBeam.write_hdf5: write files readable by from_file, storing pol as an attribute (from_file prefers it over filename parsing when present, and falls back to the filename otherwise). - pspec compute-ft-beam: reuse-or-compute an FT-beam file. The frequency grid is always taken from the pspec file itself; there is deliberately no option to specify a grid by hand, since a grid other than the data's scales the window functions' k_parallel axis (use FTBeam.from_beam for arbitrary grids). - pspec compute-window-functions: exact window functions per (spw, polpair), one HDF5 each (wf, kperp, kpara + provenance attrs), filename-based reuse, optional --workers process parallelism. The FT beam is read only over each spw's channel range. Co-authored-by: Claude <noreply@anthropic.com>

for more information, see https://pre-commit.ci

vincentmackay · 2026-07-08T19:30:20Z

@vincentmackay I'm very sorry I should have remembered that we're currently merging bigger feature changes to the release/v1 branch. There are a lot of nice updates on that branch concerning how the CLI is specified (we use cyclopts instead of typer -- they're similar but cyclopts has more features). If we merge this to main there are going to be a lot of annoying merge conflicts etc that we need to deal with.

Could you get Fable to reapply essentially the same logic but branching off of release/v1? Sorry about that!

Done, rebased onto release/v1 and rewrote the CLI with cyclopts.

steven-murray

Thanks you @vincentmackay this is looking fantastic! I just have a few small comments.

steven-murray · 2026-07-09T04:11:06Z

+    mapsize
+        Half-width of the flat-sky map the beam is projected onto (rad).


I haven't looked further down yet, but I am assuming this is not actually in radians but in direction cosines? I.e. l=sin(theta)?

This is a poorly named and defined parameter and it is totally on me! The cartesian map the beam is interpolated over has width 1 rad * sqrt(2)* mapsize (provided max_r=0.7). If mapsize=1/sqrt(2)~0.7 then the map is 1 rad wide.

steven-murray · 2026-07-09T04:14:49Z

+    out_path = Path(out_dir) / fname
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    with h5py.File(out_path, "w") as f:
+        f.create_dataset("wf", data=wf, compression="gzip")
+        f.create_dataset("kperp", data=kperp)
+        f.create_dataset("kpara", data=kpara)
+        f.attrs["algo"] = "exact"
+        f.attrs["polpair"] = polpair
+        f.attrs["spw_index"] = int(spw)
+        f.attrs["dataset_label"] = dataset_label
+        f.attrs["taper"] = str(sub.taper)
+        f.attrs["spw_freq_min_hz"] = float(spw_freqs.min())
+        f.attrs["spw_freq_max_hz"] = float(spw_freqs.max())
+        f.attrs["pspec_file"] = str(pspec_file)
+        f.attrs["ft_beam_file"] = str(ft_beam_file)
+        f.attrs["created_utc"] = datetime.datetime.now(datetime.UTC).isoformat()
+        f.attrs["producer"] = "pspec compute-window-functions"
+    logger.info("Wrote %s", out_path)


This is fine for now, but I'm flagging that in general it would be better to define external read/write methods for a "window function file" format, and then re-use these across the code, rather than having the format hard-coded inside this function.

steven-murray · 2026-07-09T04:17:42Z

+# Cartesian half-width used for the flat-sky beam projection in
+# FTBeam.from_beam. The Cartesian grid spans rx in [-_MAX_R, +_MAX_R];
+# zenith angle is za = sqrt(x^2 + y^2) in radians, so the corner of the
+# grid sits at za = sqrt(2)*_MAX_R = 1 rad (~57 deg), a sensible outer
+# bound for the beam-relevant sky.
+_MAX_R = 0.7


This is a different default than in the CLI -- is this intentional? I wonder what @adeliegorce recommends this parameter to be set to?

Yes _MAX_R should be set to 0.7. Then to pad the beam simulation and "artificially" cover smaller kperp, I have set mapsize=1 in previous computations; if mapsize=_MAX_R you just interpolate the beam simulation as is, with no padding.

Actually we could probably just set it to 1./np.sqrt(2.)

steven-murray · 2026-07-09T04:23:32Z

    def test_from_beam(self):

-        pytest.raises(NotImplementedError, uvwindow.FTBeam.from_beam, beamfile="test")
+        from pyuvdata import UVBeam


I would put this import at the top of the module

steven-murray · 2026-07-09T04:24:15Z

-        pytest.raises(NotImplementedError, uvwindow.FTBeam.from_beam, beamfile="test")
+        from pyuvdata import UVBeam
+
+        beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits")


Prefer pathlib.Path operations (i.e. beamfile = Path(DATA_PATH) / "HERA_NF_dipole_power.beamfits")

steven-murray · 2026-07-09T04:25:46Z

+        # pol fed as int
+        test2 = uvwindow.FTBeam.from_beam(
+            beamfile=beamfile, pol=-5, freq_array=freq_array, mapsize=1.0, npix=29
+        )
+        assert test2.pol == "xx"
+        assert np.allclose(test.ft_beam, test2.ft_beam)


Test functions should preferably test ONE thing. Split this second (and third) test into different, explicitly named functions.

steven-murray · 2026-07-09T04:26:12Z

+        # inplace=True modifies the object
+        test2 = uvwindow.FTBeam.from_file(ftfile=self.ft_file, spw_range=None)
+        out = test2.select_freqs(self.bandwidth[3:10], inplace=True)
+        assert out is None
+        assert np.allclose(test2.ft_beam, test.ft_beam[3:10])


same comment here.

steven-murray · 2026-07-09T04:27:09Z

+
+    def test_write_hdf5(self):
+
+        import tempfile


there is a pytest fixture that provides temporary files (tmp_path IIRC)

steven-murray · 2026-07-09T04:29:52Z

        wf_array1 = np.zeros((kperp_bins.size, test.Nfreqs))
        kpara, cyl_wf = test._get_wf_for_tau(tau, wf_array1, kperp_bins, kpara_bins)

+    def test_cylindrical_wf_matches_loop_reference(self):


I like that this test has been run, but it seems extremely convoluted for a single test. Not quite sure what to do about it, but keeping a full copy of the old code around just in the tests seems like a code smell to me. I would prefer to test against simple known examples rather than full alternate code implementations.

adeliegorce

Thanks @vincentmackay! good job spotting the frequency channelization mismatch. My main comment is that we need to coordinate merging #491 and this PR as #491 introduces major changes to the testing infrastructure and will create conflicts with the tests you've added

adeliegorce · 2026-07-16T08:52:02Z

 from astropy import units
 from pyuvdata import utils as uvutils
-from scipy.interpolate import RegularGridInterpolator
+from scipy.interpolate import RegularGridInterpolator, interp1d


This is due to my code but interp1d is legacy code and we probably should not use it

adeliegorce · 2026-07-16T09:06:22Z

            Used to convert polstr to polnum and conversely.
        """
-        raise NotImplementedError("Coming soon...")
+        from pyuvdata import UVBeam


This import could probably be moved to the top

adeliegorce · 2026-07-16T09:08:39Z

            In (start_chan, end_chan). Must be between 0 and 1024 (HERA
            bandwidth).
            If None, whole instrument bandwidth is considered.
+        freq_array : 1D array or list of floats, optional


Isn't the spw_range parameter redundant now we have freq_array?

adeliegorce · 2026-07-16T09:16:47Z

Unfortunately this will be extremely annoying to merge with #491 . It is probably worth putting any test writing on hold until this PR has been merged. Then Claude can easily merge your new tests into the new test_uvwindow.

vincentmackay changed the base branch from main to release/v1 July 8, 2026 15:19

vincentmackay and others added 4 commits July 8, 2026 11:29

style: apply pre-commit fixes from the main-based branch

0f53507

vincentmackay force-pushed the wf-cli-scripts branch from 0f64647 to 0f53507 Compare July 8, 2026 19:07

[pre-commit.ci] auto fixes from pre-commit.com hooks

43e0a3f

for more information, see https://pre-commit.ci

steven-murray reviewed Jul 9, 2026

View reviewed changes

adeliegorce reviewed Jul 16, 2026

View reviewed changes

		mapsize
		Half-width of the flat-sky map the beam is projected onto (rad).

Uh oh!

Conversation

vincentmackay commented Jul 7, 2026

Uh oh!

codecov Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

steven-murray commented Jul 8, 2026

Uh oh!

vincentmackay commented Jul 8, 2026

Uh oh!

steven-murray left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

adeliegorce left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov Bot commented Jul 7, 2026 •

edited

Loading