From db244b493659fab7fa460f1b2c91d4a4a859bf74 Mon Sep 17 00:00:00 2001 From: x10sa Date: Thu, 8 Oct 2026 12:09:42 +0200 Subject: [PATCH] updates, bug fixes etc --- pxii_bec/macros/analyse_luts.py | 583 ++++++++++++++++++++++++++++++++ 1 file changed, 583 insertions(+) create mode 100644 pxii_bec/macros/analyse_luts.py diff --git a/pxii_bec/macros/analyse_luts.py b/pxii_bec/macros/analyse_luts.py new file mode 100644 index 0000000..dfaf726 --- /dev/null +++ b/pxii_bec/macros/analyse_luts.py @@ -0,0 +1,583 @@ +"""Analyse LUT data to determine fits for DCM and undulator parameters.""" + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +from scipy.signal import savgol_filter + +from pxii_bec.scripts import beamline_context as bl + + + +def get_data(filename): + """Open the raw DCM calibration data file.""" + + filename = bl.env.lut_dir / f"{filename}.csv" + + return pd.read_csv(filename) + + +def format_polynomial(fit): + """Return polynomial coefficients as a readable equation string.""" + + terms = [] + order = len(fit) - 1 + + for i, coeff in enumerate(fit): + power = order - i + + if power > 1: + term = f"{coeff:+.3e} x^{power}" + elif power == 1: + term = f"{coeff:+.3e} x" + else: + term = f"{coeff:+.4f}" + + terms.append(term) + + return " ".join(terms).lstrip("+") + + +def plot_lut( + ax, + arr, + fit_order, + xname=None, + yname=None, + title=None, + show_equation=True, +): + """ + Fit LUT data and draw the data and fit onto an existing axes. + + Parameters + ---------- + ax : matplotlib.axes.Axes + Axes on which to draw the plot. + arr : numpy.ndarray + Two-column array containing x and y data. + fit_order : int + Polynomial order to fit. + xname : str, optional + X-axis label. + yname : str, optional + Y-axis label. + title : str, optional + Plot title. + show_equation : bool, optional + Display the fitted polynomial on the plot. + + Returns + ------- + numpy.ndarray + Polynomial coefficients. + """ + + x = arr[:, 0] + y = arr[:, 1] + + fit = np.polyfit(x, y, fit_order) + p = np.poly1d(fit) + + x_range = np.linspace(x.min(), x.max(), 100) + + ax.scatter( + x, + y, + label="Data", + color="rebeccapurple", + ) + + ax.plot( + x_range, + p(x_range), + label="Polynomial fit", + color="lightcoral", + ) + + if xname is not None: + ax.set_xlabel(xname) + + if yname is not None: + ax.set_ylabel(yname) + + if title is not None: + ax.set_title(title) + + if show_equation: + equation = format_polynomial(fit) + + ax.text( + 0.02, + 0.02, + equation, + transform=ax.transAxes, + fontsize=9, + verticalalignment="bottom", + bbox=dict( + boxstyle="round", + facecolor="whitesmoke", + edgecolor="lightgray", + alpha=0.8, + ), + ) + + return fit + + +def lut_plot( + arr, + xname, + yname, + title, + fit_order, + save=False, +): + """ + Create a standalone LUT plot. + + This is useful for calibration plots such as DCM pitch where only + one fit is required. + """ + + savefile = bl.env.lut_dir / f"{title}.png" + + fig, ax = plt.subplots() + + fit = plot_lut( + ax, + arr, + fit_order, + xname=xname, + yname=yname, + title=title, + show_equation=True, + ) + + ax.legend(fontsize=9) + + fig.tight_layout() + + if save: + fig.savefig(savefile, bbox_inches="tight") + plt.close(fig) + else: + plt.show() + + print(np.poly1d(fit)) + + return fit + + +def pitch_analysis(filename, fit_order=3, save=False): + """Fit and plot the DCM pitch calibration.""" + + df = get_data(filename) + + pitch = df[["energy", "pitch"]] + + fit = lut_plot( + pitch.to_numpy(), + xname="energy", + yname="pitch", + title="DCM_Pitch_Calibration", + fit_order=fit_order, + save=save, + ) + + return fit + +def ge_edge_analysis(filename, save=False): + """ + Analyse a Ge K-edge scan. + + Finds the approximate absorption edge from the largest negative + derivative of the BCU BPM sum signal. + + Parameters + ---------- + filename : str + CSV filename without the .csv extension. + save : bool, optional + Save the plots to the LUT directory instead of displaying them. + + Returns + ------- + dict + Measured edge energy, expected edge energy, energy offset, + and derivative data. + """ + + # Known Ge K absorption edge + expected_edge = 11103.1 # eV + + df = get_data(filename) + + energy = df["Energy"].to_numpy() + signal1 = df["bcu_bpmsum"].to_numpy() + signal = df["transmission"].to_numpy() + + # ---------------------------------------------------------- + # Calculate derivative + # ---------------------------------------------------------- + # smooth_signal = savgol_filter( + # signal, + # window_length=11, + # polyorder=2, + # ) + # derivative = np.gradient(smooth_signal, energy) + + derivative = np.gradient(signal, energy) + + # Ge absorption causes the BPM signal to fall, so locate + # the strongest negative derivative. + edge_idx = np.argmin(derivative) + + measured_edge = energy[edge_idx] + edge_derivative = derivative[edge_idx] + + energy_offset = expected_edge - measured_edge + + # ---------------------------------------------------------- + # Print results + # ---------------------------------------------------------- + + print() + print("Ge K-edge analysis") + print("------------------") + print(f"Expected edge: {expected_edge:.2f} eV") + print(f"Measured edge: {measured_edge:.2f} eV") + print(f"Energy offset: {energy_offset:+.2f} eV") + print() + + # ---------------------------------------------------------- + # Plot BPM signal + # ---------------------------------------------------------- + + fig1, ax1 = plt.subplots() + + ax1.plot( + energy, + signal, + "o-", + markersize=4, + label="Data", + ) + + # ax1.plot( + # energy, + # smooth_signal, + # "-", + # label="Smoothed", + # ) + + ax1.axvline( + measured_edge, + linestyle="--", + label=f"Measured edge: {measured_edge:.2f} eV", + ) + if energy.min() <= expected_edge <= energy.max(): + ax1.axvline( + expected_edge, + linestyle=":", + label=f"Expected edge: {expected_edge:.1f} eV", + ) + + ax1.set_xlabel("Energy (eV)") + ax1.set_ylabel("Transmission (BCU BPM / SS BPM)") + ax1.set_title("Ge K-edge scan") + ax1.legend() + + fig1.tight_layout() + + # ---------------------------------------------------------- + # Plot derivative + # ---------------------------------------------------------- + + fig2, ax2 = plt.subplots() + + ax2.plot( + energy, + derivative, + "o-", + markersize=4, + ) + + ax2.plot( + measured_edge, + edge_derivative, + "o", + markersize=8, + label=f"Edge: {measured_edge:.2f} eV", + ) + + ax2.axvline( + measured_edge, + linestyle="--", + ) + + ax2.set_xlabel("Energy (eV)") + ax2.set_ylabel("d(Transmission) / dE") + ax2.set_title("Ge K-edge derivative") + ax2.legend() + + fig2.tight_layout() + + # ---------------------------------------------------------- + # Save or display + # ---------------------------------------------------------- + + if save: + signal_file = bl.env.lut_dir / f"{filename}_signal.png" + derivative_file = bl.env.lut_dir / f"{filename}_derivative.png" + + fig1.savefig(signal_file, bbox_inches="tight", dpi=150) + fig2.savefig(derivative_file, bbox_inches="tight", dpi=150) + + plt.close(fig1) + plt.close(fig2) + + print(f"Saved: {signal_file}") + print(f"Saved: {derivative_file}") + + else: + plt.show() + + return { + "expected_edge": expected_edge, + "measured_edge": measured_edge, + "energy_offset": energy_offset, + "derivative": derivative, + } + +def perp_analysis(filename): + """Select the DCM perpendicular translation calibration data.""" + + df = get_data(filename) + sel = df[['energy','dcm_perp', 'xi_cam_x']] + savefile = bl.env.lut_dir / "perp_actual_calc.csv" + with open(savefile, 'w', encoding='utf-8') as f: + f.write("energy,calc_perp,actual_perp\n") + + energies = [8000, 10000, 12400, 15000, 18000, 20000, 22000, 25000] + for energy in energies: + + data = sel[sel['energy']== energy] + print(data) + + target = 149.4 # value of xi_cam_x at 12.4 keV + + x = data["dcm_perp"].to_numpy() + y = data["xi_cam_x"].to_numpy() + + # Linear fit: gap = slope * energy + intercept + slope, intercept = np.polyfit(x, y, 1) + + # perp_value = intercept/-slope + perp_value = (target - intercept)/slope + calc_per = calc_perp_position(energy) + + + with open(savefile, 'a', encoding='utf-8') as f: + f.write(f"{energy},{calc_per:.4f},{perp_value:.4f}\n") + + # print( + # f"Energy = {energy} eV, " + # f"Calculated perp = {calc_per:.4f}, " + # f"Actual perp = {perp_value:.4f}" + # ) + + + # fig, ax = plt.subplots(figsize = (5,7)) + # ax.scatter(x,y) + # ax.set_title(f"{energy}") + + # # Generate fitted line only over the measured range + # x_fit = np.linspace(x.min(), x.max(), 100) + # y_fit = slope * x_fit + intercept + + # # # Plot fit first and get its automatically assigned colour + # line, = ax.plot( + # x_fit, + # y_fit, + # label=energy, + # ) + # plt.show() + + # return results + + # conv_factor = 12365.75 / (2 * 3.13481) + + # df["costheta"] = np.cos(conv_factor / df["energy"]) + + # perp = df[["costheta", "bsccam_x"]] + + # return ( + # perp.to_numpy(), + # "costheta", + # "bsccam_x", + # ) + + +def roll_analysis(): + """Select the DCM roll calibration data.""" + + df = get_data() + + roll = df[["energy", "roll"]] + + # Select the useful linear range + roll = roll[ + (roll["roll"] > 0.1) + & (roll["roll"] < 9.6) + ] + + return ( + roll.to_numpy(), + "energy", + "roll", + ) + + +def gap_analysis(save=False): + """ + Fit and plot the undulator gap calibration for all harmonics. + + Produces one plot containing all harmonic data and linear fits. + Prints fit coefficients in a format that can be copied directly + into pxii_parameters.py. + """ + + df = get_data() + gap = df[["energy", "gap", "harmonic"]] + + fig, ax = plt.subplots(figsize=(10, 7)) + + results = {} + HARMONICS = [ + "H3", + "H5", + "H7", + "H9", + "H11", + "H13", + "H15", + "H17", + "H19", + ] + for harmonic in HARMONICS: + + data = gap[gap["harmonic"] == harmonic] + + if data.empty: + print(f"{harmonic}: no data") + continue + + x = data["energy"].to_numpy() + y = data["gap"].to_numpy() + + # Linear fit: gap = slope * energy + intercept + slope, intercept = np.polyfit(x, y, 1) + + # Store results + min_gap_energy = (4.5 - intercept) / slope + + results[harmonic] = { + "slope": slope, + "intercept": intercept, + "min_gap_energy": min_gap_energy, + } + + # Generate fitted line only over the measured range + x_fit = np.linspace(x.min(), x.max(), 100) + y_fit = slope * x_fit + intercept + + # Plot fit first and get its automatically assigned colour + line, = ax.plot( + x_fit, + y_fit, + label=harmonic, + ) + + # Plot measured points using the same colour + ax.scatter( + x, + y, + color=line.get_color(), + s=35, + zorder=3, + ) + + ax.set_xlabel("energy (eV)") + ax.set_ylabel("gap (mm)") + ax.set_title("Undulator Gap Calibration") + + ax.legend( + title="Harmonic", + fontsize=9, + ) + + fig.tight_layout() + + if save: + savefile = bl.env.lut_dir / "Undulator_Gap_Calibration.png" + + fig.savefig( + savefile, + bbox_inches="tight", + dpi=150, + ) + + plt.close(fig) + + else: + plt.show() + + # ---------------------------------------------------------- + # Print fit summary + # ---------------------------------------------------------- + + print() + print("Linear fits: gap = slope * energy + intercept") + print() + print( + f"{'Harmonic':<10}" + f"{'Slope':>16}" + f"{'Intercept':>16}" + f"{'Energy at 4.5 mm':>20}" + ) + + print("-" * 62) + + for harmonic, result in results.items(): + + print( + f"{harmonic:<10}" + f"{result['slope']:>16.8e}" + f"{result['intercept']:>16.8e}" + f"{result['min_gap_energy']:>16.0f} eV" + ) + + # ---------------------------------------------------------- + # Print copy/paste-ready parameters + # ---------------------------------------------------------- + + print() + print("Copy/paste into pxii_parameters:") + print() + print("harmonics = {") + + for harmonic, result in results.items(): + + print( + f' "{harmonic}": ' + f'np.array([' + f'{result["slope"]:.8e}, ' + f'{result["intercept"]:.8e}' + f']),' + ) + + print("}") + + return results +