gmode_utils.py 31.3 KB
Newer Older
Somnath, Suhas's avatar
Somnath, Suhas committed
1
2
3
4
5
6
7
# -*- coding: utf-8 -*-
"""
Created on Thu May 05 13:29:12 2016

@author: Suhas Somnath
"""

8
from __future__ import division, print_function, absolute_import
9
import sys
Somnath, Suhas's avatar
Somnath, Suhas committed
10
import itertools
11
from collections import Iterable
12
13
14
15
from multiprocessing import Pool, cpu_count
from warnings import warn
import matplotlib.pyplot as plt
import numpy as np
16
from time import time
Somnath, Suhas's avatar
Somnath, Suhas committed
17
from .fft import getNoiseFloor, noiseBandFilter, makeLPF, harmonicsPassFilter
Somnath, Suhas's avatar
Somnath, Suhas committed
18
from ..io.io_hdf5 import ioHDF5
19
from ..io.hdf_utils import getH5DsetRefs, linkRefs, getAuxData, link_as_main, copyAttributes, copy_main_attributes
20
from ..io.io_utils import getTimeStamp
Somnath, Suhas's avatar
Somnath, Suhas committed
21
from ..io.microdata import MicroDataGroup, MicroDataset
22
from ..viz.plot_utils import rainbow_plot
23
from ..io.translators.utils import build_ind_val_dsets
Somnath, Suhas's avatar
Somnath, Suhas committed
24

Unknown's avatar
Unknown committed
25

26
# TODO: Use filter_parms as a kwargs instead of a required input
27
# TODO: Phase rotation not implemented correctly. Find and use excitation frequency
Somnath, Suhas's avatar
Somnath, Suhas committed
28
29

###############################################################################
30
31


32
def test_filter(resp_wfm, filter_parms, samp_rate, show_plots=True, use_rainbow_plots=True,
33
                excit_wfm=None, central_resp_size=None, verbose=False):
Somnath, Suhas's avatar
Somnath, Suhas committed
34
35
36
37
38
    """
    Filters the provided response with the provided filters. Use this only to test filters.
    This function does not care about the file structure etc.
    
    Parameters
39
    ----------
Somnath, Suhas's avatar
Somnath, Suhas committed
40
41
42
    resp_wfm : 1D numpy float array
        Raw response waveform in the time domain
    filter_parms : dictionary
Chris Smith's avatar
Chris Smith committed
43
        Dictionary that contains all the filtering parameters, see Notes for details.
Somnath, Suhas's avatar
Somnath, Suhas committed
44
45
46
47
    samp_rate : unsigned int 
        Sampling rate in Hertz
    show_plots : (Optional) Boolean
        Whether or not to plot FFTs before and after filtering
48
    use_rainbow_plots : (Optional) Boolean
Somnath, Suhas's avatar
Somnath, Suhas committed
49
        Whether or not to plot loops whose color varied as a function of time
Somnath, Suhas's avatar
Somnath, Suhas committed
50
51
    excit_wfm : (Optional) 1D numpy float array
        Excitation waveform in the time domain. This waveform is necessary for plotting loops. 
Somnath, Suhas's avatar
Somnath, Suhas committed
52
    central_resp_size : (Optional) unsigned int
53
54
55
        Number of responce sample points from the center of the waveform to show in plots. Useful for SPORC
    verbose : (Optional) string
        Whether or not to print statements
Somnath, Suhas's avatar
Somnath, Suhas committed
56
    
Chris Smith's avatar
Chris Smith committed
57
58
59
60
61
62
63
64
65
66
    Returns
    -------
    filt_data : 1D numpy float array
        Filtered signal in the time domain

    Notes
    -----
    *Filter Parameters*

    noise_threshold : float
Somnath, Suhas's avatar
Somnath, Suhas committed
67
        0<1 eg 1E-4
Chris Smith's avatar
Chris Smith committed
68
    comb_[Hz] : Retain harmonics of frequency
Somnath, Suhas's avatar
Somnath, Suhas committed
69
70
        [first frequency, band width, number of harmonics]
    LPF_cutOff_[Hz] : float
Chris Smith's avatar
Chris Smith committed
71
72
        low pass frequency cut off frequency
    band_filt_[Hz] : 2D list
Somnath, Suhas's avatar
Somnath, Suhas committed
73
74
        [0] = center frequency, [1] = band widths
    phase_[rad] : float
Chris Smith's avatar
Chris Smith committed
75
        Compensation for instrumentation induced phase offset in radians
Somnath, Suhas's avatar
Somnath, Suhas committed
76
77
78
79
    samp_rate_[Hz] : unsigned int
        Sampling rate in Hz
    num_pix : unsigned int
        Number of pixels to filter simultaneously
Chris Smith's avatar
Chris Smith committed
80

Somnath, Suhas's avatar
Somnath, Suhas committed
81
82
    """
    num_pts = len(resp_wfm)
Unknown's avatar
Unknown committed
83

Unknown's avatar
Unknown committed
84
    show_loops = excit_wfm is not None and show_plots
Unknown's avatar
Unknown committed
85

86
87
88
    '''
    Get parameters from the dictionary.
    '''
89
90
    noise_band_filter = filter_parms.get('band_filt_[Hz]', 1)
    if isinstance(noise_band_filter, Iterable):
91
92
        noise_band_filter = noiseBandFilter(num_pts, samp_rate, noise_band_filter[0],
                                            noise_band_filter[1])
93
94
        if verbose and isinstance(noise_band_filter, Iterable):
            print('Calculated valid noise_band_filter')
95

96
    low_pass_filter = filter_parms.get('LPF_cutOff_[Hz]', -1)
Chris Smith's avatar
Chris Smith committed
97
    if low_pass_filter > 0:
98
        low_pass_filter = makeLPF(num_pts, samp_rate, low_pass_filter)
99
100
101
102
103
        if verbose and isinstance(low_pass_filter, Iterable):
            print('Calculated valid low pass filter')
    else:
        low_pass_filter = 1

Chris Smith's avatar
Chris Smith committed
104
    harmonic_filter = filter_parms.get('comb_[Hz]', 1)
105
    if isinstance(harmonic_filter, Iterable):
106
107
        harmonic_filter = harmonicsPassFilter(num_pts, samp_rate, harmonic_filter[0],
                                              harmonic_filter[1], harmonic_filter[2])
108
109
        if verbose and isinstance(harmonic_filter, Iterable):
            print('Calculated valid harmonic filter')
110

Somnath, Suhas's avatar
Somnath, Suhas committed
111
    composite_filter = noise_band_filter * low_pass_filter * harmonic_filter
112

Chris Smith's avatar
Chris Smith committed
113
    noise_floor = filter_parms.get('noise_threshold', None)
114
    fft_pix_data = np.fft.fftshift(np.fft.fft(resp_wfm))
115
116
117
    if 0 < noise_floor < 1:
        noise_floor = getNoiseFloor(fft_pix_data, noise_floor)[0]

Unknown's avatar
Unknown committed
118
119
    if show_plots:
        l_ind = int(0.5 * num_pts)
120
121
122
123
        if type(composite_filter) == np.ndarray:
            r_ind = np.max(np.where(composite_filter > 0)[0])
        else:
            r_ind = num_pts
Unknown's avatar
Unknown committed
124
        w_vec = np.linspace(-0.5 * samp_rate, 0.5 * samp_rate, num_pts) * 1E-3
Somnath, Suhas's avatar
Somnath, Suhas committed
125
        if central_resp_size:
Unknown's avatar
Unknown committed
126
127
128
            sz = int(0.5 * central_resp_size)
            l_resp_ind = -sz + l_ind
            r_resp_ind = l_ind + sz
Somnath, Suhas's avatar
Somnath, Suhas committed
129
130
131
        else:
            l_resp_ind = l_ind
            r_resp_ind = num_pts
Unknown's avatar
Unknown committed
132

Somnath, Suhas's avatar
Somnath, Suhas committed
133
134
        fig = plt.figure(figsize=(12, 8))
        lhs_colspan = 2
135
        if show_loops is False:
Somnath, Suhas's avatar
Somnath, Suhas committed
136
137
138
139
140
141
            lhs_colspan = 4
        else:
            ax_loops = plt.subplot2grid((2, 4), (0, 2), colspan=2, rowspan=2)
        ax_raw = plt.subplot2grid((2, 4), (0, 0), colspan=lhs_colspan)
        ax_filt = plt.subplot2grid((2, 4), (1, 0), colspan=lhs_colspan)
        axes = [ax_raw, ax_filt]
142
143
144
    else:
        fig = None
        axes = None
145

Somnath, Suhas's avatar
Somnath, Suhas committed
146
    if show_plots:
147
        amp = np.abs(fft_pix_data)
148
        ax_raw.semilogy(w_vec[l_ind:r_ind], amp[l_ind:r_ind])
Unknown's avatar
Unknown committed
149
        ax_raw.semilogy(w_vec[l_ind:r_ind], (composite_filter[l_ind:r_ind] + np.min(amp)) * (np.max(amp) - np.min(amp)))
150
        if noise_floor is not None:
Unknown's avatar
Unknown committed
151
            ax_raw.semilogy(w_vec[l_ind:r_ind], np.ones(r_ind - l_ind) * noise_floor)
Somnath, Suhas's avatar
Somnath, Suhas committed
152
        ax_raw.set_title('Raw Signal')
153
154
155
156
157
158

    fft_pix_data *= composite_filter

    if noise_floor is not None:
        fft_pix_data[np.abs(fft_pix_data) < noise_floor] = 1E-16  # DON'T use 0 here. ipython kernel dies

Somnath, Suhas's avatar
Somnath, Suhas committed
159
    if show_plots:
160
        ax_filt.semilogy(w_vec[l_ind:r_ind], np.abs(fft_pix_data[l_ind:r_ind]))
Somnath, Suhas's avatar
Somnath, Suhas committed
161
162
        ax_filt.set_title('Filtered Signal')
        ax_filt.set_xlabel('Frequency(kHz)')
163
        if noise_floor is not None:
164
            ax_filt.set_ylim(bottom=noise_floor)  # prevents the noise threshold from messing up plots
165
166
167

    filt_data = np.real(np.fft.ifft(np.fft.ifftshift(fft_pix_data)))

Somnath, Suhas's avatar
Somnath, Suhas committed
168
    if show_loops:
169
        if use_rainbow_plots:
170
            rainbow_plot(ax_loops, excit_wfm[l_resp_ind:r_resp_ind], filt_data[l_resp_ind:r_resp_ind] * 1E+3)
Somnath, Suhas's avatar
Somnath, Suhas committed
171
        else:
Unknown's avatar
Unknown committed
172
173
            ax_loops.plot(excit_wfm[l_resp_ind:r_resp_ind], filt_data[l_resp_ind:r_resp_ind] * 1E+3)
        ax_loops.set_title('AI vs AO')
Somnath, Suhas's avatar
Somnath, Suhas committed
174
175
176
        ax_loops.set_xlabel('Input Bias (V)')
        ax_loops.set_ylabel('Deflection (mV)')
        axes.append(ax_loops)
177
        fig.tight_layout()
Somnath, Suhas's avatar
Somnath, Suhas committed
178
179
    return filt_data, fig, axes

Unknown's avatar
Unknown committed
180

Somnath, Suhas's avatar
Somnath, Suhas committed
181
182
183
# ##############################################################################


184
def fft_filter_dataset(h5_main, filter_parms, write_filtered=True, write_condensed=False, num_cores=None):
185
    # TODO: Can simplify this function substantially. Collapse / absorb the serial and parallel functions...
Somnath, Suhas's avatar
Somnath, Suhas committed
186
187
188
189
    """
    Filters G-mode data using specified filter parameters and writes results to file.
        
    Parameters
190
    ----------
Somnath, Suhas's avatar
Somnath, Suhas committed
191
192
193
194
    h5_main : HDF5 dataset object
        Dataset containing the raw data
    filter_parms : dictionary
        Dictionary that contains all the filtering parameters
Somnath, Suhas's avatar
Somnath, Suhas committed
195
    write_filtered : (optional) Boolean - default True
Somnath, Suhas's avatar
Somnath, Suhas committed
196
        Whether or not to write filtered data to file
Somnath, Suhas's avatar
Somnath, Suhas committed
197
    write_condensed : (optional) Boolean - default False
Somnath, Suhas's avatar
Somnath, Suhas committed
198
199
200
201
202
        Whether or not to write condensed filtered data to file
    num_cores : unsigned int
        Number of cores to use for processing data in parallel
        
    Returns
203
    -------
Somnath, Suhas's avatar
Somnath, Suhas committed
204
    HDF5 group reference containing filtered dataset
Unknown's avatar
Unknown committed
205
206
    """

207
    def __max_pixel_read(h5_raw, num_cores, store_filt=True, hot_bins=None, bytes_per_bin=2, max_RAM_gb=16):
Somnath, Suhas's avatar
Somnath, Suhas committed
208
        """
209
210
        Returns the maximum number of pixels that can be stored in memory considering the output data and the
        number of cores
Somnath, Suhas's avatar
Somnath, Suhas committed
211
        
212
        Parameters
Somnath, Suhas's avatar
Somnath, Suhas committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
        ----------
        h5_raw : HDF5 Dataset reference
            Dataset containing the raw data
        num_cores : unsigned int
            Number of cores to use for processing data in parallel
        store_filt : boolean (optional)
            Write filtered data back to h5
        hot_bins : 1D numpy array (optional)
            Bins in the frequency domain to be saved to h5
        bytes_per_bin : unsigned int (optional)
            bytes per unit in the raw data - typically 2 bytes
        max_RAM_gb : unsigned int (optional)
            Maximum system memory that can be used. This is NOT respected well. Please fix this.            
        """
Unknown's avatar
Unknown committed
227

Somnath, Suhas's avatar
Somnath, Suhas committed
228
229
        # double the memory requirement if storing filtered data
        if store_filt:
Unknown's avatar
Unknown committed
230
            bytes_per_bin *= 2
Somnath, Suhas's avatar
Somnath, Suhas committed
231
232
        bytes_per_pix = h5_raw.shape[1] * bytes_per_bin
        # account for the hot bins separately
233
234
        if hot_bins is not None:
            bytes_per_pix += len(hot_bins) * 8  # complex64
Unknown's avatar
Unknown committed
235

Somnath, Suhas's avatar
Somnath, Suhas committed
236
237
        # Multiply this memory requirement per pixel by the number of cores
        bytes_per_pix *= num_cores
Unknown's avatar
Unknown committed
238
239

        max_pix = int(np.rint(max_RAM_gb * 1024 ** 3 / bytes_per_pix))
240
241
        max_pix = max(1, min(h5_raw.shape[0], max_pix))
        print('Allowed to read', max_pix, 'of', h5_raw.shape[0], 'pixels')
Unknown's avatar
Unknown committed
242

Somnath, Suhas's avatar
Somnath, Suhas committed
243
        return np.uint(max_pix)
Unknown's avatar
Unknown committed
244

245
    def __filter_chunk(raw_mat, parm_dict, recom_cores):
Somnath, Suhas's avatar
Somnath, Suhas committed
246
247
248
249
250
        """
        This function delegates the actual fitting responsibility to the
        appropriate function. Decides whether or not serial / parallel processing
        is appropriate, number of cores, number of chunks, etc.
        
251
252
        Parameters
        ----------
Somnath, Suhas's avatar
Somnath, Suhas committed
253
254
255
256
257
258
259
        raw_mat : 2D numpy array
            Raw data arranged as [repetition, points per measurement]
        parm_dict : Dictionary
            Parameters necessary for filtering
        recom_cores : unsigned int
            Number of cores to use for processing data in parallel
        
260
261
        Returns
        -------
Somnath, Suhas's avatar
Somnath, Suhas committed
262
263
264
265
266
267
268
269
        (noise_floors, filt_data, cond_data)
        
        noise_floors : 1D numpy array
            Contains the noise floors per set of measurements
        filt_data : 2D numpy array or None
            filtered data arranged as [repetition, points per measurement]
        cond_data : 2D complex numpy array
            [set of measurements, frequency bins containing data]
Unknown's avatar
Unknown committed
270

Somnath, Suhas's avatar
Somnath, Suhas committed
271
        """
Unknown's avatar
Unknown committed
272
273
274
        max_cores = max(1, cpu_count() - 2)

        recom_chunks = int(raw_mat.shape[0] / recom_cores)
275
276
277

        print('recom cores:', recom_cores, 'Total pixels:', raw_mat.shape[0], ', Recom chunks:', recom_chunks)

Somnath, Suhas's avatar
Somnath, Suhas committed
278
279
        if recom_cores > 1 and recom_chunks < 10:
            min_jobs = 20
Unknown's avatar
Unknown committed
280
            reduced_cores = int(raw_mat.shape[0] / min_jobs)
Somnath, Suhas's avatar
Somnath, Suhas committed
281
            # intelligently set the cores now. 
282
283
284
            recom_cores = min(max_cores, reduced_cores)
            print('Not enough jobs per core. Reducing cores to', recom_cores)

Somnath, Suhas's avatar
Somnath, Suhas committed
285
        if recom_cores > 1:
Somnath, Suhas's avatar
Somnath, Suhas committed
286
            return filter_chunk_parallel(raw_mat, parm_dict, recom_cores)
Somnath, Suhas's avatar
Somnath, Suhas committed
287
        else:
Somnath, Suhas's avatar
Somnath, Suhas committed
288
            return filter_chunk_serial(raw_mat, parm_dict)
289
290

    max_cores = max(1, cpu_count() - 2)
Somnath, Suhas's avatar
Somnath, Suhas committed
291
292
    if not num_cores:
        num_cores = max_cores
Unknown's avatar
Unknown committed
293

294
295
296
    if write_filtered is False and write_condensed is False:
        warn('You need to write the filtered and/or the condensed dataset to the file')
        return
297
298
299
300

    if 'num_pix' not in filter_parms:
        filter_parms['num_pix'] = 1

Unknown's avatar
Unknown committed
301
    num_effective_pix = h5_main.shape[0] * 1.0 / filter_parms['num_pix']
Somnath, Suhas's avatar
Somnath, Suhas committed
302
303
304
305
    if num_effective_pix % 1 > 0:
        warn('Number of pixels not divisible by the number of pixels to use for FFT filter')
        return
    num_effective_pix = int(num_effective_pix)
Unknown's avatar
Unknown committed
306
307

    num_pts = h5_main.shape[1] * filter_parms['num_pix']
308

309
310
    noise_band_filter = 1
    low_pass_filter = 1
311
312
313
    harmonic_filter = 1

    if 'band_filt_[Hz]' in filter_parms:
314
        if isinstance(filter_parms['band_filt_[Hz]'], Iterable):
315
316
317
318
319
320
321
            band_filt = filter_parms['band_filt_[Hz]']
            noise_band_filter = noiseBandFilter(num_pts, filter_parms['samp_rate_[Hz]'], band_filt[0], band_filt[1])
    if 'LPF_cutOff_[Hz]' in filter_parms:
        if filter_parms['LPF_cutOff_[Hz]'] > 0:
            low_pass_filter = makeLPF(num_pts, filter_parms['samp_rate_[Hz]'], filter_parms['LPF_cutOff_[Hz]'])

    if 'comb_[Hz]' in filter_parms:
322
        if isinstance(filter_parms['comb_[Hz]'], Iterable):
323
324
325
326
            harmonic_filter = harmonicsPassFilter(num_pts, filter_parms['samp_rate_[Hz]'], filter_parms['comb_[Hz]'][0],
                                                  filter_parms['comb_[Hz]'][1], filter_parms['comb_[Hz]'][2])

    composite_filter = noise_band_filter * low_pass_filter * harmonic_filter
Unknown's avatar
Unknown committed
327

Somnath, Suhas's avatar
Somnath, Suhas committed
328
    # ioHDF now handles automatic indexing
329
330
331
332
333
    grp_name = h5_main.name.split('/')[-1] + '-FFT_Filtering_'

    doing_noise_floor_filter = False

    if 'noise_threshold' in filter_parms:
Unknown's avatar
Unknown committed
334
        if 0 < filter_parms['noise_threshold'] < 1:
335
336
337
338
339
340
341
342
343
344
345
346
347
            ds_noise_floors = MicroDataset('Noise_Floors',
                                           data=np.zeros(shape=num_effective_pix, dtype=np.float32))
            doing_noise_floor_filter = True
        else:
            # Illegal inputs will be deleted from dictionary
            warn('Provided noise floor threshold: {} not within (0,1)'.format(filter_parms['noise_threshold']))
            del filter_parms['noise_threshold']

    if not doing_noise_floor_filter and not isinstance(composite_filter, np.ndarray):
        warn("No filtering being performed on this dataset. Exiting!")

    if isinstance(composite_filter, np.ndarray):
        ds_comp_filt = MicroDataset('Composite_Filter', np.float32(composite_filter))
Unknown's avatar
Unknown committed
348

349
    grp_filt = MicroDataGroup(grp_name, h5_main.parent.name)
Somnath, Suhas's avatar
Somnath, Suhas committed
350
351
352
    filter_parms['timestamp'] = getTimeStamp()
    filter_parms['algorithm'] = 'GmodeUtils-Parallel'
    grp_filt.attrs = filter_parms
353
354
355
356
357
    if isinstance(composite_filter, np.ndarray):
        grp_filt.addChildren([ds_comp_filt])
    if doing_noise_floor_filter:
        grp_filt.addChildren([ds_noise_floors])

Somnath, Suhas's avatar
Somnath, Suhas committed
358
    if write_filtered:
359
        ds_filt_data = MicroDataset('Filtered_Data', data=[], maxshape=h5_main.maxshape,
Somnath, Suhas's avatar
Somnath, Suhas committed
360
361
                                    dtype=np.float32, chunking=h5_main.chunks, compression='gzip')
        grp_filt.addChildren([ds_filt_data])
Unknown's avatar
Unknown committed
362

363
364
365
366
367
    hot_inds = None

    h5_pos_inds = getAuxData(h5_main, auxDataName=['Position_Indices'])[0]
    h5_pos_vals = getAuxData(h5_main, auxDataName=['Position_Values'])[0]

Somnath, Suhas's avatar
Somnath, Suhas committed
368
369
    if write_condensed:
        hot_inds = np.where(composite_filter > 0)[0]
Unknown's avatar
Unknown committed
370
371
        hot_inds = np.uint(hot_inds[int(0.5 * len(hot_inds)):])  # only need to keep half the data
        ds_spec_inds, ds_spec_vals = build_ind_val_dsets([int(0.5 * len(hot_inds))], is_spectral=True,
372
373
                                                         labels=['hot_frequencies'], units=[''], verbose=False)
        ds_spec_vals.data = np.atleast_2d(hot_inds)  # The data generated above varies linearly. Override.
374
375
        ds_cond_data = MicroDataset('Condensed_Data', data=[], maxshape=(num_effective_pix, len(hot_inds)),
                                    dtype=np.complex, chunking=(1, len(hot_inds)), compression='gzip')
376
377
378
379
380
381
382
383
384
385
386
387
388
        grp_filt.addChildren([ds_spec_inds, ds_spec_vals, ds_cond_data])
        if filter_parms['num_pix'] > 1:
            # need to make new position datasets by taking every n'th index / value:
            new_pos_vals = np.atleast_2d(h5_pos_vals[slice(0, None, filter_parms['num_pix']), :])
            ds_pos_inds, ds_pos_vals = build_ind_val_dsets([int(np.unique(h5_pos_inds[:, dim_ind]).size /
                                                                filter_parms['num_pix'])
                                                            for dim_ind in range(h5_pos_inds.shape[1])],
                                                           is_spectral=False,
                                                           labels=h5_pos_inds.attrs['labels'],
                                                           units=h5_pos_inds.attrs['units'], verbose=False)
            h5_pos_vals.data = np.atleast_2d(new_pos_vals)  # The data generated above varies linearly. Override.
            grp_filt.addChildren([ds_pos_inds, ds_pos_vals])

Somnath, Suhas's avatar
Somnath, Suhas committed
389
    hdf = ioHDF5(h5_main.file)
Somnath, Suhas's avatar
Somnath, Suhas committed
390
    h5_filt_refs = hdf.writeData(grp_filt)
391
392
393
394
    if isinstance(composite_filter, np.ndarray):
        h5_comp_filt = getH5DsetRefs(['Composite_Filter'], h5_filt_refs)[0]
    if doing_noise_floor_filter:
        h5_noise_floors = getH5DsetRefs(['Noise_Floors'], h5_filt_refs)[0]
Unknown's avatar
Unknown committed
395

Somnath, Suhas's avatar
Somnath, Suhas committed
396
397
398
    # Now need to link appropriately:
    if write_filtered:
        h5_filt_data = getH5DsetRefs(['Filtered_Data'], h5_filt_refs)[0]
399
        copyAttributes(h5_main, h5_filt_data, skip_refs=False)
400
401
402
403
404
        if isinstance(composite_filter, np.ndarray):
            linkRefs(h5_filt_data, [h5_comp_filt])
        if doing_noise_floor_filter:
            linkRefs(h5_filt_data, [h5_noise_floors])

405
        """link_as_main(h5_filt_data, h5_pos_inds, h5_pos_vals,
406
                     getAuxData(h5_main, auxDataName=['Spectroscopic_Indices'])[0],
407
                     getAuxData(h5_main, auxDataName=['Spectroscopic_Values'])[0])"""
Unknown's avatar
Unknown committed
408

Somnath, Suhas's avatar
Somnath, Suhas committed
409
410
    if write_condensed:
        h5_cond_data = getH5DsetRefs(['Condensed_Data'], h5_filt_refs)[0]
411
412
413
414
415
        if isinstance(composite_filter, np.ndarray):
            linkRefs(h5_cond_data, [h5_comp_filt])
        if doing_noise_floor_filter:
            linkRefs(h5_cond_data, [h5_noise_floors])

416
417
418
        if filter_parms['num_pix'] > 1:
            h5_pos_inds = getH5DsetRefs(['Position_Indices'], h5_filt_refs)[0]
            h5_pos_vals = getH5DsetRefs(['Position_Values'], h5_filt_refs)[0]
419

420
421
422
        link_as_main(h5_cond_data, h5_pos_inds, h5_pos_vals,
                     getH5DsetRefs(['Spectroscopic_Indices'], h5_filt_refs)[0],
                     getH5DsetRefs(['Spectroscopic_Values'], h5_filt_refs)[0])
Unknown's avatar
Unknown committed
423

Somnath, Suhas's avatar
Somnath, Suhas committed
424
425
426
    rot_pts = 0
    if 'phase_rot_[pts]' in filter_parms.keys():
        rot_pts = int(filter_parms['phase_rot_[pts]'])
Unknown's avatar
Unknown committed
427
428

    print('Filtering data now. Be patient, this could take a few minutes')
Somnath, Suhas's avatar
Somnath, Suhas committed
429

430
431
    max_pix = __max_pixel_read(h5_main, max(1, cpu_count() - 2), store_filt=write_filtered, hot_bins=hot_inds,
                               bytes_per_bin=2, max_RAM_gb=16)
Somnath, Suhas's avatar
Somnath, Suhas committed
432
    # Ensure that whole sets of pixels can be read.
Unknown's avatar
Unknown committed
433
434
    max_pix = np.uint(filter_parms['num_pix'] * np.floor(max_pix / filter_parms['num_pix']))

435
    parm_dict = {'filter_parms': filter_parms, 'composite_filter': composite_filter,
Somnath, Suhas's avatar
Somnath, Suhas committed
436
                 'rot_pts': rot_pts, 'hot_inds': hot_inds}
437
438

    t_start = time()
Somnath, Suhas's avatar
Somnath, Suhas committed
439
440
441
    st_pix = 0
    line_count = 0
    while st_pix < h5_main.shape[0]:
442
443
444
        en_pix = int(min(h5_main.shape[0], st_pix + max_pix))
        print('Reading pixels:', st_pix, 'to', en_pix, 'of', h5_main.shape[0])
        raw_mat = h5_main[st_pix:en_pix, :]
Somnath, Suhas's avatar
Somnath, Suhas committed
445
446
        # reshape to (set of pix, data in each set of pix)
        # print 'raw mat originally of shape:', raw_mat.shape
447
        raw_mat = raw_mat.reshape(-1, filter_parms['num_pix'] * raw_mat.shape[1])
Somnath, Suhas's avatar
Somnath, Suhas committed
448
449
        num_lines = raw_mat.shape[0]
        # print 'After collapsing pixels, raw mat now of shape:', raw_mat.shape
450
        (nse_flrs, filt_data, cond_data) = __filter_chunk(raw_mat, parm_dict, num_cores)
Somnath, Suhas's avatar
Somnath, Suhas committed
451
        # Insert things into appropriate HDF datasets
452
        print('Writing filtered data to h5')
Somnath, Suhas's avatar
Somnath, Suhas committed
453
        # print 'Noise floors of shape:', nse_flrs.shape
454
455
        if doing_noise_floor_filter:
            h5_noise_floors[line_count: line_count + num_lines] = nse_flrs
Somnath, Suhas's avatar
Somnath, Suhas committed
456
        if write_condensed:
Somnath, Suhas's avatar
Somnath, Suhas committed
457
            # print('Condensed data of shape:', cond_data.shape)
458
            h5_cond_data[line_count: line_count + num_lines, :] = cond_data
Somnath, Suhas's avatar
Somnath, Suhas committed
459
        if write_filtered:
Somnath, Suhas's avatar
Somnath, Suhas committed
460
            # print('Filtered data of shape:', filt_data.shape)
461
            h5_filt_data[st_pix:en_pix, :] = filt_data
Somnath, Suhas's avatar
Somnath, Suhas committed
462
463
        hdf.flush()
        st_pix = en_pix
464
465
466
467
468
469
470

    print('FFT filtering took {} seconds'.format(time() - t_start))

    if isinstance(composite_filter, np.ndarray):
        return h5_comp_filt.parent
    if doing_noise_floor_filter:
        return h5_noise_floors.parent
Unknown's avatar
Unknown committed
471
472


Somnath, Suhas's avatar
Somnath, Suhas committed
473
# #############################################################################
474
475


Somnath, Suhas's avatar
Somnath, Suhas committed
476
def filter_chunk_parallel(raw_data, parm_dict, num_cores):
477
    # TODO: Need to check to ensure that all cores are indeed being utilized
Somnath, Suhas's avatar
Somnath, Suhas committed
478
479
480
    """
    Filters the provided dataset in parallel
    
481
    Parameters
Somnath, Suhas's avatar
Somnath, Suhas committed
482
483
484
485
486
487
488
489
    ----------
    raw_data : 2D numpy array
        Raw data arranged as [repetition, points per measurement]
    parm_dict : Dictionary
        Parameters necessary for filtering
    num_cores : unsigned int
        Number of cores to use for processing data in parallel
    
490
491
    Returns
    -------
Somnath, Suhas's avatar
Somnath, Suhas committed
492
493
494
495
496
497
498
    (noise_floors, filt_data, cond_data)
    
    noise_floors : 1D numpy array
        Contains the noise floors per set of measurements
    filt_data : 2D numpy array or None
        filtered data arranged as [repetition, points per measurement]
    cond_data : 2D complex numpy array or None
Unknown's avatar
Unknown committed
499
500
        [set of measurements, frequency bins containing data]

Somnath, Suhas's avatar
Somnath, Suhas committed
501
502
503
504
505
    """
    # Make place-holders to hold the data:
    pix_per_set = parm_dict['filter_parms']['num_pix']
    num_sets = raw_data.shape[0]
    pts_per_set = raw_data.shape[1]
506
    # print('sending unit filter data of size', pts_per_set)
507
508
509
510
511

    noise_floors = None
    noise_thresh = None
    if 'noise_threshold' in parm_dict['filter_parms']:
        noise_thresh = parm_dict['filter_parms']['noise_threshold']
Unknown's avatar
Unknown committed
512
        if 0 < noise_thresh < 1:
513
514
515
516
            noise_floors = np.zeros(shape=num_sets, dtype=np.float32)
        else:
            noise_thresh = None

Somnath, Suhas's avatar
Somnath, Suhas committed
517
518
    filt_data = None
    if not parm_dict['rot_pts']:
Unknown's avatar
Unknown committed
519
        filt_data = np.zeros(shape=(num_sets * pix_per_set, int(pts_per_set / pix_per_set)), dtype=raw_data.dtype)
520

Somnath, Suhas's avatar
Somnath, Suhas committed
521
    cond_data = None
522
523
    if parm_dict['hot_inds'] is not None:
        cond_data = np.zeros(shape=(num_sets, parm_dict['hot_inds'].size), dtype=np.complex64)
Unknown's avatar
Unknown committed
524

Somnath, Suhas's avatar
Somnath, Suhas committed
525
    # Set up single parameter:
526
527
528
529
530
    if sys.version_info.major == 3:
        zip_fun = zip
    else:
        zip_fun = itertools.izip
    sing_parm = zip_fun(raw_data, itertools.repeat(parm_dict))
Unknown's avatar
Unknown committed
531

Somnath, Suhas's avatar
Somnath, Suhas committed
532
    # Setup parallel processing:
533
534
    # num_cores = 10
    pool = Pool(processes=num_cores, maxtasksperchild=None)
Unknown's avatar
Unknown committed
535

Somnath, Suhas's avatar
Somnath, Suhas committed
536
    # Start parallel processing:
Unknown's avatar
Unknown committed
537
    num_chunks = int(np.ceil(raw_data.shape[0] / num_cores))
Somnath, Suhas's avatar
Somnath, Suhas committed
538
    parallel_results = pool.imap(unit_filter, sing_parm, chunksize=num_chunks)
Somnath, Suhas's avatar
Somnath, Suhas committed
539
540
    pool.close()
    pool.join()
Unknown's avatar
Unknown committed
541

Somnath, Suhas's avatar
Somnath, Suhas committed
542
    print('Done parallel computing. Now extracting data and populating matrices')
Unknown's avatar
Unknown committed
543

Somnath, Suhas's avatar
Somnath, Suhas committed
544
    # Extract data for each line...
Unknown's avatar
Unknown committed
545
    print_set = np.linspace(0, num_sets - 1, 10, dtype=int)
546
    for set_ind, current_results in enumerate(parallel_results):
Somnath, Suhas's avatar
Somnath, Suhas committed
547
        if set_ind in print_set:
548
            print('Reading...', np.rint(100 * set_ind / num_sets), '% complete')
549
550
551

        temp_noise, filt_data_set, cond_data_set = current_results

552
553
        if noise_thresh is not None:
            noise_floors[set_ind] = temp_noise
554
555
556
        if parm_dict['hot_inds'] is not None:
            cond_data[set_ind, :] = cond_data_set
        if parm_dict['rot_pts'] is not None:
Unknown's avatar
Unknown committed
557
558
            filt_data[set_ind * pix_per_set:(set_ind + 1) * pix_per_set, :] = filt_data_set

559
    return noise_floors, filt_data, cond_data
Unknown's avatar
Unknown committed
560

561

Somnath, Suhas's avatar
Somnath, Suhas committed
562
def filter_chunk_serial(raw_data, parm_dict):
Somnath, Suhas's avatar
Somnath, Suhas committed
563
564
565
    """
    Filters the provided dataset serially
    
566
    Parameters
Somnath, Suhas's avatar
Somnath, Suhas committed
567
568
569
570
571
572
    ----------
    raw_data : 2D numpy array
        Raw data arranged as [repetition, points per measurement]
    parm_dict : Dictionary
        Parameters necessary for filtering
    
573
574
    Returns
    -------
Somnath, Suhas's avatar
Somnath, Suhas committed
575
576
577
578
579
580
581
582
    (noise_floors, filt_data, cond_data)
    
    noise_floors : 1D numpy array
        Contains the noise floors per set of measurements
    filt_data : 2D numpy array or None
        filtered data arranged as [repetition, points per measurement]
    cond_data : 2D complex numpy array or None
        [set of measurements, frequency bins containing data]
Unknown's avatar
Unknown committed
583

Somnath, Suhas's avatar
Somnath, Suhas committed
584
585
586
587
588
    """
    # Make place-holders to hold the data:
    pix_per_set = parm_dict['filter_parms']['num_pix']
    num_sets = raw_data.shape[0]
    pts_per_set = raw_data.shape[1]
589
    # print ('sending unit filter data of size', pts_per_set)
590
591
592
593
594

    noise_floors = None
    noise_thresh = None
    if 'noise_threshold' in parm_dict['filter_parms']:
        noise_thresh = parm_dict['filter_parms']['noise_threshold']
Unknown's avatar
Unknown committed
595
        if 0 < noise_thresh < 1:
596
597
598
599
            noise_floors = np.zeros(shape=num_sets, dtype=np.float32)
        else:
            noise_thresh = None

Somnath, Suhas's avatar
Somnath, Suhas committed
600
    filt_data = None
601
    if parm_dict['rot_pts'] is not None:
Unknown's avatar
Unknown committed
602
        filt_data = np.zeros(shape=(num_sets * pix_per_set, int(pts_per_set / pix_per_set)), dtype=raw_data.dtype)
603

Somnath, Suhas's avatar
Somnath, Suhas committed
604
    cond_data = None
605
    if parm_dict['hot_inds'] is not None:
Somnath, Suhas's avatar
Somnath, Suhas committed
606
        cond_data = np.zeros(shape=(num_sets, parm_dict['hot_inds'].size), dtype=np.complex64)
Unknown's avatar
Unknown committed
607

Somnath, Suhas's avatar
Somnath, Suhas committed
608
    # Filter each line
Unknown's avatar
Unknown committed
609
    print_set = np.linspace(0, num_sets - 1, 10, dtype=int)
610
    for set_ind in range(num_sets):
Somnath, Suhas's avatar
Somnath, Suhas committed
611
        if set_ind in print_set:
612
613
614
            print('Reading...', np.rint(100 * set_ind / num_sets), '% complete')

        # parm_dict['t_raw'] = raw_data[set_ind,:]
615
616
617
        (temp_noise, filt_data_set, cond_data_set) = unit_filter((raw_data[set_ind, :], parm_dict))
        if noise_thresh is not None:
            noise_floors[set_ind] = temp_noise
618
619
620
        if parm_dict['hot_inds'] is not None:
            cond_data[set_ind, :] = cond_data_set
        if parm_dict['rot_pts'] is not None:
Unknown's avatar
Unknown committed
621
622
            filt_data[set_ind * pix_per_set:(set_ind + 1) * pix_per_set, :] = filt_data_set

623
    return noise_floors, filt_data, cond_data
Somnath, Suhas's avatar
Somnath, Suhas committed
624

Unknown's avatar
Unknown committed
625

Somnath, Suhas's avatar
Somnath, Suhas committed
626
def unit_filter(single_parm):
Somnath, Suhas's avatar
Somnath, Suhas committed
627
628
629
630
    """
    Filters a single instance of a signal. 
    This is the function that is called in parallel
    
631
632
    Parameters
    ----------
Somnath, Suhas's avatar
Somnath, Suhas committed
633
634
635
636
    single_parm : Tuple
        Parameters and data for filtering a single data instance. 
        Constructed as (raw_data_vec, parm_dict)
    
637
638
    Returns
    -------
Somnath, Suhas's avatar
Somnath, Suhas committed
639
640
641
642
643
644
645
646
647
648
    noise_floors : float
        Noise floor
    filt_data : 1D numpy array or None
        filtered data
    cond_data : 1D complex numpy array or None
        frequency bins containing data
    """
    # unpack all the variables from the sole input
    t_raw = single_parm[0]
    parm_dict = single_parm[1]
649
    # t_raw = parm_dict['t_raw']
Somnath, Suhas's avatar
Somnath, Suhas committed
650
651
652
653
654
    filter_parms = parm_dict['filter_parms']
    composite_filter = parm_dict['composite_filter']
    rot_pts = parm_dict['rot_pts']
    hot_inds = parm_dict['hot_inds']

655
656
    noise_floor = None

Somnath, Suhas's avatar
Somnath, Suhas committed
657
    t_raw = t_raw.reshape(-1)
Somnath, Suhas's avatar
Somnath, Suhas committed
658
    f_data = np.fft.fftshift(np.fft.fft(t_raw))
659
660

    if 'noise_threshold' in filter_parms:
Unknown's avatar
Unknown committed
661
        if 0 < filter_parms['noise_threshold'] < 1:
662
663
664
            noise_floor = getNoiseFloor(f_data, filter_parms['noise_threshold'])[0]
            f_data[np.abs(f_data) < noise_floor] = 1E-16  # DON'T use 0 here. ipython kernel dies

Somnath, Suhas's avatar
Somnath, Suhas committed
665
    f_data = f_data * composite_filter
666

Somnath, Suhas's avatar
Somnath, Suhas committed
667
668
    cond_data = None
    filt_data = None
669
    if hot_inds is not None:
Somnath, Suhas's avatar
Somnath, Suhas committed
670
        cond_data = f_data[hot_inds]
671
    if rot_pts is not None:
Somnath, Suhas's avatar
Somnath, Suhas committed
672
        t_clean = np.real(np.fft.ifft(np.fft.ifftshift(f_data)))
673
        filt_mat = t_clean.reshape(filter_parms['num_pix'], -1)
Somnath, Suhas's avatar
Somnath, Suhas committed
674
675
676
677
        if rot_pts > 0:
            filt_data = np.roll(filt_mat, rot_pts, axis=1)
        else:
            filt_data = filt_mat
Unknown's avatar
Unknown committed
678

679
    return noise_floor, filt_data, cond_data
Somnath, Suhas's avatar
Somnath, Suhas committed
680

Unknown's avatar
Unknown committed
681

Somnath, Suhas's avatar
Somnath, Suhas committed
682
683
###############################################################################

684

Somnath, Suhas's avatar
Somnath, Suhas committed
685
def decompress_response(f_condensed_mat, num_pts, hot_inds):
Somnath, Suhas's avatar
Somnath, Suhas committed
686
687
688
    """
    Returns the time domain representation of waveform(s) that are compressed in the frequency space
    
689
    Parameters
Somnath, Suhas's avatar
Somnath, Suhas committed
690
    ----------
Somnath, Suhas's avatar
Somnath, Suhas committed
691
    f_condensed_mat : 1D or 2D complex numpy arrays
Somnath, Suhas's avatar
Somnath, Suhas committed
692
693
694
695
696
697
698
699
700
701
        Frequency domain signals arranged as [position, frequency]. 
        Only the positive frequncy bins must be in the compressed dataset. 
        The dataset is assumed to have been FFT shifted (such that 0 Hz is at the center).
    num_pts : unsigned int
        Number of points in the time domain signal
    hot_inds : 1D unsigned int numpy array
        Indices of the frequency bins in the compressed data. 
        This index array will be necessary to reverse map the condensed 
        FFT into its original form
        
702
703
    Returns
    -------
Somnath, Suhas's avatar
Somnath, Suhas committed
704
705
706
    time_resp : 2D numpy array
        Time domain response arranged as [position, time]
        
Chris Smith's avatar
Chris Smith committed
707
708
    Notes
    -----
Somnath, Suhas's avatar
Somnath, Suhas committed
709
710
    Memory is given higher priority here, so this function loops over the position
    instead of doing the inverse FFT on the complete data.
Chris Smith's avatar
Chris Smith committed
711

Somnath, Suhas's avatar
Somnath, Suhas committed
712
    """
Somnath, Suhas's avatar
Somnath, Suhas committed
713
    f_condensed_mat = np.atleast_2d(f_condensed_mat)
Somnath, Suhas's avatar
Somnath, Suhas committed
714
    hot_inds_mirror = np.flipud(num_pts - hot_inds)
Somnath, Suhas's avatar
Somnath, Suhas committed
715
716
717
718
    time_resp = np.zeros(shape=(f_condensed_mat.shape[0], num_pts), dtype=np.float32)
    for pos in range(f_condensed_mat.shape[0]):
        f_complete = np.zeros(shape=num_pts, dtype=np.complex)
        f_complete[hot_inds] = f_condensed_mat[pos, :]
Somnath, Suhas's avatar
Somnath, Suhas committed
719
        # Now add the mirror (FFT in negative X axis that was removed)
Somnath, Suhas's avatar
Somnath, Suhas committed
720
721
        f_complete[hot_inds_mirror] = np.flipud(f_condensed_mat[pos, :])
        time_resp[pos, :] = np.real(np.fft.ifft(np.fft.ifftshift(f_complete)))
Unknown's avatar
Unknown committed
722

Somnath, Suhas's avatar
Somnath, Suhas committed
723
    return np.squeeze(time_resp)
724
725


726
def reshape_from_lines_to_pixels(h5_main, pts_per_cycle, scan_step_x_m=1):
727
728
    """
    Breaks up the provided raw G-mode dataset into lines and pixels (from just lines)
729

730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
    Parameters
    ----------
    h5_main : h5py.Dataset object
        Reference to the main dataset that contains the raw data that is only broken up by lines
    pts_per_cycle : unsigned int
        Number of points in a single pixel
    scan_step_x_m : float
        Step in meters for pixels

    Returns
    -------
    h5_resh : h5py.Dataset object
        Reference to the main dataset that contains the reshaped data
    """
    if h5_main.shape[1] % pts_per_cycle != 0:
745
746
        warn('Error in reshaping the provided dataset to pixels. Check points per pixel')
        raise ValueError
Unknown's avatar
Unknown committed
747

748
    num_cols = int(h5_main.shape[1] / pts_per_cycle)
749

750
751
752
    h5_spec_vals = getAuxData(h5_main, auxDataName=['Spectroscopic_Values'])[0]
    h5_pos_vals = getAuxData(h5_main, auxDataName=['Position_Values'])[0]
    single_AO = h5_spec_vals[:, :pts_per_cycle]
753
754
755
756
757
758

    ds_spec_inds, ds_spec_vals = build_ind_val_dsets([single_AO.size], is_spectral=True,
                                                     labels=h5_spec_vals.attrs['labels'],
                                                     units=h5_spec_vals.attrs['units'], verbose=False)
    ds_spec_vals.data = np.atleast_2d(single_AO)  # The data generated above varies linearly. Override.

759
    ds_pos_inds, ds_pos_vals = build_ind_val_dsets([num_cols, h5_main.shape[0]], is_spectral=False,
760
761
762
                                                   steps=[scan_step_x_m, h5_pos_vals[1, 0]],
                                                   labels=['X', 'Y'], units=['m', 'm'], verbose=False)

763
    ds_reshaped_data = MicroDataset('Reshaped_Data', data=np.reshape(h5_main.value, (-1, pts_per_cycle)),
764
765
766
                                    compression='gzip', chunking=(10, pts_per_cycle))

    # write this to H5 as some form of filtered data.
767
    resh_grp = MicroDataGroup(h5_main.name.split('/')[-1] + '-Reshape_', parent=h5_main.parent.name)
768
769
    resh_grp.addChildren([ds_reshaped_data, ds_pos_inds, ds_pos_vals, ds_spec_inds, ds_spec_vals])

770
    hdf = ioHDF5(h5_main.file)
771
    print('Starting to reshape G-mode line data. Please be patient')
772
773
774
775
776
777
778
    h5_refs = hdf.writeData(resh_grp)

    h5_resh = getH5DsetRefs(['Reshaped_Data'], h5_refs)[0]
    # Link everything:
    linkRefs(h5_resh,
             getH5DsetRefs(['Position_Indices', 'Position_Values', 'Spectroscopic_Indices', 'Spectroscopic_Values'],
                           h5_refs))
779
780

    # Copy the two attributes that are really important but ignored:
781
    copy_main_attributes(h5_main, h5_resh)
782

783
    print('Finished reshaping G-mode line data to rows and columns')
784

785
    return h5_resh