orbital_tomog_3d_data_processor.py

from proxtoolbox.Problems.OrbitalTomog.planar_molecule.orbitaltomog_data_processor import support_from_autocorrelation
from proxtoolbox.Utilities.OrbitalTomog.array_tools import shifted_fft
from proxtoolbox.Utilities.OrbitalTomog.binning import bin_array
import numpy as np
from skimage.io import imread
from proxtoolbox.Problems.OrbitalTomog.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?
    if 'Ny' not in config or 'Nx' not in config or 'Nz' not in config:
        print('Setting problem dimensions based on data')
        config['Ny'], config['Nx'], config['Nz'] = ny, nx, nz
    elif ny != config['Ny'] or nx != config['Nx'] or nz != config['Nz']:
        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'])))
        else:
            config['data'] = bin_array(config['data'], (config['Ny'], config["Nx"]))
    # elif ny == config['Ny'] and nx == config['Nx'] or nz == config['Nz']:
    #     pass
    # else:
    #     raise ValueError('Incompatible values for Ny, Nx, Nz given in configuration dict')

    # Load measurement sensitivity, determine data_zeros
    sensitivity = imread(config['sensitivity_filename'])
    config['data_zeros'] = sensitivity == 0
    assert config['data_zeros'].shape == config['data'].shape, 'Non-matching sensitivity and data arrays'

    # 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.random_sample(inp.shape)
    config['u_0'] = inp * np.exp(1j * ph_init)

    if config['dataprocessor_plotting']:
        input_viewer = XYZStackViewer(inp, cmap='viridis')
        config['input viewer'] = input_viewer  # the reference keeps the matplotlib object alive
        support_viewer = XYZStackViewer(shifted_fft(inp).real)
        config['support viewer'] = support_viewer  # the reference keeps the matplotlib object alive

    # 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

    return config