work towards 3D orbital imaging

proxtoolbox/Problems/OrbitalTomog/orbital_tomog_3d_data_processor.py

@@ -10,11 +10,11 @@ def data_processor(config):
     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:
+    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 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:
+    elif (ny != config['Ny'] or nz != 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']):
             config['data'] = np.sqrt(bin_array(config['data'] ** 2, (config['Ny'], config["Nx"], config['Nz'])))
@@ -23,6 +23,11 @@ def data_processor(config):
     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)
@@ -52,7 +57,7 @@ def data_processor(config):
                                                                   binary_dilate_support=1)
 
     # Initial guess
-    ph_init = 2 * np.pi * np.random.rand(inp.shape)
+    ph_init = 2 * np.pi * np.random.random_sample(inp.shape)
     config['u_0'] = inp * np.exp(1j * ph_init)
 
     if config['dataprocessor_plotting']:
@@ -107,7 +112,7 @@ def bin_array(arr: np.ndarray, new_shape: any, pad_zeros=True) -> np.ndarray:
 
 def bin_3d_array(arr: np.ndarray, new_shape: tuple) -> np.ndarray:
     """"
-    bins a 2D numpy array
+    bins a 3D numpy array
      Args:
         arr: input array to be binned
         new_shape: shape after binning, must be an integer divisor of the original shape
@@ -117,4 +122,10 @@ def bin_3d_array(arr: np.ndarray, new_shape: tuple) -> np.ndarray:
     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))
+    if np.any(np.isnan(arr)):
+        binfactor = 1
+        for i, s in enumerate(arr.shape):
+            binfactor *= new_shape[i] / s
+        return np.nanmean(arr.reshape(shape), axis=(5, 3, 1)) * binfactor
+    else:
+        return np.sum(arr.reshape(shape), axis=(5, 3, 1))

proxtoolbox/Problems/OrbitalTomog/pentacene_3d_config.py

@@ -10,6 +10,7 @@ new_config = {
     # What is the name of the data file?
     'data_filename': 'pentacene_3d_arpes.tif',  # In the directory '../../../InputData/OrbitalTomog/'
     'from intensity data': False,  # File gives field amplitudes
+    'sensitivity_filename': 'pentacene_3d_arpes_sensitivity.tif',
 
     # What type of object are we working with?
     # Options are: 'phase', 'real', 'nonnegative', 'complex'
@@ -24,9 +25,8 @@ new_config = {
     # What type of measurements are we working with?
     # Options are: 'single diffraction', 'diversity diffraction',
     #              'ptychography', and 'complex'
-    # Options are: '2D_ARPES', '3D_ARPES',
-    #              '2D_time', and '3D_time'
-    'experiment': '3D_ARPES',
+    # Orbital imaging options are: '2D ARPES', '3D ARPES', '2D time', and '3D time'
+    'experiment': '3D ARPES',
 
     # Next we move to things that most of our users will know
     # better than we will.  Some of these may be overwritten in the
@@ -35,9 +35,11 @@ new_config = {
     # Options are: 'far field' or 'near field',
     'distance': 'far field',
 
-    # What are the dimensions of the measurements?
-    'Nx': 64,
-    'Ny': 64,
+    # What are the dimensions of the measurements to be analysed? If smaller than data dimensions, will interpolate.
+    # If not given, values are taken from the data.
+    # 'Ny': 64,
+    # 'Nx': 64,
+    # 'Nz': 64,
 
     # 'fresnel_nr' : 0,
 

proxtoolbox/Problems/OrbitalTomog/phase.py

@@ -74,8 +74,10 @@ class Phase(Problem):
             raise ValueError('Constraint not recognized')
 
         # Projector 2 (k / Fourier space)
-#         used_proxoperators[1] = 'P_M'  # 'Approx_P_FreFra_Poisson'
-        used_proxoperators[1] = 'Approx_P_FreFra_Poisson'
+        if self.config['experiment'] == '3D ARPES':
+            used_proxoperators[1] = 'P_M_masked'
+        else:
+            used_proxoperators[1] = 'Approx_P_FreFra_Poisson'
 
         self.config['proxoperators'] = []
 
@@ -110,9 +112,9 @@ class Phase(Problem):
             #     u_1[:, :, j] = proj1.work(self.config['u_0'])
             # end
         else:  # i.e. formulation=='product space'
-            proj_1 = self.config['proxoperators'][0](self.config)#, self.back_end)
+            proj_1 = self.config['proxoperators'][0](self.config)  # , self.back_end)
             u_1 = proj_1.work(self.config['u_0'])
-            proj_2 = self.config['proxoperators'][1](self.config)#, self.back_end)
+            proj_2 = self.config['proxoperators'][1](self.config)  # , self.back_end)
             u_2 = proj_2.work(u_1)
 
         # estimate the gap in the relevant metric