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
28 stars 13 forks source link

[Bug] Screen reading not working due to wrong dimension in torch.histogramdd #30

Closed cr-xu closed 1 year ago

cr-xu commented 1 year ago

Screen reading is not working currently: Input data to the torch.histogramdd should be number_datapoints*number_dimension

c.f. https://github.com/desy-ml/cheetah/blob/11c75f495a71152053a284a8e18a46de5cb7a5e2/cheetah/accelerator.py#L758

MWE:

import cheetah
import test.ARESlatticeStage3v1_9 as ares

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(
    n=int(1e5), sigma_x=175e-6, sigma_y=175e-6
)

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