├── tests ├── __init__.py ├── bmad_lattices │ ├── test_drift.bmad │ ├── test_quad.bmad │ ├── test_quad_zero.bmad │ ├── test_bend.bmad │ ├── test_quad_tilt.bmad │ ├── test_sextupole.bmad │ ├── test_sextupole_tilt.bmad │ ├── test_bend_y.bmad │ ├── test_drift_quad.bmad │ ├── test_quad_offset.bmad │ ├── test_rf_cavity.bmad │ ├── test_sextupole_offset.bmad │ ├── test_crab_cavity.bmad │ ├── test_bend_body.bmad │ ├── test_bend_hard_exit.bmad │ ├── test_bend_hard_ent.bmad │ ├── test_bend_soft_ent.bmad │ ├── test_rf_cavity_tilt.bmad │ ├── test_crab_cavity_offset.bmad │ ├── test_rf_cavity_offset.bmad │ ├── test_crab_cavity_tilt.bmad │ ├── test_bend_linear_both_ends.bmad │ └── test_bend_soft_both_ends.bmad ├── test_torch_track.py └── test_bmadx_torch.py ├── bmadx ├── bmad_torch │ ├── __init__.py │ ├── taylor_map.py │ └── track_torch.py ├── low_level │ ├── __init__.py │ ├── sqrt_one.py │ ├── particle_rf_time.py │ ├── apply_energy_kick.py │ ├── low_energy_z_correction.py │ ├── quad_mat2_calc.py │ └── offset_particle.py ├── tracking_routines │ ├── __init__.py │ ├── track_a_drift.py │ ├── track_a_rf_cavity.py │ ├── track_a_sextupole.py │ ├── track_a_crab_cavity.py │ ├── track_a_quadrupole.py │ └── track_a_bend.py ├── constants.py ├── __init__.py ├── structures.py ├── track.py ├── distgen_utils.py ├── plot.py └── pmd_utils.py ├── docs └── examples │ ├── optimization │ └── model-calibration.png │ └── numba_tests.ipynb ├── .gitignore ├── environment-macos.yml ├── scripts └── execute_notebooks.bash ├── environment.yml ├── setup.py ├── dev ├── torch_module_profiling │ ├── Untitled.ipynb │ ├── torch_module_profiling.py │ ├── torch_module_profiling_no_grad.py │ └── a └── speed.ipynb ├── .github └── workflows │ └── ci.yml ├── README.md └── LICENSE /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /bmadx/bmad_torch/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /bmadx/low_level/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /bmadx/tracking_routines/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/examples/optimization/model-calibration.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/bmad-sim/Bmad-X/HEAD/docs/examples/optimization/model-calibration.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | bmadx.egg-info 3 | .ipynb_checkpoints 4 | misc 5 | .nfs* 6 | **/*digested* 7 | **/__pycache__ 8 | *.pyc 9 | .idea/ 10 | **/*.html 11 | -------------------------------------------------------------------------------- /tests/bmad_lattices/test_drift.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_drift 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | d1: drift, L = 1.0 11 | 12 | lat: line = (d1) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_quad.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | q1: quad, L = 0.1, K1 = 10.0 11 | 12 | lat: line = (q1) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_quad_zero.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | q1: quad, L = 0.1, K1 = 0.0 11 | 12 | lat: line = (q1) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_bend.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | b1: sbend, L = 0.1, G = 0.5, fringe_type = none 11 | 12 | lat: line = (b1) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /tests/bmad_lattices/test_quad_tilt.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad_offset_tilt 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | q1: quad, L = 0.1, K1 = 10.0, tilt=0.3 11 | 12 | lat: line = (q1) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_sextupole.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_sextupole 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | sxtpl: sextupole, L = 0.1, K2 = 10.0, NUM_STEPS=5 11 | 12 | lat: line = (sxtpl) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_sextupole_tilt.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_sextupole_tilt 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | sxtpl: sextupole, L = 0.1, K2 = 10.0, tilt=0.3 11 | 12 | lat: line = (sxtpl) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_bend_y.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | b1: sbend, L = 0.1, G = 0.5, fringe_type = none, ref_tilt = pi/2 11 | 12 | lat: line = (b1) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /bmadx/low_level/sqrt_one.py: -------------------------------------------------------------------------------- 1 | def make_sqrt_one(lib): 2 | """Makes function using library lib.""" 3 | sqrt = lib.sqrt 4 | 5 | def sqrt_one(x): 6 | """Routine to calculate Sqrt[1+x] - 1 to machine precision.""" 7 | sq = sqrt(1 + x) 8 | rad = sq + 1 9 | 10 | return x/rad 11 | 12 | return sqrt_one -------------------------------------------------------------------------------- /tests/bmad_lattices/test_drift_quad.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_drift_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | d1: drift, L=1.0 11 | q1: quad, L = 0.1, K1 = 10.0 12 | 13 | lat: line = (d1, q1, d1, q1, d1) 14 | 15 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_quad_offset.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_drift_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | q1: quad, L = 0.1, K1 = 10.0, x_offset = 1.0e-3, y_offset=-2.0e-3 11 | 12 | lat: line = (q1) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_rf_cavity.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_rf_cavity 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | cav: rfcav, L = 0.2, VOLTAGE = 1.0e4, PHI0 = 0.5, RF_FREQUENCY = 1.0e9, NUM_STEPS=1 11 | 12 | lat: line = (cav) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_sextupole_offset.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_sextupole_offset 2 | beginning[beta_a] = 10 3 | beginning[beta_b] = 10 4 | 5 | beginning[p0c] = 4.0e7 6 | parameter[particle] = electron 7 | parameter[geometry] = open 8 | 9 | sxtpl: sextupole, L = 0.1, K2 = 10.0, x_offset = 1.0e-3, y_offset=-2.0e-3 10 | 11 | lat: line = (sxtpl) 12 | 13 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_crab_cavity.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_crab_cavity 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | cav: crab_cavity, L = 0.2, VOLTAGE = 1.0e4, PHI0 = 0.5, RF_FREQUENCY = 1.0e9, NUM_STEPS=1 11 | 12 | lat: line = (cav) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_bend_body.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | b1: sbend, L = 0.1, G = 0.5, fringe_at = entrance_end, fringe_type = hard_edge_only, E1 = 0.02, FINT = 0.0, HGAP = 0.03 11 | 12 | lat: line = (b1) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /tests/bmad_lattices/test_bend_hard_exit.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | b1: sbend, L = 0.1, G = 0.5, fringe_at = exit_end, fringe_type = hard_edge_only, E2 = 0.02, FINTX = 0, HGAPX = 0.03 11 | 12 | lat: line = (b1) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /environment-macos.yml: -------------------------------------------------------------------------------- 1 | name: bmadx 2 | channels: 3 | - conda-forge 4 | - pytorch 5 | dependencies: 6 | - python >=3.10 7 | - ipython 8 | - ipykernel 9 | - pytorch 10 | - numpy 11 | - matplotlib 12 | - pytest 13 | - numba 14 | - numdifftools 15 | - kornia 16 | - bmad 17 | - pytao 18 | - distgen 19 | - pip 20 | - pip: 21 | # Install the bmad-x package as well 22 | - . 23 | -------------------------------------------------------------------------------- /tests/bmad_lattices/test_bend_hard_ent.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | b1: sbend, L = 0.1, G = 0.5, fringe_at = entrance_end, fringe_type = hard_edge_only, E1 = 0.02, FINT = 0.0, HGAP = 0.0 11 | 12 | lat: line = (b1) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /tests/bmad_lattices/test_bend_soft_ent.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | b1: sbend, L = 0.1, G = 0.5, fringe_at = entrance_end, fringe_type = soft_edge_only, E1 = 0.02, FINT = 0.5, HGAP = 0.03 11 | 12 | lat: line = (b1) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /tests/bmad_lattices/test_rf_cavity_tilt.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_crab_cavity_offset_and_tilt 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | cav: rfcav, L = 0.2, VOLTAGE = 1.0e4, PHI0 = 0.5, RF_FREQUENCY = 1.0e9, TILT = 0.3, NUM_STEPS=1 11 | 12 | lat: line = (cav) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /bmadx/constants.py: -------------------------------------------------------------------------------- 1 | import scipy.constants 2 | 3 | # math constants: 4 | PI = scipy.constants.pi 5 | 6 | # physical constants: 7 | M_ELECTRON = scipy.constants.value('electron mass energy equivalent in MeV')*1e6 # mass [eV] 8 | M_PROTON = scipy.constants.value('proton mass energy equivalent in MeV')*1e6 # mass [eV] 9 | C_LIGHT = scipy.constants.c # speed of light [m/s] 10 | E_CHARGE = scipy.constants.e # elementary charge [C] -------------------------------------------------------------------------------- /tests/bmad_lattices/test_crab_cavity_offset.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_crab_cavity_offset 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | cav: crab_cavity, L = 0.2, VOLTAGE = 1.0e4, PHI0 = 0.5, RF_FREQUENCY = 1.0e9, X_OFFSET = 1e-3, Y_OFFSET = -2e-3, NUM_STEPS=1 11 | 12 | lat: line = (cav) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_rf_cavity_offset.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_rf_cavity_offset_and_tilt 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | cav: rfcav, L = 0.2, VOLTAGE = 1.0e4, PHI0 = 0.5, RF_FREQUENCY = 1.0e9, X_OFFSET = 1e-3, Y_OFFSET = -2e-3, NUM_STEPS=1 11 | 12 | lat: line = (cav) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /scripts/execute_notebooks.bash: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | NOTEBOOKS=$(find . -type f -name "*.ipynb" -not -path '*/.*') 4 | 5 | SKIP="parallel" 6 | 7 | echo $NOTEBOOKS 8 | 9 | for file in $NOTEBOOKS 10 | do 11 | if [[ "$file" == *"$SKIP"* ]]; then 12 | echo "Skipping $file" 13 | continue 14 | fi 15 | 16 | echo "Executing $file" 17 | jupyter nbconvert --to notebook --execute $file --inplace 18 | done 19 | -------------------------------------------------------------------------------- /tests/bmad_lattices/test_crab_cavity_tilt.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_crab_cavity_tilt 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | cav: crab_cavity, L = 0.2, VOLTAGE = 1.0e4, PHI0 = 0.5, RF_FREQUENCY = 1.0e9, TILT = 0.3, NUM_STEPS=1, mat6_calc_method = bmad_standard 11 | 12 | lat: line = (cav) 13 | 14 | use, lat -------------------------------------------------------------------------------- /tests/bmad_lattices/test_bend_linear_both_ends.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | b1: sbend, L = 0.1, G = 0.5, fringe_at = both_ends, fringe_type = linear_edge, E1 = 0.2, E2 = 0.3, FINT = 0.05, HGAP = 0.03, FINTX = 0.08, HGAPX = 0.01 11 | 12 | lat: line = (b1) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: bmadx 2 | channels: 3 | - conda-forge 4 | - pytorch 5 | - nvidia 6 | dependencies: 7 | - python >=3.10 8 | - ipython 9 | - ipykernel 10 | - pytorch 11 | - pytorch-cuda >=11.7 12 | - numpy 13 | - matplotlib 14 | - pytest 15 | - numba 16 | - numdifftools 17 | - kornia 18 | - bmad 19 | - pytao 20 | - distgen 21 | - pip 22 | - pip: 23 | # Install the bmad-x package as well 24 | - . 25 | -------------------------------------------------------------------------------- /tests/bmad_lattices/test_bend_soft_both_ends.bmad: -------------------------------------------------------------------------------- 1 | parameter[lattice] = test_quad 2 | 3 | beginning[beta_a] = 10 4 | beginning[beta_b] = 10 5 | 6 | beginning[p0c] = 4.0e7 7 | parameter[particle] = electron 8 | parameter[geometry] = open 9 | 10 | b1: sbend, L = 0.1, G = 0.5, fringe_at = both_ends, fringe_type = soft_edge_only, E1 = 0.02, FINT = 0.5, HGAP = 0.03, E2 = 0.03, FINTX = 0.6, HGAPX = 0.04 11 | 12 | lat: line = (b1) 13 | 14 | use, lat 15 | -------------------------------------------------------------------------------- /bmadx/low_level/particle_rf_time.py: -------------------------------------------------------------------------------- 1 | from bmadx.constants import C_LIGHT 2 | 3 | def make_particle_rf_time(lib): 4 | """Makes function given library lib.""" 5 | sqrt = lib.sqrt 6 | 7 | def particle_rf_time(p): 8 | """Returns rf time of Particle p.""" 9 | beta = (1+p.pz) * p.p0c / sqrt(((1+p.pz)*p.p0c)**2 + p.mc2**2) 10 | time = - p.z / (beta * C_LIGHT) 11 | 12 | return time 13 | 14 | return particle_rf_time -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | with open("README.md", "r", encoding="utf-8") as fh: 4 | long_description = fh.read() 5 | 6 | setuptools.setup( 7 | name="bmadx", 8 | version="0.0.1", 9 | author="J.P. Gonzalez-Aguilera", 10 | description="Experimental Bmad code transcribed in Python with Numba and Pytorch support", 11 | long_description=long_description, 12 | long_description_content_type="text/markdown", 13 | url="https://github.com/bmad-sim/Bmad-X", 14 | project_urls={"Bug Tracker": "https://github.com/bmad-sim/Bmad-X/issues"}, 15 | license="Apache-2.0", 16 | packages=[ 17 | package 18 | for package in setuptools.find_packages(".") 19 | if package not in {"test"} 20 | ], 21 | ) 22 | -------------------------------------------------------------------------------- /bmadx/bmad_torch/taylor_map.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd.functional import jacobian 3 | 4 | from bmadx.structures import Particle 5 | 6 | def get_transport_matrix(lattice, s, p0c, mc2): 7 | def f(x): 8 | return lattice(Particle(*x, s, p0c, mc2))[:6] 9 | 10 | J = jacobian(f, torch.zeros(6)) 11 | matrix_elements = [] 12 | for ele in J[:6]: 13 | if lattice.batch_shape == torch.Size(): 14 | matrix_elements += [ele.unsqueeze(0)] 15 | else: 16 | if ele.shape == torch.Size([*lattice.batch_shape, 6]): 17 | matrix_elements += [ele.unsqueeze(-2)] 18 | elif ele.shape == torch.Size([6]): 19 | matrix_elements += [ele.repeat(*lattice.batch_shape, 1, 1)] 20 | else: 21 | raise RuntimeError("unhandled shape for jacobian") 22 | 23 | return torch.cat(matrix_elements, dim=-2) 24 | -------------------------------------------------------------------------------- /dev/torch_module_profiling/Untitled.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "id": "e24c4f5b-7ff7-4b96-81cd-8a0e1f5a5ba9", 7 | "metadata": {}, 8 | "outputs": [], 9 | "source": [ 10 | "import json\n", 11 | "import os\n", 12 | "import sys\n", 13 | "import numpy as np\n", 14 | "import matplotlib.pyplot as plt" 15 | ] 16 | } 17 | ], 18 | "metadata": { 19 | "kernelspec": { 20 | "display_name": "PhaseSpaceReconstruction", 21 | "language": "python", 22 | "name": "phase_space_reconstruction" 23 | }, 24 | "language_info": { 25 | "codemirror_mode": { 26 | "name": "ipython", 27 | "version": 3 28 | }, 29 | "file_extension": ".py", 30 | "mimetype": "text/x-python", 31 | "name": "python", 32 | "nbconvert_exporter": "python", 33 | "pygments_lexer": "ipython3", 34 | "version": "3.10.8" 35 | } 36 | }, 37 | "nbformat": 4, 38 | "nbformat_minor": 5 39 | } 40 | -------------------------------------------------------------------------------- /bmadx/low_level/apply_energy_kick.py: -------------------------------------------------------------------------------- 1 | from bmadx.structures import Particle 2 | from bmadx.low_level.sqrt_one import make_sqrt_one 3 | 4 | def make_apply_energy_kick(lib): 5 | """Makes function given library lib.""" 6 | sqrt = lib.sqrt 7 | sqrt_one = make_sqrt_one(lib) 8 | 9 | def apply_energy_kick(dE, p_in): 10 | """Changes the energy of a particle by dE.""" 11 | z, pz, p0c, mc2 = p_in.z, p_in.pz, p_in.p0c, p_in.mc2 12 | 13 | pc = (1 + pz) * p0c 14 | beta_old = (1+pz) * p0c / sqrt(((1+pz)*p0c)**2 + mc2**2) 15 | E_old = pc / beta_old 16 | 17 | E_new = E_old + dE 18 | 19 | pz = pz + (1 + pz) * sqrt_one((2*E_old*dE + dE**2)/pc**2) 20 | pc_new = p0c * (1 + pz) 21 | beta_new = pc_new / E_new 22 | z = z * beta_new / beta_old 23 | 24 | return Particle(p_in.x, p_in.px, p_in.y, p_in.py, z, pz, 25 | p_in.s, p0c, mc2) 26 | 27 | return apply_energy_kick -------------------------------------------------------------------------------- /bmadx/low_level/low_energy_z_correction.py: -------------------------------------------------------------------------------- 1 | def make_low_energy_z_correction(lib): 2 | """Makes function given library lib""" 3 | sqrt = lib.sqrt 4 | 5 | def low_energy_z_correction(pz, p0c, mass, ds): 6 | """Corrects the change in z-coordinate due to speed < c_light. 7 | Input: 8 | p0c -- reference particle momentum in eV 9 | mass -- particle mass in eV 10 | Output: 11 | dz -- dz=(ds-d_particle) + ds*(beta - beta_ref)/beta_ref 12 | """ 13 | beta = (1+pz) * p0c / sqrt(((1+pz)*p0c)**2 + mass**2) 14 | beta0 = p0c / sqrt( p0c**2 + mass**2) 15 | e_tot = sqrt(p0c**2+mass**2) 16 | 17 | evaluation = mass * (beta0*pz)**2 18 | dz = (ds * pz * (1 - 3*(pz*beta0**2)/2+pz**2*beta0**2 19 | * (2*beta0**2-(mass/e_tot)**2/2) ) 20 | * (mass/e_tot)**2 21 | * (evaluation<3e-7*e_tot) 22 | + (ds*(beta-beta0)/beta0) 23 | * (evaluation>=3e-7*e_tot) ) 24 | 25 | return dz 26 | 27 | return low_energy_z_correction -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: Tests 2 | 3 | on: [push, pull_request] 4 | 5 | jobs: 6 | build: 7 | name: Test ${{ matrix.os }} (${{ matrix.python-version }}) 8 | runs-on: ${{ matrix.os }} 9 | defaults: 10 | run: 11 | shell: bash 12 | strategy: 13 | fail-fast: false 14 | matrix: 15 | os: [ubuntu-latest] 16 | python-version: ["3.10", "3.11", "3.12"] 17 | steps: 18 | - uses: actions/checkout@v4 19 | - uses: conda-incubator/setup-miniconda@v3 20 | with: 21 | python-version: ${{ matrix.python-version }} 22 | mamba-version: "*" 23 | channels: conda-forge 24 | activate-environment: bmadx 25 | environment-file: environment.yml 26 | 27 | #- name: flake8 28 | # shell: bash -l {0} 29 | # run: | 30 | # flake8 . 31 | 32 | - name: Show conda environment packages 33 | shell: bash -l {0} 34 | run: | 35 | conda list 36 | 37 | - name: Ensure importability 38 | shell: bash -l {0} 39 | run: | 40 | cd / 41 | python -c "import bmadx" 42 | 43 | - name: Run Tests 44 | shell: bash -l {0} 45 | run: | 46 | pytest tests 47 | -------------------------------------------------------------------------------- /tests/test_torch_track.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from bmadx.bmad_torch.track_torch import ( 4 | Beam, 5 | TorchCrabCavity, 6 | TorchDrift, 7 | TorchLattice, 8 | TorchQuadrupole, 9 | ) 10 | from bmadx.bmad_torch.taylor_map import get_transport_matrix 11 | 12 | 13 | def test_lattice(): 14 | q1 = TorchQuadrupole(torch.tensor(0.1), torch.tensor(1.0)) 15 | d1 = TorchDrift(torch.tensor(1.0)) 16 | c1 = TorchCrabCavity(torch.tensor(1.0), torch.tensor(100.0), torch.tensor(1.0)) 17 | 18 | lattice = TorchLattice([q1, d1, c1]) 19 | 20 | test_beam = Beam( 21 | torch.ones(6, 6), 22 | p0c=torch.tensor(10.0e6), 23 | ) 24 | print(lattice(test_beam)) 25 | 26 | 27 | def test_jacobian(): 28 | ks = [ 29 | torch.tensor(1.0), 30 | torch.linspace(0, 1, 10), 31 | torch.linspace(0, 1, 10).unsqueeze(-1), 32 | ] 33 | for k in ks: 34 | q1 = TorchQuadrupole(torch.tensor(0.1), k) 35 | assert q1.batch_shape == k.shape 36 | d1 = TorchDrift(torch.tensor(0.5)) 37 | 38 | lattice = TorchLattice([q1, d1]) 39 | assert lattice.batch_shape == k.shape 40 | 41 | J = get_transport_matrix( 42 | lattice, torch.tensor(0.0), torch.tensor(10.0e6), torch.tensor(0.511e6) 43 | ) 44 | assert J.shape == torch.Size([*k.shape, 6, 6]) 45 | 46 | 47 | 48 | test_lattice() 49 | test_jacobian() 50 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Bmad-X 2 | 3 | Experimental Bmad Code transcribed in Python with Numba and Pytorch support. 4 | 5 | ## Installation 6 | 7 | ```bash 8 | git clone https://github.com/bmad-sim/Bmad-X.git 9 | # pytorch-cuda on Windows or Linux. Use environment-macos.yml for pytorch-cpu on MacOS. 10 | conda env create -f environment.yml 11 | conda activate bmadx 12 | ``` 13 | 14 | For a development installation of Bmad-X, run the following after creating the environment: 15 | 16 | ```bash 17 | python -m pip install -e . 18 | ``` 19 | 20 | ## Cite 21 | 22 | ```bibtex 23 | @inproceedings{gonzalez-aguilera:ipac2023-wepa065, 24 | title = {Towards fully differentiable accelerator modeling}, 25 | author = {Gonzalez-Aguilera, J. and Kim, Y.-K. and Roussel, R. and Edelen, A. and Mayes, C.}, 26 | year = 2023, 27 | month = {05}, 28 | booktitle = {Proc. IPAC'23}, 29 | publisher = {JACoW Publishing, Geneva, Switzerland}, 30 | series = {IPAC'23 - 14th International Particle Accelerator Conference}, 31 | number = 14, 32 | pages = {2797--2800}, 33 | doi = {10.18429/JACoW-IPAC2023-WEPA065}, 34 | isbn = {978-3-95450-231-8}, 35 | issn = {2673-5490}, 36 | url = {https://indico.jacow.org/event/41/contributions/2122}, 37 | % booktitle = {Proc. 14th International Particle Accelerator Conference}, 38 | paper = {WEPA065}, 39 | venue = {Venice, Italy}, 40 | language = {English} 41 | } 42 | ``` 43 | -------------------------------------------------------------------------------- /bmadx/tracking_routines/track_a_drift.py: -------------------------------------------------------------------------------- 1 | from bmadx.structures import Particle 2 | from bmadx.low_level.sqrt_one import make_sqrt_one 3 | 4 | def make_track_a_drift(lib): 5 | """Makes track_a_drift given the library lib.""" 6 | sqrt = lib.sqrt 7 | sqrt_one = make_sqrt_one(lib) 8 | 9 | def track_a_drift(p_in, drift): 10 | """Tracks the incoming Particle p_in though drift element 11 | and returns the outgoing particle. 12 | See Bmad manual section 24.9 13 | """ 14 | L = drift.L 15 | 16 | s = p_in.s 17 | p0c = p_in.p0c 18 | mc2 = p_in.mc2 19 | 20 | x, px, y, py, z, pz = p_in.x, p_in.px, p_in.y, p_in.py, p_in.z, p_in.pz 21 | 22 | P = 1 + pz # Particle's total momentum over p0 23 | Px = px / P # Particle's 'x' momentum over p0 24 | Py = py / P # Particle's 'y' momentum over p0 25 | Pxy2 = Px**2 + Py**2 # Particle's transverse mometum^2 over p0^2 26 | Pl = sqrt(1-Pxy2) # Particle's longitudinal momentum over p0 27 | 28 | x = x + L * Px / Pl 29 | y = y + L * Py / Pl 30 | 31 | # z = z + L * ( beta/beta_ref - 1.0/Pl ) but numerically accurate: 32 | dz = L * (sqrt_one((mc2**2 * (2*pz+pz**2))/((p0c*P)**2 + mc2**2)) 33 | + sqrt_one(-Pxy2)/Pl) 34 | z = z + dz 35 | s = s + L 36 | 37 | return Particle(x, px, y, py, z, pz, s, p0c, mc2) 38 | 39 | return track_a_drift -------------------------------------------------------------------------------- /bmadx/low_level/quad_mat2_calc.py: -------------------------------------------------------------------------------- 1 | def make_quad_mat2_calc(lib): 2 | """Makes function given library lib""" 3 | sqrt = lib.sqrt 4 | sin = lib.sin 5 | cos = lib.cos 6 | sinh = lib.sinh 7 | cosh = lib.cosh 8 | absolute = lib.abs 9 | 10 | def quad_mat2_calc(k1, length, rel_p): 11 | """Returns 2x2 transfer matrix elements aij and the 12 | coefficients to calculate the change in z position. 13 | Input: 14 | k1_ref -- Quad strength: k1 > 0 ==> defocus 15 | length -- Quad length 16 | rel_p -- Relative momentum P/P0 17 | Output: 18 | a11, a12, a21, a22 -- transfer matrix elements 19 | c1, c2, c3 -- second order derivatives of z such that 20 | z = c1 * x_0^2 + c2 * x_0 * px_0 + c3* px_0^2 21 | **NOTE**: accumulated error due to machine epsilon. REVISIT 22 | """ 23 | eps = 2.220446049250313e-16 # machine epsilon to double precission 24 | 25 | sqrt_k = sqrt(absolute(k1)+eps) 26 | sk_l = sqrt_k * length 27 | 28 | cx = cos(sk_l) * (k1<=0) + cosh(sk_l) * (k1>0) 29 | sx = (sin(sk_l)/(sqrt_k))*(k1<=0) + (sinh(sk_l)/(sqrt_k))*(k1>0) 30 | 31 | a11 = cx 32 | a12 = sx / rel_p 33 | a21 = k1 * sx * rel_p 34 | a22 = cx 35 | 36 | c1 = k1 * (-cx * sx + length) / 4 37 | c2 = -k1 * sx**2 / (2 * rel_p) 38 | c3 = -(cx * sx + length) / (4 * rel_p**2) 39 | 40 | return [[a11, a12], [a21, a22]], [c1, c2, c3] 41 | 42 | return quad_mat2_calc -------------------------------------------------------------------------------- /bmadx/__init__.py: -------------------------------------------------------------------------------- 1 | # Constants 2 | from .constants import * 3 | 4 | # Structures (particle and elements named tuples) 5 | from .structures import Particle, Drift, Quadrupole, CrabCavity, RFCavity, SBend, \ 6 | Sextupole 7 | 8 | # Tracking routines 9 | from .tracking_routines.track_a_drift import make_track_a_drift 10 | from .tracking_routines.track_a_quadrupole import make_track_a_quadrupole 11 | from .tracking_routines.track_a_crab_cavity import make_track_a_crab_cavity 12 | from .tracking_routines.track_a_rf_cavity import make_track_a_rf_cavity 13 | from .tracking_routines.track_a_bend import make_track_a_bend 14 | from .tracking_routines.track_a_sextupole import make_track_a_sextupole 15 | 16 | # Tracking library dictionary (as of now, supports NumPy and PyTorch): 17 | import numpy as np 18 | import torch 19 | 20 | def make_tracking_dict(lib): 21 | dic = { 22 | "Drift" : make_track_a_drift(lib), 23 | "Quadrupole" : make_track_a_quadrupole(lib), 24 | "CrabCavity" : make_track_a_crab_cavity(lib), 25 | "RFCavity" : make_track_a_rf_cavity(lib), 26 | "SBend": make_track_a_bend(lib), 27 | "Sextupole": make_track_a_sextupole(lib) 28 | } 29 | return dic 30 | 31 | LIB_DICT = { 32 | np: { 33 | 'tracking_routine': make_tracking_dict(np), 34 | 'number_type': np.ndarray, 35 | 'construct_type': np.array 36 | }, 37 | torch: { 38 | 'tracking_routine': make_tracking_dict(torch), 39 | 'number_type': torch.Tensor, 40 | 'construct_type': torch.tensor 41 | } 42 | } 43 | 44 | # Tracking functions 45 | from .track import track_element, track_lattice, track_lattice_save_particles, track_lattice_save_stats -------------------------------------------------------------------------------- /bmadx/low_level/offset_particle.py: -------------------------------------------------------------------------------- 1 | from bmadx.structures import Particle 2 | 3 | def make_offset_particle(lib, f): 4 | """Makes function f given library lib. 5 | f -- set or unset 6 | """ 7 | sin = lib.sin 8 | cos = lib.cos 9 | 10 | def offset_particle_set(x_offset, y_offset, tilt, p_lab): 11 | """Transform from lab to element coords. 12 | See Bmad manual (2022-11-06) sections 5.6.1, 15.3.1 and 24.2 13 | **NOTE**: transverse only as of now. 14 | """ 15 | s = sin(tilt) 16 | c = cos(tilt) 17 | x_ele_int = p_lab.x - x_offset 18 | y_ele_int = p_lab.y - y_offset 19 | x_ele = x_ele_int*c + y_ele_int*s 20 | y_ele = -x_ele_int*s + y_ele_int*c 21 | px_ele = p_lab.px*c + p_lab.py*s 22 | py_ele = -p_lab.px*s + p_lab.py*c 23 | 24 | return Particle(x_ele, px_ele, y_ele, py_ele, p_lab.z, p_lab.pz, 25 | p_lab.s, p_lab.p0c, p_lab.mc2) 26 | 27 | def offset_particle_unset(x_offset, y_offset, tilt, p_ele): 28 | """Transforms from element body to lab coords. 29 | See Bmad manual (2022-11-06) sections 5.6.1, 15.3.1 and 24.2 30 | **NOTE**: transverse only as of now. 31 | """ 32 | s = sin(tilt) 33 | c = cos(tilt) 34 | x_lab_int = p_ele.x*c - p_ele.y*s 35 | y_lab_int = p_ele.x*s + p_ele.y*c 36 | x_lab = x_lab_int + x_offset 37 | y_lab = y_lab_int + y_offset 38 | px_lab = p_ele.px*c - p_ele.py*s 39 | py_lab = p_ele.px*s + p_ele.py*c 40 | 41 | return Particle(x_lab, px_lab, y_lab, py_lab, p_ele.z, p_ele.pz, 42 | p_ele.s, p_ele.p0c, p_ele.mc2) 43 | 44 | f_dict = { 45 | "set" : offset_particle_set, 46 | "unset" : offset_particle_unset 47 | } 48 | 49 | return f_dict[f] 50 | -------------------------------------------------------------------------------- /bmadx/structures.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | from bmadx.constants import M_ELECTRON 3 | 4 | Particle = namedtuple('Particle', 'x px y py z pz s p0c mc2') 5 | 6 | Drift = namedtuple('Drift', 'L') 7 | 8 | Quadrupole = namedtuple( 9 | 'Quadrupole', 10 | [ 11 | 'L', 12 | 'K1', 13 | 'NUM_STEPS', 14 | 'X_OFFSET', 15 | 'Y_OFFSET', 16 | 'TILT' 17 | ], 18 | defaults=(None, None, 1, 0.0, 0.0, 0.0) 19 | ) 20 | 21 | Sextupole = namedtuple( 22 | 'Sextupole', 23 | [ 24 | 'L', 25 | 'K2', 26 | 'NUM_STEPS', 27 | 'X_OFFSET', 28 | 'Y_OFFSET', 29 | 'TILT' 30 | ], 31 | defaults=(None, None, 1, 0.0, 0.0, 0.0) 32 | ) 33 | 34 | CrabCavity = namedtuple( 35 | 'CrabCavity', 36 | [ 37 | 'L', 38 | 'VOLTAGE', 39 | 'PHI0', 40 | 'RF_FREQUENCY', 41 | 'NUM_STEPS', 42 | 'X_OFFSET', 43 | 'Y_OFFSET', 44 | 'TILT' 45 | ], 46 | defaults=(None, None, None, None, 1, 0.0, 0.0, 0.0) 47 | ) 48 | 49 | RFCavity = namedtuple( 50 | 'RFCavity', 51 | [ 52 | 'L', 53 | 'VOLTAGE', 54 | 'PHI0', 55 | 'RF_FREQUENCY', 56 | 'NUM_STEPS', 57 | 'X_OFFSET', 58 | 'Y_OFFSET', 59 | 'TILT' 60 | ], 61 | defaults=(None, None, None, None, 1, 0.0, 0.0, 0.0) 62 | ) 63 | 64 | SBend = namedtuple( 65 | 'SBend', 66 | [ 67 | 'L', 68 | 'P0C', 69 | 'G', 70 | 'DG', 71 | 'E1', 72 | 'E2', 73 | 'FINT', 74 | 'HGAP', 75 | 'FINTX', 76 | 'HGAPX', 77 | 'FRINGE_AT', 78 | 'FRINGE_TYPE', 79 | 'TILT' 80 | ], 81 | defaults=( 82 | None, 83 | None, 84 | None, 85 | 0.0, 86 | 0.0, 87 | 0.0, 88 | 0.0, 89 | 0.0, 90 | 0.0, 91 | 0.0, 92 | "both_ends", 93 | "linear_edge", 94 | 0.0 95 | ) 96 | ) -------------------------------------------------------------------------------- /bmadx/tracking_routines/track_a_rf_cavity.py: -------------------------------------------------------------------------------- 1 | from bmadx.structures import Drift 2 | 3 | from bmadx.low_level.offset_particle import make_offset_particle 4 | from bmadx.low_level.particle_rf_time import make_particle_rf_time 5 | from bmadx.low_level.apply_energy_kick import make_apply_energy_kick 6 | 7 | from bmadx.tracking_routines.track_a_drift import make_track_a_drift 8 | 9 | from bmadx.constants import C_LIGHT, PI 10 | 11 | def make_track_a_rf_cavity(lib): 12 | """Makes track_a_crab_cavity given the library lib.""" 13 | sqrt = lib.sqrt 14 | sin = lib.sin 15 | track_this_drift = make_track_a_drift(lib) 16 | offset_particle_set = make_offset_particle(lib, 'set') 17 | offset_particle_unset = make_offset_particle(lib, 'unset') 18 | particle_rf_time = make_particle_rf_time(lib) 19 | apply_energy_kick = make_apply_energy_kick(lib) 20 | 21 | def track_a_rf_cavity(p_in, cav): 22 | """Tracks an incomming Particle p_in through rf cavity and 23 | returns the ourgoing particle. 24 | See Bmad manual section 4.9 25 | """ 26 | p0c = p_in.p0c 27 | mc2 = p_in.mc2 28 | 29 | l = cav.L 30 | 31 | x_off = cav.X_OFFSET 32 | y_off = cav.Y_OFFSET 33 | tilt = cav.TILT 34 | 35 | par = offset_particle_set(x_off, y_off, tilt, p_in) 36 | 37 | voltage = cav.VOLTAGE 38 | phase = 2 * PI * (cav.PHI0 - (particle_rf_time(par)*cav.RF_FREQUENCY)) 39 | 40 | dE = voltage * sin(phase) / 2 41 | 42 | par = apply_energy_kick(dE, par) 43 | 44 | z_old = par.z 45 | 46 | par = track_this_drift(par, Drift(l)) 47 | z, pz = par.z, par.pz 48 | beta_new = (1+pz) * p0c / sqrt(((1+pz)*p0c)**2 + mc2**2) 49 | phase = phase + 2 * PI * cav.RF_FREQUENCY * (z-z_old)/(C_LIGHT*beta_new) 50 | 51 | dE = voltage * sin(phase) / 2 52 | 53 | par = apply_energy_kick(dE, par) 54 | 55 | par = offset_particle_unset(x_off, y_off, tilt, par) 56 | 57 | return par 58 | 59 | return track_a_rf_cavity -------------------------------------------------------------------------------- /bmadx/tracking_routines/track_a_sextupole.py: -------------------------------------------------------------------------------- 1 | from bmadx.structures import Particle 2 | 3 | from bmadx.low_level.offset_particle import make_offset_particle 4 | from bmadx.low_level.low_energy_z_correction import make_low_energy_z_correction 5 | 6 | def make_track_a_sextupole(lib): 7 | """Makes track_a_quadrupole given the library lib.""" 8 | offset_particle_set = make_offset_particle(lib, 'set') 9 | offset_particle_unset = make_offset_particle(lib, 'unset') 10 | low_energy_z_correction = make_low_energy_z_correction(lib) 11 | 12 | def track_a_sextupole(p_in, sextupole): 13 | """Tracks the incoming Particle p_in though pure sextupole element and 14 | returns the outgoing particle. 15 | See Bmad manual section 24.15 16 | """ 17 | l = sextupole.L 18 | k2 = sextupole.K2 19 | 20 | n_step = sextupole.NUM_STEPS # number of divisions 21 | step_len = l / n_step # length of division 22 | 23 | x_off = sextupole.X_OFFSET 24 | y_off = sextupole.Y_OFFSET 25 | tilt = sextupole.TILT 26 | 27 | b2 = k2 * l 28 | 29 | s = p_in.s 30 | p0c = p_in.p0c 31 | mc2 = p_in.mc2 32 | 33 | # --- TRACKING --- : 34 | 35 | par = offset_particle_set(x_off, y_off, tilt, p_in) 36 | x, px, y, py, z, pz = par.x, par.px, par.y, par.py, par.z, par.pz 37 | 38 | for i in range(n_step): 39 | rel_p = 1 + pz # Particle's relative momentum (P/P0) 40 | k2 = b2/(l*rel_p) 41 | 42 | x_next = x + step_len * px 43 | y_next = y + step_len * py 44 | 45 | px_next = px + 0.5 * k2 * step_len * (y**2 - x**2) 46 | py_next = py + k2 * step_len * x * y 47 | 48 | x, px, y, py = x_next, px_next, y_next, py_next 49 | 50 | z = z + low_energy_z_correction(pz, p0c, mc2, step_len) 51 | 52 | s = s + l 53 | 54 | par = offset_particle_unset(x_off, y_off, tilt, 55 | Particle(x, px, y, py, z, pz, s, p0c, mc2)) 56 | 57 | return par 58 | 59 | return track_a_sextupole -------------------------------------------------------------------------------- /bmadx/tracking_routines/track_a_crab_cavity.py: -------------------------------------------------------------------------------- 1 | from bmadx.structures import Particle, Drift 2 | 3 | from bmadx.low_level.offset_particle import make_offset_particle 4 | from bmadx.low_level.particle_rf_time import make_particle_rf_time 5 | 6 | from bmadx.tracking_routines.track_a_drift import make_track_a_drift 7 | 8 | from bmadx.constants import C_LIGHT, PI 9 | 10 | def make_track_a_crab_cavity(lib): 11 | """Makes track_a_crab_cavity given the library lib.""" 12 | sqrt = lib.sqrt 13 | sin = lib.sin 14 | cos = lib.cos 15 | track_this_drift = make_track_a_drift(lib) 16 | offset_particle_set = make_offset_particle(lib, 'set') 17 | offset_particle_unset = make_offset_particle(lib, 'unset') 18 | particle_rf_time = make_particle_rf_time(lib) 19 | 20 | def track_a_crab_cavity(p_in, cav): 21 | """Tracks an incomming Particle p_in through crab cavity and 22 | returns the ourgoing particle. 23 | See Bmad manual section 4.9 24 | """ 25 | s = p_in.s 26 | p0c = p_in.p0c 27 | mc2 = p_in.mc2 28 | 29 | l = cav.L 30 | 31 | x_off = cav.X_OFFSET 32 | y_off = cav.Y_OFFSET 33 | tilt = cav.TILT 34 | 35 | par = offset_particle_set(x_off, y_off, tilt, p_in) 36 | 37 | par = track_this_drift(par, Drift(l/2)) 38 | x, px, y, py, z, pz = par.x, par.px, par.y, par.py, par.z, par.pz 39 | 40 | voltage = cav.VOLTAGE / p0c 41 | k_rf = 2 * PI * cav.RF_FREQUENCY / C_LIGHT 42 | phase = 2 * PI * (cav.PHI0 - (particle_rf_time(par)*cav.RF_FREQUENCY)) 43 | 44 | px = px + voltage * sin(phase) 45 | 46 | beta = (1+pz) * p0c / sqrt(((1+pz)*p0c)**2 + mc2**2) 47 | beta_old = beta 48 | E_old = (1+pz) * p0c / beta_old 49 | E_new = E_old + voltage * cos(phase) * k_rf * x * p0c 50 | pc = sqrt(E_new**2-mc2**2) 51 | beta = pc / E_new 52 | 53 | pz = (pc - p0c)/p0c 54 | z = z * beta / beta_old 55 | 56 | par = track_this_drift(Particle(x, px, y, py, z, pz, s, p0c, mc2), 57 | Drift(l/2)) 58 | 59 | par = offset_particle_unset(x_off, y_off, tilt, par) 60 | 61 | return par 62 | 63 | return track_a_crab_cavity -------------------------------------------------------------------------------- /dev/torch_module_profiling/torch_module_profiling.py: -------------------------------------------------------------------------------- 1 | #from scalene import scalene_profiler 2 | 3 | import sys 4 | #print(sys.argv) 5 | n_particles = int(sys.argv[1]) 6 | n_quadrupoles = int(sys.argv[2]) 7 | n_slices = int(sys.argv[3]) 8 | 9 | import torch 10 | 11 | from bmadx import Particle, Drift, Quadrupole, track_lattice, M_ELECTRON 12 | 13 | def create_gaussian_beam(n_particles, mean, cov, s, p0c, mc2): 14 | torch.manual_seed(0) 15 | dist = torch.distributions.multivariate_normal.MultivariateNormal(mean, cov) 16 | coords = dist.sample(torch.Size([n_particles])) # particles' coordinates 17 | 18 | return Particle(*coords.T, s, p0c, mc2) 19 | 20 | class BeamlineModel(torch.nn.Module): 21 | def __init__(self, k_set, l_d, l_q, beam_in, sigma_t): 22 | super().__init__() 23 | self.register_parameter('k_set', torch.nn.Parameter(k_set)) 24 | self.l_d = l_d 25 | self.l_q = l_q 26 | self.beam_in = beam_in 27 | self.sigma_t = sigma_t 28 | def forward(self): 29 | # create lattice 30 | lattice = [] 31 | half_drift = Drift( L = self.l_d/2) 32 | for k in self.k_set: 33 | lattice.append( half_drift ) 34 | lattice.append( Quadrupole(L = self.l_q, 35 | K1 = k, 36 | NUM_STEPS = n_slices) ) 37 | lattice.append( half_drift ) 38 | 39 | beam_out = track_lattice(self.beam_in, lattice) 40 | dx = beam_out.x.std() - self.sigma_t 41 | dy = beam_out.y.std() - self.sigma_t 42 | delta = torch.sqrt( dx**2 + dy**2 ) 43 | 44 | return delta 45 | 46 | # model constants 47 | l_d = 0.9 # drift length 48 | l_q = 0.1 # quad length 49 | sigma_t = 5e-3 # target beam size 50 | 51 | # incoming beam 52 | beam_in = create_gaussian_beam(n_particles, 53 | mean = torch.zeros(6)*1e-3, 54 | cov = torch.eye(6)*1e-6, 55 | s = torch.tensor(0.0), 56 | p0c = torch.tensor(4e7), 57 | mc2 = torch.tensor(M_ELECTRON)) 58 | 59 | # model evaluation and backward pass 60 | k_set = torch.zeros(n_quadrupoles, requires_grad=True) 61 | #scalene_profiler.start() 62 | model = BeamlineModel(k_set, l_d, l_q, beam_in, sigma_t) 63 | loss = model() 64 | loss.backward() 65 | #scalene_profiler.stop() -------------------------------------------------------------------------------- /dev/torch_module_profiling/torch_module_profiling_no_grad.py: -------------------------------------------------------------------------------- 1 | #from scalene import scalene_profiler 2 | 3 | import sys 4 | #print(sys.argv) 5 | n_particles = int(sys.argv[1]) 6 | n_quadrupoles = int(sys.argv[2]) 7 | n_slices = int(sys.argv[3]) 8 | 9 | import torch 10 | 11 | from bmadx import Particle, Drift, Quadrupole, track_lattice, M_ELECTRON 12 | 13 | def create_gaussian_beam(n_particles, mean, cov, s, p0c, mc2): 14 | torch.manual_seed(0) 15 | dist = torch.distributions.multivariate_normal.MultivariateNormal(mean, cov) 16 | coords = dist.sample(torch.Size([n_particles])) # particles' coordinates 17 | 18 | return Particle(*coords.T, s, p0c, mc2) 19 | 20 | class BeamlineModel(torch.nn.Module): 21 | def __init__(self, k_set, l_d, l_q, beam_in, sigma_t): 22 | super().__init__() 23 | self.register_parameter('k_set', torch.nn.Parameter(k_set)) 24 | self.l_d = l_d 25 | self.l_q = l_q 26 | self.beam_in = beam_in 27 | self.sigma_t = sigma_t 28 | def forward(self): 29 | # create lattice 30 | lattice = [] 31 | half_drift = Drift( L = self.l_d/2) 32 | for k in self.k_set: 33 | lattice.append( half_drift ) 34 | lattice.append( Quadrupole(L = self.l_q, 35 | K1 = k, 36 | NUM_STEPS = n_slices) ) 37 | lattice.append( half_drift ) 38 | with torch.no_grad(): 39 | beam_out = track_lattice(self.beam_in, lattice) 40 | dx = beam_out.x.std() - self.sigma_t 41 | dy = beam_out.y.std() - self.sigma_t 42 | delta = torch.sqrt( dx**2 + dy**2 ) 43 | 44 | return delta 45 | 46 | # model constants 47 | l_d = 0.9 # drift length 48 | l_q = 0.1 # quad length 49 | sigma_t = 5e-3 # target beam size 50 | 51 | # incoming beam 52 | beam_in = create_gaussian_beam(n_particles, 53 | mean = torch.zeros(6)*1e-3, 54 | cov = torch.eye(6)*1e-6, 55 | s = torch.tensor(0.0), 56 | p0c = torch.tensor(4e7), 57 | mc2 = torch.tensor(M_ELECTRON)) 58 | 59 | # model evaluation and backward pass 60 | k_set = torch.zeros(n_quadrupoles, requires_grad=True) 61 | #scalene_profiler.start() 62 | model = BeamlineModel(k_set, l_d, l_q, beam_in, sigma_t) 63 | loss = model() 64 | #loss.backward() 65 | #scalene_profiler.stop() -------------------------------------------------------------------------------- /bmadx/tracking_routines/track_a_quadrupole.py: -------------------------------------------------------------------------------- 1 | from bmadx.structures import Particle 2 | 3 | from bmadx.low_level.offset_particle import make_offset_particle 4 | from bmadx.low_level.low_energy_z_correction import make_low_energy_z_correction 5 | from bmadx.low_level.quad_mat2_calc import make_quad_mat2_calc 6 | 7 | def make_track_a_quadrupole(lib): 8 | """Makes track_a_quadrupole given the library lib.""" 9 | quad_mat2_calc = make_quad_mat2_calc(lib) 10 | offset_particle_set = make_offset_particle(lib, 'set') 11 | offset_particle_unset = make_offset_particle(lib, 'unset') 12 | low_energy_z_correction = make_low_energy_z_correction(lib) 13 | 14 | def track_a_quadrupole(p_in, quad): 15 | """Tracks the incoming Particle p_in though quad element and 16 | returns the outgoing particle. 17 | See Bmad manual section 24.15 18 | """ 19 | l = quad.L 20 | k1 = quad.K1 21 | n_step = quad.NUM_STEPS # number of divisions 22 | step_len = l / n_step # length of division 23 | 24 | x_off = quad.X_OFFSET 25 | y_off = quad.Y_OFFSET 26 | tilt = quad.TILT 27 | 28 | b1 = k1 * l 29 | 30 | s = p_in.s 31 | p0c = p_in.p0c 32 | mc2 = p_in.mc2 33 | 34 | # --- TRACKING --- : 35 | 36 | par = offset_particle_set(x_off, y_off, tilt, p_in) 37 | x, px, y, py, z, pz = par.x, par.px, par.y, par.py, par.z, par.pz 38 | 39 | for i in range(n_step): 40 | rel_p = 1 + pz # Particle's relative momentum (P/P0) 41 | k1 = b1/(l*rel_p) 42 | 43 | tx, dzx = quad_mat2_calc(-k1, step_len, rel_p) 44 | ty, dzy = quad_mat2_calc( k1, step_len, rel_p) 45 | 46 | z = ( z 47 | + dzx[0] * x**2 + dzx[1] * x * px + dzx[2] * px**2 48 | + dzy[0] * y**2 + dzy[1] * y * py + dzy[2] * py**2 ) 49 | 50 | x_next = tx[0][0] * x + tx[0][1] * px 51 | px_next = tx[1][0] * x + tx[1][1] * px 52 | y_next = ty[0][0] * y + ty[0][1] * py 53 | py_next = ty[1][0] * y + ty[1][1] * py 54 | 55 | x, px, y, py = x_next, px_next, y_next, py_next 56 | 57 | z = z + low_energy_z_correction(pz, p0c, mc2, step_len) 58 | 59 | s = s + l 60 | 61 | par = offset_particle_unset(x_off, y_off, tilt, 62 | Particle(x, px, y, py, z, pz, s, p0c, mc2)) 63 | 64 | return par 65 | 66 | return track_a_quadrupole -------------------------------------------------------------------------------- /bmadx/track.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from bmadx import LIB_DICT 3 | 4 | COORDS = ('x', 'px', 'y', 'py', 'z', 'pz') 5 | 6 | def track_element(p_in, ele): 7 | """Tracks Particle 'p_in' though element 'ele'. 8 | Returns outgoing Particle. 9 | """ 10 | lib = sys.modules[type(p_in.x).__module__] 11 | params = [*ele] 12 | for i, param in enumerate(params): 13 | if sys.modules[type(param).__module__] != lib and type(param) != int and type(param) != str: 14 | params[i] = LIB_DICT[lib]['construct_type'](param, dtype = p_in.x.dtype) 15 | 16 | lib_ele = ele._make(params) 17 | track_f = LIB_DICT[lib]['tracking_routine'][type(lib_ele).__name__] 18 | 19 | return track_f(p_in, lib_ele) 20 | 21 | def track_lattice(p_in, ele_list): 22 | """Tracks particle 'p_in' though list of elements 'ele_list'. 23 | Returns outgoing particle 24 | """ 25 | p_out = p_in 26 | for ele in ele_list: 27 | p_out = track_element(p_out, ele) 28 | return p_out 29 | 30 | def track_lattice_save_stats(p_in, ele_list, n_slices: int = 1): 31 | """Tracks particle 'p_in' though list of elements 'ele_list'. 32 | Returns dictionary with stats 33 | """ 34 | sliced_lattice = stub_lattice(ele_list, n_slices) 35 | n = len(sliced_lattice) + 1 36 | stats = {} 37 | stats['s'] = [p_in.s] * n 38 | stats['p0c'] = [p_in.p0c] * n 39 | for coord in COORDS: 40 | array = getattr(p_in, coord) 41 | stats['mean_'+coord] = [array.mean()] * n 42 | stats['sigma_'+coord] = [array.std()] * n 43 | stats['min_'+coord] = [array.min()] * n 44 | stats['max_'+coord] = [array.max()] * n 45 | 46 | p_out = p_in 47 | 48 | for i, ele in enumerate(sliced_lattice): 49 | p_out = track_element(p_out, ele) 50 | stats['s'][i+1] = p_out.s 51 | stats['p0c'][i+1] = p_out.p0c 52 | for coord in COORDS: 53 | array = getattr(p_out, coord) 54 | stats['mean_'+coord][i+1] = array.mean() 55 | stats['sigma_'+coord][i+1] = array.std() 56 | stats['min_'+coord][i+1] = array.min() 57 | stats['max_'+coord][i+1] = array.max() 58 | 59 | stats['p_out'] = p_out 60 | 61 | return stats 62 | 63 | def track_lattice_save_particles(p_in, ele_list, n_slices: int = 1): 64 | """Tracks particle 'p_in' though list of elements 'ele_list'. 65 | Returns list of Particles. all_p[0] is the incoming particle, 66 | and all_p[i] is the particle after the 'i'th element slice. 67 | """ 68 | sliced_lattice = stub_lattice(ele_list, n_slices) 69 | all_p = [None] * ( len(sliced_lattice) + 1 ) 70 | all_p[0] = p_in 71 | lib = sys.modules[type(p_in.x).__module__] 72 | for i, ele in enumerate(sliced_lattice): 73 | all_p[i+1] = track_element(all_p[i], ele) 74 | return all_p 75 | 76 | def make_element(ele, lib): 77 | """makes element named tuple with correct variable types 78 | given library 'lib'.""" 79 | params = [*ele] 80 | for i, param in enumerate(params): 81 | if sys.modules[type(param).__module__] != lib: 82 | params[i] = LIB_DICT[lib]['construct_type'](param, dtype=p.dtype) 83 | 84 | return ele._make(params) 85 | 86 | def stub_lattice(lattice, n): 87 | """Divides every element in the lattice into 'n' elements 88 | each and returns divided lattice. 89 | **NOTE**: only use with Drifts and Quads as of now 90 | """ 91 | stubbed_lattice = [] 92 | 93 | for ele in lattice: 94 | stubbed_lattice.extend(stub_element(ele, n)) 95 | 96 | return stubbed_lattice 97 | 98 | def stub_element(ele, n): 99 | """Divides ele into 'n' equal length elements and returns 100 | a list of these short elements. 101 | **NOTE**: only use with Drifts and Quads with no fringe as of now 102 | """ 103 | short_L = ele.L / n 104 | short_ele = ele._replace(L=short_L) 105 | lattice = [short_ele] * n 106 | 107 | return lattice -------------------------------------------------------------------------------- /bmadx/distgen_utils.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | 3 | import torch 4 | 5 | import distgen 6 | from distgen import Generator 7 | from distgen.physical_constants import unit_registry as unit 8 | 9 | from bmadx.constants import M_ELECTRON 10 | from bmadx.bmad_torch.track_torch import particle_to_beam 11 | from bmadx.pmd_utils import openpmd_to_bmadx_particles 12 | 13 | def create_pmd_particlegroup( 14 | base_yaml, 15 | transforms_yaml, 16 | n_particle 17 | ): 18 | """Creates openPMD ParticleGroup from dist and transform yaml files 19 | using distgen 20 | 21 | Params 22 | ------ 23 | base_yaml: yaml file with base distribution parameters 24 | transform_yaml: yaml file with transforms 25 | 26 | Returns 27 | ------- 28 | generated openPMD ParticleGroup 29 | """ 30 | 31 | gen = Generator(base_yaml) 32 | 33 | with open(transforms_yaml) as f: 34 | transforms_dict = yaml.safe_load(f) 35 | 36 | if distgen.__version__ >= '1.0.0': 37 | gen["transforms"] = transforms_dict 38 | if n_particle is not None: 39 | gen['n_particle'] = n_particle 40 | else: 41 | gen.input["transforms"] = transforms_dict 42 | if n_particle is not None: 43 | gen.input['n_particle'] = n_particle 44 | 45 | particle_group = gen.run() 46 | particle_group.drift_to_z(z=0) 47 | 48 | return particle_group 49 | 50 | 51 | def create_particles( 52 | base_yaml, 53 | transforms_yaml, 54 | p0c, 55 | s = 0.0, 56 | mc2 = M_ELECTRON, 57 | n_particle = None 58 | ): 59 | """ 60 | Creates bmadx Particle distribution from dist and transform yaml files 61 | using distgen 62 | 63 | Params 64 | ------ 65 | base_yaml: yaml file with base distribution parameters 66 | transform_yaml: yaml file with transforms 67 | p0c: reference momentum of Beam coordinates in eV 68 | save_as (str): if provided, saves the generated Beam 69 | 70 | Returns 71 | ------- 72 | generated Bmad-X Particle namedtuple. 73 | """ 74 | 75 | pmd_particle = create_pmd_particlegroup( 76 | base_yaml, 77 | transforms_yaml, 78 | n_particle 79 | ) 80 | 81 | particle = openpmd_to_bmadx_particles( 82 | pmd_particle, 83 | p0c, 84 | s, 85 | mc2 86 | ) 87 | 88 | return particle 89 | 90 | def create_beam( 91 | base_yaml, 92 | transforms_yaml, 93 | p0c, 94 | s = torch.tensor(0.0, dtype=torch.float32), 95 | mc2 = torch.tensor(M_ELECTRON, dtype=torch.float32), 96 | n_particle = None, 97 | save_as = None 98 | ): 99 | """ 100 | Creates bmadx torch Beam from dist and transform yaml files 101 | using distgen 102 | 103 | Params 104 | ------ 105 | base_yaml: yaml file with base distribution parameters 106 | transform_yaml: yaml file with transforms 107 | 108 | p0c: reference momentum of Beam coordinates in eV 109 | save_as (str): if provided, saves the generated Beam 110 | 111 | Returns 112 | ------- 113 | generated Bmad-X torch Beam. 114 | """ 115 | 116 | particle = create_particles( 117 | base_yaml, 118 | transforms_yaml, 119 | p0c, 120 | s, 121 | mc2, 122 | n_particle 123 | ) 124 | 125 | beam = particle_to_beam(particle) 126 | 127 | if save_as is not None: 128 | torch.save(beam, save_as) 129 | print(f'ground truth distribution saved at {save_as}') 130 | 131 | return beam 132 | 133 | 134 | """ 135 | # Generate openPMD particle group: 136 | particle_group = create_pmd_particlegroup(base_yaml, transforms_yaml, n_particle) 137 | 138 | # Transform to Bmad phase space coordinates: 139 | coords = np.array(openpmd_to_bmadx(particle_group, p0c)).T 140 | tkwargs = {"dtype": torch.float32} 141 | coords = torch.tensor(coords, **tkwargs) 142 | 143 | # create Bmad-X torch Beam: 144 | beam = Beam( 145 | coords, 146 | s=torch.tensor(0.0, **tkwargs), 147 | p0c=torch.tensor(p0c, **tkwargs), 148 | mc2=torch.tensor(M_ELECTRON, **tkwargs) 149 | ) 150 | 151 | # save ground truth beam 152 | if save_as is not None: 153 | torch.save(beam, save_as) 154 | print(f'ground truth distribution saved at {save_as}') 155 | 156 | return beam 157 | """ -------------------------------------------------------------------------------- /bmadx/plot.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | from copy import deepcopy 4 | 5 | SPACE_COORDS = ('x', 'y', 'z') 6 | MOMENTUM_COORDS = ('px', 'py', 'pz') 7 | LABELS = { 8 | 'x': 'x', 9 | 'px': 'p_x', 10 | 'y': 'y', 11 | 'py': 'p_y', 12 | 'z': 'z', 13 | 'pz': 'p_z' 14 | } 15 | 16 | 17 | def plot_projections( 18 | particles, 19 | coords = ('x', 'px', 'y', 'py', 'z', 'pz'), 20 | bins = 50, 21 | scale = 1e3, 22 | background = 0, 23 | same_lims = False, 24 | custom_lims = None, 25 | rasterized=True 26 | ): 27 | """ 28 | Plot of coordinates projected into every possible 2D plane. 29 | 30 | Parameters 31 | ---------- 32 | particles: bmadx.Particle 33 | beam to be plotted. If you have a bmadx_torch.Beam object, 34 | you can use the numpy_particle() method to get the bmadx.Particle 35 | 36 | coords: array-like 37 | coordinates that will be plotted. Should be a 38 | subset of ('x', 'px', 'y', 'py', 'z', 'pz'). 39 | Default: ('x', 'px', 'y', 'py', 'z', 'pz') 40 | 41 | bins: int 42 | number of bins in histograms. 43 | Default: 50 44 | 45 | scale: float 46 | scale factor for coordinates (except pz, which is always in %). 47 | 1e3 for milimeters and miliradians, and 1 for meters and radians. 48 | Default: 1e3 49 | 50 | background: bool 51 | if False, 0 frequency pixel of 2d histograms is converted to white. 52 | Default: False 53 | 54 | same_lims: bool 55 | if True, all coords will have the same limits given by the 56 | largest and lowest values in all coords. 57 | Default: False 58 | 59 | custom_lims: array 60 | if provided, sets the lims of histograms for each coords. 61 | if same_lims is Frue, custom lims should have shape 2 62 | providing min and max for every coord. 63 | if same_lims is False, custom lims should have shape 64 | (n_coords x 2). 65 | Default: None 66 | 67 | Returns 68 | ------- 69 | fig and ax pyplot objects with the projections 70 | 71 | """ 72 | 73 | n_coords = len(coords) 74 | 75 | fig_size = (n_coords*2,) * 2 76 | 77 | fig, ax = plt.subplots(n_coords, n_coords, figsize=fig_size) 78 | mycmap = plt.get_cmap('viridis') 79 | mycmap.set_under(color='white') # map 0 to this color 80 | 81 | all_coords = [] 82 | 83 | for coord in coords: 84 | all_coords.append(getattr(particles, coord)) 85 | 86 | all_coords = np.array(all_coords) 87 | 88 | if same_lims: 89 | if custom_lims is None: 90 | coord_min = np.ones(n_coords)*all_coords.min() 91 | coord_max = np.ones(n_coords)*all_coords.max() 92 | elif len(custom_lims) == 2: 93 | coord_min = np.ones(n_coords)*custom_lims[0] 94 | coord_max = np.ones(n_coords)*custom_lims[1] 95 | else: 96 | raise ValueError("custom lims should have shape 2 when same_lims=True") 97 | else: 98 | if custom_lims is None: 99 | coord_min = all_coords.min(axis=1) 100 | coord_max = all_coords.max(axis=1) 101 | elif custom_lims.shape == (n_coords, 2): 102 | coord_min = custom_lims[:,0] 103 | coord_max = custom_lims[:,1] 104 | else: 105 | raise ValueError("custom lims should have shape (n_coords x 2) when same_lims=False") 106 | 107 | for i in range(n_coords): 108 | 109 | x_coord = coords[i] 110 | 111 | if x_coord in SPACE_COORDS and scale==1e3: 112 | x_coord_unit = 'mm' 113 | elif x_coord in SPACE_COORDS and scale==1: 114 | x_coord_unit = 'm' 115 | elif x_coord in MOMENTUM_COORDS and scale==1e3: 116 | x_coord_unit = 'mrad' 117 | elif x_coord in MOMENTUM_COORDS and scale==1: 118 | x_coord_unit = 'rad' 119 | else: 120 | raise ValueError("""scales should be 1 or 1e3, 121 | coords should be a subset of ('x', 'px', 'y', 'py', 'z', 'pz') 122 | """) 123 | 124 | if x_coord=='pz': 125 | x_array = getattr(particles, x_coord)*100 126 | ax[n_coords-1,i].set_xlabel(f'${LABELS[x_coord]}$ (%)') 127 | min_x = coord_min[i]*100 128 | max_x = coord_max[i]*100 129 | if i>0: 130 | ax[i,0].set_ylabel(f'${LABELS[x_coord]}$ (%)') 131 | 132 | else: 133 | x_array = getattr(particles, x_coord)*scale 134 | ax[n_coords-1,i].set_xlabel(f'${LABELS[x_coord]}$ ({x_coord_unit})') 135 | min_x = coord_min[i]*scale 136 | max_x = coord_max[i]*scale 137 | if i>0: 138 | ax[i,0].set_ylabel(f'${LABELS[x_coord]}$ ({x_coord_unit})') 139 | 140 | ax[i,i].hist( 141 | x_array, 142 | bins=bins, 143 | range=([min_x, max_x]), 144 | rasterized = rasterized 145 | ) 146 | 147 | ax[i,i].yaxis.set_tick_params(left=False, labelleft=False) 148 | 149 | if i!= n_coords-1: 150 | ax[i,i].xaxis.set_tick_params(labelbottom=False) 151 | 152 | for j in range(i+1, n_coords): 153 | 154 | y_coord = coords[j] 155 | 156 | if y_coord=='pz': 157 | y_array = getattr(particles, y_coord)*100 158 | min_y = coord_min[j]*100 159 | max_y = coord_max[j]*100 160 | 161 | else: 162 | y_array = getattr(particles, y_coord)*scale 163 | min_y = coord_min[j]*scale 164 | max_y = coord_max[j]*scale 165 | 166 | ax[j,i].hist2d( 167 | x_array, 168 | y_array, 169 | bins = bins, 170 | range=[[min_x, max_x], [min_y, max_y]], 171 | cmap = mycmap, 172 | vmin = background, 173 | rasterized = rasterized 174 | ) 175 | 176 | ax[j,i].sharex(ax[i,i]) 177 | 178 | ax[i,j].set_visible(False) 179 | 180 | if i != 0: 181 | ax[j, i].yaxis.set_tick_params(labelleft=False) 182 | 183 | if j != n_coords-1: 184 | ax[j,i].xaxis.set_tick_params(labelbottom=False) 185 | 186 | fig.tight_layout() 187 | 188 | return fig, ax -------------------------------------------------------------------------------- /bmadx/pmd_utils.py: -------------------------------------------------------------------------------- 1 | from bmadx.constants import C_LIGHT, M_ELECTRON, E_CHARGE 2 | from bmadx.structures import Particle 3 | from bmadx.bmad_torch.track_torch import Beam, particle_to_beam 4 | import numpy as np 5 | import torch 6 | import sys 7 | 8 | from pmd_beamphysics import ParticleGroup 9 | 10 | 11 | def openpmd_to_bmadx_coords( 12 | pmd_particle: ParticleGroup, 13 | p0c 14 | ): 15 | """ 16 | Transforms openPMD-beamphysics ParticleGroup to 17 | bmad phase-space coordinates. 18 | 19 | Parameters: 20 | pmd_particle (pmd_beamphysics.ParticleGroup): openPMD-beamphysics ParticleGroup 21 | p0c (float): reference momentum in eV 22 | 23 | Returns: 24 | bmad_coods (list): list of bmad coords (x, px, y, py, z, pz) 25 | """ 26 | 27 | x = pmd_particle.x 28 | px = pmd_particle.px / p0c 29 | y = pmd_particle.y 30 | py = pmd_particle.py / p0c 31 | z = - pmd_particle.beta * C_LIGHT * pmd_particle.t 32 | pz = pmd_particle.p / p0c - 1.0 33 | 34 | bmad_coords = (x, px, y, py, z, pz) 35 | 36 | return bmad_coords 37 | 38 | 39 | def openpmd_to_bmadx_particles( 40 | pmd_particle: ParticleGroup, 41 | p0c: float, 42 | s : float = 0.0, 43 | mc2 : float = M_ELECTRON 44 | ): 45 | """ 46 | Transforms openPMD-beamphysics ParticleGroup to 47 | bmad phase-space Particle named tuple. 48 | 49 | Parameters: 50 | pmd_particle (pmd_beamphysics.ParticleGroup): openPMD-beamphysics ParticleGroup 51 | p0c (float): reference momentum in eV 52 | 53 | Returns: 54 | Bmadx Particle 55 | """ 56 | coords = openpmd_to_bmadx_coords(pmd_particle, p0c) 57 | particle = Particle( 58 | *coords, 59 | s = s, 60 | p0c = p0c, 61 | mc2 = mc2) 62 | return particle 63 | 64 | 65 | def openpmd_to_bmadx_beam( 66 | pmd_particle: ParticleGroup, 67 | p0c, 68 | s = torch.tensor(0.0, dtype=torch.float32), 69 | mc2 = torch.tensor(M_ELECTRON, dtype=torch.float32) 70 | ): 71 | """ 72 | Transforms openPMD-beamphysics ParticleGroup to 73 | bmad phase-space Particle named tuple. 74 | 75 | Parameters: 76 | pmd_particle (pmd_beamphysics.ParticleGroup): openPMD-beamphysics ParticleGroup 77 | p0c (float): reference momentum in eV 78 | 79 | Returns: 80 | Bmadx torch Beam 81 | """ 82 | particle = openpmd_to_bmadx_particles(pmd_particle, p0c, s, mc2) 83 | beam = particle_to_beam(particle) 84 | return beam 85 | 86 | 87 | def bmadx_particles_to_openpmd(particle: Particle): 88 | """ 89 | Transforms bmadx Particle to openPMD-beamphysics ParticleGroup. 90 | 91 | Parameters 92 | ---------- 93 | particle: bmax Particle 94 | particle to transform. 95 | 96 | Returns 97 | ------- 98 | pmd_beamphysics.ParticleGroup 99 | """ 100 | lib = sys.modules[type(particle.x).__module__] 101 | if lib == np: 102 | x = particle.x 103 | px = particle.px 104 | y = particle.y 105 | py = particle.py 106 | z = particle.z 107 | pz = particle.pz 108 | elif lib == torch: 109 | x = particle.x.detach().numpy() 110 | px = particle.px.detach().numpy() 111 | y = particle.y.detach().numpy() 112 | py = particle.py.detach().numpy() 113 | z = particle.z.detach().numpy() 114 | pz = particle.pz.detach().numpy() 115 | else: 116 | raise ValueError('Only numpy and torch Particles are supported as of now') 117 | 118 | dat = {} 119 | 120 | dat['x'] = x 121 | dat['px'] = px * particle.p0c 122 | dat['y'] = y 123 | dat['py'] = py * particle.p0c 124 | dat['z'] = pz * 0.0 125 | dat['pz'] = particle.p0c * ( (pz + 1.0)**2 - px**2 - py**2 )**0.5 126 | 127 | p = (1 + pz ) * particle.p0c 128 | beta = ( 129 | (p / M_ELECTRON)**2 / 130 | ( 1 + (p / M_ELECTRON)**2 ) 131 | )**0.5 132 | 133 | dat['t'] = - z / (C_LIGHT * beta) 134 | 135 | dat['status'] = np.ones_like(x, dtype=int) 136 | dat['weight'] = np.ones_like(x) * E_CHARGE 137 | 138 | if np.isclose(particle.mc2, M_ELECTRON): 139 | dat['species'] = 'electron' 140 | else: 141 | raise ValueError('only electrons are supported as of now') 142 | 143 | return ParticleGroup(data=dat) 144 | 145 | 146 | def bmadx_beam_to_openpmd(beam: Beam): 147 | """ 148 | Transforms bmadx torch Beam to openPMD-beamphysics ParticleGroup. 149 | 150 | Parameters 151 | ---------- 152 | beam: bmax torch Beam to transform. 153 | 154 | Returns 155 | ------- 156 | pmd_beamphysics.ParticleGroup 157 | """ 158 | particle = beam.numpy_particles() 159 | pmd_particle = bmadx_particles_to_openpmd(particle) 160 | return pmd_particle 161 | 162 | 163 | def save_particles_as_h5(particle: Particle, fname: str): 164 | """ 165 | Saves bmadx Particle as h5 file in openPMD-beamphysics 166 | ParticleGroup standard. 167 | 168 | Parameters 169 | ---------- 170 | particle: bmax Particle 171 | particle to transform. 172 | 173 | fname: str 174 | file name 175 | 176 | Returns 177 | ------- 178 | None 179 | """ 180 | pmd_particle = bmadx_particles_to_openpmd(particle) 181 | pmd_particle.write(fname) 182 | 183 | 184 | def save_beam_as_h5(beam: Beam, fname: str): 185 | """ 186 | Saves bmadx torch Beam as h5 file in openPMD-beamphysics 187 | ParticleGroup standard. 188 | 189 | Parameters 190 | ---------- 191 | beam: bmax torch Beam to transform. 192 | 193 | fname: str 194 | file name 195 | 196 | Returns 197 | ------- 198 | None 199 | """ 200 | pmd_particle = bmadx_beam_to_openpmd(beam) 201 | pmd_particle.write(fname) 202 | 203 | def opal_data_to_bmadx_particle( 204 | opal_data_file: str, 205 | p0c: float = None, 206 | mc2: float = M_ELECTRON 207 | ): 208 | """ 209 | Transforms OPAL particle coordinates in a data file to 210 | Bmad-X Particle beam. 211 | 212 | Parameters 213 | ---------- 214 | opal_data_file (str): OPAL data file with particle coordinates 215 | p0c (float): design momentum times c in eV as defined in Bmad coords 216 | mc2 (float): particle rest mass energy in eV 217 | 218 | Returns 219 | ------- 220 | particle (bmadx.Particle): Bmad-X Particle beam 221 | """ 222 | 223 | data = np.genfromtxt(opal_data_file, skip_header=1) 224 | 225 | pc = mc2 * np.sqrt(data[:,1]**2 + data[:,3]**2 + data[:,5]**2) 226 | 227 | # if not provided, reference momentum is avg p 228 | if p0c is None: 229 | p0c = pc.mean() 230 | 231 | # initial transforms 232 | x = data[:,0] 233 | y = data[:,2] 234 | px = mc2 * data[:,1] / p0c 235 | py = mc2 * data[:,3] / p0c 236 | pz = pc / p0c - 1.0 237 | 238 | # drift to z_avg (so that s value is the same) 239 | z0 = data[:,4].mean() 240 | dt = (z0 - data[:,4]) / data[:,5] 241 | x = x + data[:,1] * dt 242 | y = y + data[:,3] * dt 243 | 244 | # transform z coord 245 | beta = np.sqrt( 246 | (pc / mc2)**2 / 247 | ( 1 + (pc / mc2)**2 ) 248 | ) 249 | z = - beta * C_LIGHT * dt 250 | 251 | # return bmadx particle 252 | particle = Particle( 253 | x = x, 254 | px = px, 255 | y = y, 256 | py = py, 257 | z = z, 258 | pz = pz, 259 | s = 0.0, 260 | p0c = p0c, 261 | mc2 = mc2 262 | ) 263 | return particle -------------------------------------------------------------------------------- /bmadx/bmad_torch/track_torch.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import copy 3 | from torch import Tensor 4 | from torch.nn import Module, ModuleList, Parameter 5 | 6 | from bmadx.structures import Particle 7 | from bmadx.constants import M_ELECTRON 8 | from bmadx import LIB_DICT 9 | 10 | 11 | class Beam(torch.nn.Module): 12 | def __init__(self, data, p0c, s=torch.tensor(0.0), mc2=torch.tensor(M_ELECTRON)): 13 | super(Beam, self).__init__() 14 | self.keys = ["x", "px", "y", "py", "z", "pz"] 15 | 16 | self.register_buffer("data", data) 17 | self.register_parameter("p0c", Parameter(p0c, requires_grad=False)) 18 | self.register_parameter("s", Parameter(s, requires_grad=False)) 19 | self.register_parameter("mc2", Parameter(mc2, requires_grad=False)) 20 | 21 | @property 22 | def x(self): 23 | return self.data[..., 0] 24 | 25 | @property 26 | def px(self): 27 | return self.data[..., 1] 28 | 29 | @property 30 | def y(self): 31 | return self.data[..., 2] 32 | 33 | @property 34 | def py(self): 35 | return self.data[..., 3] 36 | 37 | @property 38 | def z(self): 39 | return self.data[..., 4] 40 | 41 | @property 42 | def pz(self): 43 | return self.data[..., 5] 44 | 45 | def to_list_of_beams(self): 46 | beams = [] 47 | for i in range(len(getattr(self, self.keys[0]))): 48 | beams += [ 49 | Particle( 50 | *[getattr(self, key)[i] for key in self.keys], **self._defaults 51 | ) 52 | ] 53 | 54 | return beams 55 | 56 | def detach_clone(self): 57 | return Beam(self.data.detach().clone(), self.p0c, self.s, self.mc2) 58 | 59 | def numpy_particles(self): 60 | return Particle(*self.data.detach().clone().numpy().T, 61 | p0c = self.p0c.detach().clone().numpy(), 62 | s = self.s.detach().clone().numpy(), 63 | mc2 = self.mc2.detach().clone().numpy()) 64 | 65 | 66 | class TorchElement(Module): 67 | def __init__(self, tracking_function): 68 | super(TorchElement, self).__init__() 69 | self.track = tracking_function 70 | 71 | def forward(self, beam): 72 | #par = self.track(beam, self) 73 | #return particle_to_beam(par) 74 | return self.track(beam, self) 75 | 76 | @property 77 | def batch_shape(self): 78 | out = torch.broadcast_tensors(*[val.data for _, val in self.named_parameters()]) 79 | return out[0].shape 80 | 81 | class TorchDrift(TorchElement): 82 | def __init__( 83 | self, 84 | L: Tensor 85 | ): 86 | super(TorchDrift, self).__init__( 87 | LIB_DICT[torch]['tracking_routine']['Drift'] 88 | ) 89 | self.register_parameter("L", Parameter(L, requires_grad=False)) 90 | 91 | 92 | class TorchQuadrupole(TorchElement): 93 | def __init__( 94 | self, 95 | L: Tensor, 96 | K1: Tensor, 97 | NUM_STEPS: int = 1, 98 | X_OFFSET: Tensor = torch.tensor(0.0), 99 | Y_OFFSET: Tensor = torch.tensor(0.0), 100 | TILT: Tensor = torch.tensor(0.0) 101 | ): 102 | super(TorchQuadrupole, self).__init__( 103 | LIB_DICT[torch]['tracking_routine']['Quadrupole'] 104 | ) 105 | self.register_parameter("L", Parameter(L, requires_grad=False)) 106 | self.register_parameter("X_OFFSET", Parameter(X_OFFSET, requires_grad=False)) 107 | self.register_parameter("Y_OFFSET", Parameter(Y_OFFSET, requires_grad=False)) 108 | self.register_parameter( 109 | "NUM_STEPS", Parameter(torch.tensor(NUM_STEPS), requires_grad=False) 110 | ) 111 | self.register_parameter("TILT", Parameter(TILT, requires_grad=False)) 112 | self.register_parameter("K1", Parameter(K1, requires_grad=False)) 113 | 114 | class TorchSextupole(TorchElement): 115 | def __init__( 116 | self, 117 | L: Tensor, 118 | K2: Tensor, 119 | NUM_STEPS: int = 1, 120 | X_OFFSET: Tensor = torch.tensor(0.0), 121 | Y_OFFSET: Tensor = torch.tensor(0.0), 122 | TILT: Tensor = torch.tensor(0.0), 123 | ): 124 | super(TorchSextupole, self).__init__( 125 | LIB_DICT[torch]['tracking_routine']['Sextupole'] 126 | ) 127 | self.register_parameter("L", Parameter(L, requires_grad=False)) 128 | self.register_parameter("X_OFFSET", Parameter(X_OFFSET, requires_grad=False)) 129 | self.register_parameter("Y_OFFSET", Parameter(Y_OFFSET, requires_grad=False)) 130 | self.register_parameter( 131 | "NUM_STEPS", Parameter(torch.tensor(NUM_STEPS), requires_grad=False) 132 | ) 133 | self.register_parameter("TILT", Parameter(TILT, requires_grad=False)) 134 | self.register_parameter("K2", Parameter(K2, requires_grad=False)) 135 | 136 | class TorchCrabCavity(TorchElement): 137 | def __init__( 138 | self, 139 | L: Tensor, 140 | VOLTAGE: Tensor, 141 | RF_FREQUENCY: Tensor, 142 | PHI0: Tensor = torch.tensor(0.0), 143 | X_OFFSET: Tensor = torch.tensor(0.0), 144 | Y_OFFSET: Tensor = torch.tensor(0.0), 145 | TILT: Tensor = torch.tensor(0.0) 146 | ): 147 | super(TorchCrabCavity, self).__init__( 148 | LIB_DICT[torch]['tracking_routine']['CrabCavity'] 149 | ) 150 | self.register_parameter("L", Parameter(L, requires_grad=False)) 151 | self.register_parameter("VOLTAGE", Parameter(VOLTAGE, requires_grad=False)) 152 | self.register_parameter("PHI0", Parameter(PHI0, requires_grad=False)) 153 | self.register_parameter( 154 | "RF_FREQUENCY", Parameter(RF_FREQUENCY, requires_grad=False) 155 | ) 156 | self.register_parameter("X_OFFSET", Parameter(X_OFFSET, requires_grad=False)) 157 | self.register_parameter("Y_OFFSET", Parameter(Y_OFFSET, requires_grad=False)) 158 | self.register_parameter("TILT", Parameter(TILT, requires_grad=False)) 159 | 160 | 161 | class TorchRFCavity(TorchElement): 162 | def __init__( 163 | self, 164 | L: Tensor, 165 | VOLTAGE: Tensor, 166 | RF_FREQUENCY: Tensor, 167 | PHI0: Tensor = torch.tensor(0.0), 168 | X_OFFSET: Tensor = torch.tensor(0.0), 169 | Y_OFFSET: Tensor = torch.tensor(0.0), 170 | TILT: Tensor = torch.tensor(0.0) 171 | ): 172 | super(TorchRFCavity, self).__init__( 173 | LIB_DICT[torch]['tracking_routine']['RFCavity'] 174 | ) 175 | self.register_parameter("L", Parameter(L, requires_grad=False)) 176 | self.register_parameter("VOLTAGE", Parameter(VOLTAGE, requires_grad=False)) 177 | self.register_parameter("PHI0", Parameter(PHI0, requires_grad=False)) 178 | self.register_parameter( 179 | "RF_FREQUENCY", Parameter(RF_FREQUENCY, requires_grad=False) 180 | ) 181 | self.register_parameter("X_OFFSET", Parameter(X_OFFSET, requires_grad=False)) 182 | self.register_parameter("Y_OFFSET", Parameter(Y_OFFSET, requires_grad=False)) 183 | self.register_parameter("TILT", Parameter(TILT, requires_grad=False)) 184 | 185 | 186 | class TorchSBend(TorchElement): 187 | def __init__( 188 | self, 189 | L: Tensor, 190 | P0C: Tensor, 191 | G: Tensor, 192 | DG: Tensor = torch.tensor(0.0), 193 | E1: Tensor = torch.tensor(0.0), 194 | E2: Tensor = torch.tensor(0.0), 195 | FINT: Tensor = torch.tensor(0.0), 196 | HGAP: Tensor = torch.tensor(0.0), 197 | FINTX: Tensor = torch.tensor(0.0), 198 | HGAPX: Tensor = torch.tensor(0.0), 199 | FRINGE_AT: str = "both_ends", 200 | FRINGE_TYPE: str = "none", 201 | TILT: Tensor = torch.tensor(0.0) 202 | ): 203 | 204 | fringe_at_dic = { 205 | "no_end": 0, 206 | "both_ends": 1, 207 | "entrance_end": 2, 208 | "exit_end":3 209 | } 210 | fringe_type_dic = { 211 | "none": 0, 212 | "soft_edge_only": 1, 213 | "hard_edge_only": 2, 214 | "full": 3, 215 | "linear_edge": 4, 216 | "basic_bend": 5 217 | } 218 | inv_fringe_at_dic = dict(map(reversed, fringe_at_dic.items())) 219 | inv_fringe_type_dic = dict(map(reversed, fringe_type_dic.items())) 220 | 221 | super(TorchSBend, self).__init__( 222 | LIB_DICT[torch]['tracking_routine']['SBend'] 223 | ) 224 | self.register_parameter("L", Parameter(L, requires_grad=False)) 225 | self.register_parameter("P0C", Parameter(P0C, requires_grad=False)) 226 | self.register_parameter("G", Parameter(G, requires_grad=False)) 227 | self.register_parameter("DG", Parameter(DG, requires_grad=False)) 228 | self.register_parameter("E1", Parameter(E1, requires_grad=False)) 229 | self.register_parameter("E2", Parameter(E2, requires_grad=False)) 230 | self.register_parameter("FINT", Parameter(FINT, requires_grad=False)) 231 | self.register_parameter("HGAP", Parameter(HGAP, requires_grad=False)) 232 | self.register_parameter("FINTX", Parameter(FINTX, requires_grad=False)) 233 | self.register_parameter("HGAPX", Parameter(HGAPX, requires_grad=False)) 234 | self.register_parameter("TILT", Parameter(TILT, requires_grad=False)) 235 | self.FRINGE_AT = FRINGE_AT 236 | self.FRINGE_TYPE = FRINGE_TYPE 237 | #self.register_buffer("FRINGE_AT", torch.tensor(fringe_at_dic[FRINGE_AT])) 238 | #self.register_buffer("FRINGE_TYPE", torch.tensor(fringe_type_dic[FRINGE_TYPE])) 239 | 240 | 241 | 242 | class TorchLattice(Module): 243 | def __init__(self, elements, only_last=True): 244 | super(TorchLattice, self).__init__() 245 | self.elements = ModuleList(elements) 246 | self.only_last = only_last 247 | 248 | def forward(self, p_in): 249 | if self.only_last: 250 | p = p_in 251 | for i in range(self.n_elements): 252 | p = self.elements[i](p) 253 | 254 | return p 255 | else: 256 | all_p = [None] * (self.n_elements + 1) 257 | all_p[0] = p_in 258 | 259 | for i in range(self.n_elements): 260 | all_p[i + 1] = self.elements[i](all_p[i]) 261 | 262 | return all_p 263 | 264 | def copy(self): 265 | return copy.deepcopy(self) 266 | 267 | @property 268 | def n_elements(self): 269 | return len(self.elements) 270 | 271 | @property 272 | def batch_shape(self): 273 | out = torch.broadcast_tensors( 274 | *[torch.ones(ele.batch_shape) for ele in self.elements] 275 | ) 276 | return out[0].shape 277 | 278 | def particle_to_beam(particle: Particle): 279 | ''' 280 | Returns Beam corresponding to Particle. 281 | ''' 282 | 283 | if type(particle.x) != torch.Tensor: 284 | n_par = len(particle.x) 285 | coords = torch.zeros((n_par, 6)) 286 | for i in range(6): 287 | coords[:,i] = torch.tensor(particle[i]) 288 | else: 289 | coords = torch.vstack(particle[:6]).T 290 | 291 | params = [0,0,0] 292 | for i in (0, 1, 2): 293 | if type(particle[6+i]) != torch.Tensor: 294 | params[i] = torch.tensor(particle[6+i]) 295 | else: 296 | params[i] = particle[6+i] 297 | 298 | 299 | return Beam( 300 | coords, 301 | s = params[0], 302 | p0c = params[1], 303 | mc2 = params[2] 304 | ) -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /bmadx/tracking_routines/track_a_bend.py: -------------------------------------------------------------------------------- 1 | from bmadx.structures import Particle 2 | from bmadx.constants import PI 3 | from bmadx.low_level.offset_particle import make_offset_particle 4 | 5 | #fringe_at = both_ends (default), no_end, entrance_end, exit_end 6 | #fringe_type = linear_edge (default), none, hard_edge_only, soft_edge_only, full 7 | 8 | #-------------------- 9 | # TRACKING ROUTINES 10 | #-------------------- 11 | 12 | def make_track_a_bend(lib): 13 | 14 | body, fringe_dict = make_track_a_sbend_parts(lib) 15 | offset_particle_set = make_offset_particle(lib, 'set') 16 | offset_particle_unset = make_offset_particle(lib, 'unset') 17 | 18 | def track_a_bend(p_in, bend): 19 | 20 | fringe_type = bend.FRINGE_TYPE.lower() 21 | fringe_at = bend.FRINGE_AT.lower() 22 | 23 | par = offset_particle_set(0.0, 0.0, bend.TILT, p_in) 24 | 25 | implemented_fringes = [ 26 | "none", 27 | "soft_edge_only", 28 | "hard_edge_only", 29 | "linear_edge" 30 | ] 31 | 32 | non_implemented_fringes = [ 33 | "full", 34 | "basic_bend", 35 | "sad_full" 36 | ] 37 | 38 | if fringe_type not in (implemented_fringes + non_implemented_fringes): 39 | raise ValueError(f"Unknown fringe_type setting {fringe_type}!!") 40 | 41 | elif fringe_type in non_implemented_fringes: 42 | raise NotImplementedError("FULL FRINGE to be implemented...") 43 | 44 | elif fringe_type == "none" : 45 | p1 = body(par, bend) 46 | return offset_particle_unset(0.0, 0.0, bend.TILT, p1) 47 | 48 | else: 49 | 50 | if fringe_at == "both_ends": 51 | p1 = fringe_dict[fringe_type]["entrance"](par, bend) 52 | p2 = body(p1, bend) 53 | p3 = fringe_dict[fringe_type]["exit"](p2, bend) 54 | return offset_particle_unset(0.0, 0.0, bend.TILT, p3) 55 | 56 | elif fringe_at == "entrance_end": 57 | p1 = fringe_dict[fringe_type]["entrance"](par, bend) 58 | p2 = body(p1, bend) 59 | return offset_particle_unset(0.0, 0.0, bend.TILT, p2) 60 | 61 | elif fringe_at == "exit_end": 62 | p1 = body(par, bend) 63 | p2 = fringe_dict[fringe_type]["exit"](p1, bend) 64 | return offset_particle_unset(0.0, 0.0, bend.TILT, p2) 65 | 66 | elif fringe_at=="no_end": 67 | p1 = body(par, bend) 68 | return offset_particle_unset(0.0, 0.0, bend.TILT, p1) 69 | 70 | else: 71 | raise ValueError(f"Unknown fringe_at setting {fringe_at}!!") 72 | 73 | return track_a_bend 74 | 75 | 76 | def make_track_a_sbend_parts(lib): 77 | """Makes track_a_quadrupole given the library lib.""" 78 | sqrt = lib.sqrt 79 | sin = lib.sin 80 | cos = lib.cos 81 | tan = lib.tan 82 | sinc = lib.sinc 83 | absolute = lib.abs 84 | arcsin = lib.arcsin 85 | arctan2 = lib.arctan2 86 | 87 | def sinc_bmad(x): 88 | return sinc(x/PI) 89 | 90 | def cosc(x): 91 | # branchless: (cos(x)-1)/x**2 = -1/2 [sinc(x/2)]**2 92 | return -sinc(x/PI/2)**2/2 93 | 94 | def track_body(p_in, bend): 95 | """Tracks the incoming Particle p_in though a sbend (k1=0) element 96 | and returns the outgoing particle. 97 | See Bmad manual section 24.9 98 | """ 99 | 100 | # WILLIAM's CODE HERE 101 | L = bend.L 102 | g = bend.G 103 | dg = bend.DG 104 | p0c = bend.P0C 105 | 106 | p0c_particle = p_in.p0c 107 | 108 | x = p_in.x 109 | px = p_in.px * p0c_particle / p0c 110 | y = p_in.y 111 | py = p_in.py * p0c_particle / p0c 112 | z = p_in.z 113 | pz = (p_in.pz + 1) * p0c_particle / p0c - 1 114 | s = p_in.s 115 | mc2 = p_in.mc2 116 | 117 | # Calculate the final 6 corrds 118 | 119 | px_norm = sqrt((1 + pz) ** 2 - py ** 2) # For simplicity 120 | 121 | phi1 = arcsin(px / px_norm) 122 | 123 | # g = theta / L 124 | theta = g*L 125 | g_tot = g + dg 126 | gp = g_tot / px_norm 127 | 128 | alpha = (2 * (1 + g * x) * sin(theta + phi1) * L * sinc_bmad(theta) 129 | - gp * ((1 + g * x) * L * sinc_bmad(theta)) ** 2) 130 | 131 | x2_t1 = x * cos(theta) + L ** 2 * g * cosc(theta) 132 | x2_t2 = sqrt((cos(theta + phi1) ** 2) + gp * alpha) 133 | x2_t3 = cos(theta + phi1) 134 | 135 | #x2 = where(absolute(theta + phi1) < pi / 2, 136 | # (x2_t1 + alpha / (x2_t2 + x2_t3)), 137 | # (x2_t1 + (x2_t2 - x2_t3) / gp)) 138 | 139 | temp = absolute(theta + phi1) 140 | c1 = (x2_t1 + alpha / (x2_t2 + x2_t3)) 141 | c2 = (x2_t1 + (x2_t2 - x2_t3) / gp) 142 | 143 | x2 = c1 * (temp < PI/2) + c2 * (temp >= PI/2) 144 | 145 | 146 | Lcu = x2 - L ** 2 * g * cosc(theta) - x * cos(theta) 147 | Lcv = -L * sinc_bmad(theta) - x * sin(theta) 148 | 149 | theta_p = 2 * (theta + phi1 - PI / 2 - arctan2(Lcv, Lcu)) 150 | 151 | Lc = sqrt(Lcu ** 2 + Lcv ** 2) 152 | Lp = Lc / sinc_bmad(theta_p / 2) 153 | 154 | P = p0c * (1 + pz) # in eV 155 | E = sqrt(P**2 + mc2**2) # in eV 156 | E0 = sqrt(p0c**2 + mc2**2) # in eV 157 | beta = P / E 158 | beta0 = p0c / E0 159 | 160 | xf = x2 161 | pxf = px_norm * sin(theta + phi1 - theta_p) 162 | yf = y + py * Lp / px_norm 163 | pyf = py 164 | zf = z + (beta * L / beta0) - ((1 + pz) * Lp / px_norm) 165 | pzf = pz 166 | 167 | s = s+L 168 | 169 | return Particle(xf, pxf, yf, pyf, zf, pzf, s, p0c, mc2) 170 | 171 | 172 | def track_entrance_hard(p_in, bend): 173 | L = bend.L 174 | g = bend.G 175 | dg = bend.DG 176 | p0c = bend.P0C 177 | e1 = bend.E1 178 | f_int = bend.FINT 179 | h_gap = bend.HGAP 180 | 181 | p0c_particle = p_in.p0c 182 | 183 | x = p_in.x 184 | px = p_in.px * p0c_particle / p0c 185 | y = p_in.y 186 | py = p_in.py * p0c_particle / p0c 187 | z = p_in.z 188 | pz = (p_in.pz + 1) * p0c_particle / p0c - 1 189 | s = p_in.s 190 | mc2 = p_in.mc2 191 | 192 | theta = g*L 193 | g_tot = g + dg 194 | 195 | sin1 = sin(e1) 196 | tan1 = tan(e1) 197 | sec1 = 1 / cos(e1) 198 | 199 | Sigma_M1 = ( 200 | (x**2 - y**2) * g_tot*tan1/2 201 | + y**2 * g_tot**2 * sec1**3 * (1 + sin1**2) * f_int*h_gap /2/(1 + pz) 202 | - x**3 * g_tot**2 * tan1**3 /12/(1 + pz) 203 | + x*y**2 * g_tot**2 * tan1 * sec1**2/4/(1 + pz) 204 | + (x**2 * px - x*y*py) * g_tot*tan1**2 /2/(1 + pz) 205 | + y**2 * px * g_tot * (1 + tan1**2) / 2 / (1 + pz) 206 | ) 207 | 208 | #Sigma_M1 = ( 209 | # (x**2 - y**2) * g_tot*tan1/2 210 | # + y**2 * g_tot**2 * sec1**3 * (1 + sin1**2) * f_int*h_gap /2/(1 + pz) 211 | # - x**3 * g_tot**2 * tan1**3 /12/(1 + pz) 212 | # + x*y**2 * g_tot**2 * tan1 * sec1**2/4/(1 + pz) 213 | # + (x**2 * px - 2*x*y*py) * g_tot*tan1**2 /2/(1 + pz) 214 | # - y**2 * px * g_tot * (1 + tan1**2) / 2 / (1 + pz) 215 | #) 216 | 217 | xf = x + g_tot / 2 / (1 + pz) * (-x**2 * tan1**2 + y**2 * sec1**2) 218 | 219 | pxf = ( px 220 | + x * g_tot * tan1 221 | + y**2 * g_tot**2 * (tan1 + 2*tan1**3) /2/(1 + pz) 222 | + (x*px - y*py) * g_tot * tan1**2 /(1 + pz) 223 | ) 224 | 225 | yf = y + x * y * g_tot * tan1 ** 2 / (1 + pz) 226 | 227 | pyf = ( py 228 | + y * ( -g_tot * tan1 + g_tot**2 * (1 + sin1**2) * sec1**3 * f_int * h_gap / (1 + pz)) 229 | - x * py * g_tot * tan1**2 / (1 + pz) 230 | - y * px * g_tot * (1 + tan1**2) / (1 + pz) 231 | ) 232 | 233 | zf = z + (Sigma_M1 - (x**2 - y**2) * g_tot * tan1/2) / (1 + pz) 234 | 235 | pzf = pz 236 | 237 | 238 | return Particle(xf, pxf, yf, pyf, zf, pzf, s, p0c, mc2) 239 | 240 | 241 | def track_exit_hard(p_in, bend): 242 | L = bend.L 243 | g = bend.G 244 | dg = bend.DG 245 | p0c = bend.P0C 246 | e2 = bend.E2 247 | f_int = bend.FINTX 248 | h_gap = bend.HGAPX 249 | 250 | p0c_particle = p_in.p0c 251 | 252 | x = p_in.x 253 | px = p_in.px * p0c_particle / p0c 254 | y = p_in.y 255 | py = p_in.py * p0c_particle / p0c 256 | z = p_in.z 257 | pz = (p_in.pz + 1) * p0c_particle / p0c - 1 258 | s = p_in.s 259 | mc2 = p_in.mc2 260 | 261 | theta = g*L 262 | g_tot = g + dg 263 | 264 | sin2 = sin(e2) 265 | tan2 = tan(e2) 266 | sec2 = 1 / cos(e2) 267 | 268 | 269 | Sigma_M2 = ( 270 | (x**2 - y**2) * g_tot * tan2 / 2 271 | + y**2 * g_tot**2 * sec2**3 * (1 + sin2**2) * f_int * h_gap / 2 / (1 + pz) 272 | - x**3 * g_tot**2 * tan2**3 / 12 / (1 + pz) 273 | + x*y**2 * g_tot**2 * tan2 * sec2**2 / 4 / (1 + pz) 274 | + (-x**2 * px + x*y*py) * g_tot * tan2**2 / 2 / (1 + pz) 275 | + y**2*px * g_tot * (1 + tan2**2) / 2 / (1 + pz) 276 | ) 277 | 278 | xf = x + g_tot / 2 / (1 + pz) * (x**2 * tan2**2 - y ** 2 * sec2**2) 279 | 280 | pxf = ( 281 | px + x * g_tot * tan2 282 | - (x**2 + y**2) * g_tot**2 * tan2**3 / 2 / (1 + pz) 283 | + (-x * px + y * py) * g_tot * tan2**2 / (1 + pz) 284 | ) 285 | 286 | yf = y - x * y * g_tot * tan2**2 / (1 + pz) 287 | 288 | pyf = ( 289 | py + y * ( -g_tot * tan2 + g_tot**2 * (1 + sin2**2) * sec2**3 * f_int * h_gap / (1 + pz) ) 290 | + x * y * g_tot**2 * sec2**2 * tan2 / (1 + pz) 291 | + x * py * g_tot * tan2**2 / (1 + pz) 292 | + y * px * g_tot * (1 + tan2**2) / (1 + pz) 293 | ) 294 | 295 | zf = z + (Sigma_M2 - (x**2 - y**2) * g_tot * tan2 / 2) / (1 + pz) 296 | 297 | pzf = pz 298 | 299 | return Particle(xf, pxf, yf, pyf, zf, pzf, s, p0c, mc2) 300 | 301 | 302 | def track_entrance_soft(p_in, bend): 303 | L = bend.L 304 | g = bend.G 305 | dg = bend.DG 306 | p0c = bend.P0C 307 | e1 = bend.E1 308 | f_int = bend.FINT 309 | h_gap = bend.HGAP 310 | 311 | p0c_particle = p_in.p0c 312 | 313 | x = p_in.x 314 | px = p_in.px * p0c_particle / p0c 315 | y = p_in.y 316 | py = p_in.py * p0c_particle / p0c 317 | z = p_in.z 318 | pz = (p_in.pz + 1) * p0c_particle / p0c - 1 319 | s = p_in.s 320 | mc2 = p_in.mc2 321 | 322 | g_tot = g + dg 323 | 324 | FH1 = f_int*h_gap 325 | c1 = 6*g_tot * FH1**2/(1+pz) 326 | c2 = 2*g_tot**2 * FH1/(1+pz) 327 | c3 = g_tot**2 /18 /FH1/(1+pz) 328 | 329 | xf = x + c1*pz 330 | 331 | pxf = px 332 | 333 | yf = y 334 | 335 | pyf = py + c2*y - c3*y**3 336 | 337 | zf = z + (c1*px + c2*y**2/2 - c3*y**4/4) / (1 + pz) 338 | 339 | pzf = pz 340 | 341 | return Particle(xf, pxf, yf, pyf, zf, pzf, s, p0c, mc2) 342 | 343 | 344 | def track_exit_soft(p_in, bend): 345 | L = bend.L 346 | g = bend.G 347 | dg = bend.DG 348 | p0c = bend.P0C 349 | e1 = bend.E1 350 | f_int = bend.FINTX 351 | h_gap = bend.HGAPX 352 | 353 | p0c_particle = p_in.p0c 354 | 355 | x = p_in.x 356 | px = p_in.px * p0c_particle / p0c 357 | y = p_in.y 358 | py = p_in.py * p0c_particle / p0c 359 | z = p_in.z 360 | pz = (p_in.pz + 1) * p0c_particle / p0c - 1 361 | s = p_in.s 362 | mc2 = p_in.mc2 363 | 364 | g_tot = -g - dg # sign reversed for soft exit fringe 365 | 366 | FH1 = f_int*h_gap 367 | c1 = 6*g_tot * FH1**2/(1+pz) 368 | c2 = 2*g_tot**2 * FH1/(1+pz) 369 | c3 = g_tot**2 /18 /FH1/(1+pz) 370 | 371 | xf = x + c1*pz 372 | 373 | pxf = px 374 | 375 | yf = y 376 | 377 | pyf = py + c2*y - c3*y**3 378 | 379 | zf = z + (c1*px + c2*y**2/2 - c3*y**4/4) / (1 + pz) 380 | 381 | pzf = pz 382 | 383 | return Particle(xf, pxf, yf, pyf, zf, pzf, s, p0c, mc2) 384 | 385 | 386 | def track_entrance_linear(p_in, bend): 387 | g = bend.G 388 | dg = bend.DG 389 | e1 = bend.E1 390 | f_int = bend.FINT 391 | h_gap = bend.HGAP 392 | 393 | g_tot = g + dg 394 | hx = g_tot * tan(e1) 395 | hy = -g_tot*tan( 396 | e1 - 397 | 2 * f_int * h_gap * g_tot * ( 1 + sin(e1)**2 ) / 398 | cos(e1) 399 | ) 400 | px_f = p_in.px + p_in.x * hx 401 | py_f = p_in.py + p_in.y * hy 402 | 403 | return Particle( 404 | p_in.x, 405 | px_f, 406 | p_in.y, 407 | py_f, 408 | p_in.z, 409 | p_in.pz, 410 | p_in.s, 411 | p_in.p0c, 412 | p_in.mc2) 413 | 414 | def track_exit_linear(p_in, bend): 415 | g = bend.G 416 | dg = bend.DG 417 | e2 = bend.E2 418 | f_int = bend.FINTX 419 | h_gap = bend.HGAPX 420 | 421 | g_tot = g + dg 422 | hx = g_tot * tan(e2) 423 | hy = -g_tot * tan( 424 | e2 - 425 | 2 * f_int * h_gap * g_tot * ( 1 + sin(e2)**2 ) / 426 | cos(e2) 427 | ) 428 | px_f = p_in.px + p_in.x * hx 429 | py_f = p_in.py + p_in.y * hy 430 | 431 | return Particle( 432 | p_in.x, 433 | px_f, 434 | p_in.y, 435 | py_f, 436 | p_in.z, 437 | p_in.pz, 438 | p_in.s, 439 | p_in.p0c, 440 | p_in.mc2) 441 | 442 | 443 | fringe_dict = { 444 | "hard_edge_only": { 445 | "entrance": track_entrance_hard, 446 | "exit": track_exit_hard 447 | }, 448 | "soft_edge_only": { 449 | "entrance": track_entrance_soft, 450 | "exit": track_exit_soft 451 | }, 452 | "linear_edge": { 453 | "entrance": track_entrance_linear, 454 | "exit": track_exit_linear 455 | } 456 | } 457 | 458 | return track_body, fringe_dict 459 | 460 | 461 | -------------------------------------------------------------------------------- /dev/torch_module_profiling/a: -------------------------------------------------------------------------------- 1 | Memory usage: ▁▁▁▂▂▃▃▃▃▄▄▅▅▆▆▆▇▇▇██████▇▇ (max: 270.385 MB, growth rate: 81%) 2 | /global/homes/j/jpga/Repositories/Bmad-X/dev/torch_module_profiling/torch_module_profiling.py: % of time = 0.00% (0.000ms) out of 1.604s. 3 | ╷ ╷ ╷ ╷ ╷ ╷ ╷ ╷ 4 | │Time │–––––– │–––––– │Memory │–––––– │––––––––––– │Copy │ 5 | Line │Python │native │system │Python │peak │timeline/% │(MB/s) │/global/homes/j/jpga/Repositories/Bmad-X/dev/torch_module_profiling/torch_module_profiling.py 6 | ╺━━━━━━┿━━━━━━━┿━━━━━━━┿━━━━━━━┿━━━━━━━━┿━━━━━━━┿━━━━━━━━━━━━━━━┿━━━━━━━┿━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸ 7 | 1 │ │ │ │ │ │ │ │#from scalene import scalene_profiler 8 | 2 │ │ │ │ │ │ │ │ 9 | 3 │ │ │ │ │ │ │ │import sys 10 | 4 │ │ │ │ │ │ │ │#print(sys.argv) 11 | 5 │ │ │ │ │ │ │ │n_particles = int(sys.argv[1]) 12 | 6 │ │ │ │ │ │ │ │n_quadrupoles = int(sys.argv[2]) 13 | 7 │ │ │ │ │ │ │ │n_slices = int(sys.argv[3]) 14 | 8 │ │ │ │ │ │ │ │ 15 | 9 │ │ │ │ 62% │ 60M │▁▁▁▁▁▁ 22% │ 63 │import torch 16 | 10 │ │ │ │ │ │ │ │ 17 | 11 │ │ │ │ │ │ │ │from bmadx import Particle, Drift, Quadrupole, track_lattice, M_ELECTRON 18 | 12 │ │ │ │ │ │ │ │ 19 | 13 │ │ │ │ │ │ │ │def create_gaussian_beam(n_particles, mean, cov, s, p0c, mc2): 20 | 14 │ │ │ │ │ │ │ │ torch.manual_seed(0) 21 | 15 │ │ │ │ │ │ │ │ dist = torch.distributions.multivariate_normal.MultivariateNormal(mean, cov) 22 | 16 │ │ │ │ │ │ │ │ coords = dist.sample(torch.Size([n_particles])) # particles' coordinates 23 | 17 │ │ │ │ │ │ │ │ 24 | 18 │ │ │ │ │ │ │ │ return Particle(*coords.T, s, p0c, mc2) 25 | 19 │ │ │ │ │ │ │ │ 26 | 20 │ │ │ │ │ │ │ │class BeamlineModel(torch.nn.Module): 27 | 21 │ │ │ │ │ │ │ │ def __init__(self, k_set, l_d, l_q, beam_in, sigma_t): 28 | 22 │ │ │ │ │ │ │ │ super().__init__() 29 | 23 │ │ │ │ │ │ │ │ self.register_parameter('k_set', torch.nn.Parameter(k_set)) 30 | 24 │ │ │ │ │ │ │ │ self.l_d = l_d 31 | 25 │ │ │ │ │ │ │ │ self.l_q = l_q 32 | 26 │ │ │ │ │ │ │ │ self.beam_in = beam_in 33 | 27 │ │ │ │ │ │ │ │ self.sigma_t = sigma_t 34 | 28 │ │ │ │ │ │ │ │ def forward(self): 35 | 29 │ │ │ │ │ │ │ │ # create lattice 36 | 30 │ │ │ │ │ │ │ │ lattice = [] 37 | 31 │ │ │ │ │ │ │ │ half_drift = Drift( L = self.l_d/2) 38 | 32 │ │ │ │ │ │ │ │ for k in self.k_set: 39 | 33 │ │ │ │ │ │ │ │ lattice.append( half_drift ) 40 | 34 │ │ │ │ │ │ │ │ lattice.append( Quadrupole(L = self.l_q, 41 | 35 │ │ │ │ │ │ │ │ K1 = k, 42 | 36 │ │ │ │ │ │ │ │ NUM_STEPS = n_slices) ) 43 | 37 │ │ │ │ │ │ │ │ lattice.append( half_drift ) 44 | 38 │ │ │ │ │ │ │ │ 45 | 39 │ │ │ │ 14% │ 210M │▃▃▃▄▅▅▆▆▆ 78% │ 109 │ beam_out = track_lattice(self.beam_in, lattice) 46 | 40 │ │ │ │ │ │ │ │ dx = beam_out.x.std() - self.sigma_t 47 | 41 │ │ │ │ │ │ │ │ dy = beam_out.y.std() - self.sigma_t 48 | 42 │ │ │ │ │ │ │ │ delta = torch.sqrt( dx**2 + dy**2 ) 49 | 43 │ │ │ │ │ │ │ │ 50 | 44 │ │ │ │ │ │ │ │ return delta 51 | 45 │ │ │ │ │ │ │ │ 52 | 46 │ │ │ │ │ │ │ │# model constants 53 | 47 │ │ │ │ │ │ │ │l_d = 0.9 # drift length 54 | 48 │ │ │ │ │ │ │ │l_q = 0.1 # quad length 55 | 49 │ │ │ │ │ │ │ │sigma_t = 5e-3 # target beam size 56 | 50 │ │ │ │ │ │ │ │ 57 | 51 │ │ │ │ │ │ │ │# incoming beam 58 | 52 │ │ │ │ │ │ │ │beam_in = create_gaussian_beam(n_particles, 59 | 53 │ │ │ │ │ │ │ │ mean = torch.zeros(6)*1e-3, 60 | 54 │ │ │ │ │ │ │ │ cov = torch.eye(6)*1e-6, 61 | 55 │ │ │ │ │ │ │ │ s = torch.tensor(0.0), 62 | 56 │ │ │ │ │ │ │ │ p0c = torch.tensor(4e7), 63 | 57 │ │ │ │ │ │ │ │ mc2 = torch.tensor(M_ELECTRON)) 64 | 58 │ │ │ │ │ │ │ │ 65 | 59 │ │ │ │ │ │ │ │# model evaluation and backward pass 66 | 60 │ │ │ │ │ │ │ │k_set = torch.zeros(n_quadrupoles, requires_grad=True) 67 | 61 │ │ │ │ │ │ │ │#scalene_profiler.start() 68 | 62 │ │ │ │ │ │ │ │model = BeamlineModel(k_set, l_d, l_q, beam_in, sigma_t) 69 | 63 │ │ │ │ │ │ │ │loss = model() 70 | 64 │ │ │ │ │ │ │ │loss.backward() 71 | 65 │ │ │ │ │ │ │ │#scalene_profiler.stop() 72 | ╵ ╵ ╵ ╵ ╵ ╵ ╵ ╵ 73 | Top PEAK memory consumption, by line: 74 | (1) 39: 210 MB 75 | (2) 9: 60 MB 76 | -------------------------------------------------------------------------------- /docs/examples/numba_tests.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "ccb25fd7-91e5-4291-826f-b9a18bbd6b29", 6 | "metadata": {}, 7 | "source": [ 8 | "# Bmad-X Numpy, Numba, Pytorch tests\n", 9 | "\n", 10 | "\n", 11 | "Note: on Cori `conda install cudatoolkit=11.0`" 12 | ] 13 | }, 14 | { 15 | "cell_type": "code", 16 | "execution_count": 1, 17 | "id": "6c8860e7-cf3f-44d4-84c6-94b67131863a", 18 | "metadata": {}, 19 | "outputs": [], 20 | "source": [ 21 | "from bmadx.track import make_track_a_drift, Particle, Drift\n", 22 | "\n", 23 | "import numpy as np\n", 24 | "\n", 25 | "import math\n", 26 | "\n", 27 | "import matplotlib.pyplot as plt" 28 | ] 29 | }, 30 | { 31 | "cell_type": "code", 32 | "execution_count": 2, 33 | "id": "a8147c00-824b-4ab1-9caf-1373fe1f71ef", 34 | "metadata": {}, 35 | "outputs": [], 36 | "source": [ 37 | "c_light = 2.99792458e8 #speed of light in m/s\n", 38 | "m_e = 0.510998950e6 #electron mass in eV" 39 | ] 40 | }, 41 | { 42 | "cell_type": "code", 43 | "execution_count": 3, 44 | "id": "767bb530-83aa-469c-8daf-bb17a707d918", 45 | "metadata": {}, 46 | "outputs": [], 47 | "source": [ 48 | "s = 0.0 #initial s\n", 49 | "p0c = 4.0E+07 #Reference particle momentum in eV\n", 50 | "mc2 = 1*m_e # electron mass in eV\n", 51 | "#pvec1 = [2e-3,3e-3,-3e-3,-1e-3,2e-3,-2e-3] " 52 | ] 53 | }, 54 | { 55 | "cell_type": "markdown", 56 | "id": "d6277d75-44f5-4cf9-9960-870756ccb90f", 57 | "metadata": {}, 58 | "source": [ 59 | "## Numpy\n", 60 | "\n", 61 | "Create 10 million test particles" 62 | ] 63 | }, 64 | { 65 | "cell_type": "code", 66 | "execution_count": 4, 67 | "id": "8b8098f6-c9ea-4da4-a7c4-9283106f4d71", 68 | "metadata": {}, 69 | "outputs": [ 70 | { 71 | "data": { 72 | "text/plain": [ 73 | "0.0010000745670385076" 74 | ] 75 | }, 76 | "execution_count": 4, 77 | "metadata": {}, 78 | "output_type": "execute_result" 79 | } 80 | ], 81 | "source": [ 82 | "N_PARTICLE = 10_000_000\n", 83 | "\n", 84 | "np.random.seed(999)\n", 85 | "\n", 86 | "pvec0 = np.random.normal( size=(N_PARTICLE, 6), scale=.001)\n", 87 | "np.std(pvec0[:, 0])" 88 | ] 89 | }, 90 | { 91 | "cell_type": "code", 92 | "execution_count": 5, 93 | "id": "2c56d4a8-f8b1-4806-9bb0-c23a718a8e2b", 94 | "metadata": {}, 95 | "outputs": [ 96 | { 97 | "data": { 98 | "text/plain": [ 99 | "array([ 0.00012716, 0.00156627, -0.00111006, ..., -0.00063516,\n", 100 | " -0.00116169, 0.00162879])" 101 | ] 102 | }, 103 | "execution_count": 5, 104 | "metadata": {}, 105 | "output_type": "execute_result" 106 | } 107 | ], 108 | "source": [ 109 | "P0 = Particle(pvec0[:,0],\n", 110 | " pvec0[:,1],\n", 111 | " pvec0[:,2],\n", 112 | " pvec0[:,3],\n", 113 | " pvec0[:,4],\n", 114 | " pvec0[:,5],\n", 115 | " s, p0c, mc2)\n", 116 | "P0.x" 117 | ] 118 | }, 119 | { 120 | "cell_type": "code", 121 | "execution_count": 6, 122 | "id": "27ae8000-80fb-4c38-82f7-54a8aa333ef1", 123 | "metadata": {}, 124 | "outputs": [], 125 | "source": [ 126 | "track_a_drift = make_track_a_drift(np)\n", 127 | "D1 = Drift(L=1)" 128 | ] 129 | }, 130 | { 131 | "cell_type": "code", 132 | "execution_count": 7, 133 | "id": "a94f2b9a-88a4-43b3-89cd-64267aa482bc", 134 | "metadata": {}, 135 | "outputs": [ 136 | { 137 | "data": { 138 | "text/plain": [ 139 | "0.0014142797576549417" 140 | ] 141 | }, 142 | "execution_count": 7, 143 | "metadata": {}, 144 | "output_type": "execute_result" 145 | } 146 | ], 147 | "source": [ 148 | "P1 = track_a_drift(P0, D1)\n", 149 | "np.std(P1.x)" 150 | ] 151 | }, 152 | { 153 | "cell_type": "code", 154 | "execution_count": 8, 155 | "id": "a4964de6-44e2-46ac-b5f2-07c0b4fdce1b", 156 | "metadata": {}, 157 | "outputs": [ 158 | { 159 | "name": "stdout", 160 | "output_type": "stream", 161 | "text": [ 162 | "297 ms ± 1.79 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" 163 | ] 164 | } 165 | ], 166 | "source": [ 167 | "%%timeit\n", 168 | "P1 = track_a_drift(P0, D1)" 169 | ] 170 | }, 171 | { 172 | "cell_type": "markdown", 173 | "id": "398f8b79-c3b1-4c31-8465-dd439b12f00f", 174 | "metadata": {}, 175 | "source": [ 176 | "## Numba CPU" 177 | ] 178 | }, 179 | { 180 | "cell_type": "code", 181 | "execution_count": 9, 182 | "id": "d10fefbf-87fe-43ea-b1c9-1f0f53b22051", 183 | "metadata": {}, 184 | "outputs": [ 185 | { 186 | "data": { 187 | "text/plain": [ 188 | "128" 189 | ] 190 | }, 191 | "execution_count": 9, 192 | "metadata": {}, 193 | "output_type": "execute_result" 194 | } 195 | ], 196 | "source": [ 197 | "import numba\n", 198 | "from numba import guvectorize, float64, jit\n", 199 | "numba.config.NUMBA_NUM_THREADS" 200 | ] 201 | }, 202 | { 203 | "cell_type": "code", 204 | "execution_count": 10, 205 | "id": "f64b78da-768b-41b2-abd7-93ade74592a4", 206 | "metadata": {}, 207 | "outputs": [], 208 | "source": [ 209 | "track_a_drift = make_track_a_drift(np)\n", 210 | "#track_a_drift = numba.njit(make_track_a_drift(np))" 211 | ] 212 | }, 213 | { 214 | "cell_type": "code", 215 | "execution_count": 11, 216 | "id": "c88fbeb4-2aa4-407b-81b1-7447a6dabc01", 217 | "metadata": {}, 218 | "outputs": [], 219 | "source": [ 220 | "#%%timeit\n", 221 | "#track_a_drift(P0, D1)" 222 | ] 223 | }, 224 | { 225 | "cell_type": "code", 226 | "execution_count": 12, 227 | "id": "6337f68f-7e35-48e9-b041-6a84cef4adf9", 228 | "metadata": {}, 229 | "outputs": [ 230 | { 231 | "data": { 232 | "text/plain": [ 233 | "0.0014142797576549417" 234 | ] 235 | }, 236 | "execution_count": 12, 237 | "metadata": {}, 238 | "output_type": "execute_result" 239 | } 240 | ], 241 | "source": [ 242 | "params = D1\n", 243 | "g = numba.njit( make_track_a_drift(np))\n", 244 | "\n", 245 | "\n", 246 | "@guvectorize([(float64[:], float64[:])], '(n)->(n)')\n", 247 | "def vg_numba(a_in, a_out):\n", 248 | " p_in = Particle(x = a_in[0],\n", 249 | " px = a_in[1],\n", 250 | " y = a_in[2],\n", 251 | " py = a_in[3],\n", 252 | " z = a_in[4],\n", 253 | " pz = a_in[5],\n", 254 | " s=s, p0c=p0c, mc2=mc2)\n", 255 | " p_out = g(p_in, params)\n", 256 | " a_out[0] = p_out.x\n", 257 | " a_out[1] = p_out.px \n", 258 | " a_out[2] = p_out.y\n", 259 | " a_out[3] = p_out.py \n", 260 | " a_out[4] = p_out.z\n", 261 | " a_out[5] = p_out.pz \n", 262 | "\n", 263 | "@guvectorize([(float64[:], float64[:])], '(n)->(n)', target='parallel')\n", 264 | "def vg_numba_parallel(a_in, a_out):\n", 265 | " p_in = Particle(x = a_in[0],\n", 266 | " px = a_in[1],\n", 267 | " y = a_in[2],\n", 268 | " py = a_in[3],\n", 269 | " z = a_in[4],\n", 270 | " pz = a_in[5],\n", 271 | " s=s, p0c=p0c, mc2=mc2)\n", 272 | " p_out = g(p_in, params)\n", 273 | " a_out[0] = p_out.x\n", 274 | " a_out[1] = p_out.px \n", 275 | " a_out[2] = p_out.y\n", 276 | " a_out[3] = p_out.py \n", 277 | " a_out[4] = p_out.z\n", 278 | " a_out[5] = p_out.pz \n", 279 | " \n", 280 | "\n", 281 | "pvec1 = np.zeros_like(pvec0) \n", 282 | "vg_numba_parallel(pvec0, pvec1)\n", 283 | "np.std(pvec1[:,0])" 284 | ] 285 | }, 286 | { 287 | "cell_type": "code", 288 | "execution_count": 13, 289 | "id": "bead6f91-6c24-44f5-9142-aa5d62c672c2", 290 | "metadata": {}, 291 | "outputs": [ 292 | { 293 | "name": "stdout", 294 | "output_type": "stream", 295 | "text": [ 296 | "144 ms ± 212 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" 297 | ] 298 | } 299 | ], 300 | "source": [ 301 | "%%timeit\n", 302 | "vg_numba(pvec0, pvec1)" 303 | ] 304 | }, 305 | { 306 | "cell_type": "code", 307 | "execution_count": 14, 308 | "id": "88300bed-cd64-4ce9-b17a-c255ead16385", 309 | "metadata": {}, 310 | "outputs": [ 311 | { 312 | "name": "stdout", 313 | "output_type": "stream", 314 | "text": [ 315 | "33.8 ms ± 940 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" 316 | ] 317 | } 318 | ], 319 | "source": [ 320 | "%%timeit\n", 321 | "vg_numba_parallel(pvec0, pvec1)" 322 | ] 323 | }, 324 | { 325 | "cell_type": "code", 326 | "execution_count": 15, 327 | "id": "b45c3e37-53fd-4590-9458-b7da161971c2", 328 | "metadata": {}, 329 | "outputs": [ 330 | { 331 | "data": { 332 | "text/plain": [ 333 | "0.0014142797576549417" 334 | ] 335 | }, 336 | "execution_count": 15, 337 | "metadata": {}, 338 | "output_type": "execute_result" 339 | } 340 | ], 341 | "source": [ 342 | "np.std(pvec1[:, 0])" 343 | ] 344 | }, 345 | { 346 | "cell_type": "markdown", 347 | "id": "6e6552c8-0af6-4795-a4dd-98a3ab825b6c", 348 | "metadata": {}, 349 | "source": [ 350 | "## Numba CUDA" 351 | ] 352 | }, 353 | { 354 | "cell_type": "code", 355 | "execution_count": 16, 356 | "id": "38a055f6-ce97-4c34-a11e-707023b50044", 357 | "metadata": {}, 358 | "outputs": [], 359 | "source": [ 360 | "from numba import cuda" 361 | ] 362 | }, 363 | { 364 | "cell_type": "code", 365 | "execution_count": 17, 366 | "id": "0e6df42a-9ab7-42f8-abd9-1acc30618138", 367 | "metadata": {}, 368 | "outputs": [], 369 | "source": [ 370 | "from numba import cuda, guvectorize, float64\n", 371 | "\n", 372 | "@guvectorize([(float64[:], float64[:])], '(n)->(n)', target='cuda')\n", 373 | "def crashes(a_in, a_out):\n", 374 | " pass" 375 | ] 376 | }, 377 | { 378 | "cell_type": "code", 379 | "execution_count": 18, 380 | "id": "43109d2a-d204-4fe3-ada8-3ac74dbaf8b3", 381 | "metadata": {}, 382 | "outputs": [], 383 | "source": [ 384 | "# Particle method\n", 385 | "params = D1\n", 386 | "g = cuda.jit(make_track_a_drift(math), device=True)\n", 387 | "\n", 388 | "@guvectorize([(float64[:], float64[:])], '(n)->(n)', target='cuda')\n", 389 | "def vg_parallel_cuda(a_in, a_out):\n", 390 | " p_in = Particle(x = a_in[0],\n", 391 | " px = a_in[1],\n", 392 | " y = a_in[2],\n", 393 | " py = a_in[3],\n", 394 | " z = a_in[4],\n", 395 | " pz = a_in[5],\n", 396 | " s=s, p0c=p0c, mc2=mc2)\n", 397 | " p_out = g(p_in, params)\n", 398 | " a_out[0] = p_out.x\n", 399 | " a_out[1] = p_out.px \n", 400 | " a_out[2] = p_out.y\n", 401 | " a_out[3] = p_out.py \n", 402 | " a_out[4] = p_out.z\n", 403 | " a_out[5] = p_out.pz " 404 | ] 405 | }, 406 | { 407 | "cell_type": "code", 408 | "execution_count": 19, 409 | "id": "0c204c5f-58fc-401d-a7b0-8a3e49064020", 410 | "metadata": {}, 411 | "outputs": [ 412 | { 413 | "name": "stdout", 414 | "output_type": "stream", 415 | "text": [ 416 | "98.5 ms ± 4.44 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" 417 | ] 418 | } 419 | ], 420 | "source": [ 421 | "%%timeit\n", 422 | "vg_parallel_cuda(pvec0, pvec1)" 423 | ] 424 | }, 425 | { 426 | "cell_type": "code", 427 | "execution_count": 20, 428 | "id": "de51fff9-178e-430d-9ac2-ac043743d7aa", 429 | "metadata": {}, 430 | "outputs": [ 431 | { 432 | "data": { 433 | "text/plain": [ 434 | "0.0014142797576549417" 435 | ] 436 | }, 437 | "execution_count": 20, 438 | "metadata": {}, 439 | "output_type": "execute_result" 440 | } 441 | ], 442 | "source": [ 443 | "np.std(pvec1[:,0])" 444 | ] 445 | }, 446 | { 447 | "cell_type": "markdown", 448 | "id": "6b65d4fd-3a43-4036-aaa9-afce30f470bc", 449 | "metadata": {}, 450 | "source": [ 451 | "## PyTorch" 452 | ] 453 | }, 454 | { 455 | "cell_type": "code", 456 | "execution_count": 21, 457 | "id": "fa4140a5-535f-45d8-baa3-d96c6a3553e0", 458 | "metadata": {}, 459 | "outputs": [], 460 | "source": [ 461 | "import torch" 462 | ] 463 | }, 464 | { 465 | "cell_type": "code", 466 | "execution_count": 22, 467 | "id": "d8eb557f-252d-48dc-b1a6-32667332b9a7", 468 | "metadata": {}, 469 | "outputs": [], 470 | "source": [ 471 | "tkwargs = {\n", 472 | " \"dtype\" : torch.double\n", 473 | "}" 474 | ] 475 | }, 476 | { 477 | "cell_type": "code", 478 | "execution_count": 23, 479 | "id": "9663e3b3-7885-4dc7-b7ea-12613867e981", 480 | "metadata": {}, 481 | "outputs": [], 482 | "source": [ 483 | "track_a_drift = make_track_a_drift(torch)" 484 | ] 485 | }, 486 | { 487 | "cell_type": "code", 488 | "execution_count": 24, 489 | "id": "ad7f2ea5-93cc-4685-9ce1-314866856016", 490 | "metadata": {}, 491 | "outputs": [], 492 | "source": [ 493 | "tvec0= torch.tensor(pvec0, requires_grad=True, **tkwargs)\n", 494 | "ts = torch.tensor(s, **tkwargs)\n", 495 | "tp0c = torch.tensor(p0c, **tkwargs)\n", 496 | "tmc2 = torch.tensor(mc2, **tkwargs)\n", 497 | "\n", 498 | "tparticles0 = Particle(tvec0[:,0],\n", 499 | " tvec0[:,1],\n", 500 | " tvec0[:,2],\n", 501 | " tvec0[:,3],\n", 502 | " tvec0[:,4],\n", 503 | " tvec0[:,5],\n", 504 | " ts, tp0c, tmc2)" 505 | ] 506 | }, 507 | { 508 | "cell_type": "code", 509 | "execution_count": 25, 510 | "id": "c14f9345-b3da-4049-a4db-ff0cb2eb35df", 511 | "metadata": {}, 512 | "outputs": [ 513 | { 514 | "data": { 515 | "text/plain": [ 516 | "tensor(0.0014, dtype=torch.float64, grad_fn=)" 517 | ] 518 | }, 519 | "execution_count": 25, 520 | "metadata": {}, 521 | "output_type": "execute_result" 522 | } 523 | ], 524 | "source": [ 525 | "tparticles1 = track_a_drift(tparticles0, params)\n", 526 | "\n", 527 | "tparticles1.x.std()" 528 | ] 529 | }, 530 | { 531 | "cell_type": "code", 532 | "execution_count": 26, 533 | "id": "38f70021-a074-445e-8e36-a75fc0a8a31b", 534 | "metadata": {}, 535 | "outputs": [ 536 | { 537 | "name": "stdout", 538 | "output_type": "stream", 539 | "text": [ 540 | "215 ms ± 20.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n", 541 | "Compiler time: 0.16 s\n" 542 | ] 543 | } 544 | ], 545 | "source": [ 546 | "%%timeit\n", 547 | "tparticles1 = track_a_drift(tparticles0, params)" 548 | ] 549 | }, 550 | { 551 | "cell_type": "markdown", 552 | "id": "7cd88f75-45a9-4197-a327-6f33177362a4", 553 | "metadata": {}, 554 | "source": [ 555 | "# System Information" 556 | ] 557 | }, 558 | { 559 | "cell_type": "code", 560 | "execution_count": 27, 561 | "id": "a50d989a-6534-4d33-bd46-f74bfae6d180", 562 | "metadata": {}, 563 | "outputs": [ 564 | { 565 | "name": "stdout", 566 | "output_type": "stream", 567 | "text": [ 568 | "perlmutter\n" 569 | ] 570 | } 571 | ], 572 | "source": [ 573 | "!echo $NERSC_HOST" 574 | ] 575 | }, 576 | { 577 | "cell_type": "code", 578 | "execution_count": 28, 579 | "id": "eb0a7e7e-2179-4008-baa1-3aa29c530844", 580 | "metadata": {}, 581 | "outputs": [ 582 | { 583 | "data": { 584 | "text/plain": [ 585 | "128" 586 | ] 587 | }, 588 | "execution_count": 28, 589 | "metadata": {}, 590 | "output_type": "execute_result" 591 | } 592 | ], 593 | "source": [ 594 | "numba.config.NUMBA_NUM_THREADS" 595 | ] 596 | }, 597 | { 598 | "cell_type": "code", 599 | "execution_count": 29, 600 | "id": "dc11c34c-22bb-4b1e-94d8-0477db77cb25", 601 | "metadata": {}, 602 | "outputs": [ 603 | { 604 | "name": "stdout", 605 | "output_type": "stream", 606 | "text": [ 607 | "Mon Oct 17 15:31:43 2022 \n", 608 | "+-----------------------------------------------------------------------------+\n", 609 | "| NVIDIA-SMI 515.48.07 Driver Version: 515.48.07 CUDA Version: 11.7 |\n", 610 | "|-------------------------------+----------------------+----------------------+\n", 611 | "| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |\n", 612 | "| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n", 613 | "| | | MIG M. |\n", 614 | "|===============================+======================+======================|\n", 615 | "| 0 NVIDIA A100-SXM... On | 00000000:03:00.0 Off | 0 |\n", 616 | "| N/A 28C P0 60W / 400W | 4080MiB / 40960MiB | 0% Default |\n", 617 | "| | | Disabled |\n", 618 | "+-------------------------------+----------------------+----------------------+\n", 619 | "| 1 NVIDIA A100-SXM... On | 00000000:41:00.0 Off | 0 |\n", 620 | "| N/A 26C P0 50W / 400W | 3MiB / 40960MiB | 0% Default |\n", 621 | "| | | Disabled |\n", 622 | "+-------------------------------+----------------------+----------------------+\n", 623 | "| 2 NVIDIA A100-SXM... On | 00000000:82:00.0 Off | 0 |\n", 624 | "| N/A 26C P0 50W / 400W | 3MiB / 40960MiB | 0% Default |\n", 625 | "| | | Disabled |\n", 626 | "+-------------------------------+----------------------+----------------------+\n", 627 | "| 3 NVIDIA A100-SXM... On | 00000000:C1:00.0 Off | 0 |\n", 628 | "| N/A 25C P0 50W / 400W | 3MiB / 40960MiB | 0% Default |\n", 629 | "| | | Disabled |\n", 630 | "+-------------------------------+----------------------+----------------------+\n", 631 | " \n", 632 | "+-----------------------------------------------------------------------------+\n", 633 | "| Processes: |\n", 634 | "| GPU GI CI PID Type Process name GPU Memory |\n", 635 | "| ID ID Usage |\n", 636 | "|=============================================================================|\n", 637 | "| 0 N/A N/A 101928 C ...nda/envs/bmadx/bin/python 4077MiB |\n", 638 | "+-----------------------------------------------------------------------------+\n" 639 | ] 640 | } 641 | ], 642 | "source": [ 643 | "!nvidia-smi" 644 | ] 645 | } 646 | ], 647 | "metadata": { 648 | "kernelspec": { 649 | "display_name": "Bmad-X", 650 | "language": "python", 651 | "name": "bmadx" 652 | }, 653 | "language_info": { 654 | "codemirror_mode": { 655 | "name": "ipython", 656 | "version": 3 657 | }, 658 | "file_extension": ".py", 659 | "mimetype": "text/x-python", 660 | "name": "python", 661 | "nbconvert_exporter": "python", 662 | "pygments_lexer": "ipython3", 663 | "version": "3.9.13" 664 | } 665 | }, 666 | "nbformat": 4, 667 | "nbformat_minor": 5 668 | } 669 | -------------------------------------------------------------------------------- /tests/test_bmadx_torch.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from pytao import Tao 3 | import torch 4 | from torch.autograd.functional import jacobian 5 | from torch.autograd.functional import hessian 6 | import warnings 7 | 8 | # bmadx imports: 9 | from bmadx import Particle, Drift, Quadrupole, CrabCavity, RFCavity 10 | from bmadx import track_element, track_lattice 11 | from bmadx import M_ELECTRON 12 | 13 | def set_tao(tao, coords): 14 | tao.cmd('set particle_start x='+str(coords[0])) 15 | tao.cmd('set particle_start px='+str(coords[1])) 16 | tao.cmd('set particle_start y='+str(coords[2])) 17 | tao.cmd('set particle_start py='+str(coords[3])) 18 | tao.cmd('set particle_start z='+str(coords[4])) 19 | tao.cmd('set particle_start pz='+str(coords[5])) 20 | 21 | class TestBmadxTorch: 22 | # create incoming particle 23 | coords = [2e-3,3e-3,-3e-3,-1e-3,2e-3,-2e-3] 24 | coords_t = torch.tensor(coords, dtype=torch.double) 25 | p_in = Particle(*coords_t, 26 | s = torch.tensor(0.0, dtype=torch.double), 27 | p0c = torch.tensor(4.0e7, dtype=torch.double), 28 | mc2 = torch.tensor(M_ELECTRON, dtype=torch.double)) 29 | 30 | # coordinates with autodiff: 31 | diff_coords = coords_t.clone().detach().requires_grad_(True) 32 | 33 | def test_drift(self): 34 | # test one particle tracking 35 | d = Drift(L=1.0) 36 | p_out = track_element(self.p_in, d) 37 | x_py = torch.hstack(p_out[:6]) 38 | 39 | tao = Tao('-lat tests/bmad_lattices/test_drift.bmad -noplot') 40 | set_tao(tao, self.coords) 41 | orbit_out=tao.orbit_at_s(ele=1) 42 | x_tao = torch.tensor([orbit_out['x'], 43 | orbit_out['px'], 44 | orbit_out['y'], 45 | orbit_out['py'], 46 | orbit_out['z'], 47 | orbit_out['pz']], dtype=torch.double) 48 | 49 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 50 | 51 | # test Jacobian 52 | f_drift = lambda x: track_element( 53 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 54 | d)[:6] 55 | 56 | J = jacobian(f_drift, self.diff_coords) 57 | mat_py = torch.vstack(J) 58 | 59 | drift_tao = tao.matrix(0,1) 60 | mat_tao = torch.tensor(drift_tao['mat6']) 61 | 62 | assert torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-14) 63 | 64 | def test_quadrupole(self): 65 | # no offset nor tilt 66 | # test one particle tracking 67 | q = Quadrupole(L=0.1, K1=10.0) 68 | p_out = track_element(self.p_in, q) 69 | x_py = torch.hstack(p_out[:6]) 70 | 71 | tao = Tao('-lat tests/bmad_lattices/test_quad.bmad -noplot') 72 | set_tao(tao, self.coords) 73 | orbit_out = tao.orbit_at_s(ele=1) 74 | x_tao = torch.tensor([orbit_out['x'], 75 | orbit_out['px'], 76 | orbit_out['y'], 77 | orbit_out['py'], 78 | orbit_out['z'], 79 | orbit_out['pz']], dtype=torch.double) 80 | 81 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 82 | 83 | # test Jacobian 84 | f_quadrupole = lambda x: track_element( 85 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 86 | q)[:6] 87 | 88 | J = jacobian(f_quadrupole, self.diff_coords) 89 | mat_py = torch.vstack(J) 90 | 91 | quad_tao = tao.matrix(0,1) 92 | mat_tao = torch.tensor(quad_tao['mat6']) 93 | 94 | assert torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-14) 95 | 96 | def test_quadrupole_offset(self): 97 | # x-y offset, no tilt 98 | # test one particle tracking 99 | q_off = Quadrupole(L = 0.1, 100 | K1 = 10.0, 101 | X_OFFSET = 1e-3, 102 | Y_OFFSET = -2e-3) 103 | 104 | p_out = track_element(self.p_in, q_off) 105 | x_py = torch.hstack(p_out[:6]) 106 | 107 | tao = Tao('-lat tests/bmad_lattices/test_quad_offset.bmad -noplot') 108 | set_tao(tao, self.coords) 109 | orbit_out = tao.orbit_at_s(ele=1) 110 | x_tao = torch.tensor([orbit_out['x'], 111 | orbit_out['px'], 112 | orbit_out['y'], 113 | orbit_out['py'], 114 | orbit_out['z'], 115 | orbit_out['pz']], dtype=torch.double) 116 | 117 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 118 | 119 | # test Jacobian 120 | f_quadrupole_off = lambda x: track_element( 121 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 122 | q_off)[:6] 123 | 124 | J = jacobian(f_quadrupole_off, self.diff_coords) 125 | mat_py = torch.vstack(J) 126 | 127 | quad_tao = tao.matrix(0,1) 128 | mat_tao = torch.tensor(quad_tao['mat6']) 129 | 130 | assert torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-14) 131 | 132 | def test_quadrupole_tilt(self): 133 | # tilt, no offset 134 | # test one particle tracking 135 | q_tilt = Quadrupole(L = 0.1, 136 | K1 = 10.0, 137 | TILT = 0.3) 138 | p_out = track_element(self.p_in, q_tilt) 139 | x_py = torch.hstack(p_out[:6]) 140 | 141 | tao = Tao('-lat tests/bmad_lattices/test_quad_tilt.bmad -noplot') 142 | set_tao(tao, self.coords) 143 | orbit_out = tao.orbit_at_s(ele=1) 144 | x_tao = torch.tensor([orbit_out['x'], 145 | orbit_out['px'], 146 | orbit_out['y'], 147 | orbit_out['py'], 148 | orbit_out['z'], 149 | orbit_out['pz']], dtype=torch.double) 150 | 151 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 152 | 153 | # test Jacobian 154 | f_quadrupole = lambda x: track_element( 155 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 156 | q_tilt)[:6] 157 | 158 | J = jacobian(f_quadrupole, self.diff_coords) 159 | mat_py = torch.vstack(J) 160 | 161 | quad_tao = tao.matrix(0,1) 162 | mat_tao = torch.tensor(quad_tao['mat6']) 163 | 164 | assert torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-14) 165 | 166 | def test_crab_cavity(self): 167 | # no offset nor tilt 168 | # test one particle tracking 169 | cav = CrabCavity(L=0.2, 170 | VOLTAGE=1e4, 171 | PHI0=0.5, 172 | RF_FREQUENCY=1e9) 173 | 174 | p_out = track_element(self.p_in, cav) 175 | x_py = torch.hstack(p_out[:6]) 176 | 177 | tao = Tao('-lat tests/bmad_lattices/test_crab_cavity.bmad -noplot') 178 | set_tao(tao, self.coords) 179 | orbit_out = tao.orbit_at_s(ele=1) 180 | x_tao = torch.tensor([orbit_out['x'], 181 | orbit_out['px'], 182 | orbit_out['y'], 183 | orbit_out['py'], 184 | orbit_out['z'], 185 | orbit_out['pz']], dtype=torch.double) 186 | 187 | #assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 188 | 189 | if torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) == False: 190 | warnings.warn('tracking not the same as Bmad to double precission') 191 | # At least should be the same up to single precission 192 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-7) 193 | 194 | # test Jacobian 195 | f_cav = lambda x: track_element( 196 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 197 | cav)[:6] 198 | 199 | J = jacobian(f_cav, self.diff_coords) 200 | mat_py = torch.vstack(J) 201 | 202 | cav_tao = tao.matrix(0,1) 203 | mat_tao = torch.tensor(cav_tao['mat6']) 204 | 205 | if torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-13) == False: 206 | warnings.warn('Jacobian not the same as Bmad to double precission') 207 | # At least should be the same up to single precission 208 | if torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-7) == False: 209 | warnings.warn('Jacobian not the same as Bmad to single precission') 210 | 211 | def test_crab_cavity_offset(self): 212 | # x-y offset, no tilt 213 | # test one particle tracking 214 | cav_off = CrabCavity(L=0.2, 215 | VOLTAGE=1e4, 216 | PHI0=0.5, 217 | RF_FREQUENCY=1e9, 218 | X_OFFSET=1e-3, 219 | Y_OFFSET=-2e-3) 220 | 221 | p_out = track_element(self.p_in, cav_off) 222 | x_py = torch.hstack(p_out[:6]) 223 | 224 | tao = Tao( 225 | '-lat tests/bmad_lattices/test_crab_cavity_offset.bmad -noplot') 226 | 227 | set_tao(tao, self.coords) 228 | orbit_out = tao.orbit_at_s(ele=1) 229 | x_tao = torch.tensor([orbit_out['x'], 230 | orbit_out['px'], 231 | orbit_out['y'], 232 | orbit_out['py'], 233 | orbit_out['z'], 234 | orbit_out['pz']], dtype=torch.double) 235 | 236 | # assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 237 | 238 | if torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) == False: 239 | warnings.warn('tracking not the same as Bmad to double precission') 240 | # At least should be the same up to single precission 241 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-7) 242 | 243 | # test Jacobian 244 | f_cav = lambda x: track_element( 245 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 246 | cav_off)[:6] 247 | 248 | J = jacobian(f_cav, self.diff_coords) 249 | mat_py = torch.vstack(J) 250 | 251 | cav_tao = tao.matrix(0,1) 252 | mat_tao = torch.tensor(cav_tao['mat6']) 253 | 254 | if torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-13) == False: 255 | warnings.warn('Jacobian not the same as Bmad to double precission') 256 | # At least should be the same up to single precission 257 | if torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-7) == False: 258 | warnings.warn('Jacobian not the same as Bmad to single precission') 259 | 260 | def test_crab_cavity_tilt(self): 261 | # tilt, no offset 262 | # test one particle tracking 263 | cav_tilt = CrabCavity(L=0.2, 264 | VOLTAGE=1e4, 265 | PHI0=0.5, 266 | RF_FREQUENCY=1e9, 267 | TILT=0.3) 268 | 269 | p_out = track_element(self.p_in, cav_tilt) 270 | x_py = torch.hstack(p_out[:6]) 271 | 272 | tao = Tao( 273 | '-lat tests/bmad_lattices/test_crab_cavity_tilt.bmad -noplot') 274 | 275 | set_tao(tao, self.coords) 276 | orbit_out = tao.orbit_at_s(ele=1) 277 | x_tao = torch.tensor([orbit_out['x'], 278 | orbit_out['px'], 279 | orbit_out['y'], 280 | orbit_out['py'], 281 | orbit_out['z'], 282 | orbit_out['pz']], dtype=torch.double) 283 | 284 | # assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 285 | 286 | if torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) == False: 287 | warnings.warn('tracking not the same as Bmad to double precission') 288 | # At least should be the same up to single precission 289 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-7) 290 | 291 | # test Jacobian 292 | f_cav = lambda x: track_element( 293 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 294 | cav_tilt)[:6] 295 | 296 | J = jacobian(f_cav, self.diff_coords) 297 | mat_py = torch.vstack(J) 298 | 299 | cav_tao = tao.matrix(0,1) 300 | mat_tao = torch.tensor(cav_tao['mat6']) 301 | 302 | if torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-13) == False: 303 | warnings.warn('Jacobian not the same as Bmad to double precission') 304 | # At least should be the same up to single precission 305 | if torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-7) == False: 306 | warnings.warn('Jacobian not the same as Bmad to single precission') 307 | 308 | def test_rf_cavity(self): 309 | # no offset nor tilt 310 | # test one particle tracking 311 | cav = RFCavity(L=0.2, 312 | VOLTAGE=1e4, 313 | PHI0=0.5, 314 | RF_FREQUENCY=1e9) 315 | 316 | p_out = track_element(self.p_in, cav) 317 | x_py = torch.hstack(p_out[:6]) 318 | 319 | tao = Tao('-lat tests/bmad_lattices/test_rf_cavity.bmad -noplot') 320 | set_tao(tao, self.coords) 321 | orbit_out = tao.orbit_at_s(ele=1) 322 | x_tao = torch.tensor([orbit_out['x'], 323 | orbit_out['px'], 324 | orbit_out['y'], 325 | orbit_out['py'], 326 | orbit_out['z'], 327 | orbit_out['pz']], dtype=torch.double) 328 | 329 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 330 | 331 | # test Jacobian 332 | f_cav = lambda x: track_element( 333 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 334 | cav)[:6] 335 | 336 | J = jacobian(f_cav, self.diff_coords) 337 | mat_py = torch.vstack(J) 338 | 339 | cav_tao = tao.matrix(0,1) 340 | mat_tao = torch.tensor(cav_tao['mat6']) 341 | 342 | assert torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-13) 343 | 344 | def test_rf_cavity_offset(self): 345 | # x-y offset, no tilt 346 | # test one particle tracking 347 | cav_off = RFCavity(L=0.2, 348 | VOLTAGE = 1e4, 349 | PHI0 = 0.5, 350 | RF_FREQUENCY = 1e9, 351 | X_OFFSET = 1e-3, 352 | Y_OFFSET = -2e-3) 353 | 354 | p_out = track_element(self.p_in, cav_off) 355 | x_py = torch.hstack(p_out[:6]) 356 | 357 | tao = Tao( 358 | '-lat tests/bmad_lattices/test_rf_cavity_offset.bmad -noplot') 359 | 360 | set_tao(tao, self.coords) 361 | orbit_out = tao.orbit_at_s(ele=1) 362 | x_tao = torch.tensor([orbit_out['x'], 363 | orbit_out['px'], 364 | orbit_out['y'], 365 | orbit_out['py'], 366 | orbit_out['z'], 367 | orbit_out['pz']], dtype=torch.double) 368 | 369 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 370 | 371 | # test Jacobian 372 | f_cav = lambda x: track_element( 373 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 374 | cav_off)[:6] 375 | 376 | J = jacobian(f_cav, self.diff_coords) 377 | mat_py = torch.vstack(J) 378 | 379 | cav_tao = tao.matrix(0,1) 380 | mat_tao = torch.tensor(cav_tao['mat6']) 381 | 382 | assert torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-13) 383 | 384 | def test_rf_cavity_tilt(self): 385 | # tilt, no offset 386 | # test one particle tracking 387 | cav_tilt = RFCavity(L=0.2, 388 | VOLTAGE=1e4, 389 | PHI0=0.5, 390 | RF_FREQUENCY=1e9, 391 | TILT=0.3) 392 | 393 | p_out = track_element(self.p_in, cav_tilt) 394 | x_py = torch.hstack(p_out[:6]) 395 | 396 | tao = Tao( 397 | '-lat tests/bmad_lattices/test_rf_cavity_tilt.bmad -noplot') 398 | 399 | set_tao(tao, self.coords) 400 | orbit_out = tao.orbit_at_s(ele=1) 401 | x_tao = torch.tensor([orbit_out['x'], 402 | orbit_out['px'], 403 | orbit_out['y'], 404 | orbit_out['py'], 405 | orbit_out['z'], 406 | orbit_out['pz']], dtype=torch.double) 407 | 408 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 409 | 410 | # test Jacobian 411 | f_cav = lambda x: track_element( 412 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 413 | cav_tilt)[:6] 414 | 415 | J = jacobian(f_cav, self.diff_coords) 416 | mat_py = torch.vstack(J) 417 | 418 | cav_tao = tao.matrix(0,1) 419 | mat_tao = torch.tensor(cav_tao['mat6']) 420 | 421 | if torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-14) == False: 422 | warnings.warn('Jacobian not the same as Bmad to double precission') 423 | # At least should be the same up to single precission 424 | assert torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-8) 425 | 426 | 427 | def test_lattice(self): 428 | # create lattice 429 | d = Drift(L=1.0) 430 | q = Quadrupole(L=0.1, K1=10) 431 | lattice = [d, q, d, q, d] 432 | 433 | # outgoing particle: 434 | p_out = track_lattice(self.p_in, lattice) 435 | x_py = torch.hstack(p_out[:6]) 436 | 437 | # Bmad lattice to compare 438 | tao = Tao('-lat tests/bmad_lattices/test_drift_quad.bmad -noplot') 439 | set_tao(tao, self.coords) 440 | orbit_out = tao.orbit_at_s(ele=5) 441 | x_tao = torch.tensor([orbit_out['x'], 442 | orbit_out['px'], 443 | orbit_out['y'], 444 | orbit_out['py'], 445 | orbit_out['z'], 446 | orbit_out['pz']], dtype=torch.double) 447 | 448 | assert torch.allclose(x_py, x_tao, atol=0, rtol=1.0e-14) 449 | 450 | # test Jacobian 451 | f_driftquadrupole = lambda x: track_lattice( 452 | Particle(*x, s=self.p_in.s, p0c=self.p_in.p0c, mc2=self.p_in.mc2), 453 | lattice)[:6] 454 | 455 | J = jacobian(f_driftquadrupole, self.diff_coords) 456 | mat_py = torch.vstack(J) 457 | 458 | lat_tao = tao.matrix(0,5) 459 | mat_tao = torch.tensor(lat_tao['mat6']) 460 | 461 | assert torch.allclose(mat_py, mat_tao, atol=0, rtol=1.0e-14) 462 | 463 | def test_hessian(self): 464 | # Particle bunch with Gaussian distribution 465 | sample_size = 1000 466 | mean = torch.zeros(6) 467 | cov = torch.diag(torch.tensor([1e-6, 2e-6, 1e-6, 2e-6, 1e-6, 2e-6])) 468 | dist = torch.distributions.multivariate_normal.MultivariateNormal( 469 | mean, cov) 470 | sample = dist.sample(torch.Size([sample_size])) 471 | p_in = Particle(*sample.T, 472 | s = torch.tensor(0.0, dtype=torch.double), 473 | p0c = torch.tensor(4.0e7, dtype=torch.double), 474 | mc2 = torch.tensor(M_ELECTRON, dtype=torch.double)) 475 | 476 | drift = Drift(L=1.0) 477 | L_q = 0.1 # Quad length 478 | 479 | def sigmax_end(k1s): 480 | """returns x beamsize after lattice composed by len(k1s)+1 481 | drifts with len(k1s) quadrupoles in between. 482 | """ 483 | lattice = [drift] 484 | 485 | for k1 in k1s: 486 | lattice.append(Quadrupole(L=L_q, K1=k1)) 487 | lattice.append(drift) 488 | 489 | p_out = track_lattice(p_in, lattice) 490 | return torch.std(p_out.x) 491 | 492 | k1s = torch.zeros(5) 493 | hessian_py = hessian(sigmax_end, k1s) 494 | isnan = torch.isnan(hessian_py) 495 | assert torch.any(~isnan) -------------------------------------------------------------------------------- /dev/speed.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "id": "b0e62b30-8b50-49b9-8c40-4dbef6535d24", 7 | "metadata": {}, 8 | "outputs": [], 9 | "source": [ 10 | "import torch\n", 11 | "import numpy as np\n", 12 | "\n", 13 | "from bmadx import Particle, Drift, Quadrupole\n", 14 | "from bmadx import track_element, track_lattice, track_lattice_save_stats, track_lattice_save_particles\n", 15 | "from bmadx import M_ELECTRON\n", 16 | "\n", 17 | "from scalene import scalene_profiler" 18 | ] 19 | }, 20 | { 21 | "cell_type": "code", 22 | "execution_count": 2, 23 | "id": "612bf0e1-6a8b-41e4-8ab6-ec208582c9c3", 24 | "metadata": {}, 25 | "outputs": [], 26 | "source": [ 27 | "# Initial beam distribution\n", 28 | "\n", 29 | "mean = torch.zeros(6)\n", 30 | "cov = torch.diag(torch.tensor([1e-6, 2e-6, 3e-6, 2e-6, 1e-6, 2e-6]))\n", 31 | "torch.manual_seed(0)\n", 32 | "dist = torch.distributions.multivariate_normal.MultivariateNormal(mean, cov)\n", 33 | "\n", 34 | "def create_beam(n_particles):\n", 35 | " coords = dist.sample(torch.Size([n_particles])) # particles' coordinates\n", 36 | "\n", 37 | " return Particle(*coords.T,\n", 38 | " s=torch.tensor(0.),\n", 39 | " p0c=torch.tensor(4e7),\n", 40 | " mc2=M_ELECTRON)" 41 | ] 42 | }, 43 | { 44 | "cell_type": "code", 45 | "execution_count": 3, 46 | "id": "c3bf23c5-94b6-4259-9bd2-114694ee1cdf", 47 | "metadata": {}, 48 | "outputs": [], 49 | "source": [ 50 | "def create_fodo(n_cells, l_d, l_q, k):\n", 51 | " q_f = Quadrupole(L = l_q, K1 = k)\n", 52 | " q_d = Quadrupole(L = l_q, K1 = -k)\n", 53 | " d = Drift(L = l_d)\n", 54 | " half_d = Drift(L = l_d/2)\n", 55 | " lat = []\n", 56 | " for i in range(n_cells):\n", 57 | " lat.extend([half_d, q_f, d, half_d, q_d])\n", 58 | " return lat" 59 | ] 60 | }, 61 | { 62 | "cell_type": "code", 63 | "execution_count": 4, 64 | "id": "96418b4e-6534-4779-a884-263ae915039a", 65 | "metadata": {}, 66 | "outputs": [], 67 | "source": [ 68 | "lat = create_fodo(10, 0.9, 0.1, 10)" 69 | ] 70 | }, 71 | { 72 | "cell_type": "code", 73 | "execution_count": 5, 74 | "id": "4530aad9-adc1-460f-ac4a-9914fac42681", 75 | "metadata": {}, 76 | "outputs": [ 77 | { 78 | "data": { 79 | "text/plain": [ 80 | "array([ 1, 10, 100, 1000, 10000, 100000])" 81 | ] 82 | }, 83 | "execution_count": 5, 84 | "metadata": {}, 85 | "output_type": "execute_result" 86 | } 87 | ], 88 | "source": [ 89 | "N = 5\n", 90 | "n_particles = np.logspace(0, N, N+1, dtype=int)\n", 91 | "n_particles" 92 | ] 93 | }, 94 | { 95 | "cell_type": "code", 96 | "execution_count": 6, 97 | "id": "47349261-98ba-41df-a62b-72b3d322d531", 98 | "metadata": {}, 99 | "outputs": [ 100 | { 101 | "name": "stdout", 102 | "output_type": "stream", 103 | "text": [ 104 | "1\n", 105 | "10\n", 106 | "100\n", 107 | "1000\n", 108 | "10000\n", 109 | "100000\n" 110 | ] 111 | } 112 | ], 113 | "source": [ 114 | "from time import perf_counter\n", 115 | "\n", 116 | "time_counter = perf_counter\n", 117 | "t_no_save = []\n", 118 | "t_save_par = []\n", 119 | "t_save_stats = []\n", 120 | "for n_par in n_particles:\n", 121 | " print(n_par)\n", 122 | " \n", 123 | " beam = create_beam(n_par)\n", 124 | " \n", 125 | " p_out = None\n", 126 | " t_0_no_save = time_counter()\n", 127 | " p_out = track_lattice(beam, lat)\n", 128 | " t_f_no_save = time_counter()\n", 129 | " t_no_save.extend([t_f_no_save - t_0_no_save])\n", 130 | " \n", 131 | " p_out = None\n", 132 | " t_0_save_par = time_counter()\n", 133 | " p_out = track_lattice_save_particles(beam, lat)\n", 134 | " t_f_save_par = time_counter()\n", 135 | " t_save_par.extend([t_f_save_par - t_0_save_par])\n", 136 | " \n", 137 | " p_out = None\n", 138 | " t_0_save_stats = time_counter()\n", 139 | " p_out = track_lattice_save_stats(beam, lat)\n", 140 | " t_f_save_stats = time_counter()\n", 141 | " t_save_stats.extend([t_f_save_stats - t_0_save_stats])" 142 | ] 143 | }, 144 | { 145 | "cell_type": "code", 146 | "execution_count": 7, 147 | "id": "b9af296b-1a33-46e0-89b9-093f903fe526", 148 | "metadata": {}, 149 | "outputs": [ 150 | { 151 | "data": { 152 | "text/plain": [ 153 | "" 154 | ] 155 | }, 156 | "execution_count": 7, 157 | "metadata": {}, 158 | "output_type": "execute_result" 159 | }, 160 | { 161 | "data": { 162 | "image/png": "iVBORw0KGgoAAAANSUhEUgAAAkIAAAG1CAYAAAAV2Js8AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8o6BhiAAAACXBIWXMAAA9hAAAPYQGoP6dpAAA+0ElEQVR4nO3de1xVdb7/8fcCAZHLVgQFlZuKJl5TaUaz1PJepmlq5nhpqhm7/BxrypPd1Hqc8WQ3L2FmzVQ6dvJeTTnHLPFSTGoaZdGUGgQpRN4AEUXZ6/cH454hEAH33gtYr+fjwUP3uu3P/rZjvf2u71pfwzRNUwAAADbkY3UBAAAAViEIAQAA2yIIAQAA2yIIAQAA2yIIAQAA2yIIAQAA2yIIAQAA22pkdQF1mdPp1JEjRxQSEiLDMKwuBwAAVINpmiosLFSrVq3k41N1nw9BqApHjhxRdHS01WUAAIBayM7OVps2barchiBUhZCQEEllDRkaGmpxNQAAoDoKCgoUHR3tOo9XhSBUhQuXw0JDQwlCAADUM9UZ1sJgaQAAYFsEIQAAYFsEIQAAYFuMEXKD0tJSnTt3zuoyUA/5+fnJ19fX6jIAwLYIQpfBNE3l5ubq5MmTVpeCeqxp06aKjIzkWVUAYAGC0GW4EIJatGihJk2acCJDjZimqdOnTysvL0+SFBUVZXFFAGA/BKFaKi0tdYWg5s2bW10O6qnAwEBJUl5enlq0aMFlMgDwMgZL19KFMUFNmjSxuBLUdxe+Q4wzAwDvIwhVIjk5WYmJiUpKSrrktlwOw+XiOwQA1iEIVeLee+9Venq69uzZY3UpAADAgwhCAADAtghCAADAErlFudqds1u5RbmW1cBdYwAAwOs2HNigef+YJ6fplI/hozl95mhMwhiv10GPUB2Rk1+s1ENHlZNfbHUpAAB4VG5RrisESZLTdGreP+ZZ0jNEEKoDVu/J0tX/s1W3vbJLV//PVq3ek+XR9xswYIBmzJihWbNmKSwsTJGRkZo7d265bbKysjRq1CgFBwcrNDRU48eP108//XTRY5aUlOi+++5TVFSUGjdurLi4OM2fP9+1/vnnn1fXrl0VFBSk6Oho3XPPPTp16pQkKT8/X4GBgfq///u/csfcsGGDgoKCXNsdPnxYEyZMULNmzdS8eXONGjVKmZmZ7mkUAIDXZBVkuULQBU7TqezCbK/XQhCyWE5+sWZv2C+nWfbaaUqPbPjK4z1Db7zxhoKCgrRr1y4tWLBATz75pLZs2SKp7InHo0eP1vHjx7V9+3Zt2bJFhw4d0oQJEy56vMWLF+vdd9/VmjVr9O233+qvf/2r4uLiXOt9fHy0ePFiffXVV3rjjTe0detWzZo1S5LkcDh0ww03aNWqVeWO+eabb7rC2OnTpzVw4EAFBwdrx44d+vjjjxUcHKxhw4appKTE/Q0EAPCYmNAY+RjlI4iP4aPokGiv18IYIYtlHC1yhaALSk1TmUdPK8oR6LH37datm+bMmSNJSkhI0IsvvqiPPvpIgwcP1ocffqgvv/xSGRkZio4u+1KuXLlSnTt31p49eyp9vlJWVpYSEhLUr18/GYah2NjYcutnzpzp+nt8fLyeeuop3X333Vq6dKkkadKkSZoyZYpOnz6tJk2aqKCgQO+//77Wr18vSXrrrbfk4+OjV1991fXcnddee01NmzbVtm3bNGTIELe3EQDAMyKDIjWnz5wKY4QigyK9XgtByGLx4UHyMVQuDPkahuLCPfvE6m7dupV7HRUV5Zrz6ptvvlF0dLQrBElSYmKimjZtqm+++abSIDRt2jQNHjxYHTt21LBhw3TjjTeWCycpKSn605/+pPT0dBUUFOj8+fM6c+aMioqKFBQUpBtuuEGNGjXSu+++q1tvvVXr169XSEiI6xh79+7VwYMHFRISUu59z5w5o0OHDrmtXQAA3jEmYYz6tuqr7MJsRYdEWxKCJC6NWS7KEaj5Y7rK91+9HL6GoT+N6eLR3iBJ8vPzK/faMAw5nWXXa03TrPRpxxdbLkk9e/ZURkaGnnrqKRUXF2v8+PG65ZZbJEk//PCDRowYoS5dumj9+vXau3evkpOTJf17Wgl/f3/dcsstevPNNyWVXRabMGGCGjUqy+pOp1O9evVSWlpauZ/vvvtOt912mxtaBADgbZFBkUqKTLIsBEn0CNUJE5JidG2HCGUePa248CYeD0GXkpiYqKysLGVnZ7t6hdLT05Wfn69OnTpddL/Q0FBNmDBBEyZM0C233KJhw4bp+PHj+uyzz3T+/Hk999xz8vEpy95r1qypsP+kSZM0ZMgQff3110pJSdFTTz3lWtezZ0+tXr1aLVq0UGhoqJs/MQDArugRqiOiHIHq06655SFIkgYNGqRu3bpp0qRJ2rdvn3bv3q0pU6aof//+6t27d6X7vPDCC3rrrbf0z3/+U999953Wrl2ryMhINW3aVO3atdP58+e1ZMkSff/991q5cqWWLVtW4Rj9+/dXy5YtNWnSJMXFxenXv/61a92kSZMUHh6uUaNGaefOncrIyND27dv1hz/8QT/++KPH2gIA0LARhFCBYRh6++231axZM1177bUaNGiQ2rZtq9WrV190n+DgYD399NPq3bu3kpKSlJmZqU2bNsnHx0c9evTQ888/r6efflpdunTRqlWryt1a/5/vO3HiRH3xxReaNGlSuXVNmjTRjh07FBMTozFjxqhTp0767W9/q+LiYnqIAAC1ZpimaV56M3sqKCiQw+FQfn5+hZPtmTNnlJGRofj4eDVu3NiiCtEQ8F0CAPeq6vz9S/QIAQAA2yIIAQAA2yIIAQAA2yIIAQAA2yIIAQAA2yIIAQAA2yIIAQAA2yIIAQAA2yIIwZYyMzNlGIbS0tKqtf20adM0evRoj9YEAPA+Jl1Fgzdt2jSdPHlSb7/9tmtZdHS0cnJyFB4ebl1hAADLEYTQYJWWlsowjErX+fr6KjIy0ssVAQDqGi6N1RX5h6WMHWV/eti6devUtWtXBQYGqnnz5ho0aJCKiookSXv27NHgwYMVHh4uh8Oh/v37a9++fa59J06cqFtvvbXc8c6dO6fw8HC99tprkiTTNLVgwQK1bdtWgYGB6t69u9atW1dlTXFxcXrqqad02223KTg4WK1atdKSJUvKbfP888+ra9euCgoKUnR0tO655x6dOnXKtf71119X06ZN9d577ykxMVEBAQG6/fbb9cYbb+idd96RYRgyDEPbtm2r9NLY119/rRtuuEGhoaEKCQnRNddco0OHDlVa76U+44kTJzRp0iRFREQoMDBQCQkJrvYBANQd9AjVBftWSH/7g2Q6JcNHGrlI6jnFI2+Vk5OjiRMnasGCBbr55ptVWFionTt36sLcu4WFhZo6daoWL14sSXruuec0YsQIHThwQCEhIZo0aZLGjx+vU6dOKTg4WJK0efNmFRUVaezYsZKkxx57TBs2bNBLL72khIQE7dixQ7/5zW8UERGh/v37X7S2Z555Ro888ojmzp2rzZs36/7779cVV1yhwYMHS5J8fHy0ePFixcXFKSMjQ/fcc49mzZqlpUuXuo5x+vRpzZ8/X6+++qqaN2+uyMhInTlzRgUFBa4gEhYWpiNHjpR778OHD+vaa6/VgAEDtHXrVoWGhuqTTz7R+fPnK631Up/x8ccfV3p6uv7+978rPDxcBw8eVHFxcW3+kwEAPMnEReXn55uSzPz8/ArriouLzfT0dLO4uPjy3uTkj6Y5t6lpzgn998/cZmXLPWDv3r2mJDMzM7Na258/f94MCQkx//a3v5mmaZolJSVmeHi4uWLFCtc2EydONMeNG2eapmmeOnXKbNy4sZmamlruOHfccYc5ceLEi75PbGysOWzYsHLLJkyYYA4fPvyi+6xZs8Zs3ry56/Vrr71mSjLT0tLKbTd16lRz1KhR5ZZlZGSYkszPP//cNE3TnD17thkfH2+WlJRU+l7/eYzqfMaRI0eat99++0Vr/09u+y4BAEzTrPr8/UtcGrPa8UNlPUH/ySyVjn/vkbfr3r27rr/+enXt2lXjxo3TK6+8ohMnTrjW5+Xlafr06erQoYMcDoccDodOnTqlrKwsSZKfn5/GjRunVatWSZKKior0zjvvaNKkSZKk9PR0nTlzRoMHD1ZwcLDrZ8WKFRe9zHRBnz59Krz+5ptvXK9TUlI0ePBgtW7dWiEhIZoyZYqOHTvmuqwnSf7+/urWrVuN2yUtLU3XXHON/Pz8LrltdT7j3Xffrbfeeks9evTQrFmzlJqaWuOaAACex6Uxq4W1K7sc9p9hyPCVwtp65O18fX21ZcsWpaam6oMPPtCSJUv06KOPateuXYqPj9e0adP0888/a+HChYqNjVVAQID69OmjkpIS1zEmTZqk/v37Ky8vT1u2bFHjxo01fPhwSZLTWfY53n//fbVu3brcewcEBNS43guDnX/44QeNGDFC06dP11NPPaWwsDB9/PHHuuOOO3Tu3DnX9oGBgRcdIF2VwMDAam9bnc84fPhw/fDDD3r//ff14Ycf6vrrr9e9996rZ599tsa1AQA8hx4hqzlal40JMnzLXhu+0siFZcs9xDAMXX311Zo3b54+//xz+fv7a+PGjZKknTt3asaMGRoxYoQ6d+6sgIAAHT16tNz+ffv2VXR0tFavXq1Vq1Zp3Lhx8vf3lyTXIOWsrCy1b9++3E90dHSVdX366acVXl9xxRWSpM8++0znz5/Xc889p1//+tfq0KFDhXE+F+Pv76/S0tIqt+nWrZt27txZLlRdTHU/Y0REhKZNm6a//vWvWrhwoZYvX16tegEA3kOPUF3Qc4rU7vqyy2FhbT0agnbt2qWPPvpIQ4YMUYsWLbRr1y79/PPP6tSpkySpffv2WrlypXr37q2CggI99NBDFXpLDMPQbbfdpmXLlum7775TSkqKa11ISIgefPBB3X///XI6nerXr58KCgqUmpqq4OBgTZ069aK1ffLJJ1qwYIFGjx6tLVu2aO3atXr//fclSe3atdP58+e1ZMkSjRw5Up988omWLVtWrc8cFxenzZs369tvv1Xz5s3lcDgqbHPfffdpyZIluvXWWzV79mw5HA59+umnuuqqq9SxY8dy21bnMz7xxBPq1auXOnfurLNnz+q9995ztTEAoA7xwpilessrg6W9LD093Rw6dKgZERFhBgQEmB06dDCXLFniWr9v3z6zd+/eZkBAgJmQkGCuXbvWjI2NNV944YVyx/n6669NSWZsbKzpdDrLrXM6neaiRYvMjh07mn5+fmZERIQ5dOhQc/v27RetKzY21pw3b545fvx4s0mTJmbLli3NhQsXltvm+eefN6OioszAwEBz6NCh5ooVK0xJ5okTJ0zTLBss7XA4Khw7Ly/PHDx4sBkcHGxKMlNSUioMljZN0/ziiy/MIUOGmE2aNDFDQkLMa665xjx06JBpmhUHXF/qMz711FNmp06dzMDAQDMsLMwcNWqU+f3331f62evrdwkA6qqaDJY2TPNf902jgoKCAjkcDuXn5ys0NLTcujNnzigjI0Px8fFq3LixRRU2HHFxcZo5c6ZmzpxpdSlex3cJANyrqvP3LzFGCAAA2BZBCAAAWCInv1iph44qJ9+6B84yWBp1QmZmptUlAAC8aPWeLM3esF9OU/IxpPljumpCUozX66BHCAAAeFVOfrErBEmS05Qe2fCVJT1DBCEAAOBVGUeLXCHoglLTVObR016vhSAEAAC8Kj48SD6/mATA1zAUF97E67UQhAAAgFdFOQI1f0xX+f5rSiRfw9CfxnRRlKP60x25C4OlAQCA101IitG1HSKUefS04sKbWBKCJIIQAACwSJQj0LIAdAGXxgAAgDXyD0sZO8r+tAhBCHXegAEDajX1xrRp0zR69Gi31wMAcIN9K6SFXaQ3Rpb9uW+FJWUQhAAAgHflH5b+9gfJdJa9Np3S32Za0jNEEKojcotytTtnt3KLcj3+XuvWrVPXrl0VGBio5s2ba9CgQSoqKpIk7dmzR4MHD1Z4eLgcDof69++vffv2ufadOHGibr311nLHO3funMLDw/Xaa69JkkzT1IIFC9S2bVsFBgaqe/fuWrduXZU1LV26VAkJCWrcuLFatmypW265RVJZr8727du1aNEiGYYhwzCUmZmp0tJS3XHHHYqPj1dgYKA6duyoRYsWuY43d+5cvfHGG3rnnXdc+23btk0lJSW67777FBUVpcaNGysuLk7z5893S7sCAKrp+KF/h6ALzFLp+PdeL4XB0nXAhgMbNO8f8+Q0nfIxfDSnzxyNSRjjkffKycnRxIkTtWDBAt18880qLCzUzp07ZZplT7YqLCzU1KlTtXjxYknSc889pxEjRujAgQMKCQnRpEmTNH78eJ06dUrBwcGSpM2bN6uoqEhjx46VJD322GPasGGDXnrpJSUkJGjHjh36zW9+o4iICPXv379CTZ999plmzJihlStXqm/fvjp+/Lh27twpSVq0aJG+++47denSRU8++aQkKSIiQk6nU23atNGaNWsUHh6u1NRU/e53v1NUVJTGjx+vBx98UN98840KCgpcAS0sLEyLFy/Wu+++qzVr1igmJkbZ2dnKzs72SFsDAC4irJ1k+JQPQ4avFNbW66UQhCyWW5TrCkGS5DSdmvePeerbqq8igyLd/n45OTk6f/68xowZo9jYWElS165dXeuvu+66ctu//PLLatasmbZv364bb7xRQ4cOVVBQkDZu3KjJkydLkt58802NHDlSoaGhKioq0vPPP6+tW7eqT58+kqS2bdvq448/1ssvv1xpEMrKylJQUJBuvPFGhYSEKDY2VldeeaUkyeFwyN/fX02aNFFk5L/bw9fXV/PmzXO9jo+PV2pqqtasWaPx48crODhYgYGBOnv2bLn9srKylJCQoH79+skwDFcbAAC8yNFaGrmo7HKYWVoWgkYuLFvuZVwas1hWQZYrBF3gNJ3KLvRML0X37t11/fXXq2vXrho3bpxeeeUVnThxwrU+Ly9P06dPV4cOHeRwOORwOHTq1CllZWVJkvz8/DRu3DitWrVKklRUVKR33nlHkyZNkiSlp6frzJkzGjx4sIKDg10/K1as0KFDhyqtafDgwYqNjVXbtm01efJkrVq1SqdPX/ox68uWLVPv3r0VERGh4OBgvfLKK646L2batGlKS0tTx44dNWPGDH3wwQfVajcAgJv1nCLN3C9Nfa/sz55TLCmDIGSxmNAY+Rjl/zP4GD6KDon2yPv5+vpqy5Yt+vvf/67ExEQtWbJEHTt2VEZGhqSyoLB3714tXLhQqampSktLU/PmzVVSUuI6xqRJk/Thhx8qLy9Pb7/9tho3bqzhw4dLkpzOslD3/vvvKy0tzfWTnp5+0XFCISEh2rdvn/73f/9XUVFReuKJJ9S9e3edPHnyop9jzZo1uv/++/Xb3/5WH3zwgdLS0nT77beXq7MyPXv2VEZGhp566ikVFxdr/PjxrvFIAAAvc7SW4q+xpCfoAoKQxSKDIjWnzxxXGLowRsgTl8UuMAxDV199tebNm6fPP/9c/v7+2rhxoyRp586dmjFjhkaMGKHOnTsrICBAR48eLbd/3759FR0drdWrV2vVqlUaN26c/P39JUmJiYkKCAhQVlaW2rdvX+4nOvri4a5Ro0YaNGiQFixYoC+//FKZmZnaunWrJMnf31+lpaXltt+5c6f69u2re+65R1deeaXat29focepsv0kKTQ0VBMmTNArr7yi1atXa/369Tp+/HjNGxIAUO8xRqgOGJMwRn1b9VV2YbaiQ6I9GoJ27dqljz76SEOGDFGLFi20a9cu/fzzz+rUqZMkqX379lq5cqV69+6tgoICPfTQQwoMLP/UT8MwdNttt2nZsmX67rvvlJKS4loXEhKiBx98UPfff7+cTqf69eungoICpaamKjg4WFOnTq1Q03vvvafvv/9e1157rZo1a6ZNmzbJ6XSqY8eOkqS4uDjt2rVLmZmZCg4OVlhYmNq3b68VK1Zo8+bNio+P18qVK7Vnzx7Fx8e7jhsXF6fNmzfr22+/VfPmzeVwOPTiiy8qKipKPXr0kI+Pj9auXavIyEg1bdrUA60NAKjzTFxUfn6+KcnMz8+vsK64uNhMT083i4uLLais9tLT082hQ4eaERERZkBAgNmhQwdzyZIlrvX79u0ze/fubQYEBJgJCQnm2rVrzdjYWPOFF14od5yvv/7alGTGxsaaTqez3Dqn02kuWrTI7Nixo+nn52dGRESYQ4cONbdv315pTTt37jT79+9vNmvWzAwMDDS7detmrl692rX+22+/NX/961+bgYGBpiQzIyPDPHPmjDlt2jTT4XCYTZs2Ne+++27z4YcfNrt37+7aLy8vzxw8eLAZHBxsSjJTUlLM5cuXmz169DCDgoLM0NBQ8/rrrzf37dt3+Q17GerrdwkA6qqqzt+/ZJjmv+6bRgUFBQVyOBzKz89XaGhouXVnzpxRRkaG4uPj1bhxY4sqREPAdwkA3Kuq8/cvMUYIAADYFkEIAADYFkEIAADYFkEIAADYli2C0M0336xmzZrx4DwAAFCOLYLQjBkztGLFCo8c+8KTlIHa4jsEANaxxQMVBw4cqG3btrn1mP7+/vLx8dGRI0cUEREhf39/GYbh1vdAw2aapkpKSvTzzz/Lx8fH9XRuAID3WB6EduzYoWeeeUZ79+5VTk6ONm7cqNGjR5fbZunSpXrmmWeUk5Ojzp07a+HChbrmmmusKfhffHx8FB8fr5ycHB05csTSWlC/NWnSRDExMfLxsUUHLQDUKZYHoaKiInXv3l233367xo4dW2H96tWrNXPmTC1dulRXX321Xn75ZQ0fPlzp6emKiYmRJPXq1Utnz56tsO8HH3ygVq1aeax2f39/xcTE6Pz585XOaQVciq+vrxo1akRvIgBYxPIgNHz4cNfM5ZV5/vnndccdd+jOO++UJC1cuFCbN2/WSy+9pPnz50uS9u7d65Zazp49Wy5QFRQUXHIfwzDk5+cnPz8/t9QAAAC8p073xZeUlGjv3r0aMmRIueVDhgxRamqq299v/vz5cjgcrp+qZksHAAD1X50OQkePHlVpaalatmxZbnnLli2Vm5tb7eMMHTpU48aN06ZNm9SmTRvt2bOn0u1mz56t/Px81092dvZl1Q8AAOo2yy+NVccvx0+YplmjMRWbN2+u1nYBAQEKCAioUW0AAKD+qtM9QuHh4fL19a3Q+5OXl1ehlwgAAKCm6nQQ8vf3V69evbRly5Zyy7ds2aK+fftaVBUAAGgoLL80durUKR08eND1OiMjQ2lpaQoLC1NMTIweeOABTZ48Wb1791afPn20fPlyZWVlafr06RZWDQAAGgLLg9Bnn32mgQMHul4/8MADkqSpU6fq9ddf14QJE3Ts2DE9+eSTysnJUZcuXbRp0ybFxsZaVTIAAGggDNM0TauLqKsKCgrkcDiUn5+v0NBQq8sBAADVUJPzd50eIwQAAOBJBCEAAGBbBKFKJCcnKzExUUlJSVaXAgAAPIgxQlVgjBAAAPUPY4QAAACqgSAEAABsiyAEAABsiyAEAABsiyAEAABsiyAEAABsiyAEAABsiyAEAABsiyAEAABsiyBUCabYAADAHphiowpMsQEAQP3DFBsAAADVQBACAAC2RRACAAC2RRACAAC2RRACAAC2RRACAAC2RRACAAC2RRACAAC2RRACAAC2RRACAAC2RRACAAC2RRCqBJOuAgBgD0y6WgUmXQUAoP5h0lUAAIBqIAgBAADbIggBAADbIggBAADbIggBAADbIggBAADbIggBAADbIggBAADbIggBAADbIggBAADbIggBAADbIggBAADbIggBAADbIghVIjk5WYmJiUpKSrK6FAAA4EGGaZqm1UXUVQUFBXI4HMrPz1doaKjV5QAAgGqoyfmbHiEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBKFKJCcnKzExUUlJSVaXAgAAPMgwTdOs6U5nz57V7t27lZmZqdOnTysiIkJXXnml4uPjPVGjZQoKCuRwOJSfn6/Q0FCrywEAANVQk/N3o5ocODU1VUuWLNHbb7+tkpISNW3aVIGBgTp+/LjOnj2rtm3b6ne/+52mT5+ukJCQy/oQAAAAnlbtS2OjRo3SLbfcotatW2vz5s0qLCzUsWPH9OOPP+r06dM6cOCAHnvsMX300Ufq0KGDtmzZ4sm6AQAALlu1e4SGDBmitWvXyt/fv9L1bdu2Vdu2bTV16lR9/fXXOnLkiNuKBAAA8IRajRGyC8YIAQBQ/9Tk/F2ru8ays7P1448/ul7v3r1bM2fO1PLly2tzOAAAAEvUKgjddtttSklJkSTl5uZq8ODB2r17tx555BE9+eSTbi0QAADAU2oVhL766itdddVVkqQ1a9aoS5cuSk1N1ZtvvqnXX3/dnfUBAAB4TK2C0Llz5xQQECBJ+vDDD3XTTTdJkq644grl5OS4rzoAAAAPqlUQ6ty5s5YtW6adO3dqy5YtGjZsmCTpyJEjat68uVsLBAAA8JRaBaGnn35aL7/8sgYMGKCJEyeqe/fukqR3333XdckMAACgrqv17fOlpaUqKChQs2bNXMsyMzPVpEkTtWjRwm0FWonb5wEAqH88NsXGf/L19S0XgiQpLi6utocDAADwumpfGhs2bJhSU1MvuV1hYaGefvppJScnX1ZhAAAAnlbtHqFx48Zp/PjxCgkJ0U033aTevXurVatWaty4sU6cOKH09HR9/PHH2rRpk2688UY988wznqwbAADgstVojFBJSYnWrVun1atXa+fOnTp58mTZQQxDiYmJGjp0qO666y517NjRU/V6FWOEAACof2py/r6sucby8/NVXFys5s2by8/Pr7aHqbMIQgAA1D9eGSwtSQ6HQw6H43IOAQAAYJlaPUcIAACgISAIAQAA2yIIAQAA2yIIVSI5OVmJiYlKSkqyuhQAAOBBtb5r7OTJk1q3bp0OHTqkhx56SGFhYdq3b59atmyp1q1bu7tOS3DXGAAA9Y/H7xr78ssvNWjQIDkcDmVmZuquu+5SWFiYNm7cqB9++EErVqyoVeEAAADeVKtLYw888ICmTZumAwcOqHHjxq7lw4cP144dO9xWHAAAgCfVKgjt2bNHv//97yssb926tXJzcy+7KAAAAG+oVRBq3LixCgoKKiz/9ttvFRERcdlFAQAAeEOtgtCoUaP05JNP6ty5c5LK5hrLysrSww8/rLFjx7q1QAAAAE+pVRB69tln9fPPP6tFixYqLi5W//791b59e4WEhOi///u/3V0jAACAR9TqrrHQ0FB9/PHH2rp1q/bt2yen06mePXtq0KBB7q4PAADAYy5r9vmGjucIAQBQ/3hl9vndu3dr27ZtysvLk9PpLLfu+eefr+1hAQAAvKZWQehPf/qTHnvsMXXs2FEtW7aUYRiudf/5dwAAgLqsVkFo0aJF+stf/qJp06a5uRwAAADvqdVdYz4+Prr66qvdXQsAAIBX1SoI3X///UpOTnZ3LQAAAF5Vq0tjDz74oG644Qa1a9dOiYmJ8vPzK7d+w4YNbikOAADAk2oVhP7f//t/SklJ0cCBA9W8eXMGSAMAgHqpVkFoxYoVWr9+vW644QZ31wMAAOA1tRojFBYWpnbt2rm7FgAAAK+qVRCaO3eu5syZo9OnT7u7HgAAAK+p1aWxxYsX69ChQ2rZsqXi4uIqDJbet2+fW4oDAADwpFoFodGjR7u5DAAAAO9j0tUqMOkqAAD1T03O37UaIwQAANAQVPvSWFhYmL777juFh4erWbNmVT476Pjx424pDgAAwJOqHYReeOEFhYSEuP7OQxQBAEB9xxihKjBGCACA+sfjY4R8fX2Vl5dXYfmxY8fk6+tbm0PWKcnJyUpMTFRSUpLVpQAAAA+qVRC6WCfS2bNn5e/vf1kF1QX33nuv0tPTtWfPHqtLAQAAHlSj5wgtXrxYkmQYhl599VUFBwe71pWWlmrHjh264oor3FshAACAh9QoCL3wwguSynqEli1bVu4ymL+/v+Li4rRs2TL3VggAAOAhNQpCGRkZkqSBAwdqw4YNatasmUeKAgAA8IZaTbGRkpLi7joAAAC8jidLAwAA2yIIAQAA2yIIAQAA2yIIAQAA26rRYOkdO3ZUutzhcKh9+/YKCgpyS1EAAADeUKMgNGDAgIuu8/X11d13363nnntOfn5+l1sXAACWyS3KVVZBlmJCYxQZFGl1OfCgGgWhEydOVLr85MmT2r17tx566CFFRkbqkUcecUtxAAB424YDGzTvH/PkNJ3yMXw0p88cjUkYY3VZ8BC3zj7/zjvv6JFHHtHXX3/trkNaitnnAcBecotyNXT9UDlNp2uZj+GjzWM30zNUj3h89vmL6d69u3744Qd3HhIAAK/JKsgqF4IkyWk6lV2YbVFF8DS3BqEjR46oRYsW7jwkAABeExMaIx+j/KnRx/BRdEi0RRXB09wWhPLy8vTYY4/puuuuc9chAQDwqsigSM3pM8cVhi6MEeKyWMNVo8HSV155pQzDqLA8Pz9fP/74ozp16qS33nrLbcUBAOBtYxLGqG+rvsouzFZ0SDQhqIGrURAaPXp0pctDQ0N1xRVXaMiQIfL19XVHXQAAWCYyKJIAZBNuvWusoeGuMQAA6p+anL9r1CMkSWvXrtXbb7+tc+fOadCgQfrd735X60IBAACsVKMgtHz5ck2fPl0JCQlq3Lix1q9fr4yMDM2fP99T9QEAAHhMje4aW7JkiR599FF9++23+uKLL/TnP/9ZL774oqdqAwAA8KgaBaHvv/9et99+u+v15MmTdfbsWeXm5rq9MAAAAE+rURAqLi5WcHCw67Wvr68CAgJ0+vRptxcGAADgaTUeLP3qq6+WC0Pnz5/X66+/rvDwcNeyGTNmuKc6AAAAD6rR7fNxcXGVPlCx3AENQ99///1lF1YXcPs8AAD1j8dun8/MzLycugAAAOoUt066CgAAUJ/UKAht3bpViYmJKigoqLAuPz9fnTt31o4dO9xWHAAAgCfVKAgtXLhQd911V6XX2xwOh37/+9/rhRdecFtxAAAAnlSjIPTFF19o2LBhF10/ZMgQ7d2797KLAgAA8IYaBaGffvpJfn5+F13fqFEj/fzzz5ddFAAAgDfUKAi1bt1a+/fvv+j6L7/8UlFRUZddFAAAgDfUKAiNGDFCTzzxhM6cOVNhXXFxsebMmaMbb7zRbcUBAAB4Uo0eqPjTTz+pZ8+e8vX11X333aeOHTvKMAx98803Sk5OVmlpqfbt26eWLVt6smav4YGKAADUPx57oGLLli2Vmpqqu+++W7Nnz9aFDGUYhoYOHaqlS5c2mBAEAAAavhrPNRYbG6tNmzbpxIkTOnjwoEzTVEJCgpo1a+aJ+gAAADymxkHogmbNmikpKcmdtQAAAHgVU2xUIjk5WYmJiQQ9ALCpnPxipR46qpz8YqtLgYfVaLC03TBYGgDsZ/WeLM3esF9OU/IxpPljumpCUozVZaEGanL+pkcIAIB/yckvdoUgSXKa0iMbvqJnqAEjCAEA8C8ZR4vkNKVIHVMfn68VqWMqNU1lHj1tdWnwkFoPlgYAoKGJDw/Srb4p+u9Gr8rXMFVqGnr0/F2KC7/O6tLgIfQIAQDwL1E6rvl+f5avUXZtzNcw9Sf/VxWl4xZXBk8hCAEAcMHxQzLkLLfIx3RKx7+3qCB4GkEIAIALwtpJxi9OjYavFNbWmnrgcQQhAAAucLSWRi4qCz9S2Z8jF5YtR4PEYGkAAP5TzylSu+vLLoeFtSUENXAEIQAAfsnRmgBkE1waAwAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAttXgg1B2drYGDBigxMREdevWTWvXrrW6JAAAUEc0sroAT2vUqJEWLlyoHj16KC8vTz179tSIESMUFBRkdWkAAMBiDT4IRUVFKSoqSpLUokULhYWF6fjx4wQhAABg/aWxHTt2aOTIkWrVqpUMw9Dbb79dYZulS5cqPj5ejRs3Vq9evbRz585avddnn30mp9Op6Ojoy6waAAA0BJb3CBUVFal79+66/fbbNXbs2ArrV69erZkzZ2rp0qW6+uqr9fLLL2v48OFKT09XTEyMJKlXr146e/ZshX0/+OADtWrVSpJ07NgxTZkyRa+++upFazl79my54xQUFFzuxwMAAHWYYZqmaXURFxiGoY0bN2r06NGuZb/61a/Us2dPvfTSS65lnTp10ujRozV//vxqHffs2bMaPHiw7rrrLk2ePPmi282dO1fz5s2rsDw/P1+hoaHV/yAAAMAyBQUFcjgc1Tp/W35prColJSXau3evhgwZUm75kCFDlJqaWq1jmKapadOm6brrrqsyBEnS7NmzlZ+f7/rJzs6ude0AAKDus/zSWFWOHj2q0tJStWzZstzyli1bKjc3t1rH+OSTT7R69Wp169bNNf5o5cqV6tq1a4VtAwICFBAQcNl1AwCA+qFOB6ELDMMo99o0zQrLLqZfv35yOp2eKAsAANRzdfrSWHh4uHx9fSv0/uTl5VXoJQIAAKipOh2E/P391atXL23ZsqXc8i1btqhv374WVQUA1sktytXunN3KLare8AAAVbP80tipU6d08OBB1+uMjAylpaUpLCxMMTExeuCBBzR58mT17t1bffr00fLly5WVlaXp06dbWDUAeN+GAxs07x/z5DSd8jF8NKfPHI1JGGN1WUC9Zvnt89u2bdPAgQMrLJ86dapef/11SWUPVFywYIFycnLUpUsXvfDCC7r22ms9XltNbr8DAE/KLcrV0PVD5TT/PebRx/DR5rGbFRkUaWFlQN1Tk/O35T1CAwYM0KWy2D333KN77rnHSxUBQN2TVZBVLgRJktN0KrswmyAEXIY6PUYIAFAmJjRGPkb5X9k+ho+iQ5gyCLgcBKFKJCcnKzExUUlJSVaXAgCSpMigSM3pM8cVhi6MEaI3CLg8lo8RqssYIwSgrsktylV2YbaiQ6IJQcBF1KsxQgCA6osMiiQAAW7EpTEAAGBbBCEAAGBbBCEAAGBbBCEAAGBbBCE0aMzLBACoCneNocFiXibvyS3KVVZBlmJCY7ijCUC9Qo+QReip8KzcolxXCJLKpiKY9495tLcHbDiwQUPXD9UdH9yhoeuHasOBDVaXBADVRhCyACcOz6tqXia4D4ETQH1HEPIyThzewbxM3kHgBFDfEYQq4cm5xjhxeAfzMnkHgRNAfcdcY1XwxFxjuUW5Grp+aLkw5GP4aPPYzZykPYB5mTyPQekA6pqanL8JQlXw1KSrnDjQ0BA4AdQlBCE38eTs85w4AADwDGafrweYQRoAAOsxWBoAANgWQQgA6pGc/GKlHjqqnPxiq0sBGgQujQFAPbF6T5Zmb9gvpyn5GNL8MV01ISnG6rKAeo0eIQCoB3Lyi10hSJKcpvTIhq/oGQIuE0EIAOqBjKNFrhB0QalpKvPoaWsKAhoIghAA1APx4UHyMcov8zUMxYU3saYgoIEgCAFAPRDlCNT8MV3V2jiuPj5fq7VxXH8a00VRjkCrSwPqNQZLVyI5OVnJyckqLS21uhQAcJngu03jG/9BhumUafjI8F0kaYrVZQH1Gk+WroInnywNADWSf1ha2EX6z0mbDV9p5n7J0dq6uoA6qCbnby6NAUB9cPxQ+RAkSWapdPx7a+oBGgiCEADUB2HtJOMXv7INXymsrTX1AA0EQQgA6gNHa2nkorLwI5X9OXIhl8WAy8RgaQCoL3pOkdpdX3Y5LKwtIQhwA4IQANQnjtYEIMCNuDQGAABsiyAEAABsiyAEAABsiyBkkZz8YqUeOsrM0QAAWIjB0hZYvSdLszfsl9OUfAxp/piumpAUY3VZAADYDj1CXpaTX+wKQZLkNKVHNnxFzxAAABYgCHlZxtEiVwi6oNQ0lXn0tDUFAQBgYwQhL4sPD5KPIUXqmPr4fK1IHZOvYSguvInVpQEAYDuMEapEcnKykpOTVVpa6vZjRzkC9VbvA+r15Vz5GqZKTUN7u81VlGOE298LAABUzTBN07z0ZvZUUFAgh8Oh/Px8hYaGuueg+YelhV3KzyJt+Eoz9/O0WA/IyS9WxtEixYcHKcoRaHU5AAAvqMn5mx4hbzt+qHwIkiSztGzuIIKQW3F3HgDgUhgj5G1h7STjF81u+JZNoAi3uXB3XguzbCxWC/MYd+d5Uv5hKWNH2Z8AUI/QI+RtjtbSyEXS32aW9QQZvtLIhfQGuVnG0SLd4pOi+Y1edY3Fmn3+TmUe/RWXyNxt3wqZf/uDDNMp0/CRMXJR2SzpAFAPMEaoCh4ZI3RB/uGyy2FhbQlBHvDTj4cU/kov+Rr//nqfN3107K7P1LJNOwsra2DyD8t8oYsM/ftyr9Pwkc/Mr/heA7BMTc7fXBqziqO1FH8NJwsPaXnucLkQJEmNDKdanjtiUUUN07Hs9HIhSJJ8TKeOZX9jUUUAUDMEITRMjMXyigxnpEpNo9yy86aPMp2RFlUEADVDEELDdGEsluFb9pqxWB7ROra9Hj1/p86bZb9Kzps+euz8nWoVy+VHAPUDg6XRcPWcIrW7nrFYHhTlCNSVo2eo/4buijZylW1GasaY/gxIB1BvMFi6Ch4dLA00IDn5xco8elpx4U0IQQAsxwMVAXhVlCOQAASgXmKMEAAAsC2CEAAAsC2CEAAAsC2CEAAAsC2CEAAAsC2CEAAAsC2CUCWSk5OVmJiopKQkq0sBAAAexAMVq8ADFQEAqH+YfR4AAKAaCEIAAMC2CEIAAMC2mGusCheGTxUUFFhcCQAAqK4L5+3qDIMmCFWhsLBQkhQdHW1xJQAAoKYKCwvlcDiq3Ia7xqrgdDp15MgRhYSEyDAMJSUlac+ePRW2q2z5pZYVFBQoOjpa2dnZXrkj7WK1e2L/6mxb1Ta0c91r58qWW9nOVdXpiX0vtX1t11fnO/3L13yn+d1xufvboZ1N01RhYaFatWolH5+qRwHRI1QFHx8ftWnTxvXa19e30v9QlS2v7rLQ0FCv/E92sdo9sX91tq1qG9q57rVzZcutbOeLvb+n9r3U9rVdX53v78X25Ttd/XX87qj5tg2hnS/VE3QBg6Vr4N5776328uou85bLfe+a7F+dbavahnZ237buaufKllvZzpf7/jXd91Lb13Z9db6/9bmda7o/vzu8s7+d27kyXBqzCA9r9A7a2TtoZ++hrb2DdvaOutDO9AhZJCAgQHPmzFFAQIDVpTRotLN30M7eQ1t7B+3sHXWhnekRAgAAtkWPEAAAsC2CEAAAsC2CEAAAsC2CEAAAsC2CEAAAsC2CUB303nvvqWPHjkpISNCrr75qdTkN2s0336xmzZrplltusbqUBis7O1sDBgxQYmKiunXrprVr11pdUoNUWFiopKQk9ejRQ127dtUrr7xidUkN2unTpxUbG6sHH3zQ6lIarEaNGqlHjx7q0aOH7rzzTo+9D7fP1zHnz59XYmKiUlJSFBoaqp49e2rXrl0KCwuzurQGKSUlRadOndIbb7yhdevWWV1Og5STk6OffvpJPXr0UF5ennr27Klvv/1WQUFBVpfWoJSWlurs2bNq0qSJTp8+rS5dumjPnj1q3ry51aU1SI8++qgOHDigmJgYPfvss1aX0yCFh4fr6NGjHn8feoTqmN27d6tz585q3bq1QkJCNGLECG3evNnqshqsgQMHKiQkxOoyGrSoqCj16NFDktSiRQuFhYXp+PHj1hbVAPn6+qpJkyaSpDNnzqi0tFT8O9czDhw4oH/+858aMWKE1aXADQhCbrZjxw6NHDlSrVq1kmEYevvttytss3TpUsXHx6tx48bq1auXdu7c6Vp35MgRtW7d2vW6TZs2Onz4sDdKr3cut61RPe5s588++0xOp1PR0dEerrr+cUc7nzx5Ut27d1ebNm00a9YshYeHe6n6+sMd7fzggw9q/vz5Xqq4fnJHOxcUFKhXr17q16+ftm/f7rFaCUJuVlRUpO7du+vFF1+sdP3q1as1c+ZMPfroo/r88891zTXXaPjw4crKypKkSv8FZxiGR2uury63rVE97mrnY8eOacqUKVq+fLk3yq533NHOTZs21RdffKGMjAy9+eab+umnn7xVfr1xue38zjvvqEOHDurQoYM3y6533PF9zszM1N69e7Vs2TJNmTJFBQUFninWhMdIMjdu3Fhu2VVXXWVOnz693LIrrrjCfPjhh03TNM1PPvnEHD16tGvdjBkzzFWrVnm81vquNm19QUpKijl27FhPl9gg1Ladz5w5Y15zzTXmihUrvFFmvXc53+cLpk+fbq5Zs8ZTJTYItWnnhx9+2GzTpo0ZGxtrNm/e3AwNDTXnzZvnrZLrJXd8n4cNG2bu2bPHI/XRI+RFJSUl2rt3r4YMGVJu+ZAhQ5SamipJuuqqq/TVV1/p8OHDKiws1KZNmzR06FAryq3XqtPWuHzVaWfTNDVt2jRdd911mjx5shVl1nvVaeeffvrJ9S/mgoIC7dixQx07dvR6rfVZddp5/vz5ys7OVmZmpp599lndddddeuKJJ6wot96qTjufOHFCZ8+elST9+OOPSk9PV9u2bT1STyOPHBWVOnr0qEpLS9WyZctyy1u2bKnc3FxJZbcLPvfccxo4cKCcTqdmzZrFXR+1UJ22lqShQ4dq3759KioqUps2bbRx40YlJSV5u9x6qzrt/Mknn2j16tXq1q2ba5zAypUr1bVrV2+XW29Vp51//PFH3XHHHTJNU6Zp6r777lO3bt2sKLfequ7vDVye6rTzN998o9///vfy8fGRYRhatGiRx+6eJghZ4JdjfkzTLLfspptu0k033eTtshqkS7U1d+S5R1Xt3K9fPzmdTivKanCqaudevXopLS3Ngqoankv93rhg2rRpXqqoYaqqnfv27av9+/d7pQ4ujXlReHi4fH19K/zLIi8vr0IyxuWhrb2DdvYO2tk7aGfvqGvtTBDyIn9/f/Xq1Utbtmwpt3zLli3q27evRVU1TLS1d9DO3kE7ewft7B11rZ25NOZmp06d0sGDB12vMzIylJaWprCwMMXExOiBBx7Q5MmT1bt3b/Xp00fLly9XVlaWpk+fbmHV9RNt7R20s3fQzt5BO3tHvWpnj9yLZmMpKSmmpAo/U6dOdW2TnJxsxsbGmv7+/mbPnj3N7du3W1dwPUZbewft7B20s3fQzt5Rn9qZucYAAIBtMUYIAADYFkEIAADYFkEIAADYFkEIAADYFkEIAADYFkEIAADYFkEIAADYFkEIAADYFkEIAADYFkEIACqxbds2GYahkydPVmv7AQMGaObMmR6tCYD7EYQA2F5lIaZv377KycmRw+GwpigAXkEQAmBb586du+g6f39/RUZGyjAML1YEwNsIQgC8ZsCAAZoxY4ZmzZqlsLAwRUZGau7cudXa1zAMvfTSSxo+fLgCAwMVHx+vtWvXltvmv/7rv9ShQwc1adJEbdu21eOPP14u7MydO1c9evTQX/7yF7Vt21YBAQGaOnWqtm/frkWLFskwDBmGoczMzEovjX3yySfq37+/mjRpombNmmno0KE6ceJEpfWWlJRo1qxZat26tYKCgvSrX/1K27Ztc63/4YcfNHLkSDVr1kxBQUHq3LmzNm3aVO22BOAejawuAIC9vPHGG3rggQe0a9cu/eMf/9C0adN09dVXa/DgwZfc9/HHH9f//M//aNGiRVq5cqUmTpyoLl26qFOnTpKkkJAQvf7662rVqpX279+vu+66SyEhIZo1a5brGAcPHtSaNWu0fv16+fr6KjY2VgcOHFCXLl305JNPSpIiIiKUmZlZ7r3T0tJ0/fXX67e//a0WL16sRo0aKSUlRaWlpZXWevvttyszM1NvvfWWWrVqpY0bN2rYsGHav3+/EhISdO+996qkpEQ7duxQUFCQ0tPTFRwcXMtWBVBrJgB4Sf/+/c1+/fqVW5aUlGT+13/91yX3lWROnz693LJf/epX5t13333RfRYsWGD26tXL9XrOnDmmn5+fmZeXV6GuP/zhD+WWpaSkmJLMEydOmKZpmhMnTjSvvvrqi77Xfx7j4MGDpmEY5uHDh8ttc/3115uzZ882TdM0u3btas6dO/eixwPgHfQIAfCqbt26lXsdFRWlvLy8au3bp0+fCq/T0tJcr9etW6eFCxfq4MGDOnXqlM6fP6/Q0NBy+8TGxioiIqLGdaelpWncuHHV2nbfvn0yTVMdOnQot/zs2bNq3ry5JGnGjBm6++679cEHH2jQoEEaO3ZshbYB4HmMEQLgVX5+fuVeG4Yhp9NZ6+NdGMz86aef6tZbb9Xw4cP13nvv6fPPP9ejjz6qkpKSctsHBQXV6n0CAwOrva3T6ZSvr6/27t2rtLQ0188333yjRYsWSZLuvPNOff/995o8ebL279+v3r17a8mSJbWqDUDtEYQA1BuffvpphddXXHGFpLKBzLGxsXr00UfVu3dvJSQk6IcffqjWcf39/S861ueCbt266aOPPqrW8a688kqVlpYqLy9P7du3L/cTGRnp2i46OlrTp0/Xhg0b9Mc//lGvvPJKtY4PwH24NAag3li7dq169+6tfv36adWqVdq9e7f+/Oc/S5Lat2+vrKwsvfXWW0pKStL777+vjRs3Vuu4cXFx2rVrlzIzMxUcHKywsLAK28yePVtdu3bVPffco+nTp8vf318pKSkaN26cwsPDy23boUMHTZo0SVOmTNFzzz2nK6+8UkePHtXWrVvVtWtXjRgxQjNnztTw4cPVoUMHnThxQlu3bnUN+gbgPfQIAag35s2bp7feekvdunXTG2+8oVWrVikxMVGSNGrUKN1///2677771KNHD6Wmpurxxx+v1nEffPBB+fr6KjExUREREcrKyqqwTYcOHfTBBx/oiy++0FVXXaU+ffronXfeUaNGlf978rXXXtOUKVP0xz/+UR07dtRNN92kXbt2KTo6WpJUWlqqe++9V506ddKwYcPUsWNHLV26tJYtA6C2DNM0TauLAIBLMQxDGzdu1OjRo60uBUADQo8QAACwLYIQAMutWrVKwcHBlf507tzZ6vIANGBcGgNgucLCQv3000+VrvPz81NsbKyXKwJgFwQhAABgW1waAwAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtkUQAgAAtvX/AesJiQWUNwW4AAAAAElFTkSuQmCC\n", 163 | "text/plain": [ 164 | "
" 165 | ] 166 | }, 167 | "metadata": {}, 168 | "output_type": "display_data" 169 | } 170 | ], 171 | "source": [ 172 | "import matplotlib.pyplot as plt\n", 173 | "plt.plot(n_particles, np.array(t_no_save), '.', label='no save')\n", 174 | "plt.plot(n_particles, np.array(t_save_par), '.', label='save particles')\n", 175 | "plt.plot(n_particles, np.array(t_save_stats), '.', label='save stats')\n", 176 | "plt.xscale('log')\n", 177 | "plt.yscale('log')\n", 178 | "plt.xlabel('n_particles')\n", 179 | "plt.ylabel('CPU time (s)')\n", 180 | "plt.legend()" 181 | ] 182 | }, 183 | { 184 | "cell_type": "code", 185 | "execution_count": null, 186 | "id": "eb1333a9-bbe2-408d-ac72-d87312c01a3a", 187 | "metadata": {}, 188 | "outputs": [], 189 | "source": [] 190 | } 191 | ], 192 | "metadata": { 193 | "kernelspec": { 194 | "display_name": "Python 3 (ipykernel)", 195 | "language": "python", 196 | "name": "python3" 197 | }, 198 | "language_info": { 199 | "codemirror_mode": { 200 | "name": "ipython", 201 | "version": 3 202 | }, 203 | "file_extension": ".py", 204 | "mimetype": "text/x-python", 205 | "name": "python", 206 | "nbconvert_exporter": "python", 207 | "pygments_lexer": "ipython3", 208 | "version": "3.9.13" 209 | } 210 | }, 211 | "nbformat": 4, 212 | "nbformat_minor": 5 213 | } 214 | --------------------------------------------------------------------------------