add projectors, propagators, and initializers

CDTools/tools/__init__.py

@@ -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

CDTools/tools/cmath.py

@@ -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
-
-

CDTools/tools/image_processing.py

@@ -0,0 +1,3 @@
+from __future__ import division, print_function, absolute_import
+import numpy as np
+import torch as t

CDTools/tools/initializers.py

@@ -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)

CDTools/tools/projectors.py

@@ -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

CDTools/tools/propagators.py

@@ -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])