more careful array sums interpolation

py4DSTEM/datacube/datacube.py

@@ -405,7 +405,9 @@ def pad_Q(self, N=None, output_size=None):
         d = pad_data_diffraction(self, pad_factor=N, output_size=output_size)
         return d
 
-    def resample_Q(self, N=None, output_size=None, method="bilinear"):
+    def resample_Q(
+        self, N=None, output_size=None, method="bilinear", conserve_array_sums=False
+    ):
         """
         Resamples the data in diffraction space by resampling factor N, or to match output_size,
         using either 'fourier' or 'bilinear' interpolation.
@@ -418,7 +420,11 @@ def resample_Q(self, N=None, output_size=None, method="bilinear"):
         from py4DSTEM.preprocess import resample_data_diffraction
 
         d = resample_data_diffraction(
-            self, resampling_factor=N, output_size=output_size, method=method
+            self,
+            resampling_factor=N,
+            output_size=output_size,
+            method=method,
+            conserve_array_sums=conserve_array_sums,
         )
         return d
 

py4DSTEM/preprocess/preprocess.py

@@ -573,7 +573,11 @@ def datacube_diffraction_shift(
 
 
 def resample_data_diffraction(
-    datacube, resampling_factor=None, output_size=None, method="bilinear"
+    datacube,
+    resampling_factor=None,
+    output_size=None,
+    method="bilinear",
+    conserve_array_sums=False,
 ):
     """
     Performs diffraction space resampling of data by resampling_factor or to match output_size.
@@ -594,7 +598,10 @@ def resample_data_diffraction(
         old_size = datacube.data.shape
 
         datacube.data = fourier_resample(
-            datacube.data, scale=resampling_factor, output_size=output_size
+            datacube.data,
+            scale=resampling_factor,
+            output_size=output_size,
+            conserve_array_sums=conserve_array_sums,
         )
 
         if not resampling_factor:
@@ -650,6 +657,9 @@ def resample_data_diffraction(
                 grid_mode=True,
             )
 
+        if conserve_array_sums:
+            output_data = output_data / resampling_factor.prod()
+
         datacube.data = output_data
         datacube.calibration.set_Q_pixel_size(
             datacube.calibration.get_Q_pixel_size() / resampling_factor[0]

py4DSTEM/process/phase/phase_base_class.py

@@ -346,34 +346,49 @@ def _preprocess_datacube_and_vacuum_probe(
 
             elif reshaping_method == "fourier":
                 datacube = datacube.resample_Q(
-                    N=resampling_factor_x, method=reshaping_method
+                    N=resampling_factor_x,
+                    method=reshaping_method,
+                    conserve_array_sums=True,
                 )
                 if vacuum_probe_intensity is not None:
                     vacuum_probe_intensity = fourier_resample(
                         vacuum_probe_intensity,
                         output_size=diffraction_intensities_shape,
                         force_nonnegative=True,
+                        conserve_array_sums=True,
                     )
                 if dp_mask is not None:
                     dp_mask = fourier_resample(
                         dp_mask,
                         output_size=diffraction_intensities_shape,
                         force_nonnegative=True,
+                        conserve_array_sums=False,
                     )
 
             elif reshaping_method == "bilinear":
                 datacube = datacube.resample_Q(
-                    N=resampling_factor_x, method=reshaping_method
+                    N=resampling_factor_x,
+                    method=reshaping_method,
+                    conserve_array_sums=True,
                 )
                 if vacuum_probe_intensity is not None:
                     vacuum_probe_intensity = zoom(
                         vacuum_probe_intensity,
                         (resampling_factor_x, resampling_factor_x),
                         order=1,
+                        mode="grid-wrap",
+                        grid_mode=True,
+                    )
+                    vacuum_probe_intensity = (
+                        vacuum_probe_intensity / resampling_factor_x**2
                     )
                 if dp_mask is not None:
                     dp_mask = zoom(
-                        dp_mask, (resampling_factor_x, resampling_factor_x), order=1
+                        dp_mask,
+                        (resampling_factor_x, resampling_factor_x),
+                        order=1,
+                        mode="grid-wrap",
+                        grid_mode=True,
                     )
 
             else:

py4DSTEM/process/phase/ptychographic_methods.py

@@ -1007,6 +1007,7 @@ def _initialize_probe(
                         vacuum_probe_intensity,
                         output_size=(tx, ty),
                         vectorized=True,
+                        conserve_array_sums=True,
                         xp=xp,
                     )
 
@@ -1568,6 +1569,7 @@ def _gradient_descent_fourier_projection(self, amplitudes, overlap, fourier_mask
                 fourier_overlap,
                 output_size=self._amplitudes_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
 
@@ -1585,6 +1587,7 @@ def _gradient_descent_fourier_projection(self, amplitudes, overlap, fourier_mask
                 fourier_modified_overlap,
                 output_size=self._region_of_interest_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
 
@@ -1660,6 +1663,7 @@ def _projection_sets_fourier_projection(
                 fourier_projected_factor,
                 output_size=self._amplitudes_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
 
@@ -1676,6 +1680,7 @@ def _projection_sets_fourier_projection(
                 fourier_projected_factor,
                 output_size=self._region_of_interest_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
 
@@ -2074,7 +2079,11 @@ def _position_correction(
         # resample to match data, note: this needs to happen in reciprocal-space
         if self._resample_exit_waves:
             overlap_fft = bilinear_resample(
-                overlap_fft, output_size=self._amplitudes_shape, vectorized=True, xp=xp
+                overlap_fft,
+                output_size=self._amplitudes_shape,
+                vectorized=True,
+                conserve_array_sums=True,
+                xp=xp,
             )
 
         # unperturbed
@@ -2114,12 +2123,14 @@ def _position_correction(
                 overlap_dx_fft,
                 output_size=self._amplitudes_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
             overlap_dy_fft = bilinear_resample(
                 overlap_dy_fft,
                 output_size=self._amplitudes_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
 
@@ -2224,6 +2235,7 @@ def _return_self_consistency_errors(
                     fourier_overlap,
                     output_size=self._amplitudes_shape,
                     vectorized=True,
+                    conserve_array_sums=True,
                     xp=xp,
                 )
 
@@ -2720,6 +2732,7 @@ def _gradient_descent_fourier_projection(self, amplitudes, overlap, fourier_mask
                 fourier_overlap,
                 output_size=self._amplitudes_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
 
@@ -2741,6 +2754,7 @@ def _gradient_descent_fourier_projection(self, amplitudes, overlap, fourier_mask
                 fourier_modified_overlap,
                 output_size=self._region_of_interest_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
 
@@ -2816,6 +2830,7 @@ def _projection_sets_fourier_projection(
                 fourier_projected_factor,
                 output_size=self._amplitudes_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
 
@@ -2833,6 +2848,7 @@ def _projection_sets_fourier_projection(
                 fourier_projected_factor,
                 output_size=self._region_of_interest_shape,
                 vectorized=True,
+                conserve_array_sums=True,
                 xp=xp,
             )
 

py4DSTEM/process/phase/utils.py

@@ -2251,6 +2251,7 @@ def bilinear_resample(
     mode="grid-wrap",
     grid_mode=True,
     vectorized=True,
+    conserve_array_sums=False,
     xp=np,
 ):
     """
@@ -2321,13 +2322,17 @@ def bilinear_resample(
 
         array = out_array.reshape(tuple(array_size[:-2]) + tuple(output_size))
 
+    if conserve_array_sums:
+        array = array / np.array(scale_output).prod()
+
     return array
 
 
 def vectorized_fourier_resample(
     array,
     scale=None,
     output_size=None,
+    conserve_array_sums=False,
     xp=np,
 ):
     """
@@ -2489,7 +2494,8 @@ def vectorized_fourier_resample(
     array_resize = xp.real(xp.fft.ifft2(array_resize)).astype(xp.float32)
 
     # Normalization
-    array_resize *= scale_output
+    if not conserve_array_sums:
+        array_resize = array_resize * scale_output
 
     return array_resize
 

py4DSTEM/process/utils/utils.py

@@ -446,6 +446,7 @@ def fourier_resample(
     bandlimit_nyquist=None,
     bandlimit_power=2,
     dtype=np.float32,
+    conserve_array_sums=False,
 ):
     """
     Resize a 2D array along any dimension, using Fourier interpolation / extrapolation.
@@ -464,6 +465,7 @@ def fourier_resample(
         bandlimit_nyquist (float): Gaussian filter information limit in Nyquist units (0.5 max in both directions)
         bandlimit_power (float): Gaussian filter power law scaling (higher is sharper)
         dtype (numpy dtype): datatype for binned array. default is single precision float
+        conserve_arrray_sums (bool): If True, the sums of the array are conserved
 
     Returns:
         the resized array (2D/4D numpy array)
@@ -661,7 +663,8 @@ def fourier_resample(
         array_resize = np.maximum(array_resize, 0)
 
     # Normalization
-    array_resize = array_resize * scale_output
+    if not conserve_array_sums:
+        array_resize = array_resize * scale_output
 
     return array_resize