search

desy-ml / cheetah

Fast and differentiable particle accelerator optics simulation for reinforcement learning and optimisation applications.
https://cheetah-accelerator.readthedocs.io
GNU General Public License v3.0
27 stars 12 forks source link

Test Cheetah speed by comparing it to Ocelot #60

Open jank324 opened 10 months ago

jank324 commented 10 months ago

Right now, test/test_speed.py only tests that Cheetah is faster than some constant and relatively long time.

def test_tracking_speed():
    """Really only tests that Cheetah isn't super slow."""
    cell = cheetah.utils.subcell_of_ocelot(ares.cell, "AREASOLA1", "AREABSCR1")
    segment = cheetah.Segment.from_ocelot(cell)
    segment.AREABSCR1.is_active = True  # Turn screen on and off

    particles = cheetah.ParticleBeam.from_parameters(
        num_particles=int(1e5), sigma_x=175e-6, sigma_y=175e-6
    )

    t1 = time.time()

    _ = segment(particles)
    _ = segment.AREABSCR1.reading

    t2 = time.time()

    assert (t2 - t1) < 0.1

This is not ideal, because it is a very long time and also because the time the computations take may vary from machine to machine. It would be better to do a tracking with Ocelot and one with Cheetah and compare that, as we advertise, Cheetah is faster than Ocelot.