Remove drift replacement for off-crest Cavity

[Hespe]

Description

The initial fix for the linked issue did accidentally remove bunch compression effects for off-crest cavities.

Motivation and Context

• I have raised an issue to propose this change (required for new features and bug fixes)
  • A cavity with a phase of -90 produces NaN values in the transfer map #548

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.

[Hespe]

I tried implementing a test comparing a 90° Cavity to Ocelot, but according to a comment in the code we implement Cavity without acceleration slightly different than Ocelot so that did not work out.

[cr-xu]

So this basically rolls back the #549, and the fix is expected to be introduced in #538, right?

The changelog seems incorrect as the division by zero is not removed in this PR.

[Hespe]

So this basically rolls back the #549, and the fix is expected to be introduced in #538, right?

The changelog seems incorrect as the division by zero is not removed in this PR.

Its not quite a rollback. In #549 I did not only introduce the drift replacement but also did some minor improvements to the numerical conditioning of Cavity. It does not produce the NaN anymore (although the division by 0 is technically still there).

We should still properly remove the division by 0 though, either by cherry-picking or waiting for #538.

[Hespe]

We decided to cherry-pick the updates of #538 into this PR. We should therefore also consider #538 (comment) here.

[jank324]

@Hespe, did you check the impact of this on speed?

[jank324]

I just ran this script on it, which seems to suggest that there is no meaningful impact on speed.

[jank324]

@Hespe, maybe you could test just the line computing alpha or something similarly small with timeit to see on a very granular level what the impact is there?

[Hespe]

@Hespe, maybe you could test just the line computing alpha or something similarly small with timeit to see on a very granular level what the impact is there?

There seems to be some performance impact, even when compared to using torch.where to get rid of the singularity at 0. But that might just be the cost of getting proper gradients in that case.

[jank324]

Hmm ... I'm not happy about this. We definitely cannot find a way of rearranging this equation to avoid the new autograd object and its performance penalty?

[Hespe]

Maybe @cr-xu has an idea, but I do not see how one could remove the singularity of $\frac{\log(1+x)}{x}$ without some kind of case distinction. And a custom autograd function seems to be the way of doing that in Torch without losing gradients.

[cr-xu]

I don't see how this function can be rewritten to get rid of the singularity.

My opinion is: to keep it simple, we can also just leave it (just use torch.where() or even remove that) and note in the docstring that the allowed phases are (-90, 90) and one should not set exactly 90 degs phase offset.

Even if someone wants exactly 90 deg RF phase offset and still gradient with respect to that, they can add a small eps as that will be changed later by gradient-descent anyway. It doesn't seem to be worth it to sacrafice the speed for all other cavity use cases.

---

Apparently, the special autograd handling for pytorch sinc function is implemented in cpp so it doesn't suffer from the speed so much. (pytorch/pytorch@5854e93). I don't think that is possible for us.
Maybe we can try to open issue in pytorch and convince them to include in the main pytorch branch, but that probably takes longer and there's no promise that it will happen.

[cr-xu]

Looks good to me. If the custom autograd doesn't impact the over all performance too much, it's the proper way to go. 👍