Add Minuit optimiser and update NRSource

flamedisx/inference.py

@@ -4,6 +4,8 @@
 import tensorflow as tf
 import tensorflow_probability as tfp
 import typing as ty
+from iminuit import Minuit
+from iminuit.util import make_func_code
 
 export, __all__ = fd.exporter()
 o = tf.newaxis
@@ -332,17 +334,19 @@ def params_to_dict(self, values):
     def bestfit(self,
                 guess=None,
                 fix=None,
-                optimizer=tfp.optimizer.bfgs_minimize,
+                optimizer= Minuit.from_array_func,
                 llr_tolerance=0.1,
                 get_lowlevel_result=False,
-                use_hessian=True,
+                use_hessian=False,
                 autograph=True,
+                precision=None,
                 **kwargs):
         """Return best-fit parameter tensor
 
         :param guess: Guess parameters: dict {param: guess} of guesses to use.
         :param fix: dict {param: value} of parameters to keep fixed
         during the minimzation.
+        :param optimizer: Minuit.from_array_func or tfp.optimizer.bfgs_minimize
         :param llr_tolerance: stop minimizer if change in -2 log likelihood
         becomes less than this (roughly: using guess to convert to
         relative tolerance threshold)
@@ -360,54 +364,85 @@ def bestfit(self,
             raise ValueError("Guess must be a dictionary")
         guess = {**self.guess(), **guess, **fix}
 
-        # Unfortunately we can only set the relative tolerance for the
-        # objective; we'd like to set the absolute one.
-        # Use the guess log likelihood to normalize;
-        if llr_tolerance is not None:
-            kwargs.setdefault('f_relative_tolerance',
-                              llr_tolerance/self.minus_ll(**guess)[0])
-
         # Create a vector of guesses for the optimizer
         # Store as temp attributes so we can access them also in objective
         self._varnames = [k for k in self.param_names if k not in fix]
         self._guess_vect = fd.np_to_tf(np.array([
             guess[k] for k in self._varnames]))
         self._fix = fix
-
-        # Minimize multipliers to the guess, rather than the guess itself
-        # This is a basic kind of standardization that helps make the gradient
-        # vector reasonable.
-        x_norm = tf.ones(len(self._varnames), dtype=fd.float_type())
-
+
         if optimizer == tfp.optimizer.bfgs_minimize and use_hessian:
             # This optimizer can use the hessian information
             # Compute the inverse hessian at the guess
             inv_hess = self.inverse_hessian(guess, omit_grads=tuple(fix.keys()))
-            # We use scaled values in the optimizer so also scale the
-            # hessian. We need to multiply the hessian with the parameter
-            # values. This is the inverse hessian so we divide.
-            inv_hess /= tf.linalg.tensordot(
-                self._guess_vect, self._guess_vect, axes=0)
             # Explicitly symmetrize the matrix
             inv_hess = fd.symmetrize_matrix(inv_hess)
         else:
             inv_hess = None
 
         self._autograph_objective = autograph
-        res = optimizer(self.objective,
-                        x_norm,
-                        initial_inverse_hessian_estimate=inv_hess,
-                        **kwargs)
-        if get_lowlevel_result:
-            return res
-        if res.failed:
-            raise ValueError(f"Optimizer failure! Result: {res}")
-        res = res.position * self._guess_vect
-        res = {k: res[i].numpy() for i, k in enumerate(self._varnames)}
-        return {**res, **fix}
-
-    def objective(self, x_norm):
-        x = x_norm * self._guess_vect
+
+        if optimizer == Minuit.from_array_func:
+            x_guess =  np.array([
+            guess[k] for k in self._varnames])
+
+            fit = optimizer(self.objective_minuit,
+                           x_guess,
+                           grad=self.grad_minuit,
+                           errordef=0.5,
+                           name = self._varnames,
+                           **kwargs)
+
+            fit.migrad(precision=precision)
+            fit_result = dict(fit.values)
+            fit_errors = dict()
+            for (k, v) in fit.errors.items():
+                fit_errors[k + '_error'] = v
+
+            return fit_result, fit_errors
+
+        else:
+            x_guess = tf.ones(len(self._varnames), dtype=fd.float_type()) \
+                * self._guess_vect
+            # Unfortunately we can only set the relative tolerance for the
+            # objective; we'd like to set the absolute one.
+            # Use the guess log likelihood to normalize;
+            if llr_tolerance is not None:
+                kwargs.setdefault('f_relative_tolerance',
+                                 llr_tolerance/self.minus_ll(**guess)[0])
+
+            res = optimizer(self.objective_tf,
+                            x_guess,
+                            initial_inverse_hessian_estimate=inv_hess,
+                            **kwargs)
+            if get_lowlevel_result:
+                return res
+            if res.failed:
+                raise ValueError(f"Optimizer failure! Result: {res}")
+            res = res.position
+            res = {k: res[i].numpy() for i, k in enumerate(self._varnames)}
+            return {**res, **fix}
+
+    def objective_tf(self, x_guess):
+        x = x_guess #tensor
+
+        # Fill in the fixed variables / convert to dict
+        params = dict(**self._fix)
+        for i, k in enumerate(self._varnames):
+            params[k] = x[i]
+
+        ll, grad = self.minus_ll(
+            **params,
+            autograph=self._autograph_objective,
+            omit_grads=tuple(self._fix.keys()))
+        if tf.math.is_nan(ll):
+            tf.print(f"Objective at {x_guess} is Nan!")
+            ll *= float('inf')
+            grad *= float('nan')
+        return ll, grad 
+
+    def objective_numpy(self, x_guess):
+        x = fd.np_to_tf(x_guess) 
 
         # Fill in the fixed variables / convert to dict
         params = dict(**self._fix)
@@ -419,10 +454,17 @@ def objective(self, x_norm):
             autograph=self._autograph_objective,
             omit_grads=tuple(self._fix.keys()))
         if tf.math.is_nan(ll):
-            tf.print(f"Objective at {x_norm} is Nan!")
+            tf.print(f"Objective at {x_guess} is Nan!")
             ll *= float('inf')
             grad *= float('nan')
-        return ll, grad * self._guess_vect
+        return ll.numpy(), grad.numpy()
+
+    #TODO: make a nice wrapper for objective and gradient
+    def objective_minuit(self, x_guess):
+        return self.objective_numpy(x_guess)[0]
+
+    def grad_minuit(self, x_guess):
+        return self.objective_numpy(x_guess)[1]
 
     def inverse_hessian(self, params, omit_grads=tuple()):
         """Return inverse hessian (square tensor)

flamedisx/x1t_sr1.py

@@ -71,7 +71,7 @@ def electron_gain_std(s2_relative_ly, *, g2=11.4/(1.-0.63)/0.96):
 
     #TODO: implement better the double_pe_fraction or photon_detection_efficiency as parameter
     @staticmethod
-    def photon_detection_eff(s1_relative_ly, g1 =0.123): 
+    def photon_detection_eff(s1_relative_ly, g1 =0.142): 
         #g1 = 0.142 from paper
         mean_eff= g1 / (1. + 0.219)
         return mean_eff * s1_relative_ly
@@ -118,6 +118,46 @@ def s2_acceptance(s2):
 class SR1NRSource(SR1Source, fd.NRSource):
     extra_needed_columns = tuple(set(
         list(SR1Source.extra_needed_columns) + 
-        list(fd.NRSource.extra_needed_columns)))
+        list(fd.NRSource.extra_needed_columns)+
+        ['x_observed', 'y_observed']))
+
+    def p_electron(self, nq, *,
+            alpha=1.280, zeta=0.045, beta=273 * .9e-4, 
+            gamma=0.0141, delta=0.062,
+            drift_field=81):
+        """Fraction of detectable NR quanta that become electrons,
+        slightly adjusted from Lenardo et al.'s global fit
+        (https://arxiv.org/abs/1412.4417).
+
+        Penning quenching is accounted in the photon detection efficiency.
+        """
+        # TODO: so to make field pos-dependent, override this entire f?
+        # could be made easier...
+
+        # prevent /0  # TODO can do better than this
+        nq = nq + 1e-9
+
+        # Note: final term depends on nq now, not energy
+        # this means beta is different from lenardo et al
+        nexni = alpha * drift_field ** -zeta * (1 - tf.exp(-beta * nq))
+        ni = nq * 1 / (1 + nexni)
+
+        # Fraction of ions NOT participating in recombination
+        squiggle = gamma * drift_field ** -delta
+        fnotr = tf.math.log(1 + ni * squiggle) / (ni * squiggle)
+
+        # Finally, number of electrons produced..
+        n_el = ni * fnotr
+
+        return fd.safe_p(n_el / nq)
+
+    #TODO: Define the proper nr spectrum
+    #def _single_spectrum(self):
+    #    """Return (energies in keV, events at these energies),
+    #    both (n_events, n_energies) tensors.
+    #    """   
+    #    e = fd.np_to_tf(spectrum_df['rec_energy'])
+    #    ev = fd.np_to_tf(spectrum_df['rate']) 
+    #    return e, ev   
 
 # TODO: Modify the SR1NRSource to fit AmBe data better and add WIMPSource

requirements.txt

@@ -6,3 +6,4 @@ wimprates
 tqdm
 tensorflow==2.0.0
 tensorflow_probability==0.8.0
+iminuit

tests/test_inference.py

@@ -2,6 +2,8 @@
 import pandas as pd
 import pytest
 import tensorflow as tf
+import tensorflow_probability as tfp
+from iminuit import Minuit
 
 import flamedisx as fd
 from flamedisx.inference import DEFAULT_DSETNAME
@@ -213,7 +215,7 @@ def test_hessian(xes: fd.ERSource):
     assert abs(a - b)/(a+b) < 1e-3
 
 
-def test_bestfit(xes):
+def test_bestfit_tf(xes):
     # Test bestfit (including hessian)
     lf = fd.LogLikelihood(
         sources=dict(er=xes.__class__),
@@ -231,7 +233,31 @@ def test_bestfit(xes):
     guess['er_rate_multiplier'] = xs[np.argmin(ys)]
     assert len(guess) == 2
 
-    bestfit = lf.bestfit(guess, use_hessian=True)
+    bestfit = lf.bestfit(guess, optimizer=tfp.optimizer.bfgs_minimize, use_hessian=True)
     assert isinstance(bestfit, dict)
     assert len(bestfit) == 2
     assert bestfit['er_rate_multiplier'].dtype == np.float32
+
+
+def test_bestfit_minuit(xes):
+    # Test bestfit (including hessian)
+    lf = fd.LogLikelihood(
+        sources=dict(er=xes.__class__),
+        elife=(100e3, 500e3, 5),
+        free_rates='er',
+        data=xes.data)
+
+    guess = lf.guess()
+    # Set reasonable rate
+    # Evaluate the likelihood curve around the minimum
+    xs_er = np.linspace(0.001, 0.004, 20)  # ER source range
+    xs_nr = np.linspace(0.04, 0.1, 20)  # NR source range
+    xs = list(xs_er) + list(xs_nr)
+    ys = np.array([-lf(er_rate_multiplier=x) for x in xs])
+    guess['er_rate_multiplier'] = xs[np.argmin(ys)]
+    assert len(guess) == 2
+
+    bestfit = lf.bestfit(guess, optimizer=Minuit.from_array_func, error = (0.0001,1000))
+    assert isinstance(bestfit[0], dict)
+    assert len(bestfit[0]) == 2
+