mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-19 17:12:10 +02:00
120 lines
4.6 KiB
Python
120 lines
4.6 KiB
Python
from __future__ import division, print_function, absolute_import
|
|
|
|
from CDTools.tools.cmath import *
|
|
import torch as t
|
|
|
|
#
|
|
# This file will host tools to turn various kinds of model information
|
|
# (probe, 2D object, 3D object, etc) into exit waves leaving the sample
|
|
# area.
|
|
#
|
|
|
|
|
|
def ptycho_2D_round(probe, obj, translations):
|
|
"""Returns a stack of exit waves without accounting for subpixel shifts
|
|
|
|
This function returns a collection of exit waves, with the first
|
|
dimension as the translation index and the final dimensions
|
|
corresponding to the detector. The exit waves are calculated by
|
|
shifting the probe by the rounded value of the translation
|
|
|
|
Args:
|
|
probe (torch.Tensor) : An MxL probe function for the exit waves
|
|
object (torch.Tensor) : The object function to be probed
|
|
translations (torch.Tensor) : The Nx2 array of (i,j) translations to simulate
|
|
|
|
Returns:
|
|
torch.Tensor : An NxMxL tensor of the calculated exit waves
|
|
"""
|
|
integer_translations = t.round(translations).to(dtype=t.int32)
|
|
selections = [obj[tr[0]:tr[0]+probe.shape[0],
|
|
tr[1]:tr[1]+probe.shape[1]]
|
|
for tr in integer_translations]
|
|
return t.stack([cmult(probe,selection) for selection in selections])
|
|
|
|
|
|
|
|
|
|
def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
|
|
"""Returns a stack of exit waves accounting for subpixel shifts
|
|
|
|
This function returns a collection of exit waves, with the first
|
|
dimension as the translation index and the final dimensions
|
|
corresponding to the detector. The exit waves are calculated by
|
|
shifting the probe with each translation in turn, using linear
|
|
interpolation to combine the results
|
|
|
|
If shift_probe is True, it applies the subpixel shift to the probe,
|
|
otherwise the subpixel shift is applied to the object
|
|
|
|
Args:
|
|
probe (torch.Tensor) : An MxL probe function for the exit waves
|
|
object (torch.Tensor) : The object function to be probed
|
|
translations (torch.Tensor) : The Nx2 array of translations to simulate
|
|
shift_probe (bool) : Whether to subpixel shift the probe or object
|
|
Returns:
|
|
torch.Tensor : An NxMxL tensor of the calculated exit waves
|
|
"""
|
|
|
|
# Separate the translations into a part that chooses the window
|
|
# And a part that defines the windowing function
|
|
integer_translations = t.floor(translations)
|
|
subpixel_translations = translations - integer_translations
|
|
integer_translations = integer_translations.to(dtype=t.int32)
|
|
|
|
exit_waves = []
|
|
if shift_probe:
|
|
for tr, sp in zip(integer_translations,
|
|
subpixel_translations):
|
|
# This isn't perfectly symmetric but I think it's okay for now
|
|
# It should get the job done
|
|
# Basically, we shift the probe's position by a subpixel (i,j),
|
|
# rolling the edges of the array, and use that to multiply
|
|
# by the object
|
|
sel00 = probe[:,:]
|
|
sel01 = t.cat((probe[:,-1:],probe[:,:-1]),dim=1)
|
|
sel10 = t.cat((probe[-1:,:],probe[:-1,:]),dim=0)
|
|
sel11 = t.cat((sel01[-1:,:],sel01[:-1,:]),dim=0)
|
|
|
|
selection = sel00 * (1-sp[0])*(1-sp[1]) + \
|
|
sel10 * sp[0]*(1-sp[1]) + \
|
|
sel01 * (1-sp[0])*sp[1] + \
|
|
sel11 * sp[0]*sp[1]
|
|
|
|
obj_slice = obj[tr[0]:tr[0]+probe.shape[0],
|
|
tr[1]:tr[1]+probe.shape[1]]
|
|
|
|
exit_waves.append(cmult(selection,obj_slice))
|
|
else:
|
|
for tr, sp in zip(integer_translations,
|
|
subpixel_translations):
|
|
#
|
|
# Here we subpixel shift the object by (-i,-j) after
|
|
# slicing out the correct translation of the probe
|
|
#
|
|
|
|
sel00 = obj[tr[0]:tr[0]+probe.shape[0],
|
|
tr[1]:tr[1]+probe.shape[1]]
|
|
|
|
sel01 = obj[tr[0]:tr[0]+probe.shape[0],
|
|
tr[1]+1:tr[1]+1+probe.shape[1]]
|
|
|
|
sel10 = obj[tr[0]+1:tr[0]+1+probe.shape[0],
|
|
tr[1]:tr[1]+probe.shape[1]]
|
|
|
|
sel11 = obj[tr[0]+1:tr[0]+1+probe.shape[0],
|
|
tr[1]+1:tr[1]+1+probe.shape[1]]
|
|
|
|
selection = sel00 * (1-sp[0])*(1-sp[1]) + \
|
|
sel01 * (1-sp[0])*sp[1] + \
|
|
sel10 * sp[0]*(1-sp[1]) + \
|
|
sel11 * sp[0]*sp[1]
|
|
|
|
exit_waves.append(cmult(probe,selection))
|
|
|
|
return t.stack(exit_waves)
|
|
|
|
|
|
|
|
#TODO: Implement a sinc-interpolated shift using a fourier space shifting op
|