From f2d27ff6b1c16a18736e9d2ac3078504e046abd6 Mon Sep 17 00:00:00 2001
From: Christian Michelsen <christianmichelsen@gmail.com>
Date: Wed, 17 Aug 2022 13:33:04 +0200
Subject: [PATCH] fix: Handle empty data and better error handling

---
 src/metaDMG/fit/fits.py        | 27 +++++++++++++++++---
 src/metaDMG/fit/frequentist.py | 45 ++++++++++++++++++++++++----------
 2 files changed, 56 insertions(+), 16 deletions(-)

diff --git a/src/metaDMG/fit/fits.py b/src/metaDMG/fit/fits.py
index 0637bfc..b92ef64 100644
--- a/src/metaDMG/fit/fits.py
+++ b/src/metaDMG/fit/fits.py
@@ -11,7 +11,7 @@
 from logger_tt import logger
 from tqdm import tqdm
 
-from metaDMG.errors import BadDataError
+from metaDMG.errors import BadDataError, FittingError
 from metaDMG.fit import bayesian, fit_utils, frequentist
 from metaDMG.utils import Config
 
@@ -141,6 +141,16 @@ def fit_single_group(
     sample = config["sample"]
     tax_id = group["tax_id"].iloc[0]
 
+    if data["N"].sum() == 0:
+        from metaDMG.fit.serial import _setup_logger
+
+        _setup_logger(config)
+
+        s = f"No data for tax_id {tax_id}, sample = {sample}, i.e. sum(N) = 0. "
+        s += "Skipping the fit."
+        logger.warning(s)
+        return None
+
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore")
 
@@ -155,7 +165,15 @@ def fit_single_group(
     add_count_information(fit_result, config, group, data)
 
     if mcmc_PMD is not None and mcmc_null is not None:
-        bayesian.make_fits(fit_result, data, mcmc_PMD, mcmc_null)
+        try:
+            bayesian.make_fits(fit_result, data, mcmc_PMD, mcmc_null)
+        except:
+            from metaDMG.fit.serial import _setup_logger
+
+            _setup_logger(config)
+            s = f"Error in bayesian.make_fits for tax_id {tax_id}, sample = {sample}."
+            s += "Skipping the Bayesian fit."
+            logger.warning(s)
 
     return fit_result
 
@@ -177,13 +195,16 @@ def compute_fits_seriel(config, df_mismatches, with_progressbar=False):
         if with_progressbar:
             groupby.set_description(f"Fitting Tax ID {tax_id}")
 
-        d_fit_results[tax_id] = fit_single_group(
+        res = fit_single_group(
             config,
             group,
             mcmc_PMD,
             mcmc_null,
         )
 
+        if res is not None:
+            d_fit_results[tax_id] = res
+
     return d_fit_results
 
 
diff --git a/src/metaDMG/fit/frequentist.py b/src/metaDMG/fit/frequentist.py
index a778ccc..2c9c193 100644
--- a/src/metaDMG/fit/frequentist.py
+++ b/src/metaDMG/fit/frequentist.py
@@ -163,50 +163,69 @@ def _setup_minuit(self, m=None):
         self.m.errordef = Minuit.LIKELIHOOD
 
     def fit(self):
+        if self.verbose:
+            print("Initial fit")
         self.m.migrad()
         self.is_fitted = True
+        if self.m.valid and self.verbose:
+            print("Valid fit")
 
         # First try to refit it
         if not self.m.valid:
             if self.verbose:
-                print("refitting A")
+                print("Refitting up to 10 times using last values as p0")
             for i in range(10):
                 self.m.migrad()
                 if self.m.valid:
-                    if self.verbose:
-                        print(f"Got out, A {i}")
                     break
+        if self.m.valid and self.verbose:
+            print("Valid fit")
 
         # Then try with a totally flat guess
         if not self.m.valid:
             if self.verbose:
-                print("refitting B")
-            p0_flat = {"q": 0.0, "A": 0.0, "c": 0.01, "phi": 100}
-            self._setup_p0(p0_flat)
+                print("Refitting using a flat p0")
+            self._setup_p0({"q": 0.0, "A": 0.0, "c": 0.01, "phi": 100})
+            self._setup_minuit()
             self.m.migrad()
+        if self.m.valid and self.verbose:
+            print("Valid fit")
+
+        # Also try with the default guess
+        if not self.m.valid:
             if self.verbose:
-                print(f"Got out, B")
+                print("Refitting using the default p0")
+            self._setup_p0({"q": 0.1, "A": 0.1, "c": 0.01, "phi": 1000})
+            self._setup_minuit()
+            self.m.migrad()
+        if self.m.valid and self.verbose:
+            print("Valid fit")
 
         # If not working, continue with new guesses
         MAX_GUESSES = 100
         if not self.m.valid:
-            self.i = 0
+
             if self.verbose:
-                print("refitting C")
+                print("Getting desperate, trying random p0's")
+
+            self.i = 0
             while True:
                 p0 = fit_utils.sample_from_param_grid(self.param_grid)
-                for key, val in p0.items():
-                    self.m.values[key] = val
+                self._setup_p0(p0)
+                self._setup_minuit()
                 self.m.migrad()
+
                 if self.m.valid or self.i >= MAX_GUESSES:
                     break
                 self.m.migrad()
+
                 if self.m.valid or self.i >= MAX_GUESSES:
                     break
+
                 self.i += 1
 
-            if self.m.valid and self.verbose:
-                print(f"Got out, C {self.i}")
+        if self.m.valid and self.verbose:
+            print(f"Valid fit, number of tries = {self.i}")
 
         self.valid = self.m.valid