adding a thin-object approximation function

            sim_t = time()-t

            self.print_verbose('Propagated %d probes in %2.4f s' % (self.probe_positions[phase].shape[0], sim_t))

        if transmit:

            with catch_warnings():

                simplefilter('ignore')

                catch_warn()

                self.probes = self.probes.astype(np.float32) # discard imaginary

        # clean up device variables

        self.pot_dev_ptr.free()

        self.vars.append(self.pot_dev_ptr) # don't de-allocate here to allow for simulations with different probe/exp. params

        mask_d.free()

        propag_d.free()

        psi_k_d.free()

    data = data.reshape((n, n) + data.shape[1:]).transpose((0, 2, 1, 3) + tuple(range(4, data.ndim + 1)))

    data = data.reshape((n * data.shape[1], n * data.shape[3]) + data.shape[4:])

    return data

def thin_object_approx(potential_slices, probes):

    """

    Simulate the exit wavefunction using a thin-object approximation

    """

    psi_pos = probes * np.exp(1.j * potential_slices.sum(0))

    psi_k = np.abs(np.fft.fft2(psi_pos))

    return psi_k

 No newline at end of file