Repository: owkin/PyDESeq2 (python)

Issue:
[BUG] lfc_shrink() generates wrong results
**Describe the bug**
lfc_shrink() generates wrong results. In most cases it acts as expected. For some genes it increases absolute value of fold change. For one differentially expressed gene with high expression (base_mean = 407277.128029) it even changes the sign (
log2FoldChange = -0.753026; shrankLog2FoldChange = 1.844100). 

**To Reproduce**
```
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

from pydeseq2.dds import DeseqDataSet
from pydeseq2.default_inference import DefaultInference
from pydeseq2.ds import DeseqStats

inference = DefaultInference(n_cpus=32)

dds = DeseqDataSet(
    counts=counts_df,
    metadata=metadata_df,
    design_factors='condition',
    inference=inference,
)
dds.deseq2()

stat_res = DeseqStats(
    dds,
    contrast=['condition', 'B', 'A'],
    inference=inference
)
stat_res.summary()
stat_res_df = stat_res.results_df.copy()

stat_res.lfc_shrink('condition_B_vs_A')
stat_res_shrinked_df = stat_res.results_df.copy()
stat_res_shrinked_df = stat_res_shrinked_df.rename(columns={'log2FoldChange': 'log2FoldChange_shrank'})

merged_df = pd.merge(stat_res_df, stat_res_shrinked_df, left_index=True, right_index=True)
merged_df = merged_df[(merged_df['padj_x'].notna())]

sns.scatterplot(
    merged_df,
    x='log2FoldChange',
    y='log2FoldChange_shrank'
)

# replace the sns.scatterplot() with code below to get rid of seaborn dependency
# plt.scatter(
#     x=merged_df['log2FoldChange'],
#     y=merged_df['log2FoldChange_shrank']
#)

plt.axline((0, 0), (1,1))
```
```
Fitting size factors...
... done in 0.02 seconds.

Fitting dispersions...
... done in 1.83 seconds.

Fitting dispersion trend curve...
... done in 0.55 seconds.

Fitting MAP dispersions...
... done in 1.79 seconds.

Fitting LFCs...
... done in 2.18 seconds.

Calculating cook's distance...
... done in 0.02 seconds.

Replacing 0 outlier genes.

Running Wald tests...
... done in 2.47 seconds.

Fitting MAP LFCs...
Log2 fold change & Wald test p-value: condition B vs A
                       baseMean  log2FoldChange     lfcSE      stat    pvalue  \
GENEID                                                                          
ENSG00000000003.15  6312.552040       -0.406453  0.092641 -4.387384  0.000011   
ENSG00000000005.6      5.337949       -5.852983  2.030001 -2.883242  0.003936   
ENSG00000000419.14  3926.067459        0.396839  0.089676  4.425236  0.000010   
ENSG00000000457.14   986.626013       -0.033945  0.132989 -0.255247  0.798532   
ENSG00000000460.17  3054.441682       -0.008417  0.079740 -0.105558  0.915933   
...                         ...             ...       ...       ...       ...   
ENSG00000289714.1      0.000000             NaN       NaN       NaN       NaN   
ENSG00000289715.1      0.000000             NaN       NaN       NaN       NaN   
ENSG00000289716.1     55.728811        1.269762  0.336892  3.769053  0.000164   
ENSG00000289718.1      0.000000             NaN       NaN       NaN       NaN   
ENSG00000289719.1     11.956688        2.075506  0.837589  2.477953  0.013214   

                        padj  
GENEID                        
ENSG00000000003.15  0.000048  
ENSG00000000005.6   0.010336  
ENSG00000000419.14  0.000041  
ENSG00000000457.14  0.863010  
ENSG00000000460.17  0.944234  
...                      ...  
ENSG00000289714.1        NaN  
ENSG00000289715.1        NaN  
ENSG00000289716.1   0.000564  
ENSG00000289718.1        NaN  
ENSG00000289719.1   0.030276  

[61125 rows x 6 columns]
/home/.../pydeseq2/utils.py:1260: RuntimeWarning: overflow encountered in exp
  counts - (counts + size) / (1 + size * np.exp(-xbeta - offset))

Shrunk log2 fold change & Wald test p-value: condition B vs A
                       baseMean  log2FoldChange     lfcSE      stat    pvalue  \
GENEID                                                                          
ENSG00000000003.15  6312.552040       -0.405783  0.092409 -4.387384  0.000011   
ENSG00000000005.6      5.337949       -6.123571  2.817030 -2.883242  0.003936   
ENSG00000000419.14  3926.067459        0.387581  0.089467  4.425236  0.000010   
ENSG00000000457.14   986.626013       -0.035274  0.131941 -0.255247  0.798532   
ENSG00000000460.17  3054.441682       -0.007822  0.079505 -0.105558  0.915933   
...                         ...             ...       ...       ...       ...   
ENSG00000289714.1      0.000000             NaN       NaN       NaN       NaN   
ENSG00000289715.1      0.000000             NaN       NaN       NaN       NaN   
ENSG00000289716.1     55.728811        1.162090  0.337499  3.769053  0.000164   
ENSG00000289718.1      0.000000             NaN       NaN       NaN       NaN   
ENSG00000289719.1     11.956688        1.440262  0.866340  2.477953  0.013214   

                        padj  
GENEID                        
ENSG00000000003.15  0.000048  
ENSG00000000005.6   0.010336  
ENSG00000000419.14  0.000041  
ENSG00000000457.14  0.863010  
ENSG00000000460.17  0.944234  
...                      ...  
ENSG00000289714.1        NaN  
ENSG00000289715.1        NaN  
ENSG00000289716.1   0.000564  
ENSG00000289718.1        NaN  
ENSG00000289719.1   0.030276  

[61125 rows x 6 columns]
... done in 1.97 seconds.
```

**Expected behavior**
No increase in absolute value of fold change (if i understand correctly what apeGLM shrinking is doing).
No change in direction of fold change in very abundant differentially expressed genes.
No instability based on the name of conditions.

**Screenshots**
![image](https://github.com/user-attachments/assets/6f9e7784-8856-4750-a0e5-31e93d71ff53)


**Desktop (please complete the following information):**
 - OS: Ubuntu 20.04.3 LTS
 - pydeseq2: 0.4.12 

**Additional context**
The behaviour is not stable based on the name of groups. If i exchange A to B the values are different and the overflow error doent pop up, however the results are still suspicious. The data i can not publicly share, but potentially i would be able to share them withthout geneids. The most weird and unstable behaviour was seen in genes with high expression.

Pull request description:
BUG `nbinomGLM` returns before convergence
#### Reference Issue or PRs

Fixes #347

#### What does your PR implement? Be specific.

In some cases (e.g., genes with very large counts, c.f. #347), `nbinomGLM` stops before the solution has actually converged. This lead to aberrant shrunk LFCs (e.g., changes of signs, or values with a larger absolute value).

This PR fixes this issue by making the convergence criteria passed to `scipy.optimize.minimize` stricter.

It also adds a new test case, corresponding to the MWE provided by @martinPasen in #347. 

Patch:
diff --git a/pydeseq2/utils.py b/pydeseq2/utils.py
index 4a686dd..db09e90 100644
--- a/pydeseq2/utils.py
+++ b/pydeseq2/utils.py
@@ -1083,7 +1083,7 @@ def nbinomGLM(
         xbeta = design_matrix @ beta
         d_neg_prior = (
             beta * no_shrink_mask / prior_no_shrink_scale**2
-            + 2 * beta * shrink_mask / (prior_scale**2 + beta[shrink_index] ** 2),
+            + 2 * beta * shrink_mask / (prior_scale**2 + beta[shrink_index] ** 2)
         )
 
         d_nll = (
@@ -1116,6 +1116,10 @@ def nbinomGLM(
         jac=df,
         hess=ddf if optimizer == "Newton-CG" else None,
         method=optimizer,
+        options={
+            "ftol": 1e-8,
+            "gtol": 1e-8,
+        },
     )
 
     beta = res.x


Test patch:
diff --git a/tests/data/large_counts/r_test_dispersions.csv b/tests/data/large_counts/r_test_dispersions.csv
new file mode 100644
index 0000000..98e8cbd
--- /dev/null
+++ b/tests/data/large_counts/r_test_dispersions.csv
@@ -0,0 +1,6 @@
+"","x"
+"g1",0.11861243039715
+"g2",0.0262784375410397
+"g3",0.0154560937196622
+"g4",0.0187562440702638
+"g5",0.029726368740505
diff --git a/tests/data/large_counts/r_test_lfc_shrink_res.csv b/tests/data/large_counts/r_test_lfc_shrink_res.csv
new file mode 100644
index 0000000..3fa16c8
--- /dev/null
+++ b/tests/data/large_counts/r_test_lfc_shrink_res.csv
@@ -0,0 +1,6 @@
+"","baseMean","log2FoldChange","lfcSE","pvalue","padj"
+"g1",22.1449196667942,0.0980467648341208,0.313138103429392,0.67467748621509,0.67467748621509
+"g2",381.001165601189,-0.25238081586499,0.166175542221483,0.102486854258627,0.170811423764379
+"g3",2655.2687326308,1.32808989591327,0.129192036932687,1.04455372148152e-25,5.22276860740759e-25
+"g4",11603.4762054462,0.0630455944543123,0.134144518362108,0.621192047119601,0.67467748621509
+"g5",392923.057672917,-0.409599894742138,0.172482439150294,0.0120579588905087,0.0301448972262718
diff --git a/tests/data/large_counts/r_test_res.csv b/tests/data/large_counts/r_test_res.csv
new file mode 100644
index 0000000..f661f78
--- /dev/null
+++ b/tests/data/large_counts/r_test_res.csv
@@ -0,0 +1,6 @@
+"","baseMean","log2FoldChange","lfcSE","stat","pvalue","padj"
+"g1",22.1449196667942,0.173017658054474,0.412204780862638,0.419737145436286,0.67467748621509,0.67467748621509
+"g2",381.001165601189,-0.283429708243346,0.173572838265613,-1.63291509821152,0.102486854258627,0.170811423764379
+"g3",2655.2687326308,1.34796297540684,0.128597274635474,10.4820493220235,1.04455372148152e-25,5.22276860740759e-25
+"g4",11603.4762054462,0.0691977287249358,0.140030562520189,0.494161613576031,0.621192047119601,0.67467748621509
+"g5",392923.057672917,-0.441570776728009,0.175893527435991,-2.51044358007261,0.0120579588905087,0.0301448972262718
diff --git a/tests/data/large_counts/r_test_size_factors.csv b/tests/data/large_counts/r_test_size_factors.csv
new file mode 100644
index 0000000..6483001
--- /dev/null
+++ b/tests/data/large_counts/r_test_size_factors.csv
@@ -0,0 +1,9 @@
+"","x"
+"A1",1.06181876155545
+"A2",1.24743230704576
+"A3",1.32173521903446
+"A4",1.09150326866504
+"B1",1.0472175142877
+"B2",0.750260803174501
+"B3",1.16725656977807
+"B4",0.703136348465353
diff --git a/tests/test_pydeseq2.py b/tests/test_pydeseq2.py
index 2279fc8..54c6119 100644
--- a/tests/test_pydeseq2.py
+++ b/tests/test_pydeseq2.py
@@ -372,6 +372,72 @@ def test_iterative_size_factors(counts_df, metadata, tol=0.02):
     ).max() < tol
 
 
+def test_lfc_shrinkage_large_counts(counts_df, metadata, tol=0.02):
+    """Test that the outputs of the lfc_shrink function match those of the original
+    R package (starting from the same inputs), up to a tolerance in relative error in
+    the presence of a gene with very large counts.
+    """
+
+    test_path = str(Path(os.path.realpath(tests.__file__)).parent.resolve())
+    r_res = pd.read_csv(
+        os.path.join(test_path, "data/large_counts/r_test_res.csv"), index_col=0
+    )
+    r_shrunk_res = pd.read_csv(
+        os.path.join(test_path, "data/large_counts/r_test_lfc_shrink_res.csv"),
+        index_col=0,
+    )
+
+    r_size_factors = pd.read_csv(
+        os.path.join(test_path, "data/large_counts/r_test_size_factors.csv"),
+        index_col=0,
+    ).squeeze()
+
+    r_dispersions = pd.read_csv(
+        os.path.join(test_path, "data/large_counts/r_test_dispersions.csv"),
+        index_col=0,
+    ).squeeze()
+
+    counts_df = pd.DataFrame(
+        data=[
+            [25, 405, 1355, 12558, 489843],
+            [28, 480, 2144, 13844, 514571],
+            [12, 690, 1919, 15632, 564106],
+            [31, 420, 1684, 11513, 556380],
+            [34, 278, 3849, 11577, 412551],
+            [19, 249, 3086, 7296, 295565],
+            [17, 491, 4089, 13805, 280945],
+            [15, 251, 2785, 10492, 214062],
+        ],
+        index=["A1", "A2", "A3", "A4", "B1", "B2", "B3", "B4"],
+        columns=["g1", "g2", "g3", "g4", "g5"],
+    )
+
+    metadata_df = pd.DataFrame(
+        data=["A", "A", "A", "A", "B", "B", "B", "B"],
+        index=["A1", "A2", "A3", "A4", "B1", "B2", "B3", "B4"],
+        columns=["condition"],
+    )
+
+    dds = DeseqDataSet(counts=counts_df, metadata=metadata_df, design="~condition")
+    dds.deseq2()
+
+    dds.obsm["size_factors"] = r_size_factors.values
+    dds.varm["dispersions"] = r_dispersions.values
+    dds.varm["LFC"].iloc[:, 1] = r_res.log2FoldChange.values * np.log(2)
+
+    res = DeseqStats(dds, contrast=["condition", "B", "A"])
+    res.summary()
+    res.SE = r_res.lfcSE * np.log(2)
+    res.lfc_shrink(coeff="condition[T.B]")
+    shrunk_res = res.results_df
+
+    # Check that the same LFC are found (up to tol)
+    assert (
+        abs(r_shrunk_res.log2FoldChange - shrunk_res.log2FoldChange)
+        / abs(r_shrunk_res.log2FoldChange)
+    ).max() < tol
+
+
 # Multi-factor tests
 @pytest.mark.parametrize("with_outliers", [True, False])
 def test_multifactor_deseq(counts_df, metadata, with_outliers, tol=0.04):