orbitaltomog_data_processor.py 6.32 KB
Newer Older
Russell Luke's avatar
Russell Luke committed
1
2
import numpy as np
from scipy.io import loadmat
jansen31's avatar
jansen31 committed
3
from scipy.ndimage import binary_dilation
Russell Luke's avatar
Russell Luke committed
4
import matplotlib.pyplot as plt
Matthijs's avatar
Matthijs committed
5
from skimage.io import imread
6
7
from proxtoolbox.Utilities.OrbitalTomog.binning import bin_2d_array
from proxtoolbox.Utilities.OrbitalTomog.array_tools import shifted_ifft, shifted_fft
Russell Luke's avatar
Russell Luke committed
8

9

10
def data_processor(config):
Matthijs's avatar
bugfix    
Matthijs committed
11
12
    if config['data_filename'][1] == ':':
        path_to_file = config['data_filename']
jansen31's avatar
jansen31 committed
13
    else:
Matthijs's avatar
Matthijs committed
14
        path_to_file = '../../../../InputData/OrbitalTomog/' + config['data_filename']
Matthijs's avatar
bugfix    
Matthijs committed
15
16
17
18
19
20

    try:
        if config['data_filename'][-4:] == '.mat':  # for the old  coronene simulations by W. Bennecke
            inp_data = loadmat(path_to_file)
            inp = inp_data["I2"]
        else:
Matthijs's avatar
Matthijs committed
21
            inp = imread(path_to_file)
Matthijs's avatar
bugfix    
Matthijs committed
22
    except FileNotFoundError:
Matthijs's avatar
Matthijs committed
23
        print("Tried path %s, found nothing. "% path_to_file)
Matthijs's avatar
bugfix    
Matthijs committed
24
        path_to_file = input('Please enter the path to the datafile: ')
Matthijs's avatar
Matthijs committed
25
        inp = imread(path_to_file)
Matthijs's avatar
bugfix    
Matthijs committed
26

27
    ny, nx = inp.shape
Russell Luke's avatar
Russell Luke committed
28
    config['data'] = abs(inp)
29

30
    # Keep the same resolution?
jansen31's avatar
jansen31 committed
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
    if config['Ny'] is None and config['Nx'] is None:
        config['Ny'], config['Nx'] = ny, nx
    elif config['Ny'] == ny and config['Nx'] == nx:
        pass
    elif ny % config['Ny'] == 0 and nx % config["Nx"] == 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_2d_array(config['data'] ** 2, (config['Ny'], config["Nx"])))
        else:
            config['data'] = bin_2d_array(config['data'], (config['Ny'], config["Nx"]))
    else:
        raise ValueError('Incompatible values for Ny and Nx 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:
Matthijs's avatar
Matthijs committed
54
        config['support'] = imread(config['support_filename'])
jansen31's avatar
jansen31 committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
    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)
jansen31's avatar
jansen31 committed
71

jansen31's avatar
jansen31 committed
72
    # Autocorrelation
jansen31's avatar
jansen31 committed
73
    config['threshold for support'] = 0.1
Matthijs's avatar
Matthijs committed
74
    autocorrelation = abs(shifted_ifft(inp))
Russell Luke's avatar
Russell Luke committed
75
    config['abs_illumination'] = np.ones_like(config['support'])
jansen31's avatar
jansen31 committed
76

jansen31's avatar
jansen31 committed
77
    # Initial guess
78
    ph_init = 2 * np.pi * np.random.rand(ny, nx)
79
    config['u_0'] = inp * np.exp(1j * ph_init)
80
    config['u_0'] = np.fft.fftn(config['u_0'])
Russell Luke's avatar
Russell Luke committed
81

jansen31's avatar
jansen31 committed
82
83
84
    if ('use_sparsity_with_support' in config
            and config['use_sparsity_with_support']
            and 'sparsity_support' not in config):
Matthijs's avatar
Matthijs committed
85
        config['sparsity_support'] = binary_dilation(autocorrelation > 0.01 * np.amax(autocorrelation)).astype(np.uint)
jansen31's avatar
jansen31 committed
86

Russell Luke's avatar
Russell Luke committed
87
    if config['dataprocessor_plotting']:
88
        plt.figure(figsize=(10, 3.5))
Russell Luke's avatar
Russell Luke committed
89
90
91
92
93
94
95
96
97
98
99
100
101
        plt.subplot(131)
        plt.imshow(inp)
        plt.colorbar()
        plt.title("Photoelectron spectrum")
        plt.subplot(132)
        plt.imshow(autocorrelation.real)
        plt.colorbar()
        plt.title("iFFT(data)")
        plt.subplot(133)
        plt.imshow(config['support'])
        plt.colorbar()
        plt.title("Initial support")
        plt.show()
jansen31's avatar
jansen31 committed
102

103
104
105
106
    # Other settings
    config['fresnel_nr'] = 0
    config['FT_conv_kernel'] = 1
    config['use_farfield_formula'] = True
Russell Luke's avatar
Russell Luke committed
107
    config['magn'] = 1
jansen31's avatar
jansen31 committed
108
    config['data_sq'] = abs(config['data']) ** 2
Russell Luke's avatar
Russell Luke committed
109
    config['data_zeros'] = 1
jansen31's avatar
jansen31 committed
110

111
    return config
jansen31's avatar
jansen31 committed
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151


def support_from_autocorrelation(input_array: np.ndarray,
                                 threshold: float = 0.1,
                                 relative_threshold: bool = True,
                                 input_in_fourier_domain: bool = True,
                                 absolute_autocorrelation: bool = True,
                                 binary_dilate_support: int = 0) -> np.ndarray:
    """
    Determine an initial support from the autocorrelation of an object.

    Args:
        input_array: either the measured diffraction (arpes pattern) or a guess of the object
        threshold: support is everywhere where the autocorrelation is higher than the threshold
        relative_threshold: If true, threshold at threshold*np.amax(autocorrelation)
        input_in_fourier_domain: False if a guess of the object is given in input_array
        absolute_autocorrelation: Take the absolute value of the autocorrelation? (Generally a good idea)
        binary_dilate_support: number of dilation operations to apply to the support.

    Returns:
        support array (same dimensions as input, dtype=np.int)
    """
    if not input_in_fourier_domain:
        kspace = shifted_fft(input_array)
    else:
        kspace = input_array

    autocorrelation = shifted_ifft(abs(kspace))  # Taking absolute value yields autocorrelation by conv. theorem)
    if absolute_autocorrelation:
        autocorrelation = abs(autocorrelation)
    maxval = np.amax(autocorrelation)
    if relative_threshold:
        threshold_val = threshold * maxval
    else:
        threshold_val = threshold
    support = (autocorrelation > threshold_val).astype(np.uint)
    if binary_dilate_support > 0:
        support = binary_dilation(support, iterations=binary_dilate_support).astype(np.uint)

    return support