pol gradient fixes to dev#306
Merged
Merged
Conversation
Maps the Stokes DOF mask to the physical gradout slots the chisq/reg kernels must fill, centralizing the mcv/vcv cross-coupling that mirrors transform_gradients' Jacobian sparsity. Not yet wired in. + unit tests.
Feed the cross-coupling-aware mask to the pol gradient kernels in both compute_chisqgrad_dict and compute_reggrad_dict. Behavior-identical for now (kernels still carry the or-patches). Guard physical_grad_slots against sub-4-wide single-pol masks (Stokes-I carries 'mcv' inertly). + regression test.
The mcv/vcv cross-coupling now lives in physical_grad_slots, so drop the 'or pol_solve[3]' patches (#296) in chisqgrad_vvis / chisqgrad_vvis_nfft; each physical slot keys on its own bit again. Note in each pol chisqgrad docstring that pol_solve flags required physical gradients, not DOFs.
Zero the back-neighbor (m2/m3) terms on the first row/column in reggrad_ptv slots 0/1/3 (the back-neighbor is the zero pad), matching reggrad_vtv/reggrad_tv. Pre-fix the whole first row+col of those slots was wrong (corner ~4x off vs FD). Add epsilon_tv to reg_ptv/reggrad_ptv denominators (default 0, byte-identical) for #295 parity. Note pol_solve = physical-gradient slots in the 8 pol reggrad docstrings. Add full-grid boundary FD regression tests for ptv, vtv, and Stokes-I tv.
polchisqgrad is a legacy shim (parity tests only); document that its pol_solve is a physical-gradient mask, not a raw DOF mask.
_pol_solve_for now returns [1,1,1,1] so the previously-blind cross- coupling slots are FD-checked: reggrad_ptv psi (3), reggrad_vflux/l1v/ l2v/vtv rho (1), and slot 0 for every pol reg. Proves the reg-grad slots are individually correct against finite differences.
New pol coverage in its final-home file: TestPolChisqGradFD checks
chisqgrad_p/m/vvis against finite differences of the chisq value in all
four physical slots (pol_solve=[1,1,1,1]) for direct+nfft, asserting
vvis slot 2 (EVPA) is identically zero. TestPolChisq{,Grad}Consistency
check direct-vs-nfft agreement. Closes the m / p-slot-3 blind spots.
TestObjectiveGradPolarimetricFD checks objgrad vs FD for IP/IV/IQUV x
{direct,nfft}, with each case bundling its pol data terms + a pol reg, so
both the chisq and reg gradient paths through physical_grad_slots are
exercised. Samples the pol DOF block (past the Stokes-I block), where the
mcv/vcv cross-coupling lives -- the existing global-sampling FD tests
missed it (the dropped IP slot-3 term is ~4% off FD at V=0.02*I, ~430% at
V=0.2*I). Comments out the now-subsumed test_fd_matches_analytic_polarimetric
(backend) and test_iv_gradient_matches_finite_difference (e2e).
Add make_asym_image (broad offset/elongated/rotated double-Gaussian) and switch the Stokes-I FD fixtures (chisq_setup, reg_setup, mfreg_setup, grad_setup) to it. Breaking the reflection/rotation/x<->y symmetry of the centered Gaussian surfaces boundary/axis-ordering bugs a symmetric image hides. Blobs kept broad (grid-filling, no dead pixels) so the |.|-kink TV gradients stay FD-well-conditioned at epsilon_tv=0; all tolerances unchanged.
chisq_setup_pol and a new asym_pol_setup build on make_asym_image and use add_random_pol (ccorr>0) so EVPA, vfrac, rho, and psi all vary spatially instead of a constant pol fraction. polreg_setup switches its Stokes I to the asymmetric image (keeping the per-pixel pol jitter that keeps TV denominators non-degenerate). TestObjectiveGradPolarimetricFD now uses the structured-pol obs. Widen the pol chisq FD check to a median+max split (median 1e-5, max 1e-3): the structured-pol imcur has sharper local curvature, so 2nd-order FD truncation pushes a few small-gradient pixels to ~2.6e-4 -- well below any real pol-gradient bug (%-level), which the tight median still catches.
Manual review pass: per-slot dR/dX labels, docstrings on the reg kernels, a module-level CONVENTIONS block (imarr = [I, rho, phi=2chi, psi]), and removal of stale TODOs. Two behavior touches, both byte-identical at the defaults: - reg_vtv / reggrad_vtv now honor epsilon_tv (kwargs, default 0) like the ptv pair, instead of the value ignoring it while the grad pinned it to 0. - reggrad_ptv masks the chi-slot back-neighbor terms (c2/c3) too, for uniformity (they already self-zero at the pad). Plus an mcv_r exception-message fix and ruff-clean whitespace.
Manual review pass: docstrings on the Stokes-I reg kernels, per-block comments, 'fourier/transform matrices' labels on the diag Amatrices unpacking, and removal of dead commented-out systematic-noise code in the bispectrum data functions (the intent is now documented in apply_systematic_noise_snrcut). Purely cosmetic; ruff-clean.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## dev #306 +/- ##
==========================================
+ Coverage 47.42% 47.50% +0.07%
==========================================
Files 54 54
Lines 26899 26928 +29
Branches 4588 4591 +3
==========================================
+ Hits 12758 12793 +35
+ Misses 12664 12648 -16
- Partials 1477 1487 +10 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
achael
changed the title
rohandahale
approved these changes
Jun 16, 2026
rohandahale
left a comment
rohandahale
left a comment
Collaborator
There was a problem hiding this comment.
I reviewed the logic and ran the gradient suites locally, looks great! Good to merge.
This was referenced Jun 16, 2026
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What & why
Stops overloading
pol_solve, which conflated two distinct masks: (1) whichsolver DOFs are optimized, and (2) which physical gradient slots (I, rho, phi,
psi) the chisq/reg kernels must fill. These differ wherever a value-transform's
Jacobian is non-diagonal (mcv, vcv): one solved pol DOF drives BOTH rho and psi.
Introduces
physical_grad_slots(pol_solve, transforms)as the single source oftruth, wired into both the chisq and regularizer gradient dicts, and reverts the
kernels to clean diagonal gating.
Correctness fixes
physical_grad_slots. This fixesIP imaging when V!=0:
mcv_grad's m' gradient needs the psiphysical gradient, which the old mask zeroed.
reggrad_ptvfirst-row/col boundary masking +epsilon_tvinreg_ptv/reggrad_ptvandreg_vtv/reggrad_vtv(default 0, byte-identical).Test coverage
physical_grad_slotsunit tests; pol chisq FD all-slots (pvis/m/vvisxdirect/nfft); pol-reg FD driven with all four slots; boundary FD regressions
for
ptv/vtv/tv; one parametrized objective-FD over {IP/mcv, IV/vcv,IQUV/polcv} x {direct,nfft} that samples the polarization DOF block (the
existing global-sampling tests missed the m'-coupling bug).
spatially-varying pol (
add_random_pol, ccorr>0) so EVPA/vfrac/rho/psi allvary -- surfaces boundary/axis bugs a centered constant-fraction Gaussian hides.
Cosmetic
imager_utils.pyandpol_imager_utils.py(no behavior change), ruff-clean.
Supersedes part of #295
The
reggrad_ptvboundary +epsilon_tvfixes and the reg-grad cross-couplingland here; #295 should drop those and stay jax-agnostic plumbing (Option B).
Testing
Full suite: 1776 passed, 2 skipped, 1 xfailed.