diff --git a/src/particles/elements/mixin/thick.H b/src/particles/elements/mixin/thick.H index aa04f5ba4..4572d2a67 100644 --- a/src/particles/elements/mixin/thick.H +++ b/src/particles/elements/mixin/thick.H @@ -27,7 +27,10 @@ namespace impactx::elements * @param ds Segment length in m * @param nslice number of slices used for the application of space charge */ - Thick(amrex::ParticleReal const ds, int const nslice ) + Thick ( + amrex::ParticleReal ds, + int nslice + ) : m_ds(ds), m_nslice(nslice) { } @@ -52,7 +55,6 @@ namespace impactx::elements return m_ds; } - protected: amrex::ParticleReal m_ds; //! segment length in m int m_nslice; //! number of slices used for the application of space charge }; diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 03012e6c5..865b21466 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -92,22 +92,61 @@ void init_elements(py::module& m) py::class_(me, "Thick") .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const + amrex::ParticleReal, + amrex::ParticleReal >(), - py::arg("ds"), py::arg("nslice") = 1, + py::arg("ds"), + py::arg("nslice") = 1, "Mixin class for lattice elements with finite length." ) - .def_property_readonly("nslice", &elements::Thick::nslice) - .def_property_readonly("ds", &elements::Thick::ds) + .def_property("ds", + [](elements::Thick & th) { return th.m_ds; }, + [](elements::Thick & th, amrex::ParticleReal ds) { th.m_ds = ds; }, + "segment length in m" + ) + .def_property("nslice", + [](elements::Thick & th) { return th.m_nslice; }, + [](elements::Thick & th, int nslice) { th.m_nslice = nslice; }, + "number of slices used for the application of space charge" + ) ; py::class_(me, "Thin") .def(py::init<>(), "Mixin class for lattice elements with zero length." ) - .def_property_readonly("nslice", &elements::Thin::nslice) - .def_property_readonly("ds", &elements::Thin::ds) + .def_property_readonly("ds", + &elements::Thin::ds, + "segment length in m" + ) + .def_property_readonly("nslice", + &elements::Thin::nslice, + "number of slices used for the application of space charge" + ) + ; + + py::class_(me, "Alignment") + .def(py::init<>(), + "Mixin class for lattice elements with horizontal/vertical alignment errors." + ) + .def_property("dx", + [](elements::Alignment & a) { return a.dx(); }, + [](elements::Alignment & a, amrex::ParticleReal dx) { a.m_dx = dx; }, + "horizontal translation error in m" + ) + .def_property("dy", + [](elements::Alignment & a) { return a.dy(); }, + [](elements::Alignment & a, amrex::ParticleReal dy) { a.m_dy = dy; }, + "vertical translation error in m" + ) + .def_property("rotation", + [](elements::Alignment & a) { return a.rotation(); }, + [](elements::Alignment & a, amrex::ParticleReal rotation_degree) + { + a.m_rotation = rotation_degree * elements::Alignment::degree2rad; + }, + "rotation error around the s axis (in the x-y plane) in degree" + ) ; // diagnostics @@ -115,7 +154,9 @@ void init_elements(py::module& m) py::class_ py_BeamMonitor(me, "BeamMonitor"); py_BeamMonitor .def(py::init(), - py::arg("name"), py::arg("backend")="default", py::arg("encoding")="g", + py::arg("name"), + py::arg("backend") = "default", + py::arg("encoding") = "g", "This element writes the particle beam out to openPMD data." ) ; @@ -123,12 +164,16 @@ void init_elements(py::module& m) // beam optics - py::class_ py_Aperture(me, "Aperture"); + py::class_ py_Aperture(me, "Aperture"); py_Aperture .def(py::init([]( amrex::ParticleReal xmax, amrex::ParticleReal ymax, - std::string shape) + std::string shape, + amrex::ParticleReal dx, + amrex::ParticleReal dy, + amrex::ParticleReal rotation_degree + ) { if (shape != "rectangular" && shape != "elliptical") throw std::runtime_error("shape must be \"rectangular\" or \"elliptical\""); @@ -136,60 +181,104 @@ void init_elements(py::module& m) Aperture::Shape s = shape == "rectangular" ? Aperture::Shape::rectangular : Aperture::Shape::elliptical; - return new Aperture(xmax, ymax, s); + return new Aperture(xmax, ymax, s, dx, dy, rotation_degree); }), - py::arg("xmax"), py::arg("ymax"), py::arg("shape") = "rectangular", + py::arg("xmax"), + py::arg("ymax"), + py::arg("shape") = "rectangular", + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, "A short collimator element applying a transverse aperture boundary." ) ; register_beamoptics_push(py_Aperture); - py::class_ py_ChrDrift(me, "ChrDrift"); + py::class_ py_ChrDrift(me, "ChrDrift"); py_ChrDrift .def(py::init< - amrex::ParticleReal const, - int const >(), - py::arg("ds"), py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "A Drift with chromatic effects included." ) ; register_beamoptics_push(py_ChrDrift); - py::class_ py_ChrQuad(me, "ChrQuad"); + py::class_ py_ChrQuad(me, "ChrQuad"); py_ChrQuad .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - int const, - int const>(), - py::arg("ds"), py::arg("k"), py::arg("units") = 0, py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + int, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("k"), + py::arg("units") = 0, + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "A Quadrupole magnet with chromatic effects included." ) ; register_beamoptics_push(py_ChrQuad); - py::class_ py_ChrAcc(me, "ChrAcc"); + py::class_ py_ChrAcc(me, "ChrAcc"); py_ChrAcc .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - amrex::ParticleReal const, - int const>(), - py::arg("ds"), py::arg("ez"), py::arg("bz"), py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("ez"), + py::arg("bz"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "A region of Uniform Acceleration, with chromatic effects included." ) ; register_beamoptics_push(py_ChrAcc); - py::class_ py_ConstF(me, "ConstF"); + py::class_ py_ConstF(me, "ConstF"); py_ConstF .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - amrex::ParticleReal const, - amrex::ParticleReal const, - int const >(), - py::arg("ds"), py::arg("kx"), py::arg("ky"), py::arg("kt"), py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("kx"), + py::arg("ky"), + py::arg("kt"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "A linear Constant Focusing element." ) .def_property("kx", @@ -210,60 +299,100 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_ConstF); - py::class_ py_DipEdge(me, "DipEdge"); + py::class_ py_DipEdge(me, "DipEdge"); py_DipEdge .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - amrex::ParticleReal const, - amrex::ParticleReal const>(), - py::arg("psi"), py::arg("rc"), py::arg("g"), py::arg("K2"), + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal + >(), + py::arg("psi"), + py::arg("rc"), + py::arg("g"), + py::arg("K2"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, "Edge focusing associated with bend entry or exit." ) ; register_beamoptics_push(py_DipEdge); - py::class_ py_Drift(me, "Drift"); + py::class_ py_Drift(me, "Drift"); py_Drift .def(py::init< - amrex::ParticleReal const, - int const >(), - py::arg("ds"), py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "A drift." ) ; register_beamoptics_push(py_Drift); - py::class_ py_ExactDrift(me, "ExactDrift"); + py::class_ py_ExactDrift(me, "ExactDrift"); py_ExactDrift .def(py::init< - amrex::ParticleReal const, - int const >(), - py::arg("ds"), py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "A Drift using the exact nonlinear map." ) ; register_beamoptics_push(py_ExactDrift); - py::class_ py_ExactSbend(me, "ExactSbend"); + py::class_ py_ExactSbend(me, "ExactSbend"); py_ExactSbend .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - amrex::ParticleReal const, - int const>(), - py::arg("ds"), py::arg("phi"), py::arg("B") = 0.0, py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("phi"), + py::arg("B") = 0.0, + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "An ideal sector bend using the exact nonlinear map. When B = 0, the reference bending radius is defined by r0 = length / (angle in rad), corresponding to a magnetic field of B = rigidity / r0; otherwise the reference bending radius is defined by r0 = rigidity / B." ) ; register_beamoptics_push(py_ExactSbend); - py::class_ py_Kicker(me, "Kicker"); + py::class_ py_Kicker(me, "Kicker"); py_Kicker .def(py::init([]( amrex::ParticleReal xkick, amrex::ParticleReal ykick, - std::string const & units) + std::string const & units, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal + ) { if (units != "dimensionless" && units != "T-m") throw std::runtime_error(R"(units must be "dimensionless" or "T-m")"); @@ -273,19 +402,33 @@ void init_elements(py::module& m) Kicker::UnitSystem::Tm; return new Kicker(xkick, ykick, u); }), - py::arg("xkick"), py::arg("ykick"), py::arg("units") = "dimensionless", + py::arg("xkick"), + py::arg("ykick"), + py::arg("units") = "dimensionless", + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, R"(A thin transverse kicker element. Kicks are for units "dimensionless" or in "T-m".)" ) ; register_beamoptics_push(py_Kicker); - py::class_ py_Multipole(me, "Multipole"); + py::class_ py_Multipole(me, "Multipole"); py_Multipole .def(py::init< - int const, - amrex::ParticleReal const, - amrex::ParticleReal const>(), - py::arg("multiple"), py::arg("K_normal"), py::arg("K_skew"), + int, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal + >(), + py::arg("multiple"), + py::arg("K_normal"), + py::arg("K_skew"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, "A general thin multipole element." ) ; @@ -299,12 +442,20 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_None); - py::class_ py_NonlinearLens(me, "NonlinearLens"); + py::class_ py_NonlinearLens(me, "NonlinearLens"); py_NonlinearLens .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const>(), - py::arg("knll"), py::arg("cnll"), + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal + >(), + py::arg("knll"), + py::arg("cnll"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, "Single short segment of the nonlinear magnetic insert element." ) ; @@ -313,7 +464,8 @@ void init_elements(py::module& m) py::class_(me, "Programmable", py::dynamic_attr()) .def(py::init< amrex::ParticleReal, - int>(), + int + >(), py::arg("ds") = 0.0, py::arg("nslice") = 1, "A programmable beam optics element." ) @@ -353,19 +505,28 @@ void init_elements(py::module& m) ) ; - py::class_ py_Quad(me, "Quad"); + py::class_ py_Quad(me, "Quad"); py_Quad .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - int const>(), - py::arg("ds"), py::arg("k"), py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("k"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "A Quadrupole magnet." ) ; register_beamoptics_push(py_Quad); - py::class_ py_RFCavity(me, "RFCavity"); + py::class_ py_RFCavity(me, "RFCavity"); py_RFCavity .def(py::init< amrex::ParticleReal, @@ -374,90 +535,155 @@ void init_elements(py::module& m) amrex::ParticleReal, std::vector, std::vector, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, int, int >(), - py::arg("ds"), py::arg("escale"), py::arg("freq"), - py::arg("phase"), py::arg("cos_coefficients"), py::arg("sin_coefficients"), - py::arg("mapsteps") = 1, py::arg("nslice") = 1, + py::arg("ds"), + py::arg("escale"), + py::arg("freq"), + py::arg("phase"), + py::arg("cos_coefficients"), + py::arg("sin_coefficients"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("mapsteps") = 1, + py::arg("nslice") = 1, "An RF cavity." ) ; register_beamoptics_push(py_RFCavity); - py::class_ py_Sbend(me, "Sbend"); + py::class_ py_Sbend(me, "Sbend"); py_Sbend .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - int const>(), - py::arg("ds"), py::arg("rc"), py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("rc"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "An ideal sector bend." ) ; register_beamoptics_push(py_Sbend); - py::class_ py_CFbend(me, "CFbend"); + py::class_ py_CFbend(me, "CFbend"); py_CFbend .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - amrex::ParticleReal const, - int const>(), - py::arg("ds"), py::arg("rc"), py::arg("k"), py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("rc"), + py::arg("k"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "An ideal combined function bend (sector bend with quadrupole component)." ) ; register_beamoptics_push(py_CFbend); - py::class_ py_Buncher(me, "Buncher"); + py::class_ py_Buncher(me, "Buncher"); py_Buncher .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const>(), - py::arg("V"), py::arg("k"), + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal + >(), + py::arg("V"), + py::arg("k"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, "A short linear RF cavity element at zero-crossing for bunching." ) ; register_beamoptics_push(py_Buncher); - py::class_ py_ShortRF(me, "ShortRF"); + py::class_ py_ShortRF(me, "ShortRF"); py_ShortRF .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - amrex::ParticleReal const>(), - py::arg("V"), py::arg("freq"), py::arg("phase") = -90.0, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal + >(), + py::arg("V"), + py::arg("freq"), + py::arg("phase") = -90.0, + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, "A short RF cavity element." ) ; register_beamoptics_push(py_ShortRF); - py::class_ py_SoftSolenoid(me, "SoftSolenoid"); + py::class_ py_SoftSolenoid(me, "SoftSolenoid"); py_SoftSolenoid .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, + amrex::ParticleReal, + amrex::ParticleReal, std::vector, std::vector, - int const, - int const + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int, + int >(), - py::arg("ds"), py::arg("bscale"), - py::arg("cos_coefficients"), py::arg("sin_coefficients"), - py::arg("mapsteps") = 1, py::arg("nslice") = 1, + py::arg("ds"), + py::arg("bscale"), + py::arg("cos_coefficients"), + py::arg("sin_coefficients"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("mapsteps") = 1, + py::arg("nslice") = 1, "A soft-edge solenoid." ) ; register_beamoptics_push(py_SoftSolenoid); - py::class_ py_Sol(me, "Sol"); + py::class_ py_Sol(me, "Sol"); py_Sol .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const, - int const>(), - py::arg("ds"), py::arg("ks"), py::arg("nslice") = 1, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + int + >(), + py::arg("ds"), + py::arg("ks"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("nslice") = 1, "An ideal hard-edge Solenoid magnet." ) ; @@ -466,38 +692,57 @@ void init_elements(py::module& m) py::class_ py_PRot(me, "PRot"); py_PRot .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const>(), - py::arg("phi_in"), py::arg("phi_out"), + amrex::ParticleReal, + amrex::ParticleReal + >(), + py::arg("phi_in"), + py::arg("phi_out"), "An exact pole-face rotation in the x-z plane. Both angles are in degrees." ) ; register_beamoptics_push(py_PRot); - py::class_ py_SoftQuadrupole(me, "SoftQuadrupole"); + py::class_ py_SoftQuadrupole(me, "SoftQuadrupole"); py_SoftQuadrupole .def(py::init< amrex::ParticleReal, amrex::ParticleReal, std::vector, std::vector, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, int, int >(), - py::arg("ds"), py::arg("gscale"), - py::arg("cos_coefficients"), py::arg("sin_coefficients"), - py::arg("mapsteps") = 1, py::arg("nslice") = 1, + py::arg("ds"), + py::arg("gscale"), + py::arg("cos_coefficients"), + py::arg("sin_coefficients"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("mapsteps") = 1, + py::arg("nslice") = 1, "A soft-edge quadrupole." ) ; register_beamoptics_push(py_SoftQuadrupole); - py::class_ py_ThinDipole(me, "ThinDipole"); + py::class_ py_ThinDipole(me, "ThinDipole"); py_ThinDipole .def(py::init< - amrex::ParticleReal const, - amrex::ParticleReal const>(), - py::arg("theta"), py::arg("rc"), + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal + >(), + py::arg("theta"), + py::arg("rc"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, "A thin kick model of a dipole bend." ) ;