cleaned up cupyx.scipy imports

py4DSTEM/process/phase/dpc.py

@@ -59,17 +59,19 @@ def __init__(
         Custom.__init__(self, name=name)
 
         if device == "cpu":
+            import scipy
+
             self._xp = np
             self._asnumpy = np.asarray
-            from scipy.ndimage import gaussian_filter
+            self._scipy = scipy
 
-            self._gaussian_filter = gaussian_filter
         elif device == "gpu":
+            from cupyx import scipy
+
             self._xp = cp
             self._asnumpy = cp.asnumpy
-            from cupyx.scipy.ndimage import gaussian_filter
+            self._scipy = scipy
 
-            self._gaussian_filter = gaussian_filter
         else:
             raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
 
@@ -516,7 +518,7 @@ def _object_gaussian_constraint(self, current_object, gaussian_filter_sigma):
         constrained_object: np.ndarray
             Constrained object estimate
         """
-        gaussian_filter = self._gaussian_filter
+        gaussian_filter = self._scipy.ndimage.gaussian_filter
 
         gaussian_filter_sigma /= self.sampling[0]
         current_object = gaussian_filter(current_object, gaussian_filter_sigma)

py4DSTEM/process/phase/magnetic_ptychographic_tomography.py

@@ -156,34 +156,19 @@ def __init__(
         Custom.__init__(self, name=name)
 
         if device == "cpu":
+            import scipy
+
             self._xp = np
             self._asnumpy = np.asarray
-            from scipy.ndimage import affine_transform, gaussian_filter, rotate, zoom
-
-            self._gaussian_filter = gaussian_filter
-            self._zoom = zoom
-            self._rotate = rotate
-            self._affine_transform = affine_transform
-            from scipy.special import erf
+            self._scipy = scipy
 
-            self._erf = erf
         elif device == "gpu":
+            from cupyx import scipy
+
             self._xp = cp
             self._asnumpy = cp.asnumpy
-            from cupyx.scipy.ndimage import (
-                affine_transform,
-                gaussian_filter,
-                rotate,
-                zoom,
-            )
-
-            self._gaussian_filter = gaussian_filter
-            self._zoom = zoom
-            self._rotate = rotate
-            self._affine_transform = affine_transform
-            from cupyx.scipy.special import erf
+            self._scipy = scipy
 
-            self._erf = erf
         else:
             raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
 
@@ -601,7 +586,8 @@ def preprocess(
                 old_rot_matrix.T,
             )
 
-            probe_overlap_3D_blurred = self._gaussian_filter(probe_overlap_3D, 1.0)
+            gaussian_filter = self._scipy.ndimage.gaussian_filter
+            probe_overlap_3D_blurred = gaussian_filter(probe_overlap_3D, 1.0)
             self._object_fov_mask = asnumpy(
                 probe_overlap_3D_blurred > 0.25 * probe_overlap_3D_blurred.max()
             )

py4DSTEM/process/phase/magnetic_ptychography.py

@@ -142,23 +142,19 @@ def __init__(
         Custom.__init__(self, name=name)
 
         if device == "cpu":
+            import scipy
+
             self._xp = np
             self._asnumpy = np.asarray
-            from scipy.ndimage import gaussian_filter
-
-            self._gaussian_filter = gaussian_filter
-            from scipy.special import erf
+            self._scipy = scipy
 
-            self._erf = erf
         elif device == "gpu":
+            from cupyx import scipy
+
             self._xp = cp
             self._asnumpy = cp.asnumpy
-            from cupyx.scipy.ndimage import gaussian_filter
-
-            self._gaussian_filter = gaussian_filter
-            from cupyx.scipy.special import erf
+            self._scipy = scipy
 
-            self._erf = erf
         else:
             raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
 
@@ -623,7 +619,8 @@ def preprocess(
 
         # initialize object_fov_mask
         if object_fov_mask is None:
-            probe_overlap_blurred = self._gaussian_filter(probe_overlap, 1.0)
+            gaussian_filter = self._scipy.ndimage.gaussian_filter
+            probe_overlap_blurred = gaussian_filter(probe_overlap, 1.0)
             self._object_fov_mask = asnumpy(
                 probe_overlap_blurred > 0.25 * probe_overlap_blurred.max()
             )

py4DSTEM/process/phase/mixedstate_multislice_ptychography.py

@@ -161,6 +161,23 @@ def __init__(
     ):
         Custom.__init__(self, name=name)
 
+        if device == "cpu":
+            import scipy
+
+            self._xp = np
+            self._asnumpy = np.asarray
+            self._scipy = scipy
+
+        elif device == "gpu":
+            from cupyx import scipy
+
+            self._xp = cp
+            self._asnumpy = cp.asnumpy
+            self._scipy = scipy
+
+        else:
+            raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
+
         if initial_probe_guess is None or isinstance(initial_probe_guess, ComplexProbe):
             if num_probes is None:
                 raise ValueError(
@@ -176,27 +193,6 @@ def __init__(
                 )
             num_probes = initial_probe_guess.shape[0]
 
-        if device == "cpu":
-            self._xp = np
-            self._asnumpy = np.asarray
-            from scipy.ndimage import gaussian_filter
-
-            self._gaussian_filter = gaussian_filter
-            from scipy.special import erf
-
-            self._erf = erf
-        elif device == "gpu":
-            self._xp = cp
-            self._asnumpy = cp.asnumpy
-            from cupyx.scipy.ndimage import gaussian_filter
-
-            self._gaussian_filter = gaussian_filter
-            from cupyx.scipy.special import erf
-
-            self._erf = erf
-        else:
-            raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
-
         for key in kwargs.keys():
             if (key not in polar_symbols) and (key not in polar_aliases.keys()):
                 raise ValueError("{} not a recognized parameter".format(key))
@@ -554,7 +550,8 @@ def preprocess(
         probe_overlap = self._sum_overlapping_patches_bincounts(probe_intensities)
 
         if object_fov_mask is None:
-            probe_overlap_blurred = self._gaussian_filter(probe_overlap, 1.0)
+            gaussian_filter = self._scipy.ndimage.gaussian_filter
+            probe_overlap_blurred = gaussian_filter(probe_overlap, 1.0)
             self._object_fov_mask = asnumpy(
                 probe_overlap_blurred > 0.25 * probe_overlap_blurred.max()
             )

py4DSTEM/process/phase/mixedstate_ptychography.py

@@ -136,6 +136,23 @@ def __init__(
     ):
         Custom.__init__(self, name=name)
 
+        if device == "cpu":
+            import scipy
+
+            self._xp = np
+            self._asnumpy = np.asarray
+            self._scipy = scipy
+
+        elif device == "gpu":
+            from cupyx import scipy
+
+            self._xp = cp
+            self._asnumpy = cp.asnumpy
+            self._scipy = scipy
+
+        else:
+            raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
+
         if initial_probe_guess is None or isinstance(initial_probe_guess, ComplexProbe):
             if num_probes is None:
                 raise ValueError(
@@ -151,27 +168,6 @@ def __init__(
                 )
             num_probes = initial_probe_guess.shape[0]
 
-        if device == "cpu":
-            self._xp = np
-            self._asnumpy = np.asarray
-            from scipy.ndimage import gaussian_filter
-
-            self._gaussian_filter = gaussian_filter
-            from scipy.special import erf
-
-            self._erf = erf
-        elif device == "gpu":
-            self._xp = cp
-            self._asnumpy = cp.asnumpy
-            from cupyx.scipy.ndimage import gaussian_filter
-
-            self._gaussian_filter = gaussian_filter
-            from cupyx.scipy.special import erf
-
-            self._erf = erf
-        else:
-            raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
-
         for key in kwargs.keys():
             if (key not in polar_symbols) and (key not in polar_aliases.keys()):
                 raise ValueError("{} not a recognized parameter".format(key))
@@ -481,7 +477,8 @@ def preprocess(
         probe_overlap = self._sum_overlapping_patches_bincounts(probe_intensities)
 
         if object_fov_mask is None:
-            probe_overlap_blurred = self._gaussian_filter(probe_overlap, 1.0)
+            gaussian_filter = self._scipy.ndimage.gaussian_filter
+            probe_overlap_blurred = gaussian_filter(probe_overlap, 1.0)
             self._object_fov_mask = asnumpy(
                 probe_overlap_blurred > 0.25 * probe_overlap_blurred.max()
             )

py4DSTEM/process/phase/multislice_ptychography.py

@@ -153,23 +153,19 @@ def __init__(
         Custom.__init__(self, name=name)
 
         if device == "cpu":
+            import scipy
+
             self._xp = np
             self._asnumpy = np.asarray
-            from scipy.ndimage import gaussian_filter
-
-            self._gaussian_filter = gaussian_filter
-            from scipy.special import erf
+            self._scipy = scipy
 
-            self._erf = erf
         elif device == "gpu":
+            from cupyx import scipy
+
             self._xp = cp
             self._asnumpy = cp.asnumpy
-            from cupyx.scipy.ndimage import gaussian_filter
-
-            self._gaussian_filter = gaussian_filter
-            from cupyx.scipy.special import erf
+            self._scipy = scipy
 
-            self._erf = erf
         else:
             raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
 
@@ -529,7 +525,8 @@ def preprocess(
         probe_overlap = self._sum_overlapping_patches_bincounts(probe_intensities)
 
         if object_fov_mask is None:
-            probe_overlap_blurred = self._gaussian_filter(probe_overlap, 1.0)
+            gaussian_filter = self._scipy.ndimage.gaussian_filter
+            probe_overlap_blurred = gaussian_filter(probe_overlap, 1.0)
             self._object_fov_mask = asnumpy(
                 probe_overlap_blurred > 0.25 * probe_overlap_blurred.max()
             )

py4DSTEM/process/phase/parallax.py

@@ -87,17 +87,19 @@ def __init__(
         Custom.__init__(self, name=name)
 
         if device == "cpu":
+            import scipy
+
             self._xp = np
             self._asnumpy = np.asarray
-            from scipy.ndimage import gaussian_filter
+            self._scipy = scipy
 
-            self._gaussian_filter = gaussian_filter
         elif device == "gpu":
+            from cupyx import scipy
+
             self._xp = cp
             self._asnumpy = cp.asnumpy
-            from cupyx.scipy.ndimage import gaussian_filter
+            self._scipy = scipy
 
-            self._gaussian_filter = gaussian_filter
         else:
             raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
 
@@ -1109,7 +1111,7 @@ def subpixel_alignment(
         """
         xp = self._xp
         asnumpy = self._asnumpy
-        gaussian_filter = self._gaussian_filter
+        gaussian_filter = self._scipy.ndimage.gaussian_filter
 
         BF_sampling = 1 / asnumpy(self._kr).max() / 2
         DF_sampling = 1 / (
@@ -1774,7 +1776,7 @@ def _kernel_density_estimate(
         """ """
 
         xp = self._xp
-        gaussian_filter = self._gaussian_filter
+        gaussian_filter = self._scipy.ndimage.gaussian_filter
 
         if lanczos_alpha is not None:
             return lanczos_kernel_density_estimate(

py4DSTEM/process/phase/phase_base_class.py

@@ -73,23 +73,19 @@ def reinitialize_parameters(self, device: str = None, verbose: bool = None):
 
         if device is not None:
             if device == "cpu":
+                import scipy
+
                 self._xp = np
                 self._asnumpy = np.asarray
-                from scipy.ndimage import gaussian_filter
-
-                self._gaussian_filter = gaussian_filter
-                from scipy.special import erf
+                self._scipy = scipy
 
-                self._erf = erf
             elif device == "gpu":
+                from cupyx import scipy
+
                 self._xp = cp
                 self._asnumpy = cp.asnumpy
-                from cupyx.scipy.ndimage import gaussian_filter
-
-                self._gaussian_filter = gaussian_filter
-                from cupyx.scipy.special import erf
+                self._scipy = scipy
 
-                self._erf = erf
             else:
                 raise ValueError(f"device must be either 'cpu' or 'gpu', not {device}")
             self._device = device

py4DSTEM/process/phase/ptychographic_constraints.py

@@ -138,7 +138,7 @@ def _object_gaussian_constraint(
             Constrained object estimate
         """
         xp = self._xp
-        gaussian_filter = self._gaussian_filter
+        gaussian_filter = self._scipy.ndimage.gaussian_filter
         gaussian_filter_sigma /= self.sampling[0]
 
         if not pure_phase_object or self._object_type == "potential":
@@ -938,7 +938,7 @@ def _probe_amplitude_constraint(
             Constrained probe estimate
         """
         xp = self._xp
-        erf = self._erf
+        erf = self._scipy.special.erf
 
         probe_intensity = xp.abs(current_probe) ** 2
         current_probe_sum = xp.sum(probe_intensity)

py4DSTEM/process/phase/ptychographic_methods.py

@@ -828,7 +828,7 @@ def _rotate_zxy_volume(
         """ """
 
         xp = self._xp
-        affine_transform = self._affine_transform
+        affine_transform = self._scipy.ndimage.affine_transform
         swap_zxy_to_xyz = self._swap_zxy_to_xyz
 
         volume = volume_array.copy()