from proxtoolbox.proxoperators.proxoperator import ProxOperator
from proxtoolbox import proxoperators
from proxtoolbox.utils.cell import Cell, isCell
import numpy as np
from numpy import zeros, ones, zeros_like
# Prox operators used by Ptychography experiment
[docs]class ExplicitPtychographyProxOperator(ProxOperator):
"""
Base class for Ptychography explicit prox opertaors
"""
def __init__(self, experiment):
self.N_pie = experiment.N_pie
self.positions = experiment.positions
self.Nx = experiment.Nx
self.Ny = experiment.Ny
self.rangeNx = np.arange(experiment.Nx, dtype=np.int)
self.rangeNy = np.arange(experiment.Ny, dtype=np.int)
if hasattr(experiment, 'overrelax'):
self.overrelax = experiment.overrelax
else:
self.overrelax = 1
[docs]class Explicit_ptychography_scan_farfield_probe_ptwise(ExplicitPtychographyProxOperator):
"""
The probe update for the PHeBIE algorithm
Ref.: "Proximal Heterogeneous Block Implicit-Explicit Method
and Application to Blind Ptychographic Diffraction Imaging",
R. Hesse, D. R. Luke, S. Sabach and M. K. Tam,
SIAM J. on Imaging Sciences, 8(1):426--457 (2015).
Based on Matlab code written by Russell Luke (Inst. Fuer
Numerische und Angewandte Mathematik, Universitaet
Gottingen) on 16th February 2019.
"""
def __init__(self, experiment):
super(Explicit_ptychography_scan_farfield_probe_ptwise, self).__init__(experiment)
[docs] def eval(self, u, prox_idx=None):
"""
Evaluate this prox operator
Parameters
----------
u : Cell
Iterate. Assumes u[0] is the probe, u[1] is the object, u[2] is the
block of variables (product space) satisfying the intensity
measurements
prox_idx : int, optional
Index of this prox operator
Returns
-------
u_new : Cell
The projection.
"""
ySum = zeros_like(u[0])
phiSum = zeros(u[0].shape, dtype=u[1].dtype) # must be complex
conjObject = np.conj(u[1])
isCell_u2 = isCell(u[2])
for yPos in range(self.N_pie):
indy = (self.positions[yPos,0] + self.rangeNy).astype(int)
indx = (self.positions[yPos,1] + self.rangeNx).astype(int)
ySum += np.abs(u[1][indy,:][:,indx])**2
if isCell_u2:
u2_values = u[2][yPos]
else:
u2_values = u[2][:,:,yPos]
phiSum += conjObject[indy,:][:,indx]*u2_values
ysum_denom = self.overrelax * np.maximum(ySum, 1e-30)
u_new = u[0] - (ySum*u[0] - phiSum)/ysum_denom
return u_new
[docs]class Explicit_ptychography_scan_farfield_object_ptwise(ExplicitPtychographyProxOperator):
"""
The object update for the PHeBIE algorithm
Ref.: "Proximal Heterogeneous Block Implicit-Explicit Method
and Application to Blind Ptychographic Diffraction Imaging",
R. Hesse, D. R. Luke, S. Sabach and M. K. Tam,
SIAM J. on Imaging Sciences, 8(1):426--457 (2015).
Based on Matlab code written by Russell Luke (Inst. Fuer
Numerische und Angewandte Mathematik, Universitaet
Gottingen) on 16th February 2019.
"""
def __init__(self, experiment):
super(Explicit_ptychography_scan_farfield_object_ptwise, self).__init__(experiment)
[docs] def eval(self, u, prox_idx=None):
"""
Evaluate this prox operator
Parameters
----------
u : Cell
Iterate. Assumes u[0] is the probe, u[1] is the object, u[2] is the
block of variables (product space) satisfying the intensity
measurements
prox_idx : int, optional
Index of this prox operator
Returns
-------
u_new : Cell
The projection.
"""
xSum = zeros_like(u[1])
phiSum = zeros(u[1].shape, dtype = np.dtype(np.complex128)) # must be complex
conjProbe = np.conj(u[0])
modSqProbe = np.real(conjProbe * u[0])
isCell_u2 = isCell(u[2])
for xPos in range(self.N_pie):
indy = (self.positions[xPos,0] + self.rangeNy).astype(int)
indx = (self.positions[xPos,1] + self.rangeNx).astype(int)
if isCell_u2:
u2_values = u[2][xPos]
else:
u2_values = u[2][:,:,xPos]
for y in range(self.Ny):
xSum[indy[y],indx] += modSqProbe[y,:]
phiSum[indy[y],indx] += conjProbe[y,:] * u2_values[y,:]
xsum_denom = self.overrelax * np.maximum(xSum, 1e-30)
u_new = u[1] - (xSum*u[1] - phiSum)/xsum_denom
return u_new
[docs]class Explicit_ptychography_image(ExplicitPtychographyProxOperator):
"""
This is the argument of the projection in Step 3 of Algorithm 4.1
of "Proximal Heterogeneous Block Implicit-Explicit Method and
Application to Blind Ptychographic Diffraction Imaging",
R. Hesse, D. R. Luke, S. Sabach and M. K. Tam,
SIAM J. on Imaging Sciences, 8(1):426--457 (2015).
Based on Matlab code written by Russell Luke (Inst. Fuer
Numerische und Angewandte Mathematik, Universitaet
Gottingen) on 16th February 2019.
"""
def __init__(self, experiment):
super(Explicit_ptychography_image, self).__init__(experiment)
if not hasattr(experiment, 'overrelax'):
self.gamma = 0
else:
self.gamma = experiment.overrelax - 1
[docs] def eval(self, u, prox_idx=None):
"""
Evaluate this prox operator
Parameters
----------
u : Cell
Iterate. Assumes u[0] is the probe, u[1] is the object, u[2] is the
block of variables (product space) satisfying the intensity
measurements
prox_idx : int, optional
Index of this prox operator
Returns
-------
u_new : Cell
The projection.
"""
isCell_u2 = isCell(u[2])
if isCell_u2:
u_new = Cell(u[2].shape)
else:
u_new = zeros_like(u[2])
for pos in range(self.N_pie):
indy = (self.positions[pos,0] + self.rangeNy).astype(int)
indx = (self.positions[pos,1] + self.rangeNx).astype(int)
if isCell_u2:
u_new[pos] = (2/(2+self.gamma))*u[0]*u[1][indy,:][:,indx] \
+ self.gamma/(2+self.gamma)*u[2][pos]
else:
u_new[:,:,pos] = (2/(2+self.gamma))*u[0]*u[1][indy,:][:,indx] \
+ self.gamma/(2+self.gamma)*u[2][:,:,pos]
return u_new
[docs]class P_ptychography_scan_farfield_probe(ProxOperator):
"""
The probe update for the PHeBIE algorithm
Ref.: "Proximal Heterogeneous Block Implicit-Explicit Method and
Application to Blind Ptychographic Diffraction Imaging",
R. Hesse, D. R. Luke, S. Sabach and M. K. Tam,
SIAM J. on Imaging Sciences, 8(1):426--457 (2015).
Based on Matlab code written by Russell Luke (Inst. Fuer
Numerische und Angewandte Mathematik, Universitaet
Gottingen) on 16th February 2019.
"""
def __init__(self, experiment):
super(P_ptychography_scan_farfield_probe, self).__init__(experiment)
self.switch_probemask = experiment.switch_probemask
self.Probe_mask = experiment.Probe_mask
[docs] def eval(self, u, prox_idx=None):
"""
Evaluate this prox operator
Parameters
----------
u : Cell
Iterate. assumes u is the probe (i.e., only part u[0]
of the whole collection of PHeBIE variable blocks)
prox_idx : int, optional
Index of this prox operator
Returns
-------
u_new : Cell
The projection.
"""
u_new = u.copy()
# Mask in the sample plane?
if self.switch_probemask:
u_new = u_new * self.Probe_mask
return u_new
[docs]class P_supp_amp_band(ProxOperator):
"""
The object update for the PHeBIE algorithm
Ref.: "Proximal Heterogeneous Block Implicit-Explicit Method and
Application to Blind Ptychographic Diffraction Imaging",
R. Hesse, D. R. Luke, S. Sabach and M. K. Tam,
SIAM J. on Imaging Sciences, 8(1):426--457 (2015).
Based on Matlab code written by Russell Luke (Inst. Fuer
Numerische und Angewandte Mathematik, Universitaet
Gottingen) on 16th February 2019.
"""
def __init__(self, experiment):
super(P_supp_amp_band, self).__init__(experiment)
self.switch_object_support_constraint = experiment.switch_object_support_constraint
self.object_support = experiment.object_support
self.trans_max_true = experiment.trans_max_true
self.trans_min_true = experiment.trans_min_true
[docs] def eval(self, u, prox_idx=None):
"""
Evaluate this prox operator
Parameters
----------
u : Cell
Iterate. assumes u is the object (i.e., only part u[1]
of the whole collection of PHeBIE variable blocks)
prox_idx : int, optional
Index of this prox operator
Returns
-------
u_new : Cell
The projection.
"""
u_new = u.copy()
# Apply the support constraint
if self.switch_object_support_constraint:
# this may be less efficient that previous code;
# however multiplication by zero keeps the sign instead
# of properly setting the value to zero which creates
# issues when taking the phase (angle)
u_new = np.where(self.object_support == 0, 0, u_new)
# oldcode: u_new = u_new * self.object_support
# Enforce that the object has values in [minTrue,maxTrue]
maxTrue = self.trans_max_true
minTrue = self.trans_min_true
absObject = np.abs(u_new)
# TODO can use np.where() to achieve the same result
high = (absObject > maxTrue).astype(absObject.dtype)
low = (absObject < minTrue).astype(absObject.dtype)
u_new *= (1.0 - low)*(1.0 - high) + (low*minTrue + high*maxTrue)/(absObject + 1e-30)
return u_new
[docs]class Prox_ptychography_ptwise_diag(ExplicitPtychographyProxOperator):
"""
The object update for the PHeBIE algorithm
Ref.: "Proximal Heterogeneous Block Implicit-Explicit Method
and Application to Blind Ptychographic Diffraction Imaging",
R. Hesse, D. R. Luke, S. Sabach and M. K. Tam,
SIAM J. on Imaging Sciences, 8(1):426--457 (2015).
Based on Matlab code written by Russell Luke (Inst. Fuer
Numerische und Angewandte Mathematik, Universitaet
Gottingen) on 16th February 2019.
"""
def __init__(self, experiment):
super(Prox_ptychography_ptwise_diag, self).__init__(experiment)
self.FT_conv_kernel = experiment.FT_conv_kernel
proxClass = getattr(proxoperators, 'P_supp_amp_band')
self.prox_supp_amp_band = proxClass(experiment)
[docs] def eval(self, u, prox_idx=None):
"""
Evaluate this prox operator
Parameters
----------
u : Cell
Cell of arrays in the domain to be projected.
Assumes u[1] is the object, u[2] is the
block of variables (product space) satisfying the intensity
measurements
prox_idx : int, optional
Index of this prox operator
Returns
-------
u_new : Cell
The projection.
"""
xSum = zeros_like(u[0])
phiSum = zeros(u[0].shape, dtype = np.dtype(np.complex128)) # must be complex
#conjProbe = np.conj(u[0])
#modSqProbe = np.real(conjProbe * u[0])
isCell_u1 = isCell(u[1])
for xPos in range(self.N_pie):
indy = (self.positions[xPos,0] + self.rangeNy).astype(int)
indx = (self.positions[xPos,1] + self.rangeNx).astype(int)
if isCell_u1:
u1_values = u[1][xPos]
else:
u1_values = u[1][:,:,xPos]
conv_kernel = self.FT_conv_kernel[xPos]
conjProbe = np.conj(conv_kernel)
modSqProbe = np.real(conjProbe * conv_kernel)
for y in range(self.Ny):
xSum[indy[y],indx] += modSqProbe[y,:]
phiSum[indy[y],indx] += conjProbe[y,:] * u1_values[y,:]
xsum_denom = self.overrelax * np.maximum(xSum, 1e-30)
u_new = Cell(u.shape)
u_new[0] = phiSum/xsum_denom
u_new[0] = self.prox_supp_amp_band.eval(u_new[0])
u_new[1] = u[1].copy()
return u_new
[docs]class Prox_ptychography_product_space(ExplicitPtychographyProxOperator):
"""
Projection subroutine onto constraints arranged in a product space
Based on Matlab code written by Russell Luke (Inst. Fuer
Numerische und Angewandte Mathematik, Universitaet
Gottingen) on Oct 26, 2017.
"""
def __init__(self, experiment):
super(Prox_ptychography_product_space, self).__init__(experiment)
self.FT_conv_kernel = experiment.FT_conv_kernel
# instantiate product prox operators
self.proxOps = []
for proxClass in experiment.productProxOperators:
if proxClass != '':
prox = proxClass(experiment)
self.proxOps.append(prox)
[docs] def eval(self, u, prox_idx=None):
"""
Evaluate this prox operator
Parameters
----------
u : Cell
Cell of arrays in the domain to be projected.
Assumes u[1] is the object, u[2] is the
block of variables (product space) satisfying the intensity
measurements.
prox_idx : int, optional
Index of this prox operator
Returns
-------
u_new : Cell
The projection, a cell of arrays in the product space
"""
isCell_u1 = isCell(u[1])
u_new = u.copy() #TODO may have to check if a deep copy is needed
for xPos in range(self.N_pie):
indy = (self.positions[xPos,0] + self.rangeNy).astype(int)
indx = (self.positions[xPos,1] + self.rangeNx).astype(int)
tmp = self.FT_conv_kernel[xPos] * u[0][indy,:][:,indx]
if isCell_u1:
u_new[1][xPos] = tmp
else:
u_new[1][:,:,xPos] = tmp
K = len(self.proxOps)
for k in range(K):
if isCell_u1:
u_new[1][k] = self.proxOps[k].eval(u[1][k], k)
else:
u_new[1][:,:,k] = self.proxOps[k].eval(u[1][:,:,k], k)
return u_new