Skip to content

Commit

Permalink
Add Python bindings for LinearMap element.
Browse files Browse the repository at this point in the history
  • Loading branch information
cemitch99 committed Oct 15, 2024
1 parent 817cc51 commit 1ef158e
Showing 1 changed file with 33 additions and 0 deletions.
33 changes: 33 additions & 0 deletions src/python/elements.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1480,6 +1480,39 @@ void init_elements(py::module& m)
;
register_beamoptics_push(py_TaperedPL);

py::class_<LinearMap, elements::Named, elements::Thin, elements::Alignment> py_LinearMap(me, "LinearMap");
py_LinearMap
.def("__repr__",
[](LinearMap const & linearmap) {
return element_name(
linearmap,
std::make_pair("R", linearmap.m_transport_map),
);
}
)
.def(py::init([](
LinearTransport::Map6x6 R,
amrex::ParticleReal dx,
amrex::ParticleReal dy,
amrex::ParticleReal rotation_degree,
std::optional<std::string> name
)
{
}),
py::arg("R"),
py::arg("unit") = "dimensionless",
py::arg("dx") = 0,
py::arg("dy") = 0,
py::arg("rotation") = 0,
py::arg("name") = py::none(),
R"(A user-provided linear map, represented as a 6x6 transport matrix.)"
)
.def_property("R",
[](LinearMap & linearmap) { return linearmap.m_transport_map; },
[](LinearMap & linearmap, amrex::ParticleReal xkick) { linearmap.m_transport_map = linearmap; },
"linear map as a 6x6 transport matrix"
)


// freestanding push function
m.def("push", &Push,
Expand Down

0 comments on commit 1ef158e

Please sign in to comment.