Ejemplo n.º 1
0
    def test_aperiodic_middle(self):
        # simulated dazzler array
        array = np.zeros((1000, ))
        idx = 500
        array[idx] = 2.5

        with pytest.raises(RuntimeError):
            # should say that we found aperiodic trigger
            nudie.detect_table_start(array)
Ejemplo n.º 2
0
    def test_truncated_wave_repeat_begin(self):
        # dazzler trigger is on while repeating the first waveform in a
        # table. Check that detection works if only part of the
        # "waveform_repeat" triggers are detected at the beginning

        period = 100
        offset = 1
        amplitude = 2.5
        waveform_repeat = 2

        reference = np.zeros((1010, ))
        for i in range(waveform_repeat):
            reference[offset + i::period * waveform_repeat] = amplitude

        array = reference[offset + int(waveform_repeat // 2):]

        found, dperiod = nudie.detect_table_start(
            array, waveform_repeat=waveform_repeat)

        assert dperiod == period * waveform_repeat, "incorrect period found"

        correct_idx = np.zeros_like(array)
        correct_offset = period * waveform_repeat - int(waveform_repeat // 2)
        correct_idx[correct_offset::period * waveform_repeat] = amplitude

        assert np.all(np.where(correct_idx) == found), "positions are wrong"
Ejemplo n.º 3
0
    def test_int_array_input(self):
        pytest.skip('implementation problem')
        # known fail, steps to reproduce
        N = 1000
        repeat = 2
        num_waveforms = 100
        analog1 = np.zeros((N, ), dtype=int)
        analog1[::repeat * num_waveforms] = 2.5

        start = nudie.detect_table_start(analog1, repeat)
Ejemplo n.º 4
0
    def test_periodic_middle(self):
        period = 200
        offset = 10
        amplitude = 2.5

        array = np.zeros((1000, ))
        array[offset::period] = amplitude

        found, dperiod = nudie.detect_table_start(array)

        assert dperiod == period, "incorrect period detected"
        assert np.all(np.where(array) == found), "positions are wrong"
Ejemplo n.º 5
0
    def test_waveform_repeat(self):
        period = 200
        offset = 10
        amplitude = 2.5
        waveform_repeat = 2
        array = np.zeros((1000, ))
        for i in range(waveform_repeat):
            array[offset + i::period * waveform_repeat] = amplitude

        with pytest.raises(RuntimeError):
            # incorrect waveform setting
            nudie.detect_table_start(array)

        found, dperiod = nudie.detect_table_start(array, waveform_repeat)

        assert dperiod == period * waveform_repeat, "incorrect period found"

        correct_idx = np.zeros_like(array)
        correct_idx[offset::period * waveform_repeat] = amplitude

        assert np.all(np.where(correct_idx) == found), "positions are wrong"
Ejemplo n.º 6
0
    def process(self):
        '''execute the loading of raw data from disk
        '''
        t2 = self.t2
        table = self.table
        loop = self.loop

        # first file requires special synchronization
        # this is the rule that determines that it is the first file
        first = 'first' if all([t2 == 0, table == 0, loop == 0]) else False

        # Load everything for the given t2, table, and loop value
        analogs = nudie.load_analogtxt(self.batch_name, self.batch_path, t2,
                                       table, loop)

        cdata = nudie.load_camera_file(self.batch_name,
                                       self.batch_path,
                                       t2,
                                       table,
                                       loop,
                                       force_uint16=True)

        # Synchronize it, trimming some frames in the beginning and end
        data, (a1, a2) = nudie.trim_all(*nudie.synchronize_daq_to_camera(\
                        cdata, analog_channels=analogs, which_file=first),
                        trim_to=self.trim_to)

        data = data.astype(float)  # convert to float64 before manipulating it!

        start_idxs, period = nudie.detect_table_start(a1)

        # determine where the shutter is
        shutter_open, shutter_closed, duty_cycle = nudie.determine_shutter_shots(
            data)
        shutter_info = {
            'last open idx': shutter_open,
            'first closed idx': shutter_closed
        }

        tags = nudie.tag_phases(start_idxs,
                                period,
                                tags=self.phase_cycles,
                                nframes=data.shape[1],
                                shutter_info=shutter_info)

        self.tags = tags
        self.camera_data = data

        achan = namedtuple('analog_channels', ['a1', 'a2'])
        self.analog_channels = achan(a1, a2)

        return self
Ejemplo n.º 7
0
def test_general_case(offset, repeat, nwave):
    ''' tests whether many waveforms, several phase case works'''
    duty_cycle = 0.5
    transition_width = 3
    nframes = 100
    waveform_repeat = repeat
    nwaveforms = nwave
    tags = ['one', 'two', 'three']
    nphases = len(tags)

    data = np.zeros((nframes, 4), dtype=object)
    analog1 = np.zeros((nframes, ), dtype=float)

    states = it.cycle([(rep, phase, waveform) \
            for waveform in range(nwaveforms) \
            for phase in tags \
            for rep in range(waveform_repeat)])

    #states, tmp = it.tee(states, 2)
    #test = list(it.islice(tmp, offset, 100))

    # set shutter
    last_on = int(nframes * duty_cycle)

    shutter_it = it.chain(it.repeat('shutter open', last_on),
                          it.repeat(None, transition_width),
                          it.repeat('shutter closed'))

    for i, (rep, phase, waveform), shut in \
            zip(range(nframes), it.islice(states, offset, None), shutter_it):
        data[i] = (rep, phase, waveform, shut)

        if waveform == 0 and phase == tags[0]: analog1[i] = 2.5

    start, period = nudie.detect_table_start(analog1, waveform_repeat)

    shutter_info = {
        'last open idx': slice(None, last_on + 1),
        'first closed idx': slice(last_on + transition_width, None)
    }

    tagged = nudie.tag_phases(start, period, tags, nframes, \
            waveform_repeat=waveform_repeat, shutter_info=shutter_info)

    for i, (r, p, w, s) in enumerate(data):
        if s is None: continue  # skip transition frames
        assert i in tagged[r][w][p][s]
Ejemplo n.º 8
0
    def test_hilo_truncated(self):
        # similarly to above, what happens if the trigger falls on the last
        # sample?

        period = 200
        offset = 1
        amplitude = 2.5

        array = np.zeros((1002, ))
        array[offset::period] = amplitude

        found, dperiod = nudie.detect_table_start(array)

        assert dperiod == period, "incorrect period found"

        correct_idx = np.zeros_like(array)
        correct_idx[offset:-period:period] = amplitude

        assert np.all(np.where(correct_idx) == found), "positions are wrong"
Ejemplo n.º 9
0
    def test_first_and_last(self):
        # what if the trigger perfectly aligns such that the trigger is
        # both on the first and last sample?

        period = 200
        offset = 0
        amplitude = 2.5

        array = np.zeros((1001, ))
        array[offset::period] = amplitude

        found, dperiod = nudie.detect_table_start(array)

        assert dperiod == period, "incorrect period found"

        correct_idx = np.zeros_like(array)
        correct_idx[period:-period:period] = amplitude

        assert np.all(np.where(correct_idx) == found), "positions are wrong"
Ejemplo n.º 10
0
    def test_lohi_truncated(self):
        # if the trigger is on on the first sample, np.diff will
        # not give a lohi transition, only a hilo. This will make
        # lohi shorter than hilo by one sample

        period = 200
        amplitude = 2.5
        offset = 0

        array = np.zeros((1010, ))
        array[offset::period] = amplitude

        found, dperiod = nudie.detect_table_start(array)

        assert dperiod == period, "incorrect period found"

        correct_idx = np.zeros_like(array)
        correct_idx[offset + period::period] = amplitude

        assert np.all(np.where(correct_idx) == found), "positions are wrong"
Ejemplo n.º 11
0
    def test_truncated_wave_repeat_end(self):
        # similar to above but truncated at end

        period = 100
        amplitude = 2.5
        waveform_repeat = 2

        reference = np.zeros((1002, ))
        for i in range(waveform_repeat):
            reference[i::period * waveform_repeat] = amplitude

        array = reference[:-int(waveform_repeat // 2)]

        found, dperiod = nudie.detect_table_start(
            array, waveform_repeat=waveform_repeat)

        rep = period * waveform_repeat
        assert dperiod == rep, "incorrect period found"

        correct_idx = np.zeros_like(array)
        correct_idx[rep:-rep:rep] = amplitude
        assert np.all(np.where(correct_idx) == found), "positions are wrong"
Ejemplo n.º 12
0
Archivo: tg.py Proyecto: antonl/nudie
def run(tg_name,
        tg_batch,
        when='today',
        wavelengths=None,
        plot=False,
        pad_to=2048,
        prd_est=850.,
        lo_width=200,
        dc_width=200,
        gaussian_power=2.,
        analysis_path='./analyzed',
        datapath=None):

    if plot:
        import matplotlib as mpl
        import matplotlib.pyplot as plt
        mpl.rcParams['figure.figsize'] = (16, 12)
        mpl.use('Qt4Agg')

    nrepeat = 1  # how many times each waveform is repeated in the camera file. Assumed to be one
    waveforms_per_table = 1
    npixels = 1340
    trim_to = 3, -3

    # change datapath
    if datapath is None:
        datapath = nudie.data_folder

    # load up tg data to use
    tg_info = next(
        nudie.load_job(job_name=tg_name,
                       batch_set=[tg_batch],
                       when=when,
                       data_path=datapath))

    phase_cycles = ['none1', 'zero', 'none2', 'pipi']

    # set current batch directory
    current_path = Path(tg_info['batch_path'])

    # generate hdf filename based on data date
    analysis_folder = Path(analysis_path)

    # create folder if it doesn't exist
    if not analysis_folder.exists():
        analysis_folder.mkdir(parents=True)

    save_path = analysis_folder / (tg_info['batch_name'] + '.h5')

    # remove data file if it exists
    if save_path.exists(): save_path.unlink()

    with h5py.File(str(save_path), 'w') as sf:
        # initialize groups
        sf.create_group('axes')

        shape = (tg_info['nt2'], npixels)
        sf.create_dataset('raw transient-grating', shape, dtype=complex)

    try:
        wl = load_wavelengths(current_path.parent / wavelengths)
    except FileNotFoundError as e:
        nudie.log.error('Could not load wavelength calibration file!')
        raise e

    # define gaussian
    def gaussian2(w, x0, x):
        c = 4 * np.log(2)
        ξ = x - x0
        return np.exp(-c * (ξ / w)**(2 * gaussian_power))

    for loop, t2, table in it.product(tg_info['loop_range'],
                                      tg_info['t2_range'],
                                      tg_info['table_range']):

        if any([loop != 0, table != 0]):
            raise NotImplementedError('Code is not setup to handle multiple ' +\
                    'loops or tables.')

        # first file requires special synchronization
        # this is the rule that determines that it is the first file
        first = 'first' if all([t2 == 0, table == 0, loop == 0]) else False

        # Load everything for the given t2, table, and loop value
        analogs = nudie.load_analogtxt(tg_info['job_name'], current_path, t2,
                                       table, loop)
        cdata = nudie.load_camera_file(tg_info['job_name'],
                                       current_path,
                                       t2,
                                       table,
                                       loop,
                                       force_uint16=True)

        # Synchronize it, trimming some frames in the beginning and end
        data, (a1, a2) = nudie.trim_all(*nudie.synchronize_daq_to_camera(\
            cdata, analog_channels=analogs, which_file=first),
            trim_to=slice(*trim_to))
        data = data.astype(float)  # convert to float64 before manipulating it!

        start_idxs, period = nudie.detect_table_start(a1)

        f, data, df = nudie.wavelen_to_freq(wl, data, ret_df=True, ax=0)

        # determine where the shutter is
        shutter_open, shutter_closed, duty_cycle = nudie.determine_shutter_shots(
            data)

        # tag the phases
        shutter_info = {
            'last open idx': shutter_open,
            'first closed idx': shutter_closed
        }
        tags = nudie.tag_phases(start_idxs,
                                period,
                                tags=phase_cycles,
                                nframes=data.shape[1],
                                shutter_info=shutter_info)
        nudie.remove_incomplete_t1_waveforms(tags, phase_cycles)
        data_t = np.zeros((data.shape[1], data.shape[0]), dtype=float)

        # subtract the average of closed shutter shots from shutter open data
        for t1, k in it.product(range(waveforms_per_table), phase_cycles):
            idx_open = tags[nrepeat - 1][t1][k]['shutter open']
            idx_closed = tags[nrepeat - 1][t1][k]['shutter closed']

            data_t[idx_open, :] = data[:, idx_open].T \
                                  - data[:, idx_closed].mean(axis=1)

        t = np.fft.fftfreq(pad_to, df)
        lo_window = gaussian2(lo_width, prd_est, t)
        dc_window = gaussian2(dc_width, 0, t)

        # fft everything
        fdata = np.fft.fft(data_t, axis=1, n=pad_to)

        if all([plot, table == 0, t2 == 0]):
            tmp_fft = fdata[tags[0][0][phase_cycles[1]]['shutter open'][0], :]
            plot_windows(t, lo_window, dc_window, tmp_fft, prd_est)

        # spectral interferometry
        rIprobe = np.fft.ifft(fdata * dc_window, axis=1)[:, :npixels]
        rIlo = np.fft.ifft(fdata * lo_window, axis=1)[:, :npixels]

        rEprobe = np.sqrt(np.abs(rIprobe))
        '''
        tag_none1 = tags[0][0]['none1']['shutter open']
        nnone1 = rEprobe[tag_none1].mean(axis=0)
        none1 = rIlo[tag_none1].mean(axis=0)
        tag_zero = tags[0][0]['zero']['shutter open']
        nzero = rEprobe[tag_zero].mean(axis=0)
        zero = rIlo[tag_zero].mean(axis=0)
        tag_none2 = tags[0][0]['none2']['shutter open']
        nnone2 = rEprobe[tag_none2].mean(axis=0)
        none2 = rIlo[tag_none2].mean(axis=0)
        tag_pipi = tags[0][0]['pipi']['shutter open']
        npipi = rEprobe[tag_pipi].mean(axis=0)
        pipi = rIlo[tag_pipi].mean(axis=0)
        '''
        tag_none1 = tags[0][0]['none1']['shutter open']
        nnone1 = rEprobe[tag_none1]
        none1 = rIlo[tag_none1]
        tag_zero = tags[0][0]['zero']['shutter open']
        nzero = rEprobe[tag_zero]
        zero = rIlo[tag_zero]
        tag_none2 = tags[0][0]['none2']['shutter open']
        nnone2 = rEprobe[tag_none2]
        none2 = rIlo[tag_none2]
        tag_pipi = tags[0][0]['pipi']['shutter open']
        npipi = rEprobe[tag_pipi]
        pipi = rIlo[tag_pipi]

        #TG = 1/(nzero + npipi)*(zero/nzero + pipi/npipi) \
        #        - (1/(nnone1 + nnone2)*(none1/nnone1 + none2/nnone2))
        #TG = ((zero + pipi - none1 - none2)/(nzero + npipi)).mean(axis=0)
        TG = ((zero + pipi - none1 - none2) / (nzero + npipi)).mean(axis=0)

        if all([plot, table == 0, t2 == 0]):
            plot_phasing_tg(f, TG)
            plot_phasing_tg(f, (nzero + npipi).mean(axis=0))

        with h5py.File(str(save_path), 'a') as sf:
            # save data at current t2
            sf['raw transient-grating'][t2] = TG

    with h5py.File(str(save_path), 'a') as sf:
        # write out meta data
        sf.attrs['batch_name'] = tg_info['batch_name']
        sf.attrs['batch_no'] = tg_info['batch_no']
        sf.attrs['batch_path'] = tg_info['batch_path']
        sf.attrs['job_name'] = tg_info['job_name']
        sf.attrs['nt2'] = tg_info['nt2']
        sf.attrs['when'] = tg_info['when']

        sf.attrs['detection axis pad to'] = pad_to
        sf.attrs['probe lo delay estimate'] = prd_est
        sf.attrs['analysis timestamp'] = arrow.now().format('DD-MM-YYYY HH:mm')
        sf.attrs['nudie version'] = nudie.version
        sf.attrs['experiment type'] = 'transient-grating'

        # write out axes
        gaxes = sf.require_group('axes')
        freq_dataset = gaxes.create_dataset('detection frequency', data=f)
        freq_dataset.attrs['df'] = df
        gaxes.create_dataset('detection wavelength',
                             data=nudie.spectrometer.speed_of_light / f)
        gaxes.create_dataset('t2', data=tg_info['t2'][:, 1])

        # add dimension scales and attach them
        rdata = sf['raw transient-grating']
        rdata.dims.create_scale(gaxes['t2'], 'population time / fs')
        rdata.dims[0].label = 'population time'
        rdata.dims[0].attach_scale(gaxes['t2'])
        rdata.dims.create_scale(gaxes['detection frequency'],
                                'frequency / 1000 THz')
        rdata.dims.create_scale(gaxes['detection wavelength'],
                                'wavelength / nm')
        rdata.dims[1].label = 'detection axis'
        rdata.dims[1].attach_scale(gaxes['detection frequency'])
        rdata.dims[1].attach_scale(gaxes['detection wavelength'])
Ejemplo n.º 13
0
def run(stark_name, stark_batch, when='today', wavelengths=None, plot=False, analysis_path='analyzed', datapath = None):
    '''run the batch job to analyze the linear-stark data
    
    Performs a shot-to-shot linear stark analysis. 
    
    ```wavelengths``` is a path to a wavelengths file
    ```plot``` is a Boolean parameter controlling whether to show matplotlib
    plots
    '''

    analysis_folder = Path(analysis_path)

    # create folder if it doesn't exist
    if not analysis_folder.exists():
        analysis_folder.mkdir()

    # change datapath
    if datapath is None:
        datapath = nudie.data_folder

    # load up pp data to use for linear stark
    stark_info = next(nudie.load_job(job_name=stark_name, batch_set=[stark_batch], when=when, data_path=datapath))

    # set current batch directory
    current_path = Path(stark_info['batch_path'])

    if wavelengths is not None:
        wl = load_wavelengths(current_path.parent / wavelengths)

    # tags for phase cycles
    phase_cycles = ['none1', 'zero', 'none2', 'pipi']
    nrepeat = 1 # how many times each waveform is repeated in the camera file. Assumed to be one
    nwaveforms = 1 # number of dazzler waveforms in pump probe is 1, assumed to be one

    # used to store data generated by iterations
    saved_data = deque()

    # loop over available data
    for t2,table,loop in it.product(stark_info['t2_range'], stark_info['table_range'], stark_info['loop_range']):
        # first file requires special synchronization
        # this is the rule that determines that it is the first file
        first = 'first' if all([t2 == 0, table == 0, loop == 0]) else False
        
        # Load everything for the given t2, table, and loop value
        analogs = nudie.load_analogtxt(stark_info['job_name'], current_path, t2, table, loop)
        cdata = nudie.load_camera_file(stark_info['job_name'], current_path, t2, table, loop, force_uint16=True)
        
        # Synchronize it, trimming some frames in the beginning and end
        data, (a1, a2) = nudie.trim_all(*nudie.synchronize_daq_to_camera(\
                        cdata, analog_channels=analogs, which_file=first), trim_to=slice(2, None))
        data = data.astype(float) # convert to float64 before manipulating it!
        
        start_idxs, period = nudie.detect_table_start(a1)

        # determine where the shutter is
        shutter_open, shutter_closed, duty_cycle = nudie.determine_shutter_shots(data)
        
        # tag the phases
        shutter_info = {'last open idx': shutter_open, 'first closed idx': shutter_closed}                
        tags = nudie.tag_phases(start_idxs, period, tags=phase_cycles, nframes=data.shape[1], shutter_info=shutter_info)
        
        # verify synchronization of applied field to Dazzler
        for i,k in enumerate(phase_cycles):
            # check that every second frame in analog channel 2 is either high field, or low field, but not both
            
            # Offset 0, every second frame has high field
            ao = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter open']][0::2] > 0.2
            # Offset 0, every second frame has low field
            bo = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter open']][0::2] <= 0.2 
            
            # Offset 1, every second frame has high field
            co = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter open']][1::2] > 0.2
            # Offset 1, every second frame has low field
            do = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter open']][1::2] <= 0.2
            
            # Offset 0, every second frame has high field
            ac = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter closed']][0::2] > 0.2
            # Offset 0, every second frame has low field
            bc = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter closed']][0::2] <= 0.2 
            
            # Offset 1, every second frame has high field
            cc = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter closed']][1::2] > 0.2
            # Offset 1, every second frame has low field
            dc = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter closed']][1::2] <= 0.2
            
            assert np.logical_xor(np.all(ao), np.all(bo)), 'detected stark synchronization error, offset 0, phase %s' % k
            assert np.logical_xor(np.all(co), np.all(do)), 'detected stark synchronization error, offset 1, phase %s' % k
            
            assert np.logical_xor(np.all(ac), np.all(bc)), 'detected stark synchronization error, offset 0, phase %s' % k
            assert np.logical_xor(np.all(cc), np.all(dc)), 'detected stark synchronization error, offset 1, phase %s' % k

        # determine stark-on indexes
        stark_idx, nostark_idx = np.where(a2 > 0.2)[0], np.where(a2 <= 0.2)[0]
            
        tmp = []
        # remove the probe scatter phase by phase
        for i,k in enumerate(phase_cycles):
            so = np.intersect1d(tags[nrepeat-1][nwaveforms-1][k]['shutter open'], stark_idx, assume_unique=True)
            sc = np.intersect1d(tags[nrepeat-1][nwaveforms-1][k]['shutter closed'], stark_idx, assume_unique=True)                   

            no = np.intersect1d(tags[nrepeat-1][nwaveforms-1][k]['shutter open'], nostark_idx, assume_unique=True)
            nc = np.intersect1d(tags[nrepeat-1][nwaveforms-1][k]['shutter closed'], nostark_idx, assume_unique=True)                   
                        
            # remove probe scatter for stark shots
            data[:, so] -= data[:, sc].mean(axis=-1)[:, np.newaxis]
            # and for no-stark shots
            data[:, no] -= data[:, nc].mean(axis=-1)[:, np.newaxis]        
            
            data = ma.masked_less(data, 1e-7*data.max())
            stark_spec = np.log10(np.mean(data[:, so], axis=-1)/np.mean(data[:, no], axis=-1))
            stark_spec = ma.filled(stark_spec, 0)
            tmp.append(stark_spec)
        
        # average phase cycles together to get slight improvement in S/N
        px = max([len(i) for i in tmp]) # num pixels
        m = np.zeros((px, len(phase_cycles)))
        for i,k in enumerate(tmp):
            m[:, i] = k
        avg_spectrum = m.mean(axis=1)
                       
        mean_voltage = a2[stark_idx].mean()

        saved_data.append((avg_spectrum, "loop {}".format(loop)))

        # interpret `when` parameter
        if when == 'today':
            timestamp = arrow.now()
        elif when is not None:
            # try to interpret when as a string
            timestamp = arrow.get(when, 'YY-MM-DD')
        else:
            raise NotImplemented('don\'t pass a direct path to the ' +\
                    'filename please')

        # Write out information generated
        filename = '{timestamp:s}-{name:s}-{voltage:1.1f}kV.hdf5'.format(
                timestamp=timestamp.format('YYYY-MM-DD'),
                name=stark_info['batch_name'], voltage=mean_voltage)

        with h5py.File(str(analysis_folder / filename), mode='w') as f:
            for data, j in saved_data:
                f[j] = data
            
            f.attrs['applied voltage'] = mean_voltage
            f.attrs['file path'] = stark_info['batch_path']
            f.attrs['job name'] = stark_info['job_name']
            f.attrs['batch_no'] = stark_batch # FIXME: workaround for nudie bug

            if wavelengths is not None:
                f['wavelength axis'] = wl

            if all([plot, loop == stark_info['loop_range'].stop-1]):
                for data, loop in saved_data:
                    pyplot.plot(wl, data, label=loop)

                pyplot.xlabel('wavelength / nm')
                pyplot.ylabel('OD')
                pyplot.show()
Ejemplo n.º 14
0
def run(tg_name,
        tg_batch,
        when='today',
        wavelengths=None,
        plot=False,
        pad_to=2048,
        prd_est=850.,
        lo_width=200,
        dc_width=200,
        gaussian_power=2.,
        analysis_path='./analyzed',
        min_field=0.2,
        datapath=None):

    if plot:
        import matplotlib as mpl
        import matplotlib.pyplot as plt
        mpl.rcParams['figure.figsize'] = (16, 12)
        mpl.use('Qt4Agg')

    # change datapath
    if datapath is None:
        datapath = nudie.data_folder

    nrepeat = 1  # how many times each waveform is repeated in the camera file. Assumed to be one
    waveforms_per_table = 1
    npixels = 1340
    trim_to = 3, -3
    nstark = 2  # number of stark items

    # load up tg data to use
    tg_info = next(
        nudie.load_job(job_name=tg_name,
                       batch_set=[tg_batch],
                       when=when,
                       data_path=datapath))

    phase_cycles = ['none1', 'zero', 'none2', 'pipi']

    # set current batch directory
    current_path = Path(tg_info['batch_path'])

    # generate hdf filename based on data date
    analysis_folder = Path(analysis_path)

    # create folder if it doesn't exist
    if not analysis_folder.exists():
        analysis_folder.mkdir(parents=True)

    save_path = analysis_folder / (tg_info['batch_name'] + '.h5')

    # remove data file if it exists
    if save_path.exists(): save_path.unlink()

    with h5py.File(str(save_path), 'w') as sf:
        # initialize groups
        sf.create_group('axes')
        sf.attrs['nstark'] = nstark
        shape = (tg_info['nt2'], nstark, npixels)
        sf.create_dataset('raw transient-grating', shape, dtype=complex)

    try:
        wl = load_wavelengths(current_path.parent / wavelengths)
    except FileNotFoundError as e:
        nudie.log.error('Could not load wavelength calibration file!')
        raise e

    # define gaussian
    def gaussian2(w, x0, x):
        c = 4 * np.log(2)
        ξ = x - x0
        return np.exp(-c * (ξ / w)**(2 * gaussian_power))

    for loop, t2, table in it.product(tg_info['loop_range'],
                                      tg_info['t2_range'],
                                      tg_info['table_range']):

        if any([loop != 0, table != 0]):
            raise NotImplementedError('Code is not setup to handle multiple ' +\
                    'loops or tables.')

        # first file requires special synchronization
        # this is the rule that determines that it is the first file
        first = 'first' if all([t2 == 0, table == 0, loop == 0]) else False

        # Load everything for the given t2, table, and loop value
        analogs = nudie.load_analogtxt(tg_info['job_name'], current_path, t2,
                                       table, loop)
        cdata = nudie.load_camera_file(tg_info['job_name'],
                                       current_path,
                                       t2,
                                       table,
                                       loop,
                                       force_uint16=True)

        # Synchronize it, trimming some frames in the beginning and end
        data, (a1, a2) = nudie.trim_all(*nudie.synchronize_daq_to_camera(\
            cdata, analog_channels=analogs, which_file=first),
            trim_to=slice(*trim_to))
        data = data.astype(float)  # convert to float64 before manipulating it!

        start_idxs, period = nudie.detect_table_start(a1)

        f, data, df = nudie.wavelen_to_freq(wl, data, ret_df=True, ax=0)

        # determine where the shutter is
        shutter_open, shutter_closed, duty_cycle = nudie.determine_shutter_shots(
            data)

        # tag the phases
        shutter_info = {
            'last open idx': shutter_open,
            'first closed idx': shutter_closed
        }
        tags = nudie.tag_phases(start_idxs,
                                period,
                                tags=phase_cycles,
                                nframes=data.shape[1],
                                shutter_info=shutter_info)
        nudie.remove_incomplete_t1_waveforms(tags, phase_cycles)
        data_t = np.zeros((data.shape[1], data.shape[0]), dtype=float)

        # verify synchronization of applied field to Dazzler
        for t1, k in it.product(range(waveforms_per_table), phase_cycles):
            # check that every second frame in analog channel 2 is either high field, or low field, but not both
            idx_open = tags[nrepeat - 1][t1][k]['shutter open']
            idx_closed = tags[nrepeat - 1][t1][k]['shutter closed']

            # Offset 0, every second frame has high field
            ao = a2[idx_open][0::2] > min_field
            # Offset 0, every second frame has low field
            bo = a2[idx_open][0::2] <= min_field

            # Offset 1, every second frame has high field
            co = a2[idx_open][1::2] > min_field
            # Offset 1, every second frame has low field
            do = a2[idx_open][1::2] <= min_field

            # Offset 0, every second frame has high field
            ac = a2[idx_closed][0::2] > min_field
            # Offset 0, every second frame has low field
            bc = a2[idx_closed][0::2] <= min_field

            # Offset 1, every second frame has high field
            cc = a2[idx_closed][1::2] > min_field
            # Offset 1, every second frame has low field
            dc = a2[idx_closed][1::2] <= min_field

            assert np.logical_xor(np.all(ao), np.all(bo)), \
                'detected stark synchronization error, offset 0, phase %s' % str(k)
            assert np.logical_xor(np.all(co), np.all(do)), \
                'detected stark synchronization error, offset 1, phase %s' % str(k)

            assert np.logical_xor(np.all(ac), np.all(bc)), \
                'detected stark synchronization error, offset 0, phase %s' % str(k)
            assert np.logical_xor(np.all(cc), np.all(dc)), \
                'detected stark synchronization error, offset 1, phase %s' % str(k)

            assert np.logical_xor(np.all(ao), np.all(co)), \
                'detected stark synchronization error, offset 0, phase %s' % str(k)
            assert np.logical_xor(np.all(bo), np.all(do)), \
                'detected stark synchronization error, offset 1, phase %s' % str(k)

        # determine stark-on indexes
        stark_idx = np.where(a2 > min_field)[0]
        nostark_idx = np.where(a2 <= min_field)[0]

        # subtract the average of closed shutter shots from shutter open data
        for t1, k in it.product(range(waveforms_per_table), phase_cycles):
            idx_open = tags[nrepeat - 1][t1][k]['shutter open']
            idx_closed = tags[nrepeat - 1][t1][k]['shutter closed']

            so = np.intersect1d(idx_open, stark_idx, assume_unique=True)
            sc = np.intersect1d(idx_closed, stark_idx, assume_unique=True)

            no = np.intersect1d(idx_open, nostark_idx, assume_unique=True)
            nc = np.intersect1d(idx_closed, nostark_idx, assume_unique=True)

            data_t[so, :] = (data[:, so].T - data[:, sc].mean(axis=1))
            data_t[no, :] = (data[:, no].T - data[:, nc].mean(axis=1))

        t = np.fft.fftfreq(pad_to, df)
        lo_window = gaussian2(lo_width, prd_est, t)
        dc_window = gaussian2(dc_width, 0, t)

        # fft everything
        fdata = np.fft.fft(data_t, axis=1, n=pad_to)

        if all([plot, table == 0, t2 == 0]):
            tmp_fft = fdata[tags[0][0][phase_cycles[1]]['shutter open'][0], :]
            plot_windows(t, lo_window, dc_window, tmp_fft, prd_est)

        # spectral interferometry
        rIprobe = np.fft.ifft(fdata * dc_window, axis=1)[:, :npixels]
        rIlo = np.fft.ifft(fdata * lo_window, axis=1)[:, :npixels]

        rEprobe = np.sqrt(np.abs(rIprobe))
        '''
        tag_none1 = tags[0][0]['none1']['shutter open']
        nnone1 = rEprobe[tag_none1].mean(axis=0)
        none1 = rIlo[tag_none1].mean(axis=0)
        tag_zero = tags[0][0]['zero']['shutter open']
        nzero = rEprobe[tag_zero].mean(axis=0)
        zero = rIlo[tag_zero].mean(axis=0)
        tag_none2 = tags[0][0]['none2']['shutter open']
        nnone2 = rEprobe[tag_none2].mean(axis=0)
        none2 = rIlo[tag_none2].mean(axis=0)
        tag_pipi = tags[0][0]['pipi']['shutter open']
        npipi = rEprobe[tag_pipi].mean(axis=0)
        pipi = rIlo[tag_pipi].mean(axis=0)
        '''
        stark_tag = lambda x: np.intersect1d(x, stark_idx, assume_unique=True)
        nostark_tag = lambda x: np.intersect1d(
            x, nostark_idx, assume_unique=True)

        tag_none1 = tags[0][0]['none1']['shutter open']
        stag_none1 = stark_tag(tag_none1)
        ntag_none1 = nostark_tag(tag_none1)

        #snone1 = rEprobe[stag_none1]
        #nnone1 = rEprobe[ntag_none1]
        snone1 = rIlo[stag_none1].mean(axis=0)
        nnone1 = rIlo[ntag_none1].mean(axis=0)

        tag_zero = tags[0][0]['zero']['shutter open']
        stag_zero = stark_tag(tag_zero)
        ntag_zero = nostark_tag(tag_zero)
        snzero = rEprobe[stag_zero].mean(axis=0)
        nnzero = rEprobe[ntag_zero].mean(axis=0)
        szero = rIlo[stag_zero].mean(axis=0)
        nzero = rIlo[ntag_zero].mean(axis=0)

        tag_none2 = tags[0][0]['none2']['shutter open']
        stag_none2 = stark_tag(tag_none2)
        ntag_none2 = nostark_tag(tag_none2)

        #nnone2 = rEprobe[tag_none2]
        snone2 = rIlo[stag_none2].mean(axis=0)
        nnone2 = rIlo[ntag_none2].mean(axis=0)

        tag_pipi = tags[0][0]['pipi']['shutter open']
        stag_pipi = stark_tag(tag_pipi)
        ntag_pipi = nostark_tag(tag_pipi)
        snpipi = rEprobe[stag_pipi].mean(axis=0)
        nnpipi = rEprobe[ntag_pipi].mean(axis=0)
        spipi = rIlo[stag_pipi].mean(axis=0)
        npipi = rIlo[ntag_pipi].mean(axis=0)

        #TG = 1/(nzero + npipi)*(zero/nzero + pipi/npipi) \
        #        - (1/(nnone1 + nnone2)*(none1/nnone1 + none2/nnone2))
        #TG = ((zero + pipi - none1 - none2)/(nzero + npipi)).mean(axis=0)

        stark_axis = np.array([a2[stark_idx].mean(), a2[nostark_idx].mean()])
        #TG = ((zero + pipi - none1 - none2)/(nzero + npipi))
        sTG = ((szero + spipi - snone1 - snone2) / (snzero + snpipi))
        nTG = ((nzero + npipi - nnone1 - nnone2) / (nnzero + nnpipi))
        '''
        if all([plot, table==0, t2==0]):
            plot_phasing_tg(f, TG)
            plot_phasing_tg(f, (nzero + npipi).mean(axis=0))
        '''

        with h5py.File(str(save_path), 'a') as sf:
            # save data at current t2

            # for some reason, I need -1j now -> 1j to make the real part the
            # absorptive and -1 to do the 2D convention
            #sf['raw transient-grating'][t2, 0] = -1j*sTG
            #sf['raw transient-grating'][t2, 1] = -1j*nTG
            sf['raw transient-grating'][t2, 0] = sTG
            sf['raw transient-grating'][t2, 1] = nTG

    with h5py.File(str(save_path), 'a') as sf:
        # write out meta data
        sf.attrs['batch_name'] = tg_info['batch_name']
        sf.attrs['batch_no'] = tg_info['batch_no']
        sf.attrs['batch_path'] = tg_info['batch_path']
        sf.attrs['job_name'] = tg_info['job_name']
        sf.attrs['nt2'] = tg_info['nt2']
        sf.attrs['when'] = tg_info['when']

        sf.attrs['detection axis pad to'] = pad_to
        sf.attrs['probe lo delay estimate'] = prd_est
        sf.attrs['analysis timestamp'] = arrow.now().format('DD-MM-YYYY HH:mm')
        sf.attrs['nudie version'] = nudie.version
        sf.attrs['experiment type'] = 'stark transient-grating'

        # write out axes
        gaxes = sf.require_group('axes')
        gaxes.create_dataset('stark axis', data=stark_axis)
        freq_dataset = gaxes.create_dataset('detection frequency', data=f)
        freq_dataset.attrs['df'] = df
        gaxes.create_dataset('detection wavelength',
                             data=nudie.spectrometer.speed_of_light / f)
        gaxes.create_dataset('t2', data=tg_info['t2'][:, 1])

        # add dimension scales and attach them
        rdata = sf['raw transient-grating']
        rdata.dims.create_scale(gaxes['t2'], 'population time / fs')
        rdata.dims[0].label = 'population time'
        rdata.dims[0].attach_scale(gaxes['t2'])
        rdata.dims.create_scale(gaxes['stark axis'], 'average field / kV')
        rdata.dims[1].label = 'average voltage'
        rdata.dims[1].attach_scale(gaxes['stark axis'])
        rdata.dims.create_scale(gaxes['detection frequency'],
                                'frequency / 1000 THz')
        rdata.dims.create_scale(gaxes['detection wavelength'],
                                'wavelength / nm')
        rdata.dims[2].label = 'detection axis'
        rdata.dims[2].attach_scale(gaxes['detection frequency'])
        rdata.dims[2].attach_scale(gaxes['detection wavelength'])
Ejemplo n.º 15
0
def run(dd_name,
        dd_batch,
        when='today',
        wavelengths=None,
        plot=False,
        pump_chop=False,
        central_wl=None,
        phaselock_wl=None,
        pad_to=2048,
        waveforms_per_table=40,
        prd_est=850.,
        lo_width=200,
        dc_width=200,
        gaussian_power=2.,
        analysis_path='./analyzed',
        detrend_t1=False,
        datapath=None):

    if plot:
        import matplotlib as mpl
        import matplotlib.pyplot as plt
        mpl.rcParams['figure.figsize'] = (16, 12)
        mpl.use('Qt5Agg')

    nrepeat = 1  # how many times each waveform is repeated in the camera file. Assumed to be one
    npixels = 1340
    trim_to = 3, -3

    # change datapath
    if datapath is None:
        datapath = nudie.data_folder

    # load up 2d data to use
    dd_info = next(
        nudie.load_job(job_name=dd_name,
                       batch_set=[dd_batch],
                       when=when,
                       data_path=datapath))

    if pump_chop == True:
        phase_cycles = [(0., 0.), (1., 1.), (0, 0.6), (1., 1.6), (0, 1.3),
                        (1., 2.3), 'chop']
    else:
        phase_cycles = [(0., 0.), (1., 1.), (0, 0.6), (1., 1.6), (0, 1.3),
                        (1., 2.3)]

    # set current batch directory
    current_path = Path(dd_info['batch_path'])

    # generate hdf filename based on data date
    analysis_folder = Path(analysis_path)

    # create folder if it doesn't exist
    if not analysis_folder.exists():
        analysis_folder.mkdir(parents=True)

    save_path = analysis_folder / (dd_info['batch_name'] + '.h5')

    # remove data file if it exists
    if save_path.exists(): save_path.unlink()

    with h5py.File(str(save_path), 'w') as sf:
        # initialize groups
        sf.create_group('axes')

        shape = (dd_info['nt2'], dd_info['nt1'], npixels)
        sf.create_dataset('raw rephasing', shape, dtype=complex)
        sf.create_dataset('raw non-rephasing', shape, dtype=complex)
        sf.create_dataset('raw transient-grating', shape, dtype=complex)

    try:
        wl = load_wavelengths(current_path.parent / wavelengths)
    except FileNotFoundError as e:
        nudie.log.error('Could not load wavelength calibration file!')
        raise e

    # define gaussian
    def gaussian2(w, x0, x):
        c = 4 * np.log(2)
        ξ = x - x0
        return np.exp(-c * (ξ / w)**(2 * gaussian_power))

    ## Make phase-cycling coefficient matrix
    # subspaces
    if pump_chop == True:
        sub1 = np.array(
            [[np.exp(1j * np.pi * (x - y)),
              np.exp(-1j * np.pi * (x - y)), 1]
             for x, y in phase_cycles[0:-1:2]])
        sub2 = np.array(
            [[np.exp(1j * np.pi * (x - y)),
              np.exp(-1j * np.pi * (x - y)), 1]
             for x, y in phase_cycles[1:-1:2]])
    else:
        sub1 = np.array(
            [[np.exp(1j * np.pi * (x - y)),
              np.exp(-1j * np.pi * (x - y)), 1]
             for x, y in phase_cycles[0::2]])
        sub2 = np.array(
            [[np.exp(1j * np.pi * (x - y)),
              np.exp(-1j * np.pi * (x - y)), 1]
             for x, y in phase_cycles[1::2]])

    perm = np.array([[1, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0,
                                          0], [0, 1, 0, 0, 0, 0],
                     [0, 0, 0, 0, 1, 0], [0, 0, 1, 0, 0, 0],
                     [0, 0, 0, 0, 0, 1]])

    C = np.kron(np.array([[1, 0], [0, 0]]), sub1) + np.kron(
        np.array([[0, 0], [0, 1]]), sub2)
    Cinv = np.linalg.inv(C)
    Cperm_inv = np.dot(Cinv, perm.T)

    for loop, t2, table in it.product(dd_info['loop_range'],
                                      dd_info['t2_range'],
                                      dd_info['table_range']):

        if loop != 0:
            raise NotImplementedError(
                'Code is not setup to handle multiple loops')

        # first file requires special synchronization
        # this is the rule that determines that it is the first file
        first = 'first' if all([t2 == 0, table == 0, loop == 0]) else False

        # Load everything for the given t2, table, and loop value
        analogs = nudie.load_analogtxt(dd_info['job_name'], current_path, t2,
                                       table, loop)
        cdata = nudie.load_camera_file(dd_info['job_name'],
                                       current_path,
                                       t2,
                                       table,
                                       loop,
                                       force_uint16=True)

        # Synchronize it, trimming some frames in the beginning and end
        data, (a1, a2) = nudie.trim_all(*nudie.synchronize_daq_to_camera(\
            cdata, analog_channels=analogs, which_file=first),
            trim_to=slice(*trim_to))
        data = data.astype(float)  # convert to float64 before manipulating it!

        start_idxs, period = nudie.detect_table_start(a1)

        f, data, df = nudie.wavelen_to_freq(wl, data, ret_df=True, ax=0)

        # determine where the shutter is
        shutter_open, shutter_closed, duty_cycle = nudie.determine_shutter_shots(
            data)

        # tag the phases
        shutter_info = {
            'last open idx': shutter_open,
            'first closed idx': shutter_closed
        }
        tags = nudie.tag_phases(start_idxs,
                                period,
                                tags=phase_cycles,
                                nframes=data.shape[1],
                                shutter_info=shutter_info)
        nudie.remove_incomplete_t1_waveforms(tags, phase_cycles)
        data_t = np.zeros((data.shape[1], data.shape[0]), dtype=float)

        # subtract the average of closed shutter shots from shutter open data
        for t1, k in it.product(range(waveforms_per_table), phase_cycles):
            idx_open = tags[nrepeat - 1][t1][k]['shutter open']
            idx_closed = tags[nrepeat - 1][t1][k]['shutter closed']

            data_t[idx_open, :] = data[:, idx_open].T \
                                  - data[:, idx_closed].mean(axis=1)

        t = np.fft.fftfreq(pad_to, df)
        lo_window = gaussian2(lo_width, prd_est, t)
        dc_window = gaussian2(dc_width, 0, t)

        # fft everything
        fdata = np.fft.fft(data_t, axis=1, n=pad_to)

        if all([plot, table == 0, t2 == 0]):
            tmp_fft = fdata[tags[0][0][phase_cycles[0]]['shutter open'][0], :]
            plot_windows(t, lo_window, dc_window, tmp_fft, prd_est)

        # spectral interferometry
        rIprobe = np.fft.ifft(fdata * dc_window, axis=1)[:, :npixels]
        rIlo = np.fft.ifft(fdata * lo_window, axis=1)[:, :npixels]

        if pump_chop:
            nudie.log.warning('pump chopping code is not well tested and will' +\
                    ' probably result in wrong values!')
            nudie.log.warning('It currently doesn\'t mask invalid reference ' +\
                    'values')
            # Pump chop before division
            rEsig = np.zeros_like(rIlo)
            for t1, k in it.product(range(waveforms_per_table),
                                    phase_cycles[:-1]):
                phase_idx = tags[nrepeat - 1][t1][k]['shutter open']
                chop_idx = tags[nrepeat - 1][t1]['chop']['shutter open']

                # FIXME: not sure about sign. in rEsig, also whether to take
                # the square of rEprobe
                rEprobe = np.sqrt(
                    np.abs(0.5 * rIprobe[phase_idx, :] +
                           0.5 * rIprobe[chop_idx, :]))
                rEsig[phase_idx, :] = -(rIlo[phase_idx, :] -
                                        rIlo[chop_idx, :]) / rEprobe
        else:
            rEprobe = np.sqrt(np.abs(rIprobe))
            masked_probe = ma.masked_less(rEprobe, 1e-7 * rEprobe.max())
            rEsig = ma.filled(rIlo / masked_probe, fill_value=0)

        # average each phase together
        if pump_chop:
            avg = np.zeros(
                (len(phase_cycles[:-1]), waveforms_per_table, npixels),
                dtype=complex)
            for t1, (ik, k) in it.product(range(waveforms_per_table),
                                          enumerate(phase_cycles[:-1])):
                avg[ik,
                    t1, :] = rEsig[tags[0][t1][k]['shutter open']].mean(axis=0)
        else:
            avg = np.zeros((len(phase_cycles), waveforms_per_table, npixels),
                           dtype=complex)
            for t1, (ik, k) in it.product(range(waveforms_per_table),
                                          enumerate(phase_cycles)):
                avg[ik,
                    t1, :] = rEsig[tags[0][t1][k]['shutter open']].mean(axis=0)

        r1, nr1, tg1, r2, nr2, tg2 = np.tensordot(Cperm_inv,
                                                  avg,
                                                  axes=[[
                                                      1,
                                                  ], [
                                                      0,
                                                  ]])
        R = 0.5 * (r1 + r2)
        NR = 0.5 * (nr1 + nr2)
        TG = 0.5 * (tg1 + tg2)

        if all([plot, table == 0, t2 == 0]):
            phasing_tg = 0.5 * (R[0] + NR[0])
            plot_phasing_tg(f, phasing_tg)
            #plot_phasing_tg(f, masked_probe)

        with h5py.File(str(save_path), 'a') as sf:
            # save data at current t2
            t1_slice = slice(table * waveforms_per_table,
                             (table + 1) * waveforms_per_table)
            sf['raw rephasing'][t2, t1_slice] = R
            sf['raw non-rephasing'][t2, t1_slice] = NR
            sf['raw transient-grating'][t2, t1_slice] = TG

        del data, fdata, data_t, avg, rIprobe, rIlo, rEprobe, rEsig

    with h5py.File(str(save_path), 'a') as sf:
        if detrend_t1:
            R = sf['raw rephasing'][:]
            NR = sf['raw non-rephasing'][:]
            TG = sf['raw transient-grating'][:]

            R = detrend(R, axis=1, type='constant')
            NR = detrend(NR, axis=1, type='constant')
            TG = detrend(TG, axis=1, type='constant')

            sf['raw rephasing'][:] = R
            sf['raw non-rephasing'][:] = NR
            sf['raw transient-grating'][:] = TG

        # write out meta data
        sf.attrs['batch_name'] = dd_info['batch_name']
        sf.attrs['batch_no'] = dd_info['batch_no']
        sf.attrs['batch_path'] = dd_info['batch_path']
        sf.attrs['job_name'] = dd_info['job_name']
        sf.attrs['nt1'] = dd_info['nt1']
        sf.attrs['nt2'] = dd_info['nt2']
        sf.attrs['when'] = dd_info['when']

        sf.attrs['detection axis pad to'] = pad_to
        sf.attrs['probe lo delay estimate'] = prd_est
        sf.attrs['analysis timestamp'] = arrow.now().format('DD-MM-YYYY HH:mm')
        sf.attrs['nudie version'] = nudie.version
        sf.attrs['experiment type'] = '2d'

        # write out axes
        gaxes = sf.require_group('axes')
        freq_dataset = gaxes.create_dataset('detection frequency', data=f)
        freq_dataset.attrs['df'] = df
        gaxes.create_dataset('detection wavelength',
                             data=nudie.spectrometer.speed_of_light / f)
        gaxes.create_dataset('t1', data=dd_info['t1'])
        gaxes.create_dataset('t2', data=dd_info['t2'][:, 1])

        # add dimension scales
        rdata = sf['raw rephasing']
        rdata.dims.create_scale(gaxes['t2'], 'population time / fs')
        rdata.dims.create_scale(gaxes['t1'], 'pump delay time / fs')
        rdata.dims.create_scale(gaxes['detection frequency'],
                                'frequency / 1000 THz')
        rdata.dims.create_scale(gaxes['detection wavelength'],
                                'wavelength / nm')

        # attach them
        for tmp in [
                sf['raw rephasing'], sf['raw non-rephasing'],
                sf['raw transient-grating']
        ]:
            tmp.dims[0].label = 'population time'
            tmp.dims[0].attach_scale(gaxes['t2'])
            tmp.dims[1].label = 'pump delay time'
            tmp.dims[1].attach_scale(gaxes['t1'])
            tmp.dims[2].label = 'detection axis'
            tmp.dims[2].attach_scale(gaxes['detection frequency'])
            tmp.dims[2].attach_scale(gaxes['detection wavelength'])
Ejemplo n.º 16
0
def run(dd_name,
        dd_batch,
        when='today',
        wavelengths=None,
        plot=False,
        central_wl=None,
        phaselock_wl=None,
        pad_to=2048,
        waveforms_per_table=40,
        prd_est=850.,
        lo_width=200,
        dc_width=200,
        gaussian_power=2,
        analysis_path='./analyzed',
        min_field=0.2,
        detrend_t1=False,
        datapath=None):

    nrepeat = 1  # how many times each waveform is repeated in the camera file. Assumed to be one
    nstark = 2
    npixels = 1340
    trim_to = 3, -3

    # change datapath
    if datapath is None:
        datapath = nudie.data_folder

    # load up 2d data to use
    dd_info = next(
        nudie.load_job(job_name=dd_name,
                       batch_set=[dd_batch],
                       when=when,
                       data_path=datapath))
    phase_cycles = [(0., 0.), (1., 1.), (0, 0.6), (1., 1.6), (0, 1.3),
                    (1., 2.3)]

    # set current batch directory
    current_path = Path(dd_info['batch_path'])

    # generate hdf filename based on data date
    analysis_folder = Path(analysis_path)

    # create folder if it doesn't exist
    if not analysis_folder.exists():
        analysis_folder.mkdir()

    save_path = analysis_folder / (dd_info['batch_name'] + '.h5')

    # remove data file if it exists
    if save_path.exists(): save_path.unlink()

    with h5py.File(str(save_path), 'w') as sf:
        # initialize groups
        sf.create_group('axes')
        sf.attrs['nstark'] = nstark
        shape = (dd_info['nt2'], nstark, dd_info['nt1'], npixels)
        sf.create_dataset('raw rephasing', shape, dtype=complex)
        sf.create_dataset('raw non-rephasing', shape, dtype=complex)
        sf.create_dataset('raw transient-grating', shape, dtype=complex)

    try:
        wl = load_wavelengths(current_path.parent / wavelengths)
    except FileNotFoundError as e:
        nudie.log.error('Could not load wavelength calibration file!')
        raise e

    def gaussian2(w, x0, x):
        c = 4 * np.log(2)
        ξ = x - x0
        return np.exp(-c * (ξ / w)**(2 * gaussian_power))

    ## Make phase-cycling coefficient matrix
    # subspaces
    sub1 = np.array(
        [[np.exp(1j * np.pi * (x - y)),
          np.exp(-1j * np.pi * (x - y)), 1] for x, y in phase_cycles[0::2]])
    sub2 = np.array(
        [[np.exp(1j * np.pi * (x - y)),
          np.exp(-1j * np.pi * (x - y)), 1] for x, y in phase_cycles[1::2]])

    perm = np.array([[1, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0,
                                          0], [0, 1, 0, 0, 0, 0],
                     [0, 0, 0, 0, 1, 0], [0, 0, 1, 0, 0, 0],
                     [0, 0, 0, 0, 0, 1]])

    C = np.kron(np.array([[1, 0], [0, 0]]), sub1) + np.kron(
        np.array([[0, 0], [0, 1]]), sub2)
    Cinv = np.linalg.inv(C)
    Cperm_inv = np.dot(Cinv, perm.T)

    for loop, t2, table in it.product(dd_info['loop_range'],
                                      dd_info['t2_range'],
                                      dd_info['table_range']):

        if loop != 0:
            raise NotImplemented('code is not setup to handle multiple loops')

        # first file requires special synchronization
        # this is the rule that determines that it is the first file
        first = 'first' if all([t2 == 0, table == 0, loop == 0]) else False

        # Load everything for the given t2, table, and loop value
        analogs = nudie.load_analogtxt(dd_info['job_name'], current_path, t2,
                                       table, loop)
        cdata = nudie.load_camera_file(dd_info['job_name'],
                                       current_path,
                                       t2,
                                       table,
                                       loop,
                                       force_uint16=True)

        # Synchronize it, trimming some frames in the beginning and end
        data, (a1, a2) = nudie.trim_all(*nudie.synchronize_daq_to_camera(\
            cdata, analog_channels=analogs, which_file=first),
            trim_to=slice(*trim_to))
        data = data.astype(float)  # convert to float64 before manipulating it!

        start_idxs, period = nudie.detect_table_start(a1)

        # data comes out [npixels, nframes]
        f, data, df = nudie.wavelen_to_freq(wl, data, ret_df=True, ax=0)

        # determine where the shutter is
        shutter_open, shutter_closed, duty_cycle = nudie.determine_shutter_shots(
            data)

        # tag the phases
        shutter_info = {
            'last open idx': shutter_open,
            'first closed idx': shutter_closed
        }
        tags = nudie.tag_phases(start_idxs,
                                period,
                                tags=phase_cycles,
                                nframes=data.shape[1],
                                shutter_info=shutter_info)

        # data_t is transposed data, [nframes, npixels]
        data_t = np.zeros((data.shape[1], data.shape[0]), dtype=float)

        # verify synchronization of applied field to Dazzler
        for t1, k in it.product(range(waveforms_per_table), phase_cycles):
            # check that every second frame in analog channel 2 is either high field, or low field, but not both
            idx_open = tags[nrepeat - 1][t1][k]['shutter open']
            idx_closed = tags[nrepeat - 1][t1][k]['shutter closed']

            # Offset 0, every second frame has high field
            ao = a2[idx_open][0::2] > min_field
            # Offset 0, every second frame has low field
            bo = a2[idx_open][0::2] <= min_field

            # Offset 1, every second frame has high field
            co = a2[idx_open][1::2] > min_field
            # Offset 1, every second frame has low field
            do = a2[idx_open][1::2] <= min_field

            # Offset 0, every second frame has high field
            ac = a2[idx_closed][0::2] > min_field
            # Offset 0, every second frame has low field
            bc = a2[idx_closed][0::2] <= min_field

            # Offset 1, every second frame has high field
            cc = a2[idx_closed][1::2] > min_field
            # Offset 1, every second frame has low field
            dc = a2[idx_closed][1::2] <= min_field

            assert np.logical_xor(np.all(ao), np.all(bo)), \
                'detected stark synchronization error, offset 0, phase %s' % str(k)
            assert np.logical_xor(np.all(co), np.all(do)), \
                'detected stark synchronization error, offset 1, phase %s' % str(k)

            assert np.logical_xor(np.all(ac), np.all(bc)), \
                'detected stark synchronization error, offset 0, phase %s' % str(k)
            assert np.logical_xor(np.all(cc), np.all(dc)), \
                'detected stark synchronization error, offset 1, phase %s' % str(k)

            assert np.logical_xor(np.all(ao), np.all(co)), \
                'detected stark synchronization error, offset 0, phase %s' % str(k)
            assert np.logical_xor(np.all(bo), np.all(do)), \
                'detected stark synchronization error, offset 1, phase %s' % str(k)

        # determine stark-on indexes
        stark_idx = np.where(a2 > min_field)[0]
        nostark_idx = np.where(a2 <= min_field)[0]

        for t1, k in it.product(range(waveforms_per_table), phase_cycles):
            idx_open = tags[nrepeat - 1][t1][k]['shutter open']
            idx_closed = tags[nrepeat - 1][t1][k]['shutter closed']

            so = np.intersect1d(idx_open, stark_idx, assume_unique=True)
            sc = np.intersect1d(idx_closed, stark_idx, assume_unique=True)

            no = np.intersect1d(idx_open, nostark_idx, assume_unique=True)
            nc = np.intersect1d(idx_closed, nostark_idx, assume_unique=True)

            data_t[so, :] = (data[:, so].T - data[:, sc].mean(axis=1))
            data_t[no, :] = (data[:, no].T - data[:, nc].mean(axis=1))

        t = np.fft.fftfreq(pad_to, df)
        lo_window = gaussian2(lo_width, prd_est, t)
        dc_window = gaussian2(dc_width, 0, t)

        # fft everything
        fdata = np.fft.fft(data_t, axis=1, n=pad_to)

        if all([plot, table == 0, t2 == 0]):
            tmp_fft = fdata[tags[0][0][phase_cycles[0]]['shutter open'][0], :]
            plot_windows(t, lo_window, dc_window, tmp_fft, prd_est)

        # spectral interferometry
        rIprobe = np.fft.ifft(fdata * dc_window)[:, :npixels]
        rIlo = np.fft.ifft(fdata * lo_window)[:, :npixels]
        rEprobe = np.sqrt(np.abs(rIprobe))

        rEsig = rIlo / rEprobe

        # average each phase together
        stark_axis = np.array([a2[stark_idx].mean(), a2[nostark_idx].mean()])
        avg = np.zeros(
            (len(phase_cycles), nstark, waveforms_per_table, npixels),
            dtype=complex)
        for t1, (ik, k) in it.product(range(waveforms_per_table),
                                      enumerate(phase_cycles)):
            # create indexes for stark shots
            idx_open = tags[nrepeat - 1][t1][k]['shutter open']

            so = np.intersect1d(idx_open, stark_idx, assume_unique=True)
            no = np.intersect1d(idx_open, nostark_idx, assume_unique=True)

            avg[ik, 0, t1, :] = rEsig[so].mean(axis=0)
            avg[ik, 1, t1, :] = rEsig[no].mean(axis=0)

        r1, nr1, tg1, r2, nr2, tg2 = np.tensordot(Cperm_inv,
                                                  avg,
                                                  axes=[[
                                                      1,
                                                  ], [
                                                      0,
                                                  ]])
        R = 0.5 * (r1 + r2)
        NR = 0.5 * (nr1 + nr2)
        TG = 0.5 * (tg1 + tg2)

        if all([plot, table == 0, t2 == 0]):
            # plot the no-stark TG
            phasing_tg = 0.5 * (R + NR)
            plot_phasing_tg(f, phasing_tg[1][0])
            plot_phasing_tg(f, phasing_tg[0][0])
            #mpl.contourf(np.abs(phasing_tg), 50)
            #mpl.show()

        with h5py.File(str(save_path), 'a') as sf:
            # save data at current t2
            t1_slice = slice(table * waveforms_per_table,
                             (table + 1) * waveforms_per_table)
            sf['raw rephasing'][t2, :, t1_slice] = R
            sf['raw non-rephasing'][t2, :, t1_slice] = NR
            sf['raw transient-grating'][t2, :, t1_slice] = TG
            nudie.log.debug('Shapes of data:')
            nudie.log.debug('R: {!s}, NR: {!s}, TG: {!s}'.format(
                R.shape, NR.shape, TG.shape))

        del data, fdata, data_t, avg, rIprobe, rIlo, rEprobe, rEsig

    with h5py.File(str(save_path), 'a') as sf:
        if detrend_t1:
            R = sf['raw rephasing'][:]
            NR = sf['raw non-rephasing'][:]
            TG = sf['raw transient-grating'][:]

            R = detrend(R, axis=2, type='constant')
            NR = detrend(NR, axis=2, type='constant')
            TG = detrend(TG, axis=2, type='constant')

            sf['raw rephasing'][:] = R
            sf['raw non-rephasing'][:] = NR
            sf['raw transient-grating'][:] = TG

        # write out meta data
        sf.attrs['batch_name'] = dd_info['batch_name']
        sf.attrs['batch_no'] = dd_info['batch_no']
        sf.attrs['batch_path'] = dd_info['batch_path']
        sf.attrs['job_name'] = dd_info['job_name']
        sf.attrs['nt1'] = dd_info['nt1']
        sf.attrs['nt2'] = dd_info['nt2']
        sf.attrs['when'] = dd_info['when']

        sf.attrs['probe lo delay estimate'] = prd_est
        sf.attrs['analysis timestamp'] = arrow.now().format('MM-DD-YYYY HH:mm')
        sf.attrs['nudie version'] = nudie.version
        sf.attrs['experiment type'] = 'stark 2d'

        # write out axes
        gaxes = sf.require_group('axes')
        freq_dataset = gaxes.create_dataset('detection frequency', data=f)
        freq_dataset.attrs['df'] = df
        gaxes.create_dataset('detection wavelength',
                             data=nudie.spectrometer.speed_of_light / f)
        gaxes.create_dataset('t1', data=dd_info['t1'])
        gaxes.create_dataset('t2', data=dd_info['t2'][:, 1])
        gaxes.create_dataset('stark axis', data=stark_axis)

        # add dimension scales
        rdata = sf['raw rephasing']
        rdata.dims.create_scale(gaxes['t2'], 'population time / fs')
        rdata.dims.create_scale(gaxes['t1'], 'pump delay time / fs')
        rdata.dims.create_scale(gaxes['detection frequency'],
                                'frequency / 1000 THz')
        rdata.dims.create_scale(gaxes['detection wavelength'],
                                'wavelength / nm')
        rdata.dims.create_scale(gaxes['stark axis'], 'average field / kV')

        # attach them
        for tmp in [
                sf['raw rephasing'], sf['raw non-rephasing'],
                sf['raw transient-grating']
        ]:
            tmp.dims[0].label = 'population time'
            tmp.dims[0].attach_scale(gaxes['t2'])
            tmp.dims[1].label = 'average voltage'
            tmp.dims[1].attach_scale(gaxes['stark axis'])
            tmp.dims[2].label = 'pump delay time'
            tmp.dims[2].attach_scale(gaxes['t1'])
            tmp.dims[3].label = 'detection axis'
            tmp.dims[3].attach_scale(gaxes['detection frequency'])
            tmp.dims[3].attach_scale(gaxes['detection wavelength'])
Ejemplo n.º 17
0
def run(pp_name,
        pp_batch,
        when='today',
        wavelengths=None,
        plot=False,
        exclude=[],
        analysis_path='analyzed',
        datapath=None):
    '''run the batch job to analyze pump-probe data
    
    Performs a shot-to-shot pump-probe analysis. 
    
    ```wavelengths``` is a path to a wavelengths file
    ```plot``` is a Boolean parameter controlling whether to show matplotlib
    plots
    '''
    # tags for phase cycles
    phase_cycles = ['none1', 'zero', 'none2', 'pipi']
    nrepeat = 1  # how many times each waveform is repeated in the camera file. Assumed to be one
    nwaveforms = 1  # number of dazzler waveforms in pump probe is 1
    npixels = 1340  # length of detection axis

    # change datapath
    if datapath is None:
        datapath = nudie.data_folder

    # load up pp data to use
    pp_info = next(
        nudie.load_job(job_name=pp_name,
                       batch_set=[pp_batch],
                       when=when,
                       data_path=datapath))

    # set current batch directory
    current_path = Path(pp_info['batch_path'])

    # generate hdf filename based on data date
    analysis_folder = Path(analysis_path)

    # create folder if it doesn't exist
    if not analysis_folder.exists():
        analysis_folder.mkdir()

    save_path = analysis_folder / (pp_info['batch_name'] + '.h5')
    # remove data file if it exists
    if save_path.exists(): save_path.unlink()

    with h5py.File(str(save_path), mode='w') as sf:
        # initialize groups
        sf.create_group('axes')

        loops = np.array(list(
            filter(lambda x: x not in exclude, pp_info['loop_range'])),
                         dtype=int)
        shape = (pp_info['nt2'], len(loops), npixels)
        raw_pp = sf.create_dataset('raw pump-probe', shape, dtype=float)
        avg_pp = sf.create_dataset('averaged pump-probe', (shape[0], shape[2]),
                                   dtype=float)
    try:
        wl = load_wavelengths(current_path.parent / wavelengths)
    except FileNotFoundError as e:
        nudie.log.error('Could not load wavelength calibration file!')
        raise e

    for t2, table, loop in it.product(pp_info['t2_range'],
                                      pp_info['table_range'], loops):

        # first file requires special synchronization
        # this is the rule that determines that it is the first file
        first = 'first' if all([t2 == 0, table == 0, loop == 0]) else False

        # Load everything for the given t2, table, and loop value
        analogs = nudie.load_analogtxt(pp_info['job_name'], current_path, t2,
                                       table, loop)
        cdata = nudie.load_camera_file(pp_info['job_name'],
                                       current_path,
                                       t2,
                                       table,
                                       loop,
                                       force_uint16=True)

        # Synchronize it, trimming some frames in the beginning and end
        data, (a1, a2) = nudie.trim_all(*nudie.synchronize_daq_to_camera(\
                        cdata, analog_channels=analogs, which_file=first))
        data = data.astype(float)  # convert to float64 before manipulating it!

        # interpolate to even frequency spacing
        f, data, df = nudie.wavelen_to_freq(wl, data, ret_df=True, ax=0)

        start_idxs, period = nudie.detect_table_start(a1)

        # determine where the shutter is
        shutter_open, shutter_closed, duty_cycle = nudie.determine_shutter_shots(
            data)

        # tag the phases
        shutter_info = {
            'last open idx': shutter_open,
            'first closed idx': shutter_closed
        }
        tags = nudie.tag_phases(start_idxs,
                                period,
                                tags=phase_cycles,
                                nframes=data.shape[1],
                                shutter_info=shutter_info)

        # the first two indexes should be zero, so ignore them
        tags = tags[0][0]

        # truncate everything to the same length
        trunc_idx = slice(
            min([len(tags[k]['shutter open']) for k in phase_cycles]))

        # subtract scatter from shutter closed shots
        zero = ((data[:, tags['zero']['shutter open']][:, trunc_idx]).T \
                - data[:, tags['zero']['shutter closed']].mean(axis=-1)).T
        pipi = ((data[:, tags['pipi']['shutter open']][:, trunc_idx]).T \
                - data[:, tags['pipi']['shutter closed']].mean(axis=-1)).T
        none1 = ((data[:, tags['none1']['shutter open']][:, trunc_idx]).T \
                - data[:, tags['none1']['shutter closed']].mean(axis=-1)).T
        none2 = ((data[:, tags['none2']['shutter open']][:, trunc_idx]).T \
                - data[:, tags['none2']['shutter closed']].mean(axis=-1)).T

        # to match frank's code, except with minus signs
        A3 = (0.5 * zero - 0.5 * none1)
        B3 = (0.5 * pipi - 0.5 * none2)

        # FIXME: unsure about the proper signs for C3!
        C3 = (0.25 * zero + 0.25 * none1 + 0.25 * pipi + 0.25 * none2)
        C3 = ma.masked_less(C3, 1e-7 * C3.max())

        S3 = np.mean((0.5 * A3 + 0.5 * B3) / ma.sqrt(C3), axis=1)
        pp = ma.filled(S3 / np.max(S3), fill_value=0)

        with h5py.File(str(save_path), mode='a') as sf:
            loop_idx = np.argmin(np.abs(loops - loop))
            sf['raw pump-probe'][t2, loop_idx] = pp

    with h5py.File(str(save_path), mode='a') as sf:
        if plot:
            for t2, i in it.product(pp_info['t2_range'], range(len(loops))):
                mpl.plot(f,
                         sf['raw pump-probe'][t2, i],
                         label='t2 {:.1f} loop {:d}'.format(
                             pp_info['t2'][t2][1], loops[i]))
            mpl.legend()
            mpl.show()

        sf['averaged pump-probe'][:] = np.array(
            sf['raw pump-probe']).mean(axis=1)

        # write out meta data
        gaxes = sf['axes']
        gaxes['detection wavelength'] = nudie.spectrometer.speed_of_light / f
        gaxes['t2'] = pp_info['t2'][:, 1]  # take only the time position
        gaxes['loop'] = loops
        gaxes['detection frequency'] = f
        gaxes['detection frequency'].attrs['df'] = df

        sf.attrs['batch_name'] = pp_info['batch_name']
        sf.attrs['batch_no'] = pp_info['batch_no']
        sf.attrs['batch_path'] = pp_info['batch_path']
        sf.attrs['job_name'] = pp_info['job_name']
        sf.attrs['nt2'] = pp_info['nt2']
        sf.attrs['when'] = pp_info['when']

        sf.attrs['nloop'] = pp_info['loop_range'].stop
        sf.attrs['analysis timestamp'] = arrow.now().format('DD-MM-YYYY HH:mm')
        sf.attrs['nudie version'] = nudie.version
        sf.attrs['experiment type'] = 'pump probe'

        # add dimension scales
        rpp = sf['raw pump-probe']
        rpp.dims.create_scale(gaxes['loop'], 'loop number')
        rpp.dims.create_scale(gaxes['t2'], 'population time / fs')
        rpp.dims.create_scale(gaxes['detection frequency'],
                              'detection frequency / 1000 THz')
        rpp.dims.create_scale(gaxes['detection wavelength'],
                              'detection wavelength / nm')

        rpp.dims[0].label = 'population time'
        rpp.dims[0].attach_scale(gaxes['t2'])
        rpp.dims[1].label = 'loop number'
        rpp.dims[1].attach_scale(gaxes['loop'])
        rpp.dims[2].label = 'detection axis'
        rpp.dims[2].attach_scale(gaxes['detection frequency'])
        rpp.dims[2].attach_scale(gaxes['detection wavelength'])

        app = sf['averaged pump-probe']
        app.dims[0].label = 'population time'
        app.dims[0].attach_scale(gaxes['t2'])
        app.dims[1].label = 'detection axis'
        app.dims[1].attach_scale(gaxes['detection frequency'])
        app.dims[1].attach_scale(gaxes['detection wavelength'])