Vectorize gamma MLE fit and thread gamma.ppf for gridded transform - #93
Open
benmsanderson wants to merge 2 commits into
Open
Vectorize gamma MLE fit and thread gamma.ppf for gridded transform#93benmsanderson wants to merge 2 commits into
benmsanderson wants to merge 2 commits into
Conversation
Two independent speedups to the 3D gamma quantile mapping used for
precipitation:
1. fit_distribution_parameters_3d
Replaced the per-gridpoint scipy.stats.gamma.fit(x, floc=0) loop with
a vectorized _vectorized_gamma_mle that runs Choi-Wette initial guess
plus Newton iterations on log(k) - psi(k) = log(mean(x)) - mean(log(x))
across all gridpoints in one shot. This is the same closed-form MLE
scipy uses internally, just batched. On synthetic gamma samples the
two agree to 1e-13; on real NorESM2-MM precipitation to ~6e-5
(a few percent of the noise model's own uncertainty).
2. apply_distribution_transform (3D gamma path)
Wrapped scipy.special.gammaincinv (which underlies stats.gamma.ppf
and releases the GIL) in a ThreadPoolExecutor that chunks the spatial
axis. Output is bitwise-identical to stats.gamma.ppf. Defaults to
min(8, cpu_count()) threads and can be overridden via the
METEOR_GAMMA_PPF_THREADS env var (=1 to disable). Malformed values
silently fall back so a typo in a batch script cannot crash METEOR at
import time.
Wins on gen_gridded (see docs/profiling_baseline.md for methodology):
N=1: 171 s -> 45 s (3.8x)
N=5: 416 s -> 154 s (2.7x)
Post-change, the fixed cost (imports, netCDF loads, pattern-scaling load,
~22 s) and the actual gammaincinv arithmetic (~64 s at N=5 across 8
threads) dominate what's left.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- black reformats a couple of function calls in both files - ruff F841: drop unused `n_obs` from _vectorized_gamma_mle - ruff F841: drop unused `n_spatial` from fit_distribution_parameters_3d (leftover from the previous per-gridpoint loop that this PR removes) - ruff F401: drop unused `import os` from tests - pylint E0611 false positive: suppress `no-name-in-module` on the scipy.special import (gammaincinv is present in the runtime module but astroid's stubs don't list it) Lint-only, no behavior change. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Comment on lines
+61
to
+65
| n_threads : int, optional | ||
| Explicit thread count. If ``None`` (default), consults the | ||
| ``METEOR_GAMMA_PPF_THREADS`` env var, falling back to | ||
| ``min(8, cpu_count())``. Bypasses threading for small problems where | ||
| thread setup would dominate. |
Collaborator
There was a problem hiding this comment.
If returns diminish above n_threads = 8, should a user choice much beyond this be overridden/yield a warning?
Comment on lines
+32
to
+36
| Reads the ``METEOR_GAMMA_PPF_THREADS`` environment variable so HPC users | ||
| can pin threads without touching Python (same pattern as ``OMP_NUM_THREADS``). | ||
| Falls back to ``min(8, cpu_count())`` — bounded because the underlying | ||
| scipy special function has diminishing returns past ~8 threads on typical | ||
| hardware. Set to ``1`` to disable threading entirely. |
Collaborator
There was a problem hiding this comment.
If returns diminish much past 8 CPUs, should values way beyond this for METEOR_GAMMA_PPF_THREADS be capped? Warned about? Maybe no big deal and all users who set this can be expected to not send nonsense?
|
|
||
| # Sum and log-sum with invalid entries zeroed out | ||
| x_masked = np.where(valid, x, 0.0) | ||
| logx_masked = np.where(valid, np.log(np.maximum(x, np.finfo(float).tiny)), 0.0) |
Collaborator
There was a problem hiding this comment.
Should this mask over x_masked and not x, so as to avoid crashes for logs of negative numbers in the evaluation?
Collaborator
There was a problem hiding this comment.
Add tests with negative values?
maritsandstad
left a comment
maritsandstad
left a comment
Collaborator
There was a problem hiding this comment.
Left a few comments with questions here. Otherwise it looks alright to me
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.
Summary
Two independent speedups to the 3D gamma quantile-mapping path used for gridded precipitation output.
1.
fit_distribution_parameters_3dReplaced the per-gridpoint
for i in range(n_spatial): scipy.stats.gamma.fit(x, floc=0)loop with a vectorized_vectorized_gamma_mlethat runs Choi-Wette + Newton iterations onlog(k) - psi(k) = log(mean(x)) - mean(log(x))across all gridpoints at once. This is the same closed-form MLE scipy uses internally, just batched.rtol=1e-6.2.
apply_distribution_transform(3D gamma path)Wrapped
scipy.special.gammaincinv(which underliesstats.gamma.ppfand releases the GIL) in aThreadPoolExecutorthat chunks the spatial axis. Output is bitwise-identical tostats.gamma.ppf.min(8, cpu_count())threads.METEOR_GAMMA_PPF_THREADSenv variable (=1disables threading).Wins on
gen_griddedPost-change, the fixed cost (imports, netCDF loads, pattern-scaling load ≈ 22 s) plus the actual
gammaincinvarithmetic (~64 s at N=5 across 8 threads) dominate the remaining time.Correctness
test_vectorized_gamma_mle_matches_scipy— synthetic gamma at 200 gridpoints × 500 samples, rtol=1e-6.test_vectorized_gamma_mle_handles_invalid_columns— all-zero / all-NaN / single-positive columns fall back to safe defaults with no NaN propagation.test_threaded_gamma_ppf_3d_matches_scipy— assertsarray_equalagainststats.gamma.ppffor a(200, 5000)grid.test_threaded_gamma_ppf_3d_single_thread_bypass—n_threads=1produces the same output as multi-threaded.test_resolve_gamma_ppf_threads_env_var— env var override + fallback semantics.Test plan
pytest tests/unit/test_precipitation_transform.pyall passMETEOR_GAMMA_PPF_THREADS=1 python -m pytest tests/unit/test_precipitation_transform.py -k ppfstill passes🤖 Generated with Claude Code