#NEWLINE# def cdi_recon( diffracted_pattern , sample_obj , sup , beta = 1.15 , start_avg = 0.8 , pi_modulus_flag = 'Complex' , sw_flag = True , sw_sigma = 0.5 , sw_threshold = 0.1 , sw_start = 0.2 , sw_end = 0.8 , sw_step = 10 , n_iterations = 1000 , cb_function = None , cb_step = 10 ) : #NEWLINE# #INDENT# '\n    Run reconstruction with difference map algorithm.\n\n    Parameters\n    ---------\n    diffracted_pattern : array\n        diffraction pattern from experiments\n    sample_obj : array\n        initial sample with phase, complex number\n    sup : array\n        initial support\n    beta : float, optional\n        feedback parameter for difference map algorithm.\n        default is 1.15.\n    start_avg : float, optional\n        define the point to start doing average.\n        default is 0.8.\n    pi_modulus_flag : str, optional\n        \'Complex\' or \'Real\', defining the way to perform pi_modulus\n        calculation.\n        default is \'Complex\'.\n    sw_flag : Bool, optional\n        flag to use shrinkwrap algorithm or not.\n        default is True.\n    sw_sigma : float, optional\n        gaussian width used in sw algorithm.\n        default is 0.5.\n    sw_threshold : float, optional\n        shreshold cut in sw algorithm.\n        default is 0.1.\n    sw_start : float, optional\n        at which point to start to do shrinkwrap.\n        defualt is 0.2\n    sw_end : float, optional\n        at which point to stop shrinkwrap.\n        defualt is 0.8\n    sw_step : float, optional\n        the frequency to perform sw algorithm.\n        defualt is 10\n    n_iterations : int, optional\n        number of iterations to run.\n        default is 1000.\n    cb_function : function, optional\n        This is a callback function that expects to receive these\n        four objects: sample_obj, obj_error, diff_error, sup_error.\n        Sample_obj is a 2D array. And obj_error, diff_error, and sup_error\n        are 1D array.\n    cb_step : int, optional\n        define plotting frequency, i.e., if plot_step = 10, plot results\n        after every 10 iterations.\n\n    Returns\n    -------\n    obj_avg : array\n        reconstructed sample object\n    error_dict : dict\n        Error information for all iterations. The dict keys include\n        obj_error, diff_error and sup_error. Obj_error is a list of\n        the relative error of sample object. Diff_error is calculated as\n        the difference between new diffraction pattern and the original\n        diffraction pattern. And sup_error stores the size of the\n        sample support.\n\n    References\n    ----------\n\n    .. [1] V. Elser, "Phase retrieval by iterated projections",\n        J. Opt. Soc. Am. A, vol. 20, No. 1, 2003\n    ' #NEWLINE# diffracted_pattern = np . array ( diffracted_pattern ) #NEWLINE# diffracted_pattern = np . fft . fftshift ( diffracted_pattern ) #NEWLINE# real_operation = False #NEWLINE# if ( pi_modulus_flag . lower ( ) == 'real' ) : #NEWLINE# #INDENT# real_operation = True #NEWLINE# #UNINDENT# gamma_1 = ( ( - 1 ) / beta ) #NEWLINE# gamma_2 = ( 1 / beta ) #NEWLINE# outside_sup_index = ( sup != 1 ) #NEWLINE# error_dict = { #NEWLINE# #INDENT# #NEWLINE# #UNINDENT# } #NEWLINE# obj_error = np . zeros ( n_iterations ) #NEWLINE# diff_error = np . zeros ( n_iterations ) #NEWLINE# sup_error = np . zeros ( n_iterations ) #NEWLINE# sup_old = np . zeros_like ( diffracted_pattern ) #NEWLINE# obj_avg = np . zeros_like ( diffracted_pattern ) . astype ( complex ) #NEWLINE# avg_i = 0 #NEWLINE# time_start = time . time ( ) #NEWLINE# for n in range ( n_iterations ) : #NEWLINE# #INDENT# obj_old = np . array ( sample_obj ) #NEWLINE# obj_a = pi_modulus ( sample_obj , diffracted_pattern ) #NEWLINE# if real_operation : #NEWLINE# #INDENT# obj_a = np . abs ( obj_a ) #NEWLINE# #UNINDENT# obj_a = ( ( ( 1 + gamma_2 ) * obj_a ) - ( gamma_2 * sample_obj ) ) #NEWLINE# obj_a [ outside_sup_index ] = 0 #NEWLINE# obj_b = np . array ( sample_obj ) #NEWLINE# obj_b [ outside_sup_index ] = 0 #NEWLINE# obj_b = ( ( ( 1 + gamma_1 ) * obj_b ) - ( gamma_1 * sample_obj ) ) #NEWLINE# obj_b = pi_modulus ( obj_b , diffracted_pattern ) #NEWLINE# if real_operation : #NEWLINE# #INDENT# obj_b = np . abs ( obj_b ) #NEWLINE# #UNINDENT# sample_obj += ( beta * ( obj_a - obj_b ) ) #NEWLINE# obj_error [ n ] = ( np . linalg . norm ( ( sample_obj - obj_old ) ) / np . linalg . norm ( obj_old ) ) #NEWLINE# diff_error [ n ] = cal_diff_error ( sample_obj , diffracted_pattern ) #NEWLINE# if sw_flag : #NEWLINE# #INDENT# if ( ( n >= ( sw_start * n_iterations ) ) and ( n <= ( sw_end * n_iterations ) ) ) : #NEWLINE# #INDENT# if ( np . mod ( n , sw_step ) == 0 ) : #NEWLINE# #INDENT# logger . info ( 'Refine support with shrinkwrap' ) #NEWLINE# sup_index = find_support ( sample_obj , sw_sigma , sw_threshold ) #NEWLINE# sup = np . zeros_like ( diffracted_pattern ) #NEWLINE# sup [ sup_index ] = 1 #NEWLINE# outside_sup_index = ( sup != 1 ) #NEWLINE# sup_error [ n ] = np . sum ( sup_old ) #NEWLINE# sup_old = np . array ( sup ) #NEWLINE# #UNINDENT# #UNINDENT# #UNINDENT# if ( cb_function and ( ( n_iterations % cb_step ) == 0 ) ) : #NEWLINE# #INDENT# cb_function ( sample_obj , obj_error , diff_error , sup_error ) #NEWLINE# #UNINDENT# if ( n > ( start_avg * n_iterations ) ) : #NEWLINE# #INDENT# obj_avg += sample_obj #NEWLINE# avg_i += 1 #NEWLINE# #UNINDENT# logger . info ( ( '%d object_chi= %f, diff_chi=%f' % ( n , obj_error [ n ] , diff_error [ n ] ) ) ) #NEWLINE# #UNINDENT# obj_avg = ( obj_avg / avg_i ) #NEWLINE# time_end = time . time ( ) #NEWLINE# logger . info ( ( '%d iterations takes %f sec' % ( n_iterations , ( time_end - time_start ) ) ) ) #NEWLINE# error_dict [ 'obj_error' ] = obj_error #NEWLINE# error_dict [ 'diff_error' ] = diff_error #NEWLINE# error_dict [ 'sup_error' ] = sup_error #NEWLINE# return ( obj_avg , error_dict )