Ngmix v2.0 (CI mirror of #740)

[cailmdaley]

What this is

A same-repo mirror of #740 (@martinkilbinger's "Ngmix v2.0"), pushed to a branch on CosmoStat/shapepipe so that CI actually runs. All 57 commits are authored by Martin (and carry Lucy Baumont's and Axel Guinot's work) — pushing the branch preserves that authorship unchanged; the only thing this PR adds is a same-repo head so GitHub Actions fires the pull_request workflow without the fork-PR approval gate. #740 received no CI runs at all for this reason.

Substance is identical to #740 — see that PR for the full description. In short: upgrade ngmix to 2.4.0 and adopt Lucy's new ngmix classes/interface; overhaul the shape-measurement module; centroid-bias fix + validation; v2.0 production-run plumbing.

Going forward, opening PRs directly on CosmoStat/shapepipe (rather than from a fork) avoids this — fork PRs don't trigger our Docker-image CI without a maintainer approval that wasn't happening.

Closes/supersedes #740 once CI is green (leaving that call to Martin).

Review

A detailed review is on its way (read against Martin's checklist plus a science-quality pass). Headline from exercising the new fitter against real ngmix 2.4.0 on candide: the metacal path runs end-to-end and shear recovery is unbiased at the few×10⁻⁴ level in m. Full notes to follow.

— Claude on behalf of Cail

[cailmdaley]

CI is green — we believe this PR is ready to merge. ✅

fast-tests, build-test-publish, and the docs deploy all pass on e57544be.

A final pass since the last review round:

Cleanups & review fixes

• Renamed the moments_fail column → mcal_types_fail / NGMIX_MCAL_TYPES_FAIL through ngmix + make_cat. The v1 name counted get_guess (moments-guess) failures that no longer exist in the v2 fitter; it now counts the metacal types with nonzero fit flags, so the name now matches the meaning.
• A /code-review pass surfaced and fixed: a latent NameError in make_cat._save_ngmix_data(moments=True) (two arrays were read only in the non-moments branch but used unconditionally); a leftover debug print loop in make_cat (plus two stray prints in ngmix folded into a ValueError, and a missing f-string prefix); clear ValueErrors in the mccd_interp / read_ext_sexcat / psfex_interp / vignetmaker runners when a MULTI-EPOCH/post-process config omits the positional WCS / exp-numbers input (previously a cryptic IndexError — no behaviour change for valid configs); and a divergence guard tying compile_results' metacal-type list to METACAL_TYPES.

CI fix (fast-tests was red, unrelated to the science changes)

• pyproject.toml had no [build-system] table, so uv resolved shapepipe as a virtual project (uv.lock recorded source = { virtual = "." }) and uv sync installed only the dependencies — every test module failed collection with No module named shapepipe. Declared the setuptools backend (lockfile now editable).
• test_mpirun_launch_imports_mpi now skips when the MPI runtime library is absent (the bare fast-tests runner; the dev-image CI still exercises the real MPI path), matching the workflow's self-skip pattern for system-dependent tests. run.py's launcher gate is unchanged and correct.

Deferred, intentionally out of scope for this PR

• The PSF-fit prior-domination and g_psfo/Tpsf provenance issue is tracked separately in [BUG] ngmix module passes the galaxy shape prior to the PSF fit, which is prior-dominated #749 (with a reproduction). Metacal shear is unaffected (ngmix deconvolves by the actual PSF image); only the PSF catalogue columns are, and those aren't currently ingested downstream.
• Non-blocking follow-ups noted during review, not done here: the make_post_process TILE_ID metadata-namespace special-case, a couple of per-object efficiency hot-spots in ngmix.py (per-epoch FSCALE header re-parse, per-object make_runners), and some numpydoc/PEP8 doc gaps on the new ngmix helpers.

— Claude, on behalf of Cail

---

Generated by Claude Code

[aguinot]

Hello again Claude/Cail,

I am sorry to insist but the changes to Tpsf are wrong and not necessary. I totally understand that you want to keep the psf prior in a separate issue but this change is introducing a bug and make it "impossible" to properly compute the selection response. I really don't understand why you are introducing it.
I will also answer to #749 with an example but as I suspected, the prior is not an issue.

[cailmdaley]

Update — consolidated to a single in-image CI path.

Following a quick discussion: rather than keep the separate fast-tests job (which ran a subset in a bare uv venv, with the binary/MPI tests only self-skipping there), I've removed fast-tests.yml. This realigns with the test-suite constitution — "the image is the environment… there is no second test workflow" — and eliminates the environment divergence. deploy-image.yml's in-container pytest -rX is a strict superset of what fast-tests ran, so no coverage is lost; it remains the gate.

Kept from the CI work: the [build-system] declaration (correct independent of CI — it makes uv sync / pip install -e . install the package, per the documented dev loop), and the MPI-launcher test's skip-when-unavailable guard (harmless in-image, where MPI is present).

All checks green on 216f407a; PR is mergeable_state: clean.

— Claude, on behalf of Cail

---

Generated by Claude Code

[martinkilbinger]

@cailmdaley we should maybe propagate both original and reconvolved PSF quantities Tpsf and Tpsf_o, following @aguinot's comment.

[cailmdaley]

Thank you @aguinot for insisting! Sorry neither I nor Claude saw your initial comment before posting that "we believe this is ready to merge" comment. Will fix and take a look at your comment on #749. cheers

[cailmdaley]

Okay @aguinot I dug into the history of this and will try to express my understanding below for the sake of educating myself, tell me if this sounds right to you.

Before this branch was opened both the original and reconvolved PSFs were fit, with the original psf_obs fit quantities appropriately assigned to *o keys:

shapepipe/src/shapepipe/modules/ngmix_package/ngmix.py

Line 754 in 0cfcf4f

psf_res = make_galsimfit(psf_obs, "gauss", psf_guess)

shapepipe/src/shapepipe/modules/ngmix_package/ngmix.py

Line 798 in 0cfcf4f

psf_res_gT["T_PSFo"] += psf_res["T"] * w_tmp

and the reconvolved obs.psf fit quantities appropriately assigned to keys with no o:

shapepipe/src/shapepipe/modules/ngmix_package/ngmix.py

Lines 888 to 896 in 0cfcf4f

	psf_res = make_galsimfit(
	obs.psf,
	psf_model,
	np.array([0.0, 0.0, 0.0, 0.0, Tguess, 1.0]),
	)
	except Exception:
	continue
	g1, g2 = psf_res["g"]
	T = psf_res["T"]

shapepipe/src/shapepipe/modules/ngmix_package/ngmix.py

Lines 907 to 908 in 0cfcf4f

	tres["gpsf"] = gpsf_sum / wsum
	tres["Tpsf"] = Tpsf_sum / wsum

(here eventually psf_res_gT and tres get merged into a single dictionary).

---

Then, in d01503d the switch to ngmix 2.4 happened and make_galsim_fit went away, replaced by a MetacalBoostraper that only fits the reconvolved PSF:

shapepipe/src/shapepipe/modules/ngmix_package/ngmix.py

Lines 1187 to 1189 in 60962ff

	resdict, obsdict = boot.go(gal_obs_list)
	# compile results to include psf information
	psf_res = average_multiepoch_psf(obsdict)

and in 60962ff, the original *o keys that were no longer being calculated got assigned the reconvolved quantities:

shapepipe/src/shapepipe/modules/ngmix_package/ngmix.py

Line 688 in 60962ff

res['T_PSFo'] = psf_res['T_psf']

Thus, both the *psf and psfo keys carried the same reconvolved PSF, and I agree that this is very dangerous. It already caused a lot of confusion for example below, where two T_psfo keys were defined:

shapepipe/src/shapepipe/modules/ngmix_package/ngmix.py

Line 410 in 60962ff

output_dict[name]["T_psfo_ngmix"].append(results[idx]["T_PSFo"])

shapepipe/src/shapepipe/modules/ngmix_package/ngmix.py

Line 422 in 60962ff

output_dict[name]["Tpsf"].append(results[idx]["T_PSFo"])

which I then deduplicated in e57544b leaving the clobbered key as you noticed above. This didn't change any behavior, but it was a shortsighted change that didn't address the underlying issue and compounded confusion.

---

Anyway, it sounds like the solution is to correctly record both the original and reconvolved PSF as you suggest, and clean up these original-reconvolved clobbers. This would probably mean adding back an ngmix fit to the original PSF in do_ngmix_metacal, which would then be assigned to the epoch-averaged NGMIX_ELL_PSFo in the output catalog alongside the reconvolved NGMIX_Tpsf.

However I'm a bit confused because there are also per-epoch moment fits to the original PSF e.g. E1_PSF_HSM, SIGMA_PSF_HSM which could in theory be inverse-variance averaged in the same way as the ngmix quantities. This creates some confusion about which product to use for what, e.g. the alpha-leakage calculation in sp_validation uses the NGMIX_ELL_PSFo whereas the rho/tau calculation uses E1_PSF_HSM. Is there any reason to use different fitting methods for the alpha vs. rho/tau stats?

Here's a concrete proposal, which may be stupid: make available inverse-variance averaged moment *_HSM fits for the original PSF, and *_NGMIX fits to both the original and reconvolved PSF, all with consistent naming and size definitions. This would give us 3/4 of the 2x2 matrix of original/reconvolved x HSM/NGMIX. Maybe this is overkill, please tell me what you think! Tagging @martinkilbinger @fabianhervaspeters @sachaguer since this touches lots of stuff.

— Cail (written by himself!)

[aguinot]

So I got confused by looking at the diff on github with this line:

output_dict[name]["Tpsf"].append(results[idx]["T_PSFo"])

Which is technically fine because before, here, you do:

res['T_PSFo'] = psf_res['T_psf']

Which is extemely misleading to say the least...
So the situation is not as bad as I thought given that you store the most important quantity which is the reconvolved PSF. Now I would strongly encourage you to make the code more "human" readable and avoid this kind of misleading naming. And of course also store the fit of the original PSF.

Now, regaring your other questions. By computing the original PSF here we get the "coadded" PSF at the location of every detections in the catalogue, stars AND galaxies. The other PSF measurement you are refering too comes from the the star ONLY catalogue. So it cannot be used to compute the leakage at galaxy location, used for your object-wise PSF correction if I am not mistaken. In addition, the per-epoch with HSM is done in image coordinate while NGMix works in sky coordinates, which is a lot better (see my note at the end). If we wanted to "coadd" those PSF we would need to read the weight image at the star locations to get the weight values. Technically possible but a bookeeping nightmare and not useful. Probably require a specific module to do that, with matching, etc.
Summary:

• NGMix PSF: all detections "coadded" PSF in sky coordinate with the same fitter as for the shape measurement
• per-epoch PSF: stars only single epoch PSF in image coordinate with adaptive moments (close-ish to gaussian fitting)

I am sorry but don't understand your last proposal, I am not sure what usage we would have of this.

NOTE about HSM:
It is now possible to return the ellipticity in sky coordinate with HSM with the keyword: use_sky_coords=True which is False by default. I would strongly encourage you to change this aswell. Probably in a separate PR. (see documentation)

[cailmdaley]

Yes extremely confusing/misleading, it took me quite a while to trace this. I totally agree there should be a single source of truth for each quantity that is consistent with its column name and docstrings, although I'll note that all the non-human-readable elements were authored by humans :)

I actually just discovered that the HSM moment fit is already calculated and saved at galaxy positions:

shapepipe/src/shapepipe/modules/psfex_interp_package/psfex_interp.py

Line 312 in 216f407

Compute shapes of PSF at galaxy positions using HSM.

shapepipe/src/shapepipe/modules/make_cat_package/make_cat.py

Lines 631 to 635 in 216f407

	e_psf = (
	gpc_data["SHAPES"]["E1_PSF_HSM"],
	gpc_data["SHAPES"]["E2_PSF_HSM"],
	)
	self._add2dict(f"PSF_ELL_{epoch + 1}", e_psf, idx)

so if I'm understanding well, with the change to sky-coordinate HSM that you suggested one could get a per-object HSM fit by weighting the per-epoch fits in the same way that they're weighted for the ngmix fits without too much plumbing. But these quantities aren't used anywhere and in fact are discarded in sp_validation, so if it's not useful then let's not waste time on it!

So I will leave the HSM fits at galaxy positions alone (unless there are performance reasons to turn off the route altogether, since it doesn't seem to be used anywhere), clean up the original/convolved ngmix PSF keying, and open a separate PR for switching HSM fits to sky coordinates. Tell me if I'm still missing anything and I'll tag you when I have an implementation. Cheers

[fabianhervaspeters]

Hello, in the previous ngmix version the way we dealt with neighbors was to fill in a noise image were SExtractor gave the segmentation of a neighbor.

shapepipe/src/shapepipe/modules/ngmix_package/ngmix.py

Line 1058 in fc656da

gal_masked[weight_map == 0] = noise_img_gal[weight_map == 0]

now we do not know if this strategy is perfect (for one thing it could vary a lot between exposures what the "neighbor" is) and we do have the undiagnosed B-modes when including FLAGS=2 galaxies, yet the idea was to propagate the same strategy in the new version for now, both because we lack a better alternative (since we are not going to implement uberseg) and because we fundamentally believe this should work. I discussed this with Eric Huff and he said that in Euclid the approach was to do a noise interpolation, so a very similar approach. The function get_noise is in ngmix.py but never used. Please let me know if I missed it. Maybe we could build a mask in sky coordinates to make sure all exposures have the same area masked and then apply a 1 or 2 pixel broadening? Since discarding FLAGS=2 objects lead to a strong selection not-metaclibratable selection bias I think this deserves some thought.

EDIT: Sorry about that, I think this is okay for now, it actually appears as:

if (~mask).any():
gal_masked[~mask] = noise_img_gal[~mask]

line

[aguinot]

@cailmdaley propagating to the coadd the single epoch HSM quantites will be a lot more difficult than you think and I don't see a good reason to do so.

[aguinot]

@fabianhervaspeters I don't know why the get_noise function is not used anymore. It should give potantially more accurate results than sigma_mad used here. The idea of get_noise is to mask the central object to estimate the noise. The stamps are only 51x51 pixels so in the case where there is a large object I wouldn't be surprise that sigma_mad is biased. But it will probably not make a big difference.
Regarding your idea of having a common mask between exposures, this is already the case. The segmentation map used is the one from the coadd. I wanted to make it slightly larger but given that the SNR is higher on the coadd I thought it would already enlarge the mask enough (maybe not..).

All that being said, I don't think we should mask and replace by noise neighboring objects on the image. A better startegy would be:

• replace by noise masked pixels (cosmic rays) as they will create a discontinuity when going to fourier space (metacal process). Replacing by noise is probably just as good as interpolation but likely faster (is the noise estimation is not to bad). AND set weights to 0.
• Leave the neighboring objects on the image as they are not an issue for metacal.
• Apply 0 weight for the neighbor and do uberseg.
• Probably do a symmetrization of the weights

For uberseg, it is really not difficult to implement. Here is a piece of code I made for another project to branch to the implementation of Erin Sheldon used in DES:

import numpy as np

from meds import _uberseg

def fast_uberseg(
    seg,
    weight,
    object_number,
):
    obj_inds = np.where(seg != 0)

    Nx, Ny = seg.shape
    Ninds = len(obj_inds[0])
    seg = seg.astype(np.int32)
    weight = weight.astype(np.float32, copy=False)
    obj_inds_x = obj_inds[0].astype(np.int32, copy=False)
    obj_inds_y = obj_inds[1].astype(np.int32, copy=False)
    _uberseg.uberseg_tree(
        seg, weight, Nx, Ny, object_number, obj_inds_x, obj_inds_y, Ninds
    )

    return weight

For the symmetrization, it is implemented here.
I hope this can help.

(This is obvious but everything above need testing)