Skip to content

Add MagSpec to the simulation/experiment comparison #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
RemiLehe merged 6 commits into from
Jul 7, 2025
Merged

Add MagSpec to the simulation/experiment comparison #26

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
bbfeee1
extract magn spec data - first attempt
Jun 26, 2025
a2bde55
train params passed as arguments to initialize the beam, added commen…
aeriforme Jun 27, 2025
bf2b6ee
finalize inclusion of MagSpec
aeriforme Jul 1, 2025
9eadc5d
merged main
aeriforme Jul 1, 2025
4442451
Remove unused file
RemiLehe Jul 7, 2025
02842db
Update environment requirement
RemiLehe Jul 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 106 additions & 14 deletions fit_to_experiments/compare_sim_exp.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,9 @@
"import sys\n",
"import os\n",
"import torch\n",
"\n",
"from cheetah import ParameterBeam, ParticleBeam, Segment, Screen, Quadrupole\n",
"\n",
"from scipy.constants import e, m_e, c\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
Expand Down Expand Up @@ -100,6 +102,43 @@
" text = f.read()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c7f8a914-9826-4e24-9694-b4b035a8bb15",
"metadata": {},
"outputs": [],
"source": [
"# Extract the MagSpec spectrum\n",
"from scan_analysis.analyzers.Undulator.mag_spec_stitcher_analysis import MagSpecStitcherAnalyzer\n",
"\n",
"min_energy = 0.04 # GeV\n",
"max_energy = 0.18 # GeV\n",
"num_points = 800 # in the original data, 0.06 (0.2) GeV corresponds to index 19 (579)\n",
"energy = np.linspace(min_energy, max_energy, num_points) # GeV\n",
"reference_energy = 100e6 * 1e-9 # GeV\n",
"energy_norm = (energy - reference_energy ) / reference_energy # dimensionless\n",
"\n",
"exp_spectrum = np.zeros_like(energy)\n",
"# The 'U_BCaveMagSpec' diagnostic exists as a function of Shotnumber only in scans 6 and 7\n",
"num_scans = 2 # number of relevant scans\n",
"for i in tqdm.tqdm((6, 7)):\n",
" tag = ScanData.get_scan_tag(year=2024, month=12, day=5, number=i, experiment='Undulator')\n",
" analyzer = MagSpecStitcherAnalyzer(tag, device_name=f'U_BCaveMagSpec', skip_plt_show=True)\n",
" # To plot an image of the charge density as a function of shot number and energy use:\n",
" #analyzer.run_analysis()\n",
" # Extract the experimental data:\n",
" # array of energies in GeV, matrix of charge density in pC/GeV\n",
" # the matrix's axis=0 corresponds to the shot number, axis=1 to the energy\n",
" exp_energy, exp_charge_matrix = analyzer.load_charge_data() # GeV/c , pC/GeV\n",
" exp_energy, exp_charge_matrix = analyzer.interpolate_data(exp_energy, exp_charge_matrix, min_energy, max_energy, num_points)\n",
" exp_spectrum += np.average(exp_charge_matrix, axis=0) # average over Shotnumber\n",
"exp_spectrum /= num_scans # average over the scans\n",
"normalization = exp_spectrum.sum() / (energy_norm[1] - energy_norm[0])\n",
"exp_spectrum /= normalization\n",
"exp_spectrum = torch.tensor(exp_spectrum)"
]
},
{
"cell_type": "markdown",
"id": "7c36ecd0-36ec-4657-a8e9-8a66cc11885c",
Expand Down Expand Up @@ -165,12 +204,12 @@
" mu_y=trainable_parameters['beam_y_mean'], # in m\n",
" mu_py=trainable_parameters['beam_py_mean'],\n",
" mu_tau=torch.tensor(0.0), # in m\n",
" mu_p=torch.tensor(0.0), # dimensionless\n",
" mu_p=trainable_parameters['beam_p_mean'], # dimensionless\n",
" sigma_x=trainable_parameters['beam_x_std'], # in m\n",
" sigma_px=trainable_parameters['beam_px_std'],\n",
" sigma_y=trainable_parameters['beam_y_std'], # in m\n",
" sigma_py=trainable_parameters['beam_py_std'],\n",
" sigma_p=torch.tensor(2.5e-2),\n",
" sigma_p=trainable_parameters['beam_p_std'], # dimensionless\n",
" cov_xpx=torch.tensor(0.0),\n",
" cov_ypy=torch.tensor(0.0),\n",
" cov_taup=torch.tensor(0.0),\n",
Expand Down Expand Up @@ -202,7 +241,7 @@
"source": [
"def train(num_steps: int, lr=1e-5):\n",
" parameter_history = []\n",
" history = {'loss': []}\n",
" history = {'loss': [], 'loss_visa': [], 'loss_magspec': []}\n",
" for key in trainable_parameters:\n",
" if key != 'visa1_offset':\n",
" history[key] = []\n",
Expand All @@ -217,8 +256,15 @@
" beam = setup_beam_and_beamline(trainable_parameters, segment)\n",
" segment.track(beam)\n",
"\n",
" sim_spectrum = beam.get_1d_histogram( dimension=\"p\",\n",
" bins=num_points,\n",
" bin_range=(energy_norm[0], energy_norm[-1]),\n",
" method=\"kde\" )\n",
" # Compute loss\n",
" loss = torch.tensor([0.])\n",
"\n",
" # - Screens\n",
" loss_visa = torch.tensor([0.])\n",
" # Loop through screens\n",
" for device_name in experiment_images:\n",
" # Extract simulation and experiment images\n",
Expand All @@ -227,10 +273,18 @@
" sim_image = device.reading/device.reading.max() # Normalize simulation image\n",
" sim_image = torch.flip(sim_image, dims=[0])\n",
" # Compute error\n",
" loss += ((exp_image - sim_image)**2).mean()\n",
" loss_visa += ((exp_image - sim_image)**2).mean()\n",
"\n",
" # - MagSpec\n",
" # weighted arbitrarily by 1e9 -> loss_visa ~ loss_magspec (OoM)\n",
" loss_magspec = ((exp_spectrum - sim_spectrum)**2).mean() * 1e9\n",
"\n",
" loss = loss_visa + loss_magspec\n",
" loss.backward()\n",
"\n",
" # Log loss and parameters\n",
" history['loss_visa'].append(loss_visa.item())\n",
" history['loss_magspec'].append(loss_magspec.item())\n",
" history['loss'].append(loss.item())\n",
" for key in trainable_parameters:\n",
" if key != 'visa1_offset':\n",
Expand All @@ -255,10 +309,12 @@
" 'beam_y_mean': torch.tensor(0., requires_grad=True),\n",
" 'beam_px_mean': torch.tensor(0., requires_grad=True),\n",
" 'beam_py_mean': torch.tensor(0., requires_grad=True),\n",
" 'beam_p_mean': torch.tensor(0., requires_grad=True),\n",
" 'beam_x_std': torch.tensor(1e-4, requires_grad=True),\n",
" 'beam_y_std': torch.tensor(1e-4, requires_grad=True),\n",
" 'beam_px_std': torch.tensor(1e-4, requires_grad=True),\n",
" 'beam_py_std': torch.tensor(1e-4, requires_grad=True),\n",
" 'beam_p_std': torch.tensor(5e-2, requires_grad=True)\n",
"}"
]
},
Expand Down Expand Up @@ -289,15 +345,14 @@
"metadata": {},
"outputs": [],
"source": [
"plt.figure()\n",
"plt.subplot(211)\n",
"plt.semilogy(history['loss'])\n",
"plt.subplot(212)\n",
"for key in history.keys():\n",
" if key != 'loss':\n",
" plt.plot(history[key], label=key)\n",
"plt.legend()\n",
"plt.show()\n"
"fig, ax = plt.subplots(ncols=1, nrows=len(history), figsize=(4,16))\n",
"for i,key in enumerate(history.keys()):\n",
" ax[i].plot(history[key])\n",
" ax[i].set_title(key)\n",
" if 'loss' in key:\n",
" ax[i].set_yscale('log')\n",
"fig.tight_layout()\n",
"plt.show()"
]
},
{
Expand Down Expand Up @@ -462,6 +517,27 @@
"update_plot('UC_VisaEBeam4')"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d376bb6a-c6a0-4a55-9933-9675ef2308fc",
"metadata": {},
"outputs": [],
"source": [
"# Plot MagSpec spectrum\n",
"fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(5, 5))\n",
"beam = setup_beam_and_beamline(trainable_parameters, segment)\n",
"sim_spectrum = beam.get_1d_histogram( dimension=\"p\",\n",
" bins=num_points,\n",
" bin_range=(energy_norm[0], energy_norm[-1]),\n",
" method=\"kde\" )\n",
"ax.plot(energy_norm, sim_spectrum.detach().numpy(), label=\"simulation\")\n",
"ax.plot(energy_norm, exp_spectrum, label=\"experiment\")\n",
"ax.set(xlabel=r\"$\\delta=(E-E_0)/E_0$\", ylabel=r\"$dN/d\\delta$\")\n",
"ax.legend()\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand All @@ -470,14 +546,30 @@
"outputs": [],
"source": [
"# TODO\n",
"# - Add Aline screens and MagSpec\n",
"# - Add Aline screens\n",
"# - Change parametrization: use 5Dx5D covariance (time dependency don't matter) + 5D beam offset + offsets and tilts for each undulator\n",
"# - Try LBFGS optimizer\n",
"\n",
"# - Why different screen have different big blue (air/vacuum) reference?\n",
"# - Should I read SpatialCalibrationY? In the case of AlineEBeam1, it is not the same as SpatialCalibrationX\n",
"# - Is MagSpec energy measurement affected by pointing jitter?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bb0dfda1",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "5d2f5b97",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
Expand Down
2 changes: 1 addition & 1 deletion fit_to_experiments/environment.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,6 @@ dependencies:
- pip:
- ipympl
- tqdm
- cheetah-accelerator
- git+https://github.com/RemiLehe/cheetah@get_energy_spectrum
# Example of a package that uses .toml:
# - package-name