mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
add projectors, propagators, and initializers
This commit is contained in:
Maddie Cain
@@ -2,3 +2,9 @@ from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools import cmath
|
||||
from CDTools.tools import losses
|
||||
from CDTools.tools import data
|
||||
from CDTools.tools import image_processing
|
||||
from CDTools.tools import initializers
|
||||
from CDTools.tools import losses
|
||||
from CDTools.tools import projectors
|
||||
from CDTools.tools import propagators
|
||||
|
||||
+14
-16
@@ -98,7 +98,7 @@ def cabs(a):
|
||||
|
||||
|
||||
def cphase(a):
|
||||
"""Returns the complex conjugate of a complex torch tensor
|
||||
"""Returns the phase of a complex torch tensor
|
||||
|
||||
Pytorch uses tensors with a final dimension of 2 to represent
|
||||
complex numbers. This calculates the elementwise complex phase
|
||||
@@ -112,8 +112,8 @@ def cphase(a):
|
||||
|
||||
"""
|
||||
return t.atan2(a[...,1],a[...,0])
|
||||
|
||||
|
||||
|
||||
|
||||
def cconj(a):
|
||||
"""Returns the complex conjugate of a complex torch tensor
|
||||
|
||||
@@ -147,7 +147,7 @@ def cmult(a,b):
|
||||
torch.Tensor : A tensor storing the elementwise product
|
||||
|
||||
"""
|
||||
|
||||
|
||||
real = a[...,0] * b[...,0] - a[...,1] * b[...,1]
|
||||
imag = a[...,0] * b[...,1] + a[...,1] * b[...,0]
|
||||
return t.stack((real,imag),dim=-1)
|
||||
@@ -171,7 +171,7 @@ def cdiv(a,b):
|
||||
return cmult(a, cconj(b)) / t.unsqueeze(cabssq(b),-1)
|
||||
|
||||
|
||||
|
||||
|
||||
#
|
||||
# Not entirely sure if these belong here, but heck with it.
|
||||
# We just need the ability to do fftshifts
|
||||
@@ -180,23 +180,23 @@ def cdiv(a,b):
|
||||
|
||||
def fftshift(array,dims=None):
|
||||
"""Drop-in torch replacement for scipy.fftpack.fftshift
|
||||
|
||||
|
||||
This maps a tensor, assumed to be the output of a fast Fourier
|
||||
transform, into a tensor whose zero-frequency element is at the
|
||||
center of the tensor instead of the start. It will by default shift
|
||||
every dimension in the tensor but the last (which is assumed to
|
||||
represent the complex number and be of dimension 2), but can shift
|
||||
any arbitrary set of dimensions.
|
||||
|
||||
|
||||
Args:
|
||||
array (torch.Tensor) : An array of data to be fftshifted
|
||||
dims (iterable) : A list of all dimensions to shift
|
||||
|
||||
Returns:
|
||||
torch.Tensor : fftshifted tensor
|
||||
|
||||
|
||||
"""
|
||||
|
||||
|
||||
if dims is None:
|
||||
dims=list(range(array.dim()))[:-1]
|
||||
for dim in dims:
|
||||
@@ -211,7 +211,7 @@ def fftshift(array,dims=None):
|
||||
|
||||
def ifftshift(array,dims=None):
|
||||
"""Drop-in torch replacement for scipy.fftpack.iftshift
|
||||
|
||||
|
||||
This maps a tensor, assumed to be the shifted output of a fast
|
||||
Fourier transform, into a tensor whose zero-frequency element is
|
||||
back at the start of the tensor instead of the center. It is the
|
||||
@@ -219,25 +219,23 @@ def ifftshift(array,dims=None):
|
||||
every dimension in the tensor but the last (which is assumed to
|
||||
represent the complex number and be of dimension 2), but can shift
|
||||
any arbitrary set of dimensions.
|
||||
|
||||
|
||||
Args:
|
||||
array (torch.Tensor) : An array of data to be ifftshifted
|
||||
dims (iterable) : A list of all dimensions to shift
|
||||
|
||||
Returns:
|
||||
torch.Tensor : ifftshifted tensor
|
||||
|
||||
|
||||
"""
|
||||
|
||||
|
||||
if dims is None:
|
||||
dims=list(range(array.dim()))[:-1]
|
||||
for dim in dims:
|
||||
length = array.size()[dim]
|
||||
cut_to = length // 2
|
||||
cut_len = length - cut_to
|
||||
|
||||
|
||||
array = t.cat((array.narrow(dim,cut_to,cut_len),
|
||||
array.narrow(dim,0,cut_to)), dim)
|
||||
return array
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
import numpy as np
|
||||
import torch as t
|
||||
@@ -0,0 +1,29 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
import numpy as np
|
||||
import torch as t
|
||||
|
||||
all = ['gaussian']
|
||||
|
||||
|
||||
def gaussian(shape, amplitude, sigma, center = None):
|
||||
"""Returns an array with a centered gaussian
|
||||
|
||||
Takes in the shape, amplitude, and standard deviation of a gaussian
|
||||
and returns an array with values corresponding to a two-dimensional gaussian function
|
||||
z = amplitude*exp(-(x-center[0])**2/sigma[0]**2+(y-center[1])**2/sigma[1]**2)
|
||||
Note that [0, 0] is taken to be at the upper left corner of the array.
|
||||
Default is centered at ((shape[0]-1)/2, (shape[1]-1)/2)) because x and y are zero-indexed.
|
||||
|
||||
Args:
|
||||
shape (array_like) : A 1x2 array-like object specifying the dimensions of the output array in the form (y shape, x shape)
|
||||
amplitude (float or int): The amplitude the gaussian to simulate
|
||||
sigma (array_like): A 1x2 array-like object specifying the x- and y- standard deviation of the gaussian in the form (y stdev, y stdev)
|
||||
center (array_like) : Optional 1x2 array-like object specifying the location of the center of the gaussian (y center, x center)
|
||||
|
||||
Returns:
|
||||
numpy.array : The real-valued gaussian array
|
||||
"""
|
||||
if center is None:
|
||||
center = ((shape[0]-1)/2, (shape[1]-1)/2)
|
||||
y, x = np.mgrid[:shape[0], :shape[1]]
|
||||
return amplitude*np.exp(-((x-center[1])/sigma[1])**2-((y-center[0])/sigma[0])**2)
|
||||
@@ -0,0 +1,63 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
from CDTools.tools.cmath import *
|
||||
import torch as t
|
||||
|
||||
__all__ = ['modulus', 'support']
|
||||
|
||||
def modulus(wavefront, intensities, mask = None):
|
||||
"""Implements the modulus constraint in torch
|
||||
|
||||
This accepts a torch tensor representing the propagated simulated wavefront(s),
|
||||
where the last dimension represents the real and imaginary components of
|
||||
the propagated wavefield(s). It projects the modulus of the diffraction pattern
|
||||
onto the modulus of the simulated wavefield.
|
||||
|
||||
It assumes that the wavefront is stored in an array
|
||||
[i,j] where i corresponds to the y-axis and j corresponds to the
|
||||
x-axis, with the origin following the CS standard of being in the
|
||||
upper right.
|
||||
|
||||
Args:
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
|
||||
intensities (torch.Tensor): The measured diffraction pattern(s) stored as an JxNxM stack of real tensors
|
||||
mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
|
||||
Returns:
|
||||
torch.Tensor : The JxNxMx2 propagated wavefield with corrected intensities
|
||||
"""
|
||||
# Calculate amplitudes from intensities
|
||||
amplitudes = intensities**.5
|
||||
# Normalize wavefront so the complex elements have modulus one
|
||||
abs = cabs(wavefront)
|
||||
wavefront[...,0]/=abs
|
||||
wavefront[...,1]/=abs
|
||||
if mask is None:
|
||||
# Replace amplitude of wavefront with measured amplitude
|
||||
wavefront[...,0]*=amplitudes
|
||||
wavefront[...,1]*=amplitudes
|
||||
return wavefront
|
||||
else:
|
||||
return wavefront[mask != 0]
|
||||
|
||||
|
||||
def support(wavefront, support):
|
||||
"""Implements the support constraint in torch
|
||||
|
||||
This accepts a torch tensor representing the propagated simulated wavefront(s),
|
||||
where the last dimension represents the real and imaginary components of
|
||||
the propagated wavefield(s). It projects the support of the imaged object
|
||||
onto the simulated wavefront via a mask.
|
||||
|
||||
It assumes that the wavefront is stored in an array
|
||||
[i,j] where i corresponds to the y-axis and j corresponds to the
|
||||
x-axis, with the origin following the CS standard of being in the
|
||||
upper right.
|
||||
|
||||
Args:
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
|
||||
mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
|
||||
Returns:
|
||||
torch.Tensor : The JxNxMx2 wavefield with the mask applied
|
||||
"""
|
||||
wavefront[...,0] *= support
|
||||
wavefront[...,1] *= support
|
||||
return wavefront
|
||||
@@ -0,0 +1,161 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
from CDTools.tools.cmath import *
|
||||
import torch as t
|
||||
|
||||
__all__ = ['far_field', 'near_field', 'inverse_far_field', 'inverse_near_field', 'get_exit_waves']
|
||||
|
||||
|
||||
def far_field(wavefront):
|
||||
"""Implements a far-field propagator in torch
|
||||
|
||||
This accepts a torch tensor, where the last dimension
|
||||
represents the real and imaginary components of the wavefield,
|
||||
and returns the far-field propagated version of it assuming it matches the
|
||||
detector dimensions. It assumes that the
|
||||
propagation is purely far-field, without checking that the geometry
|
||||
is consistent with that assumption.
|
||||
|
||||
|
||||
It also assumes that the real space wavefront is stored in an array
|
||||
[i,j] where i corresponds to the y-axis and j corresponds to the
|
||||
x-axis, with the origin following the CS standard of being in the
|
||||
upper right. The zero frequency component of the propagated wavefield is
|
||||
shifted to the center of the array.
|
||||
|
||||
Args:
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
|
||||
Returns:
|
||||
torch.Tensor : The JxNxMx2 propagated wavefield
|
||||
"""
|
||||
|
||||
return fftshift(t.fft(wavefront, 2))
|
||||
|
||||
def inverse_far_field(wavefront):
|
||||
"""Implements the inverse of the far-field propagator in torch
|
||||
|
||||
This accepts a torch tensor, where the last dimension
|
||||
represents the real and imaginary components of the propagated wavefield,
|
||||
and returns the un-propagated array.
|
||||
|
||||
It assumes that the real space wavefront is stored in an array
|
||||
[i,j] where i corresponds to the y-axis and j corresponds to the
|
||||
x-axis, with the origin following the CS standard of being in the
|
||||
upper right. The zero frequency component of the propagated wavefield is
|
||||
assumed to be the center of the array.
|
||||
|
||||
Args:
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts propagated to the far-field
|
||||
Returns:
|
||||
torch.Tensor : The JxNxMx2 exit wavefield
|
||||
"""
|
||||
return t.ifft(ifftshift(wavefront), 2)
|
||||
|
||||
|
||||
def generate_angular_spectrum_propagator(shape, spacing, wavelength, z):
|
||||
"""Generates an angular-spectrum based near-field propagator from experimental quantities
|
||||
|
||||
This function generates an angular-spectrum based near field
|
||||
propagator that will work on torch Tensors. The function is structured
|
||||
this way - to generate the propagator first - because the
|
||||
generation of the propagation mask is a bit expensive and if this
|
||||
propagator is used in a reconstruction program, then it will be best
|
||||
to calculate this mask once and close over it.
|
||||
|
||||
Args:
|
||||
shape (iterable) : The shape of the arrays to be propagated
|
||||
spacing (iterable) : The pixel size in each dimension of the arrays to be propagated
|
||||
wavelength (float) : The wavelength of light to simulate propagation of
|
||||
z (float) : The distance to simulate propagation over
|
||||
Returns:
|
||||
torch.Tensor : A propagation term which accounts for the phase change that each plane wave will undergo on its journey to the prediction plane.
|
||||
"""
|
||||
|
||||
ki = fftpack.fftfreq(shape[0],spacing[0])
|
||||
kj = fftpack.fftfreq(shape[1],spacing[1])
|
||||
Ki, Kj = np.meshgrid(ki,kj)
|
||||
propagator = np.exp(1j*np.sqrt((2*np.pi/wavelength)**2
|
||||
- Ki**2 - Kj**2) * z)
|
||||
propagator = complex_to_float(propagator).astype(np.float32)
|
||||
propagator = t.from_numpy(propagator).cuda()
|
||||
|
||||
return propagator
|
||||
|
||||
|
||||
def near_field(wavefront, angular_spectrum_propagator):
|
||||
"""This function accepts an 3d torch tensor, where the
|
||||
last dimension represents the real and imaginary components of
|
||||
the wavefield, and returns the near-field propagated version of it.
|
||||
|
||||
|
||||
Args:
|
||||
angular_spectrum_propagator (torch.Tensor) : The near field propagator
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
|
||||
Returns:
|
||||
function : The wavefront propagated to the near field
|
||||
"""
|
||||
|
||||
return t.ifft(angular_spectrum_propagator * t.fft(wavefront,2), 2)
|
||||
|
||||
|
||||
|
||||
def inverse_near_field(wavefront, angular_spectrum_propagator):
|
||||
"""This function accepts a 3d torch tensor, where the
|
||||
last dimension represents the real and imaginary components of
|
||||
the near-field propagated wavefield, and returns the exit wavefront via an inverse transformation.
|
||||
|
||||
|
||||
Args:
|
||||
angular_spectrum_propagator (torch.Tensor) : The pixel size in each dimension of the arrays to be propagated
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
|
||||
Returns:
|
||||
function : A function to propagate a torch tensor.
|
||||
"""
|
||||
return t.ifft(t.fft(wavefront,2) * angular_spectrum_propagator**-1, 2)
|
||||
|
||||
|
||||
|
||||
|
||||
def get_exit_waves(probe, object, translations):
|
||||
"""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 object with each translation in turn, using linear
|
||||
interpolation.
|
||||
Args:
|
||||
probe (torch.Tensor) : An MxM 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
|
||||
Returns:
|
||||
torch.Tensor : An NxMxM 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)
|
||||
|
||||
selections = []
|
||||
for tr, sp in zip(integer_translations,
|
||||
subpixel_translations):
|
||||
|
||||
sel00 = object[tr[0]:tr[0]+probe.shape[0],
|
||||
tr[1]:tr[1]+probe.shape[1]]
|
||||
|
||||
sel01 = object[tr[0]:tr[0]+probe.shape[0],
|
||||
tr[1]+1:tr[1]+1+probe.shape[1]]
|
||||
|
||||
sel10 = object[tr[0]+1:tr[0]+1+probe.shape[0],
|
||||
tr[1]:tr[1]+probe.shape[1]]
|
||||
|
||||
sel11 = object[tr[0]+1:tr[0]+1+probe.shape[0],
|
||||
tr[1]+1:tr[1]+1+probe.shape[1]]
|
||||
|
||||
selections.append(sel00 * (1-sp[0])*(1-sp[1]) + \
|
||||
sel01 * (1-sp[0])*sp[1] + \
|
||||
sel10 * sp[0]*(1-sp[1]) + \
|
||||
sel11 * sp[0]*sp[1])
|
||||
|
||||
return t.stack([cmult(probe,selection) for selection in selections])
|
||||
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools import tools
|
||||
@@ -0,0 +1,4 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools import cmath
|
||||
from CDTools.tools import losses
|
||||
@@ -0,0 +1,243 @@
|
||||
"""Contains basic functions for dealing with complex numbers in pytorch.
|
||||
|
||||
Since pytorch doesn't have built-in support for complex numbers, but the
|
||||
fast fourier transforms in pytorch assume a specific format for complex
|
||||
arrays, this module uses that format to store complex numbers. It exposes
|
||||
functions for converting between complex numpy arrays and torch tensors
|
||||
stored in that format, as well as basic complex math operations implemented
|
||||
on the torch tensors
|
||||
"""
|
||||
from __future__ import division, print_function, absolute_import
|
||||
import numpy as np
|
||||
import torch as t
|
||||
|
||||
|
||||
__all__ = ['complex_to_torch','torch_to_complex','cabssq','cabs','cconj',
|
||||
'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift']
|
||||
|
||||
|
||||
#
|
||||
# These define the conversions to and from this format
|
||||
#
|
||||
|
||||
def complex_to_torch(x):
|
||||
"""Maps a complex numpy array to a torch tensor
|
||||
|
||||
Pytorch uses tensors with a final dimension of 2 to represent
|
||||
complex numbers. This maps a complex type numpy array to a torch
|
||||
tensor following this convention
|
||||
|
||||
Args:
|
||||
x (array_like): A numpy array to convert
|
||||
|
||||
Returns:
|
||||
torch.Tensor : A torch tensor representation of that array
|
||||
|
||||
"""
|
||||
return t.from_numpy(np.stack((np.real(x),np.imag(x)),axis=-1))
|
||||
|
||||
|
||||
def torch_to_complex(x):
|
||||
"""Maps a torch tensor to the a complex numpy array
|
||||
|
||||
Pytorch uses tensors with a final dimension of 2 to represent
|
||||
complex numbers. This maps a torch tensor following that convention
|
||||
to the appropriate numpy complex array
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): A tensor to convert
|
||||
|
||||
Returns:
|
||||
np.array : A complex typed numpy array corresponding to the input
|
||||
|
||||
"""
|
||||
x = np.array(x)
|
||||
x = x[...,0] + x[...,1] * 1j
|
||||
return x
|
||||
|
||||
|
||||
#
|
||||
# And these define the basic operations on these arrays. Note that
|
||||
# multiplication between a complex valued and real valued pytorch
|
||||
# tensor will proceed as expected because of torch's broadcasting
|
||||
# and thus doesn't need it's own function
|
||||
#
|
||||
|
||||
def cabssq(a):
|
||||
"""Returns the square of the absolute value of a complex torch tensor
|
||||
|
||||
Pytorch uses tensors with a final dimension of 2 to represent
|
||||
complex numbers. This calculates the elementwise absolute value
|
||||
squared of any toch tensor following that standard.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): An input tensor
|
||||
|
||||
Returns:
|
||||
array_like : A tensor storing the elementwise absolute value squared
|
||||
|
||||
"""
|
||||
return a[...,0]**2 + a[...,1]**2
|
||||
|
||||
|
||||
def cabs(a):
|
||||
"""Returns the absolute value of a complex torch tensor
|
||||
|
||||
Pytorch uses tensors with a final dimension of 2 to represent
|
||||
complex numbers. This calculates the elementwise absolute value
|
||||
of any torch tensor following that standard.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): An input tensor
|
||||
|
||||
Returns:
|
||||
array_like : A tensor storing the elementwise absolute value
|
||||
|
||||
"""
|
||||
return t.sqrt(cabssq(a))
|
||||
|
||||
|
||||
def cphase(a):
|
||||
"""Returns the complex conjugate of a complex torch tensor
|
||||
|
||||
Pytorch uses tensors with a final dimension of 2 to represent
|
||||
complex numbers. This calculates the elementwise complex phase
|
||||
of any torch tensor following that standard.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): An input tensor
|
||||
|
||||
Returns:
|
||||
array_like : A tensor storing the elementwise phase
|
||||
|
||||
"""
|
||||
return t.atan2(a[...,1],a[...,0])
|
||||
|
||||
|
||||
def cconj(a):
|
||||
"""Returns the complex conjugate of a complex torch tensor
|
||||
|
||||
Pytorch uses tensors with a final dimension of 2 to represent
|
||||
complex numbers. This calculates the elementwise complex conjugate
|
||||
of any torch tensor following that standard.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): An input tensor
|
||||
|
||||
Returns:
|
||||
array_like : A tensor storing the elementwise complex conjugate
|
||||
|
||||
"""
|
||||
return t.stack((a[...,0],-a[...,1]),dim=-1)
|
||||
|
||||
|
||||
|
||||
def cmult(a,b):
|
||||
"""Returns the complex product of two torch tensors
|
||||
|
||||
Pytorch uses tensors with a final dimension of 2 to represent
|
||||
complex numbers. This calculates the elementwise product
|
||||
of two torch tensors following that standard.
|
||||
|
||||
Args:
|
||||
a (torch.Tensor): An input tensor
|
||||
b (torch.Tensor): A second input tensor
|
||||
|
||||
Returns:
|
||||
torch.Tensor : A tensor storing the elementwise product
|
||||
|
||||
"""
|
||||
|
||||
real = a[...,0] * b[...,0] - a[...,1] * b[...,1]
|
||||
imag = a[...,0] * b[...,1] + a[...,1] * b[...,0]
|
||||
return t.stack((real,imag),dim=-1)
|
||||
|
||||
|
||||
def cdiv(a,b):
|
||||
"""Returns the complex quotient of two torch tensors
|
||||
|
||||
Pytorch uses tensors with a final dimension of 2 to represent
|
||||
complex numbers. This calculates the elementwise quotient
|
||||
of two torch tensors following that standard.
|
||||
|
||||
Args:
|
||||
a (torch.Tensor): An input tensor
|
||||
b (torch.Tensor): A second input tensor
|
||||
|
||||
Returns:
|
||||
torch.Tensor : A tensor storing the elementwise complex quotient
|
||||
|
||||
"""
|
||||
return cmult(a, cconj(b)) / t.unsqueeze(cabssq(b),-1)
|
||||
|
||||
|
||||
|
||||
#
|
||||
# Not entirely sure if these belong here, but heck with it.
|
||||
# We just need the ability to do fftshifts
|
||||
#
|
||||
|
||||
|
||||
def fftshift(array,dims=None):
|
||||
"""Drop-in torch replacement for scipy.fftpack.fftshift
|
||||
|
||||
This maps a tensor, assumed to be the output of a fast Fourier
|
||||
transform, into a tensor whose zero-frequency element is at the
|
||||
center of the tensor instead of the start. It will by default shift
|
||||
every dimension in the tensor but the last (which is assumed to
|
||||
represent the complex number and be of dimension 2), but can shift
|
||||
any arbitrary set of dimensions.
|
||||
|
||||
Args:
|
||||
array (torch.Tensor) : An array of data to be fftshifted
|
||||
dims (iterable) : A list of all dimensions to shift
|
||||
|
||||
Returns:
|
||||
torch.Tensor : fftshifted tensor
|
||||
|
||||
"""
|
||||
|
||||
if dims is None:
|
||||
dims=list(range(array.dim()))[:-1]
|
||||
for dim in dims:
|
||||
length = array.size()[dim]
|
||||
cut_to = (length + 1) // 2
|
||||
cut_len = length - cut_to
|
||||
array = t.cat((array.narrow(dim,cut_to,cut_len),
|
||||
array.narrow(dim,0,cut_to)), dim)
|
||||
return array
|
||||
|
||||
|
||||
|
||||
def ifftshift(array,dims=None):
|
||||
"""Drop-in torch replacement for scipy.fftpack.iftshift
|
||||
|
||||
This maps a tensor, assumed to be the shifted output of a fast
|
||||
Fourier transform, into a tensor whose zero-frequency element is
|
||||
back at the start of the tensor instead of the center. It is the
|
||||
inverse of the fftshift operator. It will by default shift
|
||||
every dimension in the tensor but the last (which is assumed to
|
||||
represent the complex number and be of dimension 2), but can shift
|
||||
any arbitrary set of dimensions.
|
||||
|
||||
Args:
|
||||
array (torch.Tensor) : An array of data to be ifftshifted
|
||||
dims (iterable) : A list of all dimensions to shift
|
||||
|
||||
Returns:
|
||||
torch.Tensor : ifftshifted tensor
|
||||
|
||||
"""
|
||||
|
||||
if dims is None:
|
||||
dims=list(range(array.dim()))[:-1]
|
||||
for dim in dims:
|
||||
length = array.size()[dim]
|
||||
cut_to = length // 2
|
||||
cut_len = length - cut_to
|
||||
|
||||
array = t.cat((array.narrow(dim,cut_to,cut_len),
|
||||
array.narrow(dim,0,cut_to)), dim)
|
||||
return array
|
||||
|
||||
|
||||
@@ -0,0 +1,344 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import h5py
|
||||
import numpy as np
|
||||
|
||||
__all__ = ['get_entry_info',
|
||||
'get_sample_info',
|
||||
'get_wavelength',
|
||||
'get_detector_geometry',
|
||||
'get_mask',
|
||||
'get_data',
|
||||
'get_ptycho_translations']
|
||||
|
||||
#
|
||||
#
|
||||
# I will put here some thoughts about how to load data into this program.
|
||||
#
|
||||
#
|
||||
# The reconstructions should have the ability to generate datasets.
|
||||
# So you could write a reconstruction engine and then it would be
|
||||
# able to simulate data directly in the engine for you to use as a
|
||||
# reconstruction
|
||||
#
|
||||
# I don't even think there needs to be a loading tool for loading cxi files
|
||||
# because there isn't really a better method beyond just loading the
|
||||
# file into an h5py object. This file could host the simple cxi file
|
||||
# browser, perhaps. But I think the reality is that we need individual
|
||||
# loaders for each kind of experiment. Perhaps we could put some basic
|
||||
# reuseable tools for inspecting cxi-type h5 files in this file.
|
||||
#
|
||||
#
|
||||
# Then, there can be some more sophisticated tools that load data for
|
||||
# specific use cases that are common - loading data for a 2D CDI experiment,
|
||||
# loading data for a 2D Ptycho experiment, loading data for Bragg Ptycho in
|
||||
# 3D, loading data for a 3D CDI experiment, etc.
|
||||
#
|
||||
#
|
||||
# Perhaps one good way to package this is for the kind of data associated
|
||||
# with any particular experiment to have it's own kind of dataset or view.
|
||||
# So there would be a "2D Ptychography" data viewer, which would contain
|
||||
# all the measured data that comes from a 2D ptychography experiment.
|
||||
# The specialized functions would plop out these data viewers, and the
|
||||
# reconstruction classes could be designed around a particular kind of
|
||||
# viewer with the most general kind just requiring a generic data viewer.
|
||||
#
|
||||
# Data viewers could have simple tools like the ability to send themselves
|
||||
# to the GPU, CPU, change the datatype, etc. I think the most generic thing
|
||||
# is as a subclass of the torch Data objects, where they would for each slice
|
||||
# return the index, a set of defining parameters (translation, angle, energy,
|
||||
# whatever), and a diffraction pattern. They would also have a "setup"
|
||||
# attribute, or "metadata", or whatever you'd want to call it, that contain
|
||||
# the various fixed experimental parameters (energy, distance, etc.)
|
||||
#
|
||||
# And I think the cxi visualizer should really go into it's own script,
|
||||
# because it's not a reuseable component.
|
||||
#
|
||||
|
||||
|
||||
#
|
||||
# Functions to inspect the basic attributes of a cxi file represented as an
|
||||
# h5 file object
|
||||
#
|
||||
|
||||
def get_entry_info(cxi_file):
|
||||
"""Returns a dictionary with the basic metadata from the cxi file's entry_1 attribute
|
||||
|
||||
Args:
|
||||
cxi_file (h5py.File) : a file object to be read
|
||||
|
||||
Returns:
|
||||
dict : A dictionary with basic metadata defined in the cxi file
|
||||
|
||||
"""
|
||||
e1 = cxi_file['entry_1']
|
||||
metadata_attrs = ['title',
|
||||
'experiment_identifier',
|
||||
'experiment_description',
|
||||
'program_name',
|
||||
'start_time',
|
||||
'end_time']
|
||||
metadata = {attr: str(e1[attr][()].decode()) for attr in metadata_attrs
|
||||
if attr in e1}
|
||||
return metadata
|
||||
|
||||
|
||||
def get_sample_info(cxi_file):
|
||||
"""Returns a dictionary with the basic metadata from the cxi file's entry_1/sample_1 attribute
|
||||
|
||||
Args:
|
||||
cxi_file (h5py.File) : a file object to be read
|
||||
|
||||
Returns:
|
||||
dict : A dictionary with basic metadata from the sample defined in the cxi file
|
||||
|
||||
"""
|
||||
if 'entry_1/sample_1' not in cxi_file:
|
||||
return None
|
||||
|
||||
s1 = cxi_file['entry_1/sample_1']
|
||||
metadata_attrs = ['name','description','unit_cell_group']
|
||||
metadata = {attr: str(s1[attr][()].decode()) for attr in metadata_attrs
|
||||
if attr in s1}
|
||||
|
||||
float_attrs = ['concentration',
|
||||
'mass',
|
||||
'temperature',
|
||||
'thickness',
|
||||
'unit_cell_volume']
|
||||
for attr in float_attrs:
|
||||
if attr in s1:
|
||||
metadata[attr] = np.float32(s1[attr][()])
|
||||
|
||||
if 'unit_cell' in s1:
|
||||
metadata['unit_cell'] = np.array(s1['unit_cell']).astype(np.float32)
|
||||
|
||||
# TODO: Add my nonstandard "surface normal" attribute here
|
||||
|
||||
# TODO: I should add the sample geometry as a valid metadata that can
|
||||
# be copied over
|
||||
|
||||
return metadata
|
||||
|
||||
|
||||
def get_wavelength(cxi_file):
|
||||
"""Returns the wavelength of the source defined in the cxi file object, in m
|
||||
|
||||
Args:
|
||||
cxi_file (h5py.File) : a file object to be read
|
||||
|
||||
Returns:
|
||||
np.float32 : The wavelength of the source defined in the cxi file
|
||||
"""
|
||||
i1 = cxi_file['entry_1/instrument_1']
|
||||
if 'source_1/wavelength' in i1:
|
||||
wavelength = np.float32(i1['source_1/wavelength'])
|
||||
elif 'source_1/energy' in i1:
|
||||
energy = np.float32(i1['source_1/energy'])
|
||||
wavelength = 1.9864459e-25 / energy
|
||||
else:
|
||||
raise KeyError('Neither Wavelength or Energy Defined in provided .cxi File')
|
||||
|
||||
return wavelength
|
||||
|
||||
|
||||
def get_detector_geometry(cxi_file):
|
||||
"""Returns a standardized description of the detector geometry defined in the cxi file object
|
||||
|
||||
It makes intelligent assumptions based on the definitions in the cxi
|
||||
file definition. The standardized description of the geometry that it
|
||||
outputs includes the sample to detector distance, the corner location
|
||||
of the detector, and the basis vectors defining the detector. It can
|
||||
only handle detectors defined as rectangular grids of pixels.
|
||||
|
||||
The distance and corner_location values are technically overdetermining
|
||||
the detector location, but for many experiments (particularly
|
||||
transmission experiments), the distance is needed and the exact
|
||||
corner location is not. If the corner location is not reported in
|
||||
the cxi file, no attempt will be made to calculate it.
|
||||
|
||||
Args:
|
||||
cxi_file (h5py.File) : a file object to be read
|
||||
|
||||
Returns:
|
||||
distance (np.float32) : The sample to detector distance, in m
|
||||
basis_vectors (np.array) : The basis vectors for the detector
|
||||
corner_location (np.array) : The location of the (0,0) pixel in the detector
|
||||
|
||||
"""
|
||||
i1 = cxi_file['entry_1/instrument_1']
|
||||
d1 = i1['detector_1']
|
||||
|
||||
if 'detector_1/basis_vectors' in i1:
|
||||
basis_vectors = np.array(d1['basis_vectors'])
|
||||
else:
|
||||
# This whole thing just to account for all the ways people can
|
||||
# implicitly define the x or y pixel size for a detector. I've
|
||||
# seen too many of these in the wild, unfortunately...
|
||||
try:
|
||||
x_pixel_size = np.float32(d1['x_pixel_size'])
|
||||
except:
|
||||
x_pixel_size = None
|
||||
try:
|
||||
y_pixel_size = np.float32(d1['y_pixel_size'])
|
||||
except:
|
||||
y_pixel_size = None
|
||||
|
||||
if x_pixel_size is None and y_pixel_size is not None:
|
||||
x_pixel_size = y_pixel_size
|
||||
elif x_pixel_size is not None and y_pixel_size is None:
|
||||
y_pixel_size = x_pixel_size
|
||||
if x_pixel_size is None and y_pixel_size is None:
|
||||
raise KeyError('Detector pixel size not defined in file.')
|
||||
basis_vectors = np.array([[0,-y_pixel_size,0],
|
||||
[-x_pixel_size,0,0]]).transpose()
|
||||
|
||||
try:
|
||||
distance = np.float32(d1['distance'])
|
||||
except:
|
||||
distance = None
|
||||
try:
|
||||
corner_position = np.array(d1['corner_position'])
|
||||
except:
|
||||
corner_position = None
|
||||
|
||||
# Don't pretend to calculate corner position from distance if it's
|
||||
# if it's not defined, but do calculate distance from corner position
|
||||
# if distance is not defined. If neither is defined, then raise
|
||||
# an error.
|
||||
if distance is None and corner_position is not None:
|
||||
detector_normal = np.cross(basis_vectors[:,0],
|
||||
basis_vectors[:,1])
|
||||
detector_normal /= np.linalg.norm(detector_normal)
|
||||
distance = np.linalg.norm(np.dot(corner_position, detector_normal))
|
||||
|
||||
if distance is None and corner_position is not None:
|
||||
raise KeyError('Neither sample to detector distance or corner position is defined in file.')
|
||||
|
||||
return distance, basis_vectors, corner_position
|
||||
|
||||
|
||||
def get_mask(cxi_file):
|
||||
"""Returns the detector mask defined in the cxi file object
|
||||
|
||||
This function converts from the format specified in the cxi file
|
||||
definition to a simple on/off mask, where a value of 1 defines a
|
||||
good pixel (on) and a value of 0 defines a bad pixel (off).
|
||||
|
||||
If any bit is set in the mask at all, it will be defined as a bad
|
||||
pixel, with the exception of pixels marked exactly as 0x00001000,
|
||||
which is defined to mean that the pixel has signal above the
|
||||
background. These pixels are treated as on pixels
|
||||
|
||||
Args:
|
||||
cxi_file (h5py.File) : a file object to be read
|
||||
|
||||
Returns:
|
||||
np.array : An array storing the mask from the cxi file
|
||||
"""
|
||||
|
||||
i1 = cxi_file['entry_1/instrument_1']
|
||||
if 'detector_1/mask' in i1:
|
||||
mask = np.array(i1['detector_1/mask']).astype(np.uint32)
|
||||
mask_on = np.equal(mask,np.uint32(0))
|
||||
mask_has_signal = np.equal(mask,np.uint32(0x00001000))
|
||||
return np.logical_or(mask_on,mask_has_signal).astype(np.uint8)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def get_data(cxi_file):
|
||||
"""Returns an array with the full stack of detector data defined in the cxi file object
|
||||
|
||||
This function will make sure to check all the various places that it's
|
||||
okay to store the data in, to ensure that it can find the data regardless
|
||||
of whether the creator of the .cxi file has remembered to link the data
|
||||
to all the required locations.
|
||||
|
||||
It will return the data array in whatever shape it's defined in.
|
||||
|
||||
It will also read out the axes attribute of the data into a list
|
||||
of strings
|
||||
|
||||
Args:
|
||||
cxi_file (h5py.File) : a file object to be read
|
||||
|
||||
Returns:
|
||||
np.array : An array storing the data defined in the cxi file
|
||||
list : A list of the axes defined in the axes attribute, if any
|
||||
"""
|
||||
# Possible locations for the data
|
||||
#
|
||||
# entry_1/detector_1/data
|
||||
if 'entry_1/data_1/data' in cxi_file:
|
||||
pull_from = 'entry_1/data_1/data'
|
||||
elif 'entry_1/instrument_1/detector_1/data' in cxi_file:
|
||||
pull_from = 'entry_1/instrument_1/detector_1/data'
|
||||
else:
|
||||
raise KeyError('Data is not defined within cxi file')
|
||||
data = np.array(cxi_file[pull_from]).astype(np.float32)
|
||||
if 'axes' in cxi_file[pull_from].attrs:
|
||||
axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':')
|
||||
axes = [axis.strip().lower() for axis in axes]
|
||||
else:
|
||||
axes = None
|
||||
|
||||
return data, axes
|
||||
|
||||
|
||||
|
||||
def get_ptycho_translations(cxi_file):
|
||||
"""Gets an array of x,y,z translations, if such an array has been defined in the file
|
||||
|
||||
It applies two operations to the translations. First, it negates them,
|
||||
because the CXI file format is designed to specify translations of the
|
||||
samples and the CDTools code specifies translations of the optics.
|
||||
Second, it transposes the array so that the first axis is translation
|
||||
ID and the second axis is the (x,y,z) components of the translation
|
||||
|
||||
Args:
|
||||
cxi_file (h5py.File) : a file object to be read
|
||||
|
||||
Returns:
|
||||
np.array : An array storing the translations defined in the cxi file
|
||||
list : A list of the axes defined in the axes attribute, if any
|
||||
"""
|
||||
|
||||
if 'entry_1/data_1/translation' in cxi_file:
|
||||
pull_from = 'entry_1/data_1/translation'
|
||||
elif 'entry_1/sample_1/geometry_1/translation' in cxi_file:
|
||||
pull_from = 'entry_1/sample_1/geometry_1/translation'
|
||||
elif 'entry_1/instrument_1/detector_1/translation' in cxi_file:
|
||||
pull_from = 'entry_1/instrument_1/detector_1/translation'
|
||||
else:
|
||||
raise KeyError('Translations are not defined within cxi file')
|
||||
|
||||
translations = -np.array(cxi_file[pull_from]).astype(np.float32).transpose()
|
||||
return translations
|
||||
|
||||
|
||||
#
|
||||
# It might be useful to make some helper functions to help write cxi files
|
||||
#
|
||||
|
||||
#
|
||||
# A function to place the skeleton of a cxi file down
|
||||
#
|
||||
|
||||
#
|
||||
# A function to define the source attributes
|
||||
#
|
||||
|
||||
#
|
||||
# A function to define the detector geometry
|
||||
#
|
||||
|
||||
#
|
||||
# A function to save out a mask, converting it to the correct format
|
||||
#
|
||||
|
||||
#
|
||||
# Perhaps a function to store the data and link it correctly? But this might
|
||||
# have to change too much situation to situation
|
||||
#
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
import numpy as np
|
||||
import torch as t
|
||||
@@ -0,0 +1,27 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
import numpy as np
|
||||
import torch as t
|
||||
|
||||
all = ['gaussian']
|
||||
|
||||
|
||||
def gaussian(shape, amplitude, sigma, center = None):
|
||||
"""Returns an array with a centered gaussian
|
||||
|
||||
Takes in the shape, amplitude, and standard deviation of a gaussian
|
||||
and returns an array with values corresponding to a two-dimensional gaussian function
|
||||
z = amplitude*exp(-(x-center[0])**2/sigma[0]**2+(y-center[1])**2/sigma[1]**2)
|
||||
Note that [0, 0] is taken to be at the upper left corner of the array.
|
||||
Default is centered at (shape[0]/2, shape[1]/2).
|
||||
|
||||
Args:
|
||||
shape (array_like) : A 1x2 array-like object specifying the dimensions of the output array
|
||||
amplitude (float or int): The amplitude the gaussian to simulate
|
||||
sigma (array_like): A 1x2 array-like object specifying the x- and y- standard deviation of the gaussian
|
||||
center (array_like) : Optional 1x2 array-like object specifying the location of the center of the gaussian
|
||||
|
||||
Returns:
|
||||
torch.Tensor : The real-valued gaussian array
|
||||
"""
|
||||
x, y = np.meshgrid(shape)
|
||||
return x
|
||||
@@ -0,0 +1,123 @@
|
||||
"""Contains various loss functions to be used for optimization
|
||||
|
||||
It exposes three losses, one returning the mean squared amplitude error, one
|
||||
that returns the mean squared intensity error, and one that returns the
|
||||
maximum likelihood metric for a system with Poisson statistics.
|
||||
|
||||
"""
|
||||
from __future__ import division, print_function, absolute_import
|
||||
import torch as t
|
||||
|
||||
|
||||
__all__ = ['amplitude_mse', 'intensity_mse', 'poisson_nll']
|
||||
|
||||
|
||||
def amplitude_mse(intensities, sim_intensities, mask=None):
|
||||
""" Returns the mean squared error of a simulated dataset's amplitudes
|
||||
|
||||
Calculates the mean squared error between a given set of
|
||||
measured diffraction intensities and a simulated set.
|
||||
|
||||
This function calculates the mean squared error between their
|
||||
associated amplitudes. Because this is not well defined for negative
|
||||
numbers, make sure that all the intensities are >0 before using this
|
||||
loss.
|
||||
|
||||
It can accept intensity and simulated intensity tensors of any shape
|
||||
as long as their shapes match, and the provided mask array can be
|
||||
broadcast correctly along them.
|
||||
|
||||
This is empirically the most useful loss function
|
||||
|
||||
Args:
|
||||
intensities (torch.Tensor) : A tensor with measured detector values
|
||||
sim_intensities (torch.Tensor) : A tensor of simulated detector intensities
|
||||
mask (torch.Tensor) : A mask with ones for pixels to include and zeros for pixels to exclude
|
||||
|
||||
Returns:
|
||||
loss (torch.Tensor) : A single value for the summed mse
|
||||
|
||||
"""
|
||||
|
||||
# I know it would be more efficient if this function took in the
|
||||
# amplitudes instead of the intensities, but I want to be consistent
|
||||
# with all the errors working off of the same inputs
|
||||
|
||||
if mask is None:
|
||||
return t.sum((t.sqrt(sim_intensities) -
|
||||
t.sqrt(intensities))**2) / intensities.view(-1).shape[0]
|
||||
else:
|
||||
masked_intensities = intensities.masked_select(mask)
|
||||
return t.sum((t.sqrt(sim_intensities.masked_select(mask)) -
|
||||
t.sqrt(masked_intensities))**2) / masked_intensities.shape[0]
|
||||
|
||||
|
||||
|
||||
def intensity_mse(intensities, sim_intensities, mask=None):
|
||||
""" Returns the mean squared error of a simulated dataset's intensities
|
||||
|
||||
Calculates the summed mean squared error between a given set of
|
||||
diffraction intensities - the measured set of detector intensities -
|
||||
and a simulated set of diffraction intensities. This function
|
||||
calculates the mean squared error between the intensities.
|
||||
|
||||
It can accept intensity and simulated intensity tensors of any shape
|
||||
as long as their shapes match, and the provided mask array can be
|
||||
broadcast correctly along them.
|
||||
|
||||
Args:
|
||||
intensities (torch.Tensor) : A tensor with measured detector intensities.
|
||||
sim_intensities (torch.Tensor) : A tensor of simulated detector intensities
|
||||
mask (torch.Tensor) : A mask with ones for pixels to include and zeros for pixels to exclude
|
||||
|
||||
Returns:
|
||||
loss (torch.Tensor) : A single value for the summed mse
|
||||
|
||||
"""
|
||||
if mask is None:
|
||||
return t.sum((sim_intensities - intensities)**2) \
|
||||
/ intensities.view(-1).shape[0]
|
||||
else:
|
||||
masked_intensities = intensities.masked_select(mask)
|
||||
return t.sum((sim_intensities.masked_select(mask) -
|
||||
masked_intensities)**2) \
|
||||
/ masked_intensities.shape[0]
|
||||
|
||||
|
||||
|
||||
def poisson_nll(intensities, sim_intensities, mask=None):
|
||||
""" Returns the Poisson negative log likelihood for a simulated dataset's intensities
|
||||
|
||||
Calculates the overall Poisson maximum likelihood metric using
|
||||
diffraction intensities - the measured set of detector intensities -
|
||||
and a simulated set of intensities. This loss would be appropriate
|
||||
for detectors in a single-photon counting mode, with their output
|
||||
scaled to number of photons
|
||||
|
||||
Note that this calculation ignores the log(intensities!) term in the
|
||||
full expression for Poisson negative log likelihood. This term doesn't
|
||||
change the calculated gradients so isn't worth taking the time to compute
|
||||
|
||||
It can accept intensity and simulated intensity tensors of any shape
|
||||
as long as their shapes match, and the provided mask array can be
|
||||
broadcast correctly along them.
|
||||
|
||||
Args:
|
||||
intensities (torch.Tensor) : A tensor with measured detector intensities.
|
||||
sim_intensities (torch.Tensor) : A tensor of simulated detector intensities
|
||||
mask (torch.Tensor) : A mask with ones for pixels to include and zeros for pixels to exclude
|
||||
|
||||
Returns:
|
||||
loss (torch.Tensor) : A single value for the poisson ML metric
|
||||
|
||||
"""
|
||||
if mask is None:
|
||||
return t.sum(sim_intensities -
|
||||
intensities * t.log(sim_intensities)) \
|
||||
/ intensities.view(-1).shape[0]
|
||||
|
||||
else:
|
||||
masked_intensities = intensities.masked_select(mask)
|
||||
masked_sims = sim_intensities.masked_select(mask)
|
||||
return t.sum(masked_sims - masked_intensities *
|
||||
t.log(masked_sims)) / masked_intensities.shape[0]
|
||||
@@ -0,0 +1,52 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
from CDTools.tools.cmath import *
|
||||
import torch as t
|
||||
|
||||
__all__ = ['modulus', 'support']
|
||||
|
||||
def modulus(wavefront, intensities, mask = None):
|
||||
"""Implements the modulus constraint in torch
|
||||
|
||||
This accepts a torch tensor representing the propagated simulated wavefront(s),
|
||||
where the last dimension represents the real and imaginary components of
|
||||
the propagated wavefield(s). It projects the modulus of the diffraction pattern
|
||||
onto the modulus of the simulated wavefield.
|
||||
|
||||
It assumes that the wavefront is stored in an array
|
||||
[i,j] where i corresponds to the y-axis and j corresponds to the
|
||||
x-axis, with the origin following the CS standard of being in the
|
||||
upper right.
|
||||
|
||||
Args:
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
|
||||
intensities (torch.Tensor): The measured diffraction pattern(s) stored as an JxNxM stack of real tensors
|
||||
mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
|
||||
Returns:
|
||||
torch.Tensor : The JxNxMx2 propagated wavefield with corrected intensities
|
||||
"""
|
||||
if mask is None:
|
||||
return cmult(cphase(wavefront), intensities**.5)
|
||||
else:
|
||||
return cmult(cphase(wavefront), intensities.masked_select(mask)**.5)
|
||||
|
||||
|
||||
def support(wavefront, mask):
|
||||
"""Implements the support constraint in torch
|
||||
|
||||
This accepts a torch tensor representing the propagated simulated wavefront(s),
|
||||
where the last dimension represents the real and imaginary components of
|
||||
the propagated wavefield(s). It projects the support of the imaged object
|
||||
onto the simulated wavefront via a mask.
|
||||
|
||||
It assumes that the wavefront is stored in an array
|
||||
[i,j] where i corresponds to the y-axis and j corresponds to the
|
||||
x-axis, with the origin following the CS standard of being in the
|
||||
upper right.
|
||||
|
||||
Args:
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
|
||||
mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
|
||||
Returns:
|
||||
torch.Tensor : The JxNxMx2 wavefield with the mask applied
|
||||
"""
|
||||
return wavefront.masked_select(mask)
|
||||
@@ -0,0 +1,168 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
from CDTools.tools.cmath import *
|
||||
import torch as t
|
||||
|
||||
__all__ = ['far_field', 'near_field', 'inverse_far_field', 'inverse_near_field', 'get_exit_waves']
|
||||
|
||||
|
||||
def far_field(wavefront, detector_shape, detector_center=None,
|
||||
scaling=1):
|
||||
"""Implements a far-field propagator in torch
|
||||
|
||||
This accepts a torch tensor, where the last dimension
|
||||
represents the real and imaginary components of the wavefield,
|
||||
and returns the far-field propagated version of it using the provided
|
||||
geometrical information about the detector. It assumes that the
|
||||
propagation is purely far-field, without checking that the geometry
|
||||
is consistent with that assumption. Note that the pitch of the
|
||||
real space array is assumed to be consistent with the detector geometry,
|
||||
such that the pixel spacing on the detector corresponds to the full
|
||||
size covered by the wavefield array.
|
||||
|
||||
It also assumes that the real space wavefront is stored in an array
|
||||
[i,j] where i corresponds to the y-axis and j corresponds to the
|
||||
x-axis, with the origin following the CS standard of being in the
|
||||
upper right.
|
||||
|
||||
Args:
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
|
||||
detector_shape (array_like): The shape of the detector to simulate
|
||||
detector_center (array_like): Optional, the pixel (i,j) coordinates of the intersection of the detector with the forward propagation direction.
|
||||
scaling (int) : Default is 1, the downscaling to apply to the measured diffraction pattern
|
||||
Returns:
|
||||
torch.Tensor : The Jxdetector_shapex2 propagated wavefield
|
||||
"""
|
||||
|
||||
if detector_center is None:
|
||||
# Default is the exact center. This is subtly different from the
|
||||
# default for a shifted FFT, where for even-sized arrays, the
|
||||
# zero frequency pixel is placed at (shape-1)//2, not (shape-1)/2
|
||||
detector_center = (np.array(detector_shape)-np.array([1,1]))/2
|
||||
|
||||
center = np.array(detector_center)
|
||||
# Split the center into a pixel and subpixel shift
|
||||
int_center = np.floor(center).astype(int)
|
||||
subpixel_shift = (center - int_center)
|
||||
|
||||
# A selection pulling out the final area from the simulated
|
||||
# diffraction pattern
|
||||
# To be used as arr[sel[0]:sel[1],sel[2]:sel[3]]
|
||||
# Use the wavefront shape before downsampling
|
||||
wf_shape = tuple(wavefront.shape)
|
||||
sel = ((wf_shape[-3]-1)//2 - scaling*int_center[0],
|
||||
(wf_shape[-3]-1)//2 - scaling*int_center[0]
|
||||
+ detector_shape[0]*scaling,
|
||||
(wf_shape[-2]-1)//2 - scaling*int_center[1],
|
||||
(wf_shape[-2]-1)//2 - scaling*int_center[1]
|
||||
+ detector_shape[1]*scaling)
|
||||
# This generates a phase ramp to use for the final subpixel
|
||||
# shift.
|
||||
Is, Js = np.mgrid[0:wavefront.shape[-3],0:wavefront.shape[-2]]
|
||||
Is = Is - np.mean(Is)
|
||||
Js = Js - np.mean(Js)
|
||||
locs = np.stack((Is, Js), axis=-1)
|
||||
# Move from pixel to frequency units
|
||||
phase_ramp_freq = 2 * np.pi * subpixel_shift / wavefront.shape[-3:-1]
|
||||
phase_ramp = np.exp(-1j * np.dot(locs, phase_ramp_freq))
|
||||
phase_ramp = complex_to_torch(phase_ramp).to(device=wavefront.device,
|
||||
dtype=wavefront.dtype)
|
||||
|
||||
ramped_wavefront = cmult(phase_ramp[None,...], wavefront)
|
||||
|
||||
sims = fftshift(fft(ifftshift(ramped_wavefront, dims=(-2,-3)),
|
||||
2, normalized=True), dims=(-2,-3))
|
||||
|
||||
return sims[:,sel[0]:sel[1],sel[2]:sel[3]]
|
||||
|
||||
def inverse_far_field(wavefront, detector_shape, detector_center=None,
|
||||
scaling=1):
|
||||
if detector_center is None:
|
||||
# Default is the exact center. This is subtly different from the
|
||||
# default for a shifted FFT, where for even-sized arrays, the
|
||||
# zero frequency pixel is placed at (shape-1)//2, not (shape-1)/2
|
||||
detector_center = (np.array(detector_shape)-np.array([1,1]))/2
|
||||
|
||||
|
||||
def inverse_near_field():
|
||||
pass
|
||||
|
||||
|
||||
|
||||
def near_field(wavefront, spacing, wavelength, z):
|
||||
"""Implements an angular-spectrum based near-field propagator in torch
|
||||
|
||||
This function is an angular-spectrum based near field
|
||||
propagator that will work on torch Tensors. The function is structured
|
||||
this way - to generate the propagator first - because the
|
||||
generation of the propagation mask is a bit expensive and if this
|
||||
propagator is used in a reconstruction program, then it will be best
|
||||
to calculate this mask once and close over it.
|
||||
The resulting function accepts an 3d torch tensor, where the
|
||||
last dimension represents the real and imaginary components of
|
||||
the wavefield, and returns the near-field propagated version of it.
|
||||
|
||||
Args:
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
|
||||
spacing (iterable) : The pixel size in each dimension of the arrays to be propagated
|
||||
wavelength (float) : The wavelength of light to simulate propagation of
|
||||
z (float) : The distance to simulate propagation over
|
||||
Returns:
|
||||
function : A function to propagate a torch tensor.
|
||||
"""
|
||||
|
||||
ki = fftpack.fftfreq(shape[0],spacing[0])
|
||||
kj = fftpack.fftfreq(shape[1],spacing[1])
|
||||
Ki, Kj = np.meshgrid(ki,kj)
|
||||
propagator = np.exp(1j*np.sqrt((2*np.pi/wavelength)**2
|
||||
- Ki**2 - Kj**2) * z)
|
||||
propagator = complex_to_float(propagator).astype(np.float32)
|
||||
propagator = t.from_numpy(propagator).cuda()
|
||||
|
||||
return t.ifft(propagator * t.fft(wavefront,2),2)
|
||||
|
||||
|
||||
|
||||
def get_exit_waves(probe, object, translations):
|
||||
"""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 object with each translation in turn, using linear
|
||||
interpolation.
|
||||
Args:
|
||||
probe (torch.Tensor) : An MxM 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
|
||||
Returns:
|
||||
torch.Tensor : An NxMxM 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)
|
||||
|
||||
selections = []
|
||||
for tr, sp in zip(integer_translations,
|
||||
subpixel_translations):
|
||||
|
||||
sel00 = object[tr[0]:tr[0]+probe.shape[0],
|
||||
tr[1]:tr[1]+probe.shape[1]]
|
||||
|
||||
sel01 = object[tr[0]:tr[0]+probe.shape[0],
|
||||
tr[1]+1:tr[1]+1+probe.shape[1]]
|
||||
|
||||
sel10 = object[tr[0]+1:tr[0]+1+probe.shape[0],
|
||||
tr[1]:tr[1]+probe.shape[1]]
|
||||
|
||||
sel11 = object[tr[0]+1:tr[0]+1+probe.shape[0],
|
||||
tr[1]+1:tr[1]+1+probe.shape[1]]
|
||||
|
||||
selections.append(sel00 * (1-sp[0])*(1-sp[1]) + \
|
||||
sel01 * (1-sp[0])*sp[1] + \
|
||||
sel10 * sp[0]*(1-sp[1]) + \
|
||||
sel11 * sp[0]*sp[1])
|
||||
|
||||
return t.stack([cmult(probe,selection) for selection in selections])
|
||||
Vendored
BIN
Binary file not shown.
@@ -0,0 +1,24 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools import initializers
|
||||
import numpy as np
|
||||
import torch as t
|
||||
|
||||
|
||||
|
||||
def test_gaussian():
|
||||
# Generate gaussian as a numpy array (square array)
|
||||
shape = [10, 10]
|
||||
sigma = [2.5, 2.5]
|
||||
center = ((shape[0]-1)/2, (shape[1]-1)/2)
|
||||
y, x = np.mgrid[:shape[0], :shape[1]]
|
||||
np_result = 10*np.exp(-((x-center[1])/sigma[1])**2-((y-center[0])/sigma[0])**2)
|
||||
assert(np.allclose(initializers.gaussian([10, 10], 10, [2.5, 2.5]), np_result))
|
||||
|
||||
# Generate gaussian as a numpy array (rectangular array)
|
||||
shape = [10, 5]
|
||||
sigma = [2.5, 2.5]
|
||||
center = ((shape[0]-1)/2, (shape[1]-1)/2)
|
||||
y, x = np.mgrid[:shape[0], :shape[1]]
|
||||
np_result = 10*np.exp(-((x-center[1])/sigma[1])**2-((y-center[0])/sigma[0])**2)
|
||||
assert(np.allclose(initializers.gaussian([10, 5], 10, [2.5, 2.5]), np_result))
|
||||
@@ -0,0 +1,21 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools import cmath
|
||||
from CDTools.tools import projectors
|
||||
import numpy as np
|
||||
import torch as t
|
||||
from scipy.fftpack import fftshift, ifftshift
|
||||
|
||||
def test_modulus():
|
||||
# Create a complex array with modulus 12 and phase pi/4
|
||||
np_result = 6**.5*(np.ones((10,10))+1j*np.ones((10,10)))
|
||||
|
||||
assert(np.allclose(cmath.torch_to_complex(projectors.modulus(t.ones((10,10,2)), 12*t.ones((10,10)))), np_result))
|
||||
|
||||
|
||||
def test_support():
|
||||
# Test masking
|
||||
support = t.zeros((10,10))
|
||||
np_result = np.zeros((10,10))
|
||||
|
||||
assert(np.allclose(cmath.torch_to_complex(projectors.support(t.ones((10,10,2)), support)), np_result))
|
||||
@@ -0,0 +1,48 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools import cmath
|
||||
from CDTools.tools import initializers
|
||||
from CDTools.tools import propagators
|
||||
|
||||
import numpy as np
|
||||
import torch as t
|
||||
import pytest
|
||||
import scipy.misc
|
||||
from scipy.fftpack import fftshift, ifftshift
|
||||
|
||||
@pytest.fixture(scope='module')
|
||||
def exit_waves_1():
|
||||
# Import scipy test image and add a random phase
|
||||
object = scipy.misc.ascent()[0:64,0:64].astype(np.complex128)
|
||||
arr = np.random.random_sample((64,64))
|
||||
object *= (arr+(1-arr**2)**.5*1j)
|
||||
|
||||
# Construct wavefront from image
|
||||
probe = initializers.gaussian([64, 64], 1e3, [5, 5])*(1+1j)
|
||||
return cmath.complex_to_torch(probe*object)
|
||||
|
||||
|
||||
def test_far_field(exit_waves_1):
|
||||
# Far field diffraction patterns calculated by numpy with zero frequency in center
|
||||
np_result = np.fft.fftshift(np.fft.fft2(cmath.torch_to_complex(exit_waves_1)))
|
||||
|
||||
assert(np.allclose(np_result, cmath.torch_to_complex(propagators.far_field(exit_waves_1))))
|
||||
|
||||
|
||||
def test_inverse_far_field(exit_waves_1):
|
||||
# We want the inverse far field to map back to the exit waves with no intensity corrections
|
||||
np_result = exit_waves_1
|
||||
# Far field result for exit waves calculated with numpy
|
||||
far_field_np_result = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(cmath.torch_to_complex(exit_waves_1))))
|
||||
|
||||
assert(np.allclose(np_result, propagators.inverse_far_field(far_field_np_result)))
|
||||
|
||||
|
||||
|
||||
def test_near_field(exit_waves_1):
|
||||
pass
|
||||
|
||||
|
||||
|
||||
def test_get_exit_waves(exit_waves_1):
|
||||
pass
|
||||
Reference in New Issue
Block a user