Skip to content

Commit

Permalink
Py Propertes: Aperture, ChrQuad, ChrAcc
Browse files Browse the repository at this point in the history 
Expose element properties to Python for `Aperture`, `ChrQuad`, and
`ChrAcc`.
  • Loading branch information
@ax3l
ax3l committed Jan 8, 2024
1 parent dc5450d commit f865a2e
Show file tree
Hide file tree
Showing 5 changed files with 111 additions and 33 deletions.
45 changes: 45 additions & 0 deletions docs/source/usage/python.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -665,6 +665,39 @@ This module provides elements for the accelerator lattice.
:param rotation: rotation error in the transverse plane [degrees]
:param nslice: number of slices used for the application of space charge

.. py:property:: k
quadrupole strength in 1/m^2 (or T/m)

.. py:property:: units
unit specification for quad strength

.. py:class:: impactx.elements.ChrAcc(ds, ez, bz, dx=0, dy=0, rotation=0, nslice=1)
Acceleration in a uniform field Ez, with a uniform solenoidal field Bz.

The Hamiltonian is expanded through second order in the
transverse variables (x,px,y,py), with the exact pt dependence retained.

:param ds: Segment length in m.
:param ez: electric field strength in m^(-1)
= (charge * electric field Ez in V/m) / (m*c^2)
:param bz: magnetic field strength in m^(-1)
= (charge * magnetic field Bz in T) / (m*c)
:param dx: horizontal translation error in m
:param dy: vertical translation error in m
:param rotation: rotation error in the transverse plane [degrees]
:param nslice: number of slices used for the application of space charge

.. py:property:: ez
electric field strength in 1/m

.. py:property:: bz
magnetic field strength in 1/m

.. py:class:: impactx.elements.RFCavity(ds, escale, freq, phase, dx=0, dy=0, rotation=0, mapsteps=1, nslice=1)
A radiofrequency cavity.
Expand Down Expand Up @@ -797,6 +830,18 @@ This module provides elements for the accelerator lattice.
:param dy: vertical translation error in m
:param rotation: rotation error in the transverse plane [degrees]

.. py:property:: shape
aperture type (rectangular, elliptical)

.. py:property:: xmax
maximum horizontal coordinate

.. py:property:: ymax
maximum vertical coordinate

.. py:class:: impactx.elements.SoftQuadrupole(ds, gscale, cos_coefficients, sin_coefficients, dx=0, dy=0, rotation=0, mapsteps=1, nslice=1)
A soft-edge quadrupole.
Expand Down
1 change: 0 additions & 1 deletion src/particles/elements/Aperture.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,6 @@ namespace impactx
/** This pushes the reference particle. */
using Thin::operator();

private:
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: 9 additions & 10 deletions src/particles/elements/ChrQuad.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,12 +81,13 @@ namespace impactx
*/
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void operator() (
PType& AMREX_RESTRICT p,
amrex::ParticleReal & AMREX_RESTRICT px,
amrex::ParticleReal & AMREX_RESTRICT py,
amrex::ParticleReal & AMREX_RESTRICT pt,
RefPart const & refpart) const {

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

// shift due to alignment errors of the element
Expand Down Expand Up @@ -153,7 +154,6 @@ namespace impactx
q2 = x;
p1 = py;
p2 = px;

}

// advance longitudinal position and momentum
Expand Down Expand Up @@ -188,8 +188,8 @@ namespace impactx
* @param[in,out] refpart reference particle
*/
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void operator() (RefPart & AMREX_RESTRICT refpart) const {

void operator() (RefPart & AMREX_RESTRICT refpart) const
{
using namespace amrex::literals; // for _rt and _prt

// assign input reference particle values
Expand Down Expand Up @@ -219,7 +219,6 @@ namespace impactx
refpart.s = s + slice_ds;
}

private:
amrex::ParticleReal m_k; //! quadrupole strength in 1/m^2 (or T/m)
int m_unit; //! unit specification for quad strength
};
Expand Down
28 changes: 14 additions & 14 deletions src/particles/elements/ChrUniformAcc.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,12 +76,13 @@ namespace impactx
*/
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void operator() (
PType& AMREX_RESTRICT p,
amrex::ParticleReal & AMREX_RESTRICT px,
amrex::ParticleReal & AMREX_RESTRICT py,
amrex::ParticleReal & AMREX_RESTRICT pt,
RefPart const & refpart) const {

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

// shift due to alignment errors of the element
Expand Down Expand Up @@ -110,11 +111,11 @@ namespace impactx
amrex::ParticleReal const pti_tot = pti_ref + pt;
amrex::ParticleReal const ptf_tot = ptf_ref + pt;
amrex::ParticleReal const pzi_tot = sqrt(pow(pti_tot,2)-1_prt);
amrex::ParticleReal const pzf_tot = sqrt(pow(ptf_tot,2)-1_prt);
amrex::ParticleReal const pzf_tot = sqrt(pow(ptf_tot,2)-1_prt);
amrex::ParticleReal const pzi_ref = sqrt(pow(pti_ref,2)-1_prt);
amrex::ParticleReal const pzf_ref = sqrt(pow(ptf_ref,2)-1_prt);
amrex::ParticleReal const pzf_ref = sqrt(pow(ptf_ref,2)-1_prt);

amrex::ParticleReal const numer = -ptf_tot + pzf_tot;
amrex::ParticleReal const numer = -ptf_tot + pzf_tot;
amrex::ParticleReal const denom = -pti_tot + pzi_tot;

// compute focusing constant (1/m) and rotation angle (in rad)
Expand All @@ -138,8 +139,8 @@ namespace impactx

// the correct symplectic update for t
tout = t + (pzf_tot - pzf_ref - pzi_tot + pzi_ref)/m_ez;
tout = tout + (1_prt/pzi_tot - 1_prt/pzf_tot)*(pow(py-alpha*x,2)+pow(px+alpha*y,2))/(2_prt*m_ez);
ptout = pt;
tout = tout + (1_prt/pzi_tot - 1_prt/pzf_tot)*(pow(py-alpha*x,2)+pow(px+alpha*y,2))/(2_prt*m_ez);
ptout = pt;

// assign intermediate momenta
px = pxout;
Expand Down Expand Up @@ -175,8 +176,8 @@ namespace impactx
* @param[in,out] refpart reference particle
*/
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void operator() (RefPart & AMREX_RESTRICT refpart) const {

void operator() (RefPart & AMREX_RESTRICT refpart) const
{
using namespace amrex::literals; // for _rt and _prt

// assign input reference particle values
Expand Down Expand Up @@ -220,7 +221,6 @@ namespace impactx
refpart.s = s + slice_ds;
}

private:
amrex::ParticleReal m_ez; //! electric field strength in 1/m
amrex::ParticleReal m_bz; //! magnetic field strength in 1/m
};
Expand Down
51 changes: 43 additions & 8 deletions src/python/elements.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,21 @@ void init_elements(py::module& m)
py::arg("rotation") = 0,
"A short collimator element applying a transverse aperture boundary."
)
.def_property("shape",
[](Aperture & ap) { return ap.m_shape; },
[](Aperture & ap, Aperture::Shape shape) { ap.m_shape = shape; },
"aperture type (rectangular, elliptical)"
)
.def_property("xmax",
[](Aperture & ap) { return ap.m_xmax; },
[](Aperture & ap, amrex::ParticleReal xmax) { ap.m_xmax = xmax; },
"maximum horizontal coordinate"
)
.def_property("ymax",
[](Aperture & ap) { return ap.m_ymax; },
[](Aperture & ap, amrex::ParticleReal ymax) { ap.m_ymax = ymax; },
"maximum vertical coordinate"
)
;
register_beamoptics_push(py_Aperture);

Expand Down Expand Up @@ -233,6 +248,16 @@ void init_elements(py::module& m)
py::arg("nslice") = 1,
"A Quadrupole magnet with chromatic effects included."
)
.def_property("k",
[](ChrQuad & cq) { return cq.m_k; },
[](ChrQuad & cq, amrex::ParticleReal k) { cq.m_k = k; },
"quadrupole strength in 1/m^2 (or T/m)"
)
.def_property("units",
[](ChrQuad & cq) { return cq.m_unit; },
[](ChrQuad & cq, int unit) { cq.m_unit = unit; },
"unit specification for quad strength"
)
;
register_beamoptics_push(py_ChrQuad);

Expand All @@ -256,6 +281,16 @@ void init_elements(py::module& m)
py::arg("nslice") = 1,
"A region of Uniform Acceleration, with chromatic effects included."
)
.def_property("ez",
[](ChrAcc & ca) { return ca.m_ez; },
[](ChrAcc & ca, amrex::ParticleReal ez) { ca.m_ez = ez; },
"electric field strength in 1/m"
)
.def_property("bz",
[](ChrAcc & ca) { return ca.m_bz; },
[](ChrAcc & ca, amrex::ParticleReal bz) { ca.m_bz = bz; },
"magnetic field strength in 1/m"
)
;
register_beamoptics_push(py_ChrAcc);

Expand All @@ -282,19 +317,19 @@ void init_elements(py::module& m)
"A linear Constant Focusing element."
)
.def_property("kx",
[](ConstF & cf) { return cf.m_kx; },
[](ConstF & cf, amrex::ParticleReal kx) { cf.m_kx = kx; },
[](ConstF & cf) { return cf.m_kx; },
[](ConstF & cf, amrex::ParticleReal kx) { cf.m_kx = kx; },
"focusing x strength in 1/m"
)
.def_property("ky",
[](ConstF & cf) { return cf.m_ky; },
[](ConstF & cf, amrex::ParticleReal ky) { cf.m_ky = ky; },
"focusing y strength in 1/m"
[](ConstF & cf) { return cf.m_ky; },
[](ConstF & cf, amrex::ParticleReal ky) { cf.m_ky = ky; },
"focusing y strength in 1/m"
)
.def_property("kt",
[](ConstF & cf) { return cf.m_kt; },
[](ConstF & cf, amrex::ParticleReal kt) { cf.m_kt = kt; },
"focusing t strength in 1/m"
[](ConstF & cf) { return cf.m_kt; },
[](ConstF & cf, amrex::ParticleReal kt) { cf.m_kt = kt; },
"focusing t strength in 1/m"
)
;
register_beamoptics_push(py_ConstF);
Expand Down

0 comments on commit f865a2e

Please sign in to comment.