From 35d7247d1822e8d82db4a4f43cc56c6019bf85ae Mon Sep 17 00:00:00 2001
From: Christian Michelsen <christianmichelsen@gmail.com>
Date: Mon, 15 Aug 2022 09:47:45 +0200
Subject: [PATCH] fix: Improve Bayesian D_max when few reads

---
 src/metaDMG/fit/bayesian.py    | 25 +++++++++++++++++--------
 src/metaDMG/fit/frequentist.py |  7 ++-----
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/src/metaDMG/fit/bayesian.py b/src/metaDMG/fit/bayesian.py
index 135c6d9..51c31de 100644
--- a/src/metaDMG/fit/bayesian.py
+++ b/src/metaDMG/fit/bayesian.py
@@ -114,20 +114,29 @@ def compute_posterior(
 
 
 def compute_D_max(mcmc, data):
-    posterior = get_posterior_predictive(mcmc, data)
-    c = mcmc.get_samples()["c"]
-    f = posterior["obs"] / data["N"]
-    f = f[:, 0]
-    D_max_samples = f - c
-    D_max_mu = np.mean(D_max_samples).item()
-    D_max_std = np.std(D_max_samples).item()
+    # posterior = get_posterior_predictive(mcmc, data)
+    # c = mcmc.get_samples()["c"]
+    # f = posterior["obs"] / data["N"]
+    # f = f[:, 0]
+    # D_max_samples = f - c
+    # D_max_mu = np.mean(D_max_samples).item()
+    # D_max_std = np.std(D_max_samples).item()
 
     # A = np.mean(mcmc.get_samples()["A"]).item()
     # phi = np.mean(mcmc.get_samples()["phi"]).item()
     # N = data["N"][0]
     # np.sqrt(A*(1-A)*(phi+N)/((phi+1)*N))
 
-    return {"mu": D_max_mu, "std": D_max_std}
+    # New method, more similar to frequentist and better when few reads
+    A = mcmc.get_samples()["A"]
+    c = mcmc.get_samples()["c"]
+    phi = mcmc.get_samples()["phi"]
+    N = max(data["N"][0], 1)
+    mu = np.mean(A)
+    std = np.mean(np.sqrt(A * (1 - A) * (phi + N) / ((phi + 1) * N)))
+    mu.item(), std.item()
+
+    return {"mu": mu.item(), "std": std.item()}
 
 
 #%%
diff --git a/src/metaDMG/fit/frequentist.py b/src/metaDMG/fit/frequentist.py
index 9ff439b..a778ccc 100644
--- a/src/metaDMG/fit/frequentist.py
+++ b/src/metaDMG/fit/frequentist.py
@@ -289,14 +289,11 @@ def _get_D_max(self):
         c = self.c
         phi = self.phi
 
-        N = self.N[0]
+        N = max(self.N[0], 1)
 
         mu = A
 
-        if N != 0:
-            std = np.sqrt(A * (1 - A) * (phi + N) / ((phi + 1) * N))
-        else:
-            std = np.nan
+        std = np.sqrt(A * (1 - A) * (phi + N) / ((phi + 1) * N))
 
         return mu, std