Replace constant tensor initialisation with different function calls

[jank324]

Description

Sub- / replacement PR for #538.

Where possible, replace new initialisations of tensors with another_tensor.new_xxx calls. If the goal is to reproduce the shape on another tensor, torch.xxx_like is used.

Motivation and Context

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

[jank324]

[jank324]

Should element inits also adopt the new x.new_zeros() syntax or keep the old one?

[Copilot]

Pull Request Overview

This PR optimizes PyTorch tensor initialization by replacing verbose torch.tensor() and torch.zeros() calls with more efficient alternatives like tensor.new_zeros() and torch.broadcast_tensors(), which reduces overhead by reusing device and dtype from existing tensors.

• Replace torch.tensor(0.0, device=..., dtype=...) with tensor.new_zeros(())
• Replace torch.zeros(..., device=..., dtype=...) with tensor.new_zeros(shape)
• Replace torch.ones_like() with torch.broadcast_tensors() where appropriate

Reviewed Changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
cheetah/utils/bmadx.py	Replace tensor initialization with new_ones()
cheetah/track_methods.py	Replace zero tensor creation with new_zeros()
cheetah/accelerator/space_charge_kick.py	Replace multiple torch.zeros() calls with new_zeros()
cheetah/accelerator/sextupole.py	Replace zero tensor constants with new_zeros()
cheetah/accelerator/screen.py	Replace torch.zeros() with new_zeros()
cheetah/accelerator/quadrupole.py	Replace zero constants and torch.ones_like() with optimized alternatives
cheetah/accelerator/drift.py	Replace multiple zero tensor creations with single new_zeros()
cheetah/accelerator/dipole.py	Replace zero tensor constants and fix tensor shape issue
cheetah/accelerator/cavity.py	Replace torch.zeros_like() and torch.full_like() with new_zeros()
cheetah/accelerator/bpm.py	Move buffer registration after misalignment initialization
CONTRIBUTING.md	Add performance guidelines for tensor creation
CHANGELOG.md	Document performance improvements

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

The buffer registration should use self.misalignment.new_full((2,), torch.nan) instead of torch.tensor((torch.nan, torch.nan), **factory_kwargs) to be consistent with the PR's optimization goals and ensure the tensor uses the same device/dtype as other parameters.

Suggested change

	torch.tensor((torch.nan, torch.nan), **factory_kwargs),
	self.misalignment.new_full((2,), torch.nan),

[jank324]

This really depends on the answer to my earlier question. I'm thinking the answer is probably no, let's not do that.