silly gpu bugs for sim

py4dstem · Sep 6, 2024 · 4714ce7 · 4714ce7
1 parent adf112b
commit 4714ce7
Showing 1 changed file with 6 additions and 78 deletions.
diff --git a/py4DSTEM/tomography/tomography.py b/py4DSTEM/tomography/tomography.py
@@ -872,6 +872,9 @@ def _real_space_radon(
 
         """
         xp = self._xp
+        device = self._device
+
+        current_object = copy_to_device(current_object, device)
 
         s = current_object.shape
 
@@ -961,8 +964,8 @@ def _diffraction_space_slice(
 
         tilt = xp.deg2rad(tilt_deg)
 
-        line_y_diff = np.fft.fftfreq(s[-1], 1 / s[-1]) * xp.cos(tilt)
-        line_z_diff = np.fft.fftfreq(s[-1], 1 / s[-1]) * xp.sin(tilt)
+        line_y_diff = xp.fft.fftfreq(s[-1], 1 / s[-1]) * xp.cos(tilt)
+        line_z_diff = xp.fft.fftfreq(s[-1], 1 / s[-1]) * xp.sin(tilt)
 
         yF_diff = xp.floor(line_y_diff).astype("int")
         zF_diff = xp.floor(line_z_diff).astype("int")
@@ -1168,81 +1171,6 @@ def _calculate_update(
         """
         xp = self._xp
 
-        # ind = self._positions_vox_F[0][0] == x_index
-
-        # diffraction_patterns_resampled = xp.zeros(
-        #     (self._positions_vox_dF[0][0].shape[0], object_sliced.shape[-1])
-        # )
-        # diffraction_patterns_resampled[
-        #     xp.ravel_multi_index(
-        #         (
-        #             self._positions_vox_F[datacube_number][0][ind],
-        #             self._positions_vox_F[datacube_number][1][ind],
-        #         ),
-        #         self._initial_datacube_shape[0:2],
-        #         mode="clip",
-        #     )
-        # ] += (
-        #     diffraction_patterns_projected[ind]
-        #     * (
-        #         (1 - self._positions_vox_dF[datacube_number][0][ind])
-        #         * (1 - self._positions_vox_dF[datacube_number][1][ind])
-        #     )[:, None]
-        # )
-
-        # diffraction_patterns_resampled[
-        #     xp.ravel_multi_index(
-        #         (
-        #             self._positions_vox_F[datacube_number][0][ind] + 1,
-        #             self._positions_vox_F[datacube_number][1][ind],
-        #         ),
-        #         self._initial_datacube_shape[0:2],
-        #         mode="clip",
-        #     )
-        # ] += (
-        #     diffraction_patterns_projected[ind]
-        #     * (
-        #         (self._positions_vox_dF[datacube_number][0][ind])
-        #         * (1 - self._positions_vox_dF[datacube_number][1][ind])
-        #     )[:, None]
-        # )
-
-        # diffraction_patterns_resampled[
-        #     xp.ravel_multi_index(
-        #         (
-        #             self._positions_vox_F[datacube_number][0][ind],
-        #             self._positions_vox_F[datacube_number][1][ind] + 1,
-        #         ),
-        #         self._initial_datacube_shape[0:2],
-        #         mode="clip",
-        #     )
-        # ] += (
-        #     diffraction_patterns_projected[ind]
-        #     * (
-        #         (1 - self._positions_vox_dF[datacube_number][0][ind])
-        #         * (self._positions_vox_dF[datacube_number][1][ind])
-        #     )[:, None]
-        # )
-
-        # diffraction_patterns_resampled[
-        #     xp.ravel_multi_index(
-        #         (
-        #             self._positions_vox_F[datacube_number][0][ind] + 1,
-        #             self._positions_vox_F[datacube_number][1][ind] + 1,
-        #         ),
-        #         self._initial_datacube_shape[0:2],
-        #         mode="clip",
-        #     )
-        # ] += (
-        #     diffraction_patterns_projected[ind]
-        #     * (
-        #         (self._positions_vox_dF[datacube_number][0][ind])
-        #         * (self._positions_vox_dF[datacube_number][1][ind])
-        #     )[:, None]
-        # )
-        # diffraction_patterns_resampled = diffraction_patterns_resampled[ind]
-        # update = diffraction_patterns_resampled - object_sliced
-
         s = self._object_shape_6D
 
         ind0 = self._positions_vox_F[datacube_number][0] == x_index
@@ -1477,7 +1405,7 @@ def _make_test_object(
 
         diffraction_cloud = self._make_diffraction_cloud(sq, q_max, [0, 0, 0])
 
-        test_object[:, :, :, 0, 0, 0] = diffraction_cloud.sum()
+        test_object[:, :, :, 0, 0, 0] = copy_to_device(diffraction_cloud.sum(), storage)
 
         for a0 in range(num):
             s1 = xp_storage.random.randint(r, sx - r)