text:
from __future__ import print_function
import sys
import datetime
import logging
import numpy as np
from StringIO import StringIO

DEF_VERSION = '1.2.2.7-dev'     # x.y.z.* to match RASDRviewer version
DEF_DELIM   = ','
DEF_AVERAGE = 1
DEF_CALIB   = 0.0           # Paul uses (0.73278^2)/2000 as Qstep^2/ADC impedance
DEF_DT_VAR  = 0.1           # percentage of avg frame time to trap 'SKIPPED' frames

# support for statistics gathering/output
DEF_STATS_FORMAT=['%.6f','%.3f','%.4f','%.4f','%.4f','%.6f','%.6f',]
DEF_STATS_DELIM=','
DEF_STATS_HEADER='timestamp,deltaT (sec),minimum,maximum,average,standard-deviation,variance'

# http://stackoverflow.com/questions/16867347/step-by-step-debugging-with-ipython
# https://github.com/ipython/ipython/wiki/Cookbook%3a-Updating-code-for-use-with-IPython-0.11-and-later
def ipsh():
    from IPython.frontend.terminal.embed import InteractiveShellEmbed
    from IPython.config.loader import Config
    import inspect
    ipshell = InteractiveShellEmbed(config=Config(),
        banner1='*** STOP\nCTRL-D to exit', exit_msg='Resume...')
    frame = inspect.currentframe().f_back
    msg   = 'Stopped at {0.f_code.co_filename} at line {0.f_lineno}'.format(frame)
    # Go back one level!
    # This is needed because the call to ipshell is inside the function ipsh()
    ipshell(msg,stack_depth=2)

def excel2dt(et):
    # http://stackoverflow.com/questions/1703505/excel-date-to-unix-timestamp
    # http://stackoverflow.com/questions/3682748/converting-unix-timestamp-string-to-readable-date-in-python
    ts = (et - 25569)*86400 # 25569 is days since 1900-01-01
    ts = ts + (4*60*60)     # timezone EST5EDT return
    return datetime.datetime.fromtimestamp(float(ts))

def dt2excel(dt,tz=None):
    import pytz
    # http://stackoverflow.com/questions/9574793/how-to-convert-a-python-datetime-datetime-to-excel-serial-date-number
    # http://stackoverflow.com/questions/4530069/python-how-to-get-a-value-of-datetime-today-that-is-timezone-aware
    dte = datetime.datetime(1899, 12, 30, tzinfo=pytz.timezone(tz)) if tz else datetime.datetime(1899, 12, 29)
    delta = dt - dte    # set w/r/t excel2dt() date; NB: do not know why this must be set 3 days earlier
    return float(delta.days) + (float(delta.seconds) / 86400)

def translate_tz(tz):
    # Paul uses some wierd strings in RASDRviewer to represent timezones
    # FIXME: I am still unhappy with this...  But I need this to make the plots we took in June-July 2014
    tzo=''
    if tz.startswith('UT20'):
        tzo='-0500'
        tz=tz.replace('UT20','UTC')
    elif tz.startswith('US/Eastern'):
        tzo='-0500'
        tz=tz.replace('US/Eastern','UTC')
    elif not tz.startswith('UTC'):  # FIXME: weak...
        tz=tz.replace('UT','UTC')
    return tz, tzo

def open_spectrum_file(filename,opts):
    obj = {}
    if opts.localtime:
        # this is Paul's "Localtime" .csv format
        # http://stackoverflow.com/questions/13311471/skip-a-specified-number-of-columns-with-numpy-genfromtxt
        # http://stackoverflow.com/questions/466345/converting-string-into-datetime
        # http://stackoverflow.com/questions/698223/how-can-i-parse-a-time-string-containing-milliseconds-in-it-with-python

        f=open(filename,'r')
        key=np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter,dtype='str')
        if len(np.atleast_1d(key)) < 2:
            raise Exception('unable to parse %s -- check first line'%filename)
        ds=key[4].replace('"','') if len(key) > 4 else '00/00/00'
        tz=key[6].replace('"','') if len(key) > 7 else 'UTC'
        # translate strange labels from RASDRviewer_W_1_0_5
        tz=translate_tz(tz)
        dumplines = 1
        fscale = 1.0
        ## RASDRviewer 1.2.2, Paul "adds" a line with a ',' in it - just kill it
        ## RASDRviewer 1.2.2.1, we negotiate that this line will be '*** BEGIN DATA COLLECTION ***'
        if opts.format.startswith('1.2.2'):
            dumplines = dumplines + 1
            tag = f.readline()
        ## RASDRViewer 1.1.1, Paul "adds" a specifier to the 2nd line and moves the rest of the lines down
        if opts.format.startswith('1.1.') or opts.format.startswith('1.2.'):
            dumplines = dumplines + 1
            key=np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter,dtype='str')
            if len(np.atleast_1d(key)) < 2:
                tag = str(key).lower()
            else:
                tag = str(key[0]).lower()
            if tag.find('hz')>0:
                if tag.startswith('khz'):
                    fscale = 0.001
                elif tag.startswith('ghz'):
                    fscale = 1000.0
        ## second/third/fourth line: frequency bin information
        dumplines = dumplines + 1
        freq=np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter,dtype='float')
        if len(np.atleast_1d(freq)) < 2:
            raise Exception('unable to parse %s -- check frequency plan'%filename)
        freq=freq[1:]           # remove the 1st column, as it is a 'nan' when interpreted as a 'float'
        ## third/fourth/fifth line-to-end: extract first *column* as text
        t=np.genfromtxt(f,delimiter=opts.delimiter,usecols=[0],dtype='str')
        if len(np.atleast_1d(t)) < 2:                                                      # BUG: I need at least two data records
            raise Exception('unable to parse %s -- check data after frequency plan'%filename)
        time = []
        for i in range(1,len(np.atleast_1d(t))):
            # translate HH:MM:SS:ms into HH:MM:SS.ms
            if t[i].count(':') > 2:
                a,x,b = t[i].rpartition(':')
            else:
                a = t[i]
                b = '0'
            ts=a+'.'+b
            time.append(datetime.datetime.strptime(ds+'T'+ts, '%m/%d/%yT%H:%M:%S.%f'))
        f.seek(0)
        for i in range(dumplines):  # dump lines so we are at beginning of data collected
            f.readline()
        # update the object to return
        obj['file'] = f
        obj['freq'] = freq
        obj['time'] = time
        obj['opts'] = opts
##    elif opts.datetime:
##        # this is the "Excel-datetime" .csv format (proposed for RASDRviewer_W_1_0_4, but not implemented)
##        f=open(filename,'r')
##        freq = np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter,dtype='float')
##        t=np.genfromtxt(f,delimiter=opts.delimiter,usecols=[0],dtype='float')
##        time = []
##        for i in range(len(t)):
##            time.append(excel2dt(t[i]))
##        f.seek(0)               # start over
##        f.readline()            # dump the first line
##        f.readline()            # dump the second line
##        # update the object to return
##        obj['file'] = f
##        obj['freq'] = freq
##        obj['time'] = time
##        obj['opts'] = opts
    else:
        # this is the "Extended" .csv format produced by RASDRviewer_W_1_2_1
        from dateutil import parser

        f=open(filename,'r')
        ## first line: key-value pairs
        buf = f.readline()
        key=np.genfromtxt(StringIO(buf.replace('\\','/')),delimiter=opts.delimiter,dtype='str',comments='\\')
        if len(np.atleast_1d(key)) < 2:
            raise Exception('unable to parse %s -- check first line'%filename)
        ds=key[4].replace('"','') if len(key) > 4 else '00/00/00'
        tz=key[6].replace('"','') if len(key) > 7 else 'UTC'
        # translate strange labels from RASDRviewer_W_1_0_5
        tz,tzo=translate_tz(tz)
        dumplines = 1
        fscale = 1.0
        ## RASDRviewer 1.2.2, Paul "adds" a line with a ',' in it - just kill it
        ## RASDRviewer 1.2.2.1, we negotiate that this line will be '*** BEGIN DATA COLLECTION ***'
        if opts.format.startswith('1.2.2'):
            dumplines = dumplines + 1
            f.readline()
        ## RASDRViewer 1.1.1, Paul "adds" a specifier to the 2nd line and moves the rest of the lines down
        if opts.format.startswith('1.1.') or opts.format.startswith('1.2.'):
            dumplines = dumplines + 1
            key=np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter,dtype='str')
            if len(np.atleast_1d(key)) < 2:
                tag = str(key).lower()
            else:
                tag = str(key[0]).lower()
            if tag.find('hz')>0:
                if tag.startswith('khz'):
                    fscale = 0.001
                elif tag.startswith('ghz'):
                    fscale = 1000.0
        ## second/third/fourth line: frequency bin information
        dumplines = dumplines + 1
        freq=np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter,dtype='float')
        if len(np.atleast_1d(freq)) < 2:
            raise Exception('unable to parse %s -- check frequency plan'%filename)
        freq=freq[1:]           # remove the 1st column, as it is a 'nan' when interpreted as a 'float'
        ## third/fourth/fifth line-to-end: extract first *column* as text
        t=np.genfromtxt(f,delimiter=opts.delimiter,usecols=[0],dtype='str')                # 1st column are *almost* ISO-8601 datetime strings
        if len(np.atleast_1d(t)) < 2:                                                      # BUG: I need at least two data records
            raise Exception('unable to parse %s -- check data after frequency plan'%filename)
        time = []
        for i in range(1,len(np.atleast_1d(t))):
            try:
                time.append(parser.parse(t[i]))
            except ValueError:
                a,x,b = t[i].rpartition(':')  # deal with Paul's 'YYYY-MM-DDTHH:MM:SS:sssZ' format in RASDRviewer_W_1_0_5
                if a.find('T')>0:
                    time.append(parser.parse(a+'.'+b+tzo))
                else:
                    time.append(parser.parse(ds+'T'+a+'.'+b+tzo))
        f.seek(0)
        for i in range(dumplines):  # dump lines so we are at beginning of data collected
            f.readline()
        # update the object to return
        obj['file'] = f
        obj['freq'] = freq*fscale
        obj['time'] = time
        obj['opts'] = opts

    if opts.verbose:
        log = logging.getLogger(__name__)
        x,y,name = f.name.replace('\\','/').rpartition('/')
        log.info('%s.date=%s (%s)'%(name,ds+' '+str(tz),str(obj['time'][0])))
        log.info('%s.time=%d'%(name,len(obj['time'])))
        log.info('%s.freq=%s (scale=%f)'%(name,str(obj['freq'].shape),fscale))
    return obj

def read_spectrum_line(obj):
    opts = obj['opts']
    f    = obj['file']
##    if opts.datetime:
##        data = np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter)
##        data = data[1:]
##    else:
    freq = obj['freq']
    data = np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter,usecols=range(1,freq.shape[0]+1),dtype='float')
    return data         # spectral information in colums 1..N-2 (last column is junk)

def read_spectrum_array(obj):
    opts = obj['opts']
    f    = obj['file']
##    if opts.datetime:
##        data = np.genfromtxt(f,delimiter=opts.delimiter)
##        data = data[1:,1:]
##    else:
    freq = obj['freq']
    try:
        plumb = f.tell()
        data = np.genfromtxt(f,delimiter=opts.delimiter,usecols=range(1,freq.shape[0]+1),dtype='float')
    except ValueError, e:
        # oops, malformed file may contain 'garbage' lines at the end...
        log = logging.getLogger(__name__)
        log.warning('Malformed file that may contain extra lines at the end.  Recovering...')
        f.seek(plumb)
##        if opts.debug:
##            ipsh()
        data = np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter,usecols=range(1,freq.shape[0]+1),dtype='float')
        for i in range(len(obj['time'])-1):
            data = np.vstack((data,np.genfromtxt(StringIO(f.readline()),delimiter=opts.delimiter,usecols=range(1,freq.shape[0]+1),dtype='float')))
    return data         # spectral information in colums 1..N-2 (last column is junk)

def generate_spectrum_plots(filename,opts):
    log = logging.getLogger(__name__)
    if len(opts.background) > 0:
        if opts.background.lower() == 'automatic':
            bg = open_spectrum_file(filename,opts)
        else:
            bg = open_spectrum_file(opts.background,opts)
        lastrow = len(bg['time'])-1
        bkg = np.zeros(len(bg['freq']))
        acc = np.zeros(len(bg['freq']))
        n   = 0
        log.info('loading %s...'%bg['file'].name)
        if not opts.line:
            s = read_spectrum_array(bg)
        for fr in range(len(bg['time'])):
            n = n+1
            if opts.line:
                s = read_spectrum_line(bg)
                log.debug('bg %d/%d mean=%f dB'%(fr,len(bg['time']),s.mean()))
                if opts.atype.lower().startswith('linear'):
                    acc += np.power(10,s/10)
                else:
                    acc += s
            else:
                if opts.atype.lower().startswith('linear'):
                    acc += np.power(10,s[fr]/10)
                else:
                    acc += s[fr]
        bkg += acc / float(n)
        log.info('background %d samples=[%f,%f] mean=%f',n,bkg.min(),bkg.max(),bkg.mean())
        if opts.background.lower() == 'automatic':
            # distribute mean value of to all frequency bins to construct a total average over all samples
            # this is needed if doing a 'self' background correction
            bkg = np.ones(len(bg['freq'])) * bkg.mean()
        if not opts.line:
            del s
        fg = open_spectrum_file(filename,opts)
        if not len(fg['freq'])==len(bg['freq']):
            raise Exception('background file must have same number of frequency bins as the foreground file')
        if not np.array_equal(fg['freq'],bg['freq']):
            log.warn('background file does not specify the same frequency bins as the foreground file')
        del bg
    else:
        fg = open_spectrum_file(filename,opts)
        bkg = np.zeros(len(fg['freq']))

    ts_a = fg['time']
    fMHz = fg['freq']+opts.fc
    zidx = np.searchsorted(fMHz,0)
    nbin = len(fMHz)
    lastrow = len(ts_a)-1
    bw   = (fMHz.max() - fMHz.min())*1e6/nbin
    if opts.calibration > 0.0:
        dbm = 10.0 * np.log10( opts.calibration * bw * 0.001 )  # power referenced to 1mW
        if opts.atype.lower().startswith('linear'):
            dbm = np.power(10,dbm/10)
    else:
        dbm = 0.0

    log.info('range (MHz)=[%f,%f]',fMHz.min(),fMHz.max())
    log.info('zero index=%d',zidx)
    log.info('frequency bins=%d',nbin)
    log.info('bandwidth/bin (Hz)=%f',bw)
    if lastrow > 0:
        delta = ts_a[lastrow] - ts_a[lastrow-1]
        log.info('start=%s',str(ts_a[0]))
        log.info('end  =%s',str(ts_a[lastrow]))
        s = (float(delta.seconds)+(float(delta.microseconds)/1e6))
        if s > 0.0:
            if s < 10.0:
                log.info('inter-frame period=%.3f sec (%.0f fr/s)',s,1.0/float(s))
            else:
                log.info('inter-frame period=%.3f sec (%.3f fr/s)',s,1.0/float(s))

    log.info('averaging=%s',str(opts.average))

    if not opts.line:
        log.info('loading %s...',fg['file'].name)
        s_a = read_spectrum_array(fg)
    acc = np.zeros(nbin)
    hold = np.zeros(nbin)
    n   = 0
    for fr in range(len(ts_a)):
        dt = ts_a[fr]
        if opts.line:
            s = read_spectrum_line(fg)
        else:
            s = s_a[fr]
        if opts.atype.lower().startswith('linear'):
            s = np.power(10,s/10)
        s = s - bkg
        if opts.canceldc:
            log.debug('Cancel DC: frq %s',str(fMHz[zidx-4:zidx+5]))
            log.debug('Cancel DC: b4 %s',str(s[zidx-4:zidx+5]))
            a = (s[zidx-1] + s[zidx+1])/2.0
            if a < s[zidx]:
                s[zidx] = a
            log.debug('Cancel DC: ftr %s',str(s[zidx-4:zidx+5]))
        log.debug('line %d/%d max@%d val=%f',fr,len(ts_a),s.argmax(),s.max())
        if opts.hold:
            hold = np.maximum(s,hold)
        if n == 0:
            tstart = dt.strftime('%Y-%m-%dT%H:%M:%S%z')
        n    = n+1
        acc += s
        if n >= opts.average or fr == lastrow:
            s   = acc / float(n)    # vector/scalar
            s   = s - dbm           # vector-scalar

            ###
            ### Scaling and determination of min/max range
            ###
            offsets_applied = ''
            if opts.ptype.lower().startswith('linear'):
                if opts.atype.lower().startswith('log'):
                    s = np.power(10,s/10)
                std = np.std(s)
                if std < s.mean()/100.0:
                    log.warning('-> signal stdev<1%% mean; applying offset of %g to scale linear plot on screen',s.mean())
                    s = s - s.mean()
                    offsets_applied += ',%+g signal'%-s.mean()
                min = s.mean() - std*4
                max = s.mean() + std*4
                log.debug('-> plot mean=%f, stdev=%g, min=%f max=%f',s.mean(),std,min,max)
                if opts.hold:
                    std = np.std(hold)
                    if std < hold.mean()/100.0:
                        log.warning('-> hold stdev<1%% mean; applying offset of %g to scale linear plot on screen',hold.mean())
                        hold = hold - hold.mean()
                        offsets_applied += ',-%g hold'%hold.mean()
                    hmin = hold.mean() - std*4
                    hmax = hold.mean() + std*4
                    log.debug('-> hold mean=%f, stdev=%g, min=%f max=%f',hold.mean(),std,hmin,hmax)
                    min = hmin if hmin < min else min
                    max = hmax if hmax > max else max
            else:
                if opts.atype.lower().startswith('linear'):
                    delta = 0.0
                    if s.min() <= 0.0:
                        delta = 0.0 - s.min()
                        log.warning('-> applying fixed offset of %g to prevent logarithm underflow',delta)
                        s += delta
                        offsets_applied += ',%+g signal'%delta
                    s = 10 * np.log10(s)
                    if delta > 0.0:
                        # it is possible that we generated some -inf values, so raise them (see np.select())
                        # NB: -323dB was empirically determined to be the limit before -inf
                        DB_LIMIT = -323.0
                        condlist = [s<DB_LIMIT,s>=DB_LIMIT]
                        choicelist = [np.ones(len(s))*DB_LIMIT,s]
                        s = np.select(condlist,choicelist)
                        offsets_applied += ',%.0f dB limit'%DB_LIMIT
                    log.debug('-> s[] median=%f, mean=%f, stdev=%g, min=%f max=%f',s[s.argsort()[len(s)/2]],s.mean(),np.std(s),s.min(),s.max())
                std = np.std(s)
                if std == 0.0:
                    log.warning('-> signal stdev is 0dB; did you subtract the signal from itself?')
                    median = s[s.argsort()[len(s)/2]]
                    min = median - 1.0
                    max = median + 1.0
                elif std < 1.0:
                    log.warning('-> signal stdev<1.0 dB; scaling signal plot about the mean value (%g)',s.mean())
                    min = s.mean() - std*4
                    max = s.mean() + std*4
                else:
                    min = np.floor(s.min())-1.0
                    max = np.ceil(s.max())+1.0
                log.debug('-> plot mean=%f, stdev=%g, min=%f max=%f',s.mean(),std,min,max)
                if opts.hold:
                    std = np.std(hold)
                    if std == 0.0:
                        log.warning('-> hold stdev is 0dB; did you subtract the signal from itself?')
                        median = hold[hold.argsort()[len(hold)/2]]
                        hmin = median - 1.0
                        hmax = median + 1.0
                    elif std < 1.0:
                        log.warning('-> hold stdev<1.0 dB; scaling hold plot about the mean value (%g)',hold.mean())
                        hmin = hold.mean() - std*4
                        hmax = hold.mean() + std*4
                    else:
                        hmin = np.floor(hold.min())-1.0
                        hmax = np.ceil(hold.max())+1.0
                    log.debug('-> hold mean=%f, stdev=%g, min=%f max=%f',hold.mean(),std,hmin,hmax)
                    min = hmin if hmin < min else min
                    max = hmax if hmax > max else max
            log.debug('-> plot scale min=%g max=%g',min,max)

            ###
            ### Determination of Graph Title and Annotations
            ###
            tstop = dt.strftime('%Y-%m-%dT%H:%M:%S%z')
            if n > 1:
                title = 'Collected between %s and %s\n'%(tstart,tstop)
                if opts.atype.lower().startswith('linear'):
                    title += 'Linear averaged'
                else:
                    title += 'Logarithm averaged'
                title += ' over %d frames'%n
            else:
                title = 'Collected at %s'%tstart
            if len(opts.background) > 0:
                title = title + ', with background subtraction'
            if opts.background.lower() == 'automatic':
                m = bkg.mean()
                if opts.atype.lower().startswith('linear'):
                    if m <= 0.0:
                        title = title + ' (automatic, zero mean)'
                    else:
                        m = 10 * np.log10(m)
                        title = title + ' (automatic, %.2f dB)'%m
                else:
                    title = title + ' (automatic, %.2f dB)'%m
            else:
                title = title + ' (file)'
            if opts.calibration:
                title = title + ', cal=%.1f'%opts.calibration
            if opts.smooth > 0:
                title = title + ' and %d point smoothing'%opts.smooth
                ll = len(s)
                hh = opts.smooth/2
                ss = np.zeros(ll)
                for xx in range(hh):
                    ss[xx]      = s[xx]
                    ss[ll-xx-1] = s[ll-xx-1]
                for xx in range(ll-opts.smooth+1):
                    ss[hh+xx] = np.mean(s[xx:xx+opts.smooth])
                s = ss
                if opts.hold:
                    ll = len(hold)
                    hh = opts.smooth/2
                    ss = np.zeros(ll)
                    for xx in range(hh):
                        ss[xx]      = hold[xx]
                        ss[ll-xx-1] = hold[ll-xx-1]
                    for xx in range(ll-opts.smooth+1):
                        ss[hh+xx] = np.mean(hold[xx:xx+opts.smooth])
                    hold = ss

            # http://stackoverflow.com/questions/6352740/matplotlib-label-each-bin
            from matplotlib.pyplot import figure, plot, axis, xlabel, ylabel, savefig, subplots
            from matplotlib.pyplot import title as _title
            from matplotlib.ticker import FormatStrFormatter

            fig, ax = subplots()
            plot(fMHz,s,hold=True,color='b',label='foreground')
            if opts.hold:
                plot(fMHz,hold,hold=True,color='g',label='hold')
            axis([fMHz[0],fMHz[nbin-1],min,max])
            ax.xaxis.set_major_formatter(FormatStrFormatter('%4.1f'))
            xlabel('frequency (MHz)')
            txt = 'spectral power'
            if dbm != 0.0:
                if opts.ptype.lower().startswith('linear'):
                    txt += ' (amplitude relative to 1dBm/Hz)'
                else:
                    txt += ' (dBm/Hz)'
            elif len(opts.background) > 0:
                if opts.ptype.lower().startswith('linear'):
                    txt += ' (amplitude relative to background)'
                else:
                    txt += ' (dB relative to background)'
                if opts.bplot and not opts.background.lower() == 'automatic':
                    # when doing an 'automatic', we used a mean value, so there isnt any point in plotting a straight line
                    plot(fMHz,bkg,color='r',hold=True,label='background')
            else:
                txt += ' (arbitrary unit)'
            txt += offsets_applied
            ylabel(txt)
            _title(title)

            name = 'spectrum-%s.png'%dt.strftime('%Y_%b_%d_%H_%M_%S')
            savefig(name)
            log.info('Saved '+name)

            acc = np.zeros(nbin)
            n   = 0
            del fig, ax

def dump_spectrum_info(filename,opts):
    log = logging.getLogger(__name__)
    fg = open_spectrum_file(filename,opts)
    ts_a = fg['time']
    fMHz = fg['freq']
    zidx = np.searchsorted(fMHz,0)
    nbin = len(fMHz)
    lastrow = len(ts_a)-1
    bw   = (fMHz.max() - fMHz.min())*1e6/nbin

    log.info('range (MHz)=[%f,%f]',fMHz.min(),fMHz.max())
    log.info('zero index=%d',zidx)
    log.info('frequency bins=%d',nbin)
    log.info('bandwidth/bin (Hz)=%f',bw)

    tstart = ts_a[0].strftime('%Y-%m-%dT%H:%M:%S')
    if lastrow > 0:
        delta = ts_a[lastrow] - ts_a[lastrow-1]
        log.info('start=%s',str(ts_a[0]))
        log.info('end  =%s',str(ts_a[lastrow]))
        deltaT = dt = float(delta.seconds)+(float(delta.microseconds)/1e6)
        if len(ts_a) > 0:
            deltaT = ts_a[lastrow] - ts_a[0]
            deltaT = float(deltaT.seconds)+(float(deltaT.microseconds)/1e6)
            deltaT = deltaT / len(ts_a)
        tstop  = ts_a[lastrow].strftime('%Y-%m-%dT%H:%M:%S')
        log.info('seconds between records=%f (avg=%f)',dt,deltaT)
        title  = 'Collected between %s and %s'%(tstart,tstop)
        deltaT = min(dt,deltaT)
    else:
        title  = 'Collected at %s'%tstart
        deltaT = 0.0

    log.info('Analyzing records...  acceptable inter-frame time is [%f,%f] sec',deltaT-(deltaT*DEF_DT_VAR),deltaT+(deltaT*DEF_DT_VAR))
    if opts.statistics:
        sts = np.zeros((len(ts_a),7)).astype(np.float)
        if not opts.line:
            s_a = read_spectrum_array(fg)
    dt  = np.zeros(len(ts_a)).astype(np.float)
    d0  = ts_a[0]
    dX  = 0.0
    for fr in range(len(ts_a)):
        if fr == 0:
            if opts.statistics:
                if opts.line:
                    s = read_spectrum_line(fg)
                else:
                    s = s_a[fr]
                # http://stackoverflow.com/questions/796008/cant-subtract-offset-naive-and-offset-aware-datetimes
                sts[fr,0] = dt2excel(ts_a[fr].replace(tzinfo=None))
                sts[fr,1] = 0.0
                sts[fr,2] = s.min()
                sts[fr,3] = s.max()
                sts[fr,4] = s.mean()
                sts[fr,5] = s.std()
                sts[fr,6] = s.var()
            dt[fr] = deltaT
            continue
        log.debug('  d0=%s ts_a[%d]=%s',str(d0),fr,str(ts_a[fr]))
        _dt0 = ts_a[fr] - d0
        dt[fr] = float(_dt0.seconds)+(float(_dt0.microseconds)/1e6)
        log.debug('    dt=%.6f abs(dt - deltaT)=%.6f',dt[fr],abs(dt[fr] - deltaT))
        if dt[fr] < 1e-6:
            log.info('*** DUPLICATE: frame %d -> %d delta %f sec, expected %f+/-%f (@%s)',fr-1,fr,dt[fr],deltaT,(deltaT*DEF_DT_VAR),d0.strftime('%Y-%m-%dT%H:%M:%S.%f'))
        elif ts_a[fr] < d0:
            log.info('*** REGRESSION: frame %d -> %d has earlier timestamp (%s) than (@%s)',fr-1,fr,ts_a[fr].strftime('%Y-%m-%dT%H:%M:%S.%f'),d0.strftime('%Y-%m-%dT%H:%M:%S.%f'))
        elif abs(dt[fr] - deltaT) > (deltaT*DEF_DT_VAR):
            log.info('*** SKIP: frame %d -> %d delta %f sec, expected %f+/-%f (@%s)',fr-1,fr,dt[fr],deltaT,(deltaT*DEF_DT_VAR),d0.strftime('%Y-%m-%dT%H:%M:%S.%f'))
            dX = dX + (dt[fr] - deltaT)
        d0 = ts_a[fr]
        if opts.statistics:
            if opts.line:
                s = read_spectrum_line(fg)
            else:
                s = s_a[fr]
            sts[fr,0] = dt2excel(ts_a[fr].replace(tzinfo=None))
            sts[fr,1] = dt[fr]
            sts[fr,2] = s.min()
            sts[fr,3] = s.max()
            sts[fr,4] = s.mean()
            sts[fr,5] = s.std()
            sts[fr,6] = s.var()
    if dX > 0.0:
        d0 = ts_a[0]
        dN = ts_a[-1]
        _dt0 = dN - d0
        capture_length = float(_dt0.seconds)+(float(_dt0.microseconds)/1e6)
        if capture_length > 0.0:
            log.info('*** %f sec unaccounted for in %f sec capture (%.3f%%)',dX,capture_length,(dX*100.0/capture_length))
        elif capture_length < 0.0:
            log.info('*** %f sec unaccounted for but total capture length appears to be less than 0.0 (%.3f sec)',dX,capture_length)
        else:
            log.info('*** %f sec unaccounted for but total capture length appears to be 0.0',dX)
    log.info('*** inter-frame min=%.3f,max=%.3f,mean=%.3f (sec)'%(dt.min(),dt.max(),dt.mean()))
    if opts.statistics:
        np.savetxt(opts.statistics,sts,fmt=DEF_STATS_FORMAT,delimiter=DEF_STATS_DELIM,header=DEF_STATS_HEADER,comments='')

# Refactored to make argument parsing separate so as to provide a GUI option
def execute(opts):
    args = opts.file

    logging.basicConfig(format='%(message)s',level=logging.DEBUG)
    if args==[]:
        logging.getLogger(__name__).critical('Please specify a filename. Run with the -h flag to see all options.')
    else:
        # logging boilerplate (screen+log file)
        logger = logging.getLogger(__name__)
        x,y,name = args[0].replace('\\','/').rpartition('/')
        name,x,ext = name.rpartition('.')
        handler = logging.FileHandler(name+'.log')
        formatter = logging.Formatter('%(levelname)s:%(message)s')
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        logger.setLevel(logging.DEBUG if opts.verbose else logging.INFO)
        try:
            if opts.gui:
                from platform import system
                if not opts.matplotlib_use and system().startswith('Windows'):
                    from matplotlib import use
                    use('wxagg')
                    opts.matplotlib_use = True
            if opts.info:
                dump_spectrum_info(args[0],opts)
            else:
                generate_spectrum_plots(args[0],opts)
            if opts.gui:
                from pylab import show
                print('=== Display interactive graphs ===')
                show()
        except Exception, e:
            logger.error('generate_spectrum_plots', exc_info=True)
            exit(1)

if __name__ == '__main__':
    import argparse
    p = argparse.ArgumentParser(
        description='Post-Process RASDRviewer/RASDRproc spectrum data output files')
    p.add_argument(      "--version", action='version', version='%(prog)s '+DEF_VERSION)
    g = p.add_mutually_exclusive_group()
    g.add_argument("-v", "--verbose", default=False, help="Turn on verbose output", action="store_true")
    g.add_argument("-q", "--quiet", default=False, help='Suppress progress messages', action="store_true")
    # call matplotlib.use() only once
    p.set_defaults(matplotlib_use = False)
    p.add_argument('-a', '--average', type=int, metavar='N', default=DEF_AVERAGE,
        help='Specify the number of spectra to average for each plot')
    p.add_argument('-b', '--background', type=str, metavar='PATH', default='',
        help='Specify how to perform background subtraction;'+
             'if the word automatic is used, then the background will be taken'+
             'from the average of all lines in the file.  Otherwise, it is taken'+
             'as a file to process.  The file must have the same frequency plan as the foreground file.')
    p.add_argument('-c', '--cancel-dc', dest='canceldc', action='store_true', default=False,
        help='Cancel out component at frequency bin for 0Hz')
    p.add_argument('-d', '--delimiter', type=str, metavar='CHAR', default=DEF_DELIM,
        help='Specify the delimiter character to use"')
    p.add_argument('-e', '--localtime', action='store_true', default=False,
        help='Indicate that .csv file has timestamps in RASDRviewer\'s "LocalTime" format')
    p.add_argument('-k', '--calibration', type=float, metavar='CONST', default=DEF_CALIB,
        help='Specify the calibration constant for the system; 0.0=uncal')    #default=%f'%DEF_CALIB)
    p.add_argument('-l', '--line', action='store_true', default=False,
        help='Perform line-by-line processing instead of loading entire file(s); NOTE: much slower but tolerates low memory better.')
##    p.add_argument('-m', '--milliwatt', dest='dbm', action='store_true', default=False,
##        help='Plot in decibels referenced to 1mW (dBm/Hz)')
##    p.add_argument('-t', '--datetime', action='store_true', default=False,
##        help='Indicate that timestamps in the .csv file are in Excel\'s datetime format')
    p.add_argument('-i', '--info', action='store_true', default=False,
        help='Produce information about a file only; do not generate any plots')
    p.add_argument('--statistics', type=str, metavar='PATH', default=None,
        help='Dump statistical information to a file in comma-separated-values format')
    p.add_argument('--debug', action='store_true', default=False,
        help='Drop into ipython shell at predefined point(s) to debug the script')
    p.add_argument('-g', '--gui', action='store_true', default=False,
        help='Create interactive PLOTS')
    p.add_argument('-s', '--smooth', type=int, metavar='N', default=0,
        help='Smooth final plot using a sliding window of N points')
    p.add_argument('--fcenter', dest='fc', type=float, default=0.0,
        help='Define the offset for the center frequency in Hz')  #default=%f'%0.0)
    p.add_argument('--hold', action='store_true', default=False,
        help='Perform a maximum value HOLD during averaging and plot it as a second line')
    p.add_argument('--bplot', action='store_true', default=False,
        help='If using background file, choose whether to plot the background reference in a difffert color')
    p.add_argument('--ptype', type=str, metavar='TYPE', default='log',
        help='Control plot vertical scale (linear or log)')
    p.add_argument('--atype', type=str, metavar='TYPE', default='log',
        help='Control averaging method (linear or log)')
        # http://www.dtic.mil/dtic/tr/fulltext/u2/657404.pdf
    ## for handling RASDRviewer versions
    v = DEF_VERSION.split('.')
    ver = v[0]+'.'+v[1]+'.'+v[2]
    p.add_argument('--format', type=str, metavar='X.Y.Z', default=ver,
        help='Specify the RASDRviewer .csv output format to interpret')
    # http://stackoverflow.com/questions/20165843/argparse-how-to-handle-variable-number-of-arguments-nargs
    p.add_argument(      "file", nargs='*')

    opts = p.parse_args(sys.argv[1:])
    execute(opts)

max_stars_repo_name: mfkiwl/RASDR

max_issues_repo_name: mfkiwl/RASDR

max_forks_repo_name: mfkiwl/RASDR


text:
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import torch.nn.functional as F

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        weight_shape = list(m.weight.data.size())
        fan_in = np.prod(weight_shape[1: 4])
        fan_out = np.prod(weight_shape[2: 4]) * weight_shape[0]
        w_bound = np.sqrt(6. / (fan_in + fan_out))
        m.weight.data.uniform_(-w_bound, w_bound)
        m.bias.data.fill_(0.0)
    elif classname.find('Linear') != -1:
        weight_shape = list(m.weight.data.size())
        fan_in = weight_shape[1]
        fan_out = weight_shape[0]
        w_bound = np.sqrt(6.0 / (fan_in + fan_out))
        m.weight.data.uniform_(-w_bound, w_bound)
        m.bias.data.fill_(0.0)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.fill_(1.0)
        m.bias.data.fill_(0.0)
    elif classname.find('LSTMCell') != -1:
        m.bias_ih.data.fill_(0.0)
        m.bias_hh.data.fill_(0.0)

class STN3d(nn.Module):
    def __init__(self, num_points = 2500):
        super(STN3d, self).__init__()
        self.num_points = num_points
        self.conv1 = torch.nn.Conv1d(4, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.mp1 = torch.nn.MaxPool1d(num_points)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 9)
        self.relu = nn.ReLU()

        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(256)

    def forward(self, x):
        # x --> 3 * 3
        batchsize = x.shape[0]
        if batchsize > 1:
            x = F.relu(self.bn1(self.conv1(x)))
            x = F.relu(self.bn2(self.conv2(x)))
            x = F.relu(self.bn3(self.conv3(x)))
            x = self.mp1(x)
            x = x.view(-1, 1024)

            x = F.relu(self.bn4(self.fc1(x)))
            x = F.relu(self.bn5(self.fc2(x)))
        else:
            x = F.relu(self.conv1(x))
            x = F.relu(self.conv2(x))
            x = F.relu(self.conv3(x))
            x = self.mp1(x)
            x = x.view(-1, 1024)

            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))

        x = self.fc3(x)

        iden = Variable(torch.eye(3)).view(1, -1).repeat(batchsize, 1)
        if x.is_cuda:
            device = torch.device('cuda:%d' % x.get_device())
            iden = iden.to(device=device)
        x = x + iden
        x = x.view(-1, 3, 3)

        return x

class PointNetfeat(nn.Module):
    def __init__(self, num_points=2500, global_feat=True):
        super(PointNetfeat, self).__init__()
        self.stn = STN3d(num_points=num_points)
        self.conv1 = torch.nn.Conv1d(4, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.mp1 = torch.nn.MaxPool1d(num_points)
        self.num_points = num_points
        self.global_feat = global_feat

    def forward(self, x):
        trans = self.stn(x)
        x = torch.cat([torch.bmm(trans, x[:, :3, :]), x[:, 3, :].unsqueeze(1)], dim=1)

        if x.shape[0] > 1:
            x = F.relu(self.bn1(self.conv1(x)))
            pointfeat = x
            x = F.relu(self.bn2(self.conv2(x)))
            x = self.bn3(self.conv3(x))
        else:
            x = F.relu(self.conv1(x))
            pointfeat = x
            x = F.relu(self.conv2(x))
            x = self.conv3(x)

        x = self.mp1(x)
        x = x.view(-1, 1024)

        if self.global_feat:
            return x, trans
        else:
            x = x.view(-1, 1024, 1).repeat(1, 1, self.num_points)
            return torch.cat([x, pointfeat], 1), trans

class CNNfeat(nn.Module):
    def __init__(self):
        super(CNNfeat, self).__init__()
        self.conv1 = nn.Sequential(  # 960 * 640
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=10,
                      stride=3, padding=0),  # 317 * 211
            nn.ReLU(),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 64, 5, 2, 0),  # 157 * 104
            nn.ReLU(),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(64, 32, 3, 2, 0),  # 78 * 51
            nn.ReLU(),
            # nn.MaxPool2d(kernel_size=2)  # 19 * 25
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(32, 16, 3, 2, 0),  # 38 * 25
            nn.ReLU(),
        )
        self.out = nn.Sequential(
            nn.Linear(16 * 38 * 25, 4096),
            nn.ReLU(),
            nn.Linear(4096, 2048),
            nn.ReLU(),
            nn.Linear(2048, 1024)
        )
        # self.out = nn.MaxPool2d(kernel_size=(18, 25))

        self.apply(weights_init)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)

        x = x.view(x.size(0), -1)
        x = self.out(x)
        x = x.view(-1, 1024)

        return x

class end_layer(nn.Module):
    def __init__(self, in_channels=1024, out_channels=1):
        super(end_layer, self).__init__()
        self.fc1 = nn.Linear(in_channels, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, out_channels)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)

        self.apply(weights_init)

    def forward(self, x):
        if x.size()[0] == 1:
            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))
        else:
            x = F.relu(self.bn1(self.fc1(x)))
            x = F.relu(self.bn2(self.fc2(x)))
        return self.fc3(x)

class PointNetActorCritic(nn.Module):
    def __init__(self, num_points=2500, num_actions=5):
        super(PointNetActorCritic, self).__init__()
        self.num_points = num_points
        self.feat = PointNetfeat(num_points, global_feat=True)

        self.lstm = nn.LSTMCell(1024, 1024)

        self.critic_linear = end_layer(in_channels=1024, out_channels=1)
        self.actor_linear = end_layer(in_channels=1024, out_channels=num_actions)

        self.apply(weights_init)
        self.train()

    def forward(self, inputs):
        x, (hx, cx) = inputs
        x, _ = self.feat(x)
        hx, cx = self.lstm(x, (hx, cx))
        x = hx

        return self.critic_linear(x), self.actor_linear(x), (hx, cx)

class PointNetGail(nn.Module):
    def __init__(self, num_points=2500, num_actions=5):
        super(PointNetGail, self).__init__()
        self.num_points = num_points
        self.feat = PointNetfeat(num_points, global_feat=True)

        self.critic_linear = end_layer(in_channels=1024, out_channels=1)
        self.actor_linear = end_layer(in_channels=1024, out_channels=num_actions)

        self.apply(weights_init)
        self.train()

    def forward(self, inputs):
        x = inputs
        x, _ = self.feat(x)

        return self.critic_linear(x), self.actor_linear(x)

if __name__ == '__main__':
    sim_data = Variable(torch.rand(10, 4, 2500))

    # trans = STN3d()
    # out = trans(sim_data)
    # print('stn', out.size())

    # pointfeat = PointNetfeat(global_feat=True)
    # out, _ = pointfeat(sim_data)
    # print('global feat', out.size())

    # pointfeat = PointNetfeat(global_feat=False)
    # out, _ = pointfeat(sim_data)
    # print('point feat', out.size())

    cls = PointNetGail(num_actions=5)
    v, q= cls(sim_data)
    print(v.shape, q.shape)
    print(v)
    print(q)

ext: py

max_issues_repo_name: emigmo/Transferable-Active-Grasping

max_forks_repo_name: emigmo/Transferable-Active-Grasping

reason: import numpy


text:
"""Fine-tuning on the IMDB dataset"""

from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
from transformers import HfArgumentParser, Trainer, TrainingArguments, set_seed, EvalPrediction
from torch.utils.tensorboard import SummaryWriter
from datasets import load_dataset
import torch

import transformers
transformers.logging.set_verbosity_debug()

import os
import sys
import logging

from scipy.special import softmax
from sklearn.metrics import f1_score, precision_score, recall_score, accuracy_score

import dataclasses
from dataclasses import dataclass, field
from typing import Optional
# We use dataclass-based configuration objects, let's define the one related to
# which model we are going to train here:
@dataclass
class AdditionalArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_path: str = field(
        metadata={"help": "Path to pretrained model"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )

def preprocess_examples(tokenizer, examples):
    """Preprocess batched examples using a pre-trained BERT tokeniser."""
    tokenized_examples = tokenizer(examples['text'], truncation=True, padding=True)
    tokenized_examples["labels"] = [val for val in examples["label"]]
    return tokenized_examples

def load_dataset_for_fine_tuning(tokenizer, seed):
    """Load the IMDB dataset using datasets library.

    Notes:
        When using `datasets` and `Trainer`, we do not have to create a `DataLoader`.
        (i.e., do not wrap `datasets` with `DataLoader`, `Trainer` only accepts
        `torch.utils.data.dataset.Dataset`!)
        Just pass the resulting data object to the instance of `Trainer`!

    References:
        https://huggingface.co/nlp/viewer/?dataset=imdb
        https://huggingface.co/datasets/imdb
        https://huggingface.co/docs/datasets/torch_tensorflow.html
    """
    dataset = load_dataset("imdb")
    train_set = dataset["train"]
    train_set.shuffle(seed=seed)
    test_set = dataset["test"]
    test_set.shuffle(seed=seed)

    print("Preprocess training samples...")
    train_set = train_set.map(lambda examples: preprocess_examples(tokenizer, examples),
                              batched=True, batch_size=5000)
    print("Preprocess test samples...")
    test_set = test_set.map(lambda examples: preprocess_examples(tokenizer, examples),
                            batched=True, batch_size=5000)

    train_set.set_format(type='torch', columns=['input_ids','attention_mask', 'labels'])
    test_set.set_format(type='torch', columns=['input_ids','attention_mask', 'labels'])

    return train_set, test_set

def main():
    ##########
    # Configs
    ##########
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.
    parser = HfArgumentParser((AdditionalArguments, TrainingArguments))

    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        args, training_args = parser.parse_args_into_dataclasses()

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
        )

    ##########
    # Setup logging
    ##########
    logger = logging.getLogger(__name__)
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
        training_args.local_rank,
        training_args.device,
        training_args.n_gpu,
        bool(training_args.local_rank != -1),
        training_args.fp16,
    )
    logger.info("Training/evaluation parameters %s", training_args)

    # Set seed
    set_seed(training_args.seed)

    ##########
    # Load pretrained model/data and tokenizer
    ##########
    config = AutoConfig.from_pretrained(args.model_path, num_labels=2)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name)
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_path,
        from_tf=bool(".ckpt" in args.model_path),
        config=config
    )

    ##########
    # Get datasets
    ##########
    # Do not pass the resulting data to `DataLoader`!
    train_set, test_set = load_dataset_for_fine_tuning(tokenizer, training_args.seed)

    ##########
    # Set up evaluation metrics
    ##########
    def compute_metrics_fn(p: EvalPrediction):
        """
        Compute accuracy, f1, precision, and recall.

        References:
            https://huggingface.co/transformers/main_classes/trainer.html#transformers.EvalPrediction
        """
        # preprocess
        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
        preds = softmax(preds, axis=1)[:, 1] # proba
        preds = (preds >= 0.5).astype(int) # binary

        # compute each metric using scikit-learn
        acc = accuracy_score(p.label_ids, preds)
        f1 = f1_score(p.label_ids, preds)
        prec = precision_score(p.label_ids, preds)
        rec = recall_score(p.label_ids, preds)

        return {"acc": acc, "f1": f1, "precision": prec, "recall": rec}

    ##########
    # Set up a Trainer
    ##########
    tb_writer = SummaryWriter(training_args.logging_dir)
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_set,
        eval_dataset=test_set,
        tb_writer=tb_writer,
        compute_metrics=compute_metrics_fn
    )

    # training
    trainer.train()

    # final evaluation
    logger.info("*** Final evaluation ***")
    eval_result = trainer.evaluate(eval_dataset=test_set)
    for key, value in eval_result.items():
        logger.info("\t%s = %s", key, value)

if __name__ == "__main__":
    main()

reason: from scipy


text:
import numpy as np
from sklearn.model_selection import StratifiedKFold, train_test_split
import os
from min2net.utils import butter_bandpass_filter
from min2net.preprocessing.MBDEPV1 import raw
from min2net.preprocessing.config import CONSTANT

CONSTANT = CONSTANT["MBDEPV1"]
raw_path = CONSTANT["raw_path"]
n_subjs = CONSTANT["n_subjs"]
# n_trials = CONSTANT['n_trials']
n_trials_per_class = CONSTANT["n_trials_per_class"]
n_chs = CONSTANT["n_chs"]
orig_smp_freq = CONSTANT["orig_smp_freq"]
MI_len = CONSTANT["MI"]["len"]

def subject_dependent_setting(
    k_folds, pick_smp_freq, bands, order, save_path, num_class=2, sel_chs=None
):
    sel_chs = CONSTANT["sel_chs"] if sel_chs == None else sel_chs
    n_folds = k_folds
    save_path = save_path + "/MBDEPV1/time_domain/{}_class/subject_dependent".format(
        num_class
    )
    n_chs = len(sel_chs)
    n_trials = n_trials_per_class * num_class
    # X_train_all, y_train_all = np.zeros(
    #     (n_subjs, n_trials, n_chs, int(MI_len * pick_smp_freq))
    # ), np.zeros((n_subjs, n_trials))
    # X_test_all, y_test_all = np.zeros(
    #     (n_subjs, n_trials, n_chs, int(MI_len * pick_smp_freq))
    # ), np.zeros((n_subjs, n_trials))

    id_chosen_chs = raw.chanel_selection(sel_chs)

    X_train_all, y_train_all, X_test_all, y_test_all= __load_MBDEPV1(
        raw_path, 1, pick_smp_freq, num_class, id_chosen_chs
    )
    X_train_all = X_train_all.reshape(1, *X_train_all.shape)
    X_test_all = X_test_all.reshape(1, *X_test_all.shape)
    y_train_all = y_train_all.reshape(1, *y_train_all.shape)
    y_test_all = y_test_all.reshape(1, *y_test_all.shape)

    # for s in range(n_subjs):
    #     X_train, y_train, X_test, y_test = __load_MBDEPV1(
    #         raw_path, s + 1, pick_smp_freq, num_class, id_chosen_chs
    #     )
    #     X_train_all[s], y_train_all[s] = X_train, y_train
    #     X_test_all[s], y_test_all[s] = X_test, y_test

    for directory in [save_path]:
        if not os.path.exists(directory):
            os.makedirs(directory)

    # Carry out subject-dependent setting with 5-fold cross validation
    for person, (X_tr, y_tr, X_te, y_te) in enumerate(
        zip(X_train_all, y_train_all, X_test_all, y_test_all)
    ):
        if len(X_tr.shape) != 3:
            raise Exception("Dimension Error, must have 3 dimension")

        if n_folds == 1:
            X_tr_cv, X_val_cv, y_tr_cv, y_val_cv = train_test_split(
                X_tr, y_tr, test_size=0.2, random_state=42, shuffle=True
            )
            print("Band-pass filtering from {} to {} Hz.".format(bands[0], bands[1]))
            X_tr_fil = butter_bandpass_filter(
                X_tr_cv, bands[0], bands[1], pick_smp_freq, order
            )
            X_val_fil = butter_bandpass_filter(
                X_val_cv, bands[0], bands[1], pick_smp_freq, order
            )
            X_te_fil = butter_bandpass_filter(
                X_te, bands[0], bands[1], pick_smp_freq, order
            )
            print(
                "Check dimension of training data {}, val data {} and testing data {}".format(
                    X_tr_fil.shape, X_val_fil.shape, X_te_fil.shape
                )
            )
            SAVE_NAME = "S{:03d}_fold001".format(person + 1)
            __save_data_with_valset(
                save_path,
                SAVE_NAME,
                X_tr_fil,
                y_tr_cv,
                X_val_fil,
                y_val_cv,
                X_te_fil,
                y_te,
            )
            print("The preprocessing of subject {} from is DONE!!!".format(person + 1))
        else:
            skf = StratifiedKFold(n_splits=n_folds, random_state=42, shuffle=True)
            for fold, (train_index, val_index) in enumerate(skf.split(X_tr, y_tr)):
                print(
                    "FOLD:",
                    fold + 1,
                    "TRAIN:",
                    len(train_index),
                    "VALIDATION:",
                    len(val_index),
                )
                X_tr_cv, X_val_cv = X_tr[train_index], X_tr[val_index]
                y_tr_cv, y_val_cv = y_tr[train_index], y_tr[val_index]

                print(
                    "Band-pass filtering from {} to {} Hz.".format(bands[0], bands[1])
                )
                X_tr_fil = butter_bandpass_filter(
                    X_tr_cv, bands[0], bands[1], pick_smp_freq, order
                )
                X_val_fil = butter_bandpass_filter(
                    X_val_cv, bands[0], bands[1], pick_smp_freq, order
                )
                X_te_fil = butter_bandpass_filter(
                    X_te, bands[0], bands[1], pick_smp_freq, order
                )
                print(
                    "Check dimension of training data {}, val data {} and testing data {}".format(
                        X_tr_fil.shape, X_val_fil.shape, X_te_fil.shape
                    )
                )
                SAVE_NAME = "S{:03d}_fold{:03d}".format(person + 1, fold + 1)
                __save_data_with_valset(
                    save_path,
                    SAVE_NAME,
                    X_tr_fil,
                    y_tr_cv,
                    X_val_fil,
                    y_val_cv,
                    X_te_fil,
                    y_te,
                )
                print(
                    "The preprocessing of subject {} from fold {} is DONE!!!".format(
                        person + 1, fold + 1
                    )
                )

def __load_MBDEPV1(PATH, subject, new_smp_freq, num_class, id_chosen_chs):
    start = CONSTANT["MI"]["start"]  # 0
    stop = CONSTANT["MI"]["stop"]  # 2
    X_train, y_tr, X_test, y_te = raw.load_crop_data(
        PATH=PATH,
        subject=subject,
        start=start,
        stop=stop,
        new_smp_freq=new_smp_freq,
        num_class=num_class,
        id_chosen_chs=id_chosen_chs,
    )
    return X_train, y_tr, X_test, y_te

def __save_data_with_valset(
    save_path, NAME, X_train, y_train, X_val, y_val, X_test, y_test
):
    np.save(save_path + "/X_train_" + NAME + ".npy", X_train)
    np.save(save_path + "/X_val_" + NAME + ".npy", X_val)
    np.save(save_path + "/X_test_" + NAME + ".npy", X_test)
    np.save(save_path + "/y_train_" + NAME + ".npy", y_train)
    np.save(save_path + "/y_val_" + NAME + ".npy", y_val)
    np.save(save_path + "/y_test_" + NAME + ".npy", y_test)
    print("save DONE")


text:
# -*- coding: utf-8 -*-
import numpy as np
import unittest

from glemmazon import encoder as t

class TestEncoder(unittest.TestCase):
    def testDenseTag(self):
        enc_tag = t.DenseTag(['_UNK', 'A', 'B', 'C'])
        np.testing.assert_array_equal(enc_tag('_UNK'),
                                      np.array([1., 0., 0., 0.]))
        np.testing.assert_array_equal(enc_tag('B'),
                                      np.array([0., 0., 1., 0.]))
        self.assertEqual(enc_tag.output_shape, (4,))

    def testDenseWordSuffix(self):
        enc_suf = t.DenseWordSuffix(['a', 'b', 'c'], 2, '_UNK')
        np.testing.assert_array_equal(
            enc_suf('ab'), np.array([0., 1., 0., 0.,
                                     0., 0., 1., 0.]))
        # Note that '_UNK' is [1., 0. , 0., 0.]
        np.testing.assert_array_equal(
            enc_suf('dd'), np.array([1., 0., 0., 0.,
                                     1., 0., 0., 0.]))
        self.assertEqual(enc_suf.output_shape, (8,))

    def testDictFeatureEncoder(self):
        dict_enc = t.DictFeatureEncoder({
            'pos': t.DenseTag(['NOUN', 'VERB']),
            'tense': t.DenseTag(['PRES', 'FUT', 'PAST']),
            'word': t.DenseWordSuffix(['a', 'b'], 2, '_UNK'),
        })
        # Note that '_UNK' is [1., 0. , 0., 0.]
        np.testing.assert_array_equal(
            dict_enc({'pos': 'NOUN', 'tense': 'PRES', 'word': 'bb'}),
            np.array([
                0., 1., 0.,  # NOUN
                0., 1., 0., 0.,  # PRES
                0., 0., 1., 0., 0., 1.  # bb
            ]))
        self.assertEqual(dict_enc.output_shape, (13,))
        self.assertCountEqual(dict_enc.scope, {'pos', 'tense', 'word'})

    # -------------------------------------------------------------------
    # Sequence feature encoders
    # -------------------------------------------------------------------
    def testSeqWordSuffix(self):
        seq_word_enc = t.SeqWordSuffix(['a', 'b', 'c'], suffix_length=3)
        np.testing.assert_array_equal(
            seq_word_enc('abbc'),
            np.array([[0., 0., 1., 0.],  # b
                      [0., 0., 1., 0.],  # b
                      [0., 0., 0., 1.],  # c
                      ]))
        self.assertEqual(seq_word_enc.output_shape, (3, 4))

    def testSeqFeatureEncoder(self):
        dict_enc = t.SeqFeatureEncoder(
            'word', t.SeqWordSuffix(['a', 'b', 'c', 'd'],
                                    suffix_length=4),
            t.DictFeatureEncoder({
                'pos': t.DenseTag(['NOUN', 'VERB']),
                'tense': t.DenseTag(['PRES', 'FUT', 'PAST']),
            }))

        # Note that '_UNK' is [1., 0. , 0., 0.]
        np.testing.assert_array_equal(
            dict_enc({'word': 'bcc', 'tense': 'FUT', 'pos': 'NOUN'}),
            np.array([
                [
                    1., 0., 0., 0., 0.,  # _UNK
                    0., 1., 0.,  # NOUN
                    0., 0., 1., 0.  # FUT
                ], [
                    0., 0., 1., 0., 0.,  # b
                    0., 1., 0.,  # NOUN
                    0., 0., 1., 0.  # FUT
                ], [
                    0., 0., 0., 1., 0.,  # c
                    0., 1., 0.,  # NOUN
                    0., 0., 1., 0.  # FUT
                ], [
                    0., 0., 0., 1., 0.,  # c
                    0., 1., 0.,  # NOUN
                    0., 0., 1., 0.  # FUT
                ]]))
        self.assertEqual(dict_enc.output_shape, (4, 12))
        self.assertCountEqual(dict_enc.scope, {'pos', 'tense', 'word'})

    # -------------------------------------------------------------------
    # Label encoders
    # -------------------------------------------------------------------
    def testLabelEncoder(self):
        le = t.LabelEncoder(['_UNK', 'A', 'B', 'C'])
        np.testing.assert_array_equal(le('_UNK'),
                                      np.array([1., 0., 0., 0.]))
        self.assertEqual(le.decode(np.array([1., 0., 0., 0.])), '_UNK')
        self.assertEqual(le.output_shape, (4,))

    def testDictLabelEncoder(self):
        dict_le = t.DictLabelEncoder({
            'lemma_suffix': t.LabelEncoder(['ed', 's', '']),
            'lemma_index': t.LabelEncoder(['1', '2', '3']),
        })
        # Normal encoding/decoding
        np.testing.assert_array_equal(
            dict_le({'lemma_suffix': 'ed', 'lemma_index': '2'}),
            np.array([[0., 0., 1., 0.],
                      [0., 1., 0., 0.]]))
        np.testing.assert_array_equal(
            dict_le({'lemma_index': '2', 'lemma_suffix': 'ed'}),
            np.array([[0., 0., 1., 0.],
                      [0., 1., 0., 0.]]))
        self.assertEqual(dict_le.decode(
            [np.array([0., 0., 1., 0.]),
             np.array([0., 1., 0., 0.])]),
            {'lemma_suffix': 'ed', 'lemma_index': '2'})
        self.assertCountEqual(dict_le.scope,
                              {'lemma_suffix', 'lemma_index'})

        # Single vector encoding/decoding
        np.testing.assert_array_equal(
            dict_le.to_single_vec(
                {'lemma_suffix': 'ed', 'lemma_index': '2'}),
            np.array([0., 0., 1., 0.,
                      0., 1., 0., 0.]))
        self.assertEqual(dict_le.from_single_vec(
            np.array([0., 0., 1., 0.,
                      0., 1., 0., 0.])),
            {'lemma_suffix': 'ed', 'lemma_index': '2'})

if __name__ == '__main__':
    unittest.main()

max_stars_repo_name: gustavoauma/glemmazon

max_issues_repo_name: gustavoauma/glemmazon

max_forks_repo_name: gustavoauma/glemmazon


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright © 2020 Michael J. Hayford
"""Post processing functions for Zemax import

.. Created on Mon Aug 10 18:17:55 2020

.. codeauthor: Michael J. Hayford
"""
import numpy as np
import rayoptics.seq.medium as mdm
import rayoptics.elem.profiles as profiles
import rayoptics.elem.surface as surface
from rayoptics.optical.model_enums import DecenterType as dec

def apply_fct_to_sm(opt_model, fct, start=None, stop=None, step=None):
    """Iterate in reverse over seq_model.ifcs. Override if needed."""
    sm = opt_model.seq_model
    start = len(sm.ifcs)-1 if start is None else start
    stop = 0 if stop is None else stop
    step = -1 if step is None else step
    num_changes = 0
    for cur in range(start, stop, step):
        if fct(opt_model, cur):
            num_changes += 1
    return num_changes

def convert_to_bend(opt_model, cur):
    """Scan the zemax import for tilted mirrors and convert to BEND types."""
    sm = opt_model.seq_model
    ifc = sm.ifcs[cur]
    if ifc.interact_mode == 'reflect':
        ifc_p = sm.ifcs[cur-1]
        ifc_f = sm.ifcs[cur+1]
        if (ifc_p.z_type == 'COORDBRK' and ifc_f.z_type == 'COORDBRK'):
            if np.array_equal(ifc_f.decenter.euler, ifc_p.decenter.euler):
                ifc.decenter = ifc_p.decenter
                ifc.decenter.dtype = dec.BEND
                sm.remove(cur+1, prev=True)
                sm.remove(cur-1)
                return True
    return False

def convert_to_dar(opt_model, cur):
    """Scan the zemax import for tilted surfs and convert to DAR types."""
    sm = opt_model.seq_model
    if cur < len(sm.ifcs)-1:
        ifc = sm.ifcs[cur]
        ifc_p = sm.ifcs[cur-1]
        ifc_f = sm.ifcs[cur+1]
        if (ifc_p.z_type == 'COORDBRK' and ifc_f.z_type == 'COORDBRK'):
            acum_dec = ifc_f.decenter.dec + ifc_p.decenter.dec
            acum_euler = ifc_f.decenter.euler + ifc_p.decenter.euler
            if np.all(acum_dec == 0) and np.all(acum_euler == 0):
                ifc.decenter = ifc_p.decenter
                ifc.decenter.dtype = dec.DAR
                sm.remove(cur+1, prev=True)
                sm.remove(cur-1)
                return True
    return False

def collapse_coordbrk(opt_model, cur):
    """Attempt to apply the cur COORDBRK to an adjacent real interface."""
    sm = opt_model.seq_model
    ifc_cb = sm.ifcs[cur]
    if ifc_cb.z_type == 'COORDBRK':
        if ifc_cb.decenter.dtype == dec.REV:
            ifc = sm.ifcs[cur-1]
            prev = True
        else:
            ifc = sm.ifcs[cur+1]
            prev = False

        if ifc.decenter is not None:
            return False
        else:
            ifc.decenter = ifc_cb.decenter
            sm.remove(cur, prev=prev)
            return True

    return False

def remove_null_sg(opt_model, cur):
    """Remove sg with planar profile and an adjacent zero thickness air gap."""
    sm = opt_model.seq_model
    ifc = sm.ifcs[cur]
    if is_null_ifc(ifc):
        prev = None
        cur_gap = False if len(sm.gaps)-1 < cur else True
        prev_gap = True if 0 < cur else False
        if cur_gap and is_null_gap(sm.gaps[cur]):
            prev = False
        elif prev_gap and is_null_gap(sm.gaps[cur-1]):
            prev = True
        if prev is not None:
            sm.remove(cur, prev=prev)
            return True

    return False

def is_null_ifc(ifc):
    if isinstance(ifc, surface.Surface):
        if isinstance(ifc.profile, profiles.Spherical):
            if (
                    ifc.profile.cv == 0 and
                    ifc.decenter is None and
                    ifc.interact_mode == 'transmit'
                    ):
                return True
    return False

def is_null_gap(gap):
    if gap.thi == 0 and isinstance(gap.medium, mdm.Air):
        return True
    else:
        return False

max_stars_repo_name: ajeddeloh/ray-optics

max_issues_repo_name: ajeddeloh/ray-optics

max_forks_repo_name: ajeddeloh/ray-optics


text:
# #!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
This module ...

__author__ = "Magnus Kvendseth Øye"
__copyright__ = "Copyright 2019, Sparkie Quadruped Robot"
__credits__ = ["Magnus Kvendseth Øye", "Petter Drønnen", "Vegard Solheim"]
__version__ = "1.0.0"
__license__ = "MIT"
__maintainer__ = "Magnus Kvendseth Øye"
__email__ = "magnus.oye@gmail.com"
__status__ = "Development"
"""

# Importing packages
import cv2
import sys
import time
import numpy as np
import zmq
import base64
from PyQt5 import QtWidgets, uic, QtCore, QtGui

# Importing from local source
from config import *
# from .util.xbox_controller import XboxController
from .util.navigation import Path

class ManualWindow(QtWidgets.QDialog):
    """doc"""

    activate = QtCore.pyqtSignal(bool)
    stop_camera = QtCore.pyqtSignal()
    stop_pose = QtCore.pyqtSignal()
    stop_serial = QtCore.pyqtSignal()
    stop_command = QtCore.pyqtSignal()
    stop_xbox_controller = QtCore.pyqtSignal()
    change_camera = QtCore.pyqtSignal(bool)

    def __init__(self):
        super(ManualWindow, self).__init__()
        self.setAttribute(QtCore.Qt.WA_DeleteOnClose)
        self.ui = '../forms/manual.ui'
        uic.loadUi(self.ui, self)
        self.setWindowFlags(QtCore.Qt.FramelessWindowHint)

        self.mode = JOYSTICK_ONLY_MODE

        # Camera filters
        self.current_camera_topic = 0
        self.camera_topics = ('color', 'fisheye', 'depth')

        # Buttons
        self.change_mode_btn = self.findChild(
            QtWidgets.QPushButton, 'changeModeBtn')
        self.turn_left = self.findChild(QtWidgets.QToolButton, 'turnRobotLeft')
        self.turn_right = self.findChild(
            QtWidgets.QToolButton, 'turnRobotRight')
        self.stand_btn = self.findChild(QtWidgets.QPushButton, 'standBtn')
        self.walk_btn = self.findChild(QtWidgets.QPushButton, 'walkBtn')
        self.stairs_btn = self.findChild(QtWidgets.QPushButton, 'stairsBtn')
        self.exit_btn = self.findChild(QtWidgets.QPushButton, 'exitBtn')
        self.powerBtn = self.findChild(QtWidgets.QPushButton, 'powerBtn')
        self.emergency_btn = self.findChild(
            QtWidgets.QPushButton, 'emergencyBtn')
        self.slow_btn = self.findChild(QtWidgets.QPushButton, 'slowBtn')
        self.medium_btn = self.findChild(QtWidgets.QPushButton, 'mediumBtn')
        self.fast_btn = self.findChild(QtWidgets.QPushButton, 'fastBtn')

        # Status indicators
        self.signal_btn = self.findChild(QtWidgets.QPushButton, 'signalBtn')
        self.controller_battery_btn = self.findChild(
            QtWidgets.QPushButton, 'controllerBatteryBtn')
        self.battery_btn = self.findChild(QtWidgets.QPushButton, 'batteryBtn')
        self.health_btn = self.findChild(QtWidgets.QPushButton, 'healthBtn')

        # Button connections
        self.emergency_btn.clicked.connect(self.turn_robot_off)
        self.powerBtn.clicked.connect(self.power_on)
        self.exit_btn.clicked.connect(self.close_window)
        self.change_mode_btn.clicked.connect(self.change_mode)
        self.stand_btn.clicked.connect(self.set_stand_btn)
        self.walk_btn.clicked.connect(self.set_walk_btn)
        self.stairs_btn.clicked.connect(self.set_stairs_btn)
        self.slow_btn.clicked.connect(self.set_slow_btn)
        self.medium_btn.clicked.connect(self.set_medium_btn)
        self.fast_btn.clicked.connect(self.set_fast_btn)
        self.turn_right.clicked.connect(self.change_camera_output)

        # Button shortcuts
        self.exit_btn.setShortcut("Ctrl+Q")

        # Mode Layouts
        self.video_frame = self.findChild(QtWidgets.QLabel, 'videoFrame')
        self.xbox_controller_frame = self.findChild(
            QtWidgets.QLabel, 'xboxcontrollerFrame')
        self.video_frame.hide()

        # Stylesheets
        self.powerBtn.setStyleSheet(
            "QPushButton#powerBtn:checked {color:black; background-color: red;}")
        self.signal_btn.setStyleSheet(
            "QPushButton#signalBtn:checked {color:black; background-color: green;}")

        self.initUI()

    def change_mode(self):
        if self.mode < MAX_MODES:
            self.mode += 1
        else:
            self.mode = JOYSTICK_ONLY_MODE
        self.update_mode_label()

    def update_mode_label(self):
        if self.mode == 0:
            self.change_mode_btn.setText("JOYSTICK ONLY")
            self.video_frame.hide()
            self.xbox_controller_frame.show()
        elif self.mode == 1:
            self.change_mode_btn.setText("ROBOT CAMERAS")
            self.xbox_controller_frame.hide()
            self.video_frame.show()
        elif self.mode == 2:
            self.change_mode_btn.setText("WAYPOINTS")
            self.xbox_controller_frame.hide()
            self.video_frame.hide()
            #self.path_frame.show()

    def set_walk_btn(self):
        if self.walk_btn.isChecked():
            self.stand_btn.setChecked(False)
            self.stairs_btn.setChecked(False)

    def set_stand_btn(self):
        if self.stand_btn.isChecked():
            self.walk_btn.setChecked(False)
            self.stairs_btn.setChecked(False)

    def set_stairs_btn(self):
        if self.stairs_btn.isChecked():
            self.walk_btn.setChecked(False)
            self.stand_btn.setChecked(False)

    def set_slow_btn(self):
        if self.slow_btn.isChecked():
            self.fast_btn.setChecked(False)
            self.medium_btn.setChecked(False)

    def set_medium_btn(self):
        if self.medium_btn.isChecked():
            self.slow_btn.setChecked(False)
            self.fast_btn.setChecked(False)

    def set_fast_btn(self):
        if self.fast_btn.isChecked():
            self.slow_btn.setChecked(False)
            self.medium_btn.setChecked(False)

    @QtCore.pyqtSlot(QtGui.QImage)
    def set_image(self, image):
        self.video_frame.setPixmap(QtGui.QPixmap.fromImage(image))

    def power_on(self):
        active = self.powerBtn.isChecked()
        if active:
            self.activate.emit(True)
        else:
            self.activate.emit(False)

    def blink_connection(self):
        if not self.signal_btn.isChecked():
            self.signal_btn.setChecked(True)
        else:
            self.signal_btn.setChecked(False)

    def initUI(self):
        self.activate.connect(self.active_all)
        self.init_camera()
        self.init_pose()
        self.init_serial()
        self.init_xbox_controller()
        # self.init_command()
        self.show()

    def init_camera(self):
        self.camera = CameraThread(self)
        self.camera.change_pixmap.connect(self.set_image)
        self.camera.blink_signal.connect(self.blink_connection)
        self.stop_camera.connect(self.camera.close_socket)
        self.change_camera.connect(self.change_camera_output)
        self.camera.start()

    def init_pose(self):
        self.pose = PoseThread(self)
        self.stop_pose.connect(self.pose.close_socket)
        self.pose.start()

    def init_serial(self):
        self.serial = SerialThread(self)
        self.stop_serial.connect(self.serial.close_socket)
        self.serial.start()

    def init_command(self):
        self.command = CommandThread(self)
        self.stop_command.connect(self.command.close_socket)
        self.command.start()

    def init_xbox_controller(self):
        self.xbox_controller = XboxControllerThread()
        self.stop_xbox_controller.connect(self.xbox_controller.close_socket)
        self.xbox_controller.start()

    def change_camera_output(self):
        self.current_camera_topic += 1
        self.camera.change_topic(self.camera_topics[self.current_camera_topic])
        if self.current_camera_topic == 2:
            self.current_camera_topic = -1

    def active_all(self):
        self.camera.activate(self.powerBtn.isChecked())
        self.pose.activate(self.powerBtn.isChecked())
        self.serial.activate(self.powerBtn.isChecked())
        self.xbox_controller.activate(self.powerBtn.isChecked())
        # self.command.activate(self.powerBtn.isChecked())

    def close_window(self):
        self.stop_camera.emit()
        self.stop_pose.emit()
        self.stop_serial.emit()
        self.stop_xbox_controller.emit()
        # self.stop_command.emit()
        self.close()

    def turn_robot_off(self):
        pass

class CameraThread(QtCore.QThread):

    blink_signal = QtCore.pyqtSignal()
    power_on = QtCore.pyqtSignal()
    change_pixmap = QtCore.pyqtSignal(QtGui.QImage)

    IP = '10.10.10.243'
    PORT = 5555
    FILTER = 'color'

    context = zmq.Context()
    footage_socket = context.socket(zmq.SUB)
    footage_socket.connect(f'tcp://{IP}:{PORT}')
    footage_socket.setsockopt_string(zmq.SUBSCRIBE, FILTER)
    footage_socket.setsockopt(zmq.CONFLATE, 1)  # last msg only.

    running = False

    active = False
    threadactive = True

    @QtCore.pyqtSlot()
    def close_socket(self):
        self.threadactive = False
        self.footage_socket.disconnect(f'tcp://{self.IP}:{self.PORT}')
        self.running = False
        print("Closing --> Camera")

    @QtCore.pyqtSlot(str)
    def change_topic(self, topic):
        self.footage_socket.setsockopt_string(zmq.UNSUBSCRIBE, self.FILTER)
        self.footage_socket.setsockopt_string(zmq.SUBSCRIBE, topic)
        print("Changeing topic")
        self.FILTER = topic

    @QtCore.pyqtSlot(bool)
    def activate(self, power_on):
        self.active = power_on
        self.footage_socket.connect(f'tcp://{self.IP}:{self.PORT}')
        self.running = True

    def run(self):
        while self.threadactive:
            topic, frame = self.footage_socket.recv_multipart()
            img = base64.b64decode(frame)
            npimg = np.fromstring(img, dtype=np.uint8)
            source = cv2.imdecode(npimg, 1)
            if self.active:
                h, w, ch = source.shape
                bytes_per_line = ch * w
                convert_to_Qt_format = QtGui.QImage(
                    source.data, w, h, bytes_per_line, QtGui.QImage.Format_RGB888)
                p = convert_to_Qt_format.scaled(1920, 1010)
                self.change_pixmap.emit(p)
                self.blink_signal.emit()

class PoseThread(QtCore.QThread):

    power_on = QtCore.pyqtSignal()

    IP = '10.10.10.243'
    PORT = 5555

    context = zmq.Context()
    pose_socket = context.socket(zmq.SUB)
    # pose_socket.connect('tcp://10.0.0.121:5555')
    pose_socket.setsockopt_string(zmq.SUBSCRIBE, 'pose')
    pose_socket.setsockopt(zmq.CONFLATE, 1)  # last msg only.

    active = False
    threadactive = True

    @QtCore.pyqtSlot()
    def close_socket(self):
        self.threadactive = False
        self.pose_socket.disconnect(f'tcp://{self.IP}:{self.PORT}')
        print("Closing --> Pose")

    @QtCore.pyqtSlot(bool)
    def activate(self, power_on):
        self.active = power_on
        self.pose_socket.connect(f'tcp://{self.IP}:{self.PORT}')

    def run(self):
        while self.threadactive:
            if self.active:
                pose = self.pose_socket.recv_string()
                print(pose)

class SerialThread(QtCore.QThread):

    power_on = QtCore.pyqtSignal()

    IP = '10.10.10.243'
    PORT = 6000

    context = zmq.Context()
    serial_socket = context.socket(zmq.SUB)
    # serial_socket.connect(f'tcp://{IP}:{PORT}')
    serial_socket.setsockopt_string(zmq.SUBSCRIBE, str(''))
    serial_socket.setsockopt(zmq.CONFLATE, 1)  # last msg only.

    active = False
    threadactive = True

    @QtCore.pyqtSlot()
    def close_socket(self):
        self.threadactive = False
        self.serial_socket.disconnect(f'tcp://{self.IP}:{self.PORT}')
        print("Closing --> Serial")

    @QtCore.pyqtSlot(bool)
    def activate(self, power_on):
        self.active = power_on
        self.serial_socket.connect(f'tcp://{self.IP}:{self.PORT}')

    def run(self):
        while self.threadactive:
            if self.active:
                data = self.serial_socket.recv_string()
                print(data)

class CommandThread(QtCore.QThread):

    power_on = QtCore.pyqtSignal()

    context = zmq.Context()
    command_socket = context.socket(zmq.PUB)
    command_socket.bind('tcp://*:5580')

    command = None
    old_command = None

    active = False
    threadactive = True

    @QtCore.pyqtSlot()
    def close_socket(self):
        self.threadactive = False
        print("Closing --> Command")

    @QtCore.pyqtSlot(bool)
    def activate(self, power_on):
        self.active = power_on

    @QtCore.pyqtSlot(str)
    def set_command(self, command):
        self.command = command

    def run(self):
        while self.threadactive:
            if self.active:
                if self.command is not None and self.command != self.old_command:
                    self.command_socket.send_multipart(
                        [b'command', self.command.encode()])
                    print(self.command)
                    self.old_command = self.command

class XboxControllerThread(QtCore.QThread):

    power_on = QtCore.pyqtSignal()

    PORT = 5590

    context = zmq.Context()
    xbox_socket = context.socket(zmq.PUB)
    xbox_socket.bind(f'tcp://*:{PORT}')

    xbox_controller = XboxController(
        None, deadzone=10, scale=100, invertYAxis=True)
    xbox_controller.setupControlCallback(
        xbox_controller.XboxControls.LTHUMBX, None)
    xbox_controller.setupControlCallback(
        xbox_controller.XboxControls.LTHUMBY, None)

    threadactive = True
    active = False

    @QtCore.pyqtSlot()
    def close_socket(self):
        self.threadactive = False
        print("Closing --> Xbox controller")

    @QtCore.pyqtSlot(bool)
    def activate(self, power_on):
        self.active = power_on
        print("Active being called")

    def run(self):
        while self.threadactive:
            if self.active:
                if self.xbox_controller.check_events():
                    self.xbox_socket.send_multipart([b'xbox_controller', base64.b64encode(
                        str(self.xbox_controller.controlValues).encode())])
                    print(self.xbox_controller.controlValues)
                time.sleep(0.1)

class WaypointThread(QtCore.QThread):

    power_on = QtCore.pyqtSignal()

    PORT = 5592

    context = zmq.Context()
    waypoint_socket = context.socket(zmq.PUB)
    waypoint_socket.bind(f'tcp://*:{PORT}')

    path = Path()

    threadactive = True
    active = False

    @QtCore.pyqtSlot()
    def close_socket(self):
        self.threadactive = False
        print("Closing --> waypoint thread")

    @QtCore.pyqtSlot(bool)
    def activate(self, power_on):
        self.active = power_on
        print("Active being called")

    def run(self):
        while self.threadactive:
            if self.active:
                if self.waypoint_socket.check_events():
                    self.waypoint_socket.send_multipart([b'waypoint', base64.b64encode(
                        str(self.xbox_controller.controlValues).encode())])
                    print(self.xbox_controller.controlValues)
                time.sleep(0.1)

max_stars_repo_name: magnusoy/Sparkie

max_issues_repo_name: magnusoy/Sparkie

max_forks_repo_name: magnusoy/Sparkie


text:
import process_data
import pandas as pd
import random
import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
import math
from pathlib import Path
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from collections import deque
import env.Strategies as Strategies

class OhlcvEnv(gym.Env):

    def __init__(self, window_size, path, train=True, show_trade=True):
        self.maxTrade = 100.0
        self.train= train
        self.show_trade = show_trade
        self.holdFactor = 1.0
        self.path = path
        #self.actions = ["LONG", "SHORT", "FLAT"]
        self.n_strategies = len(Strategies.strategies)
        self.actions = dict(min_value=0.0, max_value=self.maxTrade, type="float", shape=(self.n_strategies,))
        self.fee = 0.0005
        self.seed()
        self.file_list = []
        # load_csv
        self.load_from_csv()

        # n_features
        self.window_size = window_size
        self.n_features = self.df.shape[1]
        self.shape = (self.window_size, self.n_features)

        # defines action space
        #self.action_space = spaces.Discrete(len(self.actions))
        self.action_space = spaces.Box(low=0, high=int(self.maxTrade), shape=(self.n_strategies,), dtype=np.float32)
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=self.shape, dtype=np.float32)

    def load_from_csv(self):
        if(len(self.file_list) == 0):
            self.file_list = [x.name for x in Path(self.path).iterdir() if x.is_file()]
            self.file_list.sort()
        self.rand_episode = self.file_list.pop()
        raw_df= pd.read_csv(self.path + self.rand_episode)
        extractor = process_data.FeatureExtractor(raw_df)
        self.df = extractor.add_bar_features() # bar features o, h, l, c ---> C(4,2) = 4*3/2*1 = 6 features

        ## selected manual fetuares
        feature_list = [
            'bar_hc',
            'bar_ho',
            'bar_hl',
            'bar_cl',
            'bar_ol',
            'bar_co', 'close']
        self.df.dropna(inplace=True) # drops Nan rows
        self.closingPrices = self.df['close'].values
        self.df = self.df[feature_list].values

    def render(self, mode='human', verbose=False):
        return None

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def normalize_frame(self, frame):
        offline_scaler = StandardScaler()
        observe = frame[..., :-4]
        observe = offline_scaler.fit_transform(observe)
        agent_state = frame[..., -4:]
        temp = np.concatenate((observe, agent_state), axis=1)
        return temp

    def step(self, action):
        s, r, d, i = self._step(action)
        self.state_queue.append(s)
        return self.normalize_frame(np.concatenate(tuple(self.state_queue))), r, d, i

    def _step(self, action):

        if self.done:
            return self.state, self.reward, self.done, {}
        self.reward = 0

        buy_value = [ w*st['function'](self.closingPrices[:self.current_tick],*st['args'],**st['kwargs']) for st, w in zip(Strategies.strategies, action)]

        # two passes: sell first, then buy; might be naive in real-world settings
        v = sum(buy_value)

        if v>3.0:
            # buy
            self.holdFactor = 1.0
            v = int(v)
            if v * self.closingPrice < self.cash_in_hand:
                pass
            else:
                v = self.cash_in_hand // self.closingPrice
            self.cash_in_hand -= int(self.closingPrice * v * (1.0 + self.fee))
            self.stock_owned += v

        elif v<-3.0:
            # sell
            self.holdFactor = 1.0
            v = int(abs(v))
            if self.stock_owned > v:
                pass
            else:
                v = self.stock_owned
            self.cash_in_hand += int(self.closingPrice * v * (1.0 - self.fee))
            self.stock_owned -= v

        else:
            # hold
            self.holdFactor *= 1.01
            pass

        temp_portfolio = self.cash_in_hand + self.stock_owned * self.closingPrice
        self.portfolio = temp_portfolio
        self.reward += (temp_portfolio - self.init_cash) * self.holdFactor
        self.current_tick += 1

        if(self.show_trade and self.current_tick%100 == 0):
             print("Tick: {0}/ Portfolio (krw-won): {1}".format(self.current_tick, self.portfolio))

        self.history.append((self.action, self.current_tick, self.closingPrice, self.portfolio, self.reward))
        self.state = self.updateState()
        info = {'portfolio':np.array([self.portfolio]), "history":self.history}
        if (self.current_tick > (self.df.shape[0]) - self.window_size-1):
            self.done = True
            if(self.train == False):
                np.array([info]).dump('info/ppo_{0}_LS.info'.format(self.portfolio))

        return self.state, self.reward, self.done, info

    def reset(self):
        # hoho
        self.stock_owned = 0
        self.holdFactor = 1.0

        # self.current_tick = random.randint(0, self.df.shape[0]-800)
        if(self.train):
            self.current_tick = random.randint(0, self.df.shape[0] - 800)
        else:
            self.current_tick = 0

        print("start episode ... {0} at {1}" .format(self.rand_episode, self.current_tick))

        # clear internal variables
        self.history = [] # keep buy, sell, hold action history
        self.init_cash = 100 * 10000
        self.cash_in_hand = self.init_cash # initial balance, u can change it to whatever u like
        self.portfolio = float(self.cash_in_hand) # (coin * current_price + current_krw_balance) == portfolio
        self.closingPrice = self.closingPrices[self.current_tick]

        self.action = np.zeros(len(Strategies.strategies))
        self.done = False

        self.state_queue = deque(maxlen=self.window_size)
        self.state = self.preheat_queue()
        return self.state

    def preheat_queue(self):
        while(len(self.state_queue) < self.window_size):
            # rand_action = random.randint(0, len(self.actions)-1)
            # rand_action = 2
            rand_action = np.random.rand(len(Strategies.strategies))*self.maxTrade
            s, r, d, i= self._step(rand_action)
            self.state_queue.append(s)
        return self.normalize_frame(np.concatenate(tuple(self.state_queue)))

    def updateState(self):
        self.closingPrice = float(self.closingPrices[self.current_tick])
        state = self.df[self.current_tick]
        return state.reshape(1,-1)

max_stars_repo_name: Rinuys/RLTradingBot

max_issues_repo_name: Rinuys/RLTradingBot

max_forks_repo_name: Rinuys/RLTradingBot


text:
import numpy as np

def D3_element(fs,a,b,c,i,j,k):
    tmp = fs[i,i] + fs[j,j] + fs[k,k]\
    - fs[a,a] - fs[b,b] - fs[c,c]
    return tmp

def t3d_element(fs,ts,spinints,a,b,c,i,j,k):
    tmp = 0
    tmp += fun_d(ts,spinints,a,b,c,i,j,k)
    tmp -= fun_d(ts,spinints,b,a,c,i,j,k)
    tmp -= fun_d(ts,spinints,c,b,a,i,j,k)
    tmp -= fun_d(ts,spinints,a,b,c,j,i,k)
    tmp += fun_d(ts,spinints,b,a,c,j,i,k)
    tmp += fun_d(ts,spinints,c,b,a,j,i,k)
    tmp -= fun_d(ts,spinints,a,b,c,k,j,i)
    tmp += fun_d(ts,spinints,b,a,c,k,j,i)
    tmp += fun_d(ts,spinints,c,b,a,k,j,i)
    tmp /= D3_element(fs,a,b,c,i,j,k)
    return tmp

def t3c_element(Nelec,dim,fs,td,spinints,a,b,c,i,j,k):
    tmp = 0
    tmp += fun_c(Nelec,dim,td,spinints,a,b,c,i,j,k)
    tmp -= fun_c(Nelec,dim,td,spinints,b,a,c,i,j,k)
    tmp -= fun_c(Nelec,dim,td,spinints,c,b,a,i,j,k)
    tmp -= fun_c(Nelec,dim,td,spinints,a,b,c,j,i,k)
    tmp += fun_c(Nelec,dim,td,spinints,b,a,c,j,i,k)
    tmp += fun_c(Nelec,dim,td,spinints,c,b,a,j,i,k)
    tmp -= fun_c(Nelec,dim,td,spinints,a,b,c,k,j,i)
    tmp += fun_c(Nelec,dim,td,spinints,b,a,c,k,j,i)
    tmp += fun_c(Nelec,dim,td,spinints,c,b,a,k,j,i)
    tmp /= D3_element(fs,a,b,c,i,j,k)
    return tmp

def fun_d (ts,spinints,a,b,c,i,j,k):
    tmp = ts[a,i]*spinints[j,k,b,c]
    return tmp

def fun_c (Nelec,dim,td,spinints,a,b,c,i,j,k):
    tmp = 0
    for e in range(Nelec,dim):
        tmp += td[a,e,j,k]*spinints[e,i,b,c]
    for m in range(0,Nelec):
        tmp -= td[b,c,i,m]*spinints[m,a,j,k]
    return tmp
################################
def spinfock(eorbitals):
    """
    """
    if type(eorbitals) is np.ndarray:
        dim = 2*len(eorbitals)
        fs = np.zeros(dim)
        for i in range(0,dim):
            fs[i] = eorbitals[i//2]
        fs = np.diag(fs) # put MO energies in diagonal array
    elif type(eorbitals) is dict:
        dim = 2*len(eorbitals['alpha'])
        fs = np.zeros(dim)
        for i in range(0,dim):
            if i%2==0:
                fs[i] = eorbitals['alpha'][i//2]
            elif i%2==0:
                fs[i] = eorbitals['beta'][i//2]
        fs = np.diag(fs) # put MO energies in diagonal array
    return fs

def initialize(fs,spinints,Nelec,dim):
    """
    Init empty T1 (ts) and T2 (td) arrays
    and make denominator arrays Dai, Dabij
    """
    # Initial guess for T1 and T2
    ts = np.zeros((dim,dim))
    td = np.zeros((dim,dim,dim,dim))
    for a in range(Nelec,dim):
        for b in range(Nelec,dim):
            for i in range(0,Nelec):
                for j in range(0,Nelec):
                    td[a,b,i,j] += spinints[i,j,a,b]/(fs[i,i] + fs[j,j] - fs[a,a] - fs[b,b])

    # Equation (12) of Stanton
    Dai = np.zeros((dim,dim))
    for a in range(Nelec,dim):
        for i in range(0,Nelec):
            Dai[a,i] = fs[i,i] - fs[a,a]

    # Stanton eq (13)
    Dabij = np.zeros((dim,dim,dim,dim))
    for a in range(Nelec,dim):
        for b in range(Nelec,dim):
            for i in range(0,Nelec):
                for j in range(0,Nelec):
                    Dabij[a,b,i,j] = fs[i,i] + fs[j,j] - fs[a,a] - fs[b,b]
    return ts,td,Dai,Dabij

# Stanton eq (9)
def taus(ts,td,a,b,i,j):
  taus = td[a,b,i,j] + 0.5*(ts[a,i]*ts[b,j] - ts[b,i]*ts[a,j])
  return taus

# Stanton eq (10)
def tau(ts,td,a,b,i,j):
  tau = td[a,b,i,j] + ts[a,i]*ts[b,j] - ts[b,i]*ts[a,j]
  return tau

# We need to update our intermediates at the beginning, and
# at the end of each iteration. Each iteration provides a new
# guess at the amplitudes T1 (ts) and T2 (td), that *hopefully*
# converges to a stable, ground-state, solution.

def updateintermediates(x,Nelec,dim,fs,spinints,ts,td):
  if x == True:
    # Stanton eq (3)
    Fae = np.zeros((dim,dim))
    for a in range(Nelec,dim):
      for e in range(Nelec,dim):
        Fae[a,e] = (1 - (a == e))*fs[a,e]
        for m in range(0,Nelec):
          Fae[a,e] += -0.5*fs[m,e]*ts[a,m]
          for f in range(Nelec,dim):
            Fae[a,e] += ts[f,m]*spinints[m,a,f,e]
            for n in range(0,Nelec):
              Fae[a,e] += -0.5*taus(ts,td,a,f,m,n)*spinints[m,n,e,f]

    # Stanton eq (4)
    Fmi = np.zeros((dim,dim))
    for m in range(0,Nelec):
      for i in range(0,Nelec):
        Fmi[m,i] = (1 - (m == i))*fs[m,i]
        for e in range(Nelec,dim):
          Fmi[m,i] += 0.5*ts[e,i]*fs[m,e]
          for n in range(0,Nelec):
            Fmi[m,i] += ts[e,n]*spinints[m,n,i,e]
            for f in range(Nelec,dim):
              Fmi[m,i] += 0.5*taus(ts,td,e,f,i,n)*spinints[m,n,e,f]

    # Stanton eq (5)
    Fme = np.zeros((dim,dim))
    for m in range(0,Nelec):
      for e in range(Nelec,dim):
        Fme[m,e] = fs[m,e]
        for n in range(0,Nelec):
          for f in range(Nelec,dim):
            Fme[m,e] += ts[f,n]*spinints[m,n,e,f]

    # Stanton eq (6)
    Wmnij = np.zeros((dim,dim,dim,dim))
    for m in range(0,Nelec):
      for n in range(0,Nelec):
        for i in range(0,Nelec):
          for j in range(0,Nelec):
            Wmnij[m,n,i,j] = spinints[m,n,i,j]
            for e in range(Nelec,dim):
              Wmnij[m,n,i,j] += ts[e,j]*spinints[m,n,i,e] - ts[e,i]*spinints[m,n,j,e]
              for f in range(Nelec,dim):
                Wmnij[m,n,i,j] += 0.25*tau(ts,td,e,f,i,j)*spinints[m,n,e,f]

    # Stanton eq (7)
    Wabef = np.zeros((dim,dim,dim,dim))
    for a in range(Nelec,dim):
      for b in range(Nelec,dim):
        for e in range(Nelec,dim):
          for f in range(Nelec,dim):
            Wabef[a,b,e,f] = spinints[a,b,e,f]
            for m in range(0,Nelec):
              Wabef[a,b,e,f] += -ts[b,m]*spinints[a,m,e,f] + ts[a,m]*spinints[b,m,e,f]
              for n in range(0,Nelec):
                Wabef[a,b,e,f] += 0.25*tau(ts,td,a,b,m,n)*spinints[m,n,e,f]

    # Stanton eq (8)
    Wmbej = np.zeros((dim,dim,dim,dim))
    for m in range(0,Nelec):
      for b in range(Nelec,dim):
        for e in range(Nelec,dim):
          for j in range(0,Nelec):
            Wmbej[m,b,e,j] = spinints[m,b,e,j]
            for f in range(Nelec,dim):
              Wmbej[m,b,e,j] += ts[f,j]*spinints[m,b,e,f]
            for n in range(0,Nelec):
              Wmbej[m,b,e,j] += -ts[b,n]*spinints[m,n,e,j]
              for f in range(Nelec,dim):
                Wmbej[m,b,e,j] += -(0.5*td[f,b,j,n] + ts[f,j]*ts[b,n])*spinints[m,n,e,f]

    return Fae, Fmi, Fme, Wmnij, Wabef, Wmbej

# makeT1 and makeT2, as they imply, construct the actual amplitudes necessary for computing
# the CCSD energy (or computing an EOM-CCSD Hamiltonian, etc)

# Stanton eq (1)
def makeT1(x,Nelec,dim,fs,spinints,ts,td,Dai,Fae,Fmi,Fme):
  if x == True:
    tsnew = np.zeros((dim,dim))
    for a in range(Nelec,dim):
      for i in range(0,Nelec):
        tsnew[a,i] = fs[i,a]
        for e in range(Nelec,dim):
          tsnew[a,i] += ts[e,i]*Fae[a,e]
        for m in range(0,Nelec):
          tsnew[a,i] += -ts[a,m]*Fmi[m,i]
          for e in range(Nelec,dim):
            tsnew[a,i] += td[a,e,i,m]*Fme[m,e]
            for f in range(Nelec,dim):
              tsnew[a,i] += -0.5*td[e,f,i,m]*spinints[m,a,e,f]
            for n in range(0,Nelec):
              tsnew[a,i] += -0.5*td[a,e,m,n]*spinints[n,m,e,i]
        for n in range(0,Nelec):
          for f in range(Nelec,dim):
            tsnew[a,i] += -ts[f,n]*spinints[n,a,i,f]
        tsnew[a,i] = tsnew[a,i]/Dai[a,i]
  return tsnew

# Stanton eq (2)
def makeT2(x,Nelec,dim,fs,spinints,ts,td,Dabij,Fae,Fmi,Fme,Wmnij,Wabef,Wmbej):
  if x == True:
    tdnew = np.zeros((dim,dim,dim,dim))
    for a in range(Nelec,dim):
      for b in range(Nelec,dim):
        for i in range(0,Nelec):
          for j in range(0,Nelec):
            tdnew[a,b,i,j] += spinints[i,j,a,b]
            for e in range(Nelec,dim):
              tdnew[a,b,i,j] += td[a,e,i,j]*Fae[b,e] - td[b,e,i,j]*Fae[a,e]
              for m in range(0,Nelec):
                tdnew[a,b,i,j] += -0.5*td[a,e,i,j]*ts[b,m]*Fme[m,e] + 0.5*td[a,e,i,j]*ts[a,m]*Fme[m,e]
                continue
            for m in range(0,Nelec):
              tdnew[a,b,i,j] += -td[a,b,i,m]*Fmi[m,j] + td[a,b,j,m]*Fmi[m,i]
              for e in range(Nelec,dim):
                tdnew[a,b,i,j] += -0.5*td[a,b,i,m]*ts[e,j]*Fme[m,e] + 0.5*td[a,b,i,m]*ts[e,i]*Fme[m,e]
                continue
            for e in range(Nelec,dim):
              tdnew[a,b,i,j] += ts[e,i]*spinints[a,b,e,j] - ts[e,j]*spinints[a,b,e,i]
              for f in range(Nelec,dim):
                tdnew[a,b,i,j] += 0.5*tau(ts,td,e,f,i,j)*Wabef[a,b,e,f]
                continue
            for m in range(0,Nelec):
              tdnew[a,b,i,j] += -ts[a,m]*spinints[m,b,i,j] + ts[b,m]*spinints[m,a,i,j]
              for e in range(Nelec,dim):
                tdnew[a,b,i,j] += td[a,e,i,m]*Wmbej[m,b,e,j] - ts[e,i]*ts[a,m]*spinints[m,b,e,j]
                tdnew[a,b,i,j] += -td[a,e,j,m]*Wmbej[m,b,e,i] + ts[e,j]*ts[a,m]*spinints[m,b,e,i]
                tdnew[a,b,i,j] += -td[b,e,i,m]*Wmbej[m,a,e,j] - ts[e,i]*ts[b,m]*spinints[m,a,e,j]
                tdnew[a,b,i,j] += td[b,e,j,m]*Wmbej[m,a,e,i] - ts[e,j]*ts[b,m]*spinints[m,a,e,i]
                continue
              for n in range(0,Nelec):
                tdnew[a,b,i,j] += 0.5*tau(ts,td,a,b,m,n)*Wmnij[m,n,i,j]
                continue
            tdnew[a,b,i,j] = tdnew[a,b,i,j]/Dabij[a,b,i,j]
    return tdnew

# Expression from Crawford, Schaefer (2000)
# DOI: 10.1002/9780470125915.ch2
# Equation (134) and (173)
# computes CCSD energy given T1 and T2
def ccsdenergy(Nelec,dim,fs,spinints,ts,td):
  ECCSD = 0.0
  for i in range(0,Nelec):
    for a in range(Nelec,dim):
      ECCSD += fs[i,a]*ts[a,i]
      for j in range(0,Nelec):
        for b in range(Nelec,dim):
          ECCSD += 0.25*spinints[i,j,a,b]*td[a,b,i,j] + 0.5*spinints[i,j,a,b]*(ts[a,i])*(ts[b,j])
  return ECCSD

def E_triples_correction(Nelec,dim,fs,ts,td,spinints):
  Et = 0.0
  for i in range(0,Nelec):
    for j in range(0,Nelec):
        for k in range(0,Nelec):
            for a in range(Nelec,dim):
                for b in range(Nelec,dim):
                    for c in range(Nelec,dim):
                        Et += t3c_element(Nelec,dim,fs,td,spinints,a,b,c,i,j,k)*D3_element(fs,a,b,c,i,j,k)\
                        *(t3c_element(Nelec,dim,fs,td,spinints,a,b,c,i,j,k)+t3d_element(fs,ts,spinints,a,b,c,i,j,k))
  Et *= 1/36
  return Et

max_stars_repo_name: ZhaoYilin/moha

max_issues_repo_name: ZhaoYilin/moha

max_forks_repo_name: ZhaoYilin/moha