be_viz_utils.py 60.6 KB
Newer Older
1
2
3
4
# -*- coding: utf-8 -*-
"""
Created on Thu Jan 05 13:29:12 2017

5
@author: Suhas Somnath, Chris R. Smith
6
"""
7
from __future__ import division, print_function, absolute_import, unicode_literals
8

9
10
import os
from warnings import warn
11

12
import ipywidgets as widgets
13
import numpy as np
14
from IPython.display import display
15
from matplotlib import pyplot as plt
16

Unknown's avatar
Unknown committed
17
18
from .plot_utils import plot_loops, plot_map_stack, get_cmap_object, single_img_cbar_plot, save_fig_filebox_button, \
    set_tick_font_size
19
from ..analysis.utils.be_loop import loop_fit_function
Unknown's avatar
Unknown committed
20
from ..analysis.utils.be_sho import SHOfunc
21
22
from ..io.hdf_utils import reshape_to_Ndims, getAuxData, get_sort_order, get_dimensionality, get_attr, \
    get_source_dataset
23
24


25
def visualize_sho_results(h5_main, save_plots=True, show_plots=True, cmap=None):
26
27
28
29
30
31
32
33
34
35
36
37
    """
    Plots some loops, amplitude, phase maps for BE-Line and BEPS datasets.\n
    Note: The file MUST contain SHO fit gusses at the very least

    Parameters
    ----------
    h5_main : HDF5 Dataset
        dataset to be plotted
    save_plots : (Optional) Boolean
        Whether or not to save plots to files in the same directory as the h5 file
    show_plots : (Optional) Boolean
        Whether or not to display the plots on the screen
38
39
    cmap : String, or matplotlib.colors.LinearSegmentedColormap object (Optional)
        Requested color map
40
41
42
43
44

    Returns
    -------
    None
    """
45
    cmap = get_cmap_object(cmap)
46
47
48
49
50
51
52
53
54
55
56
57

    def __plot_loops_maps(ac_vec, resp_mat, grp_name, win_title, spec_var_title, meas_var_title, save_plots,
                          folder_path, basename, num_rows, num_cols):
        plt_title = grp_name + '_' + win_title + '_Loops'
        fig, ax = plot_loops(ac_vec, resp_mat, evenly_spaced=True, plots_on_side=5, use_rainbow_plots=False,
                             x_label=spec_var_title, y_label=meas_var_title, subtitles='Position', title=plt_title)
        if save_plots:
            fig.savefig(os.path.join(folder_path, basename + '_' + plt_title + '.png'), format='png', dpi=300)

        plt_title = grp_name + '_' + win_title + '_Snaps'
        fig, axes = plot_map_stack(resp_mat.reshape(num_rows, num_cols, resp_mat.shape[1]),
                                   color_bar_mode="each", evenly_spaced=True, title='UDVS Step #',
58
                                   heading=plt_title, cmap=cmap)
59
60
61
62
63
64
65
66
67
        if save_plots:
            fig.savefig(os.path.join(folder_path, basename + '_' + plt_title + '.png'), format='png', dpi=300)

    plt_path = None

    print('Creating plots of SHO Results from {}.'.format(h5_main.name))

    h5_file = h5_main.file

Chris Smith's avatar
Chris Smith committed
68
    expt_type = get_attr(h5_file, 'data_type')
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    if expt_type not in ['BEPSData', 'BELineData']:
        warn('Unsupported data format')
        return
    isBEPS = expt_type == 'BEPSData'

    (folder_path, basename) = os.path.split(h5_file.filename)
    basename, _ = os.path.splitext(basename)

    sho_grp = h5_main.parent

    chan_grp = h5_file['/'.join(sho_grp.name[1:].split('/')[:2])]

    grp_name = '_'.join(chan_grp.name[1:].split('/'))
    grp_name = '_'.join([grp_name, sho_grp.name.split('/')[-1].split('-')[0], h5_main.name.split('/')[-1]])

    try:
Chris Smith's avatar
Chris Smith committed
85
        h5_pos = h5_file[get_attr(h5_main, 'Position_Indices')]
86
87
88
89
90
91
92
93
94
95
    except KeyError:
        print('No Position_Indices found as attribute of {}'.format(h5_main.name))
        print('Rows and columns will be calculated from dataset shape.')
        num_rows = int(np.floor((np.sqrt(h5_main.shape[0]))))
        num_cols = int(np.reshape(h5_main, [num_rows, -1, h5_main.shape[1]]).shape[1])
    else:
        num_rows = len(np.unique(h5_pos[:, 0]))
        num_cols = len(np.unique(h5_pos[:, 1]))

    try:
Chris Smith's avatar
Chris Smith committed
96
        h5_spec_vals = h5_file[get_attr(h5_main, 'Spectroscopic_Values')]
97
98
99
    # except KeyError:
    #     warn('No Spectrosocpic Datasets found as attribute of {}'.format(h5_main.name))
    #     raise
Unknown's avatar
Unknown committed
100
    except Exception:
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
        raise

    # Assume that there's enough memory to load all the guesses into memory
    amp_mat = h5_main['Amplitude [V]'] * 1000  # convert to mV ahead of time
    freq_mat = h5_main['Frequency [Hz]'] / 1000
    q_mat = h5_main['Quality Factor']
    phase_mat = h5_main['Phase [rad]']
    rsqr_mat = h5_main['R2 Criterion']

    if isBEPS:
        meas_type = chan_grp.parent.attrs['VS_mode']
        # basically 3 kinds for now - DC/current, AC, UDVS - lets ignore this
        if meas_type == 'load user defined VS Wave from file':
            warn('Not handling custom experiments for now')
            # h5_file.close()
            return

        # Plot amplitude and phase maps at one or more UDVS steps

        if meas_type == 'AC modulation mode with time reversal':
            center = int(h5_spec_vals.shape[1] * 0.5)
            ac_vec = np.squeeze(h5_spec_vals[h5_spec_vals.attrs['AC_Amplitude']][:, 0:center])

            forw_resp = np.squeeze(amp_mat[:, slice(0, center)])
            rev_resp = np.squeeze(amp_mat[:, slice(center, None)])

            for win_title, resp_mat in zip(['Forward', 'Reverse'], [forw_resp, rev_resp]):
                __plot_loops_maps(ac_vec, resp_mat, grp_name, win_title, 'AC Amplitude', 'Amplitude', save_plots,
                                  folder_path, basename, num_rows, num_cols)
        else:
            # plot loops at a few locations
            dc_vec = np.squeeze(h5_spec_vals[h5_spec_vals.attrs['DC_Offset']])
            if chan_grp.parent.attrs['VS_measure_in_field_loops'] == 'in and out-of-field':

                dc_vec = np.squeeze(dc_vec[slice(0, None, 2)])

                in_phase = np.squeeze(phase_mat[:, slice(0, None, 2)])
                in_amp = np.squeeze(amp_mat[:, slice(0, None, 2)])
                out_phase = np.squeeze(phase_mat[:, slice(1, None, 2)])
                out_amp = np.squeeze(amp_mat[:, slice(1, None, 2)])

                for win_title, resp_mat in zip(['In_Field', 'Out_of_Field'], [in_phase * in_amp, out_phase * out_amp]):
                    __plot_loops_maps(dc_vec, resp_mat, grp_name, win_title, 'DC Bias', 'Piezoresponse (a.u.)',
                                      save_plots, folder_path, basename, num_rows, num_cols)
            else:
                __plot_loops_maps(dc_vec, phase_mat * amp_mat, grp_name, '', 'DC Bias', 'Piezoresponse (a.u.)',
                                  save_plots, folder_path, basename, num_rows, num_cols)

    else:  # BE-Line can only visualize the amplitude and phase maps:
        amp_mat = amp_mat.reshape(num_rows, num_cols)
        freq_mat = freq_mat.reshape(num_rows, num_cols)
        q_mat = q_mat.reshape(num_rows, num_cols)
        phase_mat = phase_mat.reshape(num_rows, num_cols)
        rsqr_mat = rsqr_mat.reshape(num_rows, num_cols)

        fig_ms, ax_ms = plot_map_stack(np.dstack((amp_mat, freq_mat, q_mat, phase_mat, rsqr_mat)),
                                       num_comps=5, color_bar_mode='each', heading=grp_name,
                                       title=['Amplitude (mV)', 'Frequency (kHz)', 'Quality Factor', 'Phase (deg)',
159
                                              'R^2 Criterion'], cmap=cmap)
160
161
162
163

        if save_plots:
            plt_path = os.path.join(folder_path, basename + '_' + grp_name + 'Maps.png')
            fig_ms.savefig(plt_path, format='png', dpi=300)
164
165
166
167
168
169
170

    if show_plots:
        plt.show()

    plt.close('all')


171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def plot_loop_guess_fit(vdc, ds_proj_loops, ds_guess, ds_fit, title=''):
    """
    Plots the loop guess, fit, source projected loops for a single cycle

    Parameters
    ----------
    vdc - 1D float numpy array
        DC offset vector (unshifted)
    ds_proj_loops - 2D numpy array
        Projected loops arranged as [position, vdc]
    ds_guess - 1D compound numpy array
        Loop guesses arranged as [position]
    ds_fit - 1D compound numpy array
        Loop fits arranged as [position]
    title - (Optional) String / unicode
        Title for the figure

    Returns
    ----------
    fig - matplotlib.pyplot.figure object
        Figure handle
    axes - 2D array of matplotlib.pyplot.axis handles
        handles to axes in the 2d figure
    """
    shift_ind = int(-1 * len(vdc) / 4)
    vdc_shifted = np.roll(vdc, shift_ind)
    loops_shifted = np.roll(ds_proj_loops, shift_ind, axis=1)

    num_plots = np.min([5, int(np.sqrt(ds_proj_loops.shape[0]))])
    fig, axes = plt.subplots(nrows=num_plots, ncols=num_plots, figsize=(18, 18))
    positions = np.linspace(0, ds_proj_loops.shape[0] - 1, num_plots ** 2, dtype=np.int)
    for ax, pos in zip(axes.flat, positions):
        ax.plot(vdc_shifted, loops_shifted[pos, :], 'k', label='Raw')
        ax.plot(vdc_shifted, loop_fit_function(vdc_shifted, np.array(list(ds_guess[pos]))), 'g', label='guess')
        ax.plot(vdc_shifted, loop_fit_function(vdc_shifted, np.array(list(ds_fit[pos]))), 'r--', label='Fit')
        ax.set_xlabel('V_DC (V)')
        ax.set_ylabel('PR (a.u.)')
        ax.set_title('Position ' + str(pos))
    ax.legend()
    fig.suptitle(title)
    fig.tight_layout()

    return fig, axes


216
def jupyter_visualize_beps_sho(h5_sho_dset, step_chan, resp_func=None, resp_label='Response', cmap=None):
217
218
219
    """
    Jupyer notebook ONLY function. Sets up an interactive visualizer for viewing SHO fitted BEPS data.
    Currently, this is limited to DC and AC spectroscopy datasets.
220

221
222
223
224
225
226
227
228
229
230
231
    Parameters
    ----------
    h5_sho_dset : h5py.Dataset
        dataset to be plotted
    step_chan : string / unicode
        Name of the channel that forms the primary spectroscopic axis (eg - DC offset)
    resp_func : function (optional)
        Function to apply to the spectroscopic data. Currently, DC spectroscopy uses A*cos(phi) and AC spectroscopy
        uses A
    resp_label : string / unicode (optional)
        Label for the response (y) axis.
232
233
    cmap : String, or matplotlib.colors.LinearSegmentedColormap object (Optional)
        Requested color map
234
    """
235
236
    cmap = get_cmap_object(cmap)

237
238
239
240
    guess_3d_data, success = reshape_to_Ndims(h5_sho_dset)

    h5_sho_spec_inds = getAuxData(h5_sho_dset, 'Spectroscopic_Indices')[0]
    h5_sho_spec_vals = getAuxData(h5_sho_dset, 'Spectroscopic_Values')[0]
241
242
    spec_nd, _ = reshape_to_Ndims(h5_sho_spec_inds, h5_spec=h5_sho_spec_inds)
    sho_spec_dims = np.array(spec_nd.shape[1:])
243
    sho_spec_labels = get_attr(h5_sho_spec_inds, 'labels')
244
245

    h5_pos_inds = getAuxData(h5_sho_dset, auxDataName='Position_Indices')[-1]
246
247
    pos_nd, _ = reshape_to_Ndims(h5_pos_inds, h5_pos=h5_pos_inds)
    pos_dims = list(pos_nd.shape[:h5_pos_inds.shape[1]])
248
    pos_labels = get_attr(h5_pos_inds, 'labels')
249
250
251
252
253
254

    # reshape to X, Y, step, all others
    spec_step_dim_ind = np.where(sho_spec_labels == step_chan)[0][0]
    step_dim_ind = len(pos_dims) + spec_step_dim_ind

    # move the step dimension to be the first after all position dimensions
255
    rest_sho_dim_order = list(range(len(pos_dims), len(guess_3d_data.shape)))
256
    rest_sho_dim_order.remove(step_dim_ind)
257
    new_order = list(range(len(pos_dims))) + [step_dim_ind] + rest_sho_dim_order
258
259
260
261
262

    # Transpose the 3D dataset to this shape:
    sho_guess_Nd_1 = np.transpose(guess_3d_data, new_order)

    # Now move the step dimension to the front for the spec labels as well
263
    new_spec_order = list(range(len(sho_spec_labels)))
264
265
266
267
268
269
270
271
272
273
    new_spec_order.remove(spec_step_dim_ind)
    new_spec_order = [spec_step_dim_ind] + new_spec_order

    # new_spec_labels = sho_spec_labels[new_spec_order]
    new_spec_dims = np.array(sho_spec_dims)[new_spec_order]

    # Now collapse all additional dimensions
    final_guess_shape = pos_dims + [new_spec_dims[0]] + [-1]
    sho_dset_collapsed = np.reshape(sho_guess_Nd_1, final_guess_shape)

274
275
    # Get the bias matrix:
    bias_mat, _ = reshape_to_Ndims(h5_sho_spec_vals, h5_spec=h5_sho_spec_inds)
Unknown's avatar
Unknown committed
276
277
    bias_mat = np.transpose(bias_mat[spec_step_dim_ind],
                            new_spec_order).reshape(sho_dset_collapsed.shape[len(pos_dims):])
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307

    # This is just the visualizer:
    sho_quantity = 'Amplitude [V]'
    step_ind = 0
    row_ind = 1
    col_ind = 1

    def dc_spectroscopy_func(resp_vec):
        return resp_vec['Amplitude [V]'] * np.cos(resp_vec['Phase [rad]']) * 1E+3

    def ac_spectroscopy_func(resp_vec):
        return resp_vec['Amplitude [V]']

    if resp_func is None:
        if step_chan == 'DC_Offset':
            resp_func = dc_spectroscopy_func
            resp_label = 'A cos($\phi$) (a. u.)'
        else:
            resp_func = ac_spectroscopy_func
            resp_label = 'Amplitude (a. u.)'

    spatial_map = sho_dset_collapsed[:, :, step_ind, 0][sho_quantity]
    resp_vec = sho_dset_collapsed[row_ind, col_ind, :, :]
    resp_vec = resp_func(resp_vec)

    fig = plt.figure(figsize=(12, 8))
    ax_bias = plt.subplot2grid((3, 2), (0, 0), colspan=1, rowspan=1)
    ax_map = plt.subplot2grid((3, 2), (1, 0), rowspan=2)
    ax_loop = plt.subplot2grid((3, 2), (0, 1), colspan=1, rowspan=3)

308
    ax_bias.plot(bias_mat[:, 0])
309
310
311
312
    ax_bias.set_xlabel('Bias Step')
    ax_bias.set_ylabel(step_chan.replace('_', ' ') + ' (V)')
    bias_slider = ax_bias.axvline(x=step_ind, color='r')

Unknown's avatar
Unknown committed
313
    img_map, img_cmap = single_img_cbar_plot(ax_map, spatial_map.T, show_xy_ticks=None)
Unknown's avatar
Unknown committed
314
315

    map_title = '{} - {}={}'.format(sho_quantity, step_chan, bias_mat[step_ind][0])
316
317
    ax_map.set_xlabel('X')
    ax_map.set_ylabel('Y')
Unknown's avatar
Unknown committed
318
    ax_map.set_title(map_title)
319
    crosshair = ax_map.plot(row_ind, col_ind, 'k+')[0]
320
321
322

    ax_loop.axvline(x=0, color='gray', linestyle='--')
    ax_loop.axhline(y=0, color='gray', linestyle='--')
323
    line_handles = ax_loop.plot(bias_mat, resp_vec)
324
325
326
327
    ax_loop.set_xlabel(step_chan.replace('_', ' ') + ' (V)')
    ax_loop.set_ylabel(resp_label)
    fig.tight_layout()

328
329
    plt.show()

330
331
332
333
    def update_sho_plots(sho_quantity, **kwargs):
        step_ind = kwargs['Bias Step']
        bias_slider.set_xdata((step_ind, step_ind))
        spatial_map = sho_dset_collapsed[:, :, step_ind, 0][sho_quantity]
Unknown's avatar
Unknown committed
334
335
        map_title = '{} - {}={}'.format(sho_quantity, step_chan, bias_mat[step_ind][0])
        ax_map.set_title(map_title)
336
        img_map.set_data(spatial_map.T)
337
338
339
340
        spat_mean = np.mean(spatial_map)
        spat_std = np.std(spatial_map)
        img_map.set_clim(vmin=spat_mean - 3 * spat_std, vmax=spat_mean + 3 * spat_std)

341
342
343
344
345
346
347
348
349
    def pos_picker(event):
        if not img_map.axes.in_axes(event):
            return

        xdata = int(round(event.xdata))
        ydata = int(round(event.ydata))
        current_pos = {pos_labels[0]: xdata, pos_labels[1]: ydata}

        pos_dim_vals = list(range(len(pos_labels)))
350

351
352
353
354
355
356
357
        for pos_dim_ind, pos_dim_name in enumerate(pos_labels):
            pos_dim_vals[pos_dim_ind] = current_pos[pos_dim_name]

        crosshair.set_xdata(xdata)
        crosshair.set_ydata(ydata)

        resp_vec = sho_dset_collapsed[xdata, ydata, :, :]
358
359
360
        resp_vec = resp_func(resp_vec)
        for line_handle, data in zip(line_handles, np.transpose(resp_vec)):
            line_handle.set_ydata(data)
Unknown's avatar
Unknown committed
361
362
363

        ax_loop.relim()
        ax_loop.autoscale_view()
364
365

        fig.canvas.draw()
366
367

    slider_dict = dict()
368
    slider_dict['Bias Step'] = (0, bias_mat.shape[0] - 1, 1)
369

370
    cid = img_map.figure.canvas.mpl_connect('button_press_event', pos_picker)
Unknown's avatar
Unknown committed
371

372
    display(save_fig_filebox_button(fig, h5_sho_dset.file.filename))
373
    widgets.interact(update_sho_plots, sho_quantity=list(sho_dset_collapsed.dtype.names[:-1]), **slider_dict)
374

375
def jupyter_visualize_be_spectrograms(h5_main, cmap=None):
376
377
378
379
    """
    Jupyer notebook ONLY function. Sets up a simple visualzier for visualizing raw BE data.
    Sliders for position indices can be used to visualize BE spectrograms (frequency, UDVS step).
    In the case of 2 spatial dimensions, a spatial map will be provided as well
380

381
382
383
    Parameters
    ----------
    h5_main : h5py.Dataset
384
385
386
        Raw Band Excitation dataset
    cmap : String, or matplotlib.colors.LinearSegmentedColormap object (Optional)
        Requested color map
387
    """
Unknown's avatar
Unknown committed
388
    cmap = get_cmap_object(cmap)
389

390
391
392
    h5_pos_inds = getAuxData(h5_main, auxDataName='Position_Indices')[-1]
    pos_sort = get_sort_order(np.transpose(h5_pos_inds))
    pos_dims = get_dimensionality(np.transpose(h5_pos_inds), pos_sort)
Unknown's avatar
Unknown committed
393
    pos_labels = np.array(get_attr(h5_pos_inds, 'labels'))[pos_sort]
394
395
396
397
398

    h5_spec_vals = getAuxData(h5_main, auxDataName='Spectroscopic_Values')[-1]
    h5_spec_inds = getAuxData(h5_main, auxDataName='Spectroscopic_Indices')[-1]
    spec_sort = get_sort_order(h5_spec_inds)
    spec_dims = get_dimensionality(h5_spec_inds, spec_sort)
Unknown's avatar
Unknown committed
399
    spec_labels = np.array(get_attr(h5_spec_inds, 'labels'))[spec_sort]
400
401
402
403

    ifreq = int(np.argwhere(spec_labels == 'Frequency'))
    freqs_nd = reshape_to_Ndims(h5_spec_vals, h5_spec=h5_spec_inds)[0][ifreq].squeeze()
    freqs_2d = freqs_nd.reshape(freqs_nd.shape[0], -1) / 1000  # Convert to kHz
Chris Smith's avatar
Chris Smith committed
404
405
406

    num_udvs_steps = h5_main.parent.parent.attrs.get('num_udvs_steps', None)

Unknown's avatar
Unknown committed
407
    # New version of aquisition software renamed num_udvs_steps to num_UDVS_steps
Chris Smith's avatar
Chris Smith committed
408
409
410
411
412
    if num_udvs_steps is None:
        num_udvs_steps = h5_main.parent.parent.attrs.get('num_UDVS_steps', None)

    if num_udvs_steps is None:
        warn('Number of UDVS steps could not be determined.  Will assume single step.')
413
414

    if len(pos_dims) == 2:
415
        spatial_map = np.abs(np.reshape(h5_main[:, 0], pos_dims[::-1]))
416
        spectrogram = np.reshape(h5_main[0], (num_udvs_steps, -1))
Unknown's avatar
Unknown committed
417
        fig, axes = plt.subplots(ncols=3, figsize=(12, 4), subplot_kw={'adjustable': 'box-forced'})
Unknown's avatar
Unknown committed
418
419
        spatial_img, spatial_cbar = single_img_cbar_plot(axes[0], np.abs(spatial_map), x_size=spatial_map.shape[0],
                                                         y_size=spatial_map.shape[1], cmap=cmap)
Unknown's avatar
Unknown committed
420
        axes[0].set_aspect('equal')
421
422
        axes[0].set_xlabel('X')
        axes[0].set_ylabel('Y')
423
424

        crosshair = axes[0].plot(int(0.5 * spatial_map.shape[0]), int(0.5 * spatial_map.shape[1]), 'k+')[0]
425
426

        if len(spec_dims) > 1:
Unknown's avatar
Unknown committed
427
            amp_img, amp_cbar = single_img_cbar_plot(axes[1], np.abs(spectrogram), show_xy_ticks=None, cmap=cmap,
Unknown's avatar
Unknown committed
428
429
430
                                                     extent=[freqs_2d[0, 0], freqs_2d[-1, 0],
                                                             0, spectrogram.shape[0]])

Unknown's avatar
Unknown committed
431
            phase_img, phase_cbar = single_img_cbar_plot(axes[2], np.angle(spectrogram), show_xy_ticks=None, cmap=cmap,
Unknown's avatar
Unknown committed
432
433
434
                                                         extent=[freqs_2d[0, 0], freqs_2d[-1, 0],
                                                                 0, spectrogram.shape[0]])

435
            for axis in axes[1:3]:
Unknown's avatar
Unknown committed
436
                axis.set_ylabel('BE step')
437
                axis.axis('tight')
Unknown's avatar
Unknown committed
438
439
440
441
                x0, x1 = (freqs_2d[0, 0], freqs_2d[-1, 0])
                y0, y1 = (0, spectrogram.shape[0])
                axis.set_aspect(np.abs(x1-x0)/np.abs(y1-y0))

442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
        else:
            # BE-Line
            axes[1].set_ylabel('Amplitude (a. u.)')
            axes[2].set_ylabel('Phase (rad)')
            spectrogram = np.squeeze(spectrogram)
            amp_img = axes[1].plot(np.abs(spectrogram))[0]
            phase_img = axes[2].plot(np.angle(spectrogram))[0]
            amp_full = np.abs(h5_main[()])
            amp_mean = np.mean(amp_full)
            amp_std = np.std(amp_full)
            st_devs = 4

            axes[1].set_ylim([0, amp_mean + st_devs * amp_std])
            axes[2].set_ylim([-np.pi, np.pi])

457
458
        axes[1].set_title('Amplitude - Position ({},{})'.format(int(0.5 * spatial_map.shape[0]),
                                                                int(0.5 * spatial_map.shape[0])))
459
        axes[1].set_xlabel('Frequency (kHz)')
460

461
462
        axes[2].set_title('Phase - Position ({},{})'.format(int(0.5 * spatial_map.shape[0]),
                                                                int(0.5 * spatial_map.shape[0])))
463
        axes[2].set_xlabel('Frequency (kHz)')
464

465
466
        fig.tight_layout()

467
        plt.show()
Unknown's avatar
Unknown committed
468
469
        display(save_fig_filebox_button(fig, h5_main.file.filename))

470
471
        def index_unpacker(step):
            spatial_map = np.abs(np.reshape(h5_main[:, step], pos_dims[::-1]))
472
473
474
475
476
477
            spatial_img.set_data(spatial_map)
            spat_mean = np.mean(spatial_map)
            spat_std = np.std(spatial_map)
            spatial_img.set_clim(vmin=spat_mean - 3 * spat_std, vmax=spat_mean + 3 * spat_std)

            spec_heading = ''
Unknown's avatar
Unknown committed
478
            for dim_ind, dim_name in enumerate(spec_labels):
479
                spec_heading += dim_name + ': ' + str(h5_spec_vals[dim_ind, step]) + ', '
480
            axes[0].set_title(spec_heading[:-2])
481
482
483
484
485
486
487
488
489
            fig.canvas.draw()

        def pos_picker(event):
            if not spatial_img.axes.in_axes(event):
                return

            xdata = int(round(event.xdata))
            ydata = int(round(event.ydata))
            current_pos = {pos_labels[0]: xdata, pos_labels[1]: ydata}
490

Unknown's avatar
Unknown committed
491
            pos_dim_vals = list(range(len(pos_labels)))
492

493
            for pos_dim_ind, pos_dim_name in enumerate(pos_labels):
494
495
496
497
                pos_dim_vals[pos_dim_ind] = current_pos[pos_dim_name]

            crosshair.set_xdata(xdata)
            crosshair.set_ydata(ydata)
498

499
500
            axes[1].set_title('Amplitude - Position ({},{})'.format(xdata,ydata))
            axes[2].set_title('Phase - Position ({},{})'.format(xdata, ydata))
501
502
503
504
505

            pix_ind = pos_dim_vals[0]
            for pos_dim_ind in range(1, len(pos_labels)):
                pix_ind += pos_dim_vals[pos_dim_ind] * pos_dims[pos_dim_ind - 1]
            spectrogram = np.reshape(h5_main[pix_ind], (num_udvs_steps, -1))
506
507
508
509
510
511
512

            if len(spec_dims) > 1:
                amp_img.set_data(np.abs(spectrogram))
                phase_img.set_data(np.angle(spectrogram))
            else:
                amp_img.set_ydata(np.abs(spectrogram))
                phase_img.set_ydata(np.angle(spectrogram))
Unknown's avatar
Unknown committed
513
514
            amp_cbar.changed()
            phase_cbar.changed()
515

516
517
518
519
520
            fig.canvas.draw()

        cid = spatial_img.figure.canvas.mpl_connect('button_press_event', pos_picker)
        widgets.interact(index_unpacker, step=(0, h5_main.shape[1] - 1, 1))
        plt.show()
521
522
523
524
525

    else:
        def plot_spectrogram(data, freq_vals):
            fig, axes = plt.subplots(ncols=2, figsize=(9, 5), sharey=True)
            im_handles = list()
526
            im_handles.append(axes[0].imshow(np.abs(data), cmap=cmap,
527
                                             extent=[freqs_2d[0, 0], freqs_2d[-1, 0],
528
529
                                                     data.shape[0], 0],
                                             interpolation='none'))
530
531
            axes[0].set_title('Amplitude')
            axes[0].set_ylabel('BE step')
532
            im_handles.append(axes[1].imshow(np.angle(data), cmap=cmap,
533
534
                                             extent=[freqs_2d[0, 0], freqs_2d[-1, 0],
                                                     data.shape[0], 0],
535
                                             interpolation='none'))
Unknown's avatar
Unknown committed
536
            axes[1].set_title('Phase')
537
538
539
540
541
542
543
544
            axes[0].set_xlabel('Frequency index')
            axes[1].set_xlabel('Frequency index')
            for axis in axes:
                axis.axis('tight')
                axis.set_ylim(0, data.shape[0])
            fig.tight_layout()
            return fig, axes, im_handles

Chris Smith's avatar
Chris Smith committed
545
        fig, axes, im_handles = plot_spectrogram(np.reshape(h5_main[0], (num_udvs_steps, -1)), freqs_2d)
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563

        def position_unpacker(**kwargs):
            pos_dim_vals = range(len(pos_labels))
            for pos_dim_ind, pos_dim_name in enumerate(pos_labels):
                pos_dim_vals[pos_dim_ind] = kwargs[pos_dim_name]
            pix_ind = pos_dim_vals[0]
            for pos_dim_ind in range(1, len(pos_labels)):
                pix_ind += pos_dim_vals[pos_dim_ind] * pos_dims[pos_dim_ind - 1]
            spectrogram = np.reshape(h5_main[pix_ind], (num_udvs_steps, -1))
            im_handles[0].set_data(np.abs(spectrogram))
            im_handles[1].set_data(np.angle(spectrogram))
            display(fig)

        pos_dict = dict()
        for pos_dim_ind, dim_name in enumerate(pos_labels):
            pos_dict[dim_name] = (0, pos_dims[pos_dim_ind] - 1, 1)

        widgets.interact(position_unpacker, **pos_dict)
564
        display(fig)
565

566
def jupyter_visualize_beps_loops(h5_projected_loops, h5_loop_guess, h5_loop_fit, step_chan='DC_Offset', cmap=None):
567
568
    """
    Interactive plotting of the BE Loops
569

570
571
572
573
574
575
576
577
578
579
580
    Parameters
    ----------
    h5_projected_loops : h5py.Dataset
        Dataset holding the loop projections
    h5_loop_guess : h5py.Dataset
        Dataset holding the loop guesses
    h5_loop_fit : h5py.Dataset
        Dataset holding the loop gits
    step_chan : str, optional
        The name of the Spectroscopic dimension to plot versus.  Needs testing.
        Default 'DC_Offset'
581
582
    cmap : String, or matplotlib.colors.LinearSegmentedColormap object (Optional)
        Requested color map
583
584
585
586
587

    Returns
    -------
    None
    """
588
589
    cmap = get_cmap_object(cmap)

590
591
592
593
594
595
596
597
    # Prepare some variables for plotting loops fits and guesses
    # Plot the Loop Guess and Fit Results
    proj_nd, _ = reshape_to_Ndims(h5_projected_loops)
    guess_nd, _ = reshape_to_Ndims(h5_loop_guess)
    fit_nd, _ = reshape_to_Ndims(h5_loop_fit)

    h5_projected_loops = h5_loop_guess.parent['Projected_Loops']
    h5_proj_spec_inds = getAuxData(h5_projected_loops,
598
                                   auxDataName='Spectroscopic_Indices')[-1]
599
    h5_proj_spec_vals = getAuxData(h5_projected_loops,
600
                                   auxDataName='Spectroscopic_Values')[-1]
601
602
    h5_pos_inds = getAuxData(h5_projected_loops,
                             auxDataName='Position_Indices')[-1]
603
604
    pos_nd, _ = reshape_to_Ndims(h5_pos_inds, h5_pos=h5_pos_inds)
    pos_dims = list(pos_nd.shape[:h5_pos_inds.shape[1]])
605
    pos_labels = get_attr(h5_pos_inds, 'labels')
606
607

    # reshape the vdc_vec into DC_step by Loop
608
609
    spec_nd, _ = reshape_to_Ndims(h5_proj_spec_vals, h5_spec=h5_proj_spec_inds)
    loop_spec_dims = np.array(spec_nd.shape[1:])
610
    loop_spec_labels = get_attr(h5_proj_spec_vals, 'labels')
611
612
613

    spec_step_dim_ind = np.where(loop_spec_labels == step_chan)[0][0]

614
    # # move the step dimension to be the first after all position dimensions
615
    rest_loop_dim_order = list(range(len(pos_dims), len(proj_nd.shape)))
616
    rest_loop_dim_order.pop(spec_step_dim_ind)
617
    new_order = list(range(len(pos_dims))) + [len(pos_dims) + spec_step_dim_ind] + rest_loop_dim_order
618

619
    new_spec_order = np.array(new_order[len(pos_dims):], dtype=np.uint32) - len(pos_dims)
620

621
622
    # Also reshape the projected loops to Positions-DC_Step-Loop
    final_loop_shape = pos_dims + [loop_spec_dims[spec_step_dim_ind]] + [-1]
623
624
625
    proj_nd_3 = np.reshape(proj_nd, final_loop_shape)

    # Do the same for the guess and fit datasets
626
    guess_3d = np.reshape(guess_nd, pos_dims + [-1])
627
    fit_3d = np.reshape(fit_nd, pos_dims + [-1])
628
629

    # Get the bias vector:
630
    bias_vec = np.reshape(spec_nd[spec_step_dim_ind], final_loop_shape[len(pos_dims):])
631
632

    # Shift the bias vector and the loops by a quarter cycle
633
    shift_ind = int(-1 * bias_vec.shape[0] / 4)
634
635
636
637
638
639
640
641
642
643
644
    bias_shifted = np.roll(bias_vec, shift_ind, axis=0)
    proj_nd_shifted = np.roll(proj_nd_3, shift_ind, axis=len(pos_dims))

    # This is just the visualizer:
    loop_field_names = fit_nd.dtype.names
    loop_field = loop_field_names[0]
    loop_ind = 0
    row_ind = 0
    col_ind = 0

    # Initial plot data
645
646
    spatial_map = fit_3d[:, :, loop_ind][loop_field]
    proj_data = proj_nd_shifted[col_ind, row_ind, :, loop_ind]
647
    bias_data = bias_shifted[:, loop_ind]
648
649
    guess_data = loop_fit_function(bias_data, np.array(list(guess_3d[col_ind, row_ind, loop_ind])))
    fit_data = loop_fit_function(bias_data, np.array(list(fit_3d[col_ind, row_ind, loop_ind])))
650
651
652
653
654

    fig = plt.figure(figsize=(12, 8))
    ax_map = plt.subplot2grid((1, 2), (0, 0), colspan=1, rowspan=1)
    ax_loop = plt.subplot2grid((1, 2), (0, 1), colspan=1, rowspan=1)

Unknown's avatar
Unknown committed
655
656
    im_map, im_cbar = single_img_cbar_plot(ax_map, spatial_map.T, x_size=spatial_map.shape[0],
                                           y_size=spatial_map.shape[1], cmap=cmap)
Unknown's avatar
Unknown committed
657

658
659
    ax_map.set_xlabel('X')
    ax_map.set_ylabel('Y')
Unknown's avatar
Unknown committed
660
    ax_map.set_title('{} - Loop {}'.format(loop_field, loop_ind))
661
    crosshair = ax_map.plot(row_ind, col_ind, 'k+')[0]
662
663
664
665
666
667
668
669
670
671
672
673
674
675

    ax_loop.plot(bias_data, proj_data, 'k', label='Projection')
    ax_loop.plot(bias_data, guess_data, 'g', label='Guess')
    ax_loop.plot(bias_data, fit_data, 'r--', label='Fit')
    line_handles = ax_loop.get_lines()

    ax_loop.tick_params(labelleft=False, labelright=True)
    ax_loop.yaxis.set_label_position('right')
    ax_loop.set_ylabel('PR (a.u.)')
    ax_loop.set_xlabel('V_DC (V)')
    ax_loop.set_title('Position ({},{})'.format(col_ind, row_ind))
    ax_loop.legend()
    fig.tight_layout()

676
677
    plt.show()

Unknown's avatar
Unknown committed
678
679
    display(save_fig_filebox_button(fig, h5_projected_loops.file.filename))

Unknown's avatar
Unknown committed
680
    loop_slider = widgets.IntSlider(min=0, max=bias_vec.shape[1]-1, description='Loop Number')
Unknown's avatar
Unknown committed
681

Unknown's avatar
Unknown committed
682
    def update_loop_plots(loop_field, loop_ind):
683
684
        spatial_map = fit_3d[:, :, loop_ind][loop_field]
        im_map.set_data(spatial_map.T)
685
686
687
        spat_mean = np.mean(spatial_map)
        spat_std = np.std(spatial_map)
        im_map.set_clim(vmin=spat_mean - 3 * spat_std, vmax=spat_mean + 3 * spat_std)
Unknown's avatar
Unknown committed
688
        ax_map.set_title('{} - Loop {}'.format(loop_field, loop_ind))
689

Unknown's avatar
Unknown committed
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
        xdata, ydata = crosshair.get_xydata().flatten().astype('int')

        proj_data = proj_nd_shifted[xdata, ydata, :, loop_ind]
        bias_data = bias_shifted[:, loop_ind]
        guess_data = loop_fit_function(bias_data, np.array(list(guess_3d[xdata, ydata, loop_ind])))
        fit_data = loop_fit_function(bias_data, np.array(list(fit_3d[xdata, ydata, loop_ind])))

        for line_handle, data in zip(line_handles, [proj_data, guess_data, fit_data]):
            line_handle.set_ydata(data)

        ax_loop.set_title('Position ({},{})'.format(xdata, ydata))

        ax_loop.relim()
        ax_loop.autoscale_view()

705
706
707
708
709
710
711
712
713
714
715
716
717
718
        fig.canvas.draw()

    def pos_picker(event):
        if not im_map.axes.in_axes(event):
            return

        xdata = int(round(event.xdata))
        ydata = int(round(event.ydata))
        current_pos = {pos_labels[0]: xdata, pos_labels[1]: ydata}

        pos_dim_vals = list(range(len(pos_labels)))

        for pos_dim_ind, pos_dim_name in enumerate(pos_labels):
            pos_dim_vals[pos_dim_ind] = current_pos[pos_dim_name]
719

720
721
722
723
        crosshair.set_xdata(xdata)
        crosshair.set_ydata(ydata)

        proj_data = proj_nd_shifted[xdata, ydata, :, loop_ind]
724
        bias_data = bias_shifted[:, loop_ind]
725
726
        guess_data = loop_fit_function(bias_data, np.array(list(guess_3d[xdata, ydata, loop_ind])))
        fit_data = loop_fit_function(bias_data, np.array(list(fit_3d[xdata, ydata, loop_ind])))
727
728
        for line_handle, data in zip(line_handles, [proj_data, guess_data, fit_data]):
            line_handle.set_ydata(data)
729
730
        ax_loop.set_title('Position ({},{})'.format(xdata, ydata))

Unknown's avatar
Unknown committed
731
732
733
        ax_loop.relim()
        ax_loop.autoscale_view()

734
735
736
        fig.canvas.draw()

    cid = im_map.figure.canvas.mpl_connect('button_press_event', pos_picker)
737

Unknown's avatar
Unknown committed
738
    widgets.interact(update_loop_plots, loop_field=list(fit_nd.dtype.names), loop_ind=loop_slider)
Unknown's avatar
Unknown committed
739
740


741
def jupyter_visualize_parameter_maps(h5_loop_parameters, cmap=None, **kwargs):
Unknown's avatar
Unknown committed
742
    """
Unknown's avatar
Unknown committed
743
    Interactive plot of the spatial maps of the loop parameters for all cycles.
Unknown's avatar
Unknown committed
744
745
746

    Parameters
    ----------
Unknown's avatar
Unknown committed
747
748
    h5_loop_parameters : h5py.Dataset
        The dataset containing the loop parameters to be visualized
Unknown's avatar
Unknown committed
749
    cmap : str or matplotlib.colors.Colormap
Unknown's avatar
Unknown committed
750
751
752

    Returns
    -------
Unknown's avatar
Unknown committed
753
    None
Unknown's avatar
Unknown committed
754
755
756
757
758
759
760
761
762
763
764
765

    """
    # Get the position and spectroscopic datasets
    h5_loop_pos_inds = getAuxData(h5_loop_parameters, 'Position_Indices')[-1]
    h5_loop_spec_vals = getAuxData(h5_loop_parameters, 'Spectroscopic_Values')[-1]

    pos_dims = get_dimensionality(np.transpose(h5_loop_pos_inds))
    num_cycles = h5_loop_parameters.shape[1]

    parameter_names = h5_loop_parameters.dtype.names

    parameter_map_stack = np.reshape(h5_loop_parameters[parameter_names[0]],
Unknown's avatar
Unknown committed
766
                                     [pos_dims[0], pos_dims[1], -1])
Unknown's avatar
Unknown committed
767

Unknown's avatar
Unknown committed
768
    loop_spec_labs = get_attr(h5_loop_spec_vals, 'labels')
Unknown's avatar
Unknown committed
769

770
771
    kwargs.update({'cmap': get_cmap_object(cmap)})

Unknown's avatar
Unknown committed
772
773
    map_titles = list()
    for icycle in range(num_cycles):
Unknown's avatar
Unknown committed
774
775
776
777
778
        title_list = list()
        for label in loop_spec_labs:
            val = h5_loop_spec_vals[get_attr(h5_loop_spec_vals, label)].squeeze()[icycle]
            title_list.append('{}: {}'.format(label, val))
        map_titles.append(' - '.join(title_list))
Unknown's avatar
Unknown committed
779
780

    fig, axes = plot_map_stack(parameter_map_stack, num_comps=num_cycles, color_bar_mode='each',
781
782
                               title=map_titles, heading='Maps of Loop Parameter {}'.format(parameter_names[0]),
                               **kwargs)
Unknown's avatar
Unknown committed
783
784
785

    def update_loop_maps(parameter_name):
        parameter_map_stack = np.reshape(h5_loop_parameters[parameter_name],
786
                                         [pos_dims[0], pos_dims[1], -1])
Unknown's avatar
Unknown committed
787
788
789
790
791
        fig.suptitle('Maps of Loop Parameter {}'.format(parameter_name))
        # Loop over all axes
        for icycle, ax_cycle in enumerate(axes):
            image = ax_cycle.get_images()[0]
            image.set_data(parameter_map_stack[:, :, icycle])
Unknown's avatar
Unknown committed
792
793
            image.set_clim(vmin=np.min(parameter_map_stack[:, :, icycle]),
                           vmax=np.max(parameter_map_stack[:, :, icycle]))
Unknown's avatar
Unknown committed
794

795
796
797
        fig.canvas.draw()

    plt.show()
Unknown's avatar
Unknown committed
798
799
800

    display(save_fig_filebox_button(fig, h5_loop_parameters.file.filename))
    widgets.interact(update_loop_maps, parameter_name=list(parameter_names))
Unknown's avatar
Unknown committed
801
802


803
def jupyter_visualize_loop_sho_raw_comparison(h5_loop_parameters, cmap=None):
Unknown's avatar
Unknown committed
804
805
806
807
808
    """

    Parameters
    ----------
    h5_loop_parameters
Unknown's avatar
Unknown committed
809
    cmap : str or matplotlib.colors.Colormap
Unknown's avatar
Unknown committed
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951

    Returns
    -------

    """

    # Find the precursor datasets used to calculate these parameters
    h5_loop_grp = h5_loop_parameters.parent
    h5_loop_projections = h5_loop_grp['Projected_Loops']
    h5_loop_fit = h5_loop_grp['Fit']
    h5_loop_guess = h5_loop_grp['Guess']

    h5_sho_grp = h5_loop_grp.parent
    h5_sho_fit = h5_sho_grp['Fit']
    h5_sho_guess = h5_sho_grp['Guess']


    h5_main = get_source_dataset(h5_sho_grp)

    # Now get the needed ancillary datasets for each main dataset
    h5_pos_inds = getAuxData(h5_loop_parameters, 'Position_Indices')[0]
    h5_pos_vals = getAuxData(h5_loop_parameters, 'Position_Values')[0]
    pos_order = get_sort_order(np.transpose(h5_pos_inds))
    pos_dims = get_dimensionality(np.transpose(h5_pos_inds), pos_order)
    pos_labs = get_attr(h5_pos_inds, 'labels')

    h5_loop_spec_inds = getAuxData(h5_loop_parameters, 'Spectroscopic_Indices')[0]
    h5_loop_spec_vals = getAuxData(h5_loop_parameters, 'Spectroscopic_Values')[0]
    loop_spec_order = get_sort_order(h5_loop_spec_inds)
    loop_spec_dims = get_dimensionality(h5_loop_spec_inds, loop_spec_order)
    loop_spec_labs = get_attr(h5_loop_spec_inds, 'labels')

    h5_sho_spec_inds = getAuxData(h5_sho_fit, 'Spectroscopic_Indices')[0]
    h5_sho_spec_vals = getAuxData(h5_sho_fit, 'Spectroscopic_Values')[0]
    sho_spec_order = get_sort_order(h5_sho_spec_inds)
    sho_spec_dims = get_dimensionality(h5_sho_spec_inds, sho_spec_order)
    sho_spec_labs = get_attr(h5_sho_spec_inds, 'labels')

    h5_main_spec_inds = getAuxData(h5_main, 'Spectroscopic_Indices')[0]
    h5_main_spec_vals = getAuxData(h5_main, 'Spectroscopic_Values')[0]
    main_spec_order = get_sort_order(h5_main_spec_inds)
    main_spec_dims = get_dimensionality(h5_main_spec_inds, main_spec_order)
    main_spec_labs = get_attr(h5_main_spec_inds, 'labels')

    '''
    Select the initial plotting slices
    '''
    loop_parameter_names = h5_loop_parameters.dtype.names
    loop_num_cycles = h5_loop_parameters.shape[1]
    loop_parameter_spec_labs = get_attr(h5_loop_spec_vals, 'labels')
    sho_bias_dim = np.argwhere(sho_spec_labs[sho_spec_order] == 'DC_Offset').squeeze()
    steps_per_loop = sho_spec_dims[sho_bias_dim]
    main_bias_dim = np.argwhere(main_spec_labs[main_spec_order] == 'DC_Offset').squeeze()

    selected_loop_parm = loop_parameter_names[0]
    selected_loop_cycle = 0
    selected_loop_ndims = np.unravel_index(selected_loop_cycle, loop_spec_dims, order='F')
    selected_loop_pos = int(pos_dims[0]/2), int(pos_dims[1]/2)
    selected_step = int(steps_per_loop/2)

    '''
    Get the bias vector to be plotted against
    '''
    loop_bias_vec = h5_sho_spec_vals[get_attr(h5_sho_spec_vals, 'DC_Offset')].squeeze()
    shift_ind = int(-1 * steps_per_loop / 4)
    loop_bias_vec = loop_bias_vec.reshape(sho_spec_dims)
    loop_bias_vec = np.moveaxis(loop_bias_vec, sho_bias_dim, 0).reshape(sho_spec_dims[sho_bias_dim], -1)
    loop_bias_vec = np.roll(loop_bias_vec.reshape(steps_per_loop, -1), shift_ind, axis=0)

    '''
    Get the frequency vector to be plotted against
    '''
    full_w_vec = h5_main_spec_vals[h5_main_spec_vals.attrs['Frequency']]
    full_w_vec, _ = reshape_to_Ndims(full_w_vec, h5_spec=h5_main_spec_inds)
    full_w_vec = full_w_vec.squeeze()

    '''
    Define functions to get the data
    '''
    def _get_loop_map(selected_loop_parm, selected_loop_cycle):
        # Build the map of the chosen loop parameter
        loop_parameter_map = np.reshape(h5_loop_parameters[selected_loop_parm, :, selected_loop_cycle],
                                        [pos_dims[0], pos_dims[1]])

        # Also create the title string for the map
        loop_map_title = list()
        for label in loop_parameter_spec_labs:
            val = h5_loop_spec_vals[get_attr(h5_loop_spec_vals, label)].squeeze()[selected_loop_cycle]
            loop_map_title.append('{}: {}'.format(label, val))

        return loop_parameter_map, loop_map_title

    def _get_loops(selected_loop_cycle, selected_loop_pos):
        selected_loop_ndims = np.unravel_index(selected_loop_cycle, loop_spec_dims, order='F')
        # Now build the loop plot for the selected position in the loop map
        selected_loop_bias_vec = loop_bias_vec[:, selected_loop_cycle]

        pos_ind = np.ravel_multi_index(selected_loop_pos, pos_dims)

        loop_proj_vec = h5_loop_projections[pos_ind].reshape(sho_spec_dims[::-1])
        loop_proj_vec = np.moveaxis(loop_proj_vec, sho_bias_dim, -1)[selected_loop_ndims]
        loop_proj_vec = np.roll(loop_proj_vec, shift_ind)
        loop_guess_vec = loop_fit_function(selected_loop_bias_vec,
                                           h5_loop_guess[pos_ind, selected_loop_cycle].tolist())
        loop_fit_vec = loop_fit_function(selected_loop_bias_vec,
                                           h5_loop_fit[pos_ind, selected_loop_cycle].tolist())

        return selected_loop_bias_vec, loop_proj_vec, loop_guess_vec, loop_fit_vec

    def _get_sho(selected_loop_pos, selected_step, selected_loop_cycle):
        selected_loop_ndims = np.unravel_index(selected_loop_cycle, loop_spec_dims, order='F')
        # get the SHO Guess and Fit and Raw Data for the selected position, cycle, and step
        pos_ind = np.ravel_multi_index(selected_loop_pos, pos_dims)

        # Get the frequency vector for the selected step
        w_vec = np.moveaxis(full_w_vec, main_bias_dim - len(pos_dims), -1)[selected_step][selected_loop_ndims]

        # Get the slice of the sho guess and fit
        sho_guess = h5_sho_guess[pos_ind].reshape(sho_spec_dims[::-1])
        sho_guess = np.moveaxis(sho_guess, sho_bias_dim, -1)[selected_loop_ndims][selected_step]
        sho_guess = SHOfunc(sho_guess, w_vec)
        sho_fit = h5_sho_fit[pos_ind].reshape(sho_spec_dims[::-1])
        sho_fit = np.moveaxis(sho_fit, sho_bias_dim, -1)[selected_loop_ndims][selected_step]
        sho_fit = SHOfunc(sho_fit, w_vec)

        # Get the slice of the Raw Data
        raw_data_vec, _ = reshape_to_Ndims(np.atleast_2d(h5_main[pos_ind]), h5_spec=h5_main_spec_inds)
        raw_data_vec = np.moveaxis(raw_data_vec.squeeze(),
                                   main_bias_dim - len(pos_dims), -1)[selected_step][selected_loop_ndims]

        return w_vec, sho_guess, sho_fit, raw_data_vec

    '''
    Get the starting values
    '''
    loop_parm_map, loop_map_title = _get_loop_map(selected_loop_parm, selected_loop_cycle)
    current_loop_bias_vec, current_loop_proj_vec, current_loop_guess_vec, current_loop_fit_vec = \
        _get_loops(selected_loop_cycle, selected_loop_pos)
    current_w_vec, current_sho_guess, current_sho_fit, current_raw_data = _get_sho(selected_loop_pos,
                                                                                   selected_step,
                                                                                   selected_loop_cycle)

952
953
    cmap = get_cmap_object(cmap)

Unknown's avatar
Unknown committed
954
955
956
957
958
959
960
961
962
963
964
    '''
    Build the figure
    '''
    fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(17, 17))

    ax_loop_map = axes.flatten()[0]
    ax_loop_pgf = axes.flatten()[1]
    ax_sho_amp = axes.flatten()[2]
    ax_sho_phase = axes.flatten()[3]

    # Plot the map of the loop parameters
965
    plt_loop_map = single_img_cbar_plot(ax_loop_map, loop_parm_map, cmap=cmap)
Unknown's avatar
Unknown committed
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
    loop_vert_line = ax_loop_map.axvline(x=pos_dims[0], color='k')
    loop_horz_line = ax_loop_map.axhline(y=pos_dims[1], color='k')
    ax_loop_map.set_title(loop_map_title)

    # Plot the Loops
    plt_loop_proj = ax_loop_pgf.plot(current_loop_bias_vec, current_loop_proj_vec, 'k', label='Loop Projection')
    plt_loop_guess = ax_loop_pgf.plot(current_loop_bias_vec, current_loop_guess_vec, 'g', label='Loop Guess')
    plt_loop_fit = ax_loop_pgf.plot(current_loop_bias_vec, current_loop_fit_vec, 'r--', label='Loop Fit')
    ax_loop_pgf.legend()

    # Plot the amplitude of the SHO Guess, SHO Fit, and Raw Data
    plt_raw_data_amp = ax_sho_amp.plot(current_w_vec, np.abs(current_raw_data), 'k', label='Raw Data')
    plt_sho_guess_amp = ax_sho_amp.plot(current_w_vec, np.abs(current_sho_guess), 'g', label='SHO Guess')
    plt_sho_fit_amp = ax_sho_amp.plot(current_w_vec, np.abs(current_sho_fit), 'r--', label='SHO Fit')
    ax_sho_amp.legend()

    # Plot the phase of the SHO Guess, SHO Fit, and Raw Data
    plt_raw_data_phase = ax_sho_phase.plot(current_w_vec, np.angle(current_raw_data), 'k', label='Raw Data')
    plt_sho_guess_phase = ax_sho_phase.plot(current_w_vec, np.angle(current_sho_guess), 'g', label='SHO Guess')
    plt_sho_fit_phase = ax_sho_phase.plot(current_w_vec, np.angle(current_sho_fit), 'r--', label='SHO Fit')
    ax_sho_phase.legend()

    fig.tight_layout()

    def _update_loop_map(parameter_name, selected_cycle):
        loop_parm_map, loop_map_title = _get_loop_map(parameter_name, selected_cycle)
        plt_loop_map.set_data(loop_parm_map)
        ax_loop_map.set_title(loop_map_title)

    def _update_crosshairs(x_pos, y_pos):
        loop_vert_line.set_xdata(x_pos)
        loop_horz_line.set_ydata(y_pos)

    def _update_loop_plots(selected_loop_cycle, x_pos, y_pos):
        current_loop_bias_vec, current_loop_proj_vec, current_loop_guess_vec, current_loop_fit_vec = \
            _get_loops(selected_loop_cycle, [x_pos, y_pos])

        # Update the bias vectors
        plt_loop_proj.set_xdata(current_loop_bias_vec)
        plt_loop_guess.set_xdata(current_loop_bias_vec)
        plt_loop_fit.set_xdata(current_loop_bias_vec)

        # Update the ploted data
        plt_loop_proj.set_xdata(current_loop_proj_vec)
        plt_loop_guess.set_xdata(current_loop_guess_vec)
        plt_loop_fit.set_xdata(current_loop_fit_vec)

    def _update_sho_plots(selected_loop_cycle, x_pos, y_pos, selected_step):
        current_w_vec, current_sho_guess, current_sho_fit, current_raw_data = _get_sho([x_pos, y_pos],
                                                                                       selected_step,
                                                                                       selected_loop_cycle)
        # Update x-date with new w_vec
        plt_raw_data_amp.set_xdata(current_w_vec)
        plt_raw_data_phase.set_xdata(current_w_vec)
        plt_sho_guess_amp.set_xdata(current_w_vec)
        plt_sho_guess_phase.set_xdata(current_w_vec)
        plt_sho_fit_amp.set_xdata(current_w_vec)
        plt_sho_fit_phase.set_xdata(current_w_vec)

        # Update the y-data
        plt_raw_data_amp.set_ydata(np.amp(current_raw_data))
        plt_raw_data_phase.set_ydata(np.angle(current_raw_data))
        plt_sho_guess_amp.set_ydata(np.amp(current_sho_guess))
        plt_sho_guess_phase.set_ydata(np.angle(current_sho_guess))
        plt_sho_fit_amp.set_ydata(np.amp(current_sho_fit))
        plt_sho_fit_phase.set_ydata(np.angle(current_sho_fit))


    display(save_fig_filebox_button(fig, h5_loop_parameters.file.filename))
    display(fig)
    x_pos_widget = widgets.FloatSlider(min=0.0, max=float(pos_dims[0]), step=pos_dims[0]/100, value=selected_loop_pos[0])
    y_pos_widget = widgets.FloatSlider(min=0.0, max=float(pos_dims[1]), step=pos_dims[1] / 100, value=selected_loop_pos[1])
    loop_cycle_widget = widgets.IntSlider(min=0, max=loop_num_cycles, step=1, value=selected_loop_cycle)
    spec_step_widget = widgets.IntSlider(min=0, max=steps_per_loop, step=1, value=selected_step)
    widgets.interact(_update_loop_map,
                     parameter_name=list(loop_parameter_names),
                     selected_cycle=loop_cycle_widget)
    widgets.interact(_update_crosshairs,
                     x_pos=x_pos_widget, y_pos=y_pos_widget)
    widgets.interact(_update_loop_plots,
                     selected_loop_cycle=loop_cycle_widget,
                     x_pos=x_pos_widget, y_pos=y_pos_widget)
    widgets.interact(_update_sho_plots,
                     selected_loop_cycle=loop_cycle_widget,
                     x_pos=x_pos_widget, y_pos=y_pos_widget,
                     selected_step=spec_step_widget)


def plot_loop_sho_raw_comparison(h5_loop_parameters, selected_loop_parm=None, selected_loop_cycle=0,
Unknown's avatar
Unknown committed
1055
                                 selected_loop_pos=[0,0], selected_step=0, tick_font_size=14, cmap='viridis'):
Unknown's avatar
Unknown committed
1056
1057
1058
1059
    """

    Parameters
    ----------
Unknown's avatar
Unknown committed
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
    h5_loop_parameters : h5py.Dataset
        Dataset containing the loop parameters
    selected_loop_parm : str
        The initial loop parameter to be plotted
    selected_loop_cycle : int
        The initial loop cycle to be plotted
    selected_loop_pos : array-like of two ints
        The initial position to be plotted
    selected_step : int
        The initial bias step to be plotted
    tick_font_size : 14
    cmap : str or matplotlib.colors.Colormap
        Colormap to be used in plotting the parameter map
Unknown's avatar
Unknown committed
1073
1074
1075

    Returns
    -------
Unknown's avatar
Unknown committed
1076
    None
Unknown's avatar
Unknown committed
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126

    """

    # Find the precursor datasets used to calculate these parameters
    h5_loop_grp = h5_loop_parameters.parent
    h5_loop_projections = h5_loop_grp['Projected_Loops']
    h5_loop_fit = h5_loop_grp['Fit']
    h5_loop_guess = h5_loop_grp['Guess']

    h5_sho_grp = h5_loop_grp.parent
    h5_sho_fit = h5_sho_grp['Fit']
    h5_sho_guess = h5_sho_grp['Guess']


    h5_main = get_source_dataset(h5_sho_grp)

    # Now get the needed ancillary datasets for each main dataset
    h5_pos_inds = getAuxData(h5_loop_parameters, 'Position_Indices')[0]
    h5_pos_vals = getAuxData(h5_loop_parameters, 'Position_Values')[0]
    pos_order = get_sort_order(np.transpose(h5_pos_inds))
    pos_dims = get_dimensionality(np.transpose(h5_pos_inds), pos_order)
    pos_labs = get_attr(h5_pos_inds, 'labels')

    h5_loop_spec_inds = getAuxData(h5_loop_parameters, 'Spectroscopic_Indices')[0]
    h5_loop_spec_vals = getAuxData(h5_loop_parameters, 'Spectroscopic_Values')[0]
    loop_spec_order = get_sort_order(h5_loop_spec_inds)
    loop_spec_dims = get_dimensionality(h5_loop_spec_inds, loop_spec_order)
    loop_spec_labs = get_attr(h5_loop_spec_inds, 'labels')

    h5_sho_spec_inds = getAuxData(h5_sho_fit, 'Spectroscopic_Indices')[0]
    h5_sho_spec_vals = getAuxData(h5_sho_fit, 'Spectroscopic_Values')[0]
    sho_spec_order = get_sort_order(h5_sho_spec_inds)
    sho_spec_dims = get_dimensionality(h5_sho_spec_inds, sho_spec_order)
    sho_spec_labs = get_attr(h5_sho_spec_inds, 'labels')

    h5_main_spec_inds = getAuxData(h5_main, 'Spectroscopic_Indices')[0]
    h5_main_spec_vals = getAuxData(h5_main, 'Spectroscopic_Values')[0]
    main_spec_order = get_sort_order(h5_main_spec_inds)
    main_spec_dims = get_dimensionality(h5_main_spec_inds, main_spec_order)
    main_spec_labs = get_attr(h5_main_spec_inds, 'labels')

    '''
    Select the initial plotting slices
    '''
    loop_parameter_names = h5_loop_parameters.dtype.names
    loop_num_cycles = h5_loop_parameters.shape[1]
    loop_parameter_spec_labs = get_attr(h5_loop_spec_vals, 'labels')
    sho_bias_dim = np.argwhere(sho_spec_labs[sho_spec_order] == 'DC_Offset').squeeze()
    steps_per_loop = sho_spec_dims[sho_bias_dim]
    main_bias_dim = np.argwhere(main_spec_labs[main_spec_order] == 'DC_Offset').squeeze()
Unknown's avatar
Unknown committed
1127
    main_freq_dim = np.argwhere(main_spec_labs[main_spec_order] == 'Frequency').squeeze()
Unknown's avatar
Unknown committed
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137

    if selected_loop_parm is None:
        selected_loop_parm = loop_parameter_names[0]
    selected_loop_ndims = np.unravel_index(selected_loop_cycle, loop_spec_dims, order='F')

    '''
    Get the bias vector to be plotted against
    '''
    loop_bias_vec = h5_sho_spec_vals[get_attr(h5_sho_spec_vals, 'DC_Offset')].squeeze()
    shift_ind = int(-1 * steps_per_loop / 4)
Unknown's avatar
Unknown committed
1138
1139
1140
    loop_bias_vec = loop_bias_vec.reshape(sho_spec_dims[::-1])
    loop_bias_vec = np.moveaxis(loop_bias_vec, len(loop_bias_vec.shape)-sho_bias_dim-1, 0)
    loop_bias_vec = np.reshape(loop_bias_vec, [sho_spec_dims[sho_bias_dim], -1])
Unknown's avatar
Unknown committed
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
    loop_bias_vec = np.roll(loop_bias_vec.reshape(steps_per_loop, -1), shift_ind, axis=0)

    '''
    Get the frequency vector to be plotted against
    '''
    full_w_vec = h5_main_spec_vals[h5_main_spec_vals.attrs['Frequency']]
    full_w_vec, _ = reshape_to_Ndims(full_w_vec, h5_spec=h5_main_spec_inds)
    full_w_vec = full_w_vec.squeeze()

    '''
    Define functions to get the data
    '''
    def _get_loop_map(selected_loop_parm, selected_loop_cycle):
        # Build the map of the chosen loop parameter
        loop_parameter_map = np.reshape(h5_loop_parameters[selected_loop_parm, :, selected_loop_cycle],
                                        [pos_dims[0], pos_dims[1]])

        # Also create the title string for the map
        loop_map_title = list()
        for label in loop_parameter_spec_labs:
            val = h5_loop_spec_vals[get_attr(h5_loop_spec_vals, label)].squeeze()[selected_loop_cycle]
            loop_map_title.append('{}: {}'.format(label, val))

        loop_map_title = ' - '.join(loop_map_title)

        return loop_parameter_map, loop_map_title

    def _get_loops(selected_loop_cycle, selected_loop_pos):
        selected_loop_ndims = np.unravel_index(selected_loop_cycle, loop_spec_dims, order='F')
        # Now build the loop plot for the selected position in the loop map
        selected_loop_bias_vec = loop_bias_vec[:, selected_loop_cycle]

        pos_ind = np.ravel_multi_index(selected_loop_pos, pos_dims)

        loop_proj_vec = h5_loop_projections[pos_ind].reshape(sho_spec_dims[::-1])
Unknown's avatar
Unknown committed
1176
        loop_proj_vec = np.moveaxis(loop_proj_vec, len(sho_spec_dims)-sho_bias_dim-1, -1)[selected_loop_ndims]
Unknown's avatar
Unknown committed
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
        loop_proj_vec = np.roll(loop_proj_vec, shift_ind)
        loop_guess_vec = loop_fit_function(selected_loop_bias_vec,
                                           h5_loop_guess[pos_ind, selected_loop_cycle].tolist())
        loop_fit_vec = loop_fit_function(selected_loop_bias_vec,
                                           h5_loop_fit[pos_ind, selected_loop_cycle].tolist())

        return selected_loop_bias_vec, loop_proj_vec, loop_guess_vec, loop_fit_vec

    def _get_sho(selected_loop_pos, selected_step, selected_loop_cycle):
        selected_loop_ndims = np.unravel_index(selected_loop_cycle, loop_spec_dims, order='F')
        # get the SHO Guess and Fit and Raw Data for the selected position, cycle, and step
        pos_ind = np.ravel_multi_index(selected_loop_pos, pos_dims)

        # Get the frequency vector for the selected step
Unknown's avatar
Unknown committed
1191
        w_vec = np.moveaxis(full_w_vec, main_freq_dim, -1)[selected_step][selected_loop_ndims]
Unknown's avatar
Unknown committed
1192
1193
1194

        # Get the slice of the sho guess and fit
        sho_guess = h5_sho_guess[pos_ind].reshape(sho_spec_dims[::-1])
Unknown's avatar
Unknown committed
1195
        sho_guess = np.moveaxis(sho_guess, len(sho_spec_dims)-sho_bias_dim-1, -1)[selected_loop_ndims][selected_step]
Unknown's avatar
Unknown committed
1196
1197
        sho_guess = SHOfunc(sho_guess, w_vec)
        sho_fit = h5_sho_fit[pos_ind].reshape(sho_spec_dims[::-1])
Unknown's avatar
Unknown committed
1198
        sho_fit = np.moveaxis(sho_fit, len(sho_spec_dims)-sho_bias_dim-1, -1)[selected_loop_ndims][selected_step]
Unknown's avatar
Unknown committed
1199
1200
1201
1202
1203
        sho_fit = SHOfunc(sho_fit, w_vec)

        # Get the slice of the Raw Data
        raw_data_vec, _ = reshape_to_Ndims(np.atleast_2d(h5_main[pos_ind]), h5_spec=h5_main_spec_inds)
        raw_data_vec = np.moveaxis(raw_data_vec.squeeze(),
Unknown's avatar
Unknown committed
1204
                                   main_freq_dim, -1)[selected_step][selected_loop_ndims]
Unknown's avatar
Unknown committed
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228

        return w_vec, sho_guess, sho_fit, raw_data_vec

    '''
    Get the starting values
    '''
    loop_parm_map, loop_map_title = _get_loop_map(selected_loop_parm, selected_loop_cycle)
    current_loop_bias_vec, current_loop_proj_vec, current_loop_guess_vec, current_loop_fit_vec = \
        _get_loops(selected_loop_cycle, selected_loop_pos)
    current_w_vec, current_sho_guess, current_sho_fit, current_raw_data = _get_sho(selected_loop_pos,
                                                                                   selected_step,
                                                                                   selected_loop_cycle)

    '''
    Build the figure
    '''
    fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(17, 17))

    ax_loop_map = axes.flatten()[0]
    ax_loop_pgf = axes.flatten()[1]
    ax_sho_amp = axes.flatten()[2]
    ax_sho_phase = axes.flatten()[3]

    # Plot the map of the loop parameters
1229
1230
    loop_map, loop_map_cbar = single_img_cbar_plot(ax_loop_map, loop_parm_map)
    crosshair = ax_loop_map.plot(selected_loop_pos[0], selected_loop_pos[1], 'k+')[0]
Unknown's avatar
Unknown committed
1231
1232
1233
1234
1235
    ax_loop_map.set_title(loop_map_title)
    ax_loop_map.set_xlabel('X Position')
    ax_loop_map.set_ylabel('Y Position')

    # Plot the Loops
1236
1237
1238
    loop_proj_plot = ax_loop_pgf.plot(current_loop_bias_vec, current_loop_proj_vec, 'k', label='Loop Projection')[0]
    loop_guess_plot = ax_loop_pgf.plot(current_loop_bias_vec, current_loop_guess_vec, 'g', label='Loop Guess')[0]
    loop_fit_plot = ax_loop_pgf.plot(current_loop_bias_vec, current_loop_fit_vec, 'r--', label='Loop Fit')[0]
Unknown's avatar
Unknown committed
1239
1240
1241
1242
1243
1244
1245
1246
1247
    ax_loop_pgf.legend()
    loop_pgf_title = ' - '.join(['Position: ({}, {})'.format(selected_loop_pos[0], selected_loop_pos[1]),
                                 loop_map_title])
    ax_loop_pgf.set_title(loop_pgf_title)
    ax_loop_pgf.set_ylabel('PR (a.u.)')
    ax_loop_pgf.set_xlabel('V_DC (V)')
    ax_loop_pgf.axis('tight')

    # Plot the amplitude of the SHO Guess, SHO Fit, and Raw Data
1248
1249
1250
    raw_amp_plot = ax_sho_amp.plot(current_w_vec, np.abs(current_raw_data), 'k*', label='Raw Data Amplitude')[0]
    guess_amp_plot = ax_sho_amp.plot(current_w_vec, np.abs(current_sho_guess), 'g', label='SHO Guess Amplitude')[0]
    fit_amp_plot = ax_sho_amp.plot(current_w_vec, np.abs(current_sho_fit), 'r:', label='SHO Fit Amplitude')[0]
Unknown's avatar
Unknown committed
1251
1252
1253
1254
1255
1256
1257
1258
    ax_sho_amp.legend()
    sho_title = ' - '.join(['Bias Step: {}'.format(selected_step), loop_pgf_title])
    ax_sho_amp.set_title(sho_title)
    ax_sho_amp.set_ylabel('Amplitude (V)')
    ax_sho_amp.set_xlabel('Frequency (Hz)')
    ax_sho_amp.axis('tight')

    # Plot the phase of the SHO Guess, SHO Fit, and Raw Data
1259
1260
1261
    raw_phase_plot = ax_sho_phase.plot(current_w_vec, np.angle(current_raw_data), 'k*', label='Raw Data Phase')[0]
    guess_phase_plot = ax_sho_phase.plot(current_w_vec, np.angle(current_sho_guess), 'g', label='SHO Guess Phase')[0]
    fit_phase_plot = ax_sho_phase.plot(current_w_vec, np.angle(current_sho_fit), 'r:', label='SHO Fit Phase')[0]
Unknown's avatar
Unknown committed
1262
1263
1264
1265
1266
1267
1268
1269
    ax_sho_phase.legend()
    ax_sho_phase.set_title(sho_title)
    ax_sho_phase.set_ylabel('Amplitude (V)')
    ax_sho_phase.set_xlabel('Frequency (Hz)')
    ax_sho_phase.axis('tight')

    for axis in [ax_loop_pgf, ax_sho_amp, ax_sho_phase]:
        set_tick_font_size(axis, tick_font_size)
1270
        axis.ticklabel_format(axis='both', style='sci', scilimits=(-2, 2))
Unknown's avatar
Unknown committed
1271

1272
    plt.show()
Unknown's avatar
Unknown committed
1273

1274
1275
1276
1277
    # Make some widgets now so we can reference them in updates
    loop_parm_widget = widgets.Dropdown(options=loop_parameter_names, value=selected_loop_parm)
    loop_cycle_widget = widgets.IntSlider(min=0, max=loop_num_cycles - 1, step=1, value=selected_loop_cycle)
    spec_step_widget = widgets.IntSlider(min=0, max=steps_per_loop - 1, step=1, value=selected_step)
Unknown's avatar
Unknown committed
1278

1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
    def _update_loop_plots(selected_loop_cycle, xdata, ydata):
        # Get new data
        current_loop_bias_vec, current_loop_proj_vec, current_loop_guess_vec, current_loop_fit_vec = \
            _get_loops(selected_loop_cycle, (xdata, ydata))

        # Update loop plots
        loop_proj_plot.set_xdata(current_loop_bias_vec)
        loop_proj_plot.set_ydata(current_loop_proj_vec)
        loop_guess_plot.set_xdata(current_loop_bias_vec)
        loop_guess_plot.set_ydata(current_loop_guess_vec)
        loop_fit_plot.set_xdata(current_loop_bias_vec)
        loop_fit_plot.set_ydata(current_loop_fit_vec)
        loop_map_title = ax_loop_map.get_title()
        loop_pgf_title = ' - '.join(['Position: ({}, {})'.format(xdata, ydata),
                                     loop_map_title])
        ax_loop_pgf.set_title(loop_pgf_title)
        loop_min = np.min([current_loop_proj_vec, current_loop_guess_vec, current_loop_fit_vec])
        loop_max = np.max([current_loop_proj_vec, current_loop_guess_vec, current_loop_fit_vec])
        loop_min *= (1 - np.sign(loop_min) * 0.5)
        loop_max *= (1 + np.sign(loop_max) * 0.5)

        ax_loop_pgf.set_ylim([loop_min, loop_max])

    def _update_sho_plots(selected_loop_cycle, selected_step, xdata, ydata):
        # Get the updated data
        current_w_vec, current_sho_guess, current_sho_fit, current_raw_data = _get_sho((xdata, ydata),
                                                                                       selected_step,
                                                                                       selected_loop_cycle)
        # Update SHO/Raw plots
        loop_pgf_title = ax_loop_pgf.get_title()
        sho_title = ' - '.join(['Bias Step: {}'.format(selected_step), loop_pgf_title])
        raw_amp_plot.set_xdata(current_w_vec)
        raw_amp_plot.set_ydata(np.abs(current_raw_data))
        guess_amp_plot.set_xdata(current_w_vec)
        guess_amp_plot.set_ydata(np.abs(current_sho_guess))
        fit_amp_plot.set_xdata(current_w_vec)
        fit_amp_plot.set_ydata(np.abs(current_sho_fit))
        ax_sho_amp.set_title(sho_title)
        ax_sho_amp.set_ylim([np.min(np.abs([current_raw_data, current_sho_guess, current_sho_fit])) * 0.95,
                             np.max(np.abs([current_raw_data, current_sho_guess, current_sho_fit])) * 1.05])

        raw_phase_plot.set_xdata(current_w_vec)
        raw_phase_plot.set_ydata(np.angle(current_raw_data))
        guess_phase_plot.set_xdata(current_w_vec)
        guess_phase_plot.set_ydata(np.angle(current_sho_guess))
        fit_phase_plot.set_xdata(current_w_vec)
        fit_phase_plot.set_ydata(np.angle(current_sho_fit))
        ax_sho_phase.set_title(sho_title)

    def pos_picker(event):
        # Allow for picking a new position.
        if not ax_loop_map.axes.in_axes(event):
            return