Fitting index variables (#156)

* add ptype

* add index fiting

* add test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* rm unneccessary test

* fix bug when ptype is None

* reduce complexity...

* make pre-commit happy

* add index morpher to speed up

* use index morpher

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* shape_index -> index

* Further simplify

* Minor change

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Dacheng Xu <dx2227@columbia.edu>

alea/examples/configs/unbinned_wimp_statistical_model_index_fitting.yaml

@@ -0,0 +1,79 @@
+parameter_definition:
+  wimp_mass:
+    nominal_value: 50
+    fittable: false
+    description: WIMP mass in GeV/c^2
+
+  livetime:
+    nominal_value: 2.
+    ptype: livetime
+    fittable: false
+    description: Livetime in years
+
+  wimp_rate_multiplier:
+    nominal_value: 1.0
+    ptype: rate
+    fittable: true
+    fit_limits:
+      - 0
+      - null
+    parameter_interval_bounds:
+      - 0
+      - 50
+
+  er_rate_multiplier:
+    nominal_value: 1.0
+    ptype: rate
+    uncertainty: 0.2
+    relative_uncertainty: true
+    fittable: true
+    fit_limits:
+      - 0
+      - null
+    fit_guess: 1.0
+
+  er_band_shift:
+    nominal_value: 0
+    ptype: index
+    uncertainty: null
+    fittable: true
+    blueice_anchors:
+      - -2
+      - -1
+      - 0
+      - 1
+      - 2
+    fit_limits:
+      - -2
+      - 2
+    description: ER band shape parameter (shifts the ER band up and down)
+
+likelihood_config:
+  template_folder: null  # will try to find the templates in alea
+  likelihood_terms:
+    - name: science_run
+      default_source_class: alea.template_source.TemplateSource
+      likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
+      likelihood_config: {"morpher": "IndexMorpher"}
+      analysis_space:
+        - cs1: np.arange(0, 102, 2)
+        - cs2: np.geomspace(100, 100000, 51)
+      in_events_per_bin: true
+      livetime_parameter: livetime
+      slice_args: {}
+      sources:
+      - name: er
+        histname: er_template
+        parameters:
+          - er_rate_multiplier
+          - er_band_shift
+        named_parameters:
+          - er_band_shift
+        template_filename: er_template_{er_band_shift}.ii.h5
+
+      - name: wimp
+        histname: wimp_template
+        parameters:
+          - wimp_rate_multiplier
+          - wimp_mass
+        template_filename: wimp50gev_template.ii.h5

alea/model.py

@@ -3,6 +3,7 @@
 from pydoc import locate
 from copy import deepcopy
 from typing import List, Tuple, Callable, Optional, Union
+from itertools import product
 
 import numpy as np
 from scipy.optimize import brentq
@@ -306,7 +307,9 @@ def cost(args):
         return cost
 
     @_needs_data
-    def fit(self, verbose=False, **kwargs) -> Tuple[dict, float]:
+    def fit(
+        self, verbose=False, index_fitting=True, max_index_fitting_iter=10, **kwargs
+    ) -> Tuple[dict, float]:
         """Fit the model to the data by maximizing the likelihood. Return a dict containing best-fit
         values of each parameter, and the value of the likelihood evaluated there. While the
         optimization is a minimization, the likelihood returned is the __maximum__ of the
@@ -337,14 +340,70 @@ def fit(self, verbose=False, **kwargs) -> Tuple[dict, float]:
         for par in fixed_params:
             m.fixed[par] = True
 
-        # Call migrad to do the actual minimization
-        m.migrad()
+        # Get the index variables, which could have problem if simply using migrad
+        index_variables = [
+            p
+            for p in self.parameters.parameters.values()
+            if p.ptype == "index" and p.name not in fixed_params
+        ]
+
+        if (not index_fitting) or (len(index_variables) == 0):
+            # Call migrad to do the actual minimization
+            m = self._migrad_fit(m)
+        else:
+            m = self._migrad_index_mixing_fit(m, index_variables, max_index_fitting_iter, verbose)
+
         self.minuit_object = m
         if verbose:
             print(m)
         # alert! This gives the _maximum_ likelihood
         return m.values.to_dict(), -1 * m.fval
 
+    def _migrad_fit(self, m):
+        m.migrad()
+        return m
+
+    def _migrad_index_mixing_fit(self, m, index_variables, max_index_fitting_iter, verbose):
+        index_anchors = [var.blueice_anchors for var in index_variables]
+        index_names = [var.name for var in index_variables]
+        index_grid = [
+            {index_names[i]: anchor[i] for i in range(len(anchor))}
+            for anchor in product(*index_anchors)
+        ]
+
+        # We fix the index variables in migrad
+        for par in index_names:
+            m.fixed[par] = True
+
+        # We firstly do optimization on other parameters with index variables
+        # fixed to their initial guesses. Then we grid search over the index
+        # variables given the optimized parameters. We repeat the optimization
+        # and grid search until the optimization converges
+        for itr in range(max_index_fitting_iter):
+            m.migrad()
+
+            # Find the best-fit index variables
+            lls = np.zeros(len(index_grid))
+            for i in range(len(lls)):
+                params = m.values.to_dict()
+                params.update(index_grid[i])
+                lls[i] = self.ll(**params)
+            for var in index_names:
+                m.values[var] = index_grid[np.argmax(lls)][var]
+
+            # Calculating Hessian will update the validity of
+            # the fitting given the new index variables
+            m.hesse()
+            if m.valid:
+                break
+
+        if verbose and itr == max_index_fitting_iter - 1:
+            print(
+                "The index searching iteration times reached the maximum! "
+                "The optimization could not converge!"
+            )
+        return m
+
     def _confidence_interval_checks(
         self,
         poi_name: str,

alea/models/blueice_extended_model.py

@@ -263,15 +263,15 @@ def get_source_histograms(self, likelihood_name: str, expected_events=False, **k
 
     def _process_blueice_config(self, config, template_folder_list):
         """Process the blueice config from config."""
-        blueice_config = adapt_likelihood_config_for_blueice(config, template_folder_list)
-        blueice_config["livetime_days"] = self.parameters[
-            blueice_config["livetime_parameter"]
+        pdf_base_config = adapt_likelihood_config_for_blueice(config, template_folder_list)
+        pdf_base_config["livetime_days"] = self.parameters[
+            pdf_base_config["livetime_parameter"]
         ].nominal_value
         for p in self.parameters:
             # adding the nominal rate values will screw things up in blueice!
             # So here we're just adding the nominal values of all other parameters
             if p.ptype != "rate":
-                blueice_config[p.name] = blueice_config.get(p.name, p.nominal_value)
+                pdf_base_config[p.name] = pdf_base_config.get(p.name, p.nominal_value)
 
         # sanity checks
         for source in config["sources"]:
@@ -287,20 +287,31 @@ def _process_blueice_config(self, config, template_folder_list):
 
         # add all parameters to extra_dont_hash for each source unless it is used:
         for i, source in enumerate(config["sources"]):
-            parameters_to_ignore: List[str] = [
-                p.name
-                for p in self.parameters
-                if (p.ptype == "shape") and (p.name not in source["parameters"])
-            ]
-            # no efficiency affects PDF:
-            parameters_to_ignore += [p.name for p in self.parameters if (p.ptype == "efficiency")]
-            parameters_to_ignore += source.get("extra_dont_hash_settings", [])
-
+            parameters_to_ignore = self._get_parameters_to_ignore(source)
             # ignore all shape parameters known to this model not named specifically
             # in the source:
-            blueice_config["sources"][i]["extra_dont_hash_settings"] = parameters_to_ignore
+            pdf_base_config["sources"][i]["extra_dont_hash_settings"] = parameters_to_ignore
+
+        # get blueice likelihood_config if it's given
+        likelihood_config = config.get("likelihood_config", None)
+
+        blueice_config = {
+            "pdf_base_config": pdf_base_config,
+            "likelihood_config": likelihood_config,
+        }
         return blueice_config
 
+    def _get_parameters_to_ignore(self, source):
+        parameters_to_ignore: List[str] = [
+            p.name
+            for p in self.parameters
+            if (p.ptype in ["shape", "index"]) and (p.name not in source["parameters"])
+        ]
+        # no efficiency affects PDF:
+        parameters_to_ignore += [p.name for p in self.parameters if (p.ptype == "efficiency")]
+        parameters_to_ignore += source.get("extra_dont_hash_settings", [])
+        return parameters_to_ignore
+
     def _build_ll_from_config(
         self, likelihood_config: dict, template_path: Optional[str] = None
     ) -> "LogLikelihoodSum":
@@ -326,7 +337,7 @@ def _build_ll_from_config(
             likelihood_class = cast(Callable, locate(config["likelihood_type"]))
             if likelihood_class is None:
                 raise ValueError(f"Could not find {config['likelihood_type']}!")
-            ll = likelihood_class(blueice_config)
+            ll = likelihood_class(**blueice_config)
 
             for source in config["sources"]:
                 # set rate parameters
@@ -354,7 +365,9 @@ def _build_ll_from_config(
 
                 # set shape parameters
                 shape_parameters = [
-                    p for p in source["parameters"] if self.parameters[p].ptype == "shape"
+                    p
+                    for p in source["parameters"]
+                    if self.parameters[p].ptype in ["shape", "index"]
                 ]
                 for p in shape_parameters:
                     anchors = self.parameters[p].blueice_anchors
@@ -513,7 +526,7 @@ def _set_default_ptype(self):
         If no ptype is specified, set the default ptype "needs_reinit".
 
         """
-        allowed_ptypes = ["rate", "shape", "efficiency", "livetime", "needs_reinit"]
+        allowed_ptypes = ["rate", "shape", "index", "efficiency", "livetime", "needs_reinit"]
         default_ptype = "needs_reinit"
         for p in self.parameters:
             if p.ptype is None:

alea/utils.py

@@ -4,6 +4,7 @@
 import yaml
 import importlib_resources
 import itertools
+import blueice
 from glob import glob
 from copy import deepcopy
 from pydoc import locate
@@ -12,6 +13,8 @@
 from base64 import b32encode
 from collections.abc import Mapping
 from typing import Any, List, Dict, Tuple, Optional, Union, cast, get_args, get_origin
+from blueice.pdf_morphers import Morpher
+from itertools import product
 
 import h5py
 import matplotlib.pyplot as plt
@@ -670,3 +673,17 @@ def correction_on_multiplier(x):
         plt.xlabel("Iteration")
         plt.ylabel("x")
     return x
+
+
+class IndexMorpher(Morpher):
+    """IndexMorpher is a morpher which applies no interpolation."""
+
+    def get_anchor_points(self, bounds, n_models=None):
+        grid = [par.keys() for _, (par, _, _) in self.shape_parameters.items()]
+        return list(product(*grid))
+
+    def make_interpolator(self, f, extra_dims, anchor_models):
+        return lambda z: f(anchor_models[tuple(z)])
+
+
+blueice.pdf_morphers.MORPHERS["IndexMorpher"] = IndexMorpher

tests/test_blueice_extended_model.py

@@ -21,6 +21,7 @@ def setUp(cls):
         cls.configs = [
             load_yaml("unbinned_wimp_statistical_model.yaml"),
             load_yaml("unbinned_wimp_statistical_model_simple.yaml"),
+            load_yaml("unbinned_wimp_statistical_model_index_fitting.yaml"),
         ]
         ns = [len(c["likelihood_config"]["likelihood_terms"]) for c in cls.configs]
         cls.n_likelihood_terms = ns