import numpy as np
import scipy.odr as odr
from matplotlib.table import Table
from matplotlib.font_manager import FontProperties
import matplotlib.pyplot as plt
[docs]
def linear_func(params, x):
m, b = params
return m * x + b
[docs]
def evaluate_with_uncertainties(x_new, x_new_err, slope, intercept, slope_err, intercept_err):
# Predicted y value
y_new = slope * x_new + intercept
# Uncertainty in predicted y value (including uncertainty in x_new)
y_err_new = np.sqrt(
(x_new * slope_err)**2 + # Uncertainty due to slope
intercept_err**2 + # Uncertainty due to intercept
(slope * x_new_err)**2 # Uncertainty due to x_new
)
return y_new, y_err_new
[docs]
def make_linear_fit(x,y,y_error=None,x_error=None,plot=False,names=[],title="title",save="",add_table=False,xlabel=r"$r_{s}$(thin disk) [lday]",ylabel=r"$r_{s}$(microlensing) [lday]"):
linear_model = odr.Model(linear_func)
data = odr.RealData(x, y, sx=x_error, sy=y_error)
# Set initial guesses for slope (m) and intercept (b)
beta0 = [1.0, 1.0] # initial guess for [slope, intercept]
# Create ODR object and run the fitting
odr_fit = odr.ODR(data, linear_model, beta0=beta0)
output = odr_fit.run()
# Print the results
slope, intercept = output.beta
slope_err, intercept_err = output.sd_beta
print(f"Slope: {slope} ± {slope_err}")
print(f"Intercept: {intercept} ± {intercept_err}")
# Generate the fit line
mfc = "k"
if add_table:
mfc = 'white'
if plot:
x_fit = np.linspace(min(x)*0.95, max(x)*1.1, 100)
y_fit = linear_func((slope, intercept),x_fit)
# Propagate uncertainty across the entire fit line
y_fit_upper = linear_func((slope + slope_err, intercept + intercept_err),x_fit)
y_fit_lower = linear_func((slope - slope_err, intercept - intercept_err),x_fit)
fig, ax = plt.subplots(figsize=(20, 10))
ax.errorbar(x, y, yerr=y_error, xerr=x_error, fmt='o', markersize=10, mfc=mfc, mec='k', capsize=5,ecolor="k")
#.errorbar(x, y,xerr=x_error,yerr=y_error, fmt='o',c="k", capsize=2, s=500, facecolors='none', edgecolors='blue')
ax.plot(x_fit,y_fit, linewidth=1, color='r', label=r"Linear regression", ls='--', zorder=0)
ax.fill_between(x_fit,y_fit_lower,y_fit_upper,
color="red",label=r'$\sigma$ uncertainty',alpha=0.3)
if add_table:
if len(names) ==0:
names = [str(i+1) for i in range(len(x))]
names = {i+1:name for i,name in enumerate(names)}
for i, (num,name) in enumerate(names.items()):
ax.text(x[i], y[i], str(num), fontsize=8, ha='center', va='center_baseline',zorder=10)
# Create a table for the definitions
table_data = [[str(num), fr"{definition}|$r_{{td}}=${x[num-1]:.4f} $\pm$ {x_error[num-1]:.4f} vs $r_{{m}}=${y[num-1]:.4f} $\pm$ {y_error[num-1]:.4f}"] for num, definition in zip(names.keys(), names.values())]
# Add the table on the side
table = Table(ax, bbox=[1.1, 0., 0.5, 1]) # Adjust bbox to control position and size of the table
for i, (num, definition) in enumerate(table_data):
table.add_cell(i, 0, width=0.2, height=0.1, text=num, loc='center')
table.add_cell(i, 1, width=0.5, height=0.1, text=definition.split("|")[0], loc='left')
table.add_cell(i, 2, width=1.5, height=0.1, text=definition.split("|")[1], loc='left')
# Add the table to the plot
ax.add_table(table)
# Show plot
#ax.plot([12,17], [12,17], color='grey', label=r'$R_s=R_s$',alpha=0.8,ls='--',zorder=3)
#plt.xlim(12.6,16)
#plt.ylim(14,17.5)
#plt.text(15,13,size=20)
#
ax.set_xlabel(xlabel, fontsize=20)
ax.set_ylabel(ylabel, fontsize=20)
ax.tick_params(axis='both', which='major', labelsize=20)
ax.legend(fontsize=20)
ax.set_xlim(x_fit[[0,-1]])
ax.set_ylim(y_fit[[0,-1]])
ax.set_title(title)
plt.tight_layout()
if len(save)>0:
plt.savefig(f"{save}.png",bbox_inches='tight')
plt.show()
return slope, intercept, slope_err, intercept_err