diff --git a/python/src/elements.cc b/python/src/elements.cc index bfb68e8c..fd389de1 100644 --- a/python/src/elements.cc +++ b/python/src/elements.cc @@ -132,30 +132,41 @@ class PyFieldKick: public tse::FieldKick { #endif -template -void add_methods_field_kick(py::class_ t_mapper) +// not using the wrapper but directly the double type: +// TpsaOrComplex also understands of double of gtpsa::tpsa objects are passed +template +void field_kick_add_set_methods(py::class_ t_mapper) { - - using double_type = typename Types::double_type; - t_mapper - .def("set_dx", - [](Class &kick, const double_type dx) + .def("set_dx", [](Class &kick, const double_type& dx) {kick.getTransform()->setDx(dx);}) - .def("set_dy", - [](Class &kick, const double_type dy) + .def("set_dy", [](Class &kick, const double_type& dy) {kick.getTransform()->setDy(dy);}) - .def("set_roll", - [](Class &kick, const double_type roll) - {kick.getTransform()->setRoll(roll);}) + /* + * Todo: seems to crash the interpreter + */ + .def("set_roll", [](Class &kick, const double_type& roll) + {kick.getTransform()->setRoll(roll);}, + py::keep_alive<1, 2>()) + ; +} + +template +void field_kick_add_methods(py::class_ t_mapper) +{ + + //using double_type; + field_kick_add_set_methods(t_mapper); + t_mapper + .def("get_dx", [](Class &kick) {return kick.getTransform()->getDx();}) .def("get_dy", [](Class &kick) {return kick.getTransform()->getDy();}) - // .def("get_roll", - // [](Class &kick){kick.getTransform()->getRoll();}) + .def("get_roll", + [](Class &kick){kick.getTransform()->getRoll();}) .def("is_thick", &Class::isThick) .def("as_thick", @@ -196,12 +207,22 @@ void add_methods_field_kick(py::class_ t_mapper) } +template +void classical_magnet_add_set_methods(py::class_ t_mapper) +{ + t_mapper + .def("set_main_multipole_strength", + [](Class &inst, const complex_type v) + {inst.setMainMultipoleStrength(v);}) + ; +} template -void add_methods_classical_magnet(py::class_ t_mapper) +void classical_magnet_add_methods(py::class_ t_mapper) { - using double_type = typename Types::double_type; + // using double_type = typename Types::double_type; using complex_type = typename Types::complex_type; + classical_magnet_add_set_methods(t_mapper); t_mapper .def("get_multipoles", &Class::getMultipoles) @@ -221,9 +242,6 @@ void add_methods_classical_magnet(py::class_ t_mapper) // .def("set_main_multipole_strength", // [](Class &inst, const double_type v) // {inst.setMainMultipoleStrength(v);}) - .def("set_main_multipole_strength", - [](Class &inst, const complex_type v) - {inst.setMainMultipoleStrength(v);}) .def("propagate", py::overload_cast&>(&Class::propagate), pass_d_doc) @@ -255,9 +273,11 @@ struct TemplatedClasses // auto buildClasses // (py::class_>& cell_void) + typedef tse::FieldKickKnobbed FieldKickK; + auto buildClasses - (py::class_>& elem_type){ + (py::class_>& elem_type, + py::class_>& field_kick){ typedef tse::DriftTypeWithKnob DriftTypeK; std::string drift_type_name = "Drift" + this->m_suffix; @@ -266,11 +286,7 @@ struct TemplatedClasses drift .def(py::init()); - typedef tse::FieldKickKnobbed FieldKickK; - std::string field_kick_name = "FieldKick" + this->m_suffix; - py::class_> field_kick - (this->m_module, field_kick_name.c_str(), elem_type); - add_methods_field_kick(field_kick); + field_kick_add_methods(field_kick); field_kick /* .def("getTransform", @@ -294,8 +310,7 @@ struct TemplatedClasses PyClassicalMagnet, std::shared_ptr> > cm(this->m_module, cm_name.c_str(), mpole_type); - add_methods_classical_magnet - >(cm); + classical_magnet_add_methods>(cm); cm .def(py::init()); @@ -340,7 +355,7 @@ struct TemplatedClasses .def(py::init()); - return elem_type; + return cm; } }; @@ -390,16 +405,31 @@ void py_thor_scsi_init_elements(py::module &m) // (&tse::ElemType::propagate), pass_tpsa_doc) ; + typedef tse::FieldKickKnobbed FieldKickKD; + py::class_> field_kick + (m, "FieldKick", elem_type); + + typedef tse::FieldKickKnobbed FieldKickKT; + py::class_> field_kick_tpsa + (m, "FieldKickTpsa", elem_type); + + // make it directly callable by doubles + field_kick_add_set_methods(field_kick_tpsa); + // and gtpsa::tpsa objects + field_kick_add_set_methods(field_kick_tpsa); TemplatedClasses templated_classes_std(m, ""); // required as marker and bpm are not knobbed yet - // auto elem_type = - templated_classes_std.buildClasses(elem_type); + auto classical_magnet = templated_classes_std.buildClasses(elem_type, field_kick); // Device classes / types with knobs ... handled by these two lines TemplatedClasses templated_classes_tpsa (m, "Tpsa"); - templated_classes_tpsa.buildClasses(elem_type); + auto classical_magnet_tpsa = templated_classes_tpsa.buildClasses(elem_type, field_kick_tpsa); + classical_magnet_add_set_methods> (classical_magnet_tpsa); + classical_magnet_add_set_methods> (classical_magnet_tpsa); + + // classes without knobs follow diff --git a/python/tests/knobbable_test.py b/python/tests/knobbable_test.py index ad88eecf..64fab0c5 100644 --- a/python/tests/knobbable_test.py +++ b/python/tests/knobbable_test.py @@ -155,3 +155,61 @@ def test_quadrupole_dependence_on_offset(): ps.set_zero() quad.propagate(calc_config, ps) print(ps.x, ps.px) + + +def test_knobbable_dx_set_from_float(): + C = Config() + C.setAny("L", .5) + C.setAny("name", "q1") + C.setAny("K", 1.2) + C.setAny("N", 4) + quad = tslib.QuadrupoleTpsa(C) + + + quad.set_dx(1e-3) + quad.set_dy(.5e-3) + quad.set_roll(5e-4) + + desc = gtpsa.desc(6, 2, 3, 1) + keys = dict(x=0, px=1, y=2, py=3, delta=4, ct=5, dx=6,dy=7,roll=8) + mapping = gtpsa.IndexMapping(keys) + dx = gtpsa.tpsa(desc, 1, mapping=mapping) + dx.set_knob(1e-3, "dx") + dy = gtpsa.tpsa(desc, 1, mapping=mapping) + dy.set_knob(1e-3, "dy") + roll = gtpsa.tpsa(desc, 1, mapping=mapping) + roll.set_knob(1e-3, "roll") + quad.set_dx(dx) + quad.set_dy(dy) + + # roll crashes the process ... + rollk = gtpsa.TpsaOrDouble(roll) + quad.set_roll(rollk) + + print(type(quad.get_dx())) + + +def test_knobbalbe_muls_set_from_float(): + from thor_scsi.lib import TwoDimensionalMultipolesTpsa as Mul2DTpsa + m = Mul2DTpsa(0) + m.set_multipole(3, 1e-3+1e-2j) + + +def test_knobbalbe_muls_set_from_tpsa(): + from thor_scsi.lib import TwoDimensionalMultipolesTpsa as Mul2DTpsa + import gtpsa + + desc = gtpsa.desc(6, 2, 1, 2) + c3 = gtpsa.ctpsa(desc, 1) + c3.set_knob(1e-3+1e-2j, 6) + m = Mul2DTpsa(0) + m.set_multipole(3, c3) + + c3_check = m.get_multipole(3) + + +if __name__ == "__main__": + test_knobbable_dx_set_from_float() + + # test_knobbalbe_muls_set_from_float() + # test_knobbalbe_muls_set_from_tpsa() diff --git a/src/thor_scsi/elements/quadrupole.h b/src/thor_scsi/elements/quadrupole.h index 5a13c24d..3f47e033 100644 --- a/src/thor_scsi/elements/quadrupole.h +++ b/src/thor_scsi/elements/quadrupole.h @@ -14,14 +14,6 @@ namespace thor_scsi::elements { this->setMainMultipoleStrength(config); } - -# // J.B. 14/07/23: test. - inline QuadrupoleTypeWithKnob* Q_init(const Config &config) - { - std::cout << "\nQ_init()\n"; - return this->setMainMultipoleStrength(config); - } - inline int getMainMultipoleNumber(void) const override final { return 2; };