From 9f5147a6111f040ea135e1660d5ddd1f9b0c920e Mon Sep 17 00:00:00 2001
From: Intron7 <sdicks@nvidia.com>
Date: Wed, 23 Sep 2026 12:21:32 +0200
Subject: [PATCH] Honor the requested coefficient in shrinkage grid fallback

---
 src/pydeseq2/distributions.py |  1 +
 src/pydeseq2/grid_search.py   |  4 +++
 tests/test_grid_search.py     | 60 +++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+)
 create mode 100644 tests/test_grid_search.py

diff --git a/src/pydeseq2/distributions.py b/src/pydeseq2/distributions.py
index e2c53791..1b3cbcc0 100644
--- a/src/pydeseq2/distributions.py
+++ b/src/pydeseq2/distributions.py
@@ -302,6 +302,7 @@ def ddf(beta: np.ndarray, cnst: float = scale_cnst) -> np.ndarray:
             grid_length=60,
             min_beta=-30,
             max_beta=30,
+            shrink_index=shrink_index,
         )
 
     inv_hessian = np.linalg.inv(ddf(beta, 1))
diff --git a/src/pydeseq2/grid_search.py b/src/pydeseq2/grid_search.py
index 28e158cd..fd60344b 100644
--- a/src/pydeseq2/grid_search.py
+++ b/src/pydeseq2/grid_search.py
@@ -208,6 +208,7 @@ def grid_fit_shrink_beta(
     grid_length: int = 60,
     min_beta: float = -30,
     max_beta: float = 30,
+    shrink_index: int = 1,
 ) -> np.ndarray:
     """Find best LFC parameter.
 
@@ -235,6 +236,8 @@ def grid_fit_shrink_beta(
         Lower-bound on LFC. (default: ``30``).
     max_beta
         Upper-bound on LFC. (default: ``30``).
+    shrink_index
+        Index of the LFC coordinate to shrink. (default: ``1``).
 
     Returns
     -------
@@ -255,6 +258,7 @@ def loss(beta: np.ndarray) -> float:
                 offset,
                 prior_no_shrink_scale,
                 prior_scale,
+                shrink_index=shrink_index,
             )
             / scale_cnst
         )
diff --git a/tests/test_grid_search.py b/tests/test_grid_search.py
new file mode 100644
index 00000000..2c47ca67
--- /dev/null
+++ b/tests/test_grid_search.py
@@ -0,0 +1,60 @@
+import numpy as np
+import pytest
+from scipy.optimize import OptimizeResult
+from scipy.optimize import minimize
+from scipy.stats import nbinom
+
+from pydeseq2.distributions import nbinomGLM
+
+
+@pytest.mark.parametrize("shrink_index", [0, 1])
+def test_shrink_grid_fallback_preserves_requested_prior(monkeypatch, shrink_index):
+    design = np.column_stack([np.ones(6), [0, 0, 0, 1, 1, 1]])
+    counts = np.array([12, 18, 21, 60, 55, 90])
+    size = np.full(6, 5.0)
+    offset = np.zeros(6)
+    prior_no_shrink_scale = 15.0
+    prior_scale = 0.3
+
+    def objective(beta):
+        mu = np.exp(design @ beta + offset)
+        prior = beta[1 - shrink_index] ** 2 / (2 * prior_no_shrink_scale**2)
+        prior += np.log1p((beta[shrink_index] / prior_scale) ** 2)
+        return -nbinom.logpmf(counts, size, size / (size + mu)).sum() + prior
+
+    optimum = minimize(
+        objective,
+        np.array([2.0, 1.0]),
+        method="L-BFGS-B",
+        options={"ftol": 1e-12, "gtol": 1e-6},
+    )
+    assert optimum.success
+
+    monkeypatch.setattr(
+        "pydeseq2.distributions.minimize",
+        lambda *args, **kwargs: OptimizeResult(x=np.zeros(2), success=False),
+    )
+    beta, _, converged = nbinomGLM(
+        design,
+        counts,
+        size,
+        offset,
+        prior_no_shrink_scale,
+        prior_scale,
+        shrink_index=shrink_index,
+    )
+    assert not converged
+    # The refined grid has spacing approximately 0.0345 in each coordinate.
+    np.testing.assert_allclose(beta, optimum.x, atol=0.035, rtol=0)
+    assert objective(beta) - optimum.fun < 0.01
+
+    permuted_beta, _, _ = nbinomGLM(
+        design[:, ::-1],
+        counts,
+        size,
+        offset,
+        prior_no_shrink_scale,
+        prior_scale,
+        shrink_index=1 - shrink_index,
+    )
+    np.testing.assert_allclose(permuted_beta[::-1], beta, atol=1e-12, rtol=0)