Ejemplo n.º 1
0
        def proc_file(file):
            fobj = bu.hsDat(file, load=True)

            vperp = fobj.dat[:, 0]
            elec3 = fobj.dat[:, 1]

            elec3_filt = signal.filtfilt(b2, a2, elec3)

            if plot_raw_dat:
                fac = bu.fft_norm(nsamp, fsamp)
                plt.plot(time_vec[:10000], elec3[:10000])
                plt.figure()
                plt.plot(time_vec[:10000], vperp[:10000])
                plt.figure()
                plt.loglog(freqs, fac * np.abs(np.fft.rfft(vperp)))
                plt.figure()
                plt.loglog(freqs, fac * np.abs(np.fft.rfft(elec3)))
                plt.loglog(freqs, fac * np.abs(np.fft.rfft(elec3_filt)))
                plt.show()

                input()

            inds = np.abs(freqs - fspin) < 200.0

            elec3_fft = np.fft.rfft(elec3)
            true_fspin = freqs[np.argmax(np.abs(elec3_fft))]

            amp, phase_mod = bu.demod(vperp, true_fspin, fsamp, plot=plot_demod, \
                                  filt=True, bandwidth=bandwidth, \
                                  notch_freqs=notch_freqs, notch_qs=notch_qs, \
                                  tukey=True, tukey_alpha=5.0e-4, \
                                  detrend=detrend, detrend_order=1, harmind=2.0, \
                                  force_2pi_wrap=force_2pi_wrap)

            phase_mod_filt = signal.filtfilt(b3, a3, phase_mod)
            #phase_mod_filt = phase_mod

            amp_asd = np.abs(np.fft.rfft(amp))
            phase_asd = np.abs(np.fft.rfft(phase_mod))
            phase_asd_filt = np.abs(np.fft.rfft(phase_mod_filt))

            # popt_n, pcov_n = opti.curve_fit(gauss, freqs[notch_fit_inds], \
            #                                     phase_asd_filt[notch_fit_inds], \
            #                                     p0=[10000, notch_init, 2, 0])

            if apply_notch:
                notch = freqs[np.argmax(phase_asd_filt * notch_fit_inds)]
                notch_digital = (2.0 / fsamp) * (notch)

                for i in range(notch_nharm):
                    bn, an = signal.iirnotch(notch_digital * (i + 1), notch_q)
                    phase_mod_filt = signal.lfilter(bn, an, phase_mod_filt)

            phase_asd_filt_2 = np.abs(np.fft.rfft(phase_mod_filt))

            if correct_noise_color:
                phase_asd = phase_asd * freqs**noise_color_power
                phase_asd_filt = phase_asd_filt * freqs**noise_color_power
                phase_asd_filt_2 = phase_asd_filt_2 * freqs**noise_color_power

            if plot_phase:
                plt.plot(freqs, phase_mod)
                plt.xlabel('Frequency [Hz]')
                plt.ylabel('Phase [rad]')
                plt.tight_layout()

                plt.figure()
                plt.loglog(freqs, phase_asd)
                plt.loglog(freqs, phase_asd_filt)
                plt.loglog(freqs, phase_asd_filt_2)
                plt.xlabel('Frequency [Hz]')
                plt.ylabel('Phase ASD [arb]')
                plt.tight_layout()

                plt.show()

                input()

            freq_mask = (freqs > allowed_freqs[0]) * (freqs < allowed_freqs[1])

            max_ind = np.argmax(phase_asd_filt_2 * freq_mask)
            max_freq = freqs[max_ind]

            p0 = [10000, max_freq, 5.0, 0]

            try:
                popt, pcov = opti.curve_fit(lorentzian, freqs[max_ind-30:max_ind+30], \
                                            phase_asd_filt_2[max_ind-30:max_ind+30], p0=p0, \
                                            maxfev=10000)

                fit_max = popt[1]
                fit_std = np.abs(popt[2])
            except:
                print('bad fit...')
                fit_max = max_freq
                fit_std = 5.0 * (freqs[1] - freqs[0])
                popt = p0

            if plot_sideband_fit:
                plot_freqs = np.linspace(freqs[max_ind - 30],
                                         freqs[max_ind + 30], 100)
                plt.loglog(freqs, phase_asd_filt_2)
                plt.loglog(plot_freqs, lorentzian(plot_freqs, *popt))
                print(fit_max, fit_std)
                plt.show()

                input()

            # if fit_max < 10:
            #     return

            # if len(wobble_freq):
            #     if (np.abs(fit_max - wobble_freq[-1]) / wobble_freq[-1]) > 0.1:
            #         # plt.loglog(freqs, phase_asd)
            #         # plt.loglog(freqs, phase_asd_filt)
            #         # plt.loglog(freqs, phase_asd_filt_2)
            #         # plt.show()
            #         return

            elec3_filt_fft = np.fft.rfft(elec3_filt)

            fit_ind = 100000
            short_freqs = np.fft.rfftfreq(fit_ind, d=1.0 / fsamp)
            zeros = np.zeros(fit_ind)
            voltage = np.array([zeros, zeros, zeros, elec3_filt[:fit_ind], \
                       zeros, zeros, zeros, zeros])
            efield = bu.trap_efield(voltage * tabor_mon_fac, only_x=True)

            #efield_mag = np.linalg.norm(efield, axis=0)

            efield_asd = bu.fft_norm(fit_ind, fsamp) * np.abs(
                np.fft.rfft(efield[0]))

            # max_ind = np.argmax(np.abs(elec3_filt_fft))
            short_max_ind = np.argmax(efield_asd)
            # freq_guess = freqs[max_ind]
            # phase_guess = np.mean(np.angle(elec3_filt_fft[max_ind-2:max_ind+2]))
            # amp_guess = np.sqrt(2) * np.std(efield[0])
            # p0 = [amp_guess, freq_guess, phase_guess, 0]

            # popt_l, pcov_l = opti.curve_fit(lorentzian, short_freqs[short_max_ind-100:short_max_ind+100], \
            #                                 efield_asd[short_max_ind-100:short_max_ind+100], \
            #                                 p0=[amp_guess, freq_guess, 100, 0], maxfev=10000)

            # start_sine = time.time()
            # popt, pcov = opti.curve_fit(sine, time_vec[:fit_ind], efield[0], \
            #                                 sigma=0.01*efield[0], p0=p0)

            # print popt[0], popt_l[0] * (fsamp / fit_ind)

            #print popt[0], np.sqrt(pcov[0,0])
            # amp_fit = efield_asd[short_max_ind] * np.sqrt(2.0 * fsamp / fit_ind)
            amp_fit = np.sqrt(2) * np.std(efield[0])
            # plt.plot(efield[0])
            # plt.show()

            err_val = np.mean(np.array([efield_asd[short_max_ind-10:short_max_ind], \
                                efield_asd[short_max_ind+1:short_max_ind+11]]).flatten())
            amp_err = np.sqrt(
                (err_val * fsamp / fit_ind)**2)  # + (0.01*amp_fit)**2)
            # print amp_fit, amp_err
            # stop_sine = time.time()
            # print "Field sampling: ", stop_sine - start_sine

            return [2.0 * amp_fit, np.sqrt(2) * amp_err, fit_max, fit_std]
rhobead = bu.rhobead['val']
mbead_dic = {'val': 84.3e-15, 'sterr': 1.0e-15, 'syserr': 1.5e-15}
mbead = mbead_dic['val']
Ibead = bu.get_Ibead(mbead=mbead_dic)['val']
kappa = bu.get_kappa(mbead=mbead_dic)['val']

### Environmental constants
T = 297
P = 3.5e-6 * 100  # Pressure, converted to pascals
m0 = 18.0 * constants.atomic_mass  # residual gas particl mass, in kg
beta_rot = P * np.sqrt(m0) / kappa

### Intial electric field, and initial conditions
#drive_freq = 50000.0
drive_freq = 110000.5
drive_amp = np.abs(bu.trap_efield([0, 0, 0, 400, -400, 0, 0, 0], nsamp=1)[0])
#N_opt = 2.0 * np.pi * (6000.0) * beta_rot
N_opt = 0
xi_init = np.array([p0, 0.0, 0.0, 0.0, 0.0, 2.0 * np.pi * drive_freq])
# to avoid numerical errors, dipole moment is integrated in units of e * micron
# although anytime a torque appears, the value is cast to the correct SI units
anomaly = 1.0  # Set to 0 if you want to turn of anamolous torques from residual field
real_efield_params = pickle.load(open('./real_drive_dat.p', 'rb'))
efield_rms = real_efield_params[0.0][0]
del real_efield_params[0.0]
del real_efield_params[110000.5]

### Simulation parameters
t_release = 20.0
t_sim = 120.0
#t_sim = 4.0
    def proc_file(filename):
        try:
            # Load data, computer FFTs and plot if requested
            obj = hsDat(filename)
            dat_fft = np.fft.rfft(obj.dat[:, data_ax])

            elec_mon = obj.dat[:, drive_ax]
            drive_fft = np.fft.rfft(elec_mon)

            elec_filt = tabor_mon_fac * signal.filtfilt(b2, a2, elec_mon)

            zeros = np.zeros(nsamp)
            voltage = np.array([zeros, zeros, zeros, elec_filt, \
                       -elec_filt, zeros, zeros, zeros])
            efield = bu.trap_efield(voltage)

            # plt.loglog(freqs, np.abs(np.fft.rfft(elec_mon)))
            # plt.loglog(freqs, np.abs(np.fft.rfft(elec_filt)))
            # # for i in [0,1,2]:
            # #     plt.plot(voltage[3][:50000]-voltage[4][:50000])
            # plt.show()

            max_ind = np.argmax(np.abs(drive_fft))
            freq_guess = freqs[max_ind]
            phase_guess = np.mean(np.angle(drive_fft[max_ind - 2:max_ind + 2]))
            amp_guess = np.sqrt(2) * np.std(efield[0])
            p0 = [amp_guess, freq_guess, phase_guess, 0]

            fit_ind = int(0.01 * len(time_vec))
            popt, pcov = opti.curve_fit(sine,
                                        time_vec[:fit_ind],
                                        efield[0][:fit_ind],
                                        p0=p0)
            amp_fit = popt[0]
            amp_err = np.sqrt(pcov[0, 0])

            if plot_dat:
                plt.figure()
                plt.loglog(freqs, np.abs(dat_fft))
                plt.loglog(freqs, np.abs(drive_fft))

            # Filter data outside the window of interest and plot
            # if requested
            dat_fft[finds2] = 0.
            drive_fft[finds] = 0.
            if plot_dat:
                plt.loglog(freqs, np.abs(dat_fft))
                plt.loglog(freqs, np.abs(drive_fft))
                plt.show()

            pressures = obj.attribs['pressures']

            out_list = [amp_fit, amp_err, popt[1], np.sqrt(pcov[1,1]), \
                        np.angle(np.sum(dat_fft)), np.std(np.angle(dat_fft)), \
                        np.angle(np.sum(drive_fft)), np.std(np.angle(drive_fft)), \
                        pressures[0], pressures[1], pressures[2], obj.attribs['time']]

            return out_list

            # field_amps[i] = amp_fit
            # field_amp_errs[i] = amp_err

            # field_freqs[i] = popt[1]
            # field_freq_errs[i] = np.sqrt(pcov[1,1])

            # # Compute the raw phases of drive and response
            # #    drive: phase of frot signal
            # #    response: phase of 2*frot signal
            # phases[i] = np.angle(np.sum(dat_fft))
            # phase_errs[i] = np.std(np.angle(dat_fft))
            # dphases[i] = np.angle(np.sum(drive_fft))
            # dphase_errs[i] = np.std(np.angle(drive_fft))
            # # Convert raw pressure in torr to mbar
            # pressures[i, :] = obj.attribs["pressures"]
            # pressures_mbar[i, :] = np.array(obj.attribs["pressures"]) * 1.333
            # times[i] = obj.attribs['time']
        except:
            print("bad file")
            return
def build_uncalibrated_H(fobjs, average_first=True, dpsd_thresh = 8e-1, mfreq = 1., \
                         skip_qpd=False, plot_response=False, drop_bad_bins=True, \
                         new_trap=False, lines_to_remove=[60.0], zero_drive_phase=False):
    '''Generates a transfer function from a list of DataFile objects
           INPUTS: fobjs, list of file objects
                   average_first, boolean specifying whether to average responses
                                  for a given drive before computing H
                   dpsd_thresh, threshold above which to compute H
                   mfreq, minimum frequency to consider
                   fix_HF, boolean to specify whether to try fixing spectral
                           leakage at high frequency due to drift in a timebase

           OUTPUTS: Hout, dictionary with 3x3 complex valued matrices as values
                          and frequencies as keys'''

    print("BUILDING H...")
    sys.stdout.flush()

    Hout = {}
    Hout_noise = {}
    Hout_amp = {}
    Hout_phase = {}
    Hout_fb = {}

    Hout_counts = {}

    avg_drive_fft = {}
    avg_data_fft = {}
    avg_fb_fft = {}

    if not skip_qpd:
        avg_side_fft = {}

        avg_amp_fft = {}
        avg_phase_fft = {}

    counts = {}

    if plot_response:
        fbfig, fb_axarr = plt.subplots(3,3,sharex=True,sharey=True,figsize=(9,8))
        posfig, pos_axarr = plt.subplots(3,3,sharex=True,sharey='row',figsize=(9,8))
        drivefig, drive_axarr = plt.subplots(3,3,sharex=True,sharey=True,figsize=(9,8))
        if not skip_qpd:
            ampfig, amp_axarr = plt.subplots(5,3,sharex=True,sharey=True,figsize=(9,10))
            phasefig, phase_axarr = plt.subplots(5,3,sharex=True,sharey=True,figsize=(9,10))
            sidefig, side_axarr = plt.subplots(4,3,sharex=True,sharey=True,figsize=(9,9))

    filind = 0
    for fobj in fobjs:

        N = np.shape(fobj.pos_data)[1]#number of samples
        fsamp = fobj.fsamp
        
        fft_fac = bu.fft_norm(N, fsamp)

        drive = bu.trap_efield(fobj.electrode_data) #* constants.elementary_charge
        #drive = np.roll(drive, -10, axis=-1)

        # dfft = np.fft.rfft(fobj.electrode_data) #fft of electrode drive in daxis.
        dfft = np.fft.rfft( drive ) * fft_fac
        if new_trap: 
            data_fft = np.fft.rfft(fobj.pos_data_3) * fft_fac
        else:
            data_fft = np.fft.rfft(fobj.pos_data) * fft_fac

        fb_fft = np.fft.rfft(fobj.pos_fb) * fft_fac

        if not skip_qpd:
            amp_fft = np.fft.rfft(fobj.amp) * fft_fac
            phase_fft = np.fft.rfft(fobj.phase) * fft_fac

            left = fobj.amp[2] + fobj.amp[3]
            right = fobj.amp[0] + fobj.amp[1]
            top = fobj.amp[2] + fobj.amp[0]
            bot = fobj.amp[3] + fobj.amp[1]
            side_fft = np.fft.rfft(np.array([right, left, top, bot])) * fft_fac

        fft_freqs = np.fft.rfftfreq(N, d=1.0/fsamp)

        # for i in range(len(dfft)):
        #     plt.loglog(fft_freqs, np.abs(dfft[i]))
        # # plt.loglog(fft_freqs, np.abs(dfft[4]))
        # # plt.loglog(fft_freqs, np.abs(data_fft[0]))
        # plt.show()

        dpsd = np.abs(dfft)**2  #psd for all electrode drives
        dpsd_thresh = 0.1 * np.max(dpsd.flatten())
        inds = np.where(dpsd>dpsd_thresh)#Where the dpsd is over the threshold for being used.
        eind = np.unique(inds[0])[0]
        # print(eind)

        if eind not in avg_drive_fft:
            avg_drive_fft[eind] = np.zeros(dfft.shape, dtype=np.complex128)
            avg_data_fft[eind] = np.zeros(data_fft.shape, dtype=np.complex128)
            avg_fb_fft[eind] = np.zeros(fb_fft.shape, dtype=np.complex128)
            if not skip_qpd:
                avg_amp_fft[eind] = np.zeros(amp_fft.shape, dtype=np.complex128)
                avg_phase_fft[eind] = np.zeros(phase_fft.shape, dtype=np.complex128)
                avg_side_fft[eind] = np.zeros(side_fft.shape, dtype=np.complex128)
            counts[eind] = 0.

        # try:
        avg_drive_fft[eind] += dfft
        avg_data_fft[eind] += data_fft
        avg_fb_fft[eind] += fb_fft
        if not skip_qpd:
            avg_amp_fft[eind] += amp_fft
            avg_phase_fft[eind] += phase_fft
            avg_side_fft[eind] += side_fft

        # except:
        #     traceback.print_exc()
        #     print()
        #     print(fobj.fname)
        #     derpfig, derpax = plt.subplots(1,1)
        #     for i in [0,1,2,3,4]:
        #         derpax.loglog(np.abs(amp_fft[i]) * 10**i, label=str(i))
        #     derpax.legend()
        #     plt.show()
        #     input()

        counts[eind] += 1.

    for eind in list(counts.keys()):
        print(eind, counts[eind])
        avg_drive_fft[eind] = avg_drive_fft[eind] / counts[eind]
        avg_data_fft[eind] = avg_data_fft[eind] / counts[eind]
        avg_fb_fft[eind] = avg_fb_fft[eind] / counts[eind]
        if not skip_qpd:
            avg_amp_fft[eind] = avg_amp_fft[eind] / counts[eind]
            avg_phase_fft[eind] = avg_phase_fft[eind] / counts[eind]
            avg_side_fft[eind] = avg_side_fft[eind] / counts[eind]

    poslabs = {0: 'X', 1: 'Y', 2: 'Z'}
    sidelabs = {0: 'Right', 1: 'Left', 2: 'Top', 3: 'Bottom'}
    quadlabs = {0: 'Top Right', 1: 'Bottom Right', 2: 'Top Left', \
                3: 'Bottom Left', 4: 'Backscatter'}

    for eind in list(avg_drive_fft.keys()):
        # First find drive-frequency bins above a fixed threshold
        dpsd = np.abs(avg_drive_fft[eind])**2 
        inds = np.where(dpsd > dpsd_thresh)

        # Extract the frequency indices
        finds = inds[1]

        # Ignore DC and super low frequencies
        mfreq = 1.0
        b = finds > np.argmin(np.abs(fft_freqs - mfreq))

        tf_inds = finds[b]
        for linefreq in lines_to_remove:
            print('Ignoring response at line frequency: {:0.1f}'.format(linefreq))
            line_freq_ind = np.argmin(np.abs(fft_freqs - linefreq))
            tf_inds = tf_inds[tf_inds != line_freq_ind]

        freqs = fft_freqs[tf_inds]

        xlim = (np.min(fft_freqs[1:]), np.max(fft_freqs))
        # xlim = (45.0, 130.0)

        # plt.plot(freqs, np.angle(avg_drive_fft[eind][eind,tf_inds]) / np.pi)
        # plt.title('Raw Phase From Drive FFT')
        # plt.ylabel('Apparent Phase [$\\pi \\, rad$]')
        # plt.xlabel('Drive Frequency [Hz]')
        # plt.tight_layout()

        # plt.figure()
        # plt.plot(freqs, np.unwrap(np.angle(avg_drive_fft[eind][eind,tf_inds])) / np.pi)
        # plt.title('Unwrapped Phase From Drive FFT')
        # plt.ylabel('Apparent Phase [$\\pi \\, rad$]')
        # plt.xlabel('Drive Frequency [Hz]')
        # plt.tight_layout()
        # plt.show()

        # outind = config.elec_map[eind]
        outind = eind

        if plot_response:
            
            for elec in [0,1,2]: #,3,4,5,6,7]:
                drive_axarr[elec,outind].loglog(fft_freqs, \
                                                np.abs(avg_drive_fft[eind][elec]), alpha=1.0)
                drive_axarr[elec,outind].loglog(fft_freqs[tf_inds], \
                                                np.abs(avg_drive_fft[eind][elec])[tf_inds], alpha=1.0)
                if outind == 0:
                    drive_axarr[elec,outind].set_ylabel('Efield axis ' + str(elec) \
                                                            + '\n[(V/m)/$\\sqrt{\\rm Hz}$]')
                if elec == 2: #7:
                    drive_axarr[elec,outind].set_xlabel('Frequency [Hz]')


            for resp in [0,1,2,3,4]:

                if not skip_qpd:
                    amp_axarr[resp,outind].loglog(fft_freqs[tf_inds], \
                                                  np.abs(avg_amp_fft[eind][resp])[tf_inds], alpha=1.0)
                    phase_axarr[resp,outind].loglog(fft_freqs[tf_inds], \
                                                    np.abs(avg_phase_fft[eind][resp])[tf_inds], alpha=1.0)
                    if outind == 0:
                        amp_axarr[resp,outind].set_ylabel(quadlabs[resp])
                        phase_axarr[resp,outind].set_ylabel(quadlabs[resp])
                    if resp == 4:
                        amp_axarr[resp,outind].set_xlabel('Frequency [Hz]')
                        phase_axarr[resp,outind].set_xlabel('Frequency [Hz]')

                    if resp in [0,1,2,3]:
                        side_axarr[resp,outind].loglog(fft_freqs[tf_inds], \
                                                       np.abs(avg_side_fft[eind][resp])[tf_inds], alpha=1.0)
                        if outind == 0:
                            side_axarr[resp,outind].set_ylabel(sidelabs[resp])
                        if resp == 3:
                            side_axarr[resp,outind].set_xlabel('Frequency [Hz]')


                if resp in [0,1,2]:
                    # if resp == 2:
                    #     fac = 1000.0
                    # else:
                    #     fac = 1.0
                    pos_axarr[resp,outind].loglog(fft_freqs, np.abs(avg_data_fft[eind][resp]), alpha=1.0)
                    pos_axarr[resp,outind].loglog(fft_freqs[tf_inds], \
                                                  np.abs(avg_data_fft[eind][resp])[tf_inds], alpha=1.0)

                    fb_axarr[resp,outind].loglog(fft_freqs, np.abs(avg_fb_fft[eind][resp]), alpha=1.0)
                    fb_axarr[resp,outind].loglog(fft_freqs[tf_inds], \
                                                  np.abs(avg_fb_fft[eind][resp])[tf_inds], alpha=1.0)
                    if outind == 0:
                        pos_axarr[resp,outind].set_ylabel(poslabs[resp] + ' [arb]')
                        fb_axarr[resp,outind].set_ylabel(poslabs[resp] + ' FB\n[bits/$\\sqrt{\\rm Hz}$]')
                    if resp == 2:
                        pos_axarr[resp,outind].set_xlabel('Frequency [Hz]')
                        fb_axarr[resp,outind].set_xlabel('Frequency [Hz]')
            
            drivefig.suptitle('Drive Amplitude vs. Frequency', fontsize=16)
            posfig.suptitle('ASD of XYZ vs. Frequency', fontsize=16)
            fbfig.suptitle('ASD of XYZ Feedback vs. Frequency', fontsize=16)
            if not skip_qpd:
                ampfig.suptitle('ASD of Demod. Carrier Amp vs. Frequency', fontsize=16)
                phasefig.suptitle('ASD of Demod. Carrier Phase vs. Frequency', fontsize=16)
                sidefig.suptitle('ASD of Sum of Neighboring QPD Carrier Amplitudes', fontsize=16)

                figlist = [posfig, fbfig, drivefig, ampfig, phasefig, sidefig]
                axlist = [pos_axarr, fb_axarr, drive_axarr, amp_axarr, phase_axarr, side_axarr]

            else:
                figlist = [posfig, fbfig, drivefig]
                axlist = [pos_axarr, fb_axarr, drive_axarr]


            for axind, axarr in enumerate(axlist):
                # axarr[0,0].set_xlim(*xlim)
                axarr[0,0].set_xlim(34, 107)

                for drive in [0,1,2]:
                    axarr[0,drive].set_title(poslabs[drive] + ' Drive')

                for resp in [0,1,2]:
                    if (axind != 0) and (resp != 0):
                        continue

                    mag_major_locator = LogLocator(base=10.0, numticks=30)
                    mag_minor_locator = LogLocator(base=1.0, numticks=300)

                    axarr[resp,0].yaxis.set_major_locator(mag_major_locator)
                    axarr[resp,0].yaxis.set_minor_locator(mag_minor_locator)
                    axarr[resp,0].yaxis.set_minor_formatter(NullFormatter())

                for d in [0,1,2]:
                    for r in [0,1,2]:
                        axarr[r,d].grid(which='both')



        ### Compute FFT of each response divided by FFT of each drive.
        ### This is way more information than we need for a single drive freq
        ### and electrode pair, but it allows a nice vectorization
        Hmat = np.einsum('ij, kj -> ikj', \
                             avg_data_fft[eind][:,tf_inds], 1. / avg_drive_fft[eind][:,tf_inds])

        Hmat_fb = np.einsum('ij, kj -> ikj', \
                            avg_fb_fft[eind][:,tf_inds], 1. / avg_drive_fft[eind][:,tf_inds])

        if not skip_qpd:
            Hmat_amp = np.einsum('ij, kj -> ikj', \
                                 avg_amp_fft[eind][:,tf_inds], 1. / avg_drive_fft[eind][:,tf_inds])
            Hmat_phase = np.einsum('ij, kj -> ikj', \
                                   avg_phase_fft[eind][:,tf_inds], 1. / avg_drive_fft[eind][:,tf_inds])


        # Generate an integer by which to roll the data_fft to compute the noise
        # limit of the TF measurement
        if len(tf_inds) > 1:
            shift = int(0.5 * (tf_inds[1] - tf_inds[0]))
        else:
            shift = int(0.5 * tf_inds[0])
        randadd = np.random.choice(np.arange(-int(0.1*shift), \
                                             int(0.1*shift)+1, 1))
        shift = shift + randadd
        rolled_data_fft = np.roll(avg_data_fft[eind], shift, axis=-1)

        # Compute the Noise TF
        Hmat_noise = np.einsum('ij, kj -> ikj', \
                                 rolled_data_fft[:,tf_inds], 1. / avg_drive_fft[eind][:,tf_inds])

        # Map the 3x7xNfreq arrays to dictionaries with keys given by the drive
        # frequencies and values given by 3x3 complex-values TF matrices
        #outind = config.elec_map[eind]
        outind = eind
        for i, freq in enumerate(freqs):
            if freq not in Hout:
                if i != 0 and drop_bad_bins:
                    sep = freq - freqs[i-1]
                    # Clause to ignore this particular frequency response if an
                    # above threshold response is found not on a drive bin. Sometimes
                    # random noise components pop up or some power leaks to a 
                    # neighboring bin
                    if sep < 0.8 * (freqs[1] - freqs[0]):
                        continue
                Hout[freq] = np.zeros((3,3), dtype=np.complex128)
                Hout_noise[freq] = np.zeros((3,3), dtype=np.complex128)
                Hout_amp[freq] = np.zeros((5,3), dtype=np.complex128)
                Hout_phase[freq] = np.zeros((5,3), dtype=np.complex128)

            # Add the response from this drive freq/electrode pair to the TF matrix
            Hout[freq][:,outind] += Hmat[:,eind,i]
            Hout_noise[freq][:,outind] += Hmat_noise[:,eind,i]
            if not skip_qpd:
                Hout_amp[freq][:,outind] += Hmat_amp[:,eind,i]
                Hout_phase[freq][:,outind] += Hmat_phase[:,eind,i]

    if plot_response:
        for fig in figlist:
            fig.tight_layout()
            fig.subplots_adjust(top=0.90)
        plt.show()

    # first_mats = []
    freqs = list(Hout.keys())
    freqs.sort()


    if zero_drive_phase:
        # init_phases = np.angle(Hout[freqs[0]])

        first_mats = []
        for freq in freqs[1:3]:
            first_mats.append(Hout[freq])
        first_mats = np.array(first_mats)
        init_phases = np.mean(np.unwrap(np.angle(first_mats), axis=0), axis=0)

        # print(init_phases)
        # input()

        for drive in [0,1,2]:
            if np.abs(init_phases[drive,drive]) > 1.5:
                ### Check the second frequency to make sure the first isn't crazy
                if np.abs(np.angle(Hout[freqs[1]][drive,drive])) > 1.5:
                    print("Correcting phase shift for drive channel", drive)
                    sys.stdout.flush()
                    for freq in freqs:
                        Hout[freq][drive,:] = Hout[freq][drive,:] * (-1)
                        # Hout[freq][:,drive] = Hout[freq][:,drive] * (-1)
                else:
                    Hout[freqs[0]][drive,:] = Hout[freqs[0]][drive,:] * (-1)


    out_dict = {'Hout': Hout, 'Hout_amp': Hout_amp, 'Hout_phase': Hout_phase, \
                'Hout_noise': Hout_noise}

    return out_dict
        # Analysis nested in try/except block just in case there
        # is a corrupted file or something
        try:
            # Load data, computer FFTs and plot if requested
            obj = hsDat(f)
            dat_fft = np.fft.rfft(obj.dat[:, data_ax])

            elec_mon = obj.dat[:,drive_ax]
            drive_fft = np.fft.rfft(elec_mon)

            elec_filt = tabor_mon_fac * signal.filtfilt(b2, a2, elec_mon)

            zeros = np.zeros(nsamp)
            voltage = np.array([zeros, zeros, zeros, elec_filt, \
                       -elec_filt, zeros, zeros, zeros])
            efield = bu.trap_efield(voltage)

            # plt.loglog(freqs, np.abs(np.fft.rfft(elec_mon)))
            # plt.loglog(freqs, np.abs(np.fft.rfft(elec_filt)))
            # # for i in [0,1,2]:
            # #     plt.plot(voltage[3][:50000]-voltage[4][:50000])
            # plt.show()

            max_ind = np.argmax(np.abs(drive_fft))
            freq_guess = freqs[max_ind]
            phase_guess = np.mean(np.angle(drive_fft[max_ind-2:max_ind+2]))
            amp_guess = np.sqrt(2) * np.std(efield[0])
            p0 = [amp_guess, freq_guess, phase_guess, 0]

            fit_ind = int(0.01 * len(time_vec))
            popt, pcov = opti.curve_fit(sine, time_vec[:fit_ind], efield[0][:fit_ind], p0=p0)
Ejemplo n.º 6
0
def run_mc(params):

    ind = params[0]
    pressure = params[1]
    drive_freq = params[2]
    drive_voltage = params[3]
    drive_voltage_noise = params[4]
    drive_phase_noise = params[5]
    init_angle = params[6]
    discretized_phase = params[7]

    beta_rot = pressure * np.sqrt(m0) / kappa
    drive_amp = np.abs(bu.trap_efield([0, 0, 0, drive_voltage, -1.0*drive_voltage, \
                                       0, 0, 0], nsamp=1)[0])
    drive_amp_noise = drive_voltage_noise * (drive_amp / drive_voltage)

    seed = seed_init * (ind + 1)

    xi_0 = np.array([np.pi/2.0, 0.0, 0.0, \
                     0.0, 2.0*np.pi*drive_freq, 0.0])

    time_constant = Ibead / beta_rot

    np.random.seed(seed)

    ### If desired, set a thermalization time equal to 10x the time constant
    ### for this particular pressure and Ibead combination
    if variable_thermalization:
        t_therm = np.min([10.0 * time_constant, 300.0])
        nthermfiles = int(t_therm / out_file_length) + 1
    else:
        t_therm = user_t_therm
        nthermfiles = user_nthermfiles

    values_to_save = {}
    values_to_save['mbead'] = mbead
    values_to_save['Ibead'] = Ibead
    values_to_save['kappa'] = kappa
    values_to_save['beta_rot'] = beta_rot
    values_to_save['p0'] = p0
    values_to_save['fsamp'] = fsamp
    values_to_save['fsim'] = fsim
    values_to_save['seed'] = seed
    values_to_save['xi_0'] = xi_0
    values_to_save['init_angle'] = init_angle
    values_to_save['pressure'] = pressure
    values_to_save['m0'] = m0
    values_to_save['drive_freq'] = drive_freq
    values_to_save['drive_amp'] = drive_amp
    values_to_save['drive_amp_noise'] = drive_amp_noise
    values_to_save['drive_phase_noise'] = drive_phase_noise
    values_to_save['discretized_phase'] = discretized_phase
    values_to_save['t_therm'] = t_therm

    if not TEST:
        base_filename = os.path.join(base, 'mc_{:d}/'.format(ind))

        bu.make_all_pardirs(os.path.join(base_filename, 'derp.txt'))

        param_path = os.path.join(base_filename, 'params.p')
        pickle.dump(values_to_save, open(param_path, 'wb'))


    def E_phi_func(t, t_therm=0.0, init_angle=0.0):
        raw_val = 2.0 * np.pi * drive_freq * (t + t_therm) + init_angle
        if discretized_phase:
            n_disc = int(raw_val / discretized_phase)
            return n_disc * discretized_phase
        else:
            return raw_val

    ### Matrix for the stochastic driving processes
    torque_noise = np.sqrt(4.0 * kb * T * beta_rot)
    # B = np.array([[0, 0,   0,   0],
    #               [0, 0,   0,   0],
    #               [0, 0, 1.0,   0],
    #               [0, 0,   0, 1.0]])
    B = np.array([[0, 0, 0,   0,   0,   0],
                  [0, 0, 0,   0,   0,   0],
                  [0, 0, 0,   0,   0,   0],
                  [0, 0, 0, 1.0,   0,   0],
                  [0, 0, 0,   0, 1.0,   0],
                  [0, 0, 0,   0,   0, 1.0]])
    B *= torque_noise / Ibead

    ### Define the system such that d(xi) = f(xi, t) * dt
    # @jit()
    def f(x, t):
        torque_theta = drive_amp * p0 * np.sin(0.5 * np.pi - x[0]) \
                            - 1.0 * beta_rot * x[3]

        c_amp = drive_amp
        E_phi = E_phi_func(t)
        if fterm_noise:
            c_amp += drive_amp_noise * np.random.randn()
            E_phi += drive_phase_noise * np.random.randn()

        torque_phi = c_amp * p0 * np.sin(E_phi - x[1]) * np.sin(x[0]) \
                            - 1.0 * beta_rot * x[4]

        torque_psi = -1.0 * beta_rot * x[5]

        return np.array([x[3], x[4], x[5], \
                         torque_theta / Ibead, \
                         torque_phi / Ibead, \
                         torque_psi / Ibead])

    ### Define the stochastic portion of the system
    # @jit()
    def G(x, t):
        newB = np.zeros((6,6))

        if gterm_noise:
            E_phi = E_phi_func(t)
            amp_noise_term = drive_amp_noise * p0 * np.sin(E_phi - x[1]) * np.sin(x[0])

            E_phi_rand = drive_phase_noise * np.random.randn()
            phase_noise_term = drive_amp * p0  * np.sin(E_phi_rand) * np.sin(x[0])
            newB[4,4] += amp_noise_term + phase_noise_term

        return B + newB


    ### Thermalize
    xi_init = np.copy(xi_0)
    for i in range(nthermfiles):
        t0 = i*out_file_length
        tf = (i+1)*out_file_length

        nsim = int(out_file_length * fsim)
        tvec = np.linspace(t0, tf, nsim+1)

        result = sdeint.itoint(f, G, xi_init, tvec).T
        xi_init = np.copy(result[:,-1])


    ### Redefine the system taking into account the thermalization time
    ### and the desired phase offset
    # @jit()
    def f(x, t):
        torque_theta = drive_amp * p0 * np.sin(0.5 * np.pi - x[0]) \
                            - 1.0 * beta_rot * x[3]

        c_amp = drive_amp
        E_phi = E_phi_func(t, t_therm=t_therm, init_angle=init_angle)
        if fterm_noise:
            c_amp += drive_amp_noise * np.random.randn()
            E_phi += drive_phase_noise * np.random.randn()

        torque_phi = c_amp * p0 * np.sin(E_phi - x[1]) * np.sin(x[0]) \
                            - 1.0 * beta_rot * x[4]

        torque_psi = -1.0 * beta_rot * x[5]

        return np.array([x[3], x[4], x[5], \
                         torque_theta / Ibead, \
                         torque_phi / Ibead, \
                         torque_psi / Ibead])


    # @jit()
    # def f(x, t):
    #     torque_theta = - 1.0 * beta_rot * x[2]
    #     torque_phi = - 1.0 * beta_rot * x[3]

    #     return np.array([x[2], x[3], torque_theta / Ibead, torque_phi / Ibead])

    ### Define the stochastic portion of the system
    def G(x, t):
        newB = np.zeros((6,6))
        if gterm_noise:
            E_phi = E_phi_func(t, t_therm=t_therm, init_angle=init_angle)
            amp_noise_term = drive_amp_noise * p0 * np.sin(E_phi - x[1]) * np.sin(x[0])

            E_phi_rand = drive_phase_noise * np.random.randn()
            phase_noise_term = drive_amp * p0  * np.sin(E_phi_rand) * np.sin(x[0])

        newB[4,4] += amp_noise_term + phase_noise_term

        return B + newB




    ### Run the simulation with the thermalized solution
    for i in range(nfiles):
        # start = time.time()
        t0 = i*out_file_length
        tf = (i+1)*out_file_length

        nsim = int(out_file_length * fsim)
        tvec = np.linspace(t0, tf, nsim+1)

        ### Solve!
        # print('RUNNING SIM')
        result = sdeint.itoint(f, G, xi_init, tvec).T
        xi_init = np.copy(result[:,-1])

        tvec = tvec[:-1]
        soln = result[:,:-1]

        # print('DOWNSAMPLING')
        nsamp = int(out_file_length * fsamp)
        # soln_ds, tvec_ds = signal.resample(soln, t=tvec, \
        #                                    num=nsamp, axis=-1)
        # soln_ds = signal.decimate(soln, int(upsamp))

        tvec_ds = tvec[::int(upsamp)]
        soln_ds = soln[:,::int(upsamp)]

        # plt.plot(tvec, soln[1])
        # plt.plot(tvec_ds, soln_ds[1])
        # plt.plot(tvec_ds, soln_ds_2[1])

        # plt.show()

        if not TEST:
            out_arr = np.concatenate( (tvec_ds.reshape((1, len(tvec_ds))), soln_ds) )

            filename = os.path.join(base_filename, 'outdat_{:d}.h5'.format(i)) 

            fobj = h5py.File(filename, 'w')
            fobj.create_dataset('sim_data', data=out_arr, compression='gzip', \
                                compression_opts=9)
            fobj.close()

        # stop = time.time()
        # print('Time for one file: {:0.1f}'.format(stop-start))

    return seed
Ejemplo n.º 7
0
import hs_digitizer as hs
import configuration as config

import peakdetect as pdet

from numba import jit
from joblib import Parallel, delayed

nfiles = 5

datadir = '/data/old_trap/20191017/bead1/spinning/ringdown/110kHz_start_6/'

bw = 10.0

mon_fac = 100.0
volt_to_efield = np.abs(bu.trap_efield([0, 0, 0, 1, 0, 0, 0, 0], nsamp=1)[0])

files, lengths = bu.find_all_fnames(datadir, ext='.h5', sort_time=True)
files = files[:nfiles]

real_drive = {}
for fileind, file in enumerate(files):
    obj = hs.hsDat(file)

    fsamp = obj.attribs['fsamp']
    nsamp = obj.attribs['nsamp']

    freqs = np.fft.rfftfreq(nsamp, d=1.0 / fsamp)

    sig_filt = obj.dat[:, 1]
    for i in range(10):
def find_step_cal_response(file_obj, bandwidth=1., include_in_phase=False, \
                           using_tabor=False, tabor_ind=3, mon_fac=100, \
                           ecol=-1, pcol=-1, new_trap=False, plot=False, \
                           userphase=0.0, nearest=True):
    '''Analyze a data step-calibraiton data file, find the drive frequency,
       correlate the response to the drive

       INPUTS:   file_obj, input file object
                 bandwidth, bandpass filter bandwidth

       OUTPUTS:  H, (response / drive)'''

    if new_trap:
        using_tabor = False

    if not using_tabor:
        if pcol == -1:
            if ecol == -1:
                ecol = np.argmax(file_obj.electrode_settings['driven'])
            pcol = config.elec_map[ecol]

        efield = bu.trap_efield(file_obj.electrode_data, new_trap=new_trap)
        drive = efield[pcol]

        #drive = file_obj.electrode_data[ecol]
        if plot:
            fig, axarr = plt.subplots(2, 1, sharex=True)
            tvec = np.arange(file_obj.nsamp) * (1.0 / file_obj.fsamp)
            colors = bu.get_color_map(len(file_obj.electrode_data),
                                      cmap='plasma')
            for i in range(len(file_obj.electrode_data)):
                try:
                    if file_obj.electrode_settings['driven'][i]:
                        ext = ' - driven'
                    else:
                        ext = ''
                    axarr[0].plot(tvec,
                                  file_obj.electrode_data[i],
                                  color=colors[i],
                                  label='Elec. {:s}{:s}'.format(str(i), ext))
                except:
                    2 + 2
            axarr[0].set_title('Electrode and Efield Data')
            axarr[0].set_ylabel('Voltage [V]')
            axarr[0].legend(fontsize=10, ncol=2, loc='upper right')

            for i, ax in enumerate(['X', 'Y', 'Z']):
                axarr[1].plot(tvec, efield[i], label=ax)
            axarr[1].set_ylabel('Efield [V/m]')
            axarr[1].set_xlabel('Time [s]')
            axarr[1].legend(fontsize=10, loc='upper right')

            fig.tight_layout()
            plt.show(fig)
            input()

        # plt.plot(drive)
        # plt.show()
        #drive = efield[ecol]

    elif using_tabor:
        pcol = 0
        v3 = file_obj.other_data[tabor_ind] * mon_fac
        v4 = file_obj.other_data[tabor_ind + 1] * mon_fac
        zeros = np.zeros(len(v3))

        if plot:
            colors = bu.get_color_map(2, cmap='plasma')
            plt.figure()
            plt.plot(v3,
                     color=colors[0],
                     label='Elec. {:s}'.format(str(tabor_ind)))
            plt.plot(v4,
                     color=colors[1],
                     label='Elec. {:s}'.format(str(tabor_ind + 1)))
            plt.title('Electrode data [V]')
            plt.legend()
            plt.tight_layout()
            plt.show()

            input()

        fac = 1.0
        if np.std(v4) < 0.5 * np.std(v3):
            # print('Only one Tabor drive channel being digitized...')
            v4 = zeros
            fac = 2.0
        elif np.std(v3) < 0.5 * np.std(v4):
            # print('Only one Tabor drive channel being digitized...')
            v3 = zeros
            fac = 2.0

        voltages = []
        for i in range(8):
            if i == tabor_ind:
                voltages.append(v3)
            elif i == (tabor_ind + 1):
                voltages.append(v4)
            else:
                voltages.append(zeros)

        drive = bu.trap_efield(voltages, new_trap=new_trap)[pcol] * fac

    # try:
    #     power = np.mean(file_obj.power)
    # except Exception:
    #     power = 0.0
    #     traceback.print_exc()

    zpos = np.mean(file_obj.pos_data[2])

    #drive = bu.detrend_poly(drive, order=1.0, plot=True)
    drive_fft = np.fft.rfft(drive)

    ### Find the drive frequency
    freqs = np.fft.rfftfreq(len(drive), d=1. / file_obj.fsamp)
    drive_freq = freqs[np.argmax(np.abs(drive_fft[1:])) + 1]

    # plt.plot(drive)
    # plt.show()
    # input()

    # print(drive_freq)
    # for i in range(3):
    #     plt.plot(efield[i], label=str(i))
    # plt.legend()
    # plt.show()

    ### Extract the response and detrend
    # response = file_obj.pos_data[pcol]
    if new_trap:
        response = file_obj.pos_data_3[pcol]
    else:
        response = file_obj.pos_data[pcol]
    #response = bu.detrend_poly(response, order=1.0, plot=True)

    ### Configure a time array for plotting and fitting
    cut_samp = config.adc_params["ignore_pts"]
    N = len(drive)
    dt = 1. / file_obj.fsamp
    t = np.linspace(0, (N + cut_samp - 1) * dt, N + cut_samp)
    t = t[cut_samp:]

    # print(drive_freq)
    # if drive_freq < 10.0:
    #     print(file_obj.fname)
    #     plt.plot(t, drive)
    #     plt.figure()
    #     plt.loglog(freqs, np.abs(drive_fft))
    #     plt.show()

    if drive_freq < 0.5 * bandwidth:
        apply_filter = False
    else:
        apply_filter = True

    ### Bandpass filter the response
    if apply_filter:
        b, a = signal.butter(3, [2.*(drive_freq-bandwidth/2.)/file_obj.fsamp, \
                              2.*(drive_freq+bandwidth/2.)/file_obj.fsamp ], btype = 'bandpass')
        responsefilt = signal.filtfilt(b, a, response)
    else:
        responsefilt = np.copy(response)

    if plot:
        plt.figure()
        plt.loglog(freqs, np.abs(np.fft.rfft(drive)))
        plt.loglog(freqs, np.abs(np.fft.rfft(responsefilt)))
        plt.show()

        input()

    ### Compute the full, normalized correlation and extract amplitude
    corr_full = bu.correlation(drive, responsefilt, file_obj.fsamp, drive_freq)
    ncorr = len(corr_full)

    phase_ratio = userphase / (2.0 * np.pi)
    phase_inds = np.array(
        [np.floor(phase_ratio * ncorr),
         np.ceil(phase_ratio * ncorr)],
        dtype='int')

    response_inphase = corr_full[0]
    response_max = np.max(corr_full)
    # try:
    response_userphase = np.interp([phase_ratio * ncorr], phase_inds,
                                   corr_full[phase_inds])[0]
    # except:
    #     response_userphase = corr_full[phase_inds[0]]

    ### Compute the drive amplitude, assuming it's a sine wave
    drive_amp = np.sqrt(2) * np.std(drive)  # Assume drive is sinusoidal
    # print(drive_amp)

    outdict = {}
    outdict['inphase'] = response_inphase / drive_amp
    outdict['max'] = response_max / drive_amp
    outdict['userphase'] = response_userphase / drive_amp
    outdict['userphase_nonorm'] = response_userphase
    outdict['drive'] = drive_amp
    outdict['drive_freq'] = drive_freq
    outdict['pcol'] = pcol

    return outdict
def proc_file(file):

    fobj = bu.hsDat(file, load=True)

    vperp = fobj.dat[:, 0]
    elec3 = fobj.dat[:, 1]

    try:
        phi_dg = fobj.attribs['phi_dg']
    except:
        phi_dg = 0.0

    inds = np.abs(full_freqs - fspin) < 200.0

    cut = int(0.1 * fsamp)
    zeros = np.zeros_like(elec3[:cut])
    voltages = [
        zeros, zeros, zeros, elec3[:cut], -1.0 * elec3[:cut], zeros, zeros,
        zeros
    ]
    efield = bu.trap_efield(voltages, only_x=True)[0]
    drive_amp, drive_phase = bu.get_sine_amp_phase(efield)

    elec3_fft = np.fft.rfft(elec3)
    true_fspin = np.average(full_freqs[inds], weights=np.abs(elec3_fft)[inds])


    carrier_amp, carrier_phase_mod = \
            bu.demod(vperp, true_fspin, fsamp, plot=plot_carrier_demod, \
                     filt=True, bandwidth=bandwidth,
                     notch_freqs=notch_freqs, notch_qs=notch_qs, \
                     tukey=True, tukey_alpha=5.0e-4, \
                     detrend=True, detrend_order=1, harmind=2.0)

    # b1, a1 = signal.butter(3, np.array(libration_filt_band)*2.0/fsamp, btype='bandpass')
    sos = signal.butter(3,
                        libration_filt_band,
                        btype='bandpass',
                        fs=fsamp,
                        output='sos')
    # carrier_phase_mod_filt = signal.filtfilt(b1, a1, carrier_phase_mod)
    carrier_phase_mod_filt = signal.sosfiltfilt(sos, carrier_phase_mod)

    if len(libration_filt_band):
        libration_inds = (full_freqs > libration_filt_band[0]) \
                                * (full_freqs < libration_filt_band[1])
    else:
        libration_inds = np.abs(full_freqs -
                                libration_guess) < 0.5 * libration_bandwidth

    phase_mod_fft = np.fft.rfft(carrier_phase_mod) * fac

    lib_fit_x = full_freqs[libration_inds]
    lib_fit_y = np.abs(phase_mod_fft[libration_inds])

    try:
        try:
            peaks = bu.find_fft_peaks(lib_fit_x,
                                      lib_fit_y,
                                      delta_fac=5.0,
                                      window=50)
            ind = np.argmax(peaks[:, 1])
        except:
            peaks = bu.find_fft_peaks(lib_fit_x,
                                      lib_fit_y,
                                      delta_fac=3.0,
                                      window=100)
            ind = np.argmax(peaks[:, 1])

        true_libration_freq = peaks[ind, 0]

    except:
        true_libration_freq = lib_fit_x[np.argmax(lib_fit_y)]

    libration_amp, libration_phase = \
            bu.demod(carrier_phase_mod, true_libration_freq, fsamp, \
                     plot=plot_libration_demod, filt=True, \
                     filt_band=libration_filt_band, \
                     bandwidth=libration_bandwidth, \
                     tukey=False, tukey_alpha=5.0e-4, \
                     detrend=False, detrend_order=1.0, harmind=1.0)

    libration_ds, time_vec_ds = \
            signal.resample(carrier_phase_mod_filt, t=time_vec, num=out_nsamp)
    libration_amp_ds, time_vec_ds = \
            signal.resample(libration_amp, t=time_vec, num=out_nsamp)

    libration_ds = libration_ds[out_cut:int(-1 * out_cut)]
    libration_amp_ds = libration_amp_ds[out_cut:int(-1 * out_cut)]
    time_vec_ds = time_vec_ds[out_cut:int(-1 * out_cut)]

    if plot_downsample:
        plt.plot(time_vec,
                 carrier_phase_mod_filt,
                 color='C0',
                 label='Original')
        plt.plot(time_vec_ds,
                 libration_ds,
                 color='C0',
                 ls='--',
                 label='Downsampled')
        plt.plot(time_vec, libration_amp, color='C1')  #, label='Original')
        plt.plot(time_vec_ds, libration_amp_ds, color='C1',
                 ls='--')  #, label='Downsampled')
        plt.legend()
        plt.show()

        input()

    return (time_vec_ds, libration_ds, libration_amp_ds, \
                true_libration_freq, phi_dg, drive_amp)
inds = np.abs(full_freqs - fspin) < 200.0

elec3_fft = np.fft.rfft(elec3)
true_fspin = np.average(full_freqs[inds], weights=np.abs(elec3_fft)[inds])

if not libration_guess:

    try:
        elec3_cut = 100.0 * (50000.0 / 53000.0) * elec3[:int(fsamp)]
        zeros = np.zeros_like(elec3_cut)
        voltages = [
            zeros, zeros, zeros, elec3_cut, -1.0 * elec3_cut, zeros, zeros,
            zeros
        ]
        efield = bu.trap_efield(voltages, only_x=True)[0]
        efield_amp, _ = bu.get_sine_amp_phase(efield)

        dipole = np.load(dipole_file)[0]
        Ibead = bu.get_Ibead(date=date, rhobead=rhobead)['val']

        libration_guess = np.sqrt(efield_amp * dipole / Ibead) / (2.0 * np.pi)

    except:

        amp, phase_mod = bu.demod(vperp, true_fspin, fsamp, plot=plot_carrier_demod, \
                                  filt=True, bandwidth=bandwidth,
                                  tukey=True, tukey_alpha=5.0e-4, \
                                  detrend=True, detrend_order=1, harmind=2.0)

        phase_mod_fft = np.fft.rfft(phase_mod) * fac
def run_mc(params):

    ind = params[0]
    pressure = params[1]
    drive_freq = params[2]
    drive_voltage = params[3]
    drive_voltage_noise = params[4]
    drive_phase_noise = params[5]
    init_angle = params[6]

    beta_rot = pressure * np.sqrt(m0) / kappa
    drive_amp = np.abs(bu.trap_efield([0, 0, 0, drive_voltage, -1.0*drive_voltage, \
                                        0, 0, 0], nsamp=1)[0])
    drive_amp_noise = drive_voltage_noise * (drive_amp / drive_voltage)

    lib_freq = np.sqrt(drive_amp * p0 * dipole_units / Ibead) / (2.0 * np.pi)

    xi_0 = np.array([p0*np.cos(init_angle), p0*np.sin(init_angle), 0.0, \
                        0.0, 0.0, 2.0 * np.pi * drive_freq])

    seed = seed_init * (ind + 1)

    np.random.seed(seed)

    values_to_save = {}
    values_to_save['mbead'] = mbead
    values_to_save['Ibead'] = Ibead
    values_to_save['p0'] = p0
    values_to_save['fsamp'] = fsamp
    values_to_save['seed'] = seed
    values_to_save['xi_0'] = xi_0
    values_to_save['pressure'] = pressure
    values_to_save['drive_freq'] = drive_freq
    values_to_save['drive_amp'] = drive_amp
    values_to_save['drive_amp_noise'] = drive_amp_noise
    values_to_save['drive_phase_noise'] = drive_phase_noise

    base_filename = os.path.join(base, 'mc_{:d}/'.format(ind))

    bu.make_all_pardirs(os.path.join(base_filename, 'derp.txt'))

    param_path = os.path.join(base_filename, 'params.p')
    pickle.dump(values_to_save, open(param_path, 'wb'))

    @jit()
    def rhs(t, xi):
        '''This function represents the right-hand side of the differential equation
           d(xi)/dt = rhs(t, xi), where xi is a 6-dimensional vector representing the 
           system of a rotating microsphere: {px, py, pz, omegax, omegay, omegaz}, 
           with p the dipole moment and omega the angular velocity. The system is 
           solved in Cartesian coordinates to avoid the branch cuts inherent to 
           integrating phase angles.

           The function computes the following torques:
                thermal torque, white noise with power computed from above global
                                    parameters and fluctuation dissipation theorem 
                drag torque, computed as (- beta * omega)
                drive torque, computed as (-1.0) * {px, py, pz} (cross) {Ex, Ey, Ez}
                optical torque, constant torque about the z axis
        '''
        drag_torque = -1.0 * beta_rot * xi[3:]

        #### Construct the rotating Efield drive
        Efield = np.array([drive_amp * np.cos(2.0 * np.pi * drive_freq * t), \
                           drive_amp * np.sin(2.0 * np.pi * drive_freq * t), \
                           0.0])

        drive_torque = np.cross(xi[:3] * dipole_units, Efield)
        optical_torque = np.array([0.0, 0.0, N_opt])

        total_torque = drive_torque + drag_torque + optical_torque

        return np.concatenate(
            (-1.0 * np.cross(xi[:3], xi[3:]), total_torque / Ibead))

    @jit()
    def rhs_stochastic(t, xi):
        '''Basically the same as above rhs() function, but this only includes the 
           stochastic forcing terms. Doesn't update the dipole moment projections,
           just adds more (Delta omega)

           The function computes the following torques:
                thermal torque, white noise with power computed from above global
                                    parameters and fluctuation dissipation theorem 
                drive torque, computed as (-1.0) * {px, py, pz} (cross) {Ex, Ey, Ez}
                                where the Efield only includes noise terms
        '''
        thermal_torque = np.sqrt(
            4.0 * kb * T * beta_rot * fsim) * np.random.randn(3)

        ### Amplitude noise for all three axes
        an = drive_amp_noise * np.random.randn(3)

        ### Phase noise for the two drive axes
        pn = drive_phase_noise * np.random.randn(2)

        #### Construct the rotating Efield drive
        Efield1 = np.array([drive_amp * np.cos(2.0 * np.pi * drive_freq * t), \
                            drive_amp * np.sin(2.0 * np.pi * drive_freq * t), \
                            0.0])
        Efield2 = np.array([drive_amp * np.cos(2.0 * np.pi * drive_freq * t + pn[0]), \
                            drive_amp * np.sin(2.0 * np.pi * drive_freq * t + pn[1]), \
                            0.0])
        Efield = Efield2 - Efield1 + an

        drive_torque = np.cross(xi[:3] * dipole_units, Efield)

        total_torque = drive_torque + thermal_torque

        return np.concatenate((np.zeros(3), total_torque / Ibead))

    for i in range(nfiles):
        #bu.progress_bar(i, nfiles)

        t0 = i * out_file_length
        tf = (i + 1) * out_file_length

        tvec, soln = stepper(xi_0, t0, tf, dt_sim, upsamp, rk4, rhs, \
                             system_stochastic=rhs_stochastic)

        xi_0 = soln[:, -1]

        tvec = tvec[:-1]
        soln = soln[:, :-1]
        out_arr = np.concatenate((tvec.reshape((1, len(tvec))), soln))

        filename = os.path.join(base_filename, 'outdat_{:d}.h5'.format(i))
        fobj = h5py.File(filename, 'w')
        # group = fobj.create_group('sim_data')
        fobj.create_dataset('sim_data', data=out_arr, compression='gzip', \
                            compression_opts=9)
        fobj.close()

        # filename = os.path.join(base_filename, 'outdat_{:d}.npy'.format(i))
        # np.save(open(filename, 'wb'), out_arr)

    return seed
Ejemplo n.º 12
0
def weigh_bead_efield(files, elec_ind, pow_ind, colormap='plasma', sort='time',\
                      file_inds=(0,10000), plot=True, print_res=False, pos=False, \
                      save_mass=False, new_trap=False, correct_phase_shift=False):
    '''Loops over a list of file names, loads each file, diagonalizes,
       then plots the amplitude spectral density of any number of data
       or cantilever/electrode drive signals

       INPUTS: files, list of files names to extract data
               data_axes, list of pos_data axes to plot
               cant_axes, list of cant_data axes to plot
               elec_axes, list of electrode_data axes to plot
               diag, boolean specifying whether to diagonalize

       OUTPUTS: none, plots stuff
    '''
    date = re.search(r"\d{8,}", files[0])[0]
    suffix = files[0].split('/')[-2]

    if new_trap:
        trap_str = 'new_trap'
    else:
        trap_str = 'old_trap'

    charge_file = '/data/{:s}_processed/calibrations/charges/'.format(
        trap_str) + date
    save_filename = '/data/{:s}_processed/calibrations/masses/'.format(trap_str) \
                            + date + '_' + suffix + '.mass'
    bu.make_all_pardirs(save_filename)

    if pos:
        charge_file += '_recharge.charge'
    else:
        charge_file += '.charge'

    try:
        nq = np.load(charge_file)[0]
        found_charge = True
    except:
        found_charge = False

    if not found_charge or manual_charge:
        user_nq = input('No charge file or manual requested. Guess q: ')
        nq = int(user_nq)

    if correct_phase_shift:
        print('Correcting anomalous phase-shift during analysis.')

    # nq = -16
    print('qbead: {:d} e'.format(int(nq)))
    q_bead = nq * constants.elementary_charge

    run_index = 0

    masses = []

    nfiles = len(files)
    if not print_res:
        print("Processing %i files..." % nfiles)

    all_eforce = []
    all_power = []

    all_param = []

    mass_vec = []

    p_ac = []
    p_dc = []

    e_ac = []
    e_dc = []

    pressure_vec = []

    zamp_avg = 0
    zphase_avg = 0
    zamp_N = 0
    zfb_avg = 0
    zfb_N = 0
    power_avg = 0
    power_N = 0

    Nbad = 0

    powpsd = []

    for fil_ind, fil in enumerate(files):  # 15-65

        # 4
        # if fil_ind == 16 or fil_ind == 4:
        #     continue

        bu.progress_bar(fil_ind, nfiles)

        # Load data
        df = bu.DataFile()
        try:
            if new_trap:
                df.load_new(fil)
            else:
                df.load(fil, load_other=True)
        except Exception:
            traceback.print_exc()
            continue

        try:
            # df.calibrate_stage_position()
            df.calibrate_phase()
        except Exception:
            traceback.print_exc()
            continue

        if ('20181129' in fil) and ('high' in fil):
            pressure_vec.append(1.5)
        else:
            try:
                pressure_vec.append(df.pressures['pirani'])
            except Exception:
                pressure_vec.append(0.0)

        ### Extract electrode data
        if new_trap:
            top_elec = df.electrode_data[1]
            bot_elec = df.electrode_data[2]
        else:
            top_elec = mon_fac * df.other_data[elec_ind]
            bot_elec = mon_fac * df.other_data[elec_ind + 1]

        fac = 1.0
        if np.std(top_elec) < 0.5 * np.std(bot_elec) \
                or np.std(bot_elec) < 0.5 * np.std(top_elec):
            print(
                'Adjusting electric field since only one electrode was digitized.'
            )
            fac = 2.0

        nsamp = len(top_elec)
        zeros = np.zeros(nsamp)

        voltages = [zeros, top_elec, bot_elec, zeros, \
                    zeros, zeros, zeros, zeros]
        efield = bu.trap_efield(voltages, new_trap=new_trap)
        eforce2 = fac * sign * efield[2] * q_bead

        tarr = np.arange(0, df.nsamp / df.fsamp, 1.0 / df.fsamp)

        # fig, axarr = plt.subplots(2,1,sharex=True,figsize=(10,8))

        # axarr[0].plot(tarr, top_elec, label='Top elec.')
        # axarr[0].plot(tarr, bot_elec, label='Bottom elec.')
        # axarr[0].set_ylabel('Apparent Voltages [V]')
        # axarr[0].legend(fontsize=12, loc='upper right')

        # axarr[1].plot(tarr, efield[2])
        # axarr[1].set_xlabel('Time [s]')
        # axarr[1].set_ylabel('Apparent Electric Field [V/m]')

        # fig.tight_layout()

        # plt.show()
        # input()

        freqs = np.fft.rfftfreq(df.nsamp, d=1.0 / df.fsamp)
        drive_ind = np.argmax(np.abs(np.fft.rfft(eforce2)))
        drive_freq = freqs[drive_ind]

        zamp = np.abs( np.fft.rfft(df.zcal) * bu.fft_norm(df.nsamp, df.fsamp) * \
                       np.sqrt(freqs[1] - freqs[0]) )
        zamp *= (1064.0e-9 / 2.0) * (1.0 / (2.9 * np.pi))
        zphase = np.angle(np.fft.rfft(df.zcal))
        zamp_avg += zamp[drive_ind]
        zamp_N += 1

        #plt.loglog(freqs, zamp)
        #plt.scatter(freqs[drive_ind], zamp[drive_ind], s=10, color='r')
        #plt.show()


        zfb = np.abs(np.fft.rfft(df.pos_fb[2]) * bu.fft_norm(df.nsamp, df.fsamp) * \
                      np.sqrt(freqs[1] - freqs[0]) )
        zfb_avg += zfb[drive_ind]
        zfb_N += 1

        #eforce2 = (top_elec * e_top_func(0.0) + bot_elec * e_bot_func(0.0)) * q_bead
        if noise:
            e_dc.append(np.mean(eforce2))
            e_ac_val = np.abs(np.fft.rfft(eforce2))[drive_ind]
            e_ac.append(e_ac_val * bu.fft_norm(df.nsamp, df.fsamp) \
                        * np.sqrt(freqs[1] - freqs[0]) )

        zphase_avg += (zphase[drive_ind] - np.angle(eforce2)[drive_ind])

        if np.sum(df.power) == 0.0:
            current = np.abs(df.other_data[pow_ind]) / trans_gain
        else:
            fac = 1e-6
            current = fac * df.power / trans_gain

        power = current / pd_gain
        power = power / line_filter_trans
        power = power / bs_fac

        power_avg += np.mean(power)
        power_N += 1
        if noise:
            p_dc.append(np.mean(power))
            p_ac_val = np.abs(np.fft.rfft(power))[drive_ind]
            p_ac.append(p_ac_val * bu.fft_norm(df.nsamp, df.fsamp) \
                        * np.sqrt(freqs[1] - freqs[0]) )

        fft1 = np.fft.rfft(power)
        fft2 = np.fft.rfft(df.pos_fb[2])

        if not len(powpsd):
            powpsd = np.abs(fft1)
            Npsd = 1
        else:
            powpsd += np.abs(fft1)
            Npsd += 1

        # freqs = np.fft.rfftfreq(df.nsamp, d=1.0/df.fsamp)
        # plt.loglog(freqs, np.abs(np.fft.rfft(eforce2)))
        # plt.loglog(freqs, np.abs(np.fft.rfft(power)))
        # plt.show()
        # input()

        # fig, axarr = plt.subplots(2,1,sharex=True,figsize=(10,8))

        # axarr[0].plot(tarr, power)
        # axarr[0].set_ylabel('Measured Power [Arb.]')

        # axarr[1].plot(tarr, power)
        # axarr[1].set_xlabel('Time [s]')
        # axarr[1].set_ylabel('Measured Power [Arb.]')

        # bot, top = axarr[1].get_ylim()
        # axarr[1].set_ylim(1.05*bot, 0)

        # fig.tight_layout()

        # plt.show()
        # input()


        bins, dat, errs = bu.spatial_bin(eforce2, power, nbins=200, width=0.0, #width=0.05, \
                                         dt=1.0/df.fsamp, harms=[1], \
                                         add_mean=True, verbose=False, \
                                         correct_phase_shift=correct_phase_shift, \
                                         grad_sign=0)

        dat = dat / np.mean(dat)

        #plt.plot(bins, dat, 'o')
        #plt.show()

        popt, pcov = opti.curve_fit(line, bins*1.0e13, dat, \
                                    absolute_sigma=False, maxfev=10000)
        test_vals = np.linspace(np.min(eforce2 * 1.0e13),
                                np.max(eforce2 * 1.0e13), 100)

        fit = line(test_vals, *popt)

        lev_force = -popt[1] / (popt[0] * 1.0e13)
        mass = lev_force / (9.806)

        #umass = ulev_force / 9.806
        #lmass = llev_force / 9.806

        if mass > upper_outlier or mass < lower_outlier:
            print('Crazy mass: {:0.2f} pg.... ignoring'.format(mass * 1e15))
            # fig, axarr = plt.subplots(3,1,sharex=True)
            # axarr[0].plot(eforce2)
            # axarr[1].plot(power)
            # axarr[2].plot(df.pos_data[2])
            # ylims = axarr[1].get_ylim()
            # axarr[1].set_ylim(ylims[0], 0)
            # plt.show()
            continue

        all_param.append(popt)

        all_eforce.append(bins)
        all_power.append(dat)

        mass_vec.append(mass)

    if noise:
        print('DC power: ', np.mean(p_dc), np.std(p_dc))
        print('AC power: ', np.mean(p_ac), np.std(p_ac))
        print('DC field: ', np.mean(e_dc), np.std(e_dc))
        print('AC field: ', np.mean(e_ac), np.std(e_ac))
        return

    #plt.plot(mass_vec)

    mean_popt = np.mean(all_param, axis=0)

    mean_lev = np.mean(mass_vec) * 9.806
    plot_vec = np.linspace(np.min(all_eforce), mean_lev, 100)

    if plot:
        fig = plt.figure(dpi=200, figsize=(6, 4))
        ax = fig.add_subplot(111)
        ### Plot force (in pN / g = pg) vs power
        plt.plot(np.array(all_eforce).flatten()[::5]*1e15*(1.0/9.806), \
                 np.array(all_power).flatten()[::5], \
                 'o', alpha = 0.5)
        #for params in all_param:
        #    plt.plot(plot_vec, line(plot_vec, params[0]*1e13, params[1]), \
        #             '--', color='r', lw=1, alpha=0.05)
        plt.plot(plot_vec*1e12*(1.0/9.806)*1e3, \
                 line(plot_vec, mean_popt[0]*1e13, mean_popt[1]), \
                 '--', color='k', lw=2, \
                 label='Implied mass: %0.1f pg' % (np.mean(mass_vec)*1e15))
        left, right = ax.get_xlim()
        # ax.set_xlim((left, 500))
        ax.set_xlim(*xlim)

        bot, top = ax.get_ylim()
        ax.set_ylim((0, top))

        plt.legend()
        plt.xlabel('Applied electrostatic force/$g$ (pg)')
        plt.ylabel('Optical power (arb. units)')
        plt.grid()
        plt.tight_layout()
        if save_example:
            fig.savefig(example_filename)
            fig.savefig(example_filename[:-4] + '.pdf')
            fig.savefig(example_filename[:-4] + '.svg')

        x_plotvec = np.array(all_eforce).flatten()
        y_plotvec = np.array(all_power).flatten()

        yresid = (y_plotvec - line(x_plotvec, mean_popt[0] * 1e13,
                                   mean_popt[1])) / y_plotvec

        plt.figure(dpi=200, figsize=(3, 2))
        plt.hist(yresid * 100, bins=30)
        plt.legend()
        plt.xlabel('Resid. Power [%]')
        plt.ylabel('Counts')
        plt.grid()
        plt.tight_layout()

        plt.figure(dpi=200, figsize=(3, 2))
        plt.plot(x_plotvec * 1e15, yresid * 100, 'o')
        plt.legend()
        plt.xlabel('E-Force [pN]')
        plt.ylabel('Resid. Pow. [%]')
        plt.grid()
        plt.tight_layout()

        derpfig = plt.figure(dpi=200, figsize=(3, 2))
        #derpfig.patch.set_alpha(0.0)
        plt.hist(np.array(mass_vec) * 1e15, bins=10)
        plt.xlabel('Mass (pg)')
        plt.ylabel('Count')
        plt.grid()
        #plt.title('Implied Masses, Each from 50s Integration')
        #plt.xlim(0.125, 0.131)
        plt.tight_layout()
        if save_example:
            derpfig.savefig(example_filename[:-4] + '_hist.png')
            derpfig.savefig(example_filename[:-4] + '_hist.pdf')
            derpfig.savefig(example_filename[:-4] + '_hist.svg')

        plt.show()

    final_mass = np.mean(mass_vec)
    final_err_stat = 0.5 * np.std(mass_vec)  #/ np.sqrt(len(mass_vec))
    final_err_sys = np.sqrt((0.015**2 + 0.01**2) * final_mass**2)
    final_pressure = np.mean(pressure_vec)

    if save_mass:
        save_arr = [final_mass, final_err_stat, final_err_sys]
        np.save(open(save_filename, 'wb'), save_arr)

    print('Bad Files: %i / %i' % (Nbad, nfiles))
    if print_res:
        gresid_fac = (2.0 * np.pi * freqs[drive_ind])**2 / 9.8

        print('      mass    [pg]: {:0.1f}'.format(final_mass * 1e15))
        print('      st.err  [pg]: {:0.2f}'.format(final_err_stat * 1e15))
        print('      sys.err [pg]: {:0.2f}'.format(final_err_sys * 1e15))
        print('      qbead    [e]: {:d}'.format(
            int(round(q_bead / constants.elementary_charge))))
        print('      P     [mbar]: {:0.2e}'.format(final_pressure))
        print('      <P>    [arb]: {:0.2e}'.format(power_avg / power_N))
        print('      zresid   [g]: {:0.3e}'.format(
            (zamp_avg / zamp_N) * gresid_fac))
        print('      zphase [rad]: {:0.3e}'.format(zphase_avg / zamp_N))
        print('      zfb    [arb]: {:0.3e}'.format(zfb_avg / zfb_N))
        outarr = [ final_mass*1e15, final_err_stat*1e15, final_err_sys*1e15, \
                   q_bead/constants.elementary_charge, \
                   final_pressure, power_avg / power_N, \
                   (zamp_avg / zamp_N) * gresid_fac, \
                   zphase_avg / zamp_N, zfb_avg / zfb_N ]
        return outarr
    else:
        scaled_params = np.array(all_param)
        scaled_params[:, 0] *= 1e13

        outdic = {'eforce': all_eforce, 'power': all_power, \
                  'linear_fit_params': scaled_params, \
                  'ext_masses': mass_vec}

        return outdic