Пример #1
0
 def mse_worker(args):
     i_range, psi, det_mod = args
     indvdl_mse = []
     p = det_mod[0].data[0].shape[0]
     for i in i_range:
         psi_sum = CXModal.modal_sum(abs(fft2(psi.getat(i))))
         indvdl_mse.append(sp.sum((abs(psi_sum - det_mod[i]) ** 2.).data[0]) / sp.sum(det_mod[i].data[0] ** 2.))
     return indvdl_mse
Пример #2
0
    def error(self, psi, det_mod):
        """.. method:: error(psi, det_mod)

            Calculates the MSE at a given position given the modes at that position.

            :param CXModal psi: A list of CXData instances containing all modes at a given position.
            :param np.ndarray det_mod: Modulus of measured diffraction pattern.

        """
        mode_sum = CXModal.modal_sum(abs(fft2(psi)))
        return (sp.sum((abs(mode_sum - det_mod) ** 2.).data[0]) / sp.sum(det_mod.data[0] ** 2.))**0.5
Пример #3
0
    def update_probe_nonmodal(self, i, psi_old, psi_new):
        
        d1, d2 = self.positions.data
        id1, id2 = d1//1, d2//1
        object_intensity_max = (abs(self.object)**2.0).data[0].max()

        self.probe.modes[0] += \
            CXData.shift(conj(self.object[id1[i] - self.p2:id1[i] + self.p2, id2[i] - self.p2:id2[i] + self.p2]) *
             (psi_new-psi_old)[0] / object_intensity_max, d1[i]%1, d2[i]%1)

        self.probe.normalise()
Пример #4
0
    def update_probe(self, i, psi_old, psi_new):
        
        d1, d2 = self.positions.data
        id1, id2 = d1//1, d2//1
        object_intensity_max = (abs(self.object)**2.0).data[0].max()

        for mode in range(len(self.probe)):
            self.probe.modes[mode] += \
                CXData.shift(conj(self.object[id1[i] - self.p2:id1[i] + self.p2, id2[i] - self.p2:id2[i] + self.p2]) *
                 (psi_new-psi_old)[mode] / object_intensity_max, d1[i]%1, d2[i]%1)

        self.probe.normalise()
        
        self.probe.orthogonalise()
Пример #5
0
    def M(psi, det_mod):
        """.. method:: M(mode, psi_modes, det_mod)

            Applies modulus constraint to psi_modes(mode) for a given position.

            :param list psi_modes: A list of CXData instances containing all modes at a given position.
            :param np.ndarray det_mod: Modulus of measured diffraction pattern.

        """
        if isinstance(psi, CXData):
            return ifft2(det_mod * exp(complex(0., 1.) * angle(fft2(psi))))
        elif isinstance(psi, CXModal):
            mode_sum = CXModal.modal_sum(abs(fft2(psi))**2.0)**0.5
            return ifft2((fft2(psi)/(mode_sum))*det_mod)
Пример #6
0
    def update_object(self, i, psi_old, psi_new):
        """
        Update the object from a single ptycho position.

        """
        then=time.time()
        d1, d2 = self.positions.data
        id1, id2 = d1//1, d2//1
        probe_intensity_max = CXModal.modal_sum(abs(self.probe)**2.0).data[0].max()
        
        self.object[id1[i] - self.p2:id1[i] + self.p2, id2[i] - self.p2:id2[i] + self.p2] += \
            CXData.shift(CXModal.modal_sum(conj(self.probe) * (psi_new-psi_old)) / probe_intensity_max, 
                d1[i]%1, d2[i]%1)

        if self.total_its==0 and sp.mod(i, len(self.positions.data[0]) / 10) == 0:
            self.update_figure(i)
Пример #7
0
    def init_probe(self, *args, **kwargs):

        if CXP.io.initial_probe_guess is not '':
            probe = CXData()
            probe.load(CXP.io.initial_probe_guess)
            self.probe.modes = [CXData(data=[probe.data[0]/(i+1)]) for i in range(CXP.reconstruction.probe_modes)]
            self.probe.normalise()
        else:
            dx_s = CXP.dx_s

            p, p2 = CXP.preprocessing.desired_array_shape, CXP.preprocessing.desired_array_shape/2

            probe = sp.zeros((p, p), complex)

            if CXP.experiment.optic.lower() == 'kb':
                if len(CXP.experiment.beam_size)==1:
                    bsx=bsy=np.round(CXP.experiment.beam_size[0]/dx_s)
                elif len(CXP.experiment.beam_size)==2:
                    bsx, bsy = np.round(CXP.experiment.beam_size[0]/dx_s), np.round(CXP.experiment.beam_size[1]/dx_s)

                probe = np.sinc((np.arange(p)-p2)/bsx)[:,np.newaxis]*np.sinc((np.arange(p)-p2)/bsy)[np.newaxis,:]
                

            elif CXP.experiment.optic.lower() == 'zp':
                probe = np.sinc(sp.hypot(*sp.ogrid[-p2:p2, -p2:p2])/np.round(3.*CXP.experiment.beam_size[0]/(2*CXP.dx_s)))

            ph_func = gauss_smooth(np.random.random(probe.shape), 10)
            fwhm = p/2.0
            radker = sp.hypot(*sp.ogrid[-p/2:p/2,-p/2:p/2])
            gaussian = exp(-1.0*(fwhm/2.35)**-2. * radker**2.0 )
            gaussian /= gaussian.max()
            probe = abs(gaussian*probe)* exp(complex(0.,np.pi)*ph_func/ph_func.max())

            self.probe.modes = [CXData(data=[probe/(i+1)]) for i in range(CXP.reconstruction.probe_modes)]
            
            self.probe.normalise()
Пример #8
0
    def ptycho_mesh(self):
        """
        Generate a list of ptycho scan positions.

        Outputs
        -------
        self.data : list of 2xN arrays containing horizontal and vertical scan positions in pixels
        self.initial : initial guess at ptycho scan positions (before position correction)
        self.initial_skew : initial skew
        self.initial_rot : initial rotation
        self.initial_scl : initial scaling
        self.skew : current best guess at skew
        self.rot : current best guess at rotation
        self.scl : current best guess at scaling
        self.total : total number of ptycho positions

        [optional]
        self.correct : for simulated data this contains the correct position

        """
        CXP.log.info('Getting ptycho position mesh.')
        
        if CXP.measurement.ptycho_scan_mesh == 'generate':
            if CXP.measurement.ptycho_scan_type == 'cartesian':
                x2 = 0.5*(CXP.measurement.cartesian_scan_dims[0]-1)
                y2 = 0.5*(CXP.measurement.cartesian_scan_dims[1]-1)
                tmp = map(lambda a: CXP.measurement.cartesian_step_size*a, np.mgrid[-x2:x2+1, -y2:y2+1])
                self.positions.data = [tmp[0].flatten(), tmp[1].flatten()]
                if CXP.reconstruction.flip_mesh_lr:
                    self.log.info('Flip ptycho mesh left-right')
                    self.positions.data[0] = self.data[0][::-1]
                if CXP.reconstruction.flip_mesh_ud:
                    self.log.info('Flip ptycho mesh up-down')
                    self.positions.data[1] = self.data[1][::-1]
                if CXP.reconstruction.flip_fast_axis:
                    self.log.info('Flip ptycho mesh fast axis')
                    tmp0, tmp1 = self.data[0], self.data[1]
                    self.positions.data[0], self.positions.data[1] = tmp1, tmp0
            if CXP.measurement.ptycho_scan_type == 'round_roi':
                self.positions.data = list(round_roi(CXP.measurement.round_roi_diameter, CXP.measurement.round_roi_step_size))
            if CXP.measurement.ptycho_scan_type == 'list':
                l = np.genfromtxt(CXP.measurement.list_scan_filename)
                x_pos, y_pos = [], []
                for element in l:
                    x_pos.append(element[0])
                    y_pos.append(element[1])
                self.positions.data = [sp.array(x_pos), sp.array(y_pos)]


        elif CXP.measurement.ptycho_scan_mesh == 'supplied':
            l = np.genfromtxt(CXP.measurement.list_scan_filename)
            x_pos, y_pos = [], []
            for element in l:
                x_pos.append(element[0])
                y_pos.append(element[1])
            self.positions.data = [sp.array(x_pos), sp.array(y_pos)]

        for element in self.positions.data:
            element /= CXP.dx_s
            element += CXP.ob_p/2
        self.positions.total = len(self.positions.data[0])

        self.positions.correct = [sp.zeros((self.positions.total))]*2
        jit_pix = CXP.reconstruction.initial_position_jitter_radius
        search_pix = CXP.reconstruction.ppc_search_radius

        self.positions.data[0] += jit_pix * uniform(-1, 1, self.positions.total)
        self.positions.data[1] += jit_pix * uniform(-1, 1, self.positions.total)

        if CXP.reconstruction.probe_position_correction:
            self.positions.correct[0] = self.positions.data[0]+0.25*search_pix * uniform(-1, 1, self.positions.total)
            self.positions.correct[1] = self.positions.data[1]+0.25*search_pix * uniform(-1, 1, self.positions.total)
        else:
            self.positions.correct = [self.positions.data[0].copy(), self.positions.data[1].copy()]

        data_copy = CXData(data=list(self.positions.data))
        if not CXP.reconstruction.ptycho_subpixel_shift:
            self.positions.data = [np.round(self.positions.data[0]), np.round(self.positions.data[1])]
            self.positions.correct = [np.round(self.positions.correct[0]), np.round(self.positions.correct[1])]
        CXP.rms_rounding_error = [None]*2

        for i in range(2):
            CXP.rms_rounding_error[i] = sp.sqrt(sp.sum(abs(abs(data_copy.data[i])**2.-abs(self.positions.data[i])**2.)))

        CXP.log.info('RMS Rounding Error (Per Position, X, Y):\t {:2.2f}, {:2.2f}'.format(CXP.rms_rounding_error[0]/len(self.positions.data[0]),
                                                                                           CXP.rms_rounding_error[1]/len(self.positions.data[1])))
Пример #9
0
    def update_figure(self, i=0):
        cur_cmap = cm.RdGy_r
        self.f1.clf()
        self.init_figure()

        wh = sp.where(abs(self.object.data[0]) > 0.1 * (abs(self.object.data[0]).max()))
        try:
            x1, x2 = min(wh[0]), max(wh[0])
            y1, y2 = min(wh[1]), max(wh[1])
        except (ValueError, IndexError):
            x1, x2 = 0, self.ob_p
            y1, y2 = 0, self.ob_p

        # Plot magnitude of object
        s1 = pylab.subplot(231)
        s1_im = s1.imshow(abs(self.object).data[0][x1:x2, y1:y2], cmap=cm.Greys_r)
        s1.set_title('|object|')
        plt.axis('off')
        pylab.colorbar(s1_im)

        # Plot phase of object
        s2 = pylab.subplot(232)
        s2_im = s2.imshow(sp.angle(self.object.data[0][x1:x2, y1:y2]), cmap=cm.hsv)
        s2.set_title('phase(object)')
        plt.axis('off')
        pylab.colorbar(s2_im)

        # Complex HSV plot of object
        s3 = pylab.subplot(233)
        h = ((angle(self.object).data[0][x1:x2, y1:y2] + np.pi) / (2*np.pi)) % 1.0
        s = np.ones_like(h)
        l = abs(self.object).data[0][x1:x2, y1:y2]
        l-=l.min()
        l/=l.max()
        s3_im = s3.imshow(np.dstack(v_hls_to_rgb(h,l,s)))
        s3.set_title('Complex plot of Object')
        plt.axis('off')

        # Plot probe mode 0
        s4 = pylab.subplot(234)
        s4_im = s4.imshow(abs(self.probe.modes[0].data[0]), cmap=cur_cmap)
        s4.set_title('|probe0|')
        plt.axis('off')
        pylab.colorbar(s4_im)

        if CXP.reconstruction.probe_modes>1:
            s5 = pylab.subplot(235)
            s5_im = s5.imshow(abs(self.probe.modes[1].data[0]), cmap=cur_cmap)
            s5.set_title('|probe1|')
            plt.axis('off')
            pylab.colorbar(s5_im)
        else:
            pass
        if self.ppc:
            s6 = self.f1.add_subplot(236)
            s6_im = s6.scatter(self.positions.data[0], self.positions.data[1], s=10,
                c='b', marker='o', alpha=0.5, edgecolors='none', label='current')
            patches = []
            for m in range(self.positions.total):
                patches.append(Circle((self.positions.initial[0][m], self.positions.initial[1][m]),
                               radius=CXP.reconstruction.ppc_search_radius))
            collection = PatchCollection(patches, color='tomato', alpha=0.2, edgecolors=None)
            s4.add_collection(collection)
            if CXP.measurement.simulate_data:
                s4_im = s4.scatter(self.positions.correct[0], self.positions.correct[1], s=10,
                    c='g', marker='o', alpha=0.5, edgecolors='none', label='correct')
                CXP.log.info('RMS position deviation from correct: [x:{:3.2f},y:{:3.2f}] pixels'.format(
                            sp.sqrt(sp.mean((self.positions.data[0] - self.positions.correct[0])**2.)),
                            sp.sqrt(sp.mean((self.positions.data[1] - self.positions.correct[1])**2.))))
                lines=[]
                for m in range(self.positions.total):
                    lines.append(((self.positions.correct[0][m], self.positions.correct[1][m]),
                                  (self.positions.data[0][m], self.positions.data[1][m])))
                for element in lines:
                    x, y = zip(*element)
                    s4.plot(x, y, 'g-')
            else:
                lines = []
                for m in range(self.positions.total):
                    lines.append(((self.positions.initial[0][m], self.positions.initial[1][m]),
                                  (self.positions.data[0][m], self.positions.data[1][m])))
                for element in lines:
                    x, y = zip(*element)
                    s6.plot(x, y, 'g-')
                CXP.log.info('RMS position deviation from initial: [x:{:3.2f},y:{:3.2f}] pixels'.format(
                            sp.sqrt(sp.mean((self.positions.data[0] - self.positions.initial[0])**2.)),
                            sp.sqrt(sp.mean((self.positions.data[1] - self.positions.initial[1])**2.))))
            s6.legend(prop={'size': 6})
            s6.set_title('Position Correction')
            s6.set_aspect('equal')
            extent = s6.get_window_extent().transformed(self.f1.dpi_scale_trans.inverted())
            pylab.savefig(self._cur_sequence_dir + '/ppc_{:d}.png'.format(self.total_its), bbox_inches=extent.expanded(1.2, 1.2), dpi=100)
            s6.set_aspect('auto')
        else:
            s6 = pylab.subplot(236)
            if CXP.measurement.simulate_data:
                s6_im = s6.imshow(abs(self.input_probe[1].data[0]), cmap = cur_cmap)
                s6.set_title('|input_probe1|')
            else:
                s6_im = s6.imshow(nlog(fftshift(self.det_mod[np.mod(i,self.positions.total)])).data[0], cmap=cur_cmap)
                s6.set_title('Diff Patt: {:d}'.format(i))
            plt.axis('off')
            pylab.colorbar(s6_im)
        pylab.draw()
        pylab.savefig(self._cur_sequence_dir + '/recon_{:d}.png'.format(self.total_its), dpi=60)
Пример #10
0
    def simulate_data(self):
        CXP.log.info('Simulating diffraction patterns.')
        self.sample = CXData()
        self.sample.load(CXP.io.simulation_sample_filename[0])
        self.sample.data[0] = self.sample.data[0].astype(float)
        self.sample.normalise(val=0.8)
        self.sample.data[0]+=0.2
        self.input_probe = CXModal()
        if len(CXP.io.simulation_sample_filename)>1:
            ph = CXData()
            ph.load(CXP.io.simulation_sample_filename[1])
            ph.data[0] = ph.data[0].astype(float)
            ph.normalise(val=np.pi/3)
            self.sample.data[0] = self.sample.data[0]*exp(complex(0., 1.)*ph.data[0])
        p = self.sample.data[0].shape[0]
        ham_window = sp.hamming(p)[:,np.newaxis]*sp.hamming(p)[np.newaxis,:]
        sample_large = CXData(data=sp.zeros((CXP.ob_p, CXP.ob_p), complex))
        sample_large.data[0][CXP.ob_p/2-p/2:CXP.ob_p/2+p/2, CXP.ob_p/2-p/2:CXP.ob_p/2+p/2] = self.sample.data[0]*ham_window

        ker = sp.arange(0, p)
        fwhm = p/3.0
        radker = sp.hypot(*sp.ogrid[-p/2:p/2,-p/2:p/2])
        gaussian = exp(-1.0*(fwhm/2.35)**-2. * radker**2.0 )
        ortho_modes = lambda n1, n2 : gaussian*np.sin(n1*math.pi*ker/p)[:,np.newaxis]*np.sin(n2*math.pi*ker/p)[np.newaxis, :]
        mode_generator = lambda : sp.floor(4*sp.random.random(2))+1

        used_modes = []
        self.input_psi = CXModal()
        
        for mode in range(CXP.reconstruction.probe_modes):
            if mode==0:
                new_mode = [1,1]
            else:
                new_mode = list(mode_generator())
                while new_mode in used_modes:
                    new_mode = list(mode_generator())
            used_modes.append(new_mode)
            CXP.log.info('Simulating mode {:d}: [{:d}, {:d}]'.format(mode, int(new_mode[0]), int(new_mode[1])))
            ph_func = gauss_smooth(np.random.random((p,p)), 10)
            self.input_probe.modes.append(CXData(name='probe{:d}'.format(mode), 
                data=ortho_modes(new_mode[0], new_mode[1])*exp(complex(0.,np.pi)*ph_func/ph_func.max())))
        
        self.input_probe.normalise()
        self.input_probe.orthogonalise()

        for mode in range(CXP.reconstruction.probe_modes):
            p2 = p/2
            x, y = self.positions.correct
            self.input_psi.modes.append(CXData(name='input_psi_mode{:d}'.format(mode), data=[]))
            
            for i in xrange(len(x)):
                if i%(len(x)/10)==0.:
                    CXP.log.info('Simulating diff patt {:d}'.format(i))
                tmp = (CXData.shift(sample_large, -1.0*(x[i]-CXP.ob_p/2), -1.0*(y[i]-CXP.ob_p/2))
                        [CXP.ob_p/2-p2:CXP.ob_p/2+p2, CXP.ob_p/2-p2:CXP.ob_p/2+p2]*
                        self.input_probe[mode][0])
                self.input_psi[mode].data.append(tmp.data[0])

        # Add modes incoherently
        self.det_mod = CXModal.modal_sum(abs(fft2(self.input_psi)))
        self.det_mod.save(path=CXP.io.base_dir+'/'+CXP.io.scan_id+'/raw_data/{:s}.npy'.format('det_mod'))