feat: kspace admom (gauss only) - #274
Draft
beckermr wants to merge 57 commits into
Draft
Conversation
The idea is to merge the nice features of pgauss with adaptive moments, so we get a shape for metacal and nice matched aperture colors This is not optimized at all.
numba kernels cache grids use the half-plane symmetry folding single pass rather than separate center moment pass. Use cen update approximation that converges at the fixed point lookups to avoid sin/cos every pass Also refactored into new subdir
Internally they are still padded to square. Apodization still works correctly.
The power spectrum of this field can be used to estimate errors even when the noise is correlated.
This is exactly a sqrt(2) for a round gaussian
Measure weighted moments and apply a modified adaptive step that works for a fixed mixture. Only tested for exp so far, and not optimized. Errors still not corrected Add a star type as well, which is similar to a template fit.
If this damping is repeatedlyh failing then it might look like it converged.
Also keep full e covariance for both mixture and gauss
Also more shared public for a deblender
To be consistent with PGauss
This speed increase means that for the deblender we are dominated by number of iters, which is currently higher for the exp adaptive-like updates
This greatly reduces the number of iterations needed, puts it on par with single gauss adaptive moments. References in the code
There were cases in the deblender where we were going outside the lookup table with chi2 < 0
shape_ok = Tgal > 0 and det2(Sgal) > 0
still need to clean up the exp specific names
It is equivalent, just changed so they don't look different at a quick glance
They are equivalent
this is verified to be byte identical to pre-refactor
also save deweighted moment shapes in all cases.
this is tricky due to the strong anti correlation of the flux split with T. a covariance term needed to be calculated.
The full covariance of the per band fluxes. The off diagonal is generated by the shared family response (the band sum noises are independent across bands) and is what color errors require
Unlike model_sandwich this stays calibrated under model mismatch, where the sandwich under-predicts T errors by ~17 percent and flux errors by ~11 percent for dev truth fit with exp However, it costs a factor of 3 for the gauss fit, so pull it out only when you need it.
Using threadpoolctl to force blas onto one core
the scalar half of prep_epoch, computable without any FFTs: dims, pad phase centers, weights and jacobian factors. Used by prep_epoch itself and by device-prep paths that compute the k arrays elsewhere and only need the epoch scalars on the host (the array entries of the returned epoch are None).
memory is leaking when i raise from within jitted functions. create non-raising version of some numba functions to avoid this. If a fix go in we can remove these
The fit uses all bands at once and structure is common. take this into account in S/N
In extremely rare cases the determinant check was within precision of zero and caused an exception. Put in an explicity tolerance based on products
beckermr
commented
Sep 3, 2026
Comment on lines
+11
to
+16
| flux_var_delta evaluates the response analytically for the gaussian | ||
| implied by the converged weight (the model='gauss' fitter and the | ||
| deblender's gauss objects); model_sandwich evaluates the full | ||
| sandwich for the moment matched model families with closed form | ||
| model derivatives, and reduces exactly to flux_var_delta for a | ||
| single gaussian family. |
Collaborator
Author
There was a problem hiding this comment.
This docstring doesn't match the code in either this PR or in #271.
Comment on lines
+28
to
+35
| the covariance-aware total flux s/n sqrt(F^T C^-1 F) from the | ||
| fluxes and their cross-band covariance (the Wald significance | ||
| of the flux vector). The statistic is invariant to per-band | ||
| rescaling, so raw flux sums with the raw covariance equal | ||
| physical fluxes with the physical covariance. Returns None | ||
| when the covariance is not positive definite or the form is | ||
| not finite; the caller falls back to the independent-band | ||
| quadrature sum |
Collaborator
Author
There was a problem hiding this comment.
The docstring here is very odd. Why does it declare what the caller is doing or will do? How could this function know that?
Is there a need to mention the "Wald" significance?
| ] | ||
|
|
||
|
|
||
| def joint_flux_s2n(F, fcov): |
Collaborator
Author
There was a problem hiding this comment.
These parameter names are not great. How about flux and flux_cov?
| return np.sqrt(q) | ||
|
|
||
|
|
||
| def flux_cov_delta(Sigma, sums, cov, fsums, fvars, fmcovs): |
Collaborator
Author
There was a problem hiding this comment.
Typically, Python uses lowercase variable names. Capitalized names are reserved for classes.
Suggested change
| def flux_cov_delta(Sigma, sums, cov, fsums, fvars, fmcovs): | |
| def flux_cov_delta(sigma, sums, cov, fsums, fvars, fmcovs): |
| including the first order response of the adaptive weight to | ||
| the noise (the delta method). | ||
|
|
||
| docs for flux_var_delta (diagonal of flux_cov_delta) |
Collaborator
Author
There was a problem hiding this comment.
I copied in the docstring of flux_var_delta (a function which simply calls this one), in order to try and piece together the exact math and computation happening here.
Comment on lines
+68
to
+99
| The converged weight satisfies the fixed point conditions | ||
| M(Sigma) = Sigma / 2, where M is the measured weighted covariance, | ||
| a ratio of the joint moment sums. By the implicit function | ||
| theorem the weight fluctuation is dSigma = -J^-1 dM with | ||
| J = dM/dSigma - 1/2. For the gaussian implied by the converged | ||
| weight, dM/dSigma = 1/4 exactly at the fixed point, for any | ||
| ellipticity and smoothing, so dSigma = 4 dM. The flux | ||
| normalization (proportional to sqrt(det Sigma)) and the kernel | ||
| response of the flux sum combine to dln F = tr(Sigma^-1 dM), and | ||
| for a single band the fixed weight flux fluctuation cancels | ||
| exactly, leaving | ||
|
|
||
| dF / F = tr(Sigma^-1 dS_M) / S_F | ||
|
|
||
| in terms of the second moment sum fluctuations alone. For a | ||
| round gaussian with matched weight and white noise this doubles | ||
| the fixed weight variance. With multiple bands the band flux sum | ||
| and the joint conditions are distinct and their cross covariances | ||
| enter; the band flux sums covary only with their own epochs' | ||
| contribution to the joint sums. | ||
|
|
||
| docs for flux_cov_delta (this function) | ||
| --------------------------------------- | ||
|
|
||
| The fluctuation dF_b is proportional to | ||
| dS_Fb - r_b dS_F + r_b b . dS_M (see flux_var_delta): the | ||
| joint flux and moment sum responses are shared by all bands, | ||
| and only the band's own flux sum dS_Fb is band-specific, so | ||
| the cross-band covariance assembles in closed form from the | ||
| same scalars as the variances -- the outer products of the | ||
| band shares with the shared response, plus each band's own | ||
| cross terms. No inversions beyond the 2x2 weight |
Collaborator
Author
There was a problem hiding this comment.
The math here is very dense. I realize it is a docstring, not a paper, but you may want to consider showing a bit more work for each step.
The implicit derivative through the matrix equation for a single band, in particular, would be illuminating to spell out just a bit more.
Also, you do not explain what r_b is.
In general, the doc string here could use shorter statements in plainer language so others can understand them better.
Comment on lines
+142
to
+146
| # the diagonal via the original variance expression, keeping | ||
| # the per-band variances bitwise identical to flux_var_delta | ||
| # as it stood before the cross-band assembly | ||
| idx = np.arange(fsums.size) | ||
| fcov[idx, idx] = fvars + r ** 2 * s + 2 * r * g |
Collaborator
Author
There was a problem hiding this comment.
Why is this block of code needed at all? In the original PR (#271), flux_var_delta literally takes the diagonal of the matrix returned by this function. Thus, they will always be bitwise identical, no matter how the diagonal is computed.
Comment on lines
+136
to
+141
| r = fsums / sums[5] | ||
| bcb = b @ cmm @ b | ||
| s = cff + bcb - 2 * (b @ cmf) | ||
| g = fmcovs @ b - fvars | ||
|
|
||
| fcov = np.outer(r, r) * s + np.outer(r, g) + np.outer(g, r) |
Collaborator
Author
There was a problem hiding this comment.
The relationship between the code in this block and the math in the comment is not clear to me at all. Is there a way to unify the notation? Also, there might be a few steps missing?
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.
This PR has only the code for the Gaussian parts and removes the
full_errorsfeature. I am making this PR so I can actually understand what the code is doing.