Skip to content

Large improvement in time needed for particles creation - #848

Open
freddieknets wants to merge 5 commits into
xsuite:mainfrom
freddieknets:BetterParticles
Open

freddieknets wants to merge 5 commits into
xsuite:mainfrom
freddieknets:BetterParticles

Conversation

@freddieknets

Copy link
Copy Markdown
Contributor

Description

Particles generation can take a long time.
For instance, generating 100M particles on my machines (same time on a MacBook M3 and AMD Ryzen 7 9850X3D) takes a bit more than 30 seconds.
This PR addresses this in two steps, and in doing so, gains more than an order of magnitude in speed (from 30 to 2.2 seconds).

PDG optimisation

Numeric PDG-ID arrays were processed recursively in Python—one function call and several type checks per particle.
Ion lookup scanned and lowercased all 236 short/long element names. Initial mass-table construction triggered this hundreds of times.
The new approach:

  • Converts numeric arrays through NumPy; int64 arrays are zero-copy
  • Parses the ion’s mass number, then uses a precomputed lowercase element dictionary for O(1) lookup
  • Retains existing string, antiparticle, and ion formats

Very strong speed gain: from 30 seconds down to 10 seconds.

Particles __init__ optimisation

Key improvements:

  • Avoid sentinel-prefilling arrays when every slot will be overwritten
  • Avoid boolean masks and masked copies for all-valid particles
  • Broadcast CPU scalar inputs without allocating full-size temporary arrays
  • Avoid unnecessary NumPy copies and eager particle_id allocation
  • Skip consistency checks when input and computed value are identical
  • Replace Python-level mass-ratio iteration with vectorised checks
  • Skip unnecessary default-zeta calculations and redundant CPU reorganisation
  • Fresh CPU/CUDA buffers are already zeroed, so initialisation skips redundant zero writes for default fields, pdg_id, and random state; reused buffers, pyopencl, extra capacity, invalid particles, and signed -0.0 retain the original handling
  • _update_energy_deviations avoids unused pzeta calculations and reuses the already-computed ptau expression for rvv, eliminating a second square root on the physical CPU path; unusual/negative-energy inputs use the original formula
  • _update_zeta no longer derives tau unless supplied and returns immediately for defaults

Further speed gain: from 10 seconds down to 2.2 seconds.

All Xtrack tests pass on CPU; currently running CUDA tests overnight.

Checklist

Mandatory:

  • I have added tests to cover my changes
  • All the tests are passing, including my new ones
  • I described my changes in this PR description

Optional:

  • The code I wrote follows good style practices (see PEP 8 and PEP 20).
  • I have updated the docs in relation to my changes, if applicable
  • I have tested also GPU contexts

# in calculations.
if kwargs[field].dtype != xotype._dtype:
kwargs[field] = kwargs[field].astype(xotype._dtype)
# PyOpenCL gets fussy if types don't match in calculations.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we skip the type casting here?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants