"""Implementation of fitting algorithms."""
from __future__ import annotations
from typing import List
from iminuit import Minuit # type: ignore
from . import data, theory
[docs]
class Fitter(object):
"""Superclass for fitting procedures/algorithms."""
[docs]
def __init__(self, **kwargs) -> None:
for key in kwargs:
setattr(self, key, kwargs[key])
[docs]
class MinuitFitter(Fitter):
"""Fits using iminuit Python frontend to MINUIT2 C++ library.
Args:
fitpoints (data.DataSet): points to fit to
theory (theory.Theory): theory/model to be fitted
"""
[docs]
def __init__(self, fitpoints: data.DataSet,
theory: theory.Theory, **kwargs) -> None:
self.fitpoints = fitpoints
self.theory = theory
init_vals = [v for v in self.theory.parameters.values()]
names = [k for k in self.theory.parameters.keys()]
def fcn(p):
"""Cost function for minimization - chi-square."""
# Update model parameters.
# Note: This relies on Python>3.7 where dicts are ordered!!
self.theory.parameters.update(zip(self.theory.parameters.keys(), p))
chisq = self.theory.chisq(self.fitpoints)
return chisq
fcn.errordef = Minuit.LEAST_SQUARES
self.minuit = Minuit(fcn, init_vals, name=names)
for k, v in self.theory.parameters_fixed.items():
self.minuit.fixed[k] = v
for k, v in self.theory.parameters_limits.items():
self.minuit.limits[k] = v
Fitter.__init__(self, **kwargs)
[docs]
def covsync(self):
"""Synchronize covariance and theory parameter errors with iminuit."""
self.theory.parameters_errors = self.minuit.errors.to_dict()
# iminuit gives covariance table for all parameters, here
# we take only free ones:
free_pars = self.theory.free_parameters()
try:
self.theory.covariance = {(p1, p2): self.minuit.covariance[p1, p2]
for p1 in free_pars
for p2 in free_pars}
# determine corresponding correlation matrix, but only upper triangular part since
# it is symmetric and we use it only for eye-inspecting correlations
self.theory.correlation = {}
for i, pi in enumerate(free_pars):
for j in range(i+1, len(free_pars)):
pj = free_pars[j]
self.theory.correlation[pi, pj] = self.minuit.covariance.correlation()[pi, pj]
except (AttributeError, TypeError) as error:
print("Something's problematic. No covariances available.")
self.theory.covariance = None # otherwise covariance of previous fit can survive
self.theory.correlation = None
[docs]
def fit(self):
"""Perform simple fit.
For better control, use iminuit's functions.
"""
self.minuit.migrad()
self.covsync()
# The following methods keep status of parameters (fixed, limits)
# in sync between Theory object and minuit object
[docs]
def fix_parameters(self, *args): # noqa: D402
"""fix_parameters('p1', 'p2', ...)."""
if args[0] == 'ALL':
# fix 'em all
self.theory._fix_parameters('ALL')
for par in self.theory.model.parameters:
self.minuit.fixed[par] = True
else:
self.theory._fix_parameters(*args)
for par in args:
self.minuit.fixed[par] = True
[docs]
def release_parameters(self, *args): # noqa: D402
"""release_parameters('p1', 'p2', ...)."""
self.theory._release_parameters(*args)
for par in args:
self.minuit.fixed[par] = False
[docs]
def limit_parameters(self, dct): # noqa: D402
"""limit_parameters({'p1': (lo, hi), ...}."""
self.theory.parameters_limits.update(dct)
for k, v in dct.items():
self.minuit.limits[k] = v
[docs]
def free_parameters(self) -> List[str]:
"""Return list of names of free fitting parameters."""
self.theory.free_parameters()
[docs]
def print_parameters(self):
"""Values and errors for free parameters."""
self.theory.print_parameters()