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)")
def global_hanle_curve_fit( field_data, faraday_rotation_data, lorentzian_count, niter=100, T=100, stepsize=500, threads_for_repeats=8, constant_offset=None, penalize_offset=False, # This precludes regularization regularization=None, measured_offset=None, ): # Construct initial conditions that spread out the widths. We expect some distribution typically # so an evenly spaced distribution of widths allows the fitter to search for a range of widths # in the source data. inv_hwhm_init = 1 / np.arange(0.1, 1, 1 / lorentzian_count) amplitudes_init = np.random.rand(lorentzian_count) if constant_offset is None: # Construct a model with the background constant as a free parameter model = centered_lorentzian_mixture(lorentzian_count=lorentzian_count) init_p = np.zeros(2 * lorentzian_count + 1) init_p[:-1:2] = amplitudes_init else: # Construct a model with the background as a fixed constant (not available to be optimized) model = centered_lorentzian_mixture( lorentzian_count=lorentzian_count, constant_offset=constant_offset ) # TODO: this needs to be properly scaled or its essentially zero, this feature is useless right now init_p = np.zeros(2 * lorentzian_count) init_p[::2] = amplitudes_init init_p[1::2] = inv_hwhm_init # Create scalers to give data zero mean and unit standard deviation scaler_Field = StandardScaler().fit(field_data.reshape(-1, 1)) scaler_FR = StandardScaler().fit(faraday_rotation_data.reshape(-1, 1)) field = scaler_Field.transform(field_data.reshape(-1, 1)) fr = scaler_FR.transform(faraday_rotation_data.reshape(-1, 1)) cost_func_kwargs = {} if measured_offset is not None: measured_offset = scaler_FR.transform(np.array([measured_offset]).reshape(-1, 1))[0] print("using a measured offset penalty for scaled offset %f" % measured_offset) cost_func_kwargs["measured_offset"] = measured_offset if regularization is not None: cost_func_kwargs["regularization"] = regularization p = global_curve_fit( model, field.flatten(), fr.flatten(), init_p, basinhopping_kwargs={"niter": niter, "stepsize": stepsize, "T": T}, cost_func_kwargs=cost_func_kwargs, ) # Extract the parameters from the solution, and rescale the background if constant_offset is None: amplitudes_opt = p.x[:-1:2] offset_opt = p.x[-1] * scaler_FR.scale_[0] + scaler_FR.mean_[0] else: amplitudes_opt = p.x[::2] offset_opt = constant_offset inv_hwhm_opt = np.abs(p.x[1::2]) # Sort the results with the narrow Lorentzians appearing at the lowest indices permute = inv_hwhm_opt.argsort() inv_hwhm_opt = inv_hwhm_opt[permute] amplitudes_opt = amplitudes_opt[permute] # Remove the scaling from the parameters so they are in the units of the source data amplitudes_opt = amplitudes_opt * scaler_FR.scale_[0] inv_hwhm_opt = inv_hwhm_opt / scaler_Field.scale_[0] return {"amplitude": amplitudes_opt, "inv_hwhm": inv_hwhm_opt, "offset": offset_opt, "rms_error": p.fun}