Add the Anscombe variance-stabilizing transform and its inverse - #228
Merged
Conversation
Anscombe's transform, 2*sqrt(c + 3/8), maps Poisson counts to a variable with approximately unit variance regardless of rate, so downstream steps that assume homoscedastic Gaussian noise apply to count data. The inverse is a separate Transformer/Unit pair with three methods. The default, EXACT, is Makitalo & Foi's closed-form approximation of the exact unbiased inverse; ASYMPTOTIC is (y/2)^2 - 1/8; ALGEBRAIC is (y/2)^2 - 3/8, the strict functional inverse, which round-trips exactly but underestimates the rate at low counts. EXACT clamps its input at 2*sqrt(3/8) -- the forward transform's value at zero counts -- because the D^-3 term diverges below that. The closed form is exactly 0 at the clamp, so the output floors at zero counts rather than stepping. All of it is array-API generic and covered against MLX.
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
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.
Adds
ezmsg.sigproc.math.anscombe: the forward transform2*sqrt(c + 3/8)and, as a separate Transformer/Unit pair, its inverse.Anscombe's transform maps Poisson counts — spike counts, photon counts — to a variable with approximately unit variance regardless of the underlying rate, which lets downstream steps that assume homoscedastic Gaussian noise be applied to count data. The forward direction is settings-free, like
AbsandInvert. Inputs below-3/8produce NaN, so it expects non-negative counts; there is no clipping guard, on the theory that silently rescuing negative "counts" hides an upstream mistake.The inverse
InverseAnscombeSettings.methodpicks one of three, as an enum or its string value:EXACT(default) — Makitalo & Foi (2011), the closed-form approximation of the exact unbiased inverse.ASYMPTOTIC—(y/2)^2 - 1/8. Unbiased as the rate grows, noticeably biased below ~5 counts.ALGEBRAIC—(y/2)^2 - 3/8. The strict functional inverse, so it round-trips exactly, but it underestimates the rate at low counts.Recovered rate from
inverse(E[forward(z)]), 400k Poisson draws:EXACTis the default because it is the only one that holds up at the low rates where the transform is actually interesting.The clamp in EXACT
The closed form's
D^-3term diverges to+infas its input approaches zero, so raw or negative input would produce garbage rather than small counts. It clamps at2*sqrt(3/8)— the value the forward transform produces for zero counts. The expression evaluates to exactly 0 there, so this floors the output at zero counts instead of introducing a step. It'sxp.clip, so it stays on-device for MLX rather than syncing to host.Two things worth a second opinion
The default.
EXACTmeansInverseAnscombe(Anscombe(x)) != xexactly — a round trip is off by the bias correction.ALGEBRAICwould round-trip but is the wrong estimator for the denoising use case that motivates having an inverse at all. I picked correctness-for-the-use-case over round-trip identity; the docstrings say so plainly.InverseMethodimportsOptionsEnumfrom..spectral, matching whatactivation.pydoes. It's a dependency frommath/up to a top-level module, which is a mild layering smell — say the word and I'll inline a local enum instead.A caveat now documented on the enum
These inverses map a denoised stabilized value back to a rate: they invert
rate -> E[anscombe(counts)]. Applying one to still-noisy data returns noisy counts, but the mean of that output is not the mean of the input. My first version of the bias test asserted the wrong property (E[inverse(forward(z))] == rate) and failed by 0.26 at rate 20 before I corrected it — easy mistake, hence the docstring.Testing
tests/unit/test_math.py, 63 passing; full unit suite 4119 passed, 6 skipped. New coverage: forward values over int and float input; a variance-stabilization check (Poisson at rates 5/50/500 → per-channel variance ≈ 1); exact round trip throughALGEBRAIC; str/enum equivalence andValueErroron a bad method; the closed form against hand-evaluated reference values; the zero-floor behavior for inputs at, below, and near the clamp; the low-rate bias comparison above; and empty-time plus MLX/NumPy parity for the forward transform and each inverse method.Docs are autosummary-generated and gitignored, so nothing to review there.