add estimate_errors to chi2 analysis

This commit updates the estimate_errors() function so that it works
without a list of constraints and uses arrays of low and high bounds
passed in instead of hardcoded constraints.

I can now call this function from the chi2 analysis to get the stepsizes
before running the MCMC.

4 files changed, 32 insertions, 12 deletions

diff --git a/utils/chi2 b/utils/chi2
index 9970bfd..b1c3c17 100755
--- a/utils/chi2
+++ b/utils/chi2
@@ -31,6 +31,18 @@ from scipy.stats import iqr, norm, beta
 from scipy.special import spence
 from itertools import izip_longest
 from sddm.stats import *
+import contextlib
+from sddm.dc import estimate_errors
+
+# from https://stackoverflow.com/questions/2891790/how-to-pretty-print-a-numpy-array-without-scientific-notation-and-with-given-pre
+@contextlib.contextmanager
+def printoptions(*args, **kwargs):
+    original = np.get_printoptions()
+    np.set_printoptions(*args, **kwargs)
+    try:
+        yield
+    finally: 
+        np.set_printoptions(**original)
 
 # Uncertainty on the energy scale
 # FIXME: These are just placeholders! Should get real number from stopping
@@ -364,9 +376,14 @@ def do_fit(data,muon,data_mc,bins,steps):
     nll_xopt = nll(xopt)
     print("nll(xopt) = ", nll(xopt))
 
+    stepsizes = estimate_errors(nll,xopt,low,high,constraints)
+
+    with printoptions(precision=3, suppress=True):
+        print("Errors: ", stepsizes)
+
     pos = np.empty((20, len(x0)),dtype=np.double)
     for i in range(pos.shape[0]):
-        pos[i] = xopt + np.random.randn(len(x0))*xopt*0.1
+        pos[i] = xopt + np.random.randn(len(x0))*stepsizes
         pos[i,:] = np.clip(pos[i,:],low,high)
 
     nwalkers, ndim = pos.shape
diff --git a/utils/dc b/utils/dc
index a267090..d116c65 100755
--- a/utils/dc
+++ b/utils/dc
@@ -459,7 +459,7 @@ if __name__ == '__main__':
         print("nll(xopt) = ", nll(xopt))
         #print("n = ", opt.get_numevals())
 
-    stepsizes = estimate_errors(nll,xopt,constraints)
+    stepsizes = estimate_errors(nll,xopt,low,high,constraints)
     with printoptions(precision=3, suppress=True):
         print("Errors: ", stepsizes)
 
diff --git a/utils/dc-closure-test b/utils/dc-closure-test
index 72ae95f..bf26662 100755
--- a/utils/dc-closure-test
+++ b/utils/dc-closure-test
@@ -376,7 +376,7 @@ def fit(data, sacrifice, steps):
         print("nll(xopt) = ", nll(xopt))
         #print("n = ", opt.get_numevals())
 
-    stepsizes = estimate_errors(nll,xopt,constraints)
+    stepsizes = estimate_errors(nll,xopt,low,high,constraints)
     with printoptions(precision=3, suppress=True):
         print("Errors: ", stepsizes)
 
diff --git a/utils/sddm/dc.py b/utils/sddm/dc.py
index 42b49b1..1146213 100755
--- a/utils/sddm/dc.py
+++ b/utils/sddm/dc.py
@@ -93,7 +93,7 @@ def get_proposal_func(stepsizes, low, high):
         return x, log_p_x0_given_x - log_p_x_given_x0
     return proposal
 
-def estimate_errors(nll, xopt, constraints):
+def estimate_errors(nll, xopt, low, high, constraints):
     """
     Return approximate 1 sigma errors for each parameter assuming all
     parameters are uncorrelated.
@@ -106,20 +106,23 @@ def estimate_errors(nll, xopt, constraints):
         """
         xtest = xopt.copy()
         xtest[i] = x
-        for constraint in constraints:
-            if constraint(xtest) >= 0:
-                xtest = constraint.renormalize(xtest,i)
+        if constraints:
+            for constraint in constraints:
+                if constraint(xtest) >= 0:
+                    xtest = constraint.renormalize(xtest,i)
         return nll(xtest) - (nll_xopt + 0.5)
 
     errors = np.empty_like(xopt)
 
     for i in range(len(xopt)):
-        if i < 6:
-            xlo = brentq(root,0,xopt[i],args=(xopt,i))
-            xhi = brentq(root,xopt[i],1e9,args=(xopt,i))
+        if np.sign(root(low[i],xopt,i)) == np.sign(nll_xopt):
+            xlo = xopt[i] - low[i]
         else:
-            xlo = brentq(root,0,xopt[i],args=(xopt,i))
-            xhi = brentq(root,xopt[i],1,args=(xopt,i))
+            xlo = brentq(root,low[i],xopt[i],args=(xopt,i))
+        if np.sign(root(high[i],xopt,i)) == np.sign(nll_xopt):
+            xhi = high[i] - xopt[i]
+        else:
+            xhi = brentq(root,xopt[i],high[i],args=(xopt,i))
 
         errors[i] = max(xopt[i]-xlo,xhi-xopt[i])