Track all refpts in one track function call in acceptance calculations

[swhite2401]

The method introduced by @oscarxblanco in #773 to track all refpts in a single call of the track function is implemented in the acceptance module: all refpts are first projected to the ring start point before the multi-turn tracking.

This allows for full parallelization of the tracking and is optimized for GPU parallelization where large number of processes (i.e. particles) are required and makes this method available for all acceptance calculations (transverse, longitudinal, 1D, 2D).

It could be interesting to tracking several planes in one function call as well but this is not yet implemented as this is most useful for MA / lifetime calculations.

[oscarxblanco]

Dear @swhite2401 ,
I see no change in the verbose output, it seems to go element by element. Should I choose a particular grid mode or something else ?

[swhite2401]

Ah strange... did you include several refpts? The output changes only for multiple refpts

Below the script used for testing:

import at

path = '../lattice/'
filename = 'S28F.mat'
key = 'LOW_EMIT_RING_INJ'
latticef = path+filename
ring = at.load_lattice(latticef,key=key)
ring.disable_6d()

b,s,g = ring.get_momentum_acceptance(1e-2, 0.1, nturns = 512,
                                     grid_mode=at.GridMode.RADIAL,
                                     use_mp=True, verbose=True, refpts=[0, 1, 2])
print(b)

[oscarxblanco]

Dear @swhite2401 ,
the recursive grid does not track multiple reference points. Here is the output for the first two elements.

12 cpu found for acceptance calculation
Multi-process acceptance calculation selected...
The estimated load for grid mode is 0.8333333333333334
The estimated load for recursive mode is 8.0
GridMode.RADIAL or GridMode.CARTESIAN is recommended

Running recursive boundary search:
Element CAVRF, obspt=0
The grid mode is GridMode.RECURSIVE
The planes are ['dp']
Number of angles is 2 from 3.141592653589793 to 0.0 rad
The resolution of the search is 0.001
The initial step size is [0.01]
The initial offset is [ 4.81568793e-05  9.94540489e-09  2.53785930e-05  1.05806381e-19
 -1.55011425e-04 -3.44866135e-02] with dp=None
/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/tracking/track.py:97: AtWarning: no parallel computation for a single particle
  warn(AtWarning('no parallel computation for a single particle'))
Calculation took 4.05587363243103

Running recursive boundary search:
Element LSS, obspt=3
The grid mode is GridMode.RECURSIVE
The planes are ['dp']
Number of angles is 2 from 3.141592653589793 to 0.0 rad
The resolution of the search is 0.001
The initial step size is [0.01]
The initial offset is [ 4.81686083e-05  9.94420361e-09  2.53785930e-05  1.05800248e-19
  1.48841747e-04 -3.44866135e-02] with dp=None
Calculation took 3.767972707748413

With RADIAL grid it seems to track multiple reference points.

[oscarxblanco]

Dear @swhite2401 , there seems to be another problem when the argument offset is passed.
Here is the error output:

Traceback (most recent call last):
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/lattice/lattice_object.py", line 230, in __getitem__
    return super(Lattice, self).__getitem__(key.__index__())
                                            ^^^^^^^^^^^^^
AttributeError: 'slice' object has no attribute '__index__'. Did you mean: '__init__'?

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/oblanco/Documents/public/labs/alba/work/testdomomap_multirefpts/domomap_mod.py", line 257, in <module>
    boundary, tracked, survive = at.get_momentum_acceptance(ring,
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/acceptance.py", line 438, in get_momentum_acceptance
    return get_1d_acceptance(ring, "dp", resolution, amplitude, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/acceptance.py", line 245, in get_1d_acceptance
    b, s, g = get_acceptance(
              ^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/lattice/cavity_access.py", line 48, in wrapper
    return func(ring, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/acceptance.py", line 145, in get_acceptance
    b, s, g = boundary_search(
              ^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/boundary.py", line 534, in boundary_search
    result = grid_boundary_search(
             ^^^^^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/boundary.py", line 336, in grid_boundary_search
    newring[: len(ring) - o].track(parts, use_mp=use_mp, in_place=True, **kwargs)
    ~~~~~~~^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/lattice/lattice_object.py", line 233, in __getitem__
    rg = range(*key.indices(len(self)))
                ^^^^^^^^^^^^^^^^^^^^^^
TypeError: only integer scalar arrays can be converted to a scalar index

[swhite2401]

the recursive grid does not track multiple reference points. Here is the output for the first two elements.

Ah yes correct, this is normal the recursive search is too complicated to implement in this way, I thought I had update the docstrings but apparently not.
I will improve the documentation.

[swhite2401]

Dear @swhite2401 , there seems to be another problem when the argument offset is passed. Here is the error output:

Traceback (most recent call last):
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/lattice/lattice_object.py", line 230, in __getitem__
    return super(Lattice, self).__getitem__(key.__index__())
                                            ^^^^^^^^^^^^^
AttributeError: 'slice' object has no attribute '__index__'. Did you mean: '__init__'?

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/oblanco/Documents/public/labs/alba/work/testdomomap_multirefpts/domomap_mod.py", line 257, in <module>
    boundary, tracked, survive = at.get_momentum_acceptance(ring,
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/acceptance.py", line 438, in get_momentum_acceptance
    return get_1d_acceptance(ring, "dp", resolution, amplitude, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/acceptance.py", line 245, in get_1d_acceptance
    b, s, g = get_acceptance(
              ^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/lattice/cavity_access.py", line 48, in wrapper
    return func(ring, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/acceptance.py", line 145, in get_acceptance
    b, s, g = boundary_search(
              ^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/boundary.py", line 534, in boundary_search
    result = grid_boundary_search(
             ^^^^^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/acceptance/boundary.py", line 336, in grid_boundary_search
    newring[: len(ring) - o].track(parts, use_mp=use_mp, in_place=True, **kwargs)
    ~~~~~~~^^^^^^^^^^^^^^^^^
  File "/home/oblanco/Documents/public/progs/pyenv/pyenvat/lib/python3.11/site-packages/at/lattice/lattice_object.py", line 233, in __getitem__
    rg = range(*key.indices(len(self)))
                ^^^^^^^^^^^^^^^^^^^^^^
TypeError: only integer scalar arrays can be converted to a scalar index

Indeed... I take a look, thanks!

[oscarxblanco]

Dear @swhite2401 , the previous error when using offset has disappeared.
Now I have another problem, when launching a typical case with more than 1000 points and resolution of $10^{-3}$ the terminal closes itself after 20 min of calculations. I monitored the memory the second time it happened and it occurred at about 70% of used memory.
Is there any tool you suggest I could use to log what is happening ?

[swhite2401]

@oscarxblanco, sorry I saw that problem too, I did not think you would be so efficient in looking back at it! I just did this commit to save my ongoing work...

I have used tracemalloc in the past but for at it is a bit tricky because often it just finds that atpass is the problem...
If you want to try it out you are more than welcome otherwise I was planning to get back at it, most probably tomorrow

[oscarxblanco]

@swhite2401 ,
I think it is the same problem I have when trying to parallelize the frequency map track and analysis over the whole grid (PR #756). I have not been able to solve it.

[swhite2401]

@oscarxblanco I think I have found the problem: ring.track() has to be called with in_place=True and refpts=None to avoid saving the particles coordinates turn-by-turn in a buffer.

I ran some test on the EBS lattice with nturns=512, refpts=at.Quadrupole (514) with a resolution of 1.0e-3 and the memeory usage remained stable.

Could you rerun your script to check that it is fine now?

[oscarxblanco]

@swhite2401 , sure, I'll check it now.

[swhite2401]

By the way I started using black for code formatting since you and @lfarv seem to use it now but I really don't the way it puts everything vertically...I think it really makes the code hard to read... what do you think?

[oscarxblanco]

I like black.
It changes the style on which I would normally write the code as I don't longer care to make small line changes but just separate blocks of code that are later organized systematically by black.
When there is more than one person contributing I think is an advantage to keep a standard.
Also, it does not reorganizes the code differently on every call. If you call it twice it won't make changes, which I think is good to keep track on git.
And, flake8 is always mentioning black and isort.

The disadvantage for me is that it does not longer look like code I wrote, so, in few months I wouldn't be able to tell whether I wrote it, copied, or someone else did. And there is extra work after the code is working, I have to iterate many times black, isort, flake8.

[oscarxblanco]

Dear @swhite2401 , the terminal does not longer disappears and the function finishes without errors. I would add a verbose message stating that 'tracking particles' because the messages stop after projecting to the start.

I compared the result with an old calculation done per reference point, and they are similar but not the same. I would have to check, at least redo the calculation for the single refpt.

[lfarv]

Concerning black: I also like the fact that whoever wrote some code, it always looks the same. For alignment, its rule is deliberately simple: either it fits on one line, or it puts on item per line. Every other splitting is a matter of taste, and as such, is eliminated. This makes usually the code longer, but unlike you, I find it easier to read, in particular for argument lists.

[swhite2401]

Ok then I will go with the majority then.

However if you look at the RDT branch I started pyat_computeRDT module physics.rdt you will see that it applies the same rule to formulas and you end up often with one term per line... this is really diffcult to read and debug for me, you still like it?

[swhite2401]

I compared the result with an old calculation done per reference point, and they are similar but not the same. I would have to check, at least redo the calculation for the single refpt.

Ah this is not normal, I did a check with 1.0e-3 resolution and for the EBS lattice it was fine, but maybe I missed something.... let me redo it on my side as well

[oscarxblanco]

I might be using the recursive grid a few months ago to increase a little bit the speed.
The result in single and multirefpts is the same if I use the same RADIAL grid.
All good.

[swhite2401]

Ok great, then I will clean up and let you know for the final review

[swhite2401]

@oscarxblanco , do you want that I implement the same fix in the frequency_map module while I am at it?

[oscarxblanco]

@swhite2401 , for the frequency map I need all turns. Is there a way to limit, or clean the buffer you refer to ?

[swhite2401]

@swhite2401 , for the frequency map I need all turns. Is there a way to limit, or clean the buffer you refer to ?

Ah right... the only way is to split the tracking into chunks of particles.

The size of this buffer is however easy to determine: it is an array of double with shape: (6, nparticles, nrefpts, nturns). You could decide to automatically split the tracking if it exceeds the available memory minus some margin.

[lfarv]

@swhite2401

look at the RDT branch I started pyat_computeRDT module physics.rdt

There you are combining a large indentation leaving few chars on a line, and long arithmetic expressions. And I agree that's not so nice! So it's up to you, there is no intention to make "black" compulsory not to put it in a commit hook. You can completely ignore black for a given file, or disable formatting for a single line ( # fmt: skip) or for a code block (# fmt: off / # fmt: on).

[swhite2401]

disable formatting for a single line ( # fmt: skip) or for a code block (# fmt: off / # fmt: on).

Ah thanks for the tip, I'll do that then

[swhite2401]

I think this one is ready for the final review, I have added a note in the docstrings to explain the behavior, added prints as suggested and blacked all the modified files.

[oscarxblanco]

Tested with grid.RADIAL. All OK.

[swhite2401]

I wait for @simoneliuzzo to take a look before merging