Setting Tracking Methods for different Elements

[cr-xu]

Currently, there's no clear documentation on how to actually set different tracking methods for different elements

The explanation in the MagneticLattice docstring is actually not working: (https://github.com/ocelot-collab/ocelot/blob/1e6cb91202d02b3bdaf6d9b8d2da9adc6662c030/ocelot/cpbd/magnetic_lattice.py#L166C17-L166C17) The tracking methods can only be changed via explicitly providing the class object

See below

import ocelot
from copy import deepcopy

d1 = ocelot.Drift(l=0.1, eid="d1")
d2 = ocelot.Drift(l=0.1, eid="d2")
q1 = ocelot.Quadrupole(l=0.1, eid="q1")
undulator = ocelot.Undulator(lperiod=0.1, nperiods=10, eid="undulator")
elements = [d1, q1, d2, undulator]
# nothing
method_default = {}
# tries to set different transfermap
method_new = {
    "Quadrupole": ocelot.SecondTM,
    "Undulator": ocelot.UndulatorTestTM,
}
# this works
method_explicit = {
    ocelot.cpbd.elements.quadrupole.Quadrupole: ocelot.SecondTM,
    ocelot.cpbd.elements.undulator.Undulator: ocelot.UndulatorTestTM,
}
lattice = ocelot.MagneticLattice(deepcopy(elements), method=method_default)
lattice2 = ocelot.MagneticLattice(deepcopy(elements), method=method_new)
lattice3 = ocelot.MagneticLattice(deepcopy(elements), method=method_explicit)

# Default Tracking Methods
print("Default Tracking Methods")
print("Quad: ", lattice.sequence[1]._tm_class_type, "Undulator: ", lattice.sequence[3]._tm_class_type)
# Expects to see SecondTM, but not working
print("Setting according to Documentation")
print("Quad: ", lattice2.sequence[1]._tm_class_type, "Undulator: ", lattice2.sequence[3]._tm_class_type)
# This works
print("Setting explicitly")
print("Quad: ", lattice3.sequence[1]._tm_class_type, "Undulator: ", lattice3.sequence[3]._tm_class_type)

Actual output:

Default Tracking Methods Quad: <class 'ocelot.cpbd.transformations.transfer_map.TransferMap'> Undulator: <class 'ocelot.cpbd.transformations.transfer_map.TransferMap'> Setting according to Documentation Quad: <class 'ocelot.cpbd.transformations.transfer_map.TransferMap'> Undulator: <class 'ocelot.cpbd.transformations.transfer_map.TransferMap'> Setting explicitly Quad: <class 'ocelot.cpbd.transformations.second_order.SecondTM'> Undulator: <class 'ocelot.cpbd.transformations.undulator_test.UndulatorTestTM'>

[sergey-tomin]

Thanks for reporting this. Docstring is not correct. It should be something like that.

from ocelot import * methods = {"global": SecondTM, Undulator: UndulatorTestTM, Quadruple: SecondTM }

I will fix it.

[sergey-tomin]

fixed https://github.com/ocelot-collab/ocelot/commit/79d7d928620759f793c6c999e847de4abe46ec47