orbital_tomog_3d_data_processor.py 5.36 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
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))