Upload New File

proxtoolbox/experiments/phase/proxtoolbox_experiments_phase_Elser_Experiment.py

+
+import SetProxPythonPath
+from proxtoolbox.experiments.phase.phaseExperiment import PhaseExperiment
+from proxtoolbox import proxoperators
+import numpy as np
+from numpy import fromfile, exp, nonzero, zeros, pi, resize, real, angle
+from numpy.random import rand
+from numpy.linalg import norm
+from numpy.fft import fftshift, fft2, ifft2
+
+import proxtoolbox.utils as utils
+from proxtoolbox.utils.cell import Cell, isCell
+
+#for downloading data
+import proxtoolbox.utils.GetData as GetData
+
+import matplotlib
+import matplotlib.pyplot as plt
+from matplotlib.pyplot import subplots, show, figure
+
+
+class Elser_Experiment(PhaseExperiment):
+    '''
+    Elser's Crystallography experiment class
+    '''
+
+    @staticmethod
+    def getDefaultParameters():
+        defaultParams = {
+            'experiment_name' : 'Elser',
+            'object': 'complex',
+            'constraint': 'atom count',
+            'distance': 'far field',
+            'Nx': 128,
+            'Ny': 128,
+            'Nz': 1,
+            'Atoms':100,
+            'category': 'E',
+            'product_space_dimension': 2,
+            'MAXIT': 10000,
+            'TOL': 5e-5,
+            'lambda_0': 0.85,
+            'lambda_max': 0.85,
+            'lambda_switch': 20,
+            'data_ball': 1e-30,
+            'diagnostic': True,
+	    'iterate_monitor_name': 'FeasibilityIterateMonitor',
+            'verbose': 0,
+            'graphics': 1
+        }
+        return defaultParams
+
+
+    def __init__(self, Atoms=100, category='E', **kwargs):
+        super(Elser_Experiment, self).__init__(**kwargs)
+        self.fresnel_nr = None
+        self.use_farfield_formula = None
+        self.data_zeros = None
+        self.supp_ampl = None
+        self.indicator_ampl = None
+        self.abs_illumination = None
+        self.illumination_phase = None
+        self.FT_conv_kernel = Cell(1)
+        self.Atoms=Atoms
+        self.category=category
+    def loadData(self):
+    # def loadElserData(filename, M):
+        """
+        Load Elser dataset. Create the initial iterate.
+        """
+        #make sure input data can be found, otherwise download it
+        GetData.getData('Elser')
+	#defaultParams = getDefaultParameters()
+        # So far only Nx to set the desired
+        # resolution
+        newres = self.Nx
+
+        # The following value for snr does not work well in Python
+        # as this creates numerical issues with Poisson noise
+        # (data_ball is far too small)
+        # snr = self.data_ball
+        # Instead, we choose the following bigger value which
+        # works well.
+        snr = 1e-8
+
+        print('Loading data.')
+        fn="".join(['../InputData/Phase/Elser/data', str(self.Atoms), self.category])
+        # read text data file as a matrix of numbers
+        data = np.loadtxt(fn, delimiter="\t")
+        data = np.sqrt(data)/self.Nx
+        data = np.c_[data, np.zeros(self.Nx)] # add extra column
+        print(data)
+        data = data.flatten() # make it a one-dimensional array
+        M = fftshift(data)
+        # have to look at what Elser is actually doing with his data...
+        self.magn = 1  # magnification factor
+        self.farfield = True
+        print('Using farfield formula.')
+
+        self.rt_data = Cell(1)
+        self.rt_data[0] = M
+        # standard for the main program is that
+        # the data field is the magnitude SQUARED
+        # in Luke.m this is changed to the magnitude.
+        self.data = Cell(1)
+        self.data[0] = M**2
+        self.norm_rt_data = np.linalg.norm(ifft2(M), 'fro')
+        self.data_zeros = np.where(M == 0)
+        self.support_idx = np.nonzero(S)
+        self.sets = 2
+
+        # The support constraint is sparsity determined by the number of atoms indicated in the
+        # data file name...this needs to be installed
+
+        # initial guess: follow Elser's initialization...
+
+    def setupProxOperators(self):
+        """
+        Determine the prox operators to be used for this experiment
+        """
+        super(Elser_Experiment, self).setupProxOperators()  # call parent's method
+
+        self.propagator = 'Propagator_FreFra'
+        self.inverse_propagator = 'InvPropagator_FreFra'
+
+        # remark: self.farfield is always true for these problems
+        # self.proxOperators already contains a prox operator at slot 0.
+        # Here, we add the second one.
+        if self.constraint == 'phaselift':
+            self.proxOperators.append('P_Rank1')
+        elif self.constraint == 'phaselift2':
+            self.proxOperators.append('P_rank1_SR')
+        else:
+            self.proxOperators.append('Approx_Pphase_FreFra_Poisson')
+        self.nProx = self.sets
+
+        # Now we adjust data structures and prox operators according to the algorithm
+        # What follows wa set up for JWST and works in that context.  Not
+        # sure how much of this is trasferrable to Elsers problems (DRL 14.08.2020)
+        if self.algorithm_name == 'ADMM':
+            # set-up variables in cells for block-wise algorithms
+            # ADMM has three blocks of variables, primal, auxilliary (primal
+            # variables in the image space of some linear mapping) and dual.
+            # we sort these into a 3D cell.
+            u0 = self.u0
+            self.u0 = Cell(3)
+            self.proxOperators = []
+            self.productProxOperators = []
+            K = self.product_space_dimension - 1
+            self.product_space_dimension = K
+            self.u0[0] = u0[K]
+            prox_FreFra_Poisson_class = getattr(proxoperators, 'Approx_Pphase_FreFra_Poisson')
+            prox_FreFra_Poisson = prox_FreFra_Poisson_class(self)
+            self.u0[1] = Cell(K)
+            self.u0[2] = Cell(K)
+            for j in range(K):
+                tmp_u0 = prox_FreFra_Poisson.eval(u0[j], j)
+                self.u0[1][j] = fft2(self.FT_conv_kernel[j]*tmp_u0) / (self.Nx*self.Ny)
+                self.u0[2][j] = self.u0[1][j] / self.lambda_0
+            self.proxOperators.append('Prox_primal_ADMM_indicator')
+            self.proxOperators.append('Approx_Pphase_FreFra_ADMM_Poisson')
+        elif self.algorithm_name == 'AvP2':
+            # u_0 should be a cell...we change it into a cell of cells
+            u0 = self.u0
+            self.u0 = Cell(3)
+            prox2_class = getattr(proxoperators, self.proxOperators[1])
+            prox2 = prox2_class(self)
+            self.u0[2] = prox2.eval(u0)
+            P_Diag_class = getattr(proxoperators, 'P_diag')
+            P_Diag_prox = P_Diag_class(self)
+            tmp_u = P_Diag_prox.eval(self.u0[2])
+            self.u0[1] = tmp_u[0]
+            self.u0[0] = u0[self.product_space_dimension-1]
+        elif self.algorithm_name == 'PHeBIE':
+            # set up variables in cells for block-wise, implicit-explicit algorithms
+            u0 = self.u0
+            self.u0 = Cell(2)
+            self.product_space_dimension = self.product_space_dimension - 1
+            self.u0[0] = u0[self.product_space_dimension].copy()  # copy last element of prev u0 into first slot
+            self.u0[1] = Cell(self.product_space_dimension)
+            for j in range(self.product_space_dimension):
+                self.u0[1][j] = u0[j]
+            self.proxOperators = []
+            self.nProx = 2
+            self.proxOperators.append('P_amp_support')  # project onto the support/amplitude constraints
+            self.proxOperators.append('Prox_product_space')
+            self.productProxOperators = []
+            self.n_product_Prox = self.product_space_dimension
+            for _j in range(self.n_product_Prox):
+                self.productProxOperators.append('Approx_Pphase_FreFra_Poisson')
+        elif self.algorithm_name == 'Wirtinger':
+            # Contrary to the Matlab code, we use cells instead of 3D arrays
+            L = len(self.FT_conv_kernel)
+            self.product_space_dimension = L
+            self.masks = Cell(L)
+            for j in range(L):
+                self.masks[j] = self.indicator_ampl * self.FT_conv_kernel[j]
+            self.data_zeros = None
+            self.FT_conv_kernel = self.masks
+            self.proxOperators[1] = 'Approx_Pphase_JWST_Wirt'
+            self.use_farfield_formula = True
+            self.fresnel_nr = 0
+
+    def show(self):
+        """
+        Generate graphical output from the solution
+        """
+
+        # display plot of far field data
+        # figure(123)
+        f, (ax1) = subplots(1, 1,
+                                 figsize=(self.figure_width, self.figure_height),
+                                 dpi=self.figure_dpi)
+        im = ax1.imshow(log10(self.dp + 1e-15))
+        f.colorbar(im, ax=ax1, fraction=0.046, pad=0.04)
+        ax1.set_title('Far field data')
+        plt.subplots_adjust(wspace=0.3)  # adjust horizontal space (width)
+        # between subplots (default = 0.2)
+        f.suptitle('Elser Data')
+
+        # call parent to display the other plots
+        super(Elser_Experiment, self).show()
+
+
+        u_m = self.output['u_monitor']
+        if isCell(u_m):
+            u = u_m[0]
+            if isCell(u):
+                u = u[0]
+            u2 = u_m[len(u_m)-1]
+            if isCell(u2):
+                u2 = u2[len(u2)-1]
+        else:
+            u2 = u_m
+            if u2.ndim > 2:
+                u2 = u2[:,:,0]
+            u = self.output['u']
+            if isCell(u):
+                u = u[0]
+            elif u.ndim > 2:
+                u = u[:,:,0]
+
+        algo_desc = self.algorithm.getDescription()
+        title = "Algorithm " + algo_desc
+
+        # figure 904
+        titles = ["Best approximation amplitude - physical constraint satisfied",
+                  "Best approximation phase - physical constraint satisfied",
+                  "Best approximation amplitude - Fourier constraint satisfied",
+                  "Best approximation phase - Fourier constraint satisfied"]
+        f = self.createFourImageFigure(u, u2, titles)
+        f.suptitle(title)
+        plt.subplots_adjust(hspace = 0.3) # adjust vertical space (height) between subplots (default = 0.2)
+
+        # figure 900
+        changes = self.output['stats']['changes']
+        time = self.output['stats']['time']
+        time_str = "{:.{}f}".format(time, 5) # 5 is precision
+        label = "Iterations (time = " + time_str + " s)"
+
+        f, ((ax1, ax2), (ax3, ax4)) = subplots(2, 2, \
+                    figsize = (self.figure_width, self.figure_height),
+                    dpi = self.figure_dpi)
+        self.createImageSubFigure(f, ax1, abs(u), titles[0])
+        self.createImageSubFigure(f, ax2, real(u), titles[1])
+
+        ax3.semilogy(changes)
+        ax3.set_xlabel(label)
+        ax3.set_ylabel('Log of change in iterates')
+
+        if 'gaps' in self.output['stats']:
+            gaps = self.output['stats']['gaps']
+            ax4.semilogy(gaps)
+            ax4.set_xlabel(label)
+            ax4.set_ylabel('Log of the gap distance')
+
+        f.suptitle(title)
+        plt.subplots_adjust(hspace = 0.3) # adjust vertical space (height) between subplots (default = 0.2)
+        plt.subplots_adjust(wspace = 0.3) # adjust horizontal space (width) between subplots (default = 0.2)
+
+        show()
+
+if __name__ == "__main__":
+	Elser = Elser_Experiment(Atoms=200, category='E', algorithm='RRR', lambda_0=0.95, lambda_max=0.95, anim=True)
+	Elser.run()
+	Elser.show()