Commit 229a4f21 authored by Matthijs's avatar Matthijs
Browse files

3D compatibility (not tested!)

parent d9d8a3d8
from .Phase_graphics import Phase_graphics from .Phase_graphics import Phase_graphics
from proxtoolbox.Utilities.OrbitalTomog.array_tools import * # from proxtoolbox.Utilities.OrbitalTomog.array_tools import *
from .complex_field_visualization import complex_to_rgb from .complex_field_visualization import complex_to_rgb
from .stack_viewer import XYZStackViewer, SingleStackViewer from .stack_viewer import XYZStackViewer, SingleStackViewer
__all__ = ["Phase_graphics", "roll_to_pos", "complex_to_rgb", 'fourier_interpolate', __all__ = ["Phase_graphics", "complex_to_rgb", # 'fourier_interpolate', "roll_to_pos",
'XYZStackViewer', 'SingleStackViewer'] 'XYZStackViewer', 'SingleStackViewer']
...@@ -11,19 +11,20 @@ def data_processor(config): ...@@ -11,19 +11,20 @@ def data_processor(config):
ny, nx, nz = inp.shape ny, nx, nz = inp.shape
config['data'] = abs(inp) config['data'] = abs(inp)
# Keep the same resolution? TODO: streamline for 3D arpes data with many unknown values (NaNs) in the data set # Keep the same resolution?
if 'Ny' not in config or 'Nx' not in config or 'Nz' not in config: if 'Ny' not in config or 'Nx' not in config or 'Nz' not in config:
print('Setting problem dimensions based on data') print('Setting problem dimensions based on data')
config['Ny'], config['Nx'], config['Nz'] = ny, nx, nz config['Ny'], config['Nx'], config['Nz'] = ny, nx, nz
elif (ny != config['Ny'] or nz != config['Nx'] or nz != config['Nz']) \ elif ny != config['Ny'] or nx != config['Nx'] or nz != config['Nz']:
and 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']): if not ('from intensity data' in config and config['from intensity data']):
# binning must be done for the intensity-data, as that preserves the normalization
config['data'] = np.sqrt(bin_array(config['data'] ** 2, (config['Ny'], config["Nx"], config['Nz']))) config['data'] = np.sqrt(bin_array(config['data'] ** 2, (config['Ny'], config["Nx"], config['Nz'])))
else: else:
config['data'] = bin_2d_array(config['data'], (config['Ny'], config["Nx"])) config['data'] = bin_array(config['data'], (config['Ny'], config["Nx"]))
else: # elif ny == config['Ny'] and nx == config['Nx'] or nz == config['Nz']:
raise ValueError('Incompatible values for Ny, Nx, Nz given in configuration dict') # pass
# else:
# raise ValueError('Incompatible values for Ny, Nx, Nz given in configuration dict')
# Load measurement sensitivity, determine data_zeros # Load measurement sensitivity, determine data_zeros
sensitivity = imread(config['sensitivity_filename']) sensitivity = imread(config['sensitivity_filename'])
......
...@@ -7,6 +7,7 @@ from proxtoolbox.ProxOperators.proxoperators import ProxOperator ...@@ -7,6 +7,7 @@ from proxtoolbox.ProxOperators.proxoperators import ProxOperator
from proxtoolbox.Problems.OrbitalTomog import Graphics from proxtoolbox.Problems.OrbitalTomog import Graphics
from numpy.linalg import norm from numpy.linalg import norm
from numpy import square, sqrt from numpy import square, sqrt
from proxtoolbox.Utilities.OrbitalTomog import interpolation, array_tools, binning
class Phase(Problem): class Phase(Problem):
...@@ -151,17 +152,33 @@ class Phase(Problem): ...@@ -151,17 +152,33 @@ class Phase(Problem):
""" """
# Center the solution (since position is a degree of freedom, # Center the solution (since position is a degree of freedom,
# and if desired, interpolate the results. # and if desired, interpolate the results.
# TODO: make compatible with non-2D data (e.g. 3D) center = tuple([s//2 for s in self.config['u'].shape])
yc, xc = int(self.config['Ny'] / 2), int(self.config["Nx"] / 2)
for key in ['u', 'u1', 'u2']: for key in ['u', 'u1', 'u2']:
self.output[key] = Graphics.roll_to_pos(self.output[key], yc, xc, move_maximum=True) # first move maximum self.output[key] = array_tools.roll_to_pos(self.output[key], pos=center, move_maximum=True)
self.output[key] = Graphics.roll_to_pos(self.output[key], yc, xc) # then move center of mass. self.output[key] = array_tools.roll_to_pos(self.output[key], pos=center)
# This sequence will work for objects *with a small support* even if they lie over the edge of the array # This sequence will work for objects *with a small support* even if they lie over the edge of the array
if 'interpolate_result' in self.config and self.config['interpolate_result']: if 'interpolate_result' in self.config and self.config['interpolate_result']:
self.output[key] = Graphics.fourier_interpolate(self.output[key]) self.output[key] = interpolation.fourier_interpolate(self.output[key])
if 'zoomin_on_result' in self.config and self.config['zoomin_on_result']: if 'zoomin_on_result' in self.config and self.config['zoomin_on_result']:
zmy, zmx = int(self.config['Ny'] / 4), int(self.config["Nx"] / 4) if self.config[key].ndims == 2:
self.output[key] = self.output[key][zmy:-zmy, zmx:-zmx] zmy, zmx = self.config[key].shape # self.config['Ny'] // 4, self.config["Nx"] // 4
self.output[key] = self.output[key][zmy:-zmy, zmx:-zmx]
elif self.config[key].ndims == 3:
zmy, zmx, zmz = self.config[key].shape
# (self.config['Ny'] // 4, self.config["Nx"] // 4, self.config['Nz'] // 4)
self.output[key] = self.output[key][zmy:-zmy, zmx:-zmx, zmz:-zmz]
# # Old code, developed for 2D orbital imaging:
# yc, xc = int(self.config['Ny'] / 2), int(self.config["Nx"] / 2)
# for key in ['u', 'u1', 'u2']:
# self.output[key] = array_tools.roll_to_pos(self.output[key], yc, xc, move_maximum=True) # first move maximum
# self.output[key] = array_tools.roll_to_pos(self.output[key], yc, xc) # then move center of mass.
# # This sequence will work for objects *with a small support* even if they lie over the edge of the array
# if 'interpolate_result' in self.config and self.config['interpolate_result']:
# self.output[key] = interpolation.fourier_interpolate(self.output[key])
# if 'zoomin_on_result' in self.config and self.config['zoomin_on_result']:
# zmy, zmx = int(self.config['Ny'] / 4), int(self.config["Nx"] / 4)
# self.output[key] = self.output[key][zmy:-zmy, zmx:-zmx]
def show(self): def show(self):
""" """
......
...@@ -8,7 +8,19 @@ def shift_array(arr, dy, dx): ...@@ -8,7 +8,19 @@ def shift_array(arr, dy, dx):
return temp return temp
def roll_to_pos(arr, y=0, x=0, move_maximum=False, by_abs_val=True): def roll_to_pos(arr: np.ndarray, y: int = 0, x: int = 0, pos: tuple = None, move_maximum: bool = False,
by_abs_val: bool = True) -> np.ndarray:
"""
Shift the center of mass of an array to the given position by cyclic permutation
:param arr: 2d array, works best for well-centered feature with limited support
:param y: position parameter
:param x: position parameter for second dimension
:param pos: tuple with the new position, overriding y,x values. should be used for higher-dimensional arrays
:param move_maximum: if true, look only at max-value
:param by_abs_val: take abs value for the determination of max-val or center-of-mass
:return: array like original
"""
if move_maximum: if move_maximum:
if by_abs_val or arr.dtype in [np.complex64, np.complex128]: if by_abs_val or arr.dtype in [np.complex64, np.complex128]:
old = np.floor(measurements.maximum_position(abs(arr))) old = np.floor(measurements.maximum_position(abs(arr)))
...@@ -19,7 +31,12 @@ def roll_to_pos(arr, y=0, x=0, move_maximum=False, by_abs_val=True): ...@@ -19,7 +31,12 @@ def roll_to_pos(arr, y=0, x=0, move_maximum=False, by_abs_val=True):
old = np.floor(measurements.center_of_mass(abs(arr))) old = np.floor(measurements.center_of_mass(abs(arr)))
else: else:
old = np.floor(measurements.center_of_mass(arr)) old = np.floor(measurements.center_of_mass(arr))
temp = shift_array(arr, int(y - old[0]), int(x - old[1])) if pos is not None: # dimension-independent method
shifts = tuple([int(np.round(pos[i]-old[i])) for i in range(len(pos))])
dims = tuple([i for i in range(len(pos))])
temp = np.roll(arr, shift=shifts, axis=dims)
else: # old method
temp = shift_array(arr, int(y - old[0]), int(x - old[1]))
return temp return temp
......
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