This commit is contained in:
x10sa committed
1 parent
d4fafdc178
commit
db244b4936
1 file changed
+583
@@ -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
|
||||
|
||||
Reference in new issue
Block a user