Skip to content

Commit

Permalink
Aperture Update API: Enums & Strings
Browse files Browse the repository at this point in the history 
Update the C++ API to use an enum for the aperture shape and
the Python API & inputs file syntax to accept a string. That
makes the parameters at the call sites self-describing.
  • Loading branch information
@ax3l
ax3l committed Aug 9, 2023
1 parent 42c384e commit 0c84778
Show file tree
Hide file tree
Showing 7 changed files with 56 additions and 44 deletions.
4 changes: 1 addition & 3 deletions docs/source/usage/parameters.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -473,9 +473,7 @@ Lattice Elements

* ``<element_name>.ymax`` (``float``, in meters) maximum value of the vertical coordinate

* ``<element_name>.shape`` (``integer``, dimensionless) shape of the aperture boundary

(m = 0) rectangular, (m = 1) elliptical
* ``<element_name>.shape`` (``string``) shape of the aperture boundary: ``rectangular`` (default) or ``elliptical``

* ``beam_monitor`` a beam monitor, writing all beam particles at fixed ``s`` to openPMD files.
If the same element name is used multiple times, then an output series is created with multiple outputs.
Expand Down
4 changes: 2 additions & 2 deletions docs/source/usage/python.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -651,13 +651,13 @@ References:
:param phi_in: angle of the reference particle with respect to the longitudinal (z) axis in the original frame in degrees
:param phi_out: angle of the reference particle with respect to the longitudinal (z) axis in the rotated frame in degrees

.. py:class:: impactx.elements.Aperture(xmax, ymax, shape=0)
.. py:class:: impactx.elements.Aperture(xmax, ymax, shape="rectangular")
A thin collimator element, applying a transverse aperture boundary.

:param xmax: maximum allowed value of the horizontal coordinate (meter)
:param ymax: maximum allowed value of the vertical coordinate (meter)
:param shape: aperture boundary shape (0 rectangular, 1 elliptical)
:param shape: aperture boundary shape: ``"rectangular"`` (default) or ``"elliptical"``

.. py:class:: impactx.elements.SoftQuadrupole(ds, gscale, cos_coefficients, sin_coefficients, nslice=1)
Expand Down
2 changes: 1 addition & 1 deletion examples/aperture/input_aperture.in
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ monitor.type = beam_monitor
monitor.backend = h5

collimator.type = aperture
collimator.shape = 0
collimator.shape = rectangular
collimator.xmax = 1.0e-3
collimator.ymax = 1.5e-3

Expand Down
2 changes: 1 addition & 1 deletion examples/aperture/run_aperture.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@
sim.lattice.extend(
[
monitor,
elements.Aperture(xmax=1.0e-3, ymax=1.5e-3, shape=0),
elements.Aperture(xmax=1.0e-3, ymax=1.5e-3, shape="rectangular"),
monitor,
]
)
Expand Down
9 changes: 7 additions & 2 deletions src/initialization/InitElement.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,10 +226,15 @@ namespace detail
m_lattice.emplace_back( ChrAcc(ds, ez, bz, nslice) );
} else if (element_type == "aperture") {
amrex::Real xmax, ymax;
int shape = 0;
std::string shape_str = "rectangular";
pp_element.get("xmax", xmax);
pp_element.get("ymax", ymax);
pp_element.queryAdd("shape", shape);
pp_element.queryAdd("shape", shape_str);
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(shape_str == "rectangular" || shape_str == "elliptical",
element_name + ".shape must be \"rectangular\" or \"elliptical\"");
Aperture::Shape shape = shape_str == "rectangular" ?
Aperture::Shape::rectangular :
Aperture::Shape::elliptical;
m_lattice.emplace_back( Aperture(xmax, ymax, shape) );
} else if (element_type == "beam_monitor") {
std::string openpmd_name = element_name;
Expand Down
60 changes: 30 additions & 30 deletions src/particles/elements/Aperture.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,16 +32,22 @@ namespace impactx
static constexpr auto name = "Aperture";
using PType = ImpactXParticleContainer::ParticleType;

enum Shape
{
rectangular,
elliptical
};

/** A thin collimator element that applies a transverse aperture boundary.
* Particles outside the boundary are considered lost.
*
* @param shape aperture shape (0 rectangular, 1 elliptical)
* @param shape aperture shape
* @param xmax maximum value of horizontal coordinate (m)
* @param ymax maximum value of vertical coordinate (m)
*/
Aperture( amrex::ParticleReal const xmax,
amrex::ParticleReal const ymax,
int const shape)
Aperture (amrex::ParticleReal xmax,
amrex::ParticleReal ymax,
Shape shape)
: m_shape(shape), m_xmax(xmax), m_ymax(ymax)
{
}
Expand All @@ -60,55 +66,49 @@ namespace impactx
*/
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void operator() (
PType& AMREX_RESTRICT p,
[[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT px,
[[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT py,
[[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT pt,
[[maybe_unused]] RefPart const & refpart) const {

PType& AMREX_RESTRICT p,
[[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT px,
[[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT py,
[[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT pt,
[[maybe_unused]] RefPart const & refpart) const
{
using namespace amrex::literals; // for _rt and _prt

// access AoS data such as positions and cpu/id
amrex::ParticleReal const x = p.pos(RealAoS::x);
amrex::ParticleReal const y = p.pos(RealAoS::y);
int const id = p.id();
//uint64_t const global_id = ablastr::particles::localIDtoGlobal(p.id(), p.cpu());
auto const id = p.id();

// scale horizontal and vertical coordinates
amrex::ParticleReal const u = x/m_xmax;
amrex::ParticleReal const v = y/m_ymax;
amrex::ParticleReal const u = x / m_xmax;
amrex::ParticleReal const v = y / m_ymax;

// compare against the aperture boundary

switch(m_shape) {

// rectangular aperture (default)
case 0 :
if (pow(u,2)>1 || pow(v,2)>1){
p.pos(RealAoS::x) = 0.0_prt; //replace with id change of sign
p.pos(RealAoS::y) = 0.0_prt; //replace with id change of sign
switch (m_shape)
{
case Shape::rectangular : // default
if (pow(u,2)>1 || pow(v,2) > 1_prt) {
p.pos(RealAoS::x) = 0.0_prt; // replace with id change of sign
p.pos(RealAoS::y) = 0.0_prt; // replace with id change of sign
p.id() = -id;
}
break;

// elliptical aperture
case 1 :
if (pow(u,2)+pow(v,2)>1){
p.pos(RealAoS::x) = 0.0_prt; //replace with id change of sign
p.pos(RealAoS::y) = 0.0_prt; //replace with id change of sign
case Shape::elliptical :
if (pow(u,2)+pow(v,2) > 1_prt) {
p.pos(RealAoS::x) = 0.0_prt; // replace with id change of sign
p.pos(RealAoS::y) = 0.0_prt; // replace with id change of sign
p.id() = -id;
}
break;

}

}

/** This pushes the reference particle. */
using Thin::operator();

private:
int m_shape; //! aperture type (1 rectangular, 2 elliptical)
Shape m_shape; //! aperture type (rectangular, elliptical)
amrex::ParticleReal m_xmax; //! maximum horizontal coordinate
amrex::ParticleReal m_ymax; //! maximum vertical coordinate

Expand Down
19 changes: 14 additions & 5 deletions src/python/elements.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -360,11 +360,20 @@ void init_elements(py::module& m)
;

py::class_<Aperture, elements::Thin>(me, "Aperture")
.def(py::init<
amrex::ParticleReal const,
amrex::ParticleReal const,
int>(),
py::arg("xmax"), py::arg("ymax"), py::arg("shape") = 0,
.def(py::init([](
amrex::ParticleReal xmax,
amrex::ParticleReal ymax,
std::string shape)
{
if (shape != "rectangular" && shape != "elliptical")
throw std::runtime_error("shape must be \"rectangular\" or \"elliptical\"");

Aperture::Shape s = shape == "rectangular" ?
Aperture::Shape::rectangular :
Aperture::Shape::elliptical;
return new Aperture(xmax, ymax, s);
}),
py::arg("xmax"), py::arg("ymax"), py::arg("shape") = "rectangular",
"A short collimator element applying a transverse aperture boundary."
)
;
Expand Down

0 comments on commit 0c84778

Please sign in to comment.