lnls-fac / pyaccel

Python module for beam dynamics tracking and optics calculations
MIT License
7 stars 5 forks source link

Dev/tracking/feature/return lost pos #142

Closed fernandohds564 closed 2 months ago

fernandohds564 commented 3 months ago

Description

This PR is the companion of lnls-fac/trackcpp#75, which is a retry of PR lnls-fac/trackcpp#74, based on the discussions raised there. It adds the information of lost positions of tracked particles in line_pass and ring_pass as a separated output.

Besides this change, it also features:

Tests

Comparison of tracking results with master branch:

The following comparison was performed against the master branch of trackcpp and pyaccel: The following code was run on both, master branch of trackcpp and pyaccel and this branch in pyaccel and the branch of lnls-fac/trackcpp#75:

import numpy as np
import pyaccel as pa
from mathphys.functions import load, save
from pymodels import si

mod = si.create_accelerator()
mod = si.fitted_models.vertical_dispersion_and_coupling(mod)
mod.cavity_on = True
mod.radiation_on = 0
mod.vchamber_on = True
print(mod)

eqpar = pa.optics.EqParamsFromBeamEnvelope(mod)
np.random.seed(2030948)
parts = pa.tracking.generate_bunch(1000, eqpar.envelopes[0]*(8e-3/65e-6)**2)
out_line = pa.tracking.line_pass(
    mod, parts, indices='closed', element_offset=0, parallel=True,
)
out_ring = pa.tracking.ring_pass(
    mod, parts, nr_turns=10, turn_by_turn=True, element_offset=0, parallel=True
)
# on this branch
save(out_line, 'sirius_dev_line', overwrite=True)
save(out_ring, 'sirius_dev_ring', overwrite=True)
# on master branch
# save(out_line, 'sirius_mas_line', overwrite=True)
# save(out_ring, 'sirius_mas_ring', overwrite=True)

then, the following comparison was made (must be in the current branches of this PR):

import numpy as np
from mathphys.functions import load, save

out_dev_line = load('sirius_dev_line')
out_dev_ring = load('sirius_dev_ring')
out_mas_line = load('sirius_mas_line')
out_mas_ring = load('sirius_mas_ring')

print(out_mas_line[0].tobytes() == out_dev_line[0].tobytes())
print(out_mas_ring[0].tobytes() == out_dev_ring[0].tobytes())

and the equalities hold.

Test of output shapes:

I also performed a test to check if the shape of the tracking results and the lost positions were consistent when different parameters of the tracking functions were used. The following code was used:

from itertools import product

import numpy as np
import pyaccel as pa
from mathphys.functions import load, save
from pymodels import si

mod = si.create_accelerator()
mod = si.fitted_models.vertical_dispersion_and_coupling(mod)
mod.cavity_on = True
mod.radiation_on = 0
mod.vchamber_on = True
print(mod)

parts = np.array([
    [-11.8e-3, 1e-5, 0, 0, 0, 0],
    [1e-4, 0, 0, 0, 0, 0]
]).T

print('ring_pass')
idcs = True, False
para = True, False
nrpt = 1, 2
for idx, par, nrp in product(idcs, para, nrpt):
    out = pa.tracking.ring_pass(
        mod, parts[:, :nrp], nr_turns=2, turn_by_turn=idx, element_offset=0, parallel=par
    )
    print(
        idx,
        par,
        nrp,
        out[0].shape,
        out[1].lost_pos.shape,
    )
print('line_pass')
idcs = 'closed', None
para = True, False
nrpt = 1, 2
for idx, par, nrp in product(idcs, para, nrpt):
    out = pa.tracking.line_pass(
        mod, parts[:, :nrp], indices=idx, element_offset=0, parallel=par,
    )
    print(
        idx,
        par,
        nrp,
        out[0].shape,
        out[1].lost_pos.shape
    )

resulting in:

ring_pass
True True 1 (6, 3) (6,)
True True 2 (6, 2, 3) (6, 2)
True False 1 (6, 3) (6,)
True False 2 (6, 2, 3) (6, 2)
False True 1 (6,) (6,)
False True 2 (6, 2) (6, 2)
False False 1 (6,) (6,)
False False 2 (6, 2) (6, 2)
line_pass
closed True 1 (6, 6550) (6,)
closed True 2 (6, 2, 6550) (6, 2)
closed False 1 (6, 6550) (6,)
closed False 2 (6, 2, 6550) (6, 2)
None True 1 (6,) (6,)
None True 2 (6, 2) (6, 2)
None False 1 (6,) (6,)
None False 2 (6, 2) (6, 2)