diff options
| -rw-r--r-- | chroma/parabola.py | 89 | ||||
| -rw-r--r-- | test/test_parabola.py | 66 |
2 files changed, 155 insertions, 0 deletions
diff --git a/chroma/parabola.py b/chroma/parabola.py new file mode 100644 index 0000000..0ea9f15 --- /dev/null +++ b/chroma/parabola.py @@ -0,0 +1,89 @@ +import numpy as np +import uncertainties +from uncertainties import unumpy, ufloat +import ROOT + +def build_design_matrix(x, y): + y_invsigma = 1.0/unumpy.std_devs(y) + dims = x.shape[1] + n = 1 + dims + dims*(dims+1)/2 + + A = np.zeros(shape=(len(x),n)) + + A[:,0] = 1.0 * y_invsigma + for i in xrange(dims): + A[:,1+i] = x[:,i] * y_invsigma + + col = 1 + dims + for j in xrange(dims): + for k in xrange(j,dims): + A[:,col] = x[:,j] * x[:,k] * y_invsigma + col += 1 + return A + +def build_design_vector(y): + return unumpy.nominal_values(y)/unumpy.std_devs(y) + +def parabola_fit(points): + dims = points[0][0].shape[0] + + x = np.array([p[0] for p in points]) + f = np.array([p[1] for p in points]) + + A = build_design_matrix(x, f) + B = build_design_vector(f)[:,np.newaxis] # make column vector + + # Compute best-fit parabola coefficients using a singular value + # decomposition. + U, w, V = np.linalg.svd(A, full_matrices=False) + V = V.T # Flip to convention used by Numerical Recipies + inv_w = 1.0/w + inv_w[np.abs(w) < 1e-6] = 0.0 + # Numpy version of Eq 15.4.17 from Numerical Recipies (C edition) + coeffs = np.zeros(A.shape[1]) + for i in xrange(len(coeffs)): + coeffs += (np.dot(U[:,i], B[:,0]) * inv_w[i]) * V[:,i] + + # Chi2 and probability for best fit and quadratic coefficents + chi2_terms = np.dot(A, coeffs[:,np.newaxis]) - B + chi2 = (chi2_terms**2).sum() + ndf = len(points) - (1 + dims + dims * (dims + 1) / 2) + prob = ROOT.TMath.Prob(chi2, ndf) + + # Covariance is from Eq 15.4.20 + covariance = np.dot(V*inv_w**2, V.T) + + # Pack the coefficients into ufloats + ufloat_coeffs = uncertainties.correlated_values(coeffs, covariance.tolist()) + + # Separate coefficients into a, b, and c + a = ufloat_coeffs[0] + b = ufloat_coeffs[1:dims+1] + c = np.zeros(shape=(dims,dims), dtype=object) + index = dims + 1 + for i in xrange(dims): + for j in xrange(i, dims): + c[i,j] = ufloat_coeffs[index] + c[j,i] = ufloat_coeffs[index] + if j != i: + # We combined the redundant off-diagonal parts of c + # matrix, but now we have to divide by two to + # avoid double counting when c is used later + c[i,j] /= 2.0 + c[j,i] /= 2.0 + index += 1 + + return a, np.array(b), c, chi2, prob + +def parabola_eval(x, a, b, c): + if len(x.shape) == 1: + return a + np.dot(x, b) + np.dot(x, np.dot(c, x.T)) + else: + dims = x.shape[1] + y = np.array([a] * x.shape[0]) + + for i, xrow in enumerate(x): + y[i] += np.dot(xrow, b) + y[i] += np.dot(xrow, np.dot(c, xrow.T)) + + return y diff --git a/test/test_parabola.py b/test/test_parabola.py new file mode 100644 index 0000000..7d7e7b2 --- /dev/null +++ b/test/test_parabola.py @@ -0,0 +1,66 @@ +import chroma.parabola as parabola +import numpy +from uncertainties import ufloat, unumpy +import unittest +from numpy.testing import assert_almost_equal + +class Test1D(unittest.TestCase): + def setUp(self): + self.x = numpy.array([[-1.0], [0.0], [1.0]]) + self.y = unumpy.uarray(([2.0, 1.0, 6.0], [0.1, 0.1, 0.1])) + self.a = 1.0 + self.b = numpy.array([2.0]) + self.c = numpy.array([[3.0]]) + + def test_parabola_eval(self): + y = parabola.parabola_eval(self.x, self.a, self.b, self.c) + assert_almost_equal(y, unumpy.nominal_values(self.y)) + + def test_solve(self): + points = zip(self.x, self.y) + a, b, c, chi2, prob = parabola.parabola_fit(points) + + assert_almost_equal(a.nominal_value, 1.0) + assert_almost_equal(b[0].nominal_value, 2.0) + assert_almost_equal(c[0,0].nominal_value, 3.0) + + # Compare to ROOT TGraph fitting uncerts + assert_almost_equal(a.std_dev(), 0.1) + assert_almost_equal(b[0].std_dev(), 0.0707107) + assert_almost_equal(c[0,0].std_dev(), 0.122474, decimal=5) + + +class Test2D(unittest.TestCase): + def setUp(self): + self.x = numpy.array([[-1.0,-1.0], [-1.0, 0.0], [-1.0, 1.0], + [ 0.0,-1.0], [ 0.0, 0.0], [ 0.0, 1.0], + [ 1.0,-1.0], [ 1.0, 0.0], [ 1.0, 1.0]]) + + self.a = 1.0 + self.b = numpy.array([2.0, 3.0]) + self.c = numpy.array([[3.0, 1.0],[1.0, 4.0]]) + + self.y = numpy.zeros(len(self.x), dtype=object) + for i, (x0, x1) in enumerate(self.x): + value = self.a \ + + x0 * self.b[0] + x1 * self.b[1] \ + + x0**2 * self.c[0,0] + x0 * x1 * self.c[0,1] \ + + x1 * x0 * self.c[1,0] + x1**2 * self.c[1,1] + self.y[i] = ufloat((value, 0.1)) + + def test_parabola_eval(self): + y = parabola.parabola_eval(self.x, self.a, self.b, self.c) + assert_almost_equal(y, unumpy.nominal_values(self.y)) + + def test_solve(self): + points = zip(self.x, self.y) + a, b, c, chi2, prob = parabola.parabola_fit(points) + assert_almost_equal(chi2, 0.0) + assert_almost_equal(a.nominal_value, 1.0) + assert_almost_equal(b[0].nominal_value, 2.0) + assert_almost_equal(b[1].nominal_value, 3.0) + assert_almost_equal(c[0,0].nominal_value, 3.0) + assert_almost_equal(c[0,1].nominal_value, 1.0) + assert_almost_equal(c[1,0].nominal_value, 1.0) + assert_almost_equal(c[1,1].nominal_value, 4.0) + |