Use KDE for beam screen images

cr-xu commented 5 days ago

Description

Add an option for Screen element. Now it can use KDE to get the beam image which allows differentiation. Behavior can be switched using method="histogram" and method="kde". It contains a new feature kde_bandwidth for smoothing the beam image.

The returned image has the convention [...,Y, X] where y-axis starts from the top of the image, e.g. plt.imshow() will show the image upright


import cheetah
import torch
import matplotlib.pyplot as plt

incoming = cheetah.ParticleBeam.from_parameters(num_particles = 10000, mu_x = torch.tensor([-0.5e-3]), mu_y=torch.tensor([1.0e-3]), sigma_x = torch.tensor([1e-4]), sigma_y = torch.tensor([3e-4]))

screen = cheetah.Screen(resolution=(1200, 800), pixel_size=(5e-6, 5e-6), binning=torch.tensor(2), is_active=True)

screen.track(incoming)

plt.imshow(screen.reading[0], extent=screen.extent, cmap='gray')
plt.xlabel('x [m]')
plt.ylabel('y [m]')

Motivation and Context

[x] I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

[ ] Bug fix (non-breaking change which fixes an issue)
[x] New feature (non-breaking change which adds functionality)
[ ] Breaking change (fix or feature that would cause existing functionality to change)
[ ] Documentation (update in the documentation)

Checklist

[x] I have updated the changelog accordingly (required).
[ ] My change requires a change to the documentation.
[x] I have updated the tests accordingly (required for a bug fix or a new feature).
[ ] I have updated the documentation accordingly.
[x] I have reformatted the code and checked that formatting passes (required).
[x] I have have fixed all issues found by flake8 (required).
[x] I have ensured that all pytest tests pass (required).
[ ] I have run pytest on a machine with a CUDA GPU and made sure all tests pass (required).
[x] I have checked that the documentation builds (required).

Note: We are using a maximum length of 88 characters per line

cr-xu commented 5 days ago

Hmm... so KDE clearly makes the screen reading slower, as expected :)

@jank324 Shall we

Change the test file to make the speed limit more generous?
Still provide a basic histogramdd() reading for normal tracking, when differentiation w.r.t. images are not needed? This can be switched e.g. with use_kde=False by default?

jank324 commented 5 days ago

I think we should provide both. Im just wondering whether it should be use_kde=True or method="kde".

desy-ml / cheetah