From 4714ce7088efad7bc7e9a721916edb9e0305cd5a Mon Sep 17 00:00:00 2001 From: Stephanie Ribet Date: Fri, 6 Sep 2024 15:18:54 -0700 Subject: [PATCH] silly gpu bugs for sim --- py4DSTEM/tomography/tomography.py | 84 +++---------------------------- 1 file changed, 6 insertions(+), 78 deletions(-) diff --git a/py4DSTEM/tomography/tomography.py b/py4DSTEM/tomography/tomography.py index b34ee36e1..222e523d4 100644 --- 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)