MAINT: allow padding adjustment when replacing rmesh

ChristopherMayes · Dec 16, 2024 · 9fbfe76 · 9fbfe76
1 parent adf03ec
commit 9fbfe76
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/pmd_beamphysics/wavefront.py b/pmd_beamphysics/wavefront.py
@@ -931,14 +931,22 @@ def gaussian_pulse(
             longitudinal_axis="z",
         )
 
-    def with_rmesh(self, rmesh: np.ndarray) -> Wavefront:
+    def with_rmesh(
+        self,
+        rmesh: np.ndarray,
+        pad: Sequence[int] | int | None = None,
+        fix_pad: bool = True,
+    ) -> Wavefront:
         """Create a new Wavefront instance, replacing the `rmesh`."""
+        if pad is None:
+            pad = self.pad
         return Wavefront(
             rmesh=rmesh,
             wavelength=self.wavelength,
             metadata=self.metadata,
             longitudinal_axis=self._longitudinal_axis,
             pad=self.pad,
+            fix_pad=fix_pad,
         )
 
     def with_padding(