diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 4664428bc..53d9d4342 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -146,6 +146,30 @@ set_tests_properties(chicane.analysis PROPERTIES DEPENDS "chicane.run") set_tests_properties(chicane.plot PROPERTIES DEPENDS "chicane.run") +# Python: Chicane ##################################################################### +# +if(ImpactX_PYTHON) + file(MAKE_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/chicane.py) + add_test(NAME chicane.py.run + COMMAND ${Python_EXECUTABLE} + ${ImpactX_SOURCE_DIR}/examples/chicane/run_chicane.py + WORKING_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/chicane.py + ) + add_test(NAME chicane.py.analysis + COMMAND ${ImpactX_SOURCE_DIR}/examples/chicane/analysis_chicane.py + WORKING_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/chicane.py + ) + add_test(NAME chicane.py.plot + COMMAND ${ImpactX_SOURCE_DIR}/examples/chicane/plot_chicane.py --save-png + WORKING_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/chicane.py + ) + + set_tests_properties(chicane.py.analysis PROPERTIES DEPENDS "chicane.py.run") + set_tests_properties(chicane.py.plot PROPERTIES DEPENDS "chicane.py.run") + impactx_test_set_pythonpath(chicane.py.run) +endif() + + # Constant Focusing Channel ################################################### # file(MAKE_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/cfchannel) diff --git a/examples/chicane/README.rst b/examples/chicane/README.rst index 79c4d2ecd..e156fd9d0 100644 --- a/examples/chicane/README.rst +++ b/examples/chicane/README.rst @@ -17,12 +17,20 @@ In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y` Run --- -This example an either be run with an inputs file (``impactx input_chicane.in``): +This example an either be run as a Python script (``python3 run_chicane.py``): + +.. literalinclude:: run_chicane.py + :language: python3 + :caption: You can copy this file from ``examples/chicane/run_chicane.py``. + +Or with an inputs file (``impactx input_chicane.in``): .. literalinclude:: input_chicane.in :language: ini :caption: You can copy this file from ``examples/chicane/input_chicane.in``. +Both execution modes can also be prefixed with an MPI executor, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + Analyze ------- diff --git a/examples/chicane/run_chicane.py b/examples/chicane/run_chicane.py new file mode 100644 index 000000000..5a1815c9f --- /dev/null +++ b/examples/chicane/run_chicane.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +# +# Copyright 2022 ImpactX contributors +# Authors: Marco Garten, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import amrex +from impactx import ImpactX, RefPart, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +sim.set_particle_shape(2) # B-spline order +sim.set_space_charge(False) +#sim.set_diagnostics(False) # benchmarking +sim.set_slice_step_diagnostics(True) + +# domain decomposition & space charge mesh +sim.init_grids() + +# load a 5 GeV electron beam with an initial +# normalized transverse rms emittance of 1 um +energy_MeV = 5.0e3 # reference energy +charge_C = 0.0 # used with space charge +mass_MeV = 0.510998950 +qm_qeeV = -1.0e-6/mass_MeV # charge/mass +npart = 10000 # number of macro particles + +distr = distribution.Waterbag( + sigmaX = 2.2951017632e-5, + sigmaY = 1.3084093142e-5, + sigmaT = 5.5555553e-8, + sigmaPx = 1.598353425e-6, + sigmaPy = 2.803697378e-6, + sigmaPt = 2.000000000e-6, + muxpx = 0.933345606203060, + muypy = 0.933345606203060, + mutpt = 0.999999961419755) +sim.add_particles(qm_qeeV, charge_C, distr, npart) + +# set the energy in the reference particle +sim.particle_container().ref_particle() \ + .set_energy_MeV(energy_MeV, mass_MeV) + +# design the accelerator lattice +ns = 25 # number of slices per ds in the element +rc = 10.35 # bend radius (meters) +psi = 0.048345620280243 # pole face rotation angle (radians) + +# Drift elements +dr1 = elements.Drift(ds=5.0058489435, nslice=ns) +dr2 = elements.Drift(ds=1.0, nslice=ns) +dr3 = elements.Drift(ds=2.0, nslice=ns) + +# Bend elements +sbend1 = elements.Sbend(ds=0.50037, rc=-rc, nslice=ns) +sbend2 = elements.Sbend(ds=0.50037, rc=rc, nslice=ns) + +# Dipole Edge Focusing elements +dipedge1 = elements.DipEdge(psi=-psi, rc=-rc, g=0.0, K2=0.0) +dipedge2 = elements.DipEdge(psi=psi, rc=rc, g=0.0, K2=0.0) + +lattice_half = [ + sbend1, + dipedge1, + dr1, + dipedge2, + sbend2 +] +# assign a segment with the first half of the lattice +sim.lattice.extend(lattice_half) +sim.lattice.append(dr2) +lattice_half.reverse() +# extend the lattice by a reversed half +sim.lattice.extend(lattice_half) +sim.lattice.append(dr3) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amrex.finalize()