Vectorize gamma MLE fit and thread gamma.ppf for gridded transform

CHANGELOG.rst

@@ -19,6 +19,8 @@ The changes listed in this file are categorised as follows:
 
 ### Changed
 
+- Vectorized the gridded gamma-distribution fit in ``fit_distribution_parameters_3d``. Replaced the per-gridpoint ``scipy.stats.gamma.fit`` loop with a single-shot Choi-Wette + Newton iteration on the same closed-form estimator scipy uses internally; agreement with scipy's per-gridpoint fit is at the 1e-6 relative level.
+- Threaded the gridded ``gamma.ppf`` step in ``apply_distribution_transform`` by chunking the spatial axis and dispatching to ``scipy.special.gammaincinv`` (which releases the GIL) via ``concurrent.futures``. Output is bitwise-identical to ``scipy.stats.gamma.ppf``. Thread count defaults to ``min(8, cpu_count())`` and can be overridden with the ``METEOR_GAMMA_PPF_THREADS`` environment variable (``=1`` to disable, or any positive integer); malformed values silently fall back to the default so batch scripts with a typo cannot crash METEOR at import time. Combined with the vectorized fit above, gridded generation is 3.8× faster at N=1 and 2.7× at N=5 (see ``docs/profiling_baseline.md``).
 - Now possibly to send variable length temperature scaling timeseries, fixed noise generator for wrong ordering of base data dimensions
 
 ### Fixed

src/meteor/precipitation_transform.py

@@ -13,9 +13,139 @@
 Author: METEOR Development Team
 """
 
+import os
+from concurrent.futures import ThreadPoolExecutor
+
 import numpy as np
 import xarray as xr
 from scipy import stats
+from scipy.special import (  # pylint: disable=no-name-in-module
+    digamma,
+    gammaincinv,
+    polygamma,
+)
+
+
+def _resolve_gamma_ppf_threads():
+    """Pick thread count for :func:`_threaded_gamma_ppf_3d`.
+
+    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.
+    """
+    env = os.environ.get("METEOR_GAMMA_PPF_THREADS")
+    if env:
+        try:
+            n = int(env)
+            if n >= 1:
+                return n
+        except ValueError:
+            pass
+    return min(8, os.cpu_count() or 1)
+
+
+def _threaded_gamma_ppf_3d(u, shape_flat, scale_flat, n_threads=None):
+    """gamma.ppf across the trailing spatial axis, threaded.
+
+    scipy's ``gammaincinv`` (which backs ``stats.gamma.ppf``) releases the GIL,
+    so we chunk along the spatial axis and evaluate in parallel threads. For
+    the METEOR gridded workload this dominates :func:`apply_distribution_transform`.
+
+    Parameters
+    ----------
+    u : ndarray, shape (..., n_spatial)
+        Uniform-scale quantiles (output of ``norm.cdf``).
+    shape_flat, scale_flat : ndarray, shape (n_spatial,)
+    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.
+    """
+    n_spatial = shape_flat.shape[0]
+    if n_threads is None:
+        n_threads = _resolve_gamma_ppf_threads()
+    if n_threads <= 1 or n_spatial < 4096:
+        return scale_flat * gammaincinv(shape_flat, u)
+
+    chunks = np.array_split(np.arange(n_spatial), n_threads)
+    out = np.empty_like(u)
+
+    def _work(idx):
+        return idx, scale_flat[idx] * gammaincinv(shape_flat[idx], u[..., idx])
+
+    with ThreadPoolExecutor(n_threads) as ex:
+        for idx, res in ex.map(_work, chunks):
+            out[..., idx] = res
+    return out
+
+
+def _vectorized_gamma_mle(data, max_iter=8, tol=1e-8):
+    """MLE fit of Gamma(shape, scale) with location fixed at 0, vectorized.
+
+    Same estimator scipy.stats.gamma.fit(x, floc=0) uses internally, but
+    applied to a batch of independent samples in one call. For each column j
+    of ``data`` (shape ``(n_obs, n_series)``), solve
+
+        log(k) - psi(k) = log(mean(x)) - mean(log(x))       (Choi & Wette 1969)
+        theta = mean(x) / k
+
+    by Newton's method on k, seeded with the Choi-Wette initial guess.
+
+    Non-positive samples are treated as invalid and masked out; series with
+    fewer than 2 positive samples fall back to shape=1.0, scale=mean.
+
+    Parameters
+    ----------
+    data : ndarray (n_obs, n_series)
+    max_iter : int
+    tol : float
+        Convergence tolerance on |Δk| / k.
+
+    Returns
+    -------
+    shape, scale : ndarray (n_series,)
+    """
+    x = np.asarray(data, dtype=np.float64)
+
+    valid = (x > 0) & np.isfinite(x)
+    n_valid = valid.sum(axis=0)
+
+    # 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)
+
+    with np.errstate(invalid="ignore", divide="ignore"):
+        mean_x = x_masked.sum(axis=0) / n_valid
+        mean_logx = logx_masked.sum(axis=0) / n_valid
+        s = np.log(mean_x) - mean_logx  # >= 0 with equality iff constant
+
+    # Guard against degenerate series
+    fittable = (n_valid >= 2) & np.isfinite(s) & (s > 0)
+
+    # Choi-Wette initial guess for k (only used where fittable)
+    with np.errstate(invalid="ignore", divide="ignore"):
+        k = (3.0 - s + np.sqrt(np.maximum((s - 3.0) ** 2 + 24.0 * s, 0.0))) / (12.0 * s)
+    k = np.where(fittable, k, 1.0)
+    k = np.clip(k, 1e-6, 1e6)
+
+    # Newton iterations
+    for _ in range(max_iter):
+        f = np.log(k) - digamma(k) - s
+        fp = 1.0 / k - polygamma(1, k)  # trigamma
+        step = f / fp
+        k_new = np.clip(k - step, 1e-6, 1e6)
+        if np.max(np.abs(k_new - k) / np.maximum(k, 1e-12)) < tol:
+            k = k_new
+            break
+        k = k_new
+
+    shape = np.where(fittable, k, 1.0)
+    scale = np.where(fittable, mean_x / shape, np.where(n_valid > 0, mean_x, 0.01))
+    return shape, scale
+
 
 # =============================================================================
 # PATH 1: 1D Transform (for regional/global mean time series)
@@ -194,8 +324,6 @@ def fit_distribution_parameters_3d(spatial_data, distribution="gamma"):
             f"4D (n_ensemble, n_time, n_lat, n_lon), got shape {data_array.shape}"
         )
 
-    n_spatial = n_lat * n_lon
-
     if distribution == "gaussian":
         # Fit Gaussian: simple mean and std at each grid point
         mean_params = np.mean(data_reshaped, axis=0).reshape(n_lat, n_lon)
@@ -204,32 +332,11 @@ def fit_distribution_parameters_3d(spatial_data, distribution="gamma"):
         params = {"mean": mean_params, "std": std_params}
 
     elif distribution == "gamma":
-        # Fit Gamma distribution at each grid point
-        shape_params = np.zeros(n_spatial)
-        scale_params = np.zeros(n_spatial)
-
-        print(f"Fitting Gamma distribution to {n_spatial} grid points...")
-        for i in range(n_spatial):
-            grid_data = data_reshaped[:, i]
-            # Use the 1D fitting function for consistency
-            try:
-                grid_params = fit_distribution_parameters_1d(
-                    grid_data, distribution="gamma"
-                )
-                shape_params[i] = grid_params["shape"]
-                scale_params[i] = grid_params["scale"]
-            except (ValueError, RuntimeError, RuntimeWarning):
-                # If fit fails completely (e.g., all NaNs or invalid data),
-                # use default values
-                shape_params[i] = 1.0
-                scale_params[i] = 0.01
-
-            if (i + 1) % 5000 == 0:  # pragma no cover
-                print(f"  Processed {i + 1}/{n_spatial} grid points...")
-
+        # Vectorized MLE fit across all gridpoints in one shot.
+        shape_flat, scale_flat = _vectorized_gamma_mle(data_reshaped)
         params = {
-            "shape": shape_params.reshape(n_lat, n_lon),
-            "scale": scale_params.reshape(n_lat, n_lon),
+            "shape": shape_flat.reshape(n_lat, n_lon),
+            "scale": scale_flat.reshape(n_lat, n_lon),
         }
 
     else:
@@ -351,9 +458,7 @@ def apply_distribution_transform(
         if target_dist == "gamma":
             shape_grid = target_params["shape"].flatten()
             scale_grid = target_params["scale"].flatten()
-            transformed_flat = stats.gamma.ppf(
-                uniform, a=shape_grid[None, :], scale=scale_grid[None, :]
-            )
+            transformed_flat = _threaded_gamma_ppf_3d(uniform, shape_grid, scale_grid)
         else:
             raise ValueError(
                 f"Only 'gamma' distribution supported for 3D data, got {target_dist}"

tests/unit/test_precipitation_transform.py

@@ -5,8 +5,12 @@
 import numpy as np
 import pytest
 import xarray as xr
+from scipy import stats
 
 from meteor.precipitation_transform import (
+    _resolve_gamma_ppf_threads,
+    _threaded_gamma_ppf_3d,
+    _vectorized_gamma_mle,
     apply_distribution_transform,
     apply_empirical_quantile_mapping,
     fit_distribution_parameters_1d,
@@ -95,6 +99,110 @@ def test_fit_distribution_parameters_3d():
     assert np.allclose(params["std"].all(), 0.0)
 
 
+def test_vectorized_gamma_mle_matches_scipy():
+    """The vectorized MLE must agree with scipy's per-gridpoint MLE fit.
+
+    Both implementations use the same estimator (Choi-Wette initial guess
+    plus Newton iteration on log(k) - psi(k) = log(mean(x)) - mean(log(x))),
+    so on well-behaved synthetic gamma data they should agree to
+    floating-point precision.
+    """
+    rng = np.random.default_rng(42)
+    n_series = 200
+    n_obs = 500
+    true_shape = rng.uniform(0.5, 5.0, size=n_series)
+    true_scale = rng.uniform(0.1, 10.0, size=n_series)
+    data = np.stack(
+        [rng.gamma(k, s, size=n_obs) for k, s in zip(true_shape, true_scale)],
+        axis=1,
+    )  # (n_obs, n_series)
+
+    vec_shape, vec_scale = _vectorized_gamma_mle(data)
+
+    scipy_shape = np.empty(n_series)
+    scipy_scale = np.empty(n_series)
+    for i in range(n_series):
+        k, _, s = stats.gamma.fit(data[:, i], floc=0)
+        scipy_shape[i] = k
+        scipy_scale[i] = s
+
+    np.testing.assert_allclose(vec_shape, scipy_shape, rtol=1e-6)
+    np.testing.assert_allclose(vec_scale, scipy_scale, rtol=1e-6)
+
+
+def test_vectorized_gamma_mle_handles_invalid_columns():
+    """Columns with fewer than two positive samples must fall back cleanly
+    rather than propagating NaN or crashing.
+    """
+    n_obs = 100
+    data = np.column_stack(
+        [
+            np.random.default_rng(0).gamma(2.0, 1.0, size=n_obs),
+            np.zeros(n_obs),  # all zeros -> not fittable
+            np.full(n_obs, np.nan),  # all NaN -> not fittable
+            np.array([1.0] + [0.0] * (n_obs - 1)),  # single positive value
+        ]
+    )
+
+    shape, scale = _vectorized_gamma_mle(data)
+
+    assert np.all(np.isfinite(shape))
+    assert np.all(np.isfinite(scale))
+    assert shape[0] > 0 and scale[0] > 0
+    # Degenerate columns fall back to the safe defaults from the doc.
+    assert shape[1] == 1.0
+    assert shape[2] == 1.0
+    assert shape[3] == 1.0
+
+
+def test_threaded_gamma_ppf_3d_matches_scipy():
+    """Threaded ppf must reproduce scipy.stats.gamma.ppf bitwise."""
+    rng = np.random.default_rng(0)
+    n_time, n_spatial = 200, 5_000  # n_spatial > 4096 to trigger threading
+    shape = rng.uniform(0.5, 5.0, size=n_spatial)
+    scale = rng.uniform(0.1, 10.0, size=n_spatial)
+    u = rng.uniform(0.01, 0.99, size=(n_time, n_spatial))
+
+    expected = stats.gamma.ppf(u, a=shape[None, :], scale=scale[None, :])
+    got = _threaded_gamma_ppf_3d(u, shape, scale)
+
+    np.testing.assert_array_equal(expected, got)
+
+
+def test_threaded_gamma_ppf_3d_single_thread_bypass():
+    """n_threads=1 should return the same result via the direct (non-thread) path."""
+    rng = np.random.default_rng(1)
+    shape = rng.uniform(0.5, 5.0, size=5_000)
+    scale = rng.uniform(0.1, 10.0, size=5_000)
+    u = rng.uniform(0.01, 0.99, size=(100, 5_000))
+
+    r_direct = _threaded_gamma_ppf_3d(u, shape, scale, n_threads=1)
+    r_threaded = _threaded_gamma_ppf_3d(u, shape, scale, n_threads=4)
+    np.testing.assert_array_equal(r_direct, r_threaded)
+
+
+def test_resolve_gamma_ppf_threads_env_var(monkeypatch):
+    """METEOR_GAMMA_PPF_THREADS overrides the default; invalid values fall back."""
+    monkeypatch.delenv("METEOR_GAMMA_PPF_THREADS", raising=False)
+    default = _resolve_gamma_ppf_threads()
+    assert default >= 1
+
+    monkeypatch.setenv("METEOR_GAMMA_PPF_THREADS", "3")
+    assert _resolve_gamma_ppf_threads() == 3
+
+    monkeypatch.setenv("METEOR_GAMMA_PPF_THREADS", "1")
+    assert _resolve_gamma_ppf_threads() == 1
+
+    # Non-integer input silently falls back so batch scripts with a typo
+    # don't crash METEOR at import time.
+    monkeypatch.setenv("METEOR_GAMMA_PPF_THREADS", "garbage")
+    assert _resolve_gamma_ppf_threads() == default
+
+    # Zero (i.e. "no threads") is also invalid; fall back.
+    monkeypatch.setenv("METEOR_GAMMA_PPF_THREADS", "0")
+    assert _resolve_gamma_ppf_threads() == default
+
+
 def test_apply_distribution_transform():
     """Test applying distribution transform to data."""