PRISM?

[tpchuckles]

I uploaded the code to staging. It's still not working, but I think it's because i'm fundamentally doing fourier resampling wrong. 

Here's a basic code example for fourier downsampling:

import numpy as np
import matplotlib.pyplot as plt
xs = np.linspace(0,10,1000)
ys = np.linspace(0,10,1000)
array = np.exp(-(xs[:,None]-3)**2/.5**2)*np.exp(-(ys[None,:]-7)**2/.5**2) # gaussian at 3,7
fig, ax = plt.subplots()
ax.imshow(array.T[::-1,:],extent=(0,10,0,10),cmap="inferno")
plt.show() # confirm it appears at 3,7
fft = np.fft.fft2(array)
fft2 = fft[::20,::20] # downsampling, 1000x1000 --> 50x50 
ifft2 = np.fft.ifft2(fft2)
fig, ax = plt.subplots()
ax.imshow(np.real(ifft2.T)[::-1,:],extent=(0,10,0,10),cmap="inferno") # looks nothing like our starting gaussian
plt.show()

[tpchuckles]

ah, sparsifying is wrong, cropping is right: https://en.wikipedia.org/wiki/Downsampling_(signal_processing)

# SPARSIFY
fft2 = fft[::10,::10] # downsampling, 1000x1000 --> 100x100
ifft2 = np.fft.ifft2(fft2)
fig, ax = plt.subplots()
ax.imshow(np.real(ifft2.T)[::-1,:],extent=(0,10,0,10),cmap="inferno") # looks nothing like our starting gaussian
plt.show()

# CROP
fft2=np.fft.ifftshift(np.fft.fftshift(fft)[450:550,450:550]) # CROPPING, 1000x1000 --> 100x100, shift required to keep low-freq
ifft2 = np.fft.ifft2(fft2)
fig, ax = plt.subplots()
ax.imshow(np.real(ifft2.T)[::-1,:],extent=(0,10,0,10),cmap="inferno") # looks nothing like our starting gaussian
plt.show()

[tpchuckles]

working implementation in staging.

turns out, building the probes is actually not the bottleneck (i thought the "build the probe in reciprocal, convert to real, FFT back to get our Fourier components" was gonna be slow, but not really. but this is why i implemented stay_reciprocal option for Probe.generate_single_probe and Probe.placeProbe. I also thought generating the full-sized probe (vs a cropped version) would be slow, but not really. but this led to the crop_reciprocal option for Probe.generate_single_probe and Probe.placeProbe).

the real sauce is in einsum for the summation, which is a yuuuuuge speedup.

also most recent commit includes some witchcraft to pass PrismProbe.calculateProbesFromS a reference to calculator's memmaped wavefunction_data so we can now einsum our million (lol) probes all in one go.

test 22_prism runs for 256x256 which takes about 5 minutes on cpu

[tpchuckles]

scaling still leaves some to be desired. full-sized hBN (15x12 nm) with 1024x1024 probes is going to take several hundred hours on gpu. that's a no go.

the system is 6x the size (was cropped to 10a 10b previously), so each exit wave has 36x as many k-points. could sparsifying k help, just keeping every nth?

edit: sparsifying diffraction space gets us down to ~20 minutes for 1024x1024. i'm not sure the exactly right way to do this (i just did [::kth]), but i assume it's legal since your 6x bigger system ended up with 6x as many k points which you don't really need. so if they're unwieldy, throw them out.

[tpchuckles]

on the fly ADF summation (Issue #35) would also help with ram. if we only care about the super ridiculously huge 1024x1024xNxM cube because we're going to sum later, then how about sum it now and don't blow your RAM or disk space

[tpchuckles]

prism code "works", but some optimizations are still desirable. currently, using mindk and loop_probes can give faster performance than prism.

closing this issue though, since it's implemented at least. we can worry about optimization later if there is interest.