Propagation padding

src/hotopy/holo/propagation.py

@@ -163,7 +163,7 @@ def phase_chirp(shape, fresnel_nums, real=False, ndim=None, dtype=None):
     return chirp
 
 
-class FourierConvolutionPropagator:
+class _FourierConvolutionPropagator:
     @property
     def kernel(self):
         """Propagation kernel."""
@@ -177,10 +177,11 @@ class FourierConvolutionPropagator:
         self,
         shape,
         fresnel_numbers,
-        dtype=None,
-        device=None,
+        npad=1,
         pad_width=None,
         pad_mode=None,
+        dtype=None,
+        device=None,
         keep_type=True,
     ):
         self.shape = tuple(shape)
@@ -198,7 +199,9 @@ class FourierConvolutionPropagator:
         self._fresnel_numbers = expand_fresnel_numbers(fresnel_numbers, ndim=self._ndim)
 
         if pad_width is None:
-            pad_width = 0
+            # symmetric padding if npad option is set and no pad_width passed
+            pad_size = np.multiply(self.shape, npad - 1)
+            pad_width = tuple([(s_i + 1 // 2, s_i) for s_i in pad_size])
         if pad_mode is None:
             pad_mode = "constant"  # zero padding
         self.padding = Padder(self.shape, pad_width, mode=pad_mode)
@@ -259,7 +262,34 @@ class FourierConvolutionPropagator:
         return u1
 
 
-class FresnelTFPropagator(FourierConvolutionPropagator):
+class FresnelTFPropagator(_FourierConvolutionPropagator):
+    """
+    Callable Fresnel transfer function propagator.
+
+    Parameters
+    ----------
+    shape : tuple
+        Dimension of image to propagate.
+    fresnel_numbers: float, array_like
+        Pixel Fresnel number(s) encoding the propagation distance, wavelength and detector pixel size.
+    npad: int, float, Optional
+        Padding factor, defaults to ``1``, i.e. no padding. Only used if ``pad_width = None``.
+    pad_width: tuple, Optional
+        Alternatively amount of padding per axis in NumPy notation. Defaults to ``None`` no padding.
+    pad_mode: str, Optional
+        Padding mode. See NumPy ``np.pad`` modes. Defaults to ``None`` which is zero padding.
+    dtype: torch.dtype, Optional
+        Real data type to use for propagation kernel.
+    keep_type: bool
+        If NumPy array is given, also returns NumPy if set. Defaults to ``True``. If ``False`` always a PyTorch tensor
+        is returned on device given by device argument.
+
+    Returns
+    -------
+    F: Callable
+        Propagator, callable with ``F(x)`` to propagator ``x``.
+    """
+
     def _init_kernel(self):
         # verify sufficient padding
         if not check_fresnelTF_sampling(self._pad_shape, self._fresnel_numbers):
@@ -284,6 +314,7 @@ def simulate_hologram(
     linear=False,
     npad=1,
     pad_width=None,
+    pad_mode=None,
     propagator=None,
     device=None,
     dtype=None,
@@ -326,11 +357,14 @@ def simulate_hologram(
     shape = h.shape
 
     if propagator is None:
-        if pad_width is None:
-            pad_width = tuple(np.multiply(shape, npad).astype(int))
-
         propagator = FresnelTFPropagator(
-            shape, fresnel_nums, pad_width=pad_width, dtype=dtype, device=device
+            shape,
+            fresnel_nums,
+            pad_width=pad_width,
+            npad=npad,
+            pad_mode=pad_mode,
+            dtype=dtype,
+            device=device,
         )
     elif not callable(propagator):
         raise ValueError("Propagator needs to be callable.")