mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
337 lines
12 KiB
Python
337 lines
12 KiB
Python
""" This module contains various to simulate stages in the probe-sample interaction
|
|
|
|
All the tools here are designed to work with automatic differentiation. Each
|
|
function simulates some aspect of an interaction model that can be used
|
|
for ptychographic reconstruction.
|
|
"""
|
|
|
|
from __future__ import division, print_function, absolute_import
|
|
|
|
from CDTools.tools.cmath import *
|
|
import torch as t
|
|
import numpy as np
|
|
|
|
|
|
__all__ = ['translations_to_pixel', 'pixel_to_translations',
|
|
'ptycho_2D_round','ptycho_2D_linear','ptycho_2D_sinc']
|
|
|
|
|
|
|
|
def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])):
|
|
"""Takes real space translations and outputs them in pixel space
|
|
|
|
This works for any 2D ptychography geometry. It takes in
|
|
A set of translations in (x,y) space and outputs the same translations
|
|
in internal pixel units perpendicular to the detector.
|
|
|
|
It uses information on the wavefield basis and, if defined, the
|
|
sample normal, to perform the conversion.
|
|
|
|
The assumed geometry is incoming radiation with a wavevector parallel
|
|
to the +z axis, [0,0,1]. The default sample orientation has a surface
|
|
normal parallel to this direction
|
|
|
|
Parameters
|
|
----------
|
|
basis : torch.Tensor
|
|
The real space basis the wavefields are defined in
|
|
translations : torch.Tensor
|
|
A Jx3 stack of real-space translations, or a single translation
|
|
surface_normal : torch.Tensor
|
|
Optional, the sample's surface normal
|
|
|
|
Returns
|
|
-------
|
|
pixel_translations : torch.Tensor
|
|
A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation
|
|
"""
|
|
|
|
projection_1 = t.Tensor([[1,0,0],
|
|
[0,1,0],
|
|
[0,0,0]]).to(device=translations.device,dtype=translations.dtype)
|
|
projection_2 = t.inverse(t.Tensor([[1,0,0],
|
|
[0,1,0],
|
|
-surface_normal/
|
|
surface_normal[2]])).to(device=translations.device,dtype=translations.dtype)
|
|
basis_vectors_inv = t.pinverse(basis).to(device=translations.device,
|
|
dtype=translations.dtype)
|
|
projection = t.mm(basis_vectors_inv,
|
|
t.mm(projection_2,projection_1))
|
|
projection = projection.t()
|
|
|
|
single_translation = False
|
|
if len(translations.shape) == 1:
|
|
translations = translations[None,:]
|
|
single_translation = True
|
|
|
|
pixel_translations = t.mm(translations, projection)
|
|
|
|
if single_translation:
|
|
return pixel_translations[0]
|
|
else:
|
|
return pixel_translations
|
|
|
|
|
|
def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,0,1])):
|
|
"""Takes pixel-space translations and outputs them in real space
|
|
|
|
This works for any 2D ptychography geometry. It takes in
|
|
A set of internal pixel unit translations in (i,j) space and
|
|
outputs the same translations real (x,y) space
|
|
|
|
It uses information on the wavefield basis and, if defined, the
|
|
sample normal, to perform the conversion.
|
|
|
|
The assumed geometry is incoming radiation with a wavevector parallel
|
|
to the +z axis, [0,0,1]. The default sample orientation has a surface
|
|
normal parallel to this direction. Because of this, the z direction
|
|
translation is always set to zero in the conversion
|
|
|
|
Parameters
|
|
----------
|
|
basis : torch.Tensor
|
|
The real space basis the wavefields are defined in
|
|
translations : torch.Tensor
|
|
A Jx2 stack of pixel-space translations, or a single translation
|
|
surface_normal : torch.Tensor
|
|
Optional, the sample's surface normal
|
|
|
|
Returns
|
|
-------
|
|
real_translations : torch.Tensor
|
|
A Jx3 stack of real-space translations, or a single translation
|
|
"""
|
|
projection_1 = t.Tensor([[1,0,0],
|
|
[0,1,0],
|
|
[0,0,0]]).to(device=basis.device,dtype=basis.dtype)
|
|
projection_2 = t.inverse(t.Tensor([[1,0,0],
|
|
[0,1,0],
|
|
-surface_normal/
|
|
surface_normal[2]])).to(device=basis.device,dtype=basis.dtype)
|
|
basis_vectors_inv = t.pinverse(basis)
|
|
projection = t.mm(basis_vectors_inv,
|
|
t.mm(projection_2,projection_1))
|
|
# Literally just need the pseudoinverse of the projection we used to go
|
|
# the other way
|
|
projection = t.pinverse(projection).t()
|
|
|
|
single_translation = False
|
|
if len(pixel_translations.shape) == 1:
|
|
pixel_translations = pixel_translations[None,:]
|
|
single_translation = True
|
|
|
|
translations = t.mm(pixel_translations, projection)
|
|
|
|
if single_translation:
|
|
return translations[0]
|
|
else:
|
|
return translations
|
|
|
|
|
|
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
|
|
|
|
Parameters
|
|
----------
|
|
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
|
|
-------
|
|
exit_waves : torch.Tensor
|
|
An NxMxL tensor of the calculated exit waves
|
|
"""
|
|
single_translation = False
|
|
if translations.dim() == 1:
|
|
translations = translations[None,:]
|
|
single_translation = True
|
|
|
|
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]
|
|
if single_translation:
|
|
return [cmult(probe,selection) for selection in selections][0]
|
|
else:
|
|
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
|
|
|
|
Parameters
|
|
----------
|
|
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
|
|
Default True, Whether to subpixel shift the probe or object
|
|
|
|
Returns
|
|
-------
|
|
exit_waves : torch.Tensor
|
|
An NxMxL tensor of the calculated exit waves
|
|
"""
|
|
single_translation = False
|
|
if translations.dim() == 1:
|
|
translations = translations[None,:]
|
|
single_translation = True
|
|
|
|
# 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))
|
|
|
|
if single_translation:
|
|
return exit_waves[0]
|
|
else:
|
|
return t.stack(exit_waves)
|
|
|
|
|
|
|
|
|
|
def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10):
|
|
"""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 sinc
|
|
interpolation (done via multiplication with a complex exponential
|
|
in Fourier space)
|
|
|
|
If shift_probe is True, it applies the subpixel shift to the probe,
|
|
otherwise the subpixel shift is applied to the object
|
|
|
|
Parameters
|
|
----------
|
|
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
|
|
Default True, Whether to subpixel shift the probe or object
|
|
padding : int
|
|
Default 10, if shifting the object, the padding to apply to the object to avoid circular shift effects
|
|
|
|
Returns
|
|
-------
|
|
exit_waves : torch.Tensor
|
|
An NxMxL tensor of the calculated exit waves
|
|
"""
|
|
single_translation = False
|
|
if translations.dim() == 1:
|
|
translations = translations[None,:]
|
|
single_translation = True
|
|
|
|
# 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:
|
|
i = t.arange(probe.shape[0]) - probe.shape[0]//2
|
|
j = t.arange(probe.shape[1]) - probe.shape[1]//2
|
|
I,J = t.meshgrid(i,j)
|
|
I = 2 * np.pi * I.to(t.float32) / probe.shape[0]
|
|
J = 2 * np.pi * J.to(t.float32) / probe.shape[1]
|
|
I = I.to(dtype=probe.dtype,device=probe.device)
|
|
J = J.to(dtype=probe.dtype,device=probe.device)
|
|
|
|
for tr, sp in zip(integer_translations,
|
|
subpixel_translations):
|
|
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)
|
|
|
|
obj_slice = obj[tr[0]:tr[0]+probe.shape[0],
|
|
tr[1]:tr[1]+probe.shape[1]]
|
|
|
|
exit_waves.append(cmult(shifted_probe, obj_slice))
|
|
|
|
else:
|
|
raise NotImplementedError('Object shift not yet implemented')
|
|
|
|
if single_translation:
|
|
return exit_waves[0]
|
|
else:
|
|
return t.stack(exit_waves)
|