cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-15 07:32:38 +02:00
Code Issues Packages Projects Releases Wiki Activity

Update S matrix ptychography and also update the S matrix interaction function to use a more interpretable definition of the S matrix

Browse Source
This commit is contained in:
Abe Levitan
2020-10-30 14:52:17 -04:00
parent aea0ff9bf8
commit 38345da454
2 changed files with 21 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/s_matrix_ptycho.py
+2 -1
View File
@@ -358,8 +358,9 @@ class SMatrixPtycho(CDIModel):
s_matrix = cmath.torch_to_complex(self.s_matrix.detach().cpu())
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
wavelength = self.wavelength.cpu().numpy()
return {'basis':basis, 'translation':translations,
'probe':probe,'s_matrix':s_matrix,
'background':background,
'weights':weights}
'weights':weights, 'wavelength':wavelength}
CDTools/tools/interactions.py
+19 -8
View File
@@ -440,7 +440,8 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
format. Currently, each pixel in the latter two dimensions index a location
on the input wavefield, and the first two indexes index differences from
that pixel. It is easier to interpret the resulting matrix though if the
latter two indices index locations in the output plane.
latter two indices index locations in the output plane. NOTE: I believe
this change has now been made
Parameters
----------
@@ -489,20 +490,30 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
fft_probe = fftshift(t.fft(probe, 2))
shifted_fft_probe = cmult(fft_probe, expi(-sp[0]*I - sp[1]*J))
shifted_probe = t.ifft(ifftshift(shifted_fft_probe),2)
s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[0]+2*B,
tr[1]:tr[1]+probe.shape[1]+2*B]
output = t.zeros([s_matrix_slice.shape[2]+2*B,
s_matrix_slice.shape[3]+2*B,2]).to(
output = t.zeros([probe.shape[0]+2*B,probe.shape[1]+2*B,2]).to(
device=s_matrix_slice.device,
dtype=s_matrix_slice.dtype)
for i in range(s_matrix.shape[0]):
for j in range(s_matrix.shape[1]):
output[i:i+probe.shape[0],j:j+probe.shape[1]] += \
cmult(shifted_probe, s_matrix_slice[i,j,:,:,:])
output [i:i+probe.shape[0],j:j+probe.shape[1]] += \
cmult(shifted_probe, s_matrix_slice[i,j,i:i+probe.shape[0],j:j+probe.shape[1],:])
#output = t.zeros([s_matrix_slice.shape[2]+2*B,
# s_matrix_slice.shape[3]+2*B,2]).to(
# device=s_matrix_slice.device,
# dtype=s_matrix_slice.dtype)
#for i in range(s_matrix.shape[0]):
# for j in range(s_matrix.shape[1]):
# output[i:i+probe.shape[0],j:j+probe.shape[1]] += \
# cmult(shifted_probe, s_matrix_slice[i,j,:,:,:])
exit_waves.append(output)
#exit_waves.append(cmult(shifted_probe, obj_slice))