Commit eae38efd authored by Matthijs's avatar Matthijs
Browse files

Started work towards 3D orbital tomogaphy, working on data processor and plotting routines

parent 220d5637
......@@ -115,7 +115,7 @@ class SimpleAlgorithm:
config['u'] = u
u_new = self.evaluate(u)
if 'diagnostic' in self.config:
if 'diagnostic' in self.config and self.config['diagnostic']:
# the next prox operation only gets used in the computation of
# the size of the gap. This extra step is not
# required in alternating projections, which makes RAAR
......@@ -132,7 +132,7 @@ class SimpleAlgorithm:
# tmp_change = (norm(u - u_new, 'fro') / norm_data) ** 2
tmp_change = phase_offset_compensated_norm(u, u_new, norm_type='fro', norm_factor=norm_data) ** 2
if 'diagnostic' in self.config:
if 'diagnostic' in self.config and self.config['diagnostic']:
# For Douglas-Rachford,in general it is appropriate to monitor the
# SHADOWS of the iterates, since in the convex case these converge
# even for beta=1.
......@@ -157,24 +157,24 @@ class SimpleAlgorithm:
elif q == 1:
for j in range(self.product_space_dimension):
tmp_change = tmp_change + (norm(u[:, :, j] - u_new[:, :, j], 'fro') / norm_data) ** 2
if 'diagnostic' in self.config:
if 'diagnostic' in self.config and self.config['diagnostic']:
# compute (||P_SP_Mx-P_Mx||/norm_data)^2:
tmp_gap = tmp_gap + (norm(u1[:, :, j] - u2[:, :, j], 'fro') / norm_data) ** 2
tmp_shadow = tmp_shadow + (norm(u2[:, :, j] - shadow[:, :, j], 'fro') / norm_data) ** 2
if 'diagnostic' in self.config:
if 'diagnostic' in self.config and self.config['diagnostic']:
if hasattr(self, 'truth'):
z = u1[:, :, 0]
Relerrs[iter] = norm((self.truth - exp(-1j * angle(trace(self.truth.T.transpose() * z))) * z),'fro') / self.norm_truth
else:
if 'diagnostic' in self.config:
if 'diagnostic' in self.config and self.config['diagnostic']:
if hasattr(self, 'truth'):
Relerrs[iter] = 0
for j in range(self.product_space_dimension):
for k in range(self.Nz):
tmp_change = tmp_change + (norm(u[:, :, k, j] - u_new[:, :, k, j], 'fro') / norm_data) ** 2
if 'diagnostic' in self.config:
if 'diagnostic' in self.config and self.config['diagnostic']:
# compute (||P_Sx-P_Mx||/norm_data)^2:
tmp_gap = tmp_gap + (norm(u1[:, :, k, j] - u2[:, :, k, j], 'fro') / (norm_data)) ** 2
tmp_shadow = tmp_shadow + (
......@@ -185,7 +185,7 @@ class SimpleAlgorithm:
'fro') / self.norm_Truth
change[iter] = sqrt(tmp_change)
if 'diagnostic' in self.config:
if 'diagnostic' in self.config and self.config['diagnostic']:
gap[iter] = sqrt(tmp_gap)
shadow_change[iter] = sqrt(tmp_shadow) # this is the Euclidean norm of the gap to
# the unregularized set. To monitor the Euclidean norm of the gap to the
......@@ -196,7 +196,7 @@ class SimpleAlgorithm:
# update
u = u_new
if 'diagnostic' in self.config:
if 'diagnostic' in self.config and self.config['diagnostic']:
# For Douglas-Rachford,in general it is appropriate to monitor the
# SHADOWS of the iterates, since in the convex case these converge
# even for beta=1.
......@@ -219,7 +219,7 @@ class SimpleAlgorithm:
change = change[1:iter + 1]
output = {'u': u, 'u1': u1, 'u2': u2, 'iter': iter, 'change': change}
if 'diagnostic' in self.config:
if 'diagnostic' in self.config and self.config['diagnostic']:
gap = gap[1:iter + 1]
shadow_change = shadow_change[1:iter + 1]
output['gap'] = gap
......
import matplotlib.pyplot as plt
from numpy import max, min
from matplotlib.widgets import Slider
class SingleStackViewer:
def __init__(self, volume):
fig, ax = plt.subplots()
ax.volume = volume
ax.index = volume.shape[0] // 2
maxval = max(abs(volume))
im = ax.imshow(volume[ax.index], cmap='seismic', vmin=-maxval, vmax=maxval)
plt.colorbar(im)
ax.set_title('use [j,k] to navigate')
fig.canvas.mpl_connect('key_press_event', self.process_key)
def process_key(self, event):
"""Look for special events captured by mpl_connect, send to correct function """
fig = event.canvas.figure
ax = fig.axes[0]
if event.key == 'j':
self.previous_slice(ax)
elif event.key == 'k':
self.next_slice(ax)
fig.canvas.draw()
def previous_slice(self, ax):
"""Go to the previous slice."""
volume = ax.volume
ax.index = (ax.index - 1) % volume.shape[0] # wrap around using %
ax.images[0].set_array(volume[ax.index])
def next_slice(self, ax):
"""Go to the next slice."""
volume = ax.volume
ax.index = (ax.index + 1) % volume.shape[0]
ax.images[0].set_array(volume[ax.index])
class XYZStackViewer:
def __init__(self, volume, limit_sliders=(0, 0, 0), cmap='seismic'):
"""
Dynamic matplotlib plot showing slices out of a 3d dataset
:param volume: real-valued 3d numpy array
:param limit_sliders: sliders allow range given by [limit, N-limit], where N is the maximal length
:param cmap: color map to use (e.g. 'seismic' or 'viridis'). For 'seismic, will center on 0, else'
"""
self.volume = volume
if cmap == 'seismic':
self.maxval = max(abs(volume))
self.minval = -self.maxval
else:
self.maxval = max(volume)
self.minval = min([0, min(volume)])
self.indices = [s // 2 for s in volume.shape]
self.shape = self.volume.shape
self.fig, self.ax = plt.subplots(1, 3, figsize=(10,4))
im = self.ax[0].imshow(volume[self.indices[0], :, :], cmap=cmap, vmin=-self.minval, vmax=self.maxval)
im = self.ax[1].imshow(volume[:, self.indices[0], :], cmap=cmap, vmin=-self.minval, vmax=self.maxval)
im = self.ax[2].imshow(volume[:, :, self.indices[0]], cmap=cmap, vmin=-self.minval, vmax=self.maxval)
self.ax[1].set_yticks([])
self.ax[2].set_yticks([])
self.fig.subplots_adjust(bottom=0.1, right=0.85)
cbar_ax = self.fig.add_axes([0.87, 0.2, 0.01, 0.6])
plt.colorbar(im, cax=cbar_ax)
axcolor = None # 'lightgoldenrodyellow'
sliderlength = 0.17
subax1 = self.fig.add_axes([0.15, 0.1, sliderlength, 0.03], facecolor=axcolor)
self.ax1slider = Slider(subax1, 'Slice', 0+limit_sliders[0], self.shape[0]-limit_sliders[0],
valinit=self.indices[0], valstep=1, valfmt='%d')
self.ax1slider.on_changed(self.update_1)
subax2 = self.fig.add_axes([0.41, 0.1, sliderlength, 0.03], facecolor=axcolor)
self.ax2slider = Slider(subax2, 'Slice', 0+limit_sliders[1], self.shape[0]-limit_sliders[1],
valinit=self.indices[0], valstep=1, valfmt='%d')
self.ax2slider.on_changed(self.update_2)
subax3 = self.fig.add_axes([0.67, 0.1, sliderlength, 0.03], facecolor=axcolor)
self.ax3slider = Slider(subax3, 'Slice', 0+limit_sliders[2], self.shape[0]-limit_sliders[2],
valinit=self.indices[0], valstep=1, valfmt='%d')
self.ax3slider.on_changed(self.update_3)
self.fig.show()
# self.fig.canvas.mpl_connect('key_press_event', self.process_key)
def update_1(self, n):
"""Update subfigure 1 on change of the slider """
self.indices[0] = int(n)
self.ax[0].images[0].set_array(self.volume[int(n), :, :])
self.fig.canvas.draw_idle()
def update_2(self, n):
"""Update subfigure 2 on change of the slider """
self.indices[1] = int(n)
self.ax[1].images[0].set_array(self.volume[:, int(n), :])
self.fig.canvas.draw_idle()
def update_3(self, n):
"""Update subfigure 3 on change of the slider """
self.indices[2] = int(n)
self.ax[2].images[0].set_array(self.volume[:, :, int(n)])
self.fig.canvas.draw_idle()
# def process_key(self, event):
# """Look for special events captured by mpl_connect, send to correct function """
# fig = event.canvas.figure
# ax = fig.axes[0]
# if event.key == 'j':
# self.previous_slice(ax)
# elif event.key == 'k':
# self.next_slice(ax)
# fig.canvas.draw()
#
# def previous_slice(self, ax):
# """Go to the previous slice on all plots."""
# # do stuff: set_array for all, set self.indices, set_val for sliders
#
# def next_slice(self, ax):
# """Go to the next slice."""
# # do something
from .orbitaltomog_data_processor import shifted_fft, shifted_ifft, support_from_autocorrelation, bin_2d_array
import numpy as np
from skimage.io import imread
from .Graphics.stack_viewer import XYZStackViewer
def data_processor(config):
inp = imread(config['data_filename'])
ny, nx, nz = inp.shape
config['data'] = abs(inp)
# Keep the same resolution? TODO: streamline for 3D arpes data with many unknown values (NaNs) in the data set
if config['Ny'] is None and config['Nx'] is None and config['Nz'] is None:
config['Ny'], config['Nx'], config['Nz'] = ny, nx, nz
elif config['Ny'] == ny and config['Nx'] == nx and config['Nz'] == nz:
pass
elif ny % config['Ny'] == 0 and nx % config["Nx"] == 0 and nz % config['Nz'] == 0:
# binning must be done for the intensity-data, as that preserves the normalization
if not ('from intensity data' in config and config['from intensity data']):
config['data'] = np.sqrt(bin_array(config['data'] ** 2, (config['Ny'], config["Nx"], config['Nz'])))
else:
config['data'] = bin_2d_array(config['data'], (config['Ny'], config["Nx"]))
else:
raise ValueError('Incompatible values for Ny, Nx, Nz given in configuration dict')
# Calculate electric field
if 'from intensity data' in config and config['from intensity data']:
# avoid sqrt of negative numbers (due to background subtraction)
config['data'] = np.where(config['data'] > 0,
np.sqrt(abs(config['data'])),
np.zeros_like(config['data']))
config['norm_data'] = np.sqrt(np.sum(config['data'] ** 2))
# Object support determination
try:
config['support'] = imread(config['support_filename'])
except KeyError: # 'support filename does not exist, so define a support here'
if 'threshold for support' not in config:
config['threshold for support'] = 0.1
config['support'] = support_from_autocorrelation(config['data'],
threshold=config['threshold for support'],
absolute_autocorrelation=True,
binary_dilate_support=1)
if ('use_sparsity_with_support' in config
and config['use_sparsity_with_support']
and 'sparsity_support' not in config):
if 'threshold for support' not in config:
config['threshold for support'] = 0.1
config['sparsity_support'] = support_from_autocorrelation(config['data'],
threshold=config['threshold for support'],
binary_dilate_support=1)
# Initial guess
ph_init = 2 * np.pi * np.random.rand(inp.shape)
config['u_0'] = inp * np.exp(1j * ph_init)
if config['dataprocessor_plotting']:
XYZStackViewer(inp, cmap='viridis')
XYZStackViewer(shifted_fft(inp).real)
# Other settings
config['fresnel_nr'] = 0
config['FT_conv_kernel'] = 1
config['use_farfield_formula'] = True
config['magn'] = 1
config['data_sq'] = abs(config['data']) ** 2
config['data_zeros'] = 1 # probably here I can put in the mask with NaNs (where we have no sensitivity
return config
def bin_array(arr: np.ndarray, new_shape: any, pad_zeros=True) -> np.ndarray:
"""
Reduce the size of an array by binning
:param arr: original
:param new_shape: tuple which must be an integer divisor of the original shape, or integer to bin by that factor
:return: new array
"""
# make tuple with new shape
if type(new_shape) == int: # binning factor is given
_shape = tuple([i // new_shape for i in arr.shape])
binfactor = tuple([new_shape for i in _shape])
else:
_shape = new_shape
binfactor = tuple([s // _shape[i] for i, s in enumerate(arr.shape)])
# determine if padding is needed
padding = tuple([(0, (binfactor[i] - s % binfactor[i]) % binfactor[i]) for i, s in enumerate(arr.shape)])
if pad_zeros and np.any(np.array(padding) != 0):
_arr = np.pad(arr, padding, mode='constant', constant_values=0) # pad array
_shape = tuple([s//binfactor[i] for i, s in enumerate(_arr.shape)]) # update binned size due to padding
else:
_arr = arr # expected to fail if padding has non-zeros
# send to 2d or 3d padding functions
try:
if len(arr.shape) == 2:
out = bin_2d_array(_arr, _shape)
elif len(arr.shape) == 3:
out = bin_3d_array(_arr, _shape)
else:
raise NotImplementedError('Cannot only bin 3d or 2d arrays')
return out
except ValueError:
raise ValueError("Cannot bin data with this shape. Try setting pad_zeros=True, or change the binning.")
def bin_3d_array(arr: np.ndarray, new_shape: tuple) -> np.ndarray:
""""
bins a 2D numpy array
Args:
arr: input array to be binned
new_shape: shape after binning, must be an integer divisor of the original shape
Returns:
binned np array
"""
shape = (new_shape[0], arr.shape[0] // new_shape[0],
new_shape[1], arr.shape[1] // new_shape[1],
new_shape[2], arr.shape[2] // new_shape[2])
return np.sum(arr.reshape(shape), axis=(5, 3, 1))
from .proxoperators import ProxOperator
import numpy as np
class P_nonneg_sparsity(ProxOperator):
"""
Projection subroutine for projecting onto support constraints
......
import numpy as np
from .proxoperators import ProxOperator
class P_parallel_scaled_hyperplane(ProxOperator):
"""
hyperplane projection subroutine of the Cimmino ART method for tomographic recostruction of a density profile.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment