text:
'''
This is an unofficial Python implementation of Hearing-Aid Speech Perception Index (HASPI)
Usage:
    from pyhaspi2 import haspi_v2
    haspi(x,fx,y,fy, HL = np.zeros(6))
    haspi_v2(x, fx, y, fy, HL = np.zeros(6))

    Where x is clean reference speech, fx is sampling rate of x, and y is degraded speech, and fy is sampling rate of y. HL is 6-dim vectors represents hearing loss at the 6 audiometric frequencies, with 0 normal.

Differences between original Matlab implementation by Prof. James M.Kates:
    1). In our implementation, we did not apply built-in alignment, so please make sure x and y are already aligned before using this function
    2). In our implementation of HASPI version 2, we use sigmoid weight approximation model shown in Eq.(1) in [2] instead of neural network weights
    3). We internally normalized both x and y into rms=1, and SPL level was set to 65 dB
    4). We only confirmed output results are correct in the case of HL=[0,0,0,0,0,0] (as normal hearing condition). Although HL vector can be set arbitrary in the function, we cannot guarantee results are correct since lack of extensive test.

Reference:
    [1]. Kates, James M., and Kathryn H. Arehart. "The Hearing-Aid Speech Perception Index (HASPI)." Speech Communication 65 (2014): 75-93.
    [2]. Kates, James M., and Kathryn H. Arehart. "The Hearing-Aid Speech Perception Index (HASPI) Version 2." Speech Communication (2020).
'''

import numpy as np
import librosa
import scipy
from scipy.signal import lfilter, group_delay
from numba import jit
from scipy.interpolate import interp1d
import cProfile
import time

def haspi(x, fx, y, fy, HL = np.zeros(6), alpha = -1.0):
    Level1 = 65
    L = min(len(x), len(y))
    x = x[:L]
    y = y[:L]
    rms_x = np.sqrt(np.sum(x**2) / L)
    rms_y = np.sqrt(np.sum(y**2) / L)
    x = x / rms_x
    y = y / rms_y

    itype = 0 # Intelligibility model
    xenv, xBM, yenv, yBM, xSL, ySL, fsamp = eb_EarModel(x,fx,y,fy,HL,itype,Level1)

    # Smooth the envelope outputs: 125 Hz sub-sampling rate
    segsize = 16 # Averaging segment size in msec
    xdB = eb_EnvSmooth(xenv, segsize, fsamp)
    ydB = eb_EnvSmooth(yenv, segsize, fsamp)

    # Mel cepstrum correlation using smoothed envelopes
    # m1=ave of coefficients 2-6
    # xy=vector of coefficients 1-6
    thr=2.5  # Silence threshold: sum across bands, dB above aud threshold
    addnoise=0.0 # Additive noise in dB SL to condition cross-covariances
    CepCorr, xy = eb_melcor(xdB, ydB, thr, addnoise)

    # Temporal fine structure correlation measurements
    # Compute the time-frequency segment covariances
    segcov=16 # Segment size for the covariance calculation
    sigcov,sigMSx,sigMSy = eb_BMcovary(xBM,yBM,segcov,fsamp)

    # Three-level signal segment covariance
    # cov3 vector:   [low, mid, high] intensity region average, uniform weights
    # covSII vector: [low, mid, high] with SII frequency band weights
    cov3,covSII = eb_3LevelCovary(sigcov,sigMSx,thr)

    # Intelligibility prediction
    # Combine the cepstral correlation and three-level covariance
    bias = -9.047
    wgtcep = 14.816
    wgtcov = np.array([0, 0, 4.616]) # [low, mid, high]
    arg = bias + wgtcep * CepCorr + np.sum(wgtcov * cov3)

    # Logsig transformation
    Intel = 1.0 / (1.0 + np.exp(alpha * arg)) # Logistic (logsig) function, default alpha = -1.0

    # Raw data
    raw = np.concatenate((np.array([CepCorr]), cov3))

    return Intel, raw

def eb_3LevelCovary(sigcov,sigMSx,thr):
    cov3 = np.zeros(3)
    covSII = np.zeros(3)
    nbands = sigcov.shape[0]
    # Find the segments that lie sufficiently above the threshold.
    sigRMS = np.sqrt(sigMSx)
    sigLinear = np.power(10, (sigRMS/20))
    xsum = np.sum(sigLinear, axis=0) / nbands
    xsum = 20*np.log10(xsum)
    index = np.where(xsum > thr)[0]
    nseg = len(index)
    if nseg <= 1:
        raise Exception('Function eb_3LevelCovary: Signal below threshold, outputs set to 0.')
    # Critical band center frequencies in Hz
    cfSII = [150, 250, 350, 450, 570, 700, 840, 1000, 1170, 1370, 1600, 1850, 2150, 2500, 2900, 3400, 4000, 4800, 5800, 7000, 8500]
    # cfSII = np.array(cfSII)
    # Weights for the Speech Intelligibility Index calculation
    wgtSII=[.0103, .0261, .0419, .0577, .0577, .0577, .0577, .0577, .0577, .0577, .0577, .0577, .0577, .0577, .0577, .0577, .0577, .0460, .0343, .0226, .0110]
    cfreq=eb_CenterFreq(nbands)
    fsamp=24000
    cfSII = [0] + cfSII + [fsamp]
    wgtSII = [0] + wgtSII + [0]
    CubicSpl = interp1d(np.array(cfSII), np.array(wgtSII), kind='cubic')
    wfreq = CubicSpl(cfreq)
    wfreq[0] = 0.0
    wfreq[1] = 0.0
    wfreq = wfreq / np.sum(wfreq)
    sigcov=sigcov[:,index]
    sigRMS=sigRMS[:,index]
    xsum = xsum[index]
    # Histogram of the segment intensities in phons
    dBmin = np.min(xsum)
    dBmax = np.max(xsum)
    dBstep = 0.5 # Bin width is 0.5 dB
    bins = np.arange(dBmin, dBmax+dBstep, dBstep)
    bin_pinf = 1e8
    bin_ninf = -1e8
    bin_max = np.concatenate((bins[1:], np.array([bin_pinf])))
    bin_tmp = (bins+bin_max) / 2
    bin_tmp = np.concatenate((np.array([bin_ninf]),bin_tmp))
    xhist, _ = np.histogram(xsum, bin_tmp)
    bincenters = bins
    nbins = len(xhist)

    # Compute the cumulative histogram
    xcum=np.zeros(nbins)
    xcum[0] = xhist[0]
    for k in range(1, nbins):
        xcum[k] = xcum[k-1] + xhist[k]

    xcum = xcum / xcum[nbins-1]
    # Find the boundaries for the lower, middle, and upper thirds
    edge=np.zeros(2)
    for n in range(nbins):
        if xcum[n] < 0.333:
            edge[0] = bincenters[n]
        if xcum[n] < 0.667:
            edge[1] = bincenters[n]

    # Assign segments to the lower, middle, and upper thirds
    low = np.where(xsum < edge[0])[0] # Segment indices for lower third
    tmp1 = (xsum >= edge[0])
    tmp2 = (xsum < edge[1])
    mid = np.where(tmp1*tmp2)[0]
    up = np.where(xsum>=edge[1])[0]

    # Compute the time-frequency weights. The weight=1 if a segment in a
    # frequency band is above threshold, and weight=0 if at or below threshold.
    weight = np.zeros([nbands, nseg])
    weight[sigRMS > thr] = 1
    sigcov = weight * sigcov # Apply the weights
    # Average the covariance across segment levels as a function of frequency
    cov_ave = np.zeros(nbands)
    cov_ave_SII = np.zeros(nbands)

    # Low-level segments
    s=sigcov[:,low] # Segment covariances for the low intensity
    w=weight[:,low] # Low intensity time-freq segments above threshold
    ssum = np.sum(s,axis=1) # Sum over the low-intensity segments in each frequency band
    wsum = np.sum(w,axis=1) # Sum of the above-threshold weights in each frequency band
    ncount=0
    wgtsum=0
    for n in range(nbands):
        if wsum[n]==0:
            cov_ave[n] = 0
            cov_ave_SII[n] = 0
        else:
            cov_ave[n] = ssum[n] / wsum[n]
            cov_ave_SII[n] = cov_ave[n] * wfreq[n]
            wgtsum = wgtsum + wfreq[n]
            ncount = ncount + 1
    cov3[0] = np.sum(cov_ave) / ncount
    covSII[0] = np.sum(cov_ave_SII) / wgtsum
    # Mid-level segments
    s=sigcov[:,mid] # Segment covariances for the mid intensity
    w=weight[:,mid] # Mid intensity time-freq segments above threshold
    ssum = np.sum(s,axis=1) # Sum over the mid-intensity segments in each frequency band
    wsum = np.sum(w,axis=1) # Sum of the above-threshold weights in each frequency band
    ncount=0
    wgtsum=0
    for n in range(nbands):
        if wsum[n]==0:
            cov_ave[n] = 0
            cov_ave_SII[n] = 0
        else:
            cov_ave[n] = ssum[n] / wsum[n]
            cov_ave_SII[n] = cov_ave[n] * wfreq[n]
            wgtsum = wgtsum + wfreq[n]
            ncount = ncount + 1
    cov3[1] = np.sum(cov_ave) / ncount
    covSII[1] = np.sum(cov_ave_SII) / wgtsum
    # High-level segments
    s=sigcov[:,up] # Segment covariances for the high intensity
    w=weight[:,up] # High intensity time-freq segments above threshold
    ssum = np.sum(s,axis=1) # Sum over the high-intensity segments in each frequency band
    wsum = np.sum(w,axis=1) # Sum of the above-threshold weights in each frequency band
    ncount=0
    wgtsum=0
    for n in range(nbands):
        if wsum[n]==0:
            cov_ave[n] = 0
            cov_ave_SII[n] = 0
        else:
            cov_ave[n] = ssum[n] / wsum[n]
            cov_ave_SII[n] = cov_ave[n] * wfreq[n]
            wgtsum = wgtsum + wfreq[n]
            ncount = ncount + 1
    cov3[2] = np.sum(cov_ave) / ncount
    covSII[2] = np.sum(cov_ave_SII) / wgtsum
    return cov3, covSII

def eb_BMcovary(xBM,yBM,segsize,fsamp):
    small=1.0e-30

    # Lag for computing the cross-covariance
    lagsize=1.0 # Lag (+/-) in msec
    maxlag=round(lagsize*(0.001*fsamp)) # Lag in samples

    nwin=round(segsize*(0.001*fsamp))
    test=nwin - 2*np.floor(nwin/2)
    if test>0:
        nwin = nwin + 1
    window=np.hanning(nwin) # Raised cosine von Hann window

    # TODO: translate this matlab code: wincorr=1./xcorr(window,window,maxlag);
    wincorr = np.array([0.00714486118736300, 0.00712980452227938, 0.00711541940039143, 0.00710170170277407, 0.00708864751353705, 0.00707625311737229, 0.00706451499720106, 0.00705342983191828, 0.00704299449423103, 0.00703320604858884, 0.00702406174920299, 0.00701555903815223, 0.00700769554357245, 0.00700046907792766, 0.00699387763635983, 0.00698791939511504, 0.00698259271004346, 0.00697789611517068, 0.00697382832133786, 0.00697038821490824, 0.00696757485653739, 0.00696538748000482, 0.00696382549110426, 0.00696288846658999, 0.00696257615317668, 0.00696288846658999, 0.00696382549110427, 0.00696538748000482, 0.00696757485653739, 0.00697038821490824, 0.00697382832133786, 0.00697789611517068, 0.00698259271004346, 0.00698791939511504, 0.00699387763635983, 0.00700046907792766, 0.00700769554357245, 0.00701555903815223, 0.00702406174920299, 0.00703320604858884, 0.00704299449423103, 0.00705342983191828, 0.00706451499720106, 0.00707625311737229, 0.00708864751353705, 0.00710170170277407, 0.00711541940039143, 0.00712980452227939, 0.00714486118736301])
    winsum2 = 1.0/np.sum(window**2)

    nhalf=nwin // 2
    halfwindow = window[nhalf:]

    # TODO: translate this matlab code: halfcorr=1./xcorr(halfwindow,halfwindow,maxlag);
    halfcorr = np.array([0.0171564012932667, 0.0169783665111901, 0.0168048368330894, 0.0166356935491691, 0.0164708221425357, 0.0163101121243555, 0.0161534568765699, 0.0160007535017831, 0.0158519026799599, 0.0157068085315890, 0.0155653784869908, 0.0154275231614608, 0.0152931562359640, 0.0151621943431025, 0.0150345569581033, 0.0149101662945780, 0.0147889472048270, 0.0146708270844672, 0.0145557357811784, 0.0144436055073713, 0.0143343707565927, 0.0142279682234914, 0.0141243367271791, 0.0140234171378276, 0.0139251523063533, 0.0140234171378276, 0.0141243367271791, 0.0142279682234914, 0.0143343707565927, 0.0144436055073713, 0.0145557357811784, 0.0146708270844672, 0.0147889472048270, 0.0149101662945780, 0.0150345569581033, 0.0151621943431025, 0.0152931562359640, 0.0154275231614608, 0.0155653784869908, 0.0157068085315891, 0.0158519026799599, 0.0160007535017831, 0.0161534568765699, 0.0163101121243555, 0.0164708221425357, 0.0166356935491691, 0.0168048368330894, 0.0169783665111901, 0.0171564012932667])
    halfsum2 = 1.0 / np.sum(halfwindow**2) # MS sum normalization, first segment

    nchan=xBM.shape[0]
    npts=xBM.shape[1]
    nseg = 1 + np.floor(npts/nwin) + np.floor((npts-nwin/2)/nwin)
    nseg = int(nseg)

    sigMSx = np.zeros(nchan*nseg)
    sigMSy = np.zeros(nchan*nseg)
    sigcov = np.zeros(nchan*nseg)

    for k in range(nchan):
        x = xBM[k,:]
        y = yBM[k,:]
        nstart = 0
        segx = x[nstart:nhalf] * halfwindow
        segy = y[nstart:nhalf] * halfwindow
        segx = segx - np.mean(segx)
        segy = segy - np.mean(segy)
        MSx = np.sum(segx**2) * halfsum2
        MSy = np.sum(segy**2) * halfsum2
        my_xcorr = np.correlate(segx,segy,'full')
        if len(my_xcorr) > 2*maxlag+1:
            start = (len(my_xcorr) - (2*maxlag+1)) // 2
            my_xcorr = my_xcorr[start: start+(2*maxlag+1)]

        Mxy = np.max(np.abs(my_xcorr * halfcorr))
        if (MSx > small) and (MSy > small):
            sigcov[k*nseg+0] = Mxy / np.sqrt(MSx*MSy)
        else:
            sigcov[k*nseg+0] = 0.0
        sigMSx[k*nseg+0] = MSx
        sigMSy[k*nseg+0] = MSy

        for n in range(1, nseg-1):
            nstart = nstart + nhalf
            nstop = nstart + nwin
            segx = x[nstart:nstop] * window
            segy = y[nstart:nstop] * window
            segx = segx - np.mean(segx)
            segy = segy - np.mean(segy)
            MSx = np.sum(segx**2) * winsum2
            MSy = np.sum(segy**2) * winsum2
            my_xcorr = np.correlate(segx,segy,'full')
            if len(my_xcorr) > 2*maxlag+1:
                start = (len(my_xcorr) - (2*maxlag+1)) // 2
                my_xcorr = my_xcorr[start: start+(2*maxlag+1)]

            Mxy = np.max(np.abs(my_xcorr * wincorr))
            if (MSx > small) and (MSy > small):
                sigcov[k*nseg+n] = Mxy / np.sqrt(MSx*MSy)
            else:
                sigcov[k*nseg+n] = 0.0
            sigMSx[k*nseg+n] = MSx
            sigMSy[k*nseg+n] = MSy

        nstart = nstart + nhalf
        nstop = nstart + nhalf
        segx=x[nstart:nstop]*window[:nhalf]
        segy=y[nstart:nstop]*window[:nhalf]
        segx = segx - np.mean(segx)
        segy = segy - np.mean(segy)
        MSx = np.sum(segx**2) * halfsum2
        MSy = np.sum(segy**2) * halfsum2
        my_xcorr = np.correlate(segx,segy,'full')
        if len(my_xcorr) > 2*maxlag+1:
            start = (len(my_xcorr) - (2*maxlag+1)) // 2
            my_xcorr = my_xcorr[start: start+(2*maxlag+1)]

        Mxy = np.max(np.abs(my_xcorr * halfcorr))
        if (MSx > small) and (MSy > small):
            sigcov[k*nseg+nseg-1] = Mxy / np.sqrt(MSx*MSy)
        else:
            sigcov[k*nseg+nseg-1] = 0.0
        sigMSx[k*nseg+nseg-1] = MSx
        sigMSy[k*nseg+nseg-1] = MSy

    # Limit the cross-covariance to lie between 0 and 1
    sigcov = np.clip(sigcov, a_min=0, a_max=1)
    # Adjust the BM magnitude to correspond to the envelope in dB SL
    sigMSx=2.0*sigMSx
    sigMSy=2.0*sigMSy

    sigcov = sigcov.reshape(nchan, nseg)
    sigMSx = sigMSx.reshape(nchan, nseg)
    sigMSy = sigMSy.reshape(nchan, nseg)

    return sigcov,sigMSx,sigMSy

# def my_xcorr(x, y, lags):
#     xy = np.correlate(x,y,'full')
#     if len(xy) > 2*lags+1:
#         start = (len(xy) - (2*lags+1)) // 2
#         xy = xy[start: start+(2*lags+1)]
#     elif len(xy) < 2*lags+1:
#         pads = (2*lags+1 - len(xy))//2
#         xy = np.concatenate((np.zeros(pads), xy, np.zeros(pads)))
#     return xy

@jit(nopython=True)
def eb_EnvSmooth(env,segsize,fsamp):
    nwin=round(segsize*(0.001*fsamp))
    test=nwin - 2*np.floor(nwin/2)
    if test>0:
        nwin = nwin + 1
    window=np.hanning(nwin) # Raised cosine von Hann window
    wsum=np.sum(window) # Sum for normalization
    nhalf=nwin//2
    halfwindow = window[nhalf:]
    halfsum = np.sum(halfwindow)
    # Number of segments and assign the matrix storage
    nchan = env.shape[0]
    npts = env.shape[1]
    nseg = 1 + np.floor(npts/nwin) + np.floor((npts-nwin/2)/nwin)
    nseg = int(nseg)
    smooth = np.zeros(nchan*nseg)
    for k in range(nchan):
        r = env[k, :]
        nstart = 0
        smooth[k*nseg + 0]= np.sum(r[nstart:nhalf] * halfwindow) / halfsum
        for n in range(1, nseg-1):
            nstart = nstart + nhalf
            nstop = nstart + nwin
            smooth[k*nseg+n] = np.sum(r[nstart:nstop] * window) / wsum
        nstart = nstart + nhalf
        nstop = nstart + nhalf
        smooth[k*nseg+nseg-1] = np.sum( r[nstart:nstop] * window[:nhalf] ) / halfsum

    smooth = smooth.reshape(nchan, nseg)
    return smooth

def eb_melcor(x,y,thr,addnoise):
    nbands = x.shape[0]
    nbasis=6 # Number of cepstral coefficients to be used
    freq = np.arange(0, nbasis)
    k = np.arange(0, nbands)
    cepm = np.zeros([nbands,nbasis])
    for nb in range(nbasis):
        basis = np.cos(freq[nb]*np.pi*k/(nbands-1))
        cepm[:, nb] = basis / np.linalg.norm(basis)

    # Find the segments that lie sufficiently above the quiescent rate
    xLinear = np.power(10, (x/20)) # Convert envelope dB to linear (specific loudness)
    xsum = np.sum(xLinear,axis=0) / nbands # Proportional to loudness in sones
    xsum = 20*np.log10(xsum) # Convert back to dB (loudness in phons)
    index = np.where(xsum > thr)[0]
    nsamp=len(index) # Number of segments above threshold
    if nsamp <= 1:
        raise Exception('Function eb_melcor: Signal below threshold, outputs set to 0.')
    x=x[:,index]
    y=y[:,index]
    if addnoise != 0.0:
        print('add noise in eb_melcor func!')
        x = x + addnoise * np.random.randn(x.shape[0], x.shape[1])
        y = y + addnoise * np.random.randn(y.shape[0], y.shape[1])
    xcep = np.zeros([nbasis,nsamp]) #Input
    ycep = np.zeros([nbasis,nsamp])  #Output
    for n in range(nsamp):
        for k in range(nbasis):
            xcep[k,n] = np.sum( x[:,n] * cepm[:,k]  )
            ycep[k,n] = np.sum( y[:,n] * cepm[:,k]  )

    for k in range(nbasis):
        xcep[k, :] = xcep[k, :] - np.mean(xcep[k, :])
        ycep[k, :] = ycep[k, :] - np.mean(ycep[k, :])

    xy=np.zeros(nbasis)
    small=1.0e-30
    for k in range(nbasis):
        xsum = np.sum( xcep[k,:]**2 )
        ysum = np.sum( ycep[k,:]**2 )
        if (xsum < small) or (ysum < small):
            xy[k] = 0.0
        else:
            xy[k] = np.abs( np.sum(xcep[k,:]*ycep[k,:]) / np.sqrt(xsum*ysum)  )
    m1 = np.sum(xy[1:]) / (nbasis-1)
    return m1, xy

def eb_CenterFreq(nchan, shift=None, placeholder=None):
    # Parameters for the filter bank
    lowFreq=80.0 # Lowest center frequency
    highFreq=8000.0 # Highest center frequency

    # Moore and Glasberg ERB values
    EarQ = 9.26449
    minBW = 24.7

    if shift and placeholder:
        k=1
        A=165.4
        a=2.1
        xLow = (1/a) * np.log10(k + (lowFreq/A))
        xHigh = (1/a) * np.log10(k + (highFreq/A))
        xLow=xLow*(1 + shift)
        xHigh=xHigh*(1 + shift)
        lowFreq = A * (np.power(10,(a*xLow)) - k)
        highFreq = A * (np.power(10,(a*xHigh)) - k)

    cf = -(EarQ * minBW) + np.exp(  np.arange(1, nchan) * (-np.log(highFreq + EarQ*minBW) + \
    np.log(lowFreq + EarQ*minBW))/(nchan-1)) * (highFreq + EarQ*minBW)
    cf = np.concatenate((np.array([highFreq]), cf)) # Last center frequency is set to highFreq
    cf = np.flipud(cf) # Reorder to put the low frequencies first
    return cf

def eb_LossParameters(HL, cfreq):
    aud = [250.0, 500.0, 1000.0, 2000.0, 4000.0, 6000.0]
    nfilt = len(cfreq)
    fv = [cfreq[0]] + aud + [cfreq[-1]]
    # TODO: simply use numpy interp to replace matlab interp1, might not be accurate!!!
    loss = np.interp(cfreq, fv, np.concatenate((np.array([HL[0]]),HL,np.array([HL[-1]]))))
    loss[loss<0] = 0.0
    # ---
    CR = np.zeros(nfilt)
    for i in range(nfilt):
        CR[i] = 1.25 + 2.25 * i / (nfilt-1)
    maxOHC = 70 * (1 - (1/CR))
    thrOHC=1.25*maxOHC
    attnOHC=np.zeros(nfilt) # Default is 0 dB attenuation
    attnIHC=np.zeros(nfilt)
    for i in range(nfilt):
        if loss[i] < thrOHC[i]:
            attnOHC[i]=0.8*loss[i]
            attnIHC[i]=0.2*loss[i]
        else:
            attnOHC[i]=0.8*thrOHC[i] # Maximum OHC attenuation
            attnIHC[i]=0.2*thrOHC[i] + (loss[i]-thrOHC[i])
    BW = np.ones(nfilt)
    BW = BW + (attnOHC / 50.0) + 2.0 * (attnOHC / 50.0)**6

    lowknee = attnOHC + 30 #  Lower kneepoint
    upamp = 30 + 70/CR # Output level for an input of 100 dB SPL
    CR = (100-lowknee)/(upamp + attnOHC - lowknee) # OHC loss Compression ratio
    return attnOHC,BW,lowknee,CR,attnIHC

def eb_Resamp24kHz(x,fsampx):
    fsamp = 24000
    if fsampx == fsamp:
        y = x
    elif fsampx < fsamp:
        y = librosa.resample(x, orig_sr=fsampx, target_sr=fsamp)
        xRMS = np.sqrt(np.mean(x**2))
        yRMS = np.sqrt(np.mean(y**2))
        y = (xRMS/yRMS) * y
    else:
        raise NotImplementedError
    return y, fsamp

def eb_InputAlign(x24, y24):
    nsamp = 24000
    raise NotImplementedError
    # Align x24 and y24
    return x24, y24

def eb_NALR(HL,nfir,fsamp):
    raise NotImplementedError

def eb_MiddleEar(x, fsamp):
    assert fsamp==24000
    bLP = np.array([0.434173751206302, 0.434173751206302])
    aLP = np.array([1.000000000000000, -0.131652497587396])
    y = lfilter(bLP, aLP, x)
    bHP = np.array([0.937260390269893, -1.874520780539785, 0.937260390269893])
    aHP = np.array([1.000000000000000, -1.870580640735279, 0.878460920344291])
    xout = lfilter(bHP, aHP, y)
    return xout

@jit(nopython=True)
def eb_CosSinCF(L, fs, cf):
    tpt = 2 * np.pi / fs
    npts = L
    cn=np.cos(tpt*cf)
    sn=np.sin(tpt*cf)
    coscf=np.zeros(npts)
    sincf=np.zeros(npts)
    cold=1
    sold=0
    coscf[0]=cold
    sincf[0]=sold
    for n in range(1, npts):
        arg=cold*cn + sold*sn
        sold=sold*cn - cold*sn
        cold=arg
        coscf[n]=cold
        sincf[n]=sold
    return coscf, sincf

def eb_GammatoneBM2(x,BWx,y,BWy,fs,cf):
    earQ=9.26449
    minBW=24.7
    ERB=minBW + (cf/earQ)
    # L = min(len(x),len(y))
    # x = x[:L]
    # y = y[:L]
    tpt = 2 * np.pi / fs
    tptBW=BWx*tpt*ERB*1.019
    a=np.exp(-tptBW)
    a1=4.0*a
    a2=-6.0*a*a
    a3=4.0*a*a*a
    a4=-a*a*a*a
    a5=4.0*a*a
    gain=2.0*(1-a1-a2-a3-a4)/(1+a1+a5)

    npts = len(x)
    coscf, sincf = eb_CosSinCF(npts, fs, cf)
    # cn=np.cos(tpt*cf)
    # sn=np.sin(tpt*cf)
    # coscf=np.zeros(npts)
    # sincf=np.zeros(npts)
    # cold=1
    # sold=0
    # coscf[0]=cold
    # sincf[0]=sold
    # for n in range(1, npts):
    #     arg=cold*cn + sold*sn
    #     sold=sold*cn - cold*sn
    #     cold=arg
    #     coscf[n]=cold
    #     sincf[n]=sold

    ureal = lfilter([1, a1, a5],[1, -a1, -a2, -a3, -a4], x * coscf)
    uimag = lfilter([1, a1, a5],[1, -a1, -a2, -a3, -a4], x * sincf)
    BMx=gain*(ureal*coscf + uimag*sincf)
    envx=gain*np.sqrt(ureal*ureal + uimag*uimag)

    tptBW=BWy*tpt*ERB*1.019
    a=np.exp(-tptBW)
    a1=4.0*a
    a2=-6.0*a*a
    a3=4.0*a*a*a
    a4=-a*a*a*a
    a5=4.0*a*a
    gain=2.0*(1-a1-a2-a3-a4)/(1+a1+a5)

    ureal = lfilter([1, a1, a5],[1, -a1, -a2, -a3, -a4], y * coscf)
    uimag = lfilter([1, a1, a5],[1, -a1, -a2, -a3, -a4], y * sincf)
    BMy = gain*(ureal*coscf + uimag*sincf)
    envy = gain * np.sqrt(ureal*ureal + uimag*uimag)
    return envx,BMx,envy,BMy

def eb_GammatoneEnv2(x,BWx,y,BWy,fs,cf):
    earQ=9.26449
    minBW=24.7
    ERB=minBW + (cf/earQ)
    # L = min(len(x),len(y))
    # x = x[:L]
    # y = y[:L]
    tpt = 2 * np.pi / fs
    tptBW=BWx*tpt*ERB*1.019
    a=np.exp(-tptBW)
    a1=4.0*a
    a2=-6.0*a*a
    a3=4.0*a*a*a
    a4=-a*a*a*a
    a5=4.0*a*a
    gain=2.0*(1-a1-a2-a3-a4)/(1+a1+a5)

    npts = len(x)
    coscf, sincf = eb_CosSinCF(npts, fs, cf)
    # cn=np.cos(tpt*cf)
    # sn=np.sin(tpt*cf)
    # coscf=np.zeros(npts)
    # sincf=np.zeros(npts)
    # cold=1
    # sold=0
    # coscf[0]=cold
    # sincf[0]=sold
    # for n in range(1, npts):
    #     arg=cold*cn + sold*sn
    #     sold=sold*cn - cold*sn
    #     cold=arg
    #     coscf[n]=cold
    #     sincf[n]=sold

    ureal = lfilter([1, a1, a5],[1, -a1, -a2, -a3, -a4], x * coscf)
    uimag = lfilter([1, a1, a5],[1, -a1, -a2, -a3, -a4], x * sincf)
    envx = gain * np.sqrt(ureal*ureal + uimag*uimag)

    tptBW=BWy*tpt*ERB*1.019
    a=np.exp(-tptBW)
    a1=4.0*a
    a2=-6.0*a*a
    a3=4.0*a*a*a
    a4=-a*a*a*a
    a5=4.0*a*a
    gain=2.0*(1-a1-a2-a3-a4)/(1+a1+a5)

    ureal = lfilter([1, a1, a5],[1, -a1, -a2, -a3, -a4], y * coscf)
    uimag = lfilter([1, a1, a5],[1, -a1, -a2, -a3, -a4], y * sincf)
    envy = gain * np.sqrt(ureal*ureal + uimag*uimag)

    return envx, envy

def eb_BWadjust(control,BWmin,BWmax,Level1):
    cRMS=np.sqrt(np.mean(control**2)) # RMS signal intensity
    cdB=20*np.log10(cRMS) + Level1 # Convert to dB SPL
    if cdB < 50:
        BW = BWmin
    elif cdB > 100:
        BW=BWmax
    else:
        BW=BWmin + ((cdB-50)/50)*(BWmax-BWmin)
    return BW

def eb_EnvCompressBM(envsig,bm,control,attnOHC,thrLow,CR,fsamp,Level1):
    thrHigh=100.0
    small=1.0e-30

    logenv = np.clip(control,a_min=small,a_max=None) # Don't want to take logarithm of zero or neg
    logenv = Level1 + 20*np.log10(logenv)
    logenv = np.clip(logenv, a_min=thrLow, a_max = thrHigh)

    gain = -attnOHC - (logenv - thrLow)*(1 - (1/CR))
    gain=np.power(10, (gain/20))
    flp=800
    b = [0.095107983402496, 0.095107983402496]
    a = [1.000000000000000, -0.809784033195007]
    gain = lfilter(b, a, gain)

    y = gain * envsig
    b = gain * bm
    return y, b

@jit(nopython=True)
def eb_EnvAlign(x,y):
    fsamp=24000 # Sampling rate in Hz
    range_=100 # Range in msec for the xcorr calculation
    lags=round(0.001*range_*fsamp) # Range in samples
    npts=len(x)
    lags = min(lags,npts)  # Use min of lags, length of the sequence
    xy = np.convolve(x, y[::-1])
    # xy = np.correlate(x,y,'full')
    if len(xy) > 2*lags-1:
        start = (len(xy) - (2*lags-1)) // 2
        xy = xy[start: start+(2*lags-1)]
    elif len(xy) < 2*lags-1:
        pads = (2*lags-1 - len(xy))//2
        xy = np.concatenate((np.zeros(pads), xy, np.zeros(pads)))

    location = np.argmax(xy) + 1
    delay = lags - location
    if delay > 0:
        y=np.concatenate((y[delay:npts+1], np.zeros(delay)))
    elif delay < 0:
        delay=-delay
        y = np.concatenate( (np.zeros(delay), y[:npts-delay]) )

    return y

@jit(nopython=True)
def eb_IHCadapt(xdB,xBM,delta,fsamp):
    dsmall=1.0001
    if delta < dsmall:
        delta = dsmall
    # Initialize the adaptation time constants
    tau1=2 # Rapid adaptation in msec
    tau2=60 # Short-term adaptation in msec
    tau1=0.001*tau1
    tau2=0.001*tau2 # Convert to seconds

    # Equivalent circuit parameters
    T=1/fsamp # Sampling period
    R1=1/delta
    R2=0.5*(1 - R1)
    R3=R2
    C1=tau1*(R1 + R2)/(R1*R2)
    C2=tau2/((R1 + R2)*R3)

    # Intermediate values used for the voltage update matrix inversion
    a11=R1 + R2 + R1*R2*(C1/T)
    a12=-R1
    a21=-R3
    a22=R2 + R3 + R2*R3*(C2/T)
    denom=1.0/(a11*a22 - a21*a12)
    # Additional intermediate values
    R1inv=1.0/R1
    R12C1=R1*R2*(C1/T)
    R23C2=R2*R3*(C2/T)

    # Initialize the outputs and state of the equivalent circuit
    nsamp=len(xdB)
    gain=np.ones(xdB.shape) # Gain vector to apply to the BM motion, default is 1
    ydB=np.zeros(xdB.shape) # Assign storage
    V1=0.0
    V2=0.0
    small=1.0e-30
    for n in range(nsamp):
        V0=xdB[n]
        b1=V0*R2 + R12C1*V1
        b2=R23C2*V2
        V1=denom*(a22*b1 - a12*b2)
        V2=denom*(-a21*b1 + a11*b2)
        out=(V0 - V1)*R1inv
        out = max(out, 0.0) # Envelope can not drop below threshold
        ydB[n]=out # Envelope with IHC adaptation
        gain[n]=(out + small)/(V0 + small) # Avoid division by zero

    # gain = (ydB + small) / (xdB + small)
    yBM = gain * xBM
    return ydB, yBM

def eb_EnvSL2(env,bm,attnIHC,Level1):
    small=1.0e-30
    y=Level1 - attnIHC + 20*np.log10(env + small)
    y[y<0] = 0.0

    #Convert the linear BM motion to have a dB SL envelope
    gain=(y + small) / (env + small) # Gain that converted the env from lin to dB SPL
    b=gain*bm # Apply gain to BM motion
    return y, b

def eb_BMaddnoise(x,thr,Level1):
    gn = 10**((thr - Level1)/20.0)
    noise = gn*np.random.randn(len(x))
    y = x + noise
    return y

def eb_GroupDelayComp(xenv,BW,cfreq,fsamp):
    nchan = len(BW)
    # Filter ERB from Moore and Glasberg (1983)
    earQ=9.26449
    minBW=24.7
    ERB=minBW + (cfreq/earQ)

    # Initialize the gamatone filter coefficients
    tpt=2*np.pi/fsamp
    tptBW=tpt*1.019*BW*ERB
    a=np.exp(-tptBW)
    a1=4.0*a
    a2=-6.0*a*a
    a3=4.0*a*a*a
    a4=-a*a*a*a
    a5=4.0*a*a

    gd = np.zeros(nchan)
    for n in range(nchan):
        _, gd[n] = group_delay(( [1, a1[n], a5[n]] , [1, -a1[n], -a2[n], -a3[n], -a4[n]] ), w = 1 )
    gd = np.round(gd)
    gmin = np.min(gd)
    gd=gd - gmin
    gmax=np.max(gd)
    correct = gmax - gd

    yenv=np.zeros(xenv.shape)
    for n in range(nchan):
        r = xenv[n,:]
        npts=len(r)
        r = np.concatenate( ( np.zeros(int(correct[n])), r[0:npts-int(correct[n])] ) )
        yenv[n, :] = r

    return yenv

def eb_aveSL(env,control,attnOHC,thrLow,CR,attnIHC,Level1):
    thrHigh=100.0
    # Convert the control to dB SPL
    small=1.0e-30
    logenv = np.clip(control, a_min=small, a_max=None)
    logenv = Level1 + 20*np.log10(logenv)
    logenv = np.clip(logenv, a_min=thrLow, a_max=thrHigh)

    # Compute the compression gain in dB
    gain = -attnOHC - (logenv - thrLow) * (1 - (1/CR)) # Gain in dB

    # Convert the signal envelope to dB SPL
    logenv = np.clip(env, a_min=small, a_max=None)
    logenv = Level1 + 20*np.log10(logenv)
    logenv = np.clip(logenv, a_min=0, a_max=None)
    xdB = logenv + gain - attnIHC # Apply gain to the log spectrum
    xdB = np.clip(xdB, a_min=0.0, a_max=None) # dB SL
    return xdB

def eb_EarModel(x,fx,y,fy,HL,itype,Level1):
    # assert itype==0
    nchan = 32 # Use 32 auditory frequency bands
    mdelay = 1 # Compensate for the gammatone group delay
    cfreq = eb_CenterFreq(nchan) # Center frequencies on an ERB scale
    attnOHCy,BWminy,lowkneey,CRy,attnIHCy = eb_LossParameters(HL, cfreq)

    if itype==0:
        HLx=0*HL
    else:
        HLx=HL

    attnOHCx,BWminx,lowkneex,CRx,attnIHCx = eb_LossParameters(HLx, cfreq)
    HLmax = 100*np.ones(6)
    shift = 0.02 # Basal shift of 0.02 of the basilar membrane length
    cfreq1 = eb_CenterFreq(nchan, shift=shift) # Center frequencies for the control
    _, BW1, _, _, _ = eb_LossParameters(HLmax,cfreq1) # Maximum BW for the control
    x24, _ = eb_Resamp24kHz(x, fx)
    y24, fsamp = eb_Resamp24kHz(y, fy)
    # Bulk broadband signal alignment
    # x24, y24 = eb_InputAlign(x24, y24)
    nsamp = len(x24)
    if itype==1:
        nfir = 140
        nalr, _ = eb_NALR(HL,nfir,fsamp)
        x24 = np.convolve(x24, nalr)
        x24=x24[nfir:nfir+nsamp]

    # Cochlear model
    # Middle ear
    xmid = eb_MiddleEar(x24, fsamp)
    ymid = eb_MiddleEar(y24, fsamp)

    # Initialize storage
    # Reference and processed envelopes and BM motion
    xdB = np.zeros([nchan, nsamp])
    ydB = np.zeros([nchan, nsamp])
    xBM = np.zeros([nchan, nsamp])
    yBM = np.zeros([nchan, nsamp])
    xave = np.zeros(nchan)
    yave = np.zeros(nchan)
    xcave = np.zeros(nchan)
    ycave = np.zeros(nchan)
    BWx = np.zeros(nchan)
    BWy = np.zeros(nchan)

    for n in range(nchan):
        xcontrol, ycontrol = eb_GammatoneEnv2(xmid,BW1[n],ymid,BW1[n],fsamp,cfreq1[n])
        # Adjust the auditory filter bandwidths for the average signal level
        BWx[n]=eb_BWadjust(xcontrol,BWminx[n],BW1[n],Level1)  # Reference
        BWy[n]=eb_BWadjust(ycontrol,BWminy[n],BW1[n],Level1)  # Processed
        # Envelopes and BM motion of the reference and processed signals
        xenv,xbm,yenv,ybm = eb_GammatoneBM2(xmid,BWx[n],ymid,BWy[n],fsamp,cfreq[n])
        xave[n]=np.sqrt(np.mean(xenv**2)) # Ave signal mag in each band
        yave[n]=np.sqrt(np.mean(yenv**2))
        xcave[n]=np.sqrt(np.mean(xcontrol**2)) # Ave control signal
        ycave[n]=np.sqrt(np.mean(ycontrol**2))
        # Cochlear compression for the signal envelopes and BM motion
        xc,xb = eb_EnvCompressBM(xenv,xbm,xcontrol,attnOHCx[n],lowkneex[n], CRx[n],fsamp,Level1)
        yc,yb = eb_EnvCompressBM(yenv,ybm,ycontrol,attnOHCy[n],lowkneey[n], CRy[n],fsamp,Level1)

        # Correct for the delay between the reference and output
        # TODO: speed up Align function
        # yc=eb_EnvAlign(xc,yc)
        # yb=eb_EnvAlign(xb,yb)
        # Convert the compressed envelopes and BM vibration envelopes to dB SL

        xc,xb = eb_EnvSL2(xc,xb,attnIHCx[n],Level1)
        yc,yb = eb_EnvSL2(yc,yb,attnIHCy[n],Level1)
        # Apply the IHC rapid and short-term adaptation
        delta=2.0
        xdB_tmp,xb=eb_IHCadapt(xc,xb,delta,fsamp)
        xdB[n,:] = xdB_tmp
        ydB_tmp,yb=eb_IHCadapt(yc,yb,delta,fsamp)
        ydB[n,:] = ydB_tmp
        # Additive noise level to give the auditory threshold
        IHCthr=-10.0
        xBM_tmp = eb_BMaddnoise(xb,IHCthr,Level1)
        xBM[n,:] = xBM_tmp
        yBM_tmp = eb_BMaddnoise(yb,IHCthr,Level1)
        yBM[n,:] = yBM_tmp
        # print('finished: ',n)

    if mdelay > 0:
        xdB=eb_GroupDelayComp(xdB,BWx,cfreq,fsamp)
        ydB=eb_GroupDelayComp(ydB,BWx,cfreq,fsamp)
        xBM=eb_GroupDelayComp(xBM,BWx,cfreq,fsamp)
        yBM=eb_GroupDelayComp(yBM,BWx,cfreq,fsamp)

    # Convert average gammatone outputs to dB SL
    xSL = eb_aveSL(xave,xcave,attnOHCx,lowkneex,CRx,attnIHCx,Level1)
    ySL = eb_aveSL(yave,ycave,attnOHCy,lowkneey,CRy,attnIHCy,Level1)

    return xdB, xBM, ydB, yBM, xSL, ySL, fsamp

if __name__ == "__main__":
    x, fx = librosa.load('sig_clean.wav', sr=None)
    y, fy = librosa.load('sig_out.wav', sr=None)

    # start = time.time()
    # for i in range(50):
    #     [score, raw] = haspi(x,fx,y,fy)
    # end = time.time()
    # print((end-start)/50)

    # [score, raw] = haspi(x,fx,y,fy)
    # print(score)
    # print(raw)
    # cProfile.run('haspi(x,fx,y,fy)')
    direct_wav, fs = librosa.load('direct_wav.wav', sr=None)
    reverb_wav2, fs = librosa.load('reverb_wav2.wav', sr=None)
    # reverb_wav3, fs = librosa.load('reverb_wav3.wav', sr=None)
    # reverb_wav4, fs = librosa.load('reverb_wav4.wav', sr=None)
    n, fs = librosa.load('cafenoise.wav', sr=None)
    assert fs==16000

    [score, raw] = haspi(direct_wav,fs,direct_wav+n,fs)
    print(score)
    [score, raw] = haspi(direct_wav,fs,reverb_wav2,fs)
    print(score)

    [score, raw] = haspi(direct_wav,fs,reverb_wav2+0.5*n,fs)
    print(score)

    [score, raw] = haspi(direct_wav,fs,reverb_wav2+1.2*n,fs)
    print(score)

    [score, raw] = haspi(direct_wav,fs,reverb_wav2+2.1*n,fs)
    print(score)

max_stars_repo_name: nii-yamagishilab/NELE-GAN

max_issues_repo_name: nii-yamagishilab/NELE-GAN

max_forks_repo_name: nii-yamagishilab/NELE-GAN

reason: import numpy,import scipy,from scipy,from numba


text:
import csv
import numpy as np
from scipy import signal

# validates boundary values of an array/list (prevents errors)
def check_bounds(l_bound, u_bound, rows):
    # ensures that the lower bound is smaller than the upper bound
    if l_bound >= u_bound:
        l_bound = u_bound - 1
    # checks if bounds are valid (prevents index out of bounds error)
    if (l_bound < 0) | (l_bound >= len(rows)):
        l_bound = 0
    if (u_bound > len(rows)) | (u_bound <= 0):
        u_bound = len(rows)
    return l_bound, u_bound

# reads data in a csv file and puts them in a 2D list
def read_data(file_name, l_bound, u_bound):
    data = []
    with open(file_name, "r") as file:
        reader = csv.reader(file)
        rows = list(reader)

        # ensures bounds are valid
        [l_bound, u_bound] = check_bounds(l_bound, u_bound, rows)

        # reads through the file and puts the data in the respective lists above
        for i in range(l_bound, u_bound):
            row = rows[i]
            data.append(list(np.float_(row)))
    # reshapes the list into a 2D numpy array with each feature/label being its own sub-array
    return np.vstack(data).T

# filters the input data to estimate the intended movement
def filter_data(data, TIME_PERIOD, zero_phase=True):
    # butterworth based on IIR filter is used
    [b, a] = get_filter_coefficients(TIME_PERIOD)
    # zero phase filter is used to generate the labels (slower but no distortion)
    if zero_phase:
        filtered_data = signal.filtfilt(b, a, data)
    # has distortion in the beginning and end of the data but is usable in real-time
    elif not zero_phase:
        filtered_data = signal.lfilter(b, a, data)
    else:
        print("Error: Invalid boolean input!")
        exit()
    return np.ndarray.tolist(filtered_data)  # converts np array to list

def get_filter_coefficients(TIME_PERIOD):
    nyquist = 1 / (2 * TIME_PERIOD)
    cut_off = 5 / nyquist

    # 2nd order IIR butterworth filter is used
    return signal.butter(2, cut_off, btype='lowpass')

ext: py

reason: import numpy,from scipy


text:
# Python 2 compatibility
from __future__ import print_function
from __future__ import division

# Standard libraries
import os
import pkg_resources

import numpy as np
import pytest

try:
    from lxml import etree
except ImportError:
    try:
        import xml.etree.cElementTree as etree
    except ImportError:
        try:
            import xml.etree.ElementTree as etree
        except ImportError:
          print("Failed to import ElementTree from any known place")
          raise

# Local imports
from .. import parse_files_XML
from ..simulation import Simulation
from ..utils import units

#pytestmark = pytest.mark.skip(reason="XML converter not completed")

class TestExperimentType:
    """
    """
    def test_shock_tube_experiment(self):
        """Ensure shock tube experiment can be detected.
        """
        file_path = os.path.join('testfile_st.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        kind = parse_files_XML.get_experiment_kind(root)
        assert kind == 'ST'

    def test_RCM_experiment(self):
        """Ensure rapid compression machine experiment can be detected.
        """
        file_path = os.path.join('testfile_rcm.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        kind = parse_files_XML.get_experiment_kind(root)
        assert kind == 'RCM'

class TestCommonProperties:
    """
    """
    def test_shock_tube_common_properties(self):
        """Ensure basic common properties parsed for shock tube.
        """
        file_path = os.path.join('testfile_st.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        properties = {}
        properties['kind'] = parse_files_XML.get_experiment_kind(root)
        properties = parse_files_XML.get_common_properties(properties, root)

        # Check pressure
        assert properties['pressure'] == 2.18 * units.atm

        # Check initial composition
        assert properties['composition']['H2'] == 0.00444
        assert properties['composition']['O2'] == 0.00566
        assert properties['composition']['Ar'] == 0.9899

        # Check pressure rise
        assert 'pressure-rise' not in properties

        # Make sure no other properties present
        assert (set(properties.keys()) ==
                set(['kind', 'pressure', 'composition'])
                )

    def test_shock_tube_common_properties_pressure_rise(self):
        """Ensure basic common properties parsed for shock tube.
        """
        file_path = os.path.join('testfile_st2.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        properties = {}
        properties['kind'] = parse_files_XML.get_experiment_kind(root)
        properties = parse_files_XML.get_common_properties(properties, root)

        # Check pressure
        assert properties['pressure'] == 2.18 * units.atm

        # Check initial composition
        assert properties['composition']['H2'] == 0.00444
        assert properties['composition']['O2'] == 0.00566
        assert properties['composition']['Ar'] == 0.9899

        # Check pressure rise
        assert properties['pressure-rise'] == 0.10 / units.ms

        # Make sure no other properties present
        assert (set(properties.keys()) ==
                set(['kind', 'pressure', 'pressure-rise', 'composition'])
                )

    def test_rcm_common_properties(self):
        """Ensure basic common properties parsed for RCM.
        """
        file_path = os.path.join('testfile_rcm.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        properties = {}
        properties['kind'] = parse_files_XML.get_experiment_kind(root)
        properties = parse_files_XML.get_common_properties(properties, root)

        # Check initial composition
        assert properties['composition']['H2'] == 0.12500
        assert properties['composition']['O2'] == 0.06250
        assert properties['composition']['N2'] == 0.18125
        assert properties['composition']['Ar'] == 0.63125

        assert set(properties.keys()) == set(['kind', 'composition'])

class TestIgnitionType:
    """
    """
    def test_pressure_ignition_target(self):
        """Test pressure max derivative as target for RCM.
        """
        file_path = os.path.join('testfile_rcm.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        ignition = parse_files_XML.get_ignition_type(root)

        assert ignition['target'] == 'P'
        assert ignition['type'] == 'd/dt max'

    def test_pressure_species_target(self):
        """Test pressure max derivative as target for shock tube.
        """
        file_path = os.path.join('testfile_st.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        ignition = parse_files_XML.get_ignition_type(root)

        assert ignition['target'] == 'P'
        assert ignition['type'] == 'd/dt max'

    def test_pressure_species_target_OH(self):
        """Test species max value as targetree.
        """
        file_path = os.path.join('testfile_st2.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        ignition = parse_files_XML.get_ignition_type(root)

        assert ignition['target'] == 'OH'
        assert ignition['type'] == 'max'

class TestDataGroups:
    """
    """
    def test_shock_tube_data_points(self):
        """Test parsing of ignition delay data points for shock tube file.
        """
        file_path = os.path.join('testfile_st.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        properties = {}
        properties = parse_files_XML.get_datapoints(properties, root)

        # Ensure correct temperature and ignition delay values and units
        np.testing.assert_array_equal(properties['temperature'],
                                      [1164.48, 1164.97] * units.kelvin
                                      )
        np.testing.assert_array_equal(properties['ignition-delay'],
                                      [471.54, 448.03] * units.us
                                      )

    def test_shock_tube_data_points_pressure_rise(self):
        """Test parsing of ignition delay data points for shock tube file.
        """
        file_path = os.path.join('testfile_st2.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        properties = {}
        properties = parse_files_XML.get_datapoints(properties, root)

        # Ensure correct temperature and ignition delay values and units
        np.testing.assert_allclose(properties['temperature'],
                                   1264.2 * units.kelvin
                                   )
        np.testing.assert_allclose(properties['ignition-delay'],
                                   291.57 * units.us
                                   )

    def test_rcm_data_points(self):
        """Test parsing of ignition delay data points for RCM file.
        """
        file_path = os.path.join('testfile_rcm.xml')
        filename = pkg_resources.resource_filename(__name__, file_path)
        tree = etree.parse(filename)
        root = tree.getroot()

        properties = {}
        properties = parse_files_XML.get_datapoints(properties, root)

        # Ensure correct temperature, pressure, and ignition delay values and units
        assert properties['temperature'] == 297.4 * units.kelvin
        assert properties['pressure'] == 958. * units.torr
        assert properties['ignition-delay'] == 1. * units.ms

        # Check other data group with volume history
        np.testing.assert_allclose(properties['time'],
                                   np.arange(0, 9.7e-2, 1.e-3) * units.second
                                   )

        volumes = np.array([
            5.47669375000E+002, 5.46608789894E+002, 5.43427034574E+002,
            5.38124109043E+002, 5.30700013298E+002, 5.21154747340E+002,
            5.09488311170E+002, 4.95700704787E+002, 4.79791928191E+002,
            4.61761981383E+002, 4.41610864362E+002, 4.20399162234E+002,
            3.99187460106E+002, 3.77975757979E+002, 3.56764055851E+002,
            3.35552353723E+002, 3.14340651596E+002, 2.93128949468E+002,
            2.71917247340E+002, 2.50705545213E+002, 2.29493843085E+002,
            2.08282140957E+002, 1.87070438830E+002, 1.65858736702E+002,
            1.44647034574E+002, 1.23435332447E+002, 1.02223630319E+002,
            8.10119281915E+001, 6.33355097518E+001, 5.27296586879E+001,
            4.91943750000E+001, 4.97137623933E+001, 5.02063762048E+001,
            5.06454851923E+001, 5.10218564529E+001, 5.13374097598E+001,
            5.16004693977E+001, 5.18223244382E+001, 5.20148449242E+001,
            5.21889350372E+001, 5.23536351113E+001, 5.25157124459E+001,
            5.26796063730E+001, 5.28476160610E+001, 5.30202402028E+001,
            5.31965961563E+001, 5.33748623839E+001, 5.35527022996E+001,
            5.37276399831E+001, 5.38973687732E+001, 5.40599826225E+001,
            5.42141273988E+001, 5.43590751578E+001, 5.44947289126E+001,
            5.46215686913E+001, 5.47405518236E+001, 5.48529815402E+001,
            5.49603582190E+001, 5.50642270863E+001, 5.51660349836E+001,
            5.52670070646E+001, 5.53680520985E+001, 5.54697025392E+001,
            5.55720927915E+001, 5.56749762728E+001, 5.57777790517E+001,
            5.58796851466E+001, 5.59797461155E+001, 5.60770054561E+001,
            5.61706266985E+001, 5.62600130036E+001, 5.63449057053E+001,
            5.64254496625E+001, 5.65022146282E+001, 5.65761642150E+001,
            5.66485675508E+001, 5.67208534842E+001, 5.67944133373E+001,
            5.68703658198E+001, 5.69493069272E+001, 5.70310785669E+001,
            5.71146023893E+001, 5.71978399741E+001, 5.72779572372E+001,
            5.73517897984E+001, 5.74167271960E+001, 5.74721573687E+001,
            5.75216388520E+001, 5.75759967785E+001, 5.76575701358E+001,
            5.78058719368E+001, 5.80849611077E+001, 5.85928651155E+001,
            5.94734357453E+001, 6.09310671165E+001, 6.32487551103E+001,
            6.68100309742E+001
            ]) * units.cm3
        np.testing.assert_allclose(properties['volume'], volumes)

max_stars_repo_name: vayner-s/PyTeCK

max_issues_repo_name: vayner-s/PyTeCK

max_forks_repo_name: vayner-s/PyTeCK

reason: import numpy


text:
import tensorflow as tf
import argparse
import os
from cloning_model import CloningModel
import numpy as np
import gym
import tempfile
import pickle
from mujoco_py.generated import const

def main():
    tf.enable_eager_execution()

    parser = argparse.ArgumentParser()
    parser.add_argument('envname', type=str)
    parser.add_argument('--max_timesteps', type=int)
    parser.add_argument('--render', action='store_true')
    parser.add_argument('--num_rollouts', type=int)

    args = parser.parse_args()

    def action_obs_dims_hack():
        expert_data = None
        import pickle
        with open(os.path.join('expert_data', args.envname + '.pkl'), 'rb') as f:
            print(f)
            expert_data = pickle.load(f)

        if expert_data is None:
            raise Exception('Failed to load expert data')

        observations = np.asarray(expert_data['observations'])
        actions = np.asarray(expert_data['actions'])
        actions = np.squeeze(actions)

        # initialize params for the agent
        action_dim = actions.shape[-1]
        state_dim = observations.shape[-1]
        return (action_dim, state_dim)

    checkpoint_dir = f'./checkpoints_{args.envname}'
    print(checkpoint_dir)
    os.makedirs(checkpoint_dir, exist_ok=True)
    action_dim, state_dim = action_obs_dims_hack()
    model = CloningModel(action_dim, state_dim, np.zeros(state_dim), np.zeros(state_dim), [400, 200, 100])
    optimizer = tf.train.AdamOptimizer(learning_rate=1e-5)

    root = tf.train.Checkpoint(model=model, optimizer= optimizer)
    print(tf.train.latest_checkpoint(checkpoint_dir))
    status = root.restore(tf.train.latest_checkpoint(checkpoint_dir))

    env = gym.make(args.envname)
    max_steps = args.max_timesteps or env.spec.timestep_limit

    rewards = []
    total_steps = []
    for i in range(args.num_rollouts):
        obs = env.reset()

        done = False

        accum_r = 0
        steps = 0

        if args.render:
            env.render(mode='rgb_array')
            if 'reacher' not in args.envname.lower():
                env.unwrapped.viewer.cam.type = const.CAMERA_FIXED
                env.unwrapped.viewer.cam.fixedcamid = 0

            f = tempfile.NamedTemporaryFile()
            print('named temp file', f.name)
            filename = f.name
            f.close()
            f = open(filename, 'ab')

        while not done:
            next_action = model.predict(obs[None, :])
            obs, reward, done, _ = env.step(next_action)
            accum_r += reward
            steps += 1

            if args.render:
                pickle.dump(env.render(mode='rgb_array'), f)
            if steps % 40 == 0:
                print(f'{steps}/{max_steps}')
                #print(next_action)

        if args.render:
            f.close()

        rewards.append(accum_r)
        total_steps.append(steps)
        print('steps', steps)
        print('totalr', accum_r)
    print('rewards', rewards)
    print('steps', total_steps)
    print('std:', np.std(rewards))
    print('mean:', np.mean(rewards))

if __name__ == '__main__':
    main()

max_stars_repo_name: baek-jinoo/homework

max_issues_repo_name: baek-jinoo/homework

max_forks_repo_name: baek-jinoo/homework


text:
import numpy as np
import pandas as pd
import datetime
import time
import glob
import os

# log path
path = "logs"
# table 1 or table 2 or table 3 (untargeted)
table = 2
# Lp norm

index = []
if table == 1:
    # small networks
    for network in ["mnist"]:
        for layer in ["2", "3"]:
            for norm in ["Li", "L2", "L1"]:
                for target in ["top2", "random", "least"]:
                    index.append("{}-{}-{}-{}".format(network, layer, norm, target))
    neuron_ub = 100
    neuron_lb = 0
    method = ["ours", "lip", "lp", "lpfull"]
elif table == 2:
    # targeted table
    for network in ["mnist"]:
        for layer in ["2", "3", "4"]:
            for norm in ["Li", "L2", "L1"]:
                for target in ["top2", "random", "least"]:
                    index.append("{}-{}-{}-{}".format(network, layer, norm, target))
    for network in ["cifar"]:
        for layer in ["5", "6", "7"]:
            for norm in ["Li", "L2", "L1"]:
                for target in ["top2", "random", "least"]:
                    index.append("{}-{}-{}-{}".format(network, layer, norm, target))
    neuron_ub = 3000
    neuron_lb = 1000
    method = ["ours", "lip", "lp"]
elif table == 3:
    # untargeted table
    network = "mnist"
    layer = "3"
    target = "untargeted"
    for norm in ["Li", "L2", "L1"]:
        index.append("{}-{}-{}-{}".format(network, layer, norm, target))
    network = "cifar"
    layer = "5"
    target = "untargeted"
    for norm in ["Li", "L2", "L1"]:
        index.append("{}-{}-{}-{}".format(network, layer, norm, target))
    neuron_ub = 3000
    neuron_lb = 1000
    method = ["ours", "lip", "lp"]
else:
    raise RuntimeError("unknown table " + str(table))

# create the table
dist_df = pd.DataFrame(dtype=np.float64,index=index, columns=method+["ratio"])
time_df = pd.DataFrame(dtype=np.float64,index=index, columns=method+["speedup"])
epoch_df = pd.DataFrame(dtype=np.int64,index=index, columns=method)
# table for output
output_dist_df = pd.DataFrame(dtype=str,index=index, columns=method+["ratio"])
output_time_df = pd.DataFrame(dtype=str,index=index, columns=method+["speedup"])

# fill in tables
for filename in glob.iglob(path + '/**/*.log', recursive=True):
    fn = os.path.basename(filename).split('_')
    network = fn[0]
    layer = fn[2]
    neuron = int(fn[3])
    norm = fn[4]
    method = fn[5]
    target = fn[6]
    date = fn[8]
    ltime = fn[9].split('.')[0]
    epoch = int(datetime.datetime.strptime('18'+date+ltime, '%y%m%d%H%M%S').timestamp())
    row_name = "{}-{}-{}-{}".format(network, layer, norm, target)
    if row_name in dist_df.index and neuron > neuron_lb and neuron < neuron_ub:
        if method not in dist_df.columns:
            continue
        old_number = dist_df.loc[row_name][method]
        update = True
        if old_number == old_number:
            print("WARNING: conflicting", filename, "row name is", row_name)
            prev_epoch = epoch_df.loc[row_name][method]
            epoch_str = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(epoch))
            prev_epoch_str = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(prev_epoch))
            print("Time stamp: prev {} current {}".format(prev_epoch_str, epoch_str))
            if epoch < prev_epoch:
                print("Skipped outdated file")
                update = False
        if update:
            with open(filename) as f:
                lines = f.readlines()
                result = lines[-1]
                if result[:5] != "[L0] ":
                    print("skipping incomplete", filename, "row name is", row_name)
                else:
                    result = result.split(",")
                    robustness = float(result[1].split("=")[1])
                    numimg = int(result[2].split("=")[1])
                    runtime = float(result[3].split("=")[1])
                    per_img_time = runtime / numimg
                    print("filename = {}\n row = {}, robustness = {}, time = {}".format(filename, row_name, robustness, per_img_time))
                    dist_df.loc[row_name][method] = robustness
                    time_df.loc[row_name][method] = per_img_time
                    epoch_df.loc[row_name][method] = epoch

    else:
        # print("skipping", filename, "row name is", row_name)
        pass

pd.set_option('display.width', 1000)
print(dist_df)
print(time_df)
def format_time(t):
    if t != t:
        return "-"
    if t >= 1.0:
        if t < 100:
            return "{:#.3g} s".format(t)
        else:
            return "{:.3g} s".format(t)
    else:
        t *= 1000
        if t < 100:
            return "{:#.3g} ms".format(t)
        else:
            return "{:.3g} ms".format(t)
def format_norm(n):
    if n == n:
        return "{:.5f}".format(n)
    else:
        return "-"

# compute ratios
for idx in dist_df.index.tolist():
    ours = dist_df.loc[idx]["ours"]
    lip = dist_df.loc[idx]["lip"]
    lp = dist_df.loc[idx]["lp"]
    if "lpfull" in dist_df:
        lpfull = dist_df.loc[idx]["lpfull"]
        if lpfull > lp:
            ref = lpfull
            ref_time = time_df.loc[idx]["lpfull"]
            output_dist_df.loc[idx]["lp"] = format_norm(lp)
            output_time_df.loc[idx]["lp"] = format_time(time_df.loc[idx]["lp"])
            output_dist_df.loc[idx]["lpfull"] = "\\bf "+format_norm(lpfull)
            output_time_df.loc[idx]["lpfull"] = "\\bf "+format_time(time_df.loc[idx]["lpfull"])
        else:
            ref = lp
            ref_time = time_df.loc[idx]["lp"]
            output_dist_df.loc[idx]["lp"] = "\\bf "+format_norm(lp)
            output_time_df.loc[idx]["lp"] = "\\bf "+format_time(time_df.loc[idx]["lp"])
            output_dist_df.loc[idx]["lpfull"] = format_norm(lpfull)
            output_time_df.loc[idx]["lpfull"] = format_time(time_df.loc[idx]["lpfull"])
    else:
        ref = lp
        ref_time = time_df.loc[idx]["lp"]
        if lp == lp:
            output_dist_df.loc[idx]["lp"] = "\\bf "+format_norm(lp)
            output_time_df.loc[idx]["lp"] = "\\bf "+format_time(time_df.loc[idx]["lp"])
        else:
            output_dist_df.loc[idx]["lp"] = format_norm(lp)
            output_time_df.loc[idx]["lp"] = format_time(time_df.loc[idx]["lp"])
    if ours < lip:
        new = lip
        new_time = time_df.loc[idx]["lip"]
        output_dist_df.loc[idx]["ours"] = format_norm(ours)
        output_time_df.loc[idx]["ours"] = format_time(time_df.loc[idx]["ours"])
        output_dist_df.loc[idx]["lip"] = "\\bf "+format_norm(lip)
        output_time_df.loc[idx]["lip"] = "\\bf "+format_time(time_df.loc[idx]["lip"])
    else:
        new = ours
        new_time = time_df.loc[idx]["ours"]
        output_dist_df.loc[idx]["ours"] = "\\bf "+format_norm(ours)
        output_time_df.loc[idx]["ours"] = "\\bf "+format_time(time_df.loc[idx]["ours"])
        output_dist_df.loc[idx]["lip"] = format_norm(lip)
        output_time_df.loc[idx]["lip"] = format_time(time_df.loc[idx]["lip"])
    ratio = (new - ref) / ref
    speedup = ref_time / new_time
    dist_df.loc[idx]["ratio"] = ratio
    if ratio == ratio:
        output_dist_df.loc[idx]["ratio"] = "{:+.1%}".format(ratio).replace("%", "\%")
    else:
        output_dist_df.loc[idx]["ratio"] = "-"
    time_df.loc[idx]["speedup"] = speedup
    if speedup == speedup:
        output_time_df.loc[idx]["speedup"] = "{}X".format(int(np.around(speedup)))
    else:
        output_time_df.loc[idx]["speedup"] = "-"

# add Latex &
for idx in output_dist_df.index.tolist():
    for c in output_dist_df.columns.values.tolist():
        output_dist_df.loc[idx][c] += " &"
for idx in output_time_df.index.tolist():
    for c in output_time_df.columns.values.tolist():
        output_time_df.loc[idx][c] += " &"

pd.set_option('display.width', 1000)
print(dist_df)
print(time_df)
print(output_dist_df)
print(output_time_df)


text:
'''Combine data sources on fuzzy-matched titles to create final data set'''
from multiprocessing import Pool, cpu_count
from functools import partial
import pandas as pd
import numpy as np
from fuzzywuzzy import process

from src.my_aws import S3

KEY_OMDB_TOR = 'OMDB_Torrents.csv'
KEY_NUM = 'TheNumbers_budgets.csv'
KEY_FINAL = 'Final_Data.csv'
BUCKET = 'movie-torrents'

def fuzzy_match(search_title, match_titles):
    '''
    Perform fuzzy match on provided title (searc_title)

    Args:
        search_title (str): Movie title to match up with
    Returns:
        dict: Dictionary of title searched, title it most closely
              matches and the percent that they match up
    '''
    match_title, match_percent = process.extractOne(
        search_title, match_titles)
    return {'RowTitle': search_title,
            'FuzzTitle': match_title,
            'FuzzPercent': match_percent}

class Aggregate():
    '''
    Combine data sources on fuzzy-matched titles and
    create final data set
    '''

    def __init__(self):
        '''
        Get a S3 connection object and pull down both data sources
        upon instantiation
        '''
        self.s3_obj = S3()
        self.tor_data = self.s3_obj.get_data(KEY_OMDB_TOR, BUCKET)
        self.num_data = self.s3_obj.get_data(KEY_NUM, BUCKET)
        self.final_data = None

    def make_fuzzy_match(self):
        '''
        Fuzzy match `The Numbers` data up to Torrent data up to
        a threshold of (95%) to prep Torrent data to be merged
        with `The Numbers` data

        Args:
            none
        Returns:
            pd.dataframe: Dataframe of torrent data with fuzzy
                          matched title to join on
        '''
        torrent_titles = self.tor_data['Title'].tolist()
        numbers_titles = self.num_data['title'].tolist()

        # Map fuzz search over available cpus
        num_cpus = cpu_count()
        worker_pool = Pool(num_cpus)
        fuzz_results = worker_pool.map(
            partial(fuzzy_match, match_titles=numbers_titles), torrent_titles)

        # fuzz_results = worker_pool.map(fuzzy_match, torrent_titles)
        worker_pool.close()
        worker_pool.join()

        # Put results into a dataframe and rename columns
        fuzz_data = pd.DataFrame(fuzz_results)
        fuzz_data = fuzz_data[['RowTitle', 'FuzzTitle', 'FuzzPercent']]
        fuzz_data.columns = ['Title', 'FuzzTitle', 'FuzzPercent']

        # Append to torrent dataframe and drop duplicate titles
        merge_df = pd.merge(self.tor_data, fuzz_data, how='inner', on='Title')
        merge_df.drop_duplicates(subset='Title', inplace=True)

        # Drop rows where match was below 95%
        merge_df = merge_df[merge_df['FuzzPercent'] >= 95]

        self.tor_data = merge_df
        return merge_df

    def clean_fuzz_data(self):
        '''
        Clean up date prior to merging with `The Numbers` data

        Args:
            none
        Returns:
            pd.dataframe: Dataframe of Torrent data ready to be merged
                          with `The Numbers` data
        '''
        tor = self.tor_data
        # remove where no torrent counts were received from any source
        tor['CheckTup'] = list(zip(tor['Kat_Count'].tolist(),
                                   tor['Pirate_Count'].tolist(),
                                   tor['Extra_Count'].tolist(),
                                   tor['Torrentz_Count'].tolist(),
                                   tor['Torrentz_Ver_Count'].tolist(),
                                   tor['Zoogle_Ver_Count'].tolist()))
        tor = tor[tor['CheckTup'] != ('Fail', 'Fail', 'Fail', 'Fail',
                                      'Fail', 'Fail')].reset_index(drop=True)
        # Drop temp column for checking
        del tor['CheckTup']

        # replace Fail, None, N and NaN with 0 - remove >, and <
        int_cols = ['Metascore', 'Runtime', 'imdbRating', 'imdbVotes',
                    'Kat_Count', 'Pirate_Count', 'Extra_Count',
                    'Torrentz_Count', 'Torrentz_Ver_Count', 'Zoogle_Ver_Count']

        for col in int_cols:
            tor[col] = tor[col].replace(['Fail', 'None', 'N', 'NaN'], '0')
            tor[col] = tor[col].apply(lambda x: str(x).replace(
                '>', '').replace('<', '').replace(',', ''))
            tor[col] = tor[col].replace(np.nan, 0)
            tor[col] = tor[col].fillna(value=0)
            tor[col] = pd.to_numeric(tor[col],
                                     errors='coerce',
                                     downcast='integer')
            tor[col] = tor[col].fillna(value=0)

        # fill in remaining NaN's with blanks
        tor.fillna(value='', inplace=True)
        self.tor_data = tor
        return tor

    def merge_data_sources(self):
        '''
        Merge Torrent data set with `The Numbers` data set

        Args:
            none
        Returns:
            pd.dataframe: Dataframe of resultant data set
        '''

        self.num_data.columns = ['FuzzTitle', 'ReleaseDate',
                                 'ProductionBudget', 'DomesticBudget',
                                 'WorldGross']

        # merge data frames
        data_df = pd.merge(self.tor_data, self.num_data,
                           how='left', on='FuzzTitle')
        data_df = data_df.drop_duplicates(subset='imdbID')

        self.final_data = data_df
        return data_df

    def prepare_final_data(self):
        '''
        Final preparations to final data set prior to
        upload to S3

        Args:
            none
        Returns:
            pd.dataframe: Dataframe of final data set
        '''
        data_df = self.final_data

        # Clean up dates
        data_df['Released'] = pd.to_datetime(data_df['Released'])
        data_df['ReleaseDate'] = pd.to_datetime(data_df['ReleaseDate'])
        data_df.loc[data_df['Released'].isnull(
        ), 'Released'] = data_df['ReleaseDate']

        # Drop columns no longer needed
        del data_df['ReleaseDate']
        del data_df['FuzzTitle']
        del data_df['FuzzPercent']

        # sum torrent counts
        data_df['Total_Torrents'] = data_df[['Kat_Count',
                                             'Pirate_Count',
                                             'Extra_Count',
                                             'Torrentz_Count',
                                             'Torrentz_Ver_Count',
                                             'Zoogle_Ver_Count']].sum(axis=1)

        self.final_data = data_df
        return data_df

    def put_data_s3(self):
        '''
        Upload clean data to S3 storage

        Args:
            none
        Returns:
            none
        '''
        self.s3_obj.put_data(self.final_data, KEY_FINAL, BUCKET)

    def export_data(self, write_file):
        '''
        Export data to csv file with the provided name/location (write_file)

        Args:
            write_file (str): Full file path of where to save csv file
        Returns:
            none
        '''
        self.final_data.to_csv(write_file, sep=',', index=False)

    def aggregate_data(self):
        '''
        Aggregate both data sources, clean, and then uplpoad to S3

        Args:
            none
        Returns:
            none
        '''
        self.make_fuzzy_match()
        self.clean_fuzz_data()
        self.merge_data_sources()
        self.prepare_final_data()
        self.put_data_s3()

if __name__ == '__main__':
    AGG = Aggregate()
    AGG.aggregate_data()

max_stars_repo_name: herokukuki/nadakk

max_issues_repo_name: herokukuki/nadakk

max_forks_repo_name: herokukuki/nadakk


text:
import pandas as pd
import numpy as np
import sys

outfile = sys.argv[1]

bio_df = pd.read_csv('~/biographies/metadata/hathi_ic_biog.tsv', sep='\t')

bioindex = pd.read_csv('/media/secure_volume/index/bioindex.tsv',sep='\t')

merge_bioindex = pd.merge(
    bio_df, bioindex,
    how='inner',
    left_on='docid',
    right_on='htid')

# date functions taken from https://github.com/tedunderwood/library/blob/master/SonicScrewdriver.py
def infer_date(datetype, firstdate, seconddate, textdate):
    '''Receives a date type and three dates, as strings, with no guarantee that any
    of the dates will be numeric. The logic of the data here is defined by
    MARC standards for controlfield 008:
    http://www.loc.gov/marc/bibliographic/concise/bd008a.html
    Returns a date that represents either a shaky consensus
    about the earliest attested date for this item, or 0, indicating no
    consensus.
    '''
    try:
        intdate = int(firstdate)
    except:
        # No readable date
        if firstdate.endswith('uu'):
            # Two missing places is too many.
            intdate = 0
        elif firstdate.endswith('u'):
            # but one is okay
            try:
                decade = int(firstdate[0:3])
                intdate = decade * 10
            except:
                # something's weird. fail.
                intdate = 0
        else:
            intdate = 0

    if intdate == 0:
        try:
            intdate = int(textdate)
        except:
            intdate = 0

    try:
        intsecond = int(seconddate)
    except:
        intsecond = 0

    if intsecond - intdate > 80 and intsecond < 2100:
        # A gap of more than eighty years is too much.
        # This is usually an estimated date that could be anywhere within
        # the nineteenth century.
        # note that we specify intsecond < 2100 because otherwise things
        # dated 9999 throw an error
        intdate = 0

    if datetype == 't' and intsecond > 0 and intsecond < intdate:
        intdate = intsecond
        # This is a case where we have both a publication date and
        # a copyright date. Accept the copyright date. We're going
        # for 'the earliest attested date for the item.'

    if intdate < 1000 and intsecond > 1700 and intsecond < 2100:
        intdate = intsecond

    return intdate

def date_row(row):
    datetype = row["datetype"]
    firstdate = row["startdate"]
    secondate = row["enddate"]
    if "imprintdate" in row:
        textdate = row["imprintdate"]
    else:
        textdate = row["textdate"]
    intdate = infer_date(datetype, firstdate, secondate, textdate)
    return intdate

def make_balanced(df):
    balanced_df = pd.DataFrame()
    for yr in np.arange(1923,2001):
        for gender in ['F','M','U']:
            subset_df = df.loc[(df.infer_date == yr) & (df.authgender == gender),:]
            df_shape = subset_df.shape
            if gender == 'U':
                subset_df = subset_df.head(25)
            if df_shape[0] > 50:
                subset_df = subset_df.head(50)
            else:
                print('Only', df_shape[0], 'rows written for:')
                print(yr, '-->',gender)
            balanced_df = pd.concat([balanced_df, subset_df], axis=0)
    return balanced_df

# add infer_date to the df
merge_dict = merge_bioindex.to_dict(orient='records')
for row in merge_dict:
    row['infer_date'] = date_row(row)

new_bioindex = pd.DataFrame()
new_bioindex = new_bioindex.from_records(merge_dict)

# create gender-balanced subset df
balanced_df = make_balanced(new_bioindex)

# write to given outfile
balanced_df.to_csv(outfile, index=False, sep='\t')

print(balanced_df.shape[0], 'rows written to', outfile)

max_stars_repo_name: tedunderwood/biographies

max_issues_repo_name: afcarl/biographies

max_forks_repo_name: afcarl/biographies


text:
#! /usr/bin/env python
from __future__ import absolute_import

from numpy.testing import assert_equal
from .StreamConverter_test import __create_trace
import numpy
from geomagio import TimeseriesUtility
from obspy.core import Stream, Stats, Trace, UTCDateTime

assert_almost_equal = numpy.testing.assert_almost_equal
assert_array_equal = numpy.testing.assert_array_equal

def test_create_empty_trace():
    """TimeseriesUtility_test.test_create_empty_trace()"""
    trace1 = _create_trace([1, 1, 1, 1, 1], "H", UTCDateTime("2018-01-01"))
    trace2 = _create_trace([2, 2], "E", UTCDateTime("2018-01-01"))
    observatory = "Test"
    interval = "minute"
    network = "NT"
    location = "R0"
    trace3 = TimeseriesUtility.create_empty_trace(
        starttime=trace1.stats.starttime,
        endtime=trace1.stats.endtime,
        observatory=observatory,
        channel="F",
        type="variation",
        interval=interval,
        network=network,
        station=trace1.stats.station,
        location=location,
    )
    timeseries = Stream(traces=[trace1, trace2])
    # For continuity set stats to be same for all traces
    for trace in timeseries:
        trace.stats.observatory = observatory
        trace.stats.type = "variation"
        trace.stats.interval = interval
        trace.stats.network = network
        trace.stats.station = trace1.stats.station
        trace.stats.location = location
    timeseries += trace3
    assert_equal(len(trace3.data), trace3.stats.npts)
    assert_equal(timeseries[0].stats.starttime, timeseries[2].stats.starttime)
    TimeseriesUtility.pad_timeseries(
        timeseries=timeseries,
        starttime=trace1.stats.starttime,
        endtime=trace1.stats.endtime,
    )
    assert_equal(len(trace3.data), trace3.stats.npts)
    assert_equal(timeseries[0].stats.starttime, timeseries[2].stats.starttime)
    # Change starttime by more than 1 delta
    starttime = trace1.stats.starttime
    endtime = trace1.stats.endtime
    TimeseriesUtility.pad_timeseries(timeseries, starttime - 90, endtime + 90)
    assert_equal(len(trace3.data), trace3.stats.npts)
    assert_equal(timeseries[0].stats.starttime, timeseries[2].stats.starttime)
    # test hourly/daily starttime shift
    hour_trace = TimeseriesUtility.create_empty_trace(
        starttime=trace1.stats.starttime,
        endtime=trace1.stats.endtime,
        observatory=observatory,
        channel="F",
        type="variation",
        interval="hour",
        network=network,
        station=trace1.stats.station,
        location=location,
    )

    assert_equal(hour_trace.stats.starttime, UTCDateTime("2018-01-01T00:29:30Z"))

    day_trace = TimeseriesUtility.create_empty_trace(
        starttime=trace1.stats.starttime,
        endtime=trace1.stats.endtime,
        observatory=observatory,
        channel="F",
        type="variation",
        interval="day",
        network=network,
        station=trace1.stats.station,
        location=location,
    )

    assert_equal(day_trace.stats.starttime, UTCDateTime("2018-01-01T11:59:30Z"))

    short_trace = TimeseriesUtility.create_empty_trace(
        starttime=trace1.stats.starttime,
        endtime=trace1.stats.starttime + 1,
        observatory=observatory,
        channel="F",
        type="variation",
        interval="day",
        network=network,
        station=trace1.stats.station,
        location=location,
    )

    assert_equal(short_trace.stats.starttime, UTCDateTime("2018-01-01T11:59:30Z"))
    assert_equal(short_trace.stats.endtime, short_trace.stats.starttime)

def test_get_stream_gaps():
    """TimeseriesUtility_test.test_get_stream_gaps()

    confirms that gaps are found in a stream
    """
    stream = Stream(
        [
            __create_trace("H", [numpy.nan, 1, 1, numpy.nan, numpy.nan]),
            __create_trace("Z", [0, 0, 0, 1, 1, 1]),
        ]
    )
    for trace in stream:
        # set time of first sample
        trace.stats.starttime = UTCDateTime("2015-01-01T00:00:00Z")
        # set sample rate to 1 second
        trace.stats.delta = 1
    # find gaps
    gaps = TimeseriesUtility.get_stream_gaps(stream)
    assert_equal(len(gaps["H"]), 2)
    # gap at start of H
    gap = gaps["H"][0]
    assert_equal(gap[0], UTCDateTime("2015-01-01T00:00:00Z"))
    assert_equal(gap[1], UTCDateTime("2015-01-01T00:00:00Z"))
    # gap at end of H
    gap = gaps["H"][1]
    assert_equal(gap[0], UTCDateTime("2015-01-01T00:00:03Z"))
    assert_equal(gap[1], UTCDateTime("2015-01-01T00:00:04Z"))
    # no gaps in Z channel
    assert_equal(len(gaps["Z"]), 0)

def test_get_stream_gaps_channels():
    """TimeseriesUtility_test.test_get_stream_gaps_channels()

    test that gaps are only checked in specified channels.
    """
    stream = Stream(
        [
            __create_trace("H", [numpy.nan, 1, 1, numpy.nan, numpy.nan]),
            __create_trace("Z", [0, 0, 0, 1, 1, 1]),
        ]
    )
    for trace in stream:
        # set time of first sample
        trace.stats.starttime = UTCDateTime("2015-01-01T00:00:00Z")
        # set sample rate to 1 second
        trace.stats.delta = 1
    # find gaps
    gaps = TimeseriesUtility.get_stream_gaps(stream, ["Z"])
    assert_equal("H" in gaps, False)
    assert_equal(len(gaps["Z"]), 0)

def test_get_trace_gaps():
    """TimeseriesUtility_test.test_get_trace_gaps()

    confirm that gaps are found in a trace
    """
    trace = __create_trace("H", [1, 1, numpy.nan, numpy.nan, 0, 1])
    # set time of first sample
    trace.stats.starttime = UTCDateTime("2015-01-01T00:00:00Z")
    # set sample rate to 1 minute
    trace.stats.delta = 60
    # find gap
    gaps = TimeseriesUtility.get_trace_gaps(trace)
    assert_equal(len(gaps), 1)
    gap = gaps[0]
    assert_equal(gap[0], UTCDateTime("2015-01-01T00:02:00Z"))
    assert_equal(gap[1], UTCDateTime("2015-01-01T00:03:00Z"))

def test_get_merged_gaps():
    """TimeseriesUtility_test.test_get_merged_gaps()

    confirm that gaps are merged
    """
    merged = TimeseriesUtility.get_merged_gaps(
        {
            "H": [
                # gap for 2 seconds, that starts after next gap
                [
                    UTCDateTime("2015-01-01T00:00:01Z"),
                    UTCDateTime("2015-01-01T00:00:03Z"),
                    UTCDateTime("2015-01-01T00:00:04Z"),
                ]
            ],
            # gap for 1 second, that occurs before previous gap
            "Z": [
                [
                    UTCDateTime("2015-01-01T00:00:00Z"),
                    UTCDateTime("2015-01-01T00:00:00Z"),
                    UTCDateTime("2015-01-01T00:00:01Z"),
                ],
                [
                    UTCDateTime("2015-01-01T00:00:05Z"),
                    UTCDateTime("2015-01-01T00:00:07Z"),
                    UTCDateTime("2015-01-01T00:00:08Z"),
                ],
            ],
        }
    )
    assert_equal(len(merged), 2)
    # first gap combines H and Z gaps
    gap = merged[0]
    assert_equal(gap[0], UTCDateTime("2015-01-01T00:00:00Z"))
    assert_equal(gap[1], UTCDateTime("2015-01-01T00:00:03Z"))
    # second gap is second Z gap
    gap = merged[1]
    assert_equal(gap[0], UTCDateTime("2015-01-01T00:00:05Z"))
    assert_equal(gap[1], UTCDateTime("2015-01-01T00:00:07Z"))

def test_get_trace_values():
    """TimeseriesUtility_test.test_get_trace_values()"""
    stream = Stream(
        [
            __create_trace("H", [numpy.nan, 1, 1, numpy.nan, numpy.nan]),
            __create_trace("Z", [0, 0, 0, 1, 1, 1]),
        ]
    )
    for trace in stream:
        # set time of first sample
        trace.stats.starttime = UTCDateTime("2015-01-01T00:00:00Z")
        # set sample rate to 1 second
        trace.stats.delta = 1
        trace.stats.npts = len(trace.data)
    print(stream)
    print(stream.select(channel="H")[0].times("utcdatetime"))
    # value that doesn't exist
    assert_equal(
        TimeseriesUtility.get_trace_value(
            traces=stream.select(channel="H"), time=UTCDateTime("2015-01-01T00:00:00Z")
        ),
        None,
    )
    # value that exists
    assert_equal(
        TimeseriesUtility.get_trace_value(
            traces=stream.select(channel="Z"), time=UTCDateTime("2015-01-01T00:00:00Z")
        ),
        0,
    )
    # default for value that doesn't exist
    assert_equal(
        TimeseriesUtility.get_trace_value(
            traces=stream.select(channel="H"),
            time=UTCDateTime("2015-01-01T00:00:03Z"),
            default=4,
        ),
        4,
    )

def test_has_all_channels():
    """TimeseriesUtility_test.test_has_all_channels():"""
    nan = numpy.nan
    stream = Stream(
        [
            __create_trace("H", [nan, 1, 1, nan, nan]),
            __create_trace("Z", [0, 0, 0, 1, 1]),
            __create_trace("E", [nan, nan, nan, nan, nan]),
        ]
    )
    for trace in stream:
        # set time of first sample
        trace.stats.starttime = UTCDateTime("2015-01-01T00:00:00Z")
        # set sample rate to 1 second
        trace.stats.delta = 1
        trace.stats.npts = len(trace.data)
    # check for channels
    starttime = stream[0].stats.starttime
    endtime = stream[0].stats.endtime
    assert_equal(
        TimeseriesUtility.has_all_channels(stream, ["H", "Z"], starttime, endtime), True
    )
    assert_equal(
        TimeseriesUtility.has_all_channels(stream, ["H", "Z", "E"], starttime, endtime),
        False,
    )
    assert_equal(
        TimeseriesUtility.has_all_channels(stream, ["E"], starttime, endtime), False
    )

def test_has_any_channels():
    """TimeseriesUtility_test.test_has_any_channels():"""
    nan = numpy.nan
    stream = Stream(
        [
            __create_trace("H", [nan, 1, 1, nan, nan]),
            __create_trace("Z", [0, 0, 0, 1, 1, 1]),
            __create_trace("E", [nan, nan, nan, nan, nan]),
        ]
    )
    for trace in stream:
        # set time of first sample
        trace.stats.starttime = UTCDateTime("2015-01-01T00:00:00Z")
        # set sample rate to 1 second
        trace.stats.delta = 1
        trace.stats.npts = len(trace.data)
    # check for channels
    starttime = stream[0].stats.starttime
    endtime = stream[0].stats.endtime
    assert_equal(
        TimeseriesUtility.has_any_channels(stream, ["H", "Z"], starttime, endtime), True
    )
    assert_equal(
        TimeseriesUtility.has_any_channels(stream, ["H", "Z", "E"], starttime, endtime),
        True,
    )
    assert_equal(
        TimeseriesUtility.has_any_channels(stream, ["E"], starttime, endtime), False
    )

def test_merge_streams():
    """TimeseriesUtility_test.test_merge_streams()

    confirm merge streams treats empty channels correctly
    """
    trace1 = __create_trace("H", [1, 1, 1, 1])
    trace2 = __create_trace("E", [2, numpy.nan, numpy.nan, 2])
    trace3 = __create_trace("F", [numpy.nan, numpy.nan, numpy.nan, numpy.nan])
    trace4 = __create_trace("H", [2, 2, 2, 2])
    trace5 = __create_trace("E", [3, numpy.nan, numpy.nan, 3])
    trace6 = __create_trace("F", [numpy.nan, numpy.nan, numpy.nan, numpy.nan])
    npts1 = len(trace1.data)
    npts2 = len(trace4.data)
    timeseries1 = Stream(traces=[trace1, trace2, trace3])
    timeseries2 = Stream(traces=[trace4, trace5, trace6])
    for trace in timeseries1:
        trace.stats.starttime = UTCDateTime("2018-01-01T00:00:00Z")
        trace.stats.npts = npts1
    for trace in timeseries2:
        trace.stats.starttime = UTCDateTime("2018-01-01T00:02:00Z")
        trace.stats.npts = npts2
    merged_streams1 = TimeseriesUtility.merge_streams(timeseries1)
    # Make sure the empty 'F' was not removed from stream
    assert_equal(1, len(merged_streams1.select(channel="F")))
    # Merge multiple streams with overlapping timestamps
    timeseries = timeseries1 + timeseries2

    merged_streams = TimeseriesUtility.merge_streams(timeseries)
    assert_equal(len(merged_streams), len(timeseries1))
    assert_equal(len(merged_streams[0]), 6)
    assert_equal(len(merged_streams[2]), 6)
    assert_almost_equal(merged_streams.select(channel="H")[0].data, [1, 1, 2, 2, 2, 2])
    assert_almost_equal(
        merged_streams.select(channel="E")[0].data, [2, numpy.nan, 3, 2, numpy.nan, 3]
    )
    assert_almost_equal(merged_streams.select(channel="F")[0].data, [numpy.nan] * 6)

    trace7 = __create_trace("H", [1, 1, 1, 1])
    trace8 = __create_trace("E", [numpy.nan, numpy.nan, numpy.nan, numpy.nan])
    trace9 = __create_trace("F", [numpy.nan, numpy.nan, numpy.nan, numpy.nan])
    timeseries3 = Stream(traces=[trace7, trace8, trace9])
    npts3 = len(trace7.data)
    for trace in timeseries3:
        trace.stats.starttime = UTCDateTime("2018-01-01T00:00:00Z")
        trace.stats.npts = npts3
    merged_streams3 = TimeseriesUtility.merge_streams(timeseries3)
    assert_equal(len(timeseries3), len(merged_streams3))
    assert_almost_equal(timeseries3.select(channel="H")[0].data, [1, 1, 1, 1])
    assert_equal(numpy.isnan(timeseries3.select(channel="E")[0].data).all(), True)
    assert_equal(numpy.isnan(timeseries3.select(channel="F")[0].data).all(), True)

    trace10 = __create_trace("H", [1, 1, numpy.nan, numpy.nan, 1, 1])
    trace11 = __create_trace("H", [2, 2, 2, 2])
    trace10.stats.starttime = UTCDateTime("2018-01-01T00:00:00Z")
    trace11.stats.starttime = UTCDateTime("2018-01-01T00:01:00Z")
    timeseries4 = Stream(traces=[trace10, trace11])
    merged4 = TimeseriesUtility.merge_streams(timeseries4)
    assert_equal(len(merged4[0].data), 6)
    assert_almost_equal(merged4.select(channel="H")[0].data, [1, 2, 2, 2, 1, 1])

def test_pad_timeseries():
    """TimeseriesUtility_test.test_pad_timeseries()"""
    trace1 = _create_trace([1, 1, 1, 1, 1], "H", UTCDateTime("2018-01-01"))
    trace2 = _create_trace([2, 2], "E", UTCDateTime("2018-01-01"))
    timeseries = Stream(traces=[trace1, trace2])
    TimeseriesUtility.pad_timeseries(
        timeseries=timeseries,
        starttime=trace1.stats.starttime,
        endtime=trace1.stats.endtime,
    )
    assert_equal(len(trace1.data), len(trace2.data))
    assert_equal(trace1.stats.starttime, trace2.stats.starttime)
    assert_equal(trace1.stats.endtime, trace2.stats.endtime)
    # change starttime by less than 1 delta
    starttime = trace1.stats.starttime
    endtime = trace1.stats.endtime
    TimeseriesUtility.pad_timeseries(timeseries, starttime - 30, endtime + 30)
    assert_equal(trace1.stats.starttime, starttime)
    # Change starttime by more than 1 delta
    TimeseriesUtility.pad_timeseries(timeseries, starttime - 90, endtime + 90)
    assert_equal(trace1.stats.starttime, starttime - 60)
    assert_equal(numpy.isnan(trace1.data[0]), numpy.isnan(numpy.NaN))

def test_pad_and_trim_trace():
    """TimeseriesUtility_test.test_pad_and_trim_trace()"""
    trace = _create_trace([0, 1, 2, 3, 4], "X", UTCDateTime("2018-01-01"))
    assert_equal(trace.stats.starttime, UTCDateTime("2018-01-01T00:00:00Z"))
    assert_equal(trace.stats.endtime, UTCDateTime("2018-01-01T00:04:00Z"))
    # starttime between first and second sample
    # expect first sample to be removed, start at next sample, end at same
    TimeseriesUtility.pad_and_trim_trace(
        trace,
        starttime=UTCDateTime("2018-01-01T00:00:30Z"),
        endtime=trace.stats.endtime,
    )
    assert_equal(trace.stats.starttime, UTCDateTime("2018-01-01T00:01:00Z"))
    assert_equal(trace.stats.endtime, UTCDateTime("2018-01-01T00:04:00Z"))
    assert_array_equal(trace.data, [1, 2, 3, 4])
    # endtime between last and second to last samples
    TimeseriesUtility.pad_and_trim_trace(
        trace,
        starttime=UTCDateTime("2018-01-01T00:00:30Z"),
        endtime=UTCDateTime("2018-01-01T00:03:50Z"),
    )
    assert_equal(trace.stats.starttime, UTCDateTime("2018-01-01T00:01:00Z"))
    assert_equal(trace.stats.endtime, UTCDateTime("2018-01-01T00:03:00Z"))
    assert_array_equal(trace.data, [1, 2, 3])
    # pad outward
    TimeseriesUtility.pad_and_trim_trace(
        trace,
        starttime=UTCDateTime("2018-01-01T00:00:00Z"),
        endtime=UTCDateTime("2018-01-01T00:05:00Z"),
    )
    assert_equal(trace.stats.starttime, UTCDateTime("2018-01-01T00:00:00Z"))
    assert_equal(trace.stats.endtime, UTCDateTime("2018-01-01T00:05:00Z"))
    assert_array_equal(trace.data, [numpy.nan, 1, 2, 3, numpy.nan, numpy.nan])
    # remove exactly one sample
    TimeseriesUtility.pad_and_trim_trace(
        trace,
        starttime=UTCDateTime("2018-01-01T00:00:00Z"),
        endtime=UTCDateTime("2018-01-01T00:04:00Z"),
    )
    assert_equal(trace.stats.starttime, UTCDateTime("2018-01-01T00:00:00Z"))
    assert_equal(trace.stats.endtime, UTCDateTime("2018-01-01T00:04:00Z"))
    assert_array_equal(trace.data, [numpy.nan, 1, 2, 3, numpy.nan])
    # pad start and trim end
    TimeseriesUtility.pad_and_trim_trace(
        trace,
        starttime=UTCDateTime("2017-12-31T23:58:59Z"),
        endtime=UTCDateTime("2018-01-01T00:03:00Z"),
    )
    assert_equal(trace.stats.starttime, UTCDateTime("2017-12-31T23:59:00Z"))
    assert_equal(trace.stats.endtime, UTCDateTime("2018-01-01T00:03:00Z"))
    assert_array_equal(trace.data, [numpy.nan, numpy.nan, 1, 2, 3])
    # pad end and trim start
    TimeseriesUtility.pad_and_trim_trace(
        trace,
        starttime=UTCDateTime("2018-01-01T00:00:00Z"),
        endtime=UTCDateTime("2018-01-01T00:04:00Z"),
    )
    assert_equal(trace.stats.starttime, UTCDateTime("2018-01-01T00:00:00Z"))
    assert_equal(trace.stats.endtime, UTCDateTime("2018-01-01T00:04:00Z"))
    assert_array_equal(trace.data, [numpy.nan, 1, 2, 3, numpy.nan])

def test_pad_and_trim_trace_fixfloat():
    """TimeseriesUtility_test.test_pad_and_trim_trace_fixfloat()
    This tests whether pad_and_trim_trace() handles a known
    floating point precision trouble-maker encountered with
    UTCDateTimes correctly.
    """
    trace = _create_trace(
        numpy.linspace(1, 56, 56), "X", UTCDateTime("2020-05-28T15:53:50.7"), delta=0.1
    )
    assert_equal(trace.stats.starttime, UTCDateTime("2020-05-28T15:53:50.7"))
    assert_equal(trace.stats.endtime, UTCDateTime("2020-05-28T15:53:56.2"))
    # This should create a 70 sample trace of 10 Hz data, inclusive of the
    # known startime and endtime. Early versions of pad_and_trim_trace() did
    # not handle this scenario correctly, returning 68 samples, and missing
    # both the 1st and last of the expected samples.
    TimeseriesUtility.pad_and_trim_trace(
        trace,
        starttime=UTCDateTime("2020-05-28T15:53:50.0"),
        endtime=UTCDateTime("2020-05-28T15:53:56.9"),
    )
    assert_equal(trace.stats.starttime, UTCDateTime("2020-05-28T15:53:50.0"))
    assert_equal(trace.stats.endtime, UTCDateTime("2020-05-28T15:53:56.9"))
    assert_array_equal(
        trace.data,
        numpy.concatenate(
            ([numpy.nan] * 7, numpy.linspace(1, 56, 56), [numpy.nan] * 7)
        ),
    )

def test_round_usecs():
    """TimeseriesUtility_test.test_round_usecs()
    This tests whether microsecond values are rounded or
    not depending on residual microsecond values
    """
    # test case with no residual microseconds
    time = TimeseriesUtility.round_usecs(UTCDateTime("2020-10-07T00:00:00Z"))
    assert_equal(time, UTCDateTime("2020-10-07T00:00:00Z"))
    # test case with residual microseconds
    time = TimeseriesUtility.round_usecs(UTCDateTime("2020-10-07T00:00:00.995600Z"))
    assert_equal(time, UTCDateTime("2020-10-07T00:00:00.996000Z"))
    # test case with rounding to next second
    time = TimeseriesUtility.round_usecs(UTCDateTime("2020-10-07T00:00:00.9995Z"))
    assert_equal(time, UTCDateTime("2020-10-07T00:00:01.000Z"))

def _create_trace(data, channel, starttime, delta=60.0):
    stats = Stats()
    stats.channel = channel
    stats.delta = delta
    stats.starttime = starttime
    stats.npts = len(data)
    data = numpy.array(data, dtype=numpy.float64)
    return Trace(data, stats)

max_stars_repo_name: usgs/geomag-algorithms

max_issues_repo_name: usgs/geomag-algorithms

max_forks_repo_name: alejandrodelcampillo/geomag-algorithms

reason: import numpy,from numpy


text:
#!/usr/bin/python
# coding: utf8

"""
Simplified version of stat_t_test (without multiprocessing)
used to compute Hotelling statistic

threshold for independant testing:
  rv = scipy.stats.t(df=(n1+n2-2))
  th = rv.ppf(0.95)
"""

import numpy as np

from .statistics import *

################################################################################

def zmap_1perm_2samp(X, cat1, cat2=None, rand_seed=-1, fstat=None, name=None):
    """ une permutation
    X (D, N, P) K points, N subjects, D dim
    return:
    Y (D,) zvalue at each point
    """
    if fstat is None:
        fstat = hotelling_2samples
        #name = "MP-Hotelling"

    if cat2 is None:
        cat2 = np.logical_not(cat1)

    # Données
    if rand_seed < 0:
        # Sans permutation (on peut remplacer cat par idx[cat])
        ix1 = cat1
        ix2 = cat2
    else:
        # Avec permutation
        np.random.seed(rand_seed)
        idx = np.arange(X.shape[1])[cat1 | cat2]
        per = np.random.permutation(idx.size)
        nsplit = cat1.sum()
        ix1 = idx[per][:nsplit]
        ix2 = idx[per][nsplit:]

    # Run
    Y = fstat(X[:, ix1, :], X[:, ix2, :])

    if name is not None:
        print(name + " {0}, {1}\n".format(Y.min(), Y.max()))
    return Y

def zval_kperm_2samp(X, cat1, cat2=None, nperm=100, fstat=None):
    """
    simple loop, no optimization
    return the extrema sorted
    return:
    lmin, lmax
    """
    if cat2 is None:
        cat2 = np.logical_not(cat1)

    lmin = np.zeros((nperm,))
    lmax = np.zeros((nperm,))
    for i in range(nperm):
        Y0 = zmap_1perm_2samp(X, cat1, cat2, rand_seed=i, fstat=fstat)
        lmax[i] = Y0.max()
        lmin[i] = Y0.min()
    lmax.sort()
    lmin.sort()
    return lmin, lmax

################################################################################

def zmap_1perm_2pairedsamp(X, rand_seed=-1, name=None):
    """ une permutation (changement de signe)
    X (D, N, P) D points, N subjects, P dim
    return:
    Y (D,) zvalue at each point
    """
    (d, n, p) = X.shape

    # Données
    if rand_seed < 0:
        # Sans permutation
        sign_swap = np.ones((1, n, 1))
    else:
        # Avec permutation
        sign_swap = np.random.randint(0, 2, size=(1, n, 1))
        sign_swap = 2*sign_swap - 1

    # Run
    Y = hotelling_1sample(X * sign_swap)

    if name is not None:
        print(name + " {0}, {1}\n".format(Y.min(), Y.max()))
    return Y

def zval_kperm_2pairedsamp(X, nperm):
    """
    simple loop, no optimization
    return the extrema sorted
    return:
    lmin, lmax
    """
    lmin = np.zeros((nperm,))
    lmax = np.zeros((nperm,))
    for i in range(nperm):
        Y0 = zmap_1perm_2pairedsamp(X, rand_seed=i)
        lmax[i] = Y0.max()
        lmin[i] = Y0.min()
    lmax.sort()
    lmin.sort()
    return lmin, lmax

################################################################################

def zmap_kperm_llh(X, y, nperm, nvar0=1):
    """
    WIP
    log-likelihood testing for non parametric permutation regression testing
    nvar0 is used to set the reference model

    X (n, d, p) deformation
    y (n, q)    clinicals

    return
    L (d, nperm+1)
    L[:, 0],                   no-permutation map
    np.max(L[:, 1:], axis=0)   permutation maxs
    """
    n, d, p = X.shape
    K = np.zeros((d, 1))
    L = np.zeros((d, nperm+1))

    # reference model and no permutation
    K[:, 0] = regression_loglikelihood(X, y[:, :nvar0])
    L[:, 0] = regression_loglikelihood(X, y)

    # with perm
    for k in range(nperm):
        #np.random.seed(k)
        per = np.random.permutation(n)
        yp = y.copy()
        yp[:, nvar0:] = y[per, nvar0:]
        L[:, k+1] = regression_loglikelihood(X, yp)

    L = L - K
    return L, L[:, 0], np.max(L[:, 1:], axis=0)

################################################################################
def compute_pvalues(zmap, zsamp, alpha_threshold=1., do_sort=True):
    """
    Compute p-values on the right only using:
    - zmap   (vdim,)   z-values
    - zsamp  (nsamp,)  empirical z samples

    return smallest k such that !(zmap[i] > zsamp[k])

    can perform asymetric search on right and left

    (todo voir np.searchsorted)
    """

    vdim = zmap.size
    nsamp = zsamp.size
    if nsamp == 1:
        zsamp.shape = (1,)

    if do_sort:
        zsamp.sort()

    k_threshold = (1. - alpha_threshold)*nsamp

    pval = np.zeros(vdim, dtype="uint16")
    for i in range(vdim):
        k = 0
        while k < nsamp and zmap[i] > zsamp[k]:
            k += 1
        if k >= k_threshold:
            pval[i] = k
    return pval

max_stars_repo_name: ClinicalCardiovascEngGroup/SSM

max_issues_repo_name: ClinicalCardiovascEngGroup/SSM

max_forks_repo_name: ClinicalCardiovascEngGroup/SSM


text:
import numpy as np
import loupe

def test_optimize_scalar_quadratic():
    x = loupe.rand()
    y = 4*(x-5)**2
    res = loupe.optimize(y, x)
    assert np.allclose(res.x, 5)

def test_optimize_scalar_quadratic_finite_difference():
    x = loupe.rand()
    y = 4*(x-5)**2
    res = loupe.optimize(y, x, analytic_grad=False)
    assert np.allclose(res.x, 5)

max_stars_repo_name: andykee/loupe

max_issues_repo_name: andykee/loupe

max_forks_repo_name: andykee/loupe