def double_lorentzian_fig(p, xs, ys, xm=[], ym=[], title=''):
"""
Creates a figure that plots the components of the centered double Lorentzian model, the sum of
the components, the data and a single Lorentzian model for comparison.
:param p: The centered double Lorentzian to plot
:param xs: x values of the data to plot and fit
:param ys: y values of the data to plot and fit
:param title: the title of the plot
:return: a figure that can be saved or shown
"""
ps, conv = opt.curve_fit(lorentzian, xs, ys, [p[0], p[1], 0, p[4]], maxfev=4000)
plt.ioff()
f2, ax2 = plt.subplots()
ax2.plot(xs, ys, 'ok', rasterized=True, alpha=0.3)
if (len(xm) > 2 ) and (len(ym) > 2):
ax2.plot(xm, ym, '-w')
ax2.plot(xs, [double_lorentzian_centered(x, *p) for x in xs], linewidth=3)
ax2.plot(xs, [lorentzian(x, *[p[0], p[1], 0, 0]) for x in xs], linewidth=3, label='peak 1')
ax2.plot(xs, [lorentzian(x, *[p[2], p[3], 0, 0]) for x in xs], linewidth=3, label='peak 2')
ax2.plot(xs, [lorentzian(x, *ps) for x in xs], linewidth=3, label='single')
plt.legend()
ax2.set_yticklabels(ax2.get_yticks()/1e-6)
ax2.set_ylabel('Faraday Rotation ($\mu rad$)')
ax2.set_xlabel('Field (Gauss)')
ax2.set_title(title)
plt.ion()
return f2
def plot_hanle_curve(hanle_curve_data, hanle_curve_fit_params, plot_data=True, plot_mean=True, plot_residual=True):
f2, ax2 = plt.subplots()
# Plot the data
if plot_data:
ax2.plot(
hanle_curve_data.Field,
hanle_curve_data.FR,
"o",
color=sb.xkcd_rgb["black"],
rasterized=True,
alpha=0.3,
label="Raw Data",
)
# Plot the curve mean
if plot_mean:
if plot_data:
color = sb.xkcd_rgb["mango"]
label = "Average"
else:
color = sb.xkcd_rgb["black"]
label = None
dm = hanle_curve_data.groupby("Field")
ax2.plot(dm.Field.mean(), dm.FR.mean(), "o", color=color, alpha=0.5, markersize=4, label=label)
field_grid = np.linspace(np.min(hanle_curve_data.Field), np.max(hanle_curve_data.Field), 1000)
# Plot the multiple lorentzian
count = len(list(hanle_curve_fit_params["amplitude"]))
model = centered_lorentzian_mixture(count)
params = np.zeros(2 * count + 1)
params[:-1:2] = hanle_curve_fit_params["amplitude"]
params[1:-1:2] = hanle_curve_fit_params["inv_hwhm"]
params[-1] = hanle_curve_fit_params["offset"]
if plot_residual:
ax2.plot(
dm.Field.mean(),
10 * (np.array([model(x, *params) for x in dm.Field.mean()]).flatten() - dm.FR.mean()),
color=sb.xkcd_rgb["dark yellow"],
linewidth=3,
label="Residual (10x)",
)
ax2.plot(
field_grid, [model(x, *params) for x in field_grid], color=sb.xkcd_rgb["tomato red"], linewidth=3, label="Fit"
)
colors = [sb.xkcd_rgb["cobalt"], sb.xkcd_rgb["azure"]] + list(sb.xkcd_rgb.values())
for i in range(0, count):
ax2.plot(
field_grid,
[lorentzian(x, *[params[0 + 2 * i], params[1 + 2 * i], 0, 0]) for x in field_grid],
color=colors[i],
linewidth=2,
label=("peak %d" % (i + 1)),
)
ax2.set_yticklabels(ax2.get_yticks() / 1e-6)
plt.legend()
ax2.set_ylabel("Faraday Rotation ($\mu rad$)")
ax2.set_xlabel("Field (Gauss)")