modifiedPSD.py 5.02 KB
Newer Older
Anchal's avatar
Anchal committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
'''
Authors: Anchal Gupta - anchal@caltech.edu
         Scott Aronson - scott.aronson@ligo.org
Created:  Aug 23, 2019

This script has two modified functions.
First is a modified welch function which keeps all functionality of welch
but also provides a measure of uncertainty in the returned PSD.
The second is a modified PSD calculating script, which ensures same number
of points in each decade by running modifiedWelch with different nperseg and
imporving the amount of averaging done for higher frequencies.
'''
import numpy as np
from scipy import signal

#Welch Function wrapped to also output standard deviation of data
def modwelch(data, fs=1.0, window='hann', nperseg=None, noverlap=None,
             nfft=None, detrend='constant', return_onesided=True,
             scaling='density', axis=-1, average='mean'):
    if noverlap is None:
        noverlap = nperseg//2

    # Detend the data
    if detrend == 'linear' or detrend == 'constant':
        data = signal.detrend(data, type = detrend)

    #Sort data into rows with number in each row = to nperseg
    # length of columns
    col_len=0
    end_edge=nperseg
    while(end_edge<len(data)):
        end_edge=end_edge+nperseg-noverlap
        col_len=col_len+1

    # make array of row length nperseg and column length however much data
    # will fit fully in row
    sorted_data = np.zeros((col_len,nperseg))

    # make array for PSD will have same col len but
    # row len = max freq/freq_spacing + 1
    # or (samplerate/2) /(samplerate/nperseg) + 1 which = nperseg/2 +1
    welch_array = np.zeros((col_len, nperseg//2 + 1))

    #Move through signal data array populating array
    next_start=0
    for ii in range(col_len):
        sorted_data[ii, :] = data[next_start: next_start+nperseg]
        next_start = next_start+nperseg-noverlap

    # Do row-wise welch
    for k in range(col_len):
        ff, welch_array[k, :] = signal.welch(sorted_data[k, :], fs,
                                             window = window,
                                             nperseg = nperseg,
                                             nfft = nfft, detrend = False,
                                             return_onesided = return_onesided,
                                             scaling = scaling, axis = axis,
                                             average = average)

    mean_array = np.mean(welch_array,axis=0)
    std_array = np.std(welch_array,axis=0)
    median_array = np.median(welch_array,axis=0)
    # Lower percentile equivalent to mean - 1 sigma deviation
    # (1-0.6827)/2
    lowBound_array = np.percentile(welch_array, 15.865, axis=0,)
    # Upper percentile equivalent to mean + 1 sigma deviation
    # (1-0.6827)/2 + 0.6827
    uppBound_array = np.percentile(welch_array, 84.135, axis=0,)

    if average is 'mean':
        return ff, mean_array, std_array, std_array
    elif average is 'median':
        return ff, median_array, lowBound_array, uppBound_array
    else:
        return print('Choose mean or median')

#Modified Power Spectral Density Function
def modPSD(time, data, detrend = 'linear', average = 'mean'):
    # Extract timeseries and calculated Sampling Rate
    timeSeries = time
    SampleRate = 1/(timeSeries[1] - timeSeries[0])
    # Detend the data
    if detrend == 'linear' or detrend == 'constant':
Anchal's avatar
Anchal committed
84
        NoiseSig = signal.detrend(data, type = detrend)
Anchal's avatar
Anchal committed
85
    else:
86
87
        raise RuntimeWarning('detrend can be set to linear or constant only. '
                             'Using raw data without detrending.')
Anchal's avatar
Anchal committed
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
        NoiseSig = data
    maxLWpow = int(np.floor(np.log10(SampleRate/2)))-2
    minLWpow = int(np.ceil(np.log10(SampleRate/len(NoiseSig))))
    lineWidthArray = np.logspace(minLWpow, maxLWpow, maxLWpow-minLWpow+1,
                                 base=10)
    Freq = np.zeros(0)
    PSD = np.zeros(0)
    stdPSD1 = np.zeros(0)
    stdPSD2 = np.zeros(0)
    for lw in lineWidthArray:
        nperseg = int(np.floor(SampleRate/lw))
        subFreq, subPSD, substdPSD1, substdPSD2 = modwelch(NoiseSig,
                                                           window='hann',
                                                           fs=SampleRate,
                                                           nperseg=nperseg,
                                                           detrend = False,
                                                           average = average)
        if len(Freq) == 0:
            Freq = subFreq[:10]
            PSD = subPSD[:10]
            stdPSD1 = substdPSD1[:10]
            stdPSD2 = substdPSD2[:10]
110
111
112
113
        Freq = np.concatenate((Freq,subFreq[10:100]))
        PSD = np.concatenate((PSD,subPSD[10:100]))
        stdPSD1 = np.concatenate((stdPSD1,substdPSD1[10:100]))
        stdPSD2 = np.concatenate((stdPSD2,substdPSD2[10:100]))
Anchal's avatar
Anchal committed
114

115
116
117
118
    Freq = np.concatenate((Freq,subFreq[100:]))
    PSD = np.concatenate((PSD,subPSD[100:]))
    stdPSD1 = np.concatenate((stdPSD1,substdPSD1[100:]))
    stdPSD2 = np.concatenate((stdPSD2,substdPSD2[100:]))
Anchal's avatar
Anchal committed
119
120
121
122
123

    if average == 'mean':
        return Freq, PSD, stdPSD1
    else:
        return Freq, PSD, stdPSD1, stdPSD2