Merge pull request #308 from Hespe/fix-clone-warnings

Fix warnings in test suite

Author: jank324

CHANGELOG.md

@@ -47,6 +47,7 @@ This is a major release with significant upgrades under the hood of Cheetah. Des
 - Fix plotting for segments that contain tensors with `require_grad=True` (see #288) (@hespe)
 - Fix bug where `Element.length` could not be set as a `torch.nn.Parameter` (see #301) (@jank324, @hespe)
 - Fix registration of `torch.nn.Parameter` at initilization for elements and beams (see #303) (@hespe)
+- Fix warnings about NumPy deprecations and unintentional tensor clones (see #308) (@hespe)
 
 ### 🐆 Other
 

cheetah/accelerator/dipole.py

@@ -1,7 +1,6 @@
 from typing import Literal, Optional
 
 import matplotlib.pyplot as plt
-import numpy as np
 import torch
 from matplotlib.patches import Rectangle
 from scipy.constants import physical_constants
@@ -104,14 +103,14 @@ def __init__(
         if fringe_integral is not None:
             self.fringe_integral = torch.as_tensor(fringe_integral, **factory_kwargs)
         self.fringe_integral_exit = (
-            torch.tensor(fringe_integral_exit, **factory_kwargs)
+            torch.as_tensor(fringe_integral_exit, **factory_kwargs)
             if fringe_integral_exit is not None
             else self.fringe_integral
         )
         if gap is not None:
             self.gap = torch.as_tensor(gap, **factory_kwargs)
         self.gap_exit = (
-            torch.tensor(gap_exit, **factory_kwargs)
+            torch.as_tensor(gap_exit, **factory_kwargs)
             if gap_exit is not None
             else self.gap
         )
@@ -493,7 +492,7 @@ def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> No
         plot_angle = self.angle[vector_idx] if self.angle.dim() > 0 else self.angle
 
         alpha = 1 if self.is_active else 0.2
-        height = 0.8 * (np.sign(plot_angle) if self.is_active else 1)
+        height = 0.8 * (torch.sign(plot_angle) if self.is_active else 1)
 
         patch = Rectangle(
             (plot_s, 0), plot_length, height, color="tab:green", alpha=alpha, zorder=2

cheetah/accelerator/horizontal_corrector.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 import matplotlib.pyplot as plt
-import numpy as np
 import torch
 from matplotlib.patches import Rectangle
 
@@ -88,7 +87,7 @@ def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> No
         plot_angle = self.angle[vector_idx] if self.angle.dim() > 0 else self.angle
 
         alpha = 1 if self.is_active else 0.2
-        height = 0.8 * (np.sign(plot_angle) if self.is_active else 1)
+        height = 0.8 * (torch.sign(plot_angle) if self.is_active else 1)
 
         patch = Rectangle(
             (plot_s, 0), plot_length, height, color="tab:blue", alpha=alpha, zorder=2

cheetah/accelerator/quadrupole.py

@@ -1,7 +1,6 @@
 from typing import Literal, Optional
 
 import matplotlib.pyplot as plt
-import numpy as np
 import torch
 from matplotlib.patches import Rectangle
 from scipy.constants import physical_constants
@@ -215,7 +214,7 @@ def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> No
         plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
 
         alpha = 1 if self.is_active else 0.2
-        height = 0.8 * (np.sign(plot_k1) if self.is_active else 1)
+        height = 0.8 * (torch.sign(plot_k1) if self.is_active else 1)
         patch = Rectangle(
             (plot_s, 0), plot_length, height, color="tab:red", alpha=alpha, zorder=2
         )

cheetah/accelerator/vertical_corrector.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 import matplotlib.pyplot as plt
-import numpy as np
 import torch
 from matplotlib.patches import Rectangle
 from scipy.constants import physical_constants
@@ -91,7 +90,7 @@ def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> No
         plot_angle = self.angle[vector_idx] if self.angle.dim() > 0 else self.angle
 
         alpha = 1 if self.is_active else 0.2
-        height = 0.8 * (np.sign(plot_angle) if self.is_active else 1)
+        height = 0.8 * (torch.sign(plot_angle) if self.is_active else 1)
 
         patch = Rectangle(
             (plot_s, 0), plot_length, height, color="tab:cyan", alpha=alpha, zorder=2