From 7487a94c45855d06571b8168139159aa571ca76c Mon Sep 17 00:00:00 2001 From: Eric Bezzam Date: Tue, 4 Jun 2024 07:33:48 -0700 Subject: [PATCH] Add vertical and horizontal shift options. --- waveprop/simulation.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/waveprop/simulation.py b/waveprop/simulation.py index e96d5ff..3f32270 100644 --- a/waveprop/simulation.py +++ b/waveprop/simulation.py @@ -40,6 +40,8 @@ def __init__( is_torch=False, quantize=True, return_float=True, + vertical_shift=None, + horizontal_shift=None, **kwargs ): """ @@ -90,6 +92,8 @@ def __init__( self.sensor = sensor_dict[sensor] self.random_shift = random_shift self.quantize = quantize + self.vertical_shift = vertical_shift + self.horizontal_shift = horizontal_shift # for convolution if psf is not None: @@ -165,6 +169,17 @@ def propagate(self, obj, return_object_plane=False): random_shift=self.random_shift, ) + if self.is_torch: + if self.vertical_shift is not None: + object_plane = torch.roll(object_plane, self.vertical_shift, dims=1) + if self.horizontal_shift is not None: + object_plane = torch.roll(object_plane, self.horizontal_shift, dims=2) + else: + if self.vertical_shift is not None: + object_plane = np.roll(object_plane, self.vertical_shift, axis=1) + if self.horizontal_shift is not None: + object_plane = np.roll(object_plane, self.horizontal_shift, axis=2) + if self.fft_shape is not None: # 2) Convolve with PSF if torch.is_tensor(object_plane):