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

proxtoolbox/Algorithms/SimpleAlgortihm.py

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

proxtoolbox/Problems/OrbitalTomog/Graphics/stack_viewer.py

+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

proxtoolbox/Problems/OrbitalTomog/orbital_tomog_3d_data_processor.py

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

proxtoolbox/ProxOperators/P_nonneg_sparsity.py

 from .proxoperators import ProxOperator
 import numpy as np
 
+
 class P_nonneg_sparsity(ProxOperator):
     """
     Projection subroutine for projecting onto support constraints

proxtoolbox/ProxOperators/P_parallel_scaled_hyperplane.py

 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.