method.py (1178B)
1 #!/usr/bin/env python3.9 2 3 import numpy as np 4 import sympy as sp 5 6 global x; x = sp.symbols("x") 7 8 def t0_fit(model, var_str, x_data, y_data,\ 9 p0, cov_stat, cov_relsyst, way=0, iterations=100, alpha=0.1): 10 11 var = sp.Matrix(sp.symbols(var_str)) 12 13 jacobi = sp.diff(model(x, *var), var.T) 14 A_func = sp.lambdify((x, *var), jacobi, modules=['numpy', {'Heaviside': np.heaviside}]) 15 16 f2_func = sp.lambdify((x, *var), model(x, *var), modules=['numpy', {'Heaviside': np.heaviside}]) 17 18 chi_sq = [] 19 20 for i in range(iterations): 21 A = A_func(x_data, *p0)[0].T 22 f2 = f2_func(x_data, *p0) 23 24 if way == 0: # consider D'Agostini bias 25 P = np.linalg.inv(cov_relsyst @ np.outer(f2, f2) + cov_stat) 26 else: # don't consider D'Agostini bias 27 P = np.linalg.inv(cov_relsyst * (y_data@y_data) + cov_stat) 28 29 delta_p = np.linalg.inv(A.T @ P @ A) @ A.T @ P @ (y_data - f2) 30 p0 = p0 + alpha*delta_p 31 32 chi_sq.append((y_data - f2).T @ P @ (y_data - f2)) 33 34 print(f"{i} Iterations completed") 35 36 print('Done\n') 37 38 dp = np.sqrt(np.diag(np.linalg.inv(A.T@P@A))) 39 return p0, dp, np.array(chi_sq)