pyBAOfit/mychi2.py at master · Andrei-EPFL/pyBAOfit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import configparser
import numpy as np
from scipy.interpolate import interp1d

from mydata import get_index

def fwd_subst(R, b):
    '''Forward substitution.'''
    n = b.size
    x = np.zeros(n)
    for i in range(n):
        x[i] = b[i]
        for j in range(i):
            x[i] -= R[j, i] * x[j]
        x[i] /= R[i, i]
    return x

def bwd_subst(R, b):
    '''Backward substitution.'''
    n = b.size
    x = np.zeros(n)
    for i in range(n-1, -1, -1):
        x[i] = b[i]
        for j in range(i+1, n):
            x[i] = x[i] - R[i][j] * x[j]
        x[i] /= R[i][i]
    return x

def init_lstsq(sd, Rcov, npoly, imin, imax):
    '''Initialize the least square fitting for nuisance parameters.
    Refs:
      Numerical Recipes 3rd edition, section 15.4.
      (But the matrix manipulations are simplified here.)
    Arguments:
      sd: the s array for data;
      Rcov: the upper triangular matrix from QR decomp of cov;
      npoly: the number of nuisance parameters;
      imin: the minimum index for the fitting range;
      imax: the maximum index for the fitting range.
    Return: [basis function, design matrix, RHS matrix for LS].'''
    # The basis function and design maxtrix
    nd = sd.size
    basis = np.zeros([npoly, nd])
    A = np.zeros([npoly, nd])
    for i in range(npoly):
        basis[i, imin:imax] = sd[imin:imax]**(i - 2)
        A[i, :] = fwd_subst(Rcov, basis[i, :])

    # Construct U from the QR decomposition of A
    U = np.linalg.qr(np.transpose(A), mode='r')
    # Construct M = U^{-T} . A^T
    M = np.zeros([npoly, nd])
    col = np.zeros(npoly)
    for i in range(nd):
        col = fwd_subst(U, A[:, i])
        M[:, i] = col
    # Construct M . R^{-T}
    for i in range(npoly):
        M[i, :] = bwd_subst(Rcov, M[i, :])

    return [basis, U, M]

class Chi2Class():
    def __init__(self, xim_var, xid_var, covmat_var):

        self.covmat = covmat_var
        self.xim = xim_var
        self.xid = xid_var

        self.list_params = self.xim.list_params
        self.chi2_norm = self.covmat.nmock - self.xid.nidx - 2

        print('STATUS: Initialize matrices for least square fitting of the nuisance params.')
        self.basis, self.A, self.M = init_lstsq(self.xid.sd, self.covmat.Rcov, self.xim.npoly, self.xid.imin, self.xid.imax)

    def chi2_func(self, alpha, params):
        '''Define the (log) likelihood.'''
        B = params[0]
        imin = self.xid.imin
        imax = self.xid.imax
        # Compute the model 2PCF with a given alpha
        fxim = interp1d(self.xim.sm, self.xim.xi_model(params), kind='cubic')
        xi = fxim(self.xid.sd[imin:imax] * alpha)

        # Least square fitting of nuisance parameters
        dxi = self.xid.xid[imin:imax] - xi * (B**2)
        poly = np.dot(self.M[:, imin:imax], dxi)

        a_poly = bwd_subst(self.A, poly)

        # Compute chi-squared
        diff = np.zeros(self.xid.ndbin)
        diff[imin:imax] = dxi - np.dot(np.transpose(self.basis[:, imin:imax]), a_poly)
        chisq = np.sum((fwd_subst(self.covmat.Rcov, diff))**2)
        return chisq * self.chi2_norm * self.covmat.rescale_chi2_by

    def best_fit(self, alpha, params):
        '''Compute the best-fit theoretical Xi curve.'''
        B = params[0]
        imin = self.xid.imin
        imax = self.xid.imax
        sm = self.xim.sm

        # Compute the model 2PCF with a given alpha
        fxim = interp1d(sm, self.xim.xi_model(params), kind='cubic')
        fxim_nw = interp1d(sm, self.xim.xi_model_nw(params), kind='cubic')
        xi = fxim(self.xid.sd[imin:imax] * alpha)
        xi_nw = fxim_nw(self.xid.sd[imin:imax] * alpha)

        # Least square fitting of nuisance parameters
        dxi = self.xid.xid[imin:imax] - xi * (B**2)
        poly = np.dot(self.M[:, imin:imax], dxi)
        a_poly = bwd_subst(self.A, poly)

        # Compute best-fit
        if alpha >= 1:
            imin0, imax0, nidx0 = get_index(sm, sm[1], sm[-2] / alpha)
        else:
            imin0, imax0, nidx0 = get_index(sm, sm[1] / alpha, sm[-2])

        best = B**2 * fxim(sm[imin0:imax0] * alpha)
        best_nw = B**2 * fxim_nw(sm[imin0:imax0] * alpha)
        best_broadband = np.zeros(nidx0)
        for i in range(self.xim.npoly):
            best += a_poly[i] * sm[imin0:imax0]**(i - 2)
            best_nw += a_poly[i] * sm[imin0:imax0]**(i - 2)
            best_broadband += a_poly[i] * sm[imin0:imax0]**(i - 2)
        return [sm[imin0:imax0], best, best_nw, best_broadband]