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.

• 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)
• 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

• I have updated the changelog accordingly (required).
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have reformatted the code and checked that formatting passes (required).
• I have have fixed all issues found by flake8 (required).
• 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).
• 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)
• 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
• Check if we want and can do n-dimensional batches (supported by PyTorch) ... values[..., i]
• 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.

[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 (?)

[cr-xu]

Looks good to me.
I'm not so sure about pushing .vscode/ltex... stuff though, but it's not a big issue.

[jank324]

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