ocelot-collab / ocelot

OCELOT is a multiphysics simulation toolkit designed for studying FEL and storage ring-based light sources.
GNU General Public License v3.0
85 stars 58 forks source link

Matching/optimization of arbitrary matrix elements and their propagation #147

Open jonasbjorklundsvensson opened 2 years ago

jonasbjorklundsvensson commented 2 years ago

Hi,

I'm trying to set up some lattice matching/optimization, and I have some questions (and, I guess, feature requests) after looking at this for a little while. Sorry for mixing topics a bit, but I feel that they are related.

One of the things I want to do is to make sure that a section of my lattice is second-order achromatic, i.e. have T166 = T266 = 0 at the section end. However, it seems to me that only "standard" Twiss parameters, such as betas, alphas, and first-order dispersions, are available for matching. Would I need to write my own optimization routine for this, like in the more general accelerator optimization tutorial (https://nbviewer.org/github/ocelot-collab/ocelot/blob/master/demos/ipython_tutorials/accelerator_optim.ipynb)?

As a follow-up question/request: it would also be great to have matrix elements available for plotting along with the first-order optics. It is quite informative to view the propagation of, for example, T166 and T566 along with the beta functions. Also, the first-order angle-to-space matrix elements R12 and R34 would be useful to have - we are often dealing with quite divergent beams, which can also jitter quite a bit in angle, so limiting the maximum of such parameters along certain lattice sections would also be very useful.

It would also be nice to have more beam moments available along the lattice, for example I can easily plot the mean beam energy along the accelerator but the energy spread is, to my knowledge, not there.

Best regards, Jonas

st-walker commented 2 years ago

I suppose Sergey knows more about this than I do, but I will try to express my thoughts.

One of the things I want to do is to make sure that a section of my lattice is second-order achromatic, i.e. have T166 = T266 = 0 at the section end. However, it seems to me that only "standard" Twiss parameters, such as betas, alphas, and first-order dispersions, are available for matching. Would I need to write my own optimization routine for this, like in the more general accelerator optimization tutorial (https://nbviewer.org/github/ocelot-collab/ocelot/blob/master/demos/ipython_tutorials/accelerator_optim.ipynb)?

I guess you can write your own yes, but if you look in cell 14 of that notebook, you can for your lattice call this function repeatedly and try to minimise arbitrary elements of this matrix using whatever optimisation routine you would like. Actually this function is quite odd because if you look at its definition at https://github.com/ocelot-collab/ocelot/blob/fc4938a45aff3688b68dc675654461415b457088/ocelot/cpbd/optics.py#L75-L94 then you can see that whilst the function only returns the R matrix, it sets T, R and B on the MagneticLattice instance. So this is how you can get the Ts from it. But if you just do lattice.transfer_maps(energy) then it neatly returns B, R and T. I would prefer the latter solution because its tidier.

As a follow-up question/request: it would also be great to have matrix elements available for plotting along with the first-order optics. It is quite informative to view the propagation of, for example, T166 and T566 along with the beta functions. Also, the first-order angle-to-space matrix elements R12 and R34 would be useful to have - we are often dealing with quite divergent beams, which can also jitter quite a bit in angle, so limiting the maximum of such parameters along certain lattice sections would also be very useful.

Yes I agree. It is not difficult to implement something like this though. Look in the above function to see how to loop the lattice and get the maps and yield the individual maps rather than combining them all into a single net one. If it doesn't already exist then perhaps you can write this and a test and send a pull request. MagneticLattice.transfer_maps should then call this per-element version to prevent code duplication.

It would also be nice to have more beam moments available along the lattice, for example I can easily plot the mean beam energy along the accelerator but the energy spread is, to my knowledge, not there.

The solution to this is to attach CopyBeam or SaveBeam physics processes to markers along the beamline. Then at each marker your particle array will be saved either in memory or written to file (depending on which of the two processes you use). Personally I prefer CopyBeam, I think it is tidier. Look in cpbd.magnetic_lattice for helpful functions for inserting markers into lattices. Then just read back the copied/saved ParticleArray instance, remembering that it also stores its location in s, and plot parray.s versus parray.p().std(), for example.

jonasbjorklundsvensson commented 2 years ago

Hi,

Thanks for your thoughts, @st-walker ! Sorry for the late reply, it took me some time to manage to figure things out.

I tried the suggested route of implementing my own routine using scipy.optimize and defining a minimization function wherein I update the transfer maps. After some headaches I now (soon-ish?) have a relatively general function that seems quite robust and solves the (admittedly relatively simple) problems I throw at it within a few seconds. I can easily extend this to also include Twiss functions, beam distributions, etc. So far so good!

(small note - the "cleaner" format you suggested for extracting the R and T matrices was a quite new addition, just had to update my local files for it to run. Might I ask what's in the B-matrix?)

As for the matrix element evolution, I hope to get to that eventually... It seems a bit inefficient to me to loop through the lattice again just to get this information, and it might be more efficient to do it as part of the current/usual matrix multiplication - at least if the option to do is called. I don't really know much about proper Python programming or actually using Github, but I can give it a shot. Worst case, I guess I could send one of you developers my code?

Similar opinion on using CopyBeam or SaveBeam - I already do this in several locations to get the full files for detailed analysis and plotting, but it seems inefficient and potentially slow to do this hundreds of times. I haven't looked it up, but for example the rms bunch length is directly available in the tws_out one can get from the track() function. What I mean is to just include for example the energy spread along the propagation there as well, as tws.EE or whatever.

Have a nice weekend!

st-walker commented 2 years ago

Sorry for late reply, I was at a conference and then on holiday.

I'm not sure what the B matrix is actually.

If you want to send any code then that is fine. I am not sure what you have in mind regarding the matrix multiplication or an alternative.

OK I see what you mean, dumping the whole beam is too much, maybe you want just a single aggregate property of the beam. Yes you can do that, just define your own physics procwss by inheriting PhysicsProc and write something else to file. E.g. parray.p().std() to get the energy spread, and write or save that instead of the whole ParticleArray instance. So look at CopyBeam for example, you could adapt it and do something very simple like this:

class CopyEnergySpread(PhysProc):
    def __init__(self, name: str = ""):
        super().__init__()
        self.name = name
        self.espread = None

    def apply(self, parray: ParticleArray, dz: float) -> None:
        self.espread = parray.p().std()

and then add the physics process to one or more Marker instances just like with Copy/SaveBeam

sergey-tomin commented 2 years ago

Hi, B-matrix it is zero order matrix. It can be used if lattice has misalignments, for instance.

twiss parameters which is returned from track function has many parameters like first and second moments. I just checked and it seems it does not have only tws.pp :) I will implement it in next commit.

Cheers, Sergey.

jonasbjorklundsvensson commented 1 year ago

Hi @st-walker ,

So, I finally got back into this topic. Regarding the optimization, I ended up writing my own, generalized optimizer based on the beam-based optimization in the example notebook. It's not perfect but it can take any input I want.

On the topic of calculating the propagation of the matrix elements, I looked into magnetic_lattice.py and saw that the functionality is more or less there already. I added a small copy of the existing code to store the resulting transfer matrix at each point in the loop over the sequence, like this:

def transfer_maps(self, energy, calc_prop=False):
        """
        Function calculates transfer maps, the first and second orders (R, T), for the whole lattice.

        :param energy: the initial electron beam energy [GeV]
        :return: B, R, T - matrices
        """
        Ra = np.eye(6)
        Ta = np.zeros((6, 6, 6))
        Ba = np.zeros((6, 1))
        E = energy
        for elem in self.sequence:
            for Rb, Bb, Tb, tm in zip(elem.R(E), elem.B(E), elem.T(E), elem.tms):
                Ba, Ra, Ta = transfer_maps_mult(Ba, Ra, Ta, Bb, Rb, Tb)
                E += tm.get_delta_e()
                print(Ra[0,0])

        if calc_prop:
            l_seq = len( self.sequence )

            Ras = np.zeros( (l_seq, 6, 6) )
            Ras[:,:,:] = np.eye(6) # generate l_seq eye-matrices
            Tas = np.zeros( (l_seq, 6, 6, 6) )
            Bas = np.zeros( (l_seq, 6, 1) )
            for elem in self.sequence:
                i = self.sequence.index( elem )

                for Rb, Bb, Tb, tm in zip(elem.R(E), elem.B(E), elem.T(E), elem.tms):
                    Bas[i], Ras[i], Tas[i] = transfer_maps_mult(Bas[i-1], Ras[i-1], Tas[i-1], Bb, Rb, Tb)
                    E += tm.get_delta_e()
                    print(Ras[i,0,0])

        return Ba, Ra, Ta, Bas, Ras, Tas

It's definitely not a "final" implementation, but it does almost what I want. I'd imagine that one could keep only the last loop (the one I'm trying to add) and then just output for example Ra = Ras[-1, :, :] in parallel with the propagation stuff - when timing there is no noticeable difference between the two calculations (as I'd expect).

The print lines in the above code are for diagnostic purposes, because the code doesn't work entirely correctly yet. I have a small lattice with some quads followed by a dipole and a screen, making up an imaging spectrometer. Ra[0,0] and Ras[i,0,0] are identical for each step in the sequence loop up until the dipole, after which they differ (also true for the rest of the matrix elements).

Can you tell what might be going wrong here? Best regards, Jonas

st-walker commented 1 year ago

Although it might not be the cause of your problem this is definitely wrong:

i = self.sequence.index( elem )

list.index checks for equality. If you reuse elements, and then look back to get the previous element's transfer map, you will in general get the wrong map. One solution:

for i, (Rb, Bb, Tb, tm) in enumerate(zip(elem.R(E), elem.B(E), elem.T(E), elem.tms)):
    ...

If it's not what's causing your bug (though I think it probably is), then you will have to send me something to run, because it's difficult to debug purely by eye.

jonasbjorklundsvensson commented 1 year ago

Hi,

The suggested fix produces something which is completely different whereas it was almost correct before, up until the dipole in the lattice.

I'll send you a run file by email!

st-walker commented 1 year ago

OK yes my mistake, I meant something like this:

            for i, elem in enumerate(self.sequence):

but my brain failed and I introduced the enumerate into the inner loop for some reason. That will indeed be a bug for arbitrary lattices, but in your case it doesn't make a difference because you have no duplicated elements.

Anyway it's tricky to debug, my advice would be to try to create the smallest/simplest lattice that recreates this problem and then try to debug, and I can try to help you.

jonasbjorklundsvensson commented 1 year ago

Hi,

I tried making a very short lattice, it still seems something weird happens at the dipole but before then it's all ok with the code I posted above. The code I use is pasted below.

from ocelot import *
from ocelot import Drift, SBend, Quadrupole, Marker
from scipy.constants import pi

#%% define lattice

l_q = 0.2 # (m)
l_dip = 1.085 # (m)
angle_dip = 7 /180 * pi # (rad)
e1_dip = 2.5 / 180 * pi # (rad)
e2_dip = 9.5 / 180 * pi # (rad)

## drifts
l1 = Drift( l=1, eid='L1' ) # cell exit to chamber center
l2 = Drift( l=1, eid='L2' ) # chamber center to chamber exit
l3 = Drift( l=1, eid='L3' )
l4 = Drift( l=1, eid='L4' )

## dipoles
d1 = SBend( l=l_dip, angle=angle_dip, tilt=pi/2, e1=e1_dip, e2=e2_dip, gap=0.03*2, fint=1.13235, eid='D1' ) 
# set quads
q1 = Quadrupole( l=l_q, k1=10, eid='Q1' )

## markers
startpoint = Marker(eid='STARTPOINT')
endpoint = Marker(eid='ENDPOINT')

sequence = (startpoint, l1, q1, l2, d1, l3, endpoint)

## specify tracking method 
method = {'global': SecondTM, 'SBend': SecondTM}

## generate magnetic lattice
start_id = 'STARTPOINT' # ID of the starting element
stop_id = 'ENDPOINT' # ID of the starting element
start_elem = [elem for elem in sequence if elem.id == start_id][0] # returns the element in sequence if the ID matches start_id 
stop_elem = [elem for elem in sequence if elem.id == stop_id][0] # returns the element in sequence if the ID matches start_id 
# start_elem = None # None tracks from sequence start
# stop_elem = None # None tracks to sequence end

## extract start/stop elements for use later
if start_elem is None: # find element
    start_elem = sequence[0]
if stop_elem is None: # find element
    stop_elem = sequence[-1]

lattice = MagneticLattice( sequence, start=start_elem, stop=stop_elem, method=method ) # start/stop = None uses sequence ends

#%% calculate matrix elements

E = 880e-3
B, R, T, Bs, Rs, Ts = lattice.transfer_maps( E, calc_prop=True )
For this code, the "normal" calculation of R11 at each lattice position and my implementation yield, respectively "normal" Jonas' hack Position
1.0 1.0 STARTPOINT
1.0 1.0 L1
0.8065784098850755 0.8065784098850755 Q1
-1.0627296672045818 -1.0627296672045818 L2
-1.0627296672045818 -1.0627296672045818 D1
-3.0919222396607844 -3.09092893084686 D1
-3.0919222396607844 -1.0627296672045818 D1
-4.964809355406868 -2.9129915350705753 L3
-4.964809355406868 -2.9129915350705753 ENDPOINT

Turning dipole edge angles and fint off did not solve it, but it's clearly related to the dipole. The third entry for the dipole is in my case the same as the first one, rather than the previous one.

Hope this help clarify the problem. Best regards, Jonas

jonasbjorklundsvensson commented 1 year ago

Hi @st-walker,

Did you have time to check the code above?

Best regards, Jonas

st-walker commented 1 year ago

Yes I had a look at it but it was not immediately clear to me why it is happening. I was stepping through the code line by line with a debugger, I spent maybe about half an hour on this. Did you try stepping through it with a debugger? I think the example is still too complicated. Does the bug exist with one drift and one dipole for example? If you can make the bug go away then this will help, if not, then at least your test case is simpler so that when you step through it with a debugger then it will be much easier for you to fix it.

jonasbjorklundsvensson commented 1 year ago

Hi,

My Spyder refuses to recognize the breakpoints, so I haven't managed to step through it the way I would like. I can maybe try again, or with a different method, once I get more time to return to these simulations.

Regardless, the bug looks exactly the same for the above example as for my originally more complicated setup, i.e. the discrepancy starts inside the dipole.

PauGC commented 1 year ago

Hi Jonas, Stuart, all

I was looking at the code you were discussing in this thread.

First of all, I think that it would make sense to extend the function "lattice_transfer_map()" with an additional kwarg and not to implement a new method in the "MagneticLattice" class.

I would propose the following, which I think is quite simple, since it requires few lines of code:

def lattice_transfer_map(lattice, energy, output_at_each_step: bool = False):     """     Function calculates transfer maps, the first and second orders (R, T), for the whole lattice.     Second order matrices are attached to lattice object:     lattice.T_sym - symmetric second order matrix     lattice.T - second order matrix     lattice.R - linear R matrix

    :param lattice: MagneticLattice     :param energy: the initial electron beam energy [GeV]     :param output_at_each_step: option to calculate the transfer matrix elements along the beamline     :return: R - matrix     """

    Ba = np.zeros((6, 1))     Ra = np.eye(6)     Ta = np.zeros((6, 6, 6))     Bs, Rs, Ts = [Ba], [Ra], [Ta]     E = energy     for elem in lattice.sequence:         for Rb, Bb, Tb, tm in zip(elem.R(E), elem.B(E), elem.T(E), elem.tms):             Ba, Ra, Ta = transfer_maps_mult(Ba, Ra, Ta, Bb, Rb, Tb)             #Ba = np.dot(Rb, Ba) + Bb             E += tm.get_delta_e()         Bs.append(Ba)         Rs.append(Ra)         Ts.append(Ta)

    # TODO: Adding Attributes at runtime should be avoided     lattice.E = E     lattice.T_sym = Ta     lattice.T = Ta #unsym_matrix(deepcopy(Ta))     lattice.R = Ra     lattice.B = Ba     if output_at_each_step:         return Bs, Rs, Ts     return Ra

It actually seems to work properly if I compare it with the dispersion calculated with the "twiss()" function and also seems to agree with the "normal" results of the test lattice that Jonas proposed.

It's a bit weird that the kind of output changes so much if "output_at_each_step" is selected (i.e. three lists instead of only the total Ra matrix). But maybe this is better than requiring additional kwargs to select the function's output...

What do you think?

Cheers, Pau

On 10/26/22 20:08, jonasbjorklundsvensson wrote:

Hi,

My Spyder refuses to recognize the breakpoints, so I haven't managed to step through it the way I would like. I can maybe try again, or with a different method, once I get more time to return to these simulations.

Regardless, the bug looks exactly the same for the above example as for my originally more complicated setup, i.e. the discrepancy starts inside the dipole.

  • Jonas

— Reply to this email directly, view it on GitHub https://github.com/ocelot-collab/ocelot/issues/147#issuecomment-1292415325, or unsubscribe https://github.com/notifications/unsubscribe-auth/AIC5DCAOC4IYWBBZL66JLLLWFFXSHANCNFSM56ZBXJFA. You are receiving this because you are subscribed to this thread.Message ID: @.***>

-- Pau Gonzalez Caminal FLASHForward - Deutsches Elektronen Synchrotron (DESY) Notkestr. 85, 22607 Hamburg Bldg. 1E / Room 02.514 Phone: +49 40 8998 4960 @.***

PauGC commented 1 year ago

Hi again,

just one more comment. @Stuart: in your first answer to Jonas, you sent a link to a particular line of code. Might it be that the commit you were referring to was detached from the "official" branches? If I look at master or dev_2021 I don't see that the method MagneticLattice.transfer_maps() exists...

Am I wrong? Or is there a pull request on its way that I am not aware of?

Cheers, Pau

On 11/8/22 14:40, Pau Gonzalez Caminal wrote:

Hi Jonas, Stuart, all

I was looking at the code you were discussing in this thread.

First of all, I think that it would make sense to extend the function "lattice_transfer_map()" with an additional kwarg and not to implement a new method in the "MagneticLattice" class.

I would propose the following, which I think is quite simple, since it requires few lines of code:

def lattice_transfer_map(lattice, energy, output_at_each_step: bool = False):     """     Function calculates transfer maps, the first and second orders (R, T), for the whole lattice.     Second order matrices are attached to lattice object:     lattice.T_sym - symmetric second order matrix     lattice.T - second order matrix     lattice.R - linear R matrix

    :param lattice: MagneticLattice     :param energy: the initial electron beam energy [GeV]     :param output_at_each_step: option to calculate the transfer matrix elements along the beamline     :return: R - matrix     """

    Ba = np.zeros((6, 1))     Ra = np.eye(6)     Ta = np.zeros((6, 6, 6))     Bs, Rs, Ts = [Ba], [Ra], [Ta]     E = energy     for elem in lattice.sequence:         for Rb, Bb, Tb, tm in zip(elem.R(E), elem.B(E), elem.T(E), elem.tms):             Ba, Ra, Ta = transfer_maps_mult(Ba, Ra, Ta, Bb, Rb, Tb)             #Ba = np.dot(Rb, Ba) + Bb             E += tm.get_delta_e()         Bs.append(Ba)         Rs.append(Ra)         Ts.append(Ta)

    # TODO: Adding Attributes at runtime should be avoided     lattice.E = E     lattice.T_sym = Ta     lattice.T = Ta #unsym_matrix(deepcopy(Ta))     lattice.R = Ra     lattice.B = Ba     if output_at_each_step:         return Bs, Rs, Ts     return Ra

It actually seems to work properly if I compare it with the dispersion calculated with the "twiss()" function and also seems to agree with the "normal" results of the test lattice that Jonas proposed.

It's a bit weird that the kind of output changes so much if "output_at_each_step" is selected (i.e. three lists instead of only the total Ra matrix). But maybe this is better than requiring additional kwargs to select the function's output...

What do you think?

Cheers, Pau

On 10/26/22 20:08, jonasbjorklundsvensson wrote:

Hi,

My Spyder refuses to recognize the breakpoints, so I haven't managed to step through it the way I would like. I can maybe try again, or with a different method, once I get more time to return to these simulations.

Regardless, the bug looks exactly the same for the above example as for my originally more complicated setup, i.e. the discrepancy starts inside the dipole.

  • Jonas

— Reply to this email directly, view it on GitHub https://github.com/ocelot-collab/ocelot/issues/147#issuecomment-1292415325, or unsubscribe https://github.com/notifications/unsubscribe-auth/AIC5DCAOC4IYWBBZL66JLLLWFFXSHANCNFSM56ZBXJFA. You are receiving this because you are subscribed to this thread.Message ID: @.***>

-- Pau Gonzalez Caminal FLASHForward - Deutsches Elektronen Synchrotron (DESY) Notkestr. 85, 22607 Hamburg Bldg. 1E / Room 02.514 Phone: +49 40 8998 4960 @.***

-- Pau Gonzalez Caminal FLASHForward - Deutsches Elektronen Synchrotron (DESY) Notkestr. 85, 22607 Hamburg Bldg. 1E / Room 02.514 Phone: +49 40 8998 4960 @.***

sergey-tomin commented 1 year ago

Hi Pau,

it is in the "dev" branch. "dev_2021" is merged with "dev" and not in development any more. I will remove it soon. "dev" and "master" will be also merged in about 1-2 months. As for your suggestion. It's really easy to implement, and it hasn't been implemented because we never needed it. If it can be useful to anyone, I will do it in the next few days.

Cheers, Sergey.

PauGC commented 1 year ago

Hi Sergey,

Ups, I see. Thanks for the clarification! I didn't check the dates on the history and, for some reason, I mistakenly thought that "dev_2021" was the most up-to-date branch (despite of its name...).

Cheers, Pau

On 11/8/22 15:44, sergey-tomin wrote:

Hi Pau,

it is in the "dev" branch. "dev_2021" is merged with "dev" and not in development any more. I will remove it soon. "dev" and "master" will be also merged in about 1-2 months. As for your suggestion. It's really easy to implement, and it hasn't been implemented because we never needed it. If it can be useful to anyone, I will do it in the next few days.

Cheers, Sergey.

— Reply to this email directly, view it on GitHub https://github.com/ocelot-collab/ocelot/issues/147#issuecomment-1307330247, or unsubscribe https://github.com/notifications/unsubscribe-auth/AIC5DCDJQFGM76IHKODF2KDWHJRMHANCNFSM56ZBXJFA. You are receiving this because you commented.Message ID: @.***>

-- Pau Gonzalez Caminal FLASHForward - Deutsches Elektronen Synchrotron (DESY) Notkestr. 85, 22607 Hamburg Bldg. 1E / Room 02.514 Phone: +49 40 8998 4960 @.***

sergey-tomin commented 1 year ago

Added output_at_each_step to MagneticLattice.transfer_maps(). see in dev branch https://github.com/ocelot-collab/ocelot/commit/a9edf7e90708978d226fbd310f8b4d3f29111404

PauGC commented 1 year ago

Hi guys,

I needed to do something similar to what Jonas was asking for in the beginning of this thread and I worked out a possible way to do it by enabling the inclusion of matrix elements as a constraint in the "match" function. The way to define the constraints would be, for instance:

constraints = {'global': {'beta_x': ["<", 300], 'beta_y': ["<", 300]},                          {obj: {'R11': -10, 'R12': 0}}

I tested it with some use cases for our beamline and it works fine for me.

One of the problems is that the calculation of the transfer maps requires to pass the energy of the beam, so I included it as an extra argument in the "match" function. What do you think?

Here the link to the code: https://github.com/PauGC/ocelot/blob/0ba27db1e106d219b0d949afe6038510555a7a0f/ocelot/cpbd/match.py#L153

Best, Pau

On 11/9/22 12:55, sergey-tomin wrote:

Added output_at_each_step to MagneticLattice.transfer_maps(). see in dev branch a9edf7e https://github.com/ocelot-collab/ocelot/commit/a9edf7e90708978d226fbd310f8b4d3f29111404

— Reply to this email directly, view it on GitHub https://github.com/ocelot-collab/ocelot/issues/147#issuecomment-1308639144, or unsubscribe https://github.com/notifications/unsubscribe-auth/AIC5DCEYPLR7GE7OSATZO73WHOGM7ANCNFSM56ZBXJFA. You are receiving this because you commented.Message ID: @.***>

-- Pau Gonzalez Caminal FLASHForward - Deutsches Elektronen Synchrotron (DESY) Notkestr. 85, 22607 Hamburg Bldg. 1E / Room 02.514 Phone: +49 40 8998 4960 @.***

PauGC commented 1 year ago

Sorry, I forgot to push after adding the "energy" as an argument... Here the new link: https://github.com/PauGC/ocelot/blob/f6dc2e9334369475ee18babaf26bb9cc6c46dac7/ocelot/cpbd/match.py#L153

On 11/11/22 13:04, Pau Gonzalez Caminal wrote:

Hi guys,

I needed to do something similar to what Jonas was asking for in the beginning of this thread and I worked out a possible way to do it by enabling the inclusion of matrix elements as a constraint in the "match" function. The way to define the constraints would be, for instance:

constraints = {'global': {'beta_x': ["<", 300], 'beta_y': ["<", 300]},                          {obj: {'R11': -10, 'R12': 0}}

I tested it with some use cases for our beamline and it works fine for me.

One of the problems is that the calculation of the transfer maps requires to pass the energy of the beam, so I included it as an extra argument in the "match" function. What do you think?

Here the link to the code: https://github.com/PauGC/ocelot/blob/0ba27db1e106d219b0d949afe6038510555a7a0f/ocelot/cpbd/match.py#L153

Best, Pau

On 11/9/22 12:55, sergey-tomin wrote:

Added output_at_each_step to MagneticLattice.transfer_maps(). see in dev branch a9edf7e https://github.com/ocelot-collab/ocelot/commit/a9edf7e90708978d226fbd310f8b4d3f29111404

— Reply to this email directly, view it on GitHub https://github.com/ocelot-collab/ocelot/issues/147#issuecomment-1308639144, or unsubscribe https://github.com/notifications/unsubscribe-auth/AIC5DCEYPLR7GE7OSATZO73WHOGM7ANCNFSM56ZBXJFA. You are receiving this because you commented.Message ID: @.***>

-- Pau Gonzalez Caminal FLASHForward - Deutsches Elektronen Synchrotron (DESY) Notkestr. 85, 22607 Hamburg Bldg. 1E / Room 02.514 Phone: +49 40 8998 4960 @.***

-- Pau Gonzalez Caminal FLASHForward - Deutsches Elektronen Synchrotron (DESY) Notkestr. 85, 22607 Hamburg Bldg. 1E / Room 02.514 Phone: +49 40 8998 4960 @.***