lnls-fac / trackcpp

Particle tracking code
GNU General Public License v3.0
3 stars 3 forks source link

Return lost position #74

Closed fernandohds564 closed 2 months ago

fernandohds564 commented 2 months ago

I need to have access to the position (rx, px, ry, py) of the lost particles, in my simulations to know exactly at which transverse position the particle was lost. Currently this is not possible with trackcpp, because it only returns the element, the plane and the turn where the particle was lost.

I tried to add this info making minimal changes to the code, but it comes with some disadvantages and I would like to know your opinion on the matter. I added the position of the lost particle in the part_out output of line_pass/ring_pass at the position of the next element/turn. This changes the interpretation of the results of the tracking.

Prior to this PR, this sequence of commands:

import pyaccel as pa
from pymodels import si

mod = si.create_accelerator()
mod.vchamber_on = True

rin = np.array([0.011, 0.1, 0, 0, 0.0, 0])
out1 = pa.tracking.ring_pass(mod, rin, nr_turns=5, turn_by_turn=True)
out2 = pa.tracking.line_pass(mod, rin, indices='closed')

print(out1[0][:, :5], out1[1:])
print(out2[0][:, :5], out2[1:])

would return:

[[0.011   nan   nan   nan   nan]
 [0.1     nan   nan   nan   nan]
 [0.      nan   nan   nan   nan]
 [0.      nan   nan   nan   nan]
 [0.      nan   nan   nan   nan]
 [0.      nan   nan   nan   nan]] (True, 0, 2, 'x')
[[0.011 0.011 0.011   nan   nan]
 [0.1   0.1   0.1     nan   nan]
 [0.    0.    0.      nan   nan]
 [0.    0.    0.      nan   nan]
 [0.    0.    0.      nan   nan]
 [0.    0.    0.      nan   nan]] (True, 2, 'x')

After this PR, we have the following output:

[[0.011      0.20248475        nan        nan        nan]
 [0.1        0.1               nan        nan        nan]
 [0.         0.                nan        nan        nan]
 [0.         0.                nan        nan        nan]
 [0.         0.                nan        nan        nan]
 [0.         0.00957424        nan        nan        nan]] (True, 0, 2, 'x')
[[0.011      0.011      0.011      0.20248475        nan]
 [0.1        0.1        0.1        0.1               nan]
 [0.         0.         0.         0.                nan]
 [0.         0.         0.         0.                nan]
 [0.         0.         0.         0.                nan]
 [0.         0.         0.         0.00957424        nan]] (True, 2, 'x')

Note that ring_pass returns the position of the lost particle as if it were the position of the particle at the beginning of the second turn, even tough the particle was lost at the beginning of the first turn. Additionally, line_pass returns a position clearly larger than the vacuum chamber limits at the entrance of the next element.

Notice that if we call ring_pass with turn_by_turn=False, there will be no nan's in part_out and the most simple way to check if a given particle was lost is by checking if the output lost_plane is equal to no_plane. For an example, check the code below:

rin = np.zeros((6, 2), dtype=float)
rin[:, 1] = np.array([0.011, 0.1, 0, 0, 0.0, 0])
part_out, lost_flag, lost_turn, lost_element, lost_plane = pa.tracking.ring_pass(
    mod, rin, nr_turns=5, turn_by_turn=False
)

lost_part = np.array(lost_plane) != None
print(part_out)
print(lost_part)
[[0.         0.20248475]
 [0.         0.1       ]
 [0.         0.        ]
 [0.         0.        ]
 [0.         0.        ]
 [0.         0.00957424]]
[False  True]

What are your opinions on this, guys? There are other ways of implementing this feature. We could, for instance, add another output to line_pass and ring_pass.

xresende commented 2 months ago

@fernandohds564 , I do NOT foresee problems with how you implemented the fix. But having separate outputs for the two functions, with lost positions, is much more natural to me.

ps: sorry, forgot the negative that I added above now !

fernandohds564 commented 2 months ago

@fernandohds564 , I do foresee problems with how you implemented the fix. But having separate outputs for the two functions, with lost positions, is much more natural to me.

@xresende , the disadvantage of that solution is that the memory used by the output of the tracking would increase a lot. Is it worth it? do you have any other ideas?

xresende commented 2 months ago

@fernandohds564 , I do foresee problems with how you implemented the fix. But having separate outputs for the two functions, with lost positions, is much more natural to me.

@xresende , the disadvantage of that solution is that the memory used by the output of the tracking would increase a lot. Is it worth it? do you have any other ideas?

I dont get it. why would it increase? i will look for you for an offline discussion...

vellosok75 commented 2 months ago

I guess I could get used to this new convention but I also agree that the additional outputs solution seems more natural.

VitorSouzaLNLS commented 2 months ago

Like Matheus and Murilo, I also think that I can get used to this change and the new interpretation of the tracking output. But having two functions with extended "lost_pos" output sounds more natural.

VitorSouzaLNLS commented 2 months ago

I had an idead: with this change being approved and applied to trackcpp, we can deal with the output directly on pyaccel.

The current implementation on pyaccel for line_pass is:

def _line_pass(accelerator, p_in, indices, element_offset, set_seed=False):
    # store only final position?
    args = _trackcpp.LinePassArgs()
    for idx in indices:
        args.indices.push_back(int(idx))
    args.element_offset = int(element_offset)

    n_part = p_in.shape[1]
    p_out = _np.zeros((6, n_part * len(indices)), dtype=float)

    # re-seed pseudo-random generator
    if set_seed:
        _set_random_seed()

    # tracking
    lost_flag = bool(_trackcpp.track_linepass_wrapper(
        accelerator.trackcpp_acc, p_in, p_out, args))

    p_out = p_out.reshape(6, n_part, -1)
    p_out = _np.squeeze(p_out)

    # fills vectors with info about particle loss
    lost_element = list(args.lost_element)
    lost_plane = [LOST_PLANES[lp] for lp in args.lost_plane]

    return p_out, lost_flag, lost_element, lost_plane

Now, considering the approval of the present PR, we can do something like:

def _line_pass(accelerator, p_in, indices, element_offset, set_seed=False):
    # store only final position?
    args = _trackcpp.LinePassArgs()
    for idx in indices:
        args.indices.push_back(int(idx))
    args.element_offset = int(element_offset)

    n_part = p_in.shape[1]
    p_out = _np.zeros((6, n_part * len(indices)), dtype=float)

    # re-seed pseudo-random generator
    if set_seed:
        _set_random_seed()

    # tracking
    lost_flag = bool(_trackcpp.track_linepass_wrapper(
        accelerator.trackcpp_acc, p_in, p_out, args))

    p_out = p_out.reshape(6, n_part, -1)
    p_out = _np.squeeze(p_out)

    # fills vectors with info about particle loss
    lost_element = list(args.lost_element)
    lost_plane = [LOST_PLANES[lp] for lp in args.lost_plane]

    # create lost particles positions array
    lost_pos = np.ones_like(p_out)*np.nan

    # overwrite p_out in case of lost particles
    for i, elem_idx in enumerate(lost_element):
        if elem_idx != -1: # loss condition
            lost_pos[:,i] = p_out[:,i]
            p_out[:,i] *= np.nan

    return p_out, lost_flag, lost_element, lost_plane, lost_pos

I'm not if this implementation is indeed correct, but is just an idea. Also I don't know how much this would impact in performance.

fernandohds564 commented 2 months ago

Hi @VitorSouzaLNLS , thanks for the suggestions. I had thought of something similar as you suggested, but having different meanings in rout for pyaccel and trackcpp is not a good idea. Talking offline with @xresende, we decided to go with the additional output. I'll implement that today. I'll also change lost_flat to be a list the same size as the number of tracked particles, so that checking for particle loss becomes simpler.

fernandohds564 commented 2 months ago

Guys, I'll close this PR and open another one with the changes we agreed upon.