better memory management of unused stuff

[tpchuckles]

I can run tests/15_molecular_dynamics.py on a v100 GPU, which has 16 GB RAM, and it runs MD, then OOMs on the multislice.

I then re-run the exact same code, and now since it is importing the trajectory instead of generating it in MD, the multislice runs fine. (then OOMs on TACAW).

in the mean time, I have added (commit 590323b) WFData.crop just to reduce the size of WFData._array so we can FFT it and pass the tests for PR 31. but anything we can aggressively throw out in the future would help.

[tpchuckles]

WFData.crop does not appear to "free up" any ram (it simply gives you a smaller area before you convert to TACAW). https://pythonspeed.com/articles/numpy-memory-views/

a = torch.rand(size=(1000,1000,1000))

uses 3.88 GB ram on my macbook

a=a[200:-200,200:-200,200:-200]

crops to 20% the volume (0.6^3), but 3.88 GB is still used.

a=a[200:-200,200:-200,200:-200]

again takes us to 0.8% of the starting volume (0.2^3), but 3.88 GB is still used.

del a

frees up that 3.88 GB

the same thing happens with numpy by the way. I have no idea how torch or numpy allocate memory, nor how garbage collection is supposed to work here! my guess is, we're not really cropping the dataset in memory, but re-referencing a smaller portion, so the whole (uncropped) array/tensor is still kept in ram.

a = torch.zeros(size=(600,600,600))+a[200:-200,200:-200,200:-200]

appears to work and RAM usage drops to 977 MB (a new smaller tensor is created and populated from the old one. it takes longer to create and populate (perceptible delay, rather than instant), but RAM is freed)

[tpchuckles]

Here are the results using memory_profiler for Calculators and 04_haadf

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
   100  352.223 MiB  352.223 MiB           1       @profile
   101                                             def setup(
   102                                                 self,
   103                                                 trajectory: Trajectory,
   104                                                 aperture: float = 0.0,
   105                                                 voltage_eV: float = 60e3,
   106                                                 defocus: float = 0.0,
   107                                                 slice_thickness: float = 0.5,
   108                                                 sampling: float = 0.1,
   109                                                 probe_xs: Optional[List[float]] = None,
   110                                                 probe_ys: Optional[List[float]] = None,
   111                                                 probe_positions: Optional[List[Tuple[float, float]]] = None,
   112                                                 batch_size: int = 10,
   113                                                 save_path: Optional[Path] = None,
   114                                                 cleanup_temp_files: bool = False,
   115                                                 slice_axis: int = 2,
   116                                                 cache_levels: list = ["exitwaves"], # options include: exitwaves, slices, potentials (this replaces store_all_slices)
   117                                                 max_kx = np.inf,
   118                                                 max_ky = np.inf,
   119                                             ):
   120                                                 """
   121                                                 Set up multislice simulation using PyTorch acceleration.
   122                                                 
   123                                                 Args:
   124                                                     trajectory: Input trajectory data
   125                                                     aperture: Objective aperture semi-angle in mrad
   126                                                     voltage_eV: Accelerating voltage in eV
   127                                                     defocus: Defocus in Angstroms (not implemented yet)
   128                                                     slice_thickness: Thickness of each slice in Angstroms
   129                                                     sampling: Sampling rate in Angstroms per pixel
   130                                                     probe_positions: List of (x,y) probe positions in Angstroms
   131                                                     batch_size: Number of frames to process at once
   132                                                     save_path: Optional path to save wave function data
   133                                                     cleanup_temp_files: Whether to delete temp files after loading
   134                                                     store_all_slices: If True, store wavefunction at each slice for 3D visualization
   135                                                 """
   136                                         
   137  352.230 MiB    0.008 MiB           1           self.trajectory = trajectory
   138  352.230 MiB    0.000 MiB           1           self.aperture = aperture
   139  352.234 MiB    0.004 MiB           1           self.voltage_eV = voltage_eV
   140  352.234 MiB    0.000 MiB           1           self.defocus = defocus
   141  352.234 MiB    0.000 MiB           1           self.slice_thickness = slice_thickness
   142  352.234 MiB    0.000 MiB           1           self.sampling = sampling
   143  352.234 MiB    0.000 MiB           1           self.probe_xs = probe_xs
   144  352.234 MiB    0.000 MiB           1           self.probe_ys = probe_ys
   145  352.234 MiB    0.000 MiB           1           self.probe_positions = probe_positions
   146  352.234 MiB    0.000 MiB           1           self.save_path = save_path
   147  352.234 MiB    0.000 MiB           1           self.cleanup_temp_files = cleanup_temp_files
   148  352.234 MiB    0.000 MiB           1           self.slice_axis = slice_axis
   149  352.234 MiB    0.000 MiB           1           self.cache_levels = cache_levels
   150  352.234 MiB    0.000 MiB           1           self.max_kx = max_kx
   151  352.234 MiB    0.000 MiB           1           self.max_ky = max_ky
   152                                         
   153                                                 # Generate cache key and setup output directory
   154  352.273 MiB    0.039 MiB           2           cache_key = self._generate_cache_key(trajectory, aperture, voltage_eV,
   155  352.234 MiB    0.000 MiB           1                                              slice_thickness, sampling, probe_positions)
   156                                                 #print(cache_key)
   157  352.273 MiB    0.000 MiB           1           self.output_dir = Path("psi_data/" + ("torch" if TORCH_AVAILABLE else "numpy") + "_"+cache_key)
   158  352.273 MiB    0.000 MiB           1           self.output_dir.mkdir(parents=True, exist_ok=True)
   159                                                         
   160                                                 # Set up spatial grids
   161  352.281 MiB    0.008 MiB           1           xs,ys,zs,lx,ly,lz=gridFromTrajectory(trajectory,sampling=sampling,slice_thickness=slice_thickness)
   162  352.281 MiB    0.000 MiB           1           nx=len(xs) ; ny=len(ys) ; nz=len(zs)
   163  352.281 MiB    0.000 MiB           1           self.xs = xs ; self.ys = ys ; self.zs = zs
   164  352.281 MiB    0.000 MiB           1           self.lx = lx ; self.ly = ly ; self.lz = lz
   165  352.281 MiB    0.000 MiB           1           self.nx = nx ; self.ny = ny ; self.nz = nz
   166  352.281 MiB    0.000 MiB           1           self.dx = xs[1]-xs[0] ; self.dy = ys[1]-ys[0] ; self.dy = ys[1]-ys[0]
   167                                         
   168                                                 # calculate kxs kys here, so we can crop them, since we'll pre-allocate wavefunction_data below
   169  353.109 MiB    0.828 MiB           1           self.kxs = xp.fft.fftshift(xp.fft.fftfreq(self.nx, self.sampling))  # k-space in 1/Å
   170  353.113 MiB    0.004 MiB           1           self.kys = xp.fft.fftshift(xp.fft.fftfreq(self.ny, self.sampling))  # k-space in 1/Å
   171  353.387 MiB    0.273 MiB           1           self.i1 = xp.argwhere(self.kxs >= -max_kx)[0][0]   # first element >=
   172  353.457 MiB    0.070 MiB           1           self.i2 = xp.argwhere(self.kxs <= max_kx)[-1][0]+1 # last element <=, +1, so i1:i2 includes i2
   173  353.457 MiB    0.000 MiB           1           self.j1 = xp.argwhere(self.kys >= -max_ky)[0][0]
   174  353.457 MiB    0.000 MiB           1           self.j2 = xp.argwhere(self.kys <= max_ky)[-1][0]+1
   175  353.512 MiB    0.055 MiB           1           self.kxs = self.kxs[self.i1:self.i2]
   176  353.516 MiB    0.004 MiB           1           self.kys = self.kys[self.j1:self.j2]
   177  353.547 MiB    0.031 MiB           1           self.nx = self.i2 - self.i1 ; self.ny = self.j2 - self.j1 ; nx = self.nx ; ny = self.ny
   178                                         
   179                                                 # Preferred to pass probe_xs and probe_ys from which we will define a grid
   180  353.547 MiB    0.000 MiB           1           if self.probe_xs is not None and self.probe_ys is not None:
   181  353.574 MiB    0.027 MiB           1               x,y = np.meshgrid(self.probe_xs,self.probe_ys,indexing='ij')
   182  353.609 MiB    0.035 MiB           1               self.probe_positions = np.asarray(list(zip(x.flat,y.flat)))
   183  353.609 MiB    0.000 MiB           1               lx = len(self.probe_xs) ; ly = len(self.probe_ys)
   184                                         
   185                                                 # Set up default probe position if not provided
   186  353.609 MiB    0.000 MiB           1           if self.probe_positions is None:
   187                                                     self.probe_positions = [(lx/2, ly/2)]  # Center probe
   188                                                     self.probe_xs = [lx/2] ; self.probe_ys = [ly/2]
   189                                         
   190                                                 # Create probe on the correct device from the start
   191  356.594 MiB    2.984 MiB           1           self.base_probe = Probe(xs, ys, self.aperture, self.voltage_eV, device=self.device)
   192                                         
   193                                                 # Initialize storage for results
   194  356.594 MiB    0.000 MiB           1           self.n_frames = trajectory.n_frames
   195  356.594 MiB    0.000 MiB           1           self.n_probes = len(self.probe_positions)
   196                                         
   197                                                 # Set dtype based on the actual device we're using
   198  356.594 MiB    0.000 MiB           1           if TORCH_AVAILABLE and self.device is not None:
   199  356.594 MiB    0.000 MiB           1               if self.device.type == 'mps':
   200                                                         self.complex_dtype = torch.complex64
   201                                                         self.float_dtype = torch.float32
   202                                                     else:
   203  356.594 MiB    0.000 MiB           1                   self.complex_dtype = torch.complex128
   204  356.594 MiB    0.000 MiB           1                   self.float_dtype = torch.float64
   205                                                 else:
   206                                                     self.complex_dtype = np.complex128
   207                                                     self.float_dtype = np.float64
   208                                         
   209                                                 # Storage: [probe, frame, x, y, layer] - matches WFData expected format
   210  356.594 MiB    0.000 MiB           1           self.n_layers = nz if "slices" in cache_levels else 1
   211  356.594 MiB    0.000 MiB           1           if TORCH_AVAILABLE and self.device is not None:
   212  910.312 MiB  553.719 MiB           2               self.wavefunction_data = torch.zeros((self.n_probes, self.n_frames, nx, ny, self.n_layers),
   213  356.594 MiB    0.000 MiB           1                                                      dtype=self.complex_dtype, device=self.device)
   214                                                 else:
   215                                                     self.wavefunction_data = np.zeros((self.n_probes, self.n_frames, nx, ny, self.n_layers),
   216                                                                                        dtype=self.complex_dtype)


Filename: /Volumes/Alexandria/ORNL/Various Code/multislice/PySlice/asdf/tests/../src/pyslice/multislice/calculators.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
   250  910.320 MiB  910.320 MiB           1       @profile
   251                                             def run(self) -> WFData:
   252                                         
   253                                                 # Process frames with caching and multiprocessing
   254  910.320 MiB    0.000 MiB           1           total_start_time = time.time()
   255  910.320 MiB    0.000 MiB           1           frames_computed = 0
   256  910.320 MiB    0.000 MiB           1           frames_cached = 0
   257                                         
   258                                                 # quality of life sanity checks: user may have set things (e.g. probe array) with the wrong data type (e.g. numpy instead of tensor). let's try to catch and correct those here
   259  910.324 MiB    0.004 MiB           1           if isinstance(self.base_probe._array,np.ndarray) and TORCH_AVAILABLE:
   260                                                     self.base_probe._array = xp.tensor(self.base_probe._array)
   261                                         
   262                                                 # Process frames one at a time with tqdm progress tracking
   263 2501.934 MiB    0.555 MiB           2           with tqdm(total=self.n_frames, desc="Processing frames", unit="frame") as pbar:
   264 2501.934 MiB    0.000 MiB           4               for frame_idx in range(self.n_frames):
   265 2445.090 MiB    0.000 MiB           3                   cache_file = self.output_dir / f"frame_{frame_idx}.npy"
   266 2445.090 MiB    0.000 MiB           3                   positions = self.trajectory.positions[frame_idx]
   267 2445.090 MiB    0.000 MiB           3                   atom_types = self.trajectory.atom_types
   268                                                         
   269 2445.090 MiB    0.000 MiB           6                   args = [ frame_idx, positions, atom_types, self.xs, self.ys, self.zs,
   270 2445.090 MiB    0.000 MiB           3                          self.aperture, self.voltage_eV, self.base_probe, self.probe_positions, self.element_map,
   271 2445.090 MiB    0.000 MiB           3                          cache_file, self.cache_levels, self.slice_axis, self.device ]
   272                                                         
   273                                                         # Process frame
   274 2445.090 MiB    0.000 MiB           3                   if frame_idx == 0 and self.n_frames == 1:
   275                                                             args[0] = -1
   276                                         
   277 2501.934 MiB 7165.824 MiB           3                   frame_idx_result, frame_data, was_cached = _process_frame_worker_torch(args)
   278                                                         
   279                                                         # crop frame's diffraction image
   280 2501.934 MiB    0.020 MiB           3                   frame_data = frame_data[:,self.i1:self.i2,self.j1:self.j2,:,:]
   281                                         
   282                                                         # Store result
   283 2501.934 MiB    0.000 MiB         675                   for probe_idx in range(self.n_probes):
   284 2501.934 MiB    0.000 MiB         672                       if "slices" in self.cache_levels:
   285                                                                 # frame_data shape: (n_probes, nx, ny, n_slices, 1)
   286                                                                 self.wavefunction_data[probe_idx, frame_idx, :, :, :] = frame_data[probe_idx, :, :, :, 0]
   287                                                             else:
   288 2501.934 MiB    0.000 MiB         672                           self.wavefunction_data[probe_idx, frame_idx, :, :, 0] = frame_data[probe_idx, :, :, 0, 0]
   289                                                         
   290 2501.934 MiB    0.000 MiB           3                   if was_cached:
   291                                                             frames_cached += 1
   292                                                         else:
   293 2501.934 MiB    0.000 MiB           3                       frames_computed += 1
   294                                                         
   295                                                         # Update progress bar for this frame
   296 2501.934 MiB    0.008 MiB           3                   pbar.update(1)
   297                                                 
   298 2501.934 MiB    0.000 MiB           1           total_time = time.time() - total_start_time
   299 2501.934 MiB    0.000 MiB           1           logger.info(f"Simulation completed in {total_time:.2f}s ({frames_computed} computed, {frames_cached} cached)")
   300                                                 
   301                                                 # Create metadata
   302 2501.934 MiB    0.000 MiB           1           params = {
   303 2501.934 MiB    0.000 MiB           1               'aperture': self.aperture,
   304 2501.934 MiB    0.000 MiB           1               'voltage_eV': self.voltage_eV,
   305 2501.934 MiB    0.000 MiB           1               'defocus': self.defocus,
   306 2501.934 MiB    0.000 MiB           1               'slice_thickness': self.slice_thickness,
   307 2501.934 MiB    0.000 MiB           1               'sampling': self.sampling,
   308 2501.934 MiB    0.000 MiB           1               'grid_shape': (self.nx, self.ny, self.nz),
   309 2501.934 MiB    0.000 MiB           1               'box_size': (self.lx, self.ly, self.lz),
   310 2501.934 MiB    0.000 MiB           1               'n_atoms': self.trajectory.n_atoms,
   311 2501.934 MiB    0.000 MiB           1               'calculator': 'MultisliceCalculator'
   312                                                 }
   313                                                 
   314                                                 # Create coordinate arrays for output
   315                                                 # Note: WFData expects (probe_positions, time, kx, ky, layer) format
   316                                                 # Create k-space coordinates to match expected format (same as AbTem)
   317                                                 # TWP: If we're not going to also provide a shifted/etc reciprocal_array, we shouldn't shift the kxs
   318                                                 #kxs = xp.fft.fftfreq(self.nx, d=self.dx)
   319                                                 #kys = xp.fft.fftfreq(self.ny, d=self.dy)
   320                                                 #kxs = xp.fft.fftshift(xp.fft.fftfreq(self.nx, self.sampling))  # k-space in 1/Å MOVING TO INIT SO WE CAN CROP ON-THE-FLY
   321                                                 #kys = xp.fft.fftshift(xp.fft.fftfreq(self.ny, self.sampling))  # k-space in 1/Å
   322 2501.934 MiB    0.000 MiB           1           time_array = np.arange(self.n_frames) * self.trajectory.timestep  # Time array in ps
   323 2501.934 MiB    0.000 MiB           1           layer_array = np.arange(self.nz) if "slices" in self.cache_levels else np.array([0])  # Layer indices
   324                                                 
   325                                                 # Package results
   326 2501.941 MiB    0.008 MiB           2           wf_data = WFData(
   327 2501.934 MiB    0.000 MiB           1               probe_positions=self.probe_positions,
   328 2501.934 MiB    0.000 MiB           1               probe_xs=self.probe_xs,
   329 2501.934 MiB    0.000 MiB           1               probe_ys=self.probe_ys,
   330 2501.934 MiB    0.000 MiB           1               time=time_array,
   331 2501.934 MiB    0.000 MiB           1               kxs=self.kxs,
   332 2501.934 MiB    0.000 MiB           1               kys=self.kys,
   333 2501.934 MiB    0.000 MiB           1               xs=self.xs,
   334 2501.934 MiB    0.000 MiB           1               ys=self.ys,
   335 2501.934 MiB    0.000 MiB           1               layer=layer_array,
   336 2501.934 MiB    0.000 MiB           1               array=self.wavefunction_data,
   337 2501.934 MiB    0.000 MiB           1               probe=self.base_probe,
   338 2501.934 MiB    0.000 MiB           1               cache_dir=self.output_dir
   339                                                 )
   340                                                 
   341                                                 # Handle cleanup
   342 2501.941 MiB    0.000 MiB           1           if self.cleanup_temp_files:
   343                                                     logger.info("Cleaning up cache files...")
   344                                                     for frame_idx in range(n_frames):
   345                                                         cache_file = self.output_dir / f"frame_{frame_idx}.npy"
   346                                                         if cache_file.exists():
   347                                                             cache_file.unlink()
   348                                                     try:
   349                                                         self.output_dir.rmdir()
   350                                                     except OSError:
   351                                                         pass
   352                                                 else:
   353 2501.941 MiB    0.000 MiB           1               logger.info(f"Cache files saved in: {self.output_dir}")
   354                                                 
   355                                                 # Save if requested - psi files already saved during processing
   356                                                 
   357 2501.941 MiB    0.000 MiB           1           return wf_data


Filename: /Volumes/Alexandria/ORNL/Various Code/multislice/PySlice/asdf/tests/../src/pyslice/multislice/calculators.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
   361 2445.090 MiB 5574.797 MiB           3   @profile
   362                                         def _process_frame_worker_torch(args):
   363 2445.090 MiB    0.000 MiB           3       frame_idx, positions, atom_types, xs, ys, zs, aperture, eV, probe, probe_positions, element_map, cache_file, cache_levels , slice_axis, device = args
   364                                         
   365 2445.090 MiB    0.004 MiB           3       if cache_file.exists() and "exitwaves" in cache_levels:
   366                                                 global logging_tracker
   367                                                 if "cache_exists" not in logging_tracker:
   368                                                     logging_tracker.append("cache_exists")
   369                                                     logging.warning("One or more frames reloaded from cache: "+str(cache_file.parent))
   370                                                 return frame_idx, xp.asarray(np.load(cache_file)), True # if always saving as numpy, then must cast to torch array if re-reading cache file back in
   371                                         
   372                                             # Use the device passed from the calculator, or auto-detect if None
   373 2445.090 MiB    0.000 MiB           3       if TORCH_AVAILABLE:
   374 2445.090 MiB    0.000 MiB           3           if device is not None:
   375 2445.090 MiB    0.000 MiB           3               worker_device = device
   376                                                 else:
   377                                                     worker_device = torch.device('cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu'))
   378                                         
   379                                                 # Set dtype based on worker device
   380 2445.090 MiB    0.000 MiB           3           if worker_device.type == 'mps':
   381                                                     worker_complex_dtype = torch.complex64
   382                                                     worker_float_dtype = torch.float32
   383                                                 else:
   384 2445.090 MiB    0.000 MiB           3               worker_complex_dtype = torch.complex128
   385 2445.090 MiB    0.000 MiB           3               worker_float_dtype = torch.float64
   386                                             else:
   387                                                 worker_device = None
   388                                                 worker_complex_dtype = np.complex128
   389                                                 worker_float_dtype = np.float64
   390                                         
   391 2445.090 MiB    0.000 MiB           3       atom_type_names = []
   392 2445.090 MiB    0.000 MiB         603       for atom_type in atom_types:
   393 2445.090 MiB    0.000 MiB         600           if atom_type in element_map:
   394 2445.090 MiB    0.000 MiB         600               atom_type_names.append(element_map[atom_type])
   395                                                 else:
   396                                                     atom_type_names.append(atom_type)
   397                                             
   398                                             #try:
   399 2445.090 MiB    4.609 MiB           3       potential = Potential(xs, ys, zs, positions, atom_type_names, kind="kirkland", device=worker_device, slice_axis=slice_axis, progress=(frame_idx==-1), cache_dir=cache_file.parent if "potentials" in cache_levels else None, frame_idx = frame_idx)
   400                                         
   401 2445.090 MiB    0.000 MiB           3       n_probes = len(probe_positions)
   402 2445.090 MiB    0.000 MiB           3       nx, ny = len(xs), len(ys)
   403 2445.090 MiB    0.000 MiB           3       n_slices = len(zs)
   404                                         
   405 2501.621 MiB  477.195 MiB           3       batched_probes = create_batched_probes(probe, probe_positions, worker_device)
   406 2501.934 MiB  740.047 MiB           3       exit_waves_batch = Propagate(batched_probes, potential, worker_device, progress=(frame_idx==-1), onthefly=True, store_all_slices = ("slices" in cache_levels) )
   407                                         
   408 2501.934 MiB    0.000 MiB           3       if "slices" in cache_levels:
   409                                                 # exit_waves_batch shape: (n_slices, n_probes, nx, ny)
   410                                                 if TORCH_AVAILABLE and worker_device is not None:
   411                                                     frame_data = torch.zeros((n_probes, nx, ny, n_slices, 1), dtype=worker_complex_dtype, device=worker_device)
   412                                                 else:
   413                                                     frame_data = np.zeros((n_probes, nx, ny, n_slices, 1), dtype=worker_complex_dtype)
   414                                         
   415                                                 # Convert all slices to k-space
   416                                                 for slice_idx in range(n_slices):
   417                                                     slice_waves = exit_waves_batch[slice_idx, :, :, :]  # (n_probes, nx, ny)
   418                                                     kwarg = {"dim":(-2,-1)} if TORCH_AVAILABLE else {"axes":(-2,-1)}
   419                                                     waves_k = xp.fft.fft2(slice_waves, **kwarg)
   420                                                     diffraction_patterns = xp.fft.fftshift(waves_k, **kwarg)
   421                                         
   422                                                     # Store in frame_data
   423                                                     for i in range(n_probes):
   424                                                         frame_data[i, :, :, slice_idx, 0] = diffraction_patterns[i, :, :]
   425                                             else:
   426                                                 # exit_waves_batch shape: (n_probes, nx, ny)
   427 2501.934 MiB    0.000 MiB           3           if TORCH_AVAILABLE and worker_device is not None:
   428 2501.934 MiB    0.000 MiB           3               frame_data = torch.zeros((n_probes, nx, ny, 1, 1), dtype=worker_complex_dtype, device=worker_device)
   429                                                 else:
   430                                                     frame_data = np.zeros((n_probes, nx, ny, 1, 1), dtype=worker_complex_dtype)
   431                                         
   432                                                 # Convert all exit waves to k-space
   433 2501.934 MiB    0.000 MiB           3           kwarg = {"dim":(-2,-1)} if TORCH_AVAILABLE else {"axes":(-2,-1)}
   434 2501.934 MiB    0.000 MiB           3           exit_waves_k = xp.fft.fft2(exit_waves_batch, **kwarg)
   435 2501.934 MiB  369.141 MiB           3           diffraction_patterns = xp.fft.fftshift(exit_waves_k, **kwarg)
   436                                         
   437                                                 # Store results
   438 2501.934 MiB    0.004 MiB           3           frame_data[:, :, :, 0, 0] = diffraction_patterns #.cpu().numpy()
   439                                             #else:
   440                                             #    # Fallback to individual processing
   441                                             #    for probe_idx, (px, py) in enumerate(probe_positions):
   442                                             #        shifted_probe = probe.copy()
   443                                             #        
   444                                             #        probe_k = torch.fft.fft2(shifted_probe.array)
   445                                             #        
   446                                             #        kx_shift = torch.exp(2j * torch.pi * shifted_probe.kxs[:, None] * px)
   447                                             #        ky_shift = torch.exp(2j * torch.pi * shifted_probe.kys[None, :] * py)
   448                                             #       probe_k_shifted = probe_k * kx_shift * ky_shift
   449                                             #        
   450                                             #        shifted_probe.array = torch.fft.ifft2(probe_k_shifted)
   451                                             #        
   452                                             #        exit_wave_torch = PropagateTorch(shifted_probe, potential, worker_device)
   453                                             #        
   454                                             #        exit_wave_k = torch.fft.fft2(exit_wave_torch)
   455                                             #       diffraction_pattern = torch.fft.fftshift(exit_wave_k)
   456                                             #     
   457                                             #       frame_data[probe_idx, :, :, 0, 0] = diffraction_pattern.cpu().numpy()
   458                                             
   459                                             # Convert to CPU numpy array for saving
   460 2501.934 MiB    0.000 MiB           3       if TORCH_AVAILABLE and hasattr(frame_data, 'cpu'):
   461 2501.934 MiB    0.027 MiB           3           frame_data_cpu = frame_data.cpu().numpy()
   462                                             else:
   463                                                 frame_data_cpu = frame_data
   464                                         
   465 2501.934 MiB    0.000 MiB           3       if "exitwaves" in cache_levels:
   466                                                 np.save(cache_file, frame_data_cpu)
   467                                         
   468 2501.934 MiB    0.000 MiB           3       return frame_idx, frame_data, False
   469                                                 
   470                                             #except Exception as e:
   471                                             #    logger.error(f"Error processing frame {frame_idx} with PyTorch: {e}")
   472                                             #    from .potential import Potential
   473                                             #    from .multislice_npy import Probe, Propagate
   474                                             #    
   475                                             #    potential = Potential(xs, ys, zs, positions, atom_type_names, kind="kirkland")
   476                                             #    probe = Probe(xs, ys, aperture, eV)
   477                                             #    
   478                                             #    n_probes = len(probe_positions)
   479                                             #    nx, ny = len(xs), len(ys)
   480                                             ##    frame_data = np.zeros((n_probes, nx, ny, 1, 1), dtype=complex)
   481                                             #    
   482                                             #    for probe_idx, (px, py) in enumerate(probe_positions):
   483                                             #        exit_wave = Propagate(probe, potential)
   484                                             #        diffraction_pattern = np.fft.fftshift(np.fft.fft2(exit_wave))
   485                                             #        frame_data[probe_idx, :, :, 0, 0] = diffraction_pattern
   486                                             #    
   487                                             #    np.save(cache_file, frame_data)
   488                                             #    return frame_idx, frame_data, False

it's time to be honest with myself. i apparently don't know how memory management in python works.

but there may be some opportunity for improvement in here:
~ batched_probes is nprobes,nx,ny (where nprobes might be a lot).
~ exit_waves_batch, frame_data, exit_waves_k, diffraction_pattern are holding all the same data and more or less each used once. i don't know if there are any "in place" operations we could be doing instead of spawning copies of the data.