@@ -55,7 +55,7 @@ def chisq_BAO_alam(self,theta):
5555
5656 return np .dot (dist - distance_theory , np .dot (np .linalg .inv (cov ), dist - distance_theory ))
5757
58- def chisq_Pantheon_plus (self ,theta ):
58+ def chisq_Pantheon_plus (self ,theta , residual = False ):
5959 cmb_z , mb , covariance = self .db .get_pantheon_plus ()
6060 helio_z = self .db .get_pantheon_plus (Zhel = True )
6161 distance_theory = self .theory .distance_modulus (theta , cmb_z , helio_z )
@@ -66,9 +66,12 @@ def chisq_Pantheon_plus(self,theta):
6666 icov = np .linalg .inv (covariance )
6767 self .inv_covariance ['Pantheon_plus' ] = icov
6868
69- return np .dot (delta , np .dot (icov , delta ))
69+ if residual :
70+ return delta , np .sqrt (np .diag (covariance ))
71+ else :
72+ return np .dot (delta , np .dot (icov , delta ))
7073
71- def chisq_Pantheon_old (self ,theta ):
74+ def chisq_Pantheon_old (self ,theta , residual = False ):
7275 cmb_z , mb , covariance = self .db .get_pantheon_old ()
7376 helio_z = self .db .get_pantheon_old (Zhel = True )
7477 distance_theory = self .theory .distance_modulus (theta , cmb_z , helio_z )
@@ -78,17 +81,22 @@ def chisq_Pantheon_old(self,theta):
7881 else :
7982 icov = np .linalg .inv (covariance )
8083 self .inv_covariance ['Pantheon_old' ] = icov
81-
82- return np .dot (delta , np .dot (icov , delta ))
84+ if residual :
85+ return delta , np .sqrt (np .diag (covariance ))
86+ else :
87+ return np .dot (delta , np .dot (icov , delta ))
8388
84- def chisq_QSO_dm (self ,theta ):
89+ def chisq_QSO (self ,theta , residual = False ):
8590 p = self .params (theta )
8691 #pdb.set_trace()
8792 z ,dm ,cov = self .db .get_QSO ()
88- distance_theory = self .theory .distance_modulus (theta , z ,z )
93+ distance_theory = self .theory .distance_modulus (theta , z ,z ) - p . M_b
8994 delta = dm - distance_theory
9095 var = np .diag (cov )
91- return np .sum (delta ** 2. / (var + p .qso_sigma ** 2 ))
96+ if residual :
97+ return delta , np .sqrt (var )
98+ else :
99+ return np .sum (delta ** 2. / (var + p .qso_sigma ** 2 ))
92100
93101 def chisq_QSO_full (self ,theta ):
94102 p = self .params (theta )
@@ -147,13 +155,20 @@ def __init__(self, data, GP_filename=None, cosmo_filename=None):
147155 self .nmodel = self .nTasks
148156 self .self_scale = MTGP ['self_scale' ]
149157 self .scatter = MTGP ['sigma_int' ]
158+ self .offset = MTGP ['offset' ]
159+
160+ # extra paramter per dataset to be added to the theta
161+ self .ind_scatter = MTGP ['ind_scatter' ]
150162
151163 # To define the priors for the GP hyperparameters using the the GPconfig file
152164 self .priors = []
153165 for i in range (self .nTasks ):
154166 i = i + 1 # To start from 1
155167 self .priors .append ([MTGP [f'Task_{ i } ]['sigma_f_min' ], MTGP [f'Task_{ i } ]['sigma_f_max' ]])
156168 self .priors .append ([MTGP [f'Task_{ i } ]['l_min' ], MTGP [f'Task_{ i } ]['l_max' ]])
169+ if self .ind_scatter :
170+ self .priors .append ([MTGP [f'Task_{ i } ]['sig_int_min' ], MTGP [f'Task_{ i } ]['sig_int_max' ]])
171+
157172
158173 self .priors = np .array (self .priors )
159174
@@ -184,44 +199,68 @@ def set_theta(self, theta):
184199 # and possibly with intrinsic scatter at the end [..., sigma_int]
185200 # Final form is [[sigma_f_1, l_1], [sigma_f_2, l_2], ...]
186201
187- if self .self_scale == False and len (theta ) != 2 * self .nTasks and self .scatter == False :
188- raise ValueError ('The number of hyperparameters does not match the number of tasks' )
189- elif self .self_scale == False and len (theta ) != 2 * self .nTasks + 1 and self .scatter == True :
190- pass
202+
203+ if self .self_scale == False :
204+ params_GP = []
205+ if len (theta )>= 2 * self .nTasks :
206+ for i in range (self .nTasks ):
207+ params_GP .append ([theta [2 * i ], theta [2 * i + 1 ]])
208+ else :
209+ raise ValueError (f'The number of hyperparameters does not match the number of tasks, there must be atleast { 2 * self .nTasks } )
210+
211+ elif self .self_scale == True :
212+ params_GP = []
213+ if len (theta )>= self .nTasks + 1 :
214+ for i in range (self .nTasks ):
215+ params_GP .append ([theta [i ], theta [self .nTasks ]])
216+ else :
217+ raise ValueError (f'The number of hyperparameters does not match the number of tasks, there must be atleast { self .nTasks + 1 } )
218+
219+
220+ if self .ind_scatter == True and len (theta )>= self .nTasks + 2 :
221+ params_scatter = theta [- self .nTasks :]
191222 else :
192- # rearrange the theta values into a list of lists
193- params = []
194- for i in range (self .nTasks ):
195- params .append ([theta [2 * i ], theta [2 * i + 1 ]])
223+ params_scatter = [self .scatter ]* self .nTasks
196224
197- return params
225+ return params_GP , params_scatter
198226
199227 def logPrior (self , theta ):
200228 # To set the priors for the GP hyperparameters
201229 for (lower , upper ), value in zip (self .priors , theta ):
202230 if not lower < value < upper :
203- # print(value, lower, upper)
204- # sys.exit()
205231 return - np .inf
206232 return 0.0
207233
208- def logLike (self ,theta ):
209- chi2 = self .chisq (theta )
234+ def logLike (self ,theta_GP , theta_scatter ):
235+ chi2 = self .chisq (theta_GP , theta_scatter )
210236 return - 0.5 * chi2
211237
212238 def logProb (self , theta ):
239+
213240 lp = self .logPrior (theta )
214- theta = self .set_theta (theta )
241+ theta_GP , theta_scatter = self .set_theta (theta )
215242 if not np .isfinite (lp ):
216243 return - np .inf
217- return lp + self .logLike (theta )
244+ return lp + self .logLike (theta_GP , theta_scatter )
218245
219- def chisq (self ,theta ):
246+ def chisq (self , theta_GP , theta_scatter ):
220247 # We define the GP likelihood here
221248 chi2 = 0
249+
250+ # To set the scatter covariance matrix which is a diagonal matrix of size len of ntasks
251+ scatter_matrix = self .offset * np .eye (len (self .D_covmat ))
252+ scatter_diag = []
253+ for i in range (self .nTasks ):
254+ # create arrays of length of the data for each task
255+ scatter_diag = np .append (scatter_diag , len (np .array (self .d .x [f'x{ i } ]))* [theta_scatter [i ]** 2. ])
256+
257+ int_scatter_covmat = scatter_matrix + np .diag (scatter_diag )
258+
222259 # To call the GP class to get the covariance matrix
223- GP_covmat = self .GP_kernels .Cov_Mat (theta )
224- Total_covmat = GP_covmat + self .D_covmat + 1e-6 * np .eye (len (self .D_covmat ))
260+ GP_covmat = self .GP_kernels .Cov_Mat (theta_GP )
261+
262+ # To calculate the total covariance matrix
263+ Total_covmat = GP_covmat + self .D_covmat + int_scatter_covmat
225264
226265 # To perform cholesky decomposition
227266 L_inv = np .linalg .inv (Total_covmat )
@@ -235,6 +274,4 @@ def chisq(self,theta):
235274 y_inv = np .dot (L_inv , y )
236275 chi2 = np .dot (y , y_inv .T ) + log_det - offset
237276
238- return chi2
239-
240-
277+ return chi2
0 commit comments