Vectorised simulations

[jank324]

‼️ Do not merge before v0.6.2 has been released! This must go in a version after that.

Description

Makes all simulations in Cheetah vectorised, as you would expect in normal neural network model implementations.

Motivation and Context

Vectorised simulations should allow for much faster parallel simulation of beams and settings. This is very similar to how you might run a neural network with batches, and helps both speed-based and gradient-based applications.

Closes #100.

• [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)
• [x] 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).
• [x] My change requires a change to the documentation.
• [x] I have updated the tests accordingly (required for a bug fix or a new feature).
• [x] 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).
• [x] 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

[jank324]

Here is a tricky thing to consider: What effect does it have on gradients when we conditionally set some value in a computation to a constant. Is that effect bad? If so, how do we fix it?

E.g. igamma2 in Drift.transfer_map.

[jank324]

We also need to consider input dimension flexibility. For now, most of the reworked code probably assumes one input dimension for beams and one for settings. Maybe the final code can be more flexible, i.e. allow any number of input dimensions for both ... kind of like having n-dimensional batches.

[jank324]

I've made the decision now, that the batch dimensions of the incoming beam and the settings must match. We should add a way to broadcast them, i.e. if a beam only has a single value or if a some elements only have single value, they are broadcast to the dimensions of the batched input. This should not be done if any of them has a different batch dimension, which should not be allowed and hence fail.

[jank324]

Looking at the optimize_speed.ipynb example Notebook, it seems like this change currently slows down Cheetah between 2x and 4x for single samples. On the other hand, if you run with 1,000 samples, we are looking at a 40x to 55x speed improvement on an M1 Pro CPU. This would likely be more with more samples and/or on GPU.

[jank324]

Some outstanding todos:

• [ ] Make example Notebook on batched compuations (decided that speed optimisations example is enough)
• [x] Add batching to speed optimisations example Notebook (cleaning it up)
• [ ] Change gradient-based Notebook to use batches (might not really make sense for the specific gradient-based example notebook)
• [ ] Thorough testing because this change affects pretty much all computations
• [ ] Modify cheetah-demos repository to work with this, but retain paper version on branch for v6.2
• [x] Check if we want and can do n-dimensional batches (supported by PyTorch) ... values[..., i]
• [x] Adapt documentation, README, docstrings, etc.

• [ ] Check that slightly iffy conditional constant replacements don't negatively affect gradients more than they need to

  A further todo from 27.03.2024:

  • [ ] Add at least the functionality to accept just defining an element with dimensionless tensor,so that it's backward-compatible

[jank324]

This will probably be a v0.7 feature.

[jank324]

The last commit solves #127.

@cr-xu we might need to look at this at some point. Small rounding errors sometimes result in the emittance computation taking sqrt of a negative number, or in an emittance of 0.0. I think both is not physically possible. For now this is solved with the clamping, but I wonder if there is another reason for this and the clamping only fixes the symptoms. Not at all urgent though.

[cr-xu]

The last commit solves #127.

@cr-xu we might need to look at this at some point. Small rounding errors sometimes result in the emittance computation taking sqrt of a negative number, or in an emittance of 0.0. I think both is not physically possible. For now this is solved with the clamping, but I wonder if there is another reason for this and the clamping only fixes the symptoms. Not at all urgent though.

At first sight, this is a price we pay for staying in the SI-Units. Emittances are just more commonly very small numbers, usually expressed in mm*mrad, etc.

Since now we can choose the dtype freely, could it be resolved with double-precision...? Can you provide a MWE where the emittance goes to zero or negative numbers?

[jank324]

@cr-xu I think I messed up trying to merge your cavity fixes. Could you please have a look?

[cr-xu]

@cr-xu I think I messed up trying to merge your cavity fixes. Could you please have a look?

I believe it's properly fixed now.

[jank324]

@cr-xu I have two things you might be able to answer:

• I removed the voltage <= 0.0 case from the cavity code because I'm not sure if it makes sense to catch it in the vectorised example. Do you see any potential issue with this?
  • Do we have any tests that check for physical correctness of the cavity against known correct results or Ocelot or whatever else?

[jank324]

Okay ... what else needs to be done?

Should we update the cheetah-demos repo before merging this ... kind of as a further test before putting this on master or should we just go for it?

[cr-xu]

@cr-xu I have two things you might be able to answer:

• I removed the voltage <= 0.0 case from the cavity code because I'm not sure if it makes sense to catch it in the vectorised example. Do you see any potential issue with this?

I think it's safe to do it, as one normally expects an accelerating voltage. For the edge cases (ERL..?) the voltage<=0 case here wouldn't be accurate anyway.

• Do we have any tests that check for physical correctness of the cavity against known correct results or Ocelot or whatever else?

I don't think of one yet. But yes, it's generally helpful to bring in more tests in the future, instead of mostly relying on ARES-EA section.

[cr-xu]

I just updated the BO with Cheetah prior example in the cheetah-demo in a separate branch (see commit 9940072) The vectorized version works without any problem and shows a nice speed-up :) I think the demos can be merged later back to the main branch, once this PR is merged.

The broadcast functionality is really handy to use. 👍

[jank324]

@cr-xu I'm getting the Merging is blocked message again. I would like to figure out this time what is causing it. The only restriction I could find in the settings is the one I set a long time ago that you cannot push directly to master without a PR. I think it has something to do with changes requested in a review not being resolved, so I requested a re-review from you to test this.

[jank324]

Nevermind ... I fixed it. The branch locking option seems to have been the issue

[jank324]

Okay ... I didn't really fix it, but I also don't get it.

The branch locking option seems to be what causes it to tell us that we cannot merge this PR. At first I didn't want to disable it because then I was able to push to master without a PR. It turns out though, that I can do that even when branch locking is turn on, presumably because I'm an admin. So I turned it off and now we can merge PRs again.

I don't like that I'm able to accidentally push to master and as far as I understand the current settings should prevent an admin to push to master ... but clearly they don't. I also googled this and it keeps pointing to an option "Include administrators", but I cannot seem to find that one. I feels a little bit like it may have been renamed or removed (?)

[jank324]

Just for reference: Even with your approval, reactivating the branch lock blocks the merge.