zdmp

[cwjames1983]

Pull request to allow p(xopt,xrad) to be simultaneously estimated. Contains new features, and also a few bug fixes. Pulling in now to allow Claude to work its magic

[profxj]

This is a review of the full zdmp vs main diff. There's a lot of good restructuring here (the survey rewrite, figures.py split, states.py presets), but the sweep turned up a set of confirmed breakages — mostly stale calls against renamed/deleted code and copy-paste slips in the new likelihood code — that I think need fixing before merge. Line numbers refer to the branch head.

1. Crashes / wrong results in core paths (blocking)

• grid.py:1169 — Grid.update()'s DMhalo branch still calls get_efficiency_from_wlist(DMlist, wlist, wplist, model=...), but the new signature is (wlist, plist, model=..., ...). Any grid.update({'DMhalo': ...}) (MCMC/cube scans with DMhalo free) raises TypeError: got multiple values for argument 'model'. Same stale pattern in zdm/scripts/MCMC/est_pz.py:245 and average_grids.py:257.
• grid.py:436,439 — calc_pdv() now accumulates with += into self.b_fractions/self.w_fractions, but they're only zeroed when None or when beam_b is passed. Every grid.update() therefore adds fractions on top of the old ones: rates ~double after the first update, triple after the second, corrupting every subsequent likelihood. The old code assigned per-slice on j==0.
• grid.py:555 — calc_rates() applies self.sfr *= self.fz in place, so any update path that re-runs calc_rates without set_evolution (e.g. host-DM param updates) applies the z-fraction correction repeatedly, suppressing rates by fz^N.
• states.py:310-311, 378-379 — lRmin/lRmax are swapped in both HoffmannRepeaters26Pn (lRmax=-4.51, lRmin=3.78) and HoffmannRepeaters26Prep (lRmax=-3.92, lRmin=2.47), giving Rmin >> Rmax and a negative repeater constant in repeat_grid.calc_constant.
• states.py:66 — elif case == "HoffmannRepeaters26": never matches (the valid cases are ...Pn/...Prep), so the else branch calls set_chime_scat and silently overwrites the scattering/width parameters that set_fit_params just set for those cases.
• iteration.py:592 — in calc_likelihoods_1D the penalty check is if Pn==0: — Pn is the boolean function argument, the computed probability is probN. The penalty branch is unreachable and an underflowing probN goes to np.log10(0) = -inf.
• iteration.py:1318 — flg_baddm is only assigned in the sigmaDMG==0 and sigmaHalo==0 branch; with Galactic DM uncertainty enabled, if flg_baddm: raises NameError.
• iteration.py:~1270-1441 — calc_likelihoods_2D contains two stitched-together copies of the p(z,DM) computation. The first (~90 lines) is computed and then discarded, except the stray re-interpolation at 1363-1366 overwrites the MW-uncertainty-weighted pvals from 1303-1309, and the orphaned else: at 1367 references izs1/dkzs1 that only exist in the if pdmz: branch (NameError when pdmz=False). Deleting the first copy restores a single code path.
• iteration.py:730-750 — the ptauw block: tomult[:,:inoztaulist] uses a list as a slice bound (TypeError); the result is named tz_tomult but lines 749-750 multiply by undefined zt_tomult; and the interpolation terms reuse dkis1/dktaus1 where dkis2/dktaus2 belong, so weights don't sum to 1. (Same dkis1-for-dkis2 copy-paste in the 2D block at ~1508-1518.)
• iteration.py:154-155 — get_joint_path_zdm_likelihoods calls get_PATH_lls(..., return_all=True), but get_PATH_lls has no such parameter; its early return (empty ilist) also yields a single float that callers unpack as two values.
• survey.py:1186-1192 — adjust_widths: undefined name DM in the smearing computation, self.WIDTHS (capital S) typo, and the chained fancy-index assignment self.WIDTHs[ismear][bad] = ... writes to a temporary copy, so the bad-value floor never applies (NaNs at the **0.5).
• survey.py:1244-1249 — fix_coordinates writes self.frbs[i,'Gb'] = ... on a DataFrame, which creates a junk tuple-named column instead of setting the cell (needs .loc[i,'Gb']), so missing coordinates are never filled; also line 1226 parses Gb strings with hr=True, scaling latitudes by 15x.
• survey.py:1133 — intrinsic widths compute WIDTH**2 - (tscale*TAU)**2 over all FRBs before the TAU != -1 mask is applied, so the -1 "missing" sentinel is squared in as if it were a real 1 ms scattering time, corrupting IWIDTHs for every FRB without a tau measurement.
• optical.py:473, 492 — loudas_model.give_p_mr_mass/give_p_mr_sfr reference undefined n_redshift_bins (should be self.nzbins): NameError on every call.
• optical.py:622-629 — duplicated condition elif self.opstate.AbsPriorMeth==0: (should be ==1) makes the method-1 branch unreachable; if reached it would also hit unbound AbsPrior.
• optical.py:1779, 1791 — SimplekAbsoluteMags array path tests np.isscalar(Abs) (parameter is App) and returns np.log10(ApparentMags) — both undefined names.
• optical_numerics.py:~304 — in calc_path_priors the if imatch is None: ... continue continues the inner survey loop, not the FRB loop, so an FRB found in no survey falls through: nfitted is still incremented and s.DMEGs[imatch] is evaluated with imatch=None (or the previous FRB's stale survey), corrupting the priors instead of skipping.
• MCMC.py:177 — the try/except ValueError → -inf guard around the likelihood evaluation is commented out, so any math error in an impossible parameter region now kills the whole emcee run instead of returning -inf (the comment above still says it should).
• grid.py:958, 990 — GenMCFRB: lowest-z-bin sample uses self.dz/0.5 (= 2·dz; should be /2.), stretching low-z MC redshifts ~2.5x; top-DM-bin branch uses dDM (a cumulative-probability increment) where the DM bin width self.ddm is intended.
• grid.py:209 — get_dm_coeffs clamps only the lower bound; an FRB with DMEG ≥ dmvals[-1] gives idms2 == ndm and an IndexError in calc_likelihoods_1D/2D (get_z_coeffs has the top-bin clamp; this one needs it too).
• pcosmic.py:266 — get_mean_DM calls exit() (killing the interpreter) when DMbar[0]/DMbar[1] deviates from 1/3 by > 0.01, but the deviation scales as ~0.09·dz, so legitimately coarse z-grids terminate the whole process (e.g. a batch MCMC job) with no catchable exception. Raise a ValueError instead — and the tolerance probably wants to scale with dz.

2. References to deleted code (decide: restore or delete the callers)

• setup.cfg:82-83 — console scripts zdm_build_cube and zdm_pzdm point at zdm/scripts/build_cube.py (deleted in this PR) and zdm/scripts/pzdm.py (doesn't exist); the new CLAUDE.md (lines 91-92) documents both.
• zdm/analyze_cube.py:42-43 — calls iteration.set_orders/set_cube_shape, deleted with the cube machinery; breaks all cube post-processing (slurp_cube and the papers/ scripts that use it).
• zdm/cube.py:155/167/187 — still invokes it.cube_likelihoods and process_pfile, both deleted.
• zdm/scripts/run_slice_benchmark.py:11 — from zdm import MCMC2, deleted on this branch (same dead import in tests/deprecated_test_update.py:10); line 60 also passes sig= to get_log_likelihood, which has no such kwarg.
• zdm/MC_sample/mc.py:67 — new file calls misc_functions.plot_grid_2, deleted in this same PR (figures.py only has plot_grid).
• zdm/tests/test_scat_methods.py:110 — calls survey.geometric_lognormals, deleted from survey.py; the test suite fails.
• zdm/misc_functions.py:770 — width_test still calls the removed module-level survey.make_widths(s, g.state) (now a Survey method with a different signature).

3. Silent behavior changes worth an explicit decision

• survey.py:208 — width-method codes were renumbered: old Wmethod=3 (single-FRB exact width) is now 4, and 3 is a new z-dependent mode. Existing configs carrying 3 silently change meaning. Suggest keeping the old numbering or erroring on the ambiguous value.
• grid.py:130 — wdist changed from a bool to a None-sentinel test while the docstring still says bool; callers passing wdist=False now get the enabled branch. Rename the parameter or test explicitly.
• survey.py:606 — beam value B == -1 silently gets uniform weights across all bins (adjacent comment says "this is incorrect"). Prefer NaN/None + an explicit mask so unknown beam values fail loudly rather than flowing into PATH weights.
• survey.py:1376 — init_halo_coeffs flips the sign of the 116.72 Yamasaki–Totani coefficient with no comment — intentional?

4. Performance notes (non-blocking, but they're in the MCMC hot path)

• optical_numerics.run_path re-reads FRB JSON, the PATH candidate CSV, and load_std_priors() per FRB per posterior evaluation; calc_path_priors also rebuilds op.Field() each call. All parameter-independent — cache once keyed by TNS name.
• grid.calc_rates → construct_fz np.loads two .npy files and rebuilds an interp1d per grid construction; survey.process_dm_mask re-loads the DMMASK .npy per reinit. Both cacheable.
• grid.smear_z convolves per-DM-column in a Python loop; scipy.ndimage.convolve1d(..., axis=0) (already imported) does it in one vectorized call.
• calc_pdv runs a full-grid np.isnan scan per (beam, width) pair — gate behind a debug flag or check the final sum once.

5. Duplication / cleanup (non-blocking)

• figures.py:633 get_alevels is a verbatim copy of find_Alevels 26 lines below in the same file.
• optical.py Simple{k}{Apparent,Absolute}Mags are four copies of one distance-modulus computation (the broken SimplekAbsoluteMags above is a symptom); one helper with k=0 wrappers would remove ~120 lines and the latent crash.
• Four near-identical bin/weight interpolation helpers (survey.get_internal_coeffs, survey.get_w_coeffs, grid.get_dm_coeffs, grid.get_z_coeffs) with subtly different edge handling — one shared helper would make the convention testable once.
• MCMC.calc_log_posterior carries dead code in the emcee hot path: unused wrappers/ident/sident, an unreachable debug=False block, and a shadowing import multiprocessing re-import.

---

Review prepared with Claude Code: an 8-agent scan of the full diff, with findings verified against the branch and the highest-impact items manually spot-checked.

[cwjames1983]

I have addressed the above. I only note things I have not responded to:

survey.py:1133 — intrinsic widths compute WIDTH**2 - (tscale*TAU)**2 over all FRBs before the TAU != -1 mask is applied, so the -1 "missing" sentinel is squared in as if it were a real 1 ms scattering time, corrupting IWIDTHs for every FRB without a tau measurement.
This is fine - the code requires an "OKTAU' value which identifies which TAU measurements to use. No matter what we do here, some nonsense will occur if these tau values are ever used. Potentially, I could set nans?

MCMC.py:177 — the try/except ValueError → -inf guard around the likelihood evaluation is commented out, so any math error in an impossible parameter region now kills the whole emcee run instead of returning -inf (the comment above still says it should).
I want this to happen. Currently, the code should itself identify impossible regions of the parameter space, and assign large likelihood penalties to that, rather than producing nans. If a nan is produced, I want the program to crash - it's a sign of something more sinister going on. Having a try-except at this stage causes the MCMC to continue, potentially allowing bugs to propagate.

I can't for the life of me work out what:
zdm/tests/test_scat_methods.py:110 — calls survey.geometric_lognormals, deleted from survey.py; the test suite fails.
did. So I'm just commenting this out from the test.

zdm/misc_functions.py:770 — width_test still calls the removed module-level survey.make_widths(s, g.state) (now a Survey method with a different signature).
is entirely outdated - I'm deleting this function

[cwjames1983]

Remaining questions are:

• delete grid.update. It's so deprecated now, it's dangerous to have. Agreed? I think getting Claude to re-write this would be better - it's actually got useful functionality, because when running a slice, it'd be nice to have things go faster, but in reality, I think too many things are now correlated and dependent that this is too much of a danger to exist. Then I can delete tests/deprecated_test_update.py, and run_slice_benchmark.py

• Am I able to just delete the lines:

[options.entry_points]
console_scripts =
#
    zdm_build_cube = zdm.scripts.build_cube:run
    zdm_pzdm = zdm.scripts.pzdm:run

from setup.cfg?

Can we just delete zdm/cube.py and zdm/analyze_cube.py?

[cwjames1983]

On silent changes:
Would be good to find all Wmethod=3 code, and check it has correct behaviour

• survey.py:606 — beam value B == -1 silently gets uniform weights across all bins (adjacent comment says "this is incorrect"). Prefer NaN/None + an explicit mask so unknown beam values fail loudly rather than flowing into PATH weights.
  Turns out behaviour is correct - I have since thought about it more. Removed statements saying behaviour is incorrect

survey.py:1376 — init_halo_coeffs flips the sign of the 116.72 Yamasaki–Totani coefficient with no comment — intentional?
I've checked, -116.72 is correct, must have been a typo that Jordan fixed

[cwjames1983]

Regarding duplication / performance: I'd like to ask Claude to implement these, if possible, as a test. Or otherwise, keep thee recommendations around for the future

[profxj]

giddy up

[profxj]

Correction to my review above: the item flagging zdm/tests/test_scat_methods.py:110 as calling the deleted survey.geometric_lognormals is wrong — that line is commented out, and the test passes. Please disregard that one item; the rest of the review stands.

[cwjames1983]

OK, cut extra chaff, will wait for tests to complete, then merge
Now working with Jordan on his changes