feat: add polarized sky model support via eigendecomposition M matrix by kartikmandar · Pull Request #114 · HERA-Team/matvis

kartikmandar · 2026-04-04T18:31:17Z

Summary

Implements full Stokes (I, Q, U, V) sky model support via eigendecomposition of the coherency matrix C = M M† (arXiv:2312.09763, Section 3.2 & Appendix B)
Sign-split approach for EOR/negative-flux scenarios (V = V_pos - V_neg)
Fully backward compatible — omitting stokes param runs existing code with zero overhead

New API: simulate_vis(..., stokes=np.array([I,Q,U,V]), negative_flux="raise"|"split")

Test plan

20 unit tests for M matrix math
7 accuracy tests (old vs new path matches to atol=1e-12)
6 runtime benchmarks (~1.1x overhead for polarized sky)
All 103 existing tests pass
GPU path untested locally (mirrors CPU changes)

Copilot

Pull request overview

Adds full-Stokes polarized sky-model support to matvis by decomposing the 2×2 coherency matrix into an M matrix (including a sign-split option for negative-flux/EoR-like scenarios), while keeping the existing unpolarized code path unchanged when stokes is omitted.

Changes:

Introduces coherency/Stokes conversion utilities and M-matrix construction (eigendecomposition + sign-split).
Plumbs stokes and negative_flux through the public wrapper into CPU/GPU simulators and updates Z-matrix construction to accept an M matrix.
Adds extensive unit/accuracy tests and a new benchmark test module.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
tests/test_polarized_sky.py	Accuracy/back-compat tests comparing old vs new polarized sky paths
tests/test_polarized_benchmark.py	Runtime benchmark-style tests (currently collected by pytest)
tests/test_coherency.py	Unit tests for coherency ↔ Stokes and M-matrix math
src/matvis/wrapper.py	Public API additions: `stokes`, `negative_flux`; per-frequency stokes slicing
src/matvis/cpu/cpu.py	CPU implementation of polarized-sky path + sign-split support
src/matvis/gpu/gpu.py	GPU implementation of polarized-sky path + sign-split support
src/matvis/core/getz.py	Adds `m_matrix` option to build Z for polarized sky
src/matvis/core/coherency.py	New coherency + M-matrix math module
src/matvis/core/init.py	Extends input validation to include `stokes`

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

codecov · 2026-04-06T11:20:24Z

Codecov Report

❌ Patch coverage is 89.86667% with 38 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.34%. Comparing base (af58859) to head (b8074c9).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/matvis/core/coherency.py	91.15%	11 Missing and 9 partials ⚠️
src/matvis/cpu/cpu.py	85.41%	3 Missing and 4 partials ⚠️
src/matvis/gpu/gpu.py	86.00%	3 Missing and 4 partials ⚠️
src/matvis/core/__init__.py	75.00%	1 Missing and 1 partial ⚠️
src/matvis/wrapper.py	85.71%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #114      +/-   ##
==========================================
- Coverage   98.88%   96.34%   -2.55%     
==========================================
  Files          22       23       +1     
  Lines         989     1341     +352     
  Branches      103      158      +55     
==========================================
+ Hits          978     1292     +314     
- Misses          6       25      +19     
- Partials        5       24      +19

Flag	Coverage Δ
unittests	`84.34% <66.40%> (+5.87%)`	⬆️

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

Thank you so much @kartikmandar this is looking really great!

I do have some comments that I think will make it tighter. I also think that another test that could be added is to run some set of random sources through the new polarized visibility simulator and then run their negative through, and check that the visibilities are the negative of each other.

Also, should it not be true that if a source is "physical" (i.e. positive eigenvalues) at one time, it should be physical at all other times (assuming a physical beam)? I think this must be true. In that case, we should be able to do the masking once in the sim, raise an error if anything is negative and negative_flux="raise", and then move on.

…-split M matrix Implements full Stokes (I, Q, U, V) sky model support for the matvis visibility simulator. The coherency matrix C = M M† is computed via eigendecomposition (λ± = I ± T), enabling polarized sky simulations through the existing V = Z Z† framework. Two approaches: - Eigendecomposition: for physical sources where eigenvalues are non-negative - Sign-split: for EOR/negative-flux scenarios, splits into V_pos - V_neg New API: simulate_vis(..., stokes=array(4, Nsrc, Nfreq), negative_flux="raise"|"split") Backward compatible: omitting stokes runs the existing sqrt(I) path unchanged. New files: - src/matvis/core/coherency.py: M matrix math (eigendecomp, sign-split, Stokes↔coherency) - tests/test_coherency.py: 20 unit tests for coherency math - tests/test_polarized_sky.py: 7 accuracy tests (old vs new path, sign-split, edge cases) - tests/test_polarized_benchmark.py: 6 runtime benchmarks Modified files: - core/getz.py: vectorized k-contraction for full 2x2 M matrix - core/__init__.py: stokes validation - cpu/cpu.py, gpu/gpu.py: stokes/negative_flux params, M matrix pipeline - wrapper.py: stokes/negative_flux API

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

…reservation - Implement negative_flux="ignore" with explicit branching in cpu/gpu backends - Add ValueError for unknown negative_flux values - Fix has_neg: use xp.any(a | b) instead of Python or on cupy 0-d arrays - Convert eigenvalue check to Python floats upfront to reduce GPU→CPU syncs - Preserve input dtype (float32→complex64) instead of forcing float64

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

- Remove unused xp parameter from coherency_to_stokes (duck typing works) - Update gpu.py call site that passed xp=cp - Use xp.finfo(norm_sq.dtype).tiny instead of xp.finfo(float).tiny in both compute_m_matrix_eigen and compute_m_matrix_sign_split to avoid upcasting float32 intermediates to float64

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

- Add E741 per-file-ignore for coherency.py and tests (Stokes I notation) - Add docstrings to test classes (D101) - Fix whitespace around operators in f-strings (E226) - Remove f-prefix from strings without placeholders (F541)

- Fix cupy/numpy index mismatch in coords.py when GPU polarized sky path indexes astropy skycoords with cupy above_horizon array - Fix nsrc_alloc mismatch in gpu.py by deriving it from coords.nsrc_alloc (matching CPU pattern) instead of calculating independently, which caused shape errors with source_buffer < 1.0 - Add 5 GPU tests for polarized sky eigendecomposition covering all negative_flux branches (raise, split, ignore, invalid) - Update beam_spline_opts from {"kx": 1, "ky": 1} to {"order": 1} for pyuvdata 3.2.1 compatibility (map_coordinates API) - Change test_cpu_vs_gpu.py import guard from pycuda to cupy

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

…anup tests - Fix citation to Kittiwisit et al. (2025) - Replace negative_flux: str with raise_on_negative_flux: bool - Remove "ignore" option for negative eigenvalues - Make fluxes optional when stokes is provided (derived as stokes[0]) - Make simulate_vis keyword-only for cleaner API - Extract process_polarized_chunk() to core/coherency.py to deduplicate identical CPU/GPU polarized sky loop bodies - Always use compute_m_matrix_sign_split (per Steven's suggestion), check has_neg after instead of branching on three code paths - Simplify getz.py beam slicing: bm[:, :, :1, :] instead of bm[:, :, 0, :][:, :, xp.newaxis, :] - Clarify getz.py exptau mutation comment - Add stokes and raise_on_negative_flux to cpu.py docstring - Fix docstring "intensity map" -> "sky model" - Register benchmark marker and skip benchmarks unless RUN_BENCHMARKS=1 - Parametrize precision in backward compatibility test - Rename TestEigendecompAccuracy -> TestPolarizedSkySanity - Remove redundant fluxes= from test calls where stokes is provided - Update test_matvis_cpu.py to use keyword arguments

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

steven-murray

Thank you @kartikmandar -- I think this is very close now!

I have one more comment, and also we need to make sure we add tests against pyuvsim for the polarized skies. Currently it looks like we have sanity checks, which is good, but nothing testing specific more complicated scenarios.

BTW in the future if you address review comments it is easiest if you leave the comments open rather than resolving (unless they are super duper straight forward) because it means on my review I can check how you implemented specific requests.

steven-murray · 2026-04-30T17:04:17Z

    coord_method = CoordinateRotation._methods[coord_method]

    # Determine if we have a polarized sky model
    polarized_sky = stokes is not None and polarized


I think the logic should be that if stokes is provided, polarized_sky is True. The user shouldn't have to both pass stokes AND polarized_sky=True. That would easily lead to people passing a full stokes array, assuming they will get a polarized sky, but then not getting what they expected.

got it. will do that.

…Handling

Addresses Steven's PR HERA-Team#114 review (2026-04-30) on src/matvis/cpu/cpu.py:189: passing stokes implies polarized=True, so the user shouldn't have to set both. polarized now defaults to None and is inferred from stokes; explicit polarized=False with stokes is rejected with ValueError instead of being silently flipped, so accidental misuse is caught.

Addresses Steven's PR HERA-Team#114 review (2026-04-30) calling for cross-checks against pyuvsim covering more than the existing Stokes-I sanity tests. The new test feeds Q=0.2*I, U=0.3*I, V=0.4*I through the matvis stokes path and pyuvsim, and reuses the existing compare_sims helper. The analytic-beam case passes cleanly at rtol=2e-4. The uvbeam case xfails strictly: matvis stores J[feed, ax=0]=theta-hat (matching pyuvdata's documented Naxes_vec convention and pyradiosky's coherency basis) while pyuvsim explicitly swaps to J[feed, ax=0]=phi-hat (pyuvsim/antenna.py:145-149). Gaussian e-field has identical theta and phi components so the swap is a no-op there; UVBeam doesn't, so the two RIME conventions diverge by O(10%) on diagonals and worse on XY/YX. Resolution requires a coordinated fix across matvis/pyuvsim/ pyradiosky and is out of scope for this PR — flagging via xfail keeps the cross-check on the record. Helper: get_standard_sim_params now swaps fluxes for stokes when use_polarized_sky=True, extracting the (4, Nsrc, Nfreq) Stokes array from sky_model.stokes after at_frequencies() so matvis sees bit- identical Stokes to what pyuvsim uses. Drive-by: corrects max-imag bookkeeping in compare_sims (np.abs(np.max(...)) -> np.max(np.abs(...))). Cosmetic — affects only the printed err message, not the assertion.

Web research turned up specific citations supporting the (theta, phi) convention as the consensus across the radio-astronomy stack: - pyuvdata documents Naxes_vec axis 0 as theta-hat - pyradiosky's coherency rotation is in the theta/phi basis - Kittiwisit et al. 2025 (the paper matvis cites for the eigendecomp M-matrix) explicitly defines C with theta first - matvis's three open polarization-convention issues (HERA-Team#15, HERA-Team#28, HERA-Team#41) - pyuvsim's open issue #395 conceding the antenna-to-Jones path is suspect; pyuvsim issue #196 confirming no analytic ground-truth test exists for polarized output - The pyuvsim swap was added in 2018 (commit 78ca8bc4 'fix tests') with no design memo and validated only against a Stokes-I zenith source — invariant under the swap and unable to detect a sign error The diagnostic itself surfaced the Q/U asymmetry: Q-only disagrees ~7-9%, U-only ~12-16%, V-only matches because Stokes V lives in the imaginary off-diagonal which is invariant under a real axis swap.

Integrates ~40 upstream commits (through PR HERA-Team#123) onto the polarized-sky branch. Conflicts resolved in 4 files: - core/getz.py: kept the polarized m_matrix Z-path; adopted upstream's memory-efficient unpolarized branch (PR HERA-Team#120, beam_idx loop instead of fancy indexing) and the reshape-rebind idiom in both branches. - gpu/gpu.py: kept both upstream's memory_buffer param and our stokes/raise_on_negative_flux params. - wrapper.py: kept both upstream's coord_method/coord_method_params/ matprod_method params and our stokes/raise_on_negative_flux params, in the signature, docstring, and the per-frequency backend call. - setup.cfg: accepted upstream's deletion (config migrated to pyproject.toml). Ported our additions into pyproject.toml: the `benchmark` pytest marker, and the E741 lint ignores for coherency.py and tests (flake8 -> ruff migration; E741 is in ruff's selected "E" group and our Stokes I/Q/U/V variable names trigger it).

steven-murray · 2026-06-11T20:38:15Z

Thanks @kartikmandar. Looks like just a couple more tests are needed.

Copilot AI review requested due to automatic review settings April 4, 2026 18:31

Copilot started reviewing on behalf of kartikmandar April 4, 2026 18:31 View session

Copilot AI reviewed Apr 4, 2026

View reviewed changes

kartikmandar requested a review from Copilot April 4, 2026 21:46

Copilot started reviewing on behalf of kartikmandar April 4, 2026 21:46 View session

Copilot AI reviewed Apr 4, 2026

View reviewed changes

Comment thread tests/test_polarized_benchmark.py

Comment thread src/matvis/core/coherency.py Outdated

Comment thread src/matvis/core/coherency.py Outdated

Comment thread src/matvis/core/coherency.py Outdated

steven-murray reviewed Apr 7, 2026

View reviewed changes

kartikmandar and others added 16 commits April 16, 2026 21:03

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

2df3964

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

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

f29dcb9

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

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

dc6e285

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

style: fix flake8 lint errors from pre-commit.ci

3cec78e

- Add E741 per-file-ignore for coherency.py and tests (Stokes I notation) - Add docstrings to test classes (D101) - Fix whitespace around operators in f-strings (E226) - Remove f-prefix from strings without placeholders (F541)

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

2d0f31d

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

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

f72f2f3

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

docs: clarify the two-pass exptau-sharing rule in getz.py

0b90781

docs: generalize I_sky description and add sky-model modes Notes

6100663

refactor: hoist physicality check to simulate-start

ca23bf2

feat: source-partition sign-split for single-pass eigen

3465497

refactor: wrapper API — one of fluxes/stokes, context-aware default

29ec38f

steven-murray reviewed Apr 30, 2026

View reviewed changes

kartikmandar added 5 commits May 4, 2026 09:47

refactor: parametrize negative-flux tests, rename to TestNegativeFlux…

e3d3bdd

…Handling

kartikmandar added 2 commits June 11, 2026 18:25

build: restore benchmark marker and E741 ignores after upstream merge

98e3e5a

style: fix ruff lint and format findings in polarized-sky code

ce58643

kartikmandar force-pushed the feat/polarized-sky-eigendecomp branch from e1a7eee to ce58643 Compare June 11, 2026 12:56

Merge branch 'HERA-Team:main' into feat/polarized-sky-eigendecomp

b8074c9

kartikmandar requested a review from steven-murray June 11, 2026 14:26

steven-murray approved these changes Jun 11, 2026

View reviewed changes

Conversation

kartikmandar commented Apr 4, 2026

Summary

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

steven-murray left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

steven-murray left a comment

Choose a reason for hiding this comment

Uh oh!

steven-murray Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

kartikmandar May 3, 2026

Choose a reason for hiding this comment

Uh oh!

steven-murray commented Jun 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov Bot commented Apr 6, 2026 •

edited

Loading