diff --git a/pxii_bec/bec_ipython_client/plugins/alignment_fit.py b/pxii_bec/bec_ipython_client/plugins/alignment_fit.py new file mode 100644 index 0000000..5cd0418 --- /dev/null +++ b/pxii_bec/bec_ipython_client/plugins/alignment_fit.py @@ -0,0 +1,49 @@ +import numpy as np +from scipy.ndimage import gaussian_filter1d +from lmfit.models import GaussianModel + + +def alignment_fit_and_plot( + history_index: int, + device_name: str, + signal_name: str | None = None, + smoothing_sigma: float = 2.0, +): + """ + Get data for a completed scan from the BEC history, apply smoothing, gaussian fit, + gradient, and plot all the results. + + Args: + history_index (int): scan to fetch, e.g. -1 for the most recent scan + device_name (str): the device for which to get the monitoring data + + """ + # Fetch scan data from the history + # by default, signal = device name, unless otherwise specified + signal = signal_name or device_name + scan = bec.history[history_index] + md = scan.metadata["bec"] + data = scan.devices[device_name][signal].read()["value"] + # motor name is a bytes object in the metadata, so make a string + motor_name = md["scan_motors"][0].decode() + + # Create a plot and a text box to display results + dock_area = bec.gui.new() + wf = dock_area.new().new(bec.gui.available_widgets.Waveform) + wf.title = f"Scan {md['scan_number']}: {md['scan_name']} of {motor_name}" + text = dock_area.new(position="right").new(widget=bec.gui.available_widgets.TextBox) + + # Calculate some processed data and add everything to the plot + wf.plot(data, label="Raw data") + smoothed_data = gaussian_filter1d(data, smoothing_sigma) + wf.plot(smoothed_data, label="Smoothed") + gradient = np.gradient(smoothed_data) + wf.plot(gradient, label="gradient") + + # Fit a Gaussian model to the smoothed data and show the fitting parameters in the textbox + x_data = scan.devices[motor_name][motor_name].read()["value"] + model = GaussianModel() + result = model.fit(smoothed_data, x=x_data) + text.set_plain_text(f"Fit parameters: \n{result.params.pretty_repr()}") + + return result