text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#   wwz.py
#   6 June 15:17
#   11 September 12:16 semicolon fixes
#   Copyright 2014 M. Emre Aydin <emre.m.aydin@gmail.com>
#   http://about.me/emre.aydin
#   http://github.com/eaydin

"""
    A Weighted Wavelet Z-Transformation Application for Python.
    Translated by M. Emre Aydin - emre.m.aydin@gmail.com
    http://about.me/emre.aydin
    Available at http://github.com/eaydin

For details, read the docstrings of WWZ class.
"""

import math
import sys
import os
from datetime import datetime

# Checking if Python is 2.7
if not sys.version_info[:2] == (2, 7):
    print "You need Python 2.7 for this script to run."
    raise SystemExit

# Checking if we have numpy
try:
    import numpy
except:
    print "Please Install Numpy"
    raise SystemExit

# Checking if we have argparse (which is default in 2.7 so this is
# probably a useless check since we already check if Python is 2.7)
try:
    import argparse
except:
    print "Argparse Not Available?! You sure this is Python 2.7?"
    raise SystemExit

# The WWZ Class Begins

class WWZ(object):
    """The Main class object.
    This object class does not get any arguments.
    Available methods are:
        readfile()
        roundtau()
        maketau()
        makefreq()
        matrix_inv()
        wwt()
        writefile()
        writegnu()

    Arguments:
        fileName: the input filename, should be the lightcurve.
        outputfileName: the output filename.
        flo: the low frequency value. Float.
        fhi: the high frequency value. Float.
        df: the frequency step. Float.
        dcon: the C Window constant. Float.
        timedivisions: the The Divisions value, choose 50.0 if not
            sure, that is the default used by Templeton.
        max_periods: set True if you want the second output file.
            It outputs the tau's with maximum period values for easier
            estimation.
        gnuplot_compatible: splits tau values by a blank line if set
            True, so that pm3d of gnuplot easily maps the plot.
    """

    def __init__(self):
        """Initializing the object"""

    def readfile(self, fileName):
        """Read the input file.
        The argument is the file pointer, not the filename as a string.
        The values in file should be delimited with spaces or tabs.
        Ignores lines starting with # and %, as if they're comment lines.

        Returns two arrays:
            Time value, read from the first column of input file.
            Magnitude value, read from the second column of input file.
        """

        time = []
        magnitude = []

        for line in fileName:
            # Check if it's a comment line
            if line.strip()[0] != "%" and line.strip()[0] != "#":
                line_time = float(line.split()[0])
                line_mag = float(line.split()[1])
                time.append(line_time)
                magnitude.append(line_mag)

        fileName.close()

        # Just a routine check for parameter number equality
        # This should be cleaned up a bit
        if len(time) != len(magnitude):
            print "Number of Time and Magnitude input do not match. \
                  Please check the input file."
            raise SystemExit

        # Return two arrays
        return time, magnitude

    def roundtau(self, darg):
        """Rounds the tau's. from G. Foster's Code.
        This is actually called by the maketau method.
        The input is dtspan/timedivisions,
        where dtspan is the entire timespan of the lightcurve.
        so dtspan = time[-1] - time[0]

        Returns the round value.
        """

        dex = math.log(darg, 10)
        nex = int(dex)

        darg = darg / math.pow(10, nex)

        if darg >= 5:
            darg = 5.0
        elif darg >= 2:
            darg = 2.0
        else:
            darg = 1.0

        darg = darg * math.pow(10, nex)
        return darg

    def matrix_inv(self, inputMatrix):
        """The Matrix Inversion Method.
        Arguments are:
            inputMatrix: the input matrix_inv

        Returns the inverted matrix.
        """
        # Lines 202 - 252 Fortran

        ndim = 2
        dsol = numpy.zeros(shape=(3, 3))

        for i in range(0, 3):
            for j in range(0, 3):
                dsol[i][j] = 0.0
            dsol[i][i] = 1.0

        for i in range(0, 3):
            if inputMatrix[i][i] == 0.0:
                if i == ndim:
                    return
                for j in range(0, 3):
                    if inputMatrix[j][i] != 0.0:
                        for k in range(0, 3):
                            inputMatrix[i][k] = inputMatrix[i][k] + \
                                                inputMatrix[j][k]
                            dsol[i][j] = dsol[i][j] + dsol[j][k]

            dfac = inputMatrix[i][i]

            for j in range(0, 3):
                inputMatrix[i][j] = inputMatrix[i][j] / dfac
                dsol[i][j] = dsol[i][j] / dfac

            for j in range(0, 3):
                if j != i:
                    dfac = inputMatrix[j][i]
                    for k in range(0, 3):
                        inputMatrix[j][k] = inputMatrix[j][k] - \
                                          (inputMatrix[i][k] * dfac)
                        dsol[j][k] = dsol[j][k] - (dsol[i][k] * dfac)

        # The unnecessery loop
        #for i in range(0, 3):
            #for j in range(0, 3):
                #self.dmat[i][j] = dsol[i][j]

        return dsol

    def maketau(self, time, timedivisions):
        """The maketau method.
        Arguments are:
            time = The array of time values
            timedivisions = The value of timedivisions to create tau values

            Returns an array of calculated tau values.
        """
        # The Maketau section
        # Lines 90 - 122 Fortran

        dtauhi = time[-1]
        dtaulo = time[0]
        # the java translation uses 1 for dtaulo,
        # but that should be a bug.

        dtspan = dtauhi - dtaulo
        dtstep = self.roundtau(dtspan / timedivisions)

        dtaulo = dtstep * int(dtaulo / dtstep)
        dtauhi = dtstep * int((dtauhi / dtstep) + 0.5)

        tau = []
        dtau = dtaulo
        while dtau <= dtauhi:
            tau.append(dtau)
            dtau = dtau + dtstep
        return tau

        ### End of Maketau

    def makefreq(self, flo, fhi, df):
        """The Makefreq section.
        Arguments are:
            flo = Low Frequency
            fhi = High Frequency
            df = Frequency Step
        """
        # Lines 149 - 181 Fortran
        # Lines 356 - 370 Java
        freq = []
        freq.append(flo)
        nfreq = int((fhi - flo) / df) + 1

        # These lines seem skeptical!
        for i in range(1,nfreq+1):
            freq.append(flo + ((i - 1) * df))

        return freq

        ### End of Makefreq

    def wwt(self, time, magnitude, flo, fhi, df, dcon, timedivisions):
        """The WWZ Algorithm
        Arguments are:
            time = The time values as an array
            magnitude = The magnitude values as an array
            flo = The Low Frequency
            fhi = The High Frequency
            df = The Frequency Step
            dcon = The C constant of WWZ Window
            timedivisions = The TAU steps

            Returns a NumPy Array
        """

        dave = numpy.mean(magnitude)
        dvar = numpy.var(magnitude)

        freq = self.makefreq(flo, fhi, df)
        nfreq = len(freq)
        dmat = numpy.zeros(shape=(3,3))

        ### End of Initializing

        tau = self.maketau(time, timedivisions)
        ntau = len(tau)

        ### WWT Stars Here

        dvec = [0,0,0] # length is 3
        dcoef = [0,0,0] # length is 3

        itau = 0
        ifreq = 0
        idat = 0

        domega = 0.0
        dweight2 = 0.0
        dz = 0.0
        dweight = 0.0

        dcc = 0.0
        dcw = 0.0
        dss = 0.0
        dsw = 0.0
        dxw = 0.0
        dvarw = 0.0

        dtau = 0.0

        dpower = 0.0
        dpowz = 0.0
        damp = 0.0
        dneff = 0.0
        davew = 0.0

        dfre = 0.0

        n1 = 0
        n2 = 0

        dmz = 0.0
        dmzfre = 0.0
        dmzamp = 0.0
        dmcon = 0.0
        dmneff = 0.0

        twopi = 2.0 * math.pi

        ndim = 2
        itau1 = 0 # --> 1 or 0 ??
        itau2 = ntau # --> ???
        ifreq1 = 1
        ifreq2 = nfreq
        nstart = 1

        # Creating output arrays
        output = numpy.empty((ntau*(nfreq-1), 6))
        numdat=len(time)
        index = 0

        # Use for itau in range(itau1,itau2) for parallel
        for itau in range(0, itau2):
            nstart = 1
            dtau = tau[itau]

            dmfre = 0.0
            dmamp = 0.0
            dmcon = 0.0
            dmneff = 0.0
            dmz = -1.0 # less than the smallest WWZ

            for ifreq in range(ifreq1, ifreq2):
                dfre = freq[ifreq]
                domega = dfre * twopi

                for i in range(0, ndim + 1):
                    dvec[i] = 0.0
                    for j in range(0, ndim + 1):
                        dmat[i][j] = 0.0

                dweight2 = 0.0

                for idat in range(nstart, numdat):

                    dz = domega * (time[idat] - dtau)
                    dweight = math.exp(-1.0 * dcon * dz * dz)

                    if (dweight > 10**(-9)):
                        dcc = math.cos(dz)
                        dcw = dweight * dcc
                        dss = math.sin(dz)
                        dsw = dweight * dss
                        dmat[0][0] = dmat[0][0] + dweight
                        dweight2 = dweight2 + (dweight**2)
                        dmat[0][1] = dmat[0][1] + dcw
                        dmat[0][2] = dmat[0][2] + dsw
                        dmat[1][1] = dmat[1][1] + (dcw * dcc)
                        dmat[1][2] = dmat[1][2] + (dcw * dss)
                        dmat[2][2] = dmat[2][2] + (dsw * dss)

                        dxw = dweight * magnitude[idat]
                        dvec[0] = dvec[0] + dxw
                        dvarw = dvarw + (dxw * magnitude[idat])
                        dvec[1] = dvec[1] + (dcw * magnitude[idat])
                        dvec[2] = dvec[2] + (dsw * magnitude[idat])

                    elif dz > 0.0:
                        break
                    else:
                        nstart = idat + 1

                dpower = 0.0
                damp = 0.0

                for n1 in range(0, ndim + 1):
                    dcoef[n1] = 0.0

                if (dweight2 > 0.0):
                    dneff = (dmat[0][0] * dmat[0][0]) / dweight2
                else:
                    dneff = 0.0

                if (dneff > 3.0):

                    for n1 in range(0, ndim + 1):
                        dvec[n1] = dvec[n1] / dmat[0][0]

                        for n2 in range(1, ndim + 1 ):
                            dmat[n1][n2] = dmat[n1][n2] / \
                                                dmat[0][0]

                    if (dmat[0][0] > 0.005):
                        dvarw = dvarw / dmat[0][0]
                    else:
                        dvarw = 0.0

                    dmat[0][0] = 1.0
                    davew = dvec[0]
                    dvarw = dvarw - (davew ** 2)

                    if (dvarw <= 0.0):
                        dvarw = 10**-12

                    for n1 in range(1, ndim + 1):
                        for n2 in range(0, n1):
                            dmat[n1][n2] = dmat[n2][n1]

                    dmat = self.matrix_inv(dmat)

                    for n1 in range(0, ndim + 1):
                        for n2 in range(0, ndim + 1):
                            dcoef[n1] = dcoef[n1] + dmat[n1][n2] * \
                                                    dvec[n2]

                        dpower = dpower + (dcoef[n1] * dvec[n1])

                    dpower = dpower - (davew ** 2)
                    dpowz = (dneff - 3.0) * dpower / (dvarw - dpower) / 2.0
                    dpower = (dneff - 1.0) * dpower / dvarw / 2.0
                    damp = math.sqrt(dcoef[1] * dcoef[1] + \
                                     dcoef[2] * dcoef[2])

                else:

                    dpowz = 0.0
                    dpower = 0.0
                    damp = 0.0

                    if (dneff < (10**(-9))):
                        dneff = 0.0

                if (damp < (10**(-9))):
                    damp = 0.0

                if (dpower < (10**(-9))):
                    dpower = 0.0

                if (dpowz < (10**(-9))):
                    dpowz = 0.0

                # Let's write everything out.
                output[index] = [dtau, dfre, dpowz, damp, dcoef[0], dneff]

                index = index + 1

                if (dpowz > dmz):
                    dmz = dpowz
                    dmfre = dfre
                    dmamp = damp
                    dmcon = dcoef[0]
                    dmneff = dneff

        return output

    def writefile(self, wwz_output, outputFile, no_headers, max_periods):
        """The write file method.
        Arguments are:
            wwz_output = The NumPy array of WWZ values to write
            outputFile = The output file pointer, not the filename
            no_headers = If true, will not write headers to the output
            max_periods = If true, will create a file with period values
                          of maximum WWZ statistics
        """

        numpy.set_printoptions(precision=5)
        numpy.set_printoptions(suppress=True)
        numpy.set_printoptions(threshold='nan')

        if no_headers:
            numpy.savetxt(outputFile, wwz_output, delimiter="\t", \
                          fmt="%10.4f")
        else:
            numpy.savetxt(outputFile, wwz_output, delimiter="\t", \
                        fmt="%10.4f", comments="#", \
                        header="%9s %10s %10s %10s %10s %10s" % \
                        ("TAU", "FREQ", "WWZ", "AMP", "COEF", "NEFF"))

    def writegnu(self, wwz_output, outputFile, no_headers, \
                max_periods, ntau):
        """The write file method, adapted to work with GnuPlot.
        Arguments are:
            wwz_output = The NumPy array of WWZ values to write
            outputFile = The output file pointer, not the filename
            no_headers = If true, will not write headers to the output
            max_periods = If true, will create a file with period values
                          of maximum WWZ statistics
            ntau = The number of tau values, this is needed in order to
                   split the wwz_output equally
                To calculate ntau, use the equation below:
                    len(wwz_output) /
                      int(((freq_high - freq_low) / freq_step) + 1)
        """

        numpy.set_printoptions(precision=5)
        numpy.set_printoptions(suppress=True)
        numpy.set_printoptions(threshold='nan')

        splitArray = numpy.vsplit(wwz_output, ntau)

        # check if the file is in append mode
        # if not, reopen it
        if outputFile.mode != 'a':
            outputFile.close()
            outputFile = open(outputFile.name, "a")

        # write headers if expected
        if not no_headers:
            header="%9s %10s %10s %10s %10s %10s" % \
                   ("TAU", "FREQ", "WWZ", "AMP", "COEF", "NEFF")
            outputFile.write("#" + header + "\n")

        # split the array and add newlines in between tau values,
        # write the output
        for i in range(0,ntau):
            numpy.savetxt(outputFile, splitArray[i], delimiter="\t", \
                          fmt="%10.4f")
            if i != ntau-1:
                outputFile.write("\n")

# The WWZPAR Class Begins

class WWZPAR(object):
    """This is for Parallel Processing only.
    This works different than the WWZ class, it takes arguments directly.

    It gets input as a filepointer, NOT as arrays!
    Arguments:
        fileName: the input filename, should be the lightcurve.
        outputfileName: the output filename.
        flo: the low frequency value. Float.
        fhi: the high frequency value. Float.
        df: the frequency step. Float.
        dcon: the C Window constant. Float.
        timedivisions: the The Divisions value, choose 50.0 if not sure,
                        that is the default used by Templeton.
        max_periods: set True if you want the second output file.
                      It outputs the tau's with maximum period values
                      for easier estimation.
        gnuplot_compatible: splits tau values by a blank line if set True,
                             so that pm3d of gnuplot easily maps the plot.
    """

    def __init__(self, fileName, outputfileName, flo, fhi, df, dcon, \
                timedivisions, max_periods, gnuplot_compatible):
        """Initializing the object"""

        self.inputfile = fileName

        self.outputfilename1 = outputfileName
        self.outputfilename2 = outputfileName.name + ".max_periods"

        self.max_periods = max_periods
        self.gnuplot_compatible = gnuplot_compatible

        self.timedivisions = timedivisions
        # This (50) is an assumption by Templeton.
        # VStars leaves this optional
        # but keeps the default value

        self.fhi = fhi
        self.flo = flo
        self.df = df

        # dcon is the Window Constant "c" in Foster's equations.
        self.dcon = dcon

        self.fileName = fileName

        self.time = [] # Input time, first column in the file
        self.magnitude = [] # Input magnitude, second column in the file

        self.dave = 0.0 # average
        self.dvar = 0.0 # variance

        self.nfreq = int((self.fhi - self.flo) / self.df) + 1

        self.freq = []

        self.dmat = numpy.zeros(shape=(3, 3))

    def readfile(self):
        """Read the input file"""

        #read_file = open(self.inputfilename, "r")
        for line in self.inputfile:
            if line.strip()[0] != "%" and line.strip()[0] != "#":
                line_time = float(line.split()[0])
                line_mag = float(line.split()[1])
                self.time.append(line_time)
                self.magnitude.append(line_mag)
                self.dave = self.dave + line_mag
                self.dvar = self.dvar + (line_mag ** 2)

        self.inputfile.close()
        if len(self.time) != len(self.magnitude):
            print "Number of Time and Magnitude input do not match. \
                   Please check the input file."
            raise SystemExit
            # Just a routine check for parameter number equality.
            # This should be cleaned up a bit.

        # Calculating Header Values
        self.numdat = len(self.time)
        self.dave = self.dave / self.numdat
        self.dvar = (self.dvar / self.numdat) - (self.dave ** 2)
        self.dsig = math.sqrt((self.dvar * self.numdat) / (self.numdat - 1))

    def roundtau(self, darg):
        """Rounds the tau's. from G. Foster's Code."""

        dex = math.log(darg, 10)
        nex = int(dex)

        darg = darg / math.pow(10, nex)

        if darg >= 5:
            darg = 5.0
        elif darg >= 2:
            darg = 2.0
        else:
            darg = 1.0

        darg = darg * math.pow(10, nex)
        return darg

    def matrix_inv(self,input_matrix):
        """The Matrix Inversion Function"""
        # Lines 202 - 252 Fortran

        ndim = 2
        dsol = numpy.zeros(shape=(3, 3))

        for i in range(0, 3):
            for j in range(0, 3):
                dsol[i][j] = 0.0
            dsol[i][i] = 1.0

        for i in range(0, 3):
            if input_matrix[i][i] == 0.0:
                if i == ndim:
                    return
                for j in range(0, 3):
                    if input_matrix[j][i] != 0.0:
                        for k in range(0, 3):
                            input_matrix[i][k] = input_matrix[i][k] + \
                                                 input_matrix[j][k]
                            dsol[i][j] = dsol[i][j] + dsol[j][k]

            dfac = input_matrix[i][i]

            for j in range(0, 3):
                input_matrix[i][j] = input_matrix[i][j] / dfac
                dsol[i][j] = dsol[i][j] / dfac

            for j in range(0, 3):
                if j != i:
                    dfac = input_matrix[j][i]
                    for k in range(0, 3):
                        input_matrix[j][k] = input_matrix[j][k] - \
                                            (input_matrix[i][k] * dfac)
                        dsol[j][k] = dsol[j][k] - (dsol[i][k] * dfac)

        return dsol

    def maketau(self):
        """The Maketau section"""
        # Lines 90 - 122 Fortran

        dtaulo = self.time[0] # the java translation uses 1 for this,
        # but that should be a bug.
        dtauhi = self.time[-1]

        dtspan = dtauhi - dtaulo
        dtstep = self.roundtau(dtspan / self.timedivisions)

        dtaulo = dtstep * int(dtaulo / dtstep)
        dtauhi = dtstep * int((dtauhi / dtstep) + 0.5)

        self.tau = []

        dtau = dtaulo
        while dtau <= dtauhi:
            #print dtau
            self.tau.append(dtau)
            dtau = dtau + dtstep

        self.ntau = len(self.tau)

    def makefreq(self):
        """The Makefreq section"""
        # Lines 149 - 181 Fortran
        # Lines 356 - 370 Java

        self.freq.append(self.flo)

        # These lines seem skeptical!
        for i in range(1,self.nfreq+1):
            self.freq.append(self.flo + ((i - 1) * self.df))

    def wwt(self, output1_par, itau1_par, itau2_par):
        """The WWZ Algorithm in Parallel Mode"""

        output1 = open(output1_par, "w")
        max_periods = self.max_periods
        gnuplot_compatible = self.gnuplot_compatible

        if max_periods == True:
            output2_par = output1_par + ".max_periods.par"
            output2 = open(output2_par, "w")

        dvec = [0,0,0] # length is 3
        dcoef = [0,0,0] # length is 3

        itau = 0
        ifreq = 0
        idat = 0

        domega = 0.0
        dweight2 = 0.0
        dz = 0.0
        dweight = 0.0

        dcc = 0.0
        dcw = 0.0
        dss = 0.0
        dsw = 0.0
        dxw = 0.0
        dvarw = 0.0

        dtau = 0.0

        dpower = 0.0
        dpowz = 0.0
        damp = 0.0
        dneff = 0.0
        davew = 0.0

        dfre = 0.0

        n1 = 0
        n2 = 0

        dmz = 0.0
        dmzfre = 0.0
        dmzamp = 0.0
        dmcon = 0.0
        dmneff = 0.0

        twopi = 2.0 * math.pi

        ndim = 2
        itau1 = 0 # ----> 1 or 0 ??
        itau2 = self.ntau # -------> ???
        ifreq1 = 1
        ifreq2 = self.nfreq
        nstart = 1

        dmat_par = numpy.zeros(shape=(3, 3))

        for itau in range(itau1_par, itau2_par):
            nstart = 1
            dtau = self.tau[itau]

            dmfre = 0.0
            dmamp = 0.0
            dmcon = 0.0
            dmneff = 0.0
            dmz = -1.0 # less than the smallest WWZ

            for ifreq in range(ifreq1, ifreq2 + 1):
                dfre = self.freq[ifreq]
                domega = dfre * twopi

                for i in range(0, ndim + 1):
                    dvec[i] = 0.0
                    for j in range(0, ndim + 1):
                        dmat_par[i][j] = 0.0

                dweight2 = 0.0

                for idat in range(nstart, self.numdat):

                    dz = domega * (self.time[idat] - dtau)
                    dweight = math.exp(-1.0 * self.dcon * dz * dz)

                    if (dweight > 10**(-9)):
                        dcc = math.cos(dz)
                        dcw = dweight * dcc
                        dss = math.sin(dz)
                        dsw = dweight * dss
                        dmat_par[0][0] = dmat_par[0][0] + dweight
                        dweight2 = dweight2 + (dweight**2)
                        dmat_par[0][1] = dmat_par[0][1] + dcw
                        dmat_par[0][2] = dmat_par[0][2] + dsw
                        dmat_par[1][1] = dmat_par[1][1] + (dcw * dcc)
                        dmat_par[1][2] = dmat_par[1][2] + (dcw * dss)
                        dmat_par[2][2] = dmat_par[2][2] + (dsw * dss)

                        dxw = dweight * self.magnitude[idat]
                        dvec[0] = dvec[0] + dxw
                        dvarw = dvarw + (dxw * self.magnitude[idat])
                        dvec[1] = dvec[1] + (dcw * self.magnitude[idat])
                        dvec[2] = dvec[2] + (dsw * self.magnitude[idat])

                    elif dz > 0.0:
                        break
                    else:
                        nstart = idat + 1

                dpower = 0.0
                damp = 0.0

                for n1 in range(0, ndim + 1):
                    dcoef[n1] = 0.0

                if (dweight2 > 0.0):
                    dneff = (dmat_par[0][0] * dmat_par[0][0]) / dweight2
                else:
                    dneff = 0.0

                if (dneff > 3.0):

                    for n1 in range(0, ndim + 1):
                        dvec[n1] = dvec[n1] / dmat_par[0][0]

                        for n2 in range(1, ndim + 1 ):
                            dmat_par[n1][n2] = dmat_par[n1][n2] / \
                                               dmat_par[0][0]

                    if (dmat_par[0][0] > 0.0):
                        dvarw = dvarw / dmat_par[0][0]
                    else:
                        dvarw = 0.0

                    dmat_par[0][0] = 1.0
                    davew = dvec[0]
                    dvarw = dvarw - (davew ** 2)

                    if (dvarw <= 0.0):
                        dvarw = 10**-12

                    for n1 in range(1, ndim + 1):
                        for n2 in range(0, n1):
                            dmat_par[n1][n2] = dmat_par[n2][n1]

                    dmat_par = self.matrix_inv(dmat_par)

                    for n1 in range(0, ndim + 1):
                        for n2 in range(0, ndim + 1):
                            dcoef[n1] = dcoef[n1] + \
                                        dmat_par[n1][n2] * dvec[n2]

                        dpower = dpower + (dcoef[n1] * dvec[n1])

                    dpower = dpower - (davew ** 2)
                    dpowz = (dneff - 3.0) * dpower / (dvarw - dpower) / 2.0
                    dpower = (dneff - 1.0) * dpower / dvarw / 2.0
                    damp = math.sqrt(dcoef[1] * dcoef[1] + \
                                     dcoef[2] * dcoef[2])

                else:

                    dpowz = 0.0
                    dpower = 0.0
                    damp = 0.0

                    if (dneff < (10**(-9))):
                        dneff = 0.0

                if (damp < (10**(-9))):
                    damp = 0.0

                if (dpower < (10**(-9))):
                    dpower = 0.0

                if (dpowz < (10**(-9))):
                    dpowz = 0.0

                # Let's write everything out.

                output1.write("%s \t %s \t %s \t %s \t %s \t %s\n" % \
                             (str(dtau),str(dfre),str(dpowz),str(damp), \
                             str(dcoef[0]),str(dneff)))

                if (dpowz > dmz):
                    dmz = dpowz
                    dmfre = dfre
                    dmamp = damp
                    dmcon = dcoef[0]
                    dmneff = dneff

            #
            if max_periods == True:
                # writes the max_periods output if specified
                output2.write("%f \t %f \t %f \t %f \t %f \t %f\n" % \
                             (dtau, dmfre, dmz,  dmamp, dmcon, dmneff))

            if gnuplot_compatible == True:
                # added so that gnuplot reads out of the box
                output1.write("\n")

# If the script runs as a standalone, below is triggered

if __name__ == '__main__':

    # Parsing the arguments

    description = """
    A Weighted Wavelet Z-Transformation Application for Python.
    Translated by M. Emre Aydin - emre.m.aydin@gmail.com
    http://about.me/emre.aydin
    Available at http://github.com/eaydin

    Input arguments can be read from a file. The file descriptor
        prefix is '@'.
    In order to read argument from a file named args.txt,
        the argument @args.txt should be passed.
    An example for args.txt:

        -f=myinputfile.txt
        -o=theoutputfile.output
        -m
        --freq-step=0.001
        -l=0.001
        -hi=0.01
        -c=0.001
        -p=0

    You can pass arguments from file and commandline at the same time.
    If two same arguments passed by this method, the latter will be
    used. So if you want to override some arguments in a an argument
    file, specify the file first.
    An example usage for our earlier @args.txt is as:

        python wwz.py @args.txt -c=0.0125

    The above command will use the settings in args.txt but will
    use c=0.0125 instead of c=0.001

    Comments and blank lines are NOT allowed in argument files.

    Import this script via Python to use it as a module, rather than
    a standalone script. (import wwz)

    """

    parser = argparse.ArgumentParser(prog='wwz.py', \
                 formatter_class=argparse.RawDescriptionHelpFormatter,\
                 fromfile_prefix_chars="@", description=description)

    parser.add_argument("-f", "--file", type=argparse.FileType("r"),\
                        default=sys.stdin, required=True,\
                        help="the Input File, Raw Lightcurve")
    parser.add_argument("-o", "--output", type=argparse.FileType('w'),\
                        default=sys.stdout, required=True,\
                        help="the Output File Name")
    parser.add_argument("-l", "--freq-low", type=float, required=True,\
                        help="the Low Frequency Value")
    parser.add_argument("-hi", "--freq-high", type=float, required=True,\
                        help="the High Frequency Value")
    parser.add_argument("-d", "--freq-step", type=float, required=True,\
                        help="the dF value, incremental step for Frequency")
    parser.add_argument("-c", "--dcon", type=float, required=True,\
                        help="the C constant for the Window Function")
    parser.add_argument("-g", "--gnuplot-compatible", action="store_true",\
                        default=False, help="the Output file is GNUPlot \
                        compatible, which means the tau's will be grouped \
                        so that pm3d can easily map. Default value is \
                        'False'.")
    parser.add_argument("-m", "--max-periods", action="store_true", \
                        default=False, help="Creates a secondary \
                        output with the maximum Periods for each single \
                        tau. This can be drawn in 2D. The output filename \
                        is derived from the -o option, added 'max_periods'. \
                        Default value is 'False'.")
    parser.add_argument("-t", "--time-divisions", type=float, default=50.0, \
                        help="The Time Divisions value. Templeton assumes \
                        this as 50. VStars from AAVSO leaves this optional \
                        contrary to Templeton, yet it's default value is \
                        also 50.")
    parser.add_argument("--time", action="store_true", default=False, \
                        help="Calculate the time of operation in seconds \
                        and print to standard output.")
    parser.add_argument("--no-headers", action="store_true", default=False, \
                        help="Doesn't print headers to output files if set. \
                        Default is 'False'.")
    parser.add_argument("-p", "--parallel", help="Created threads to speed \
                        up the process. Default value is '1', which means \
                        single thread. '0' means number of detected CPUs,\
                        can be overridden.", type=int, default=1)

    args = parser.parse_args()

    # Check if user asks for time calculation
    if args.time:
        starttime = datetime.now()

    # Get the process number
    cpu=args.parallel

    if cpu != 1:
        try:
            import multiprocessing
        except (ImportError, NotImplementedError):
            print "Multiprocessing not available, using single CPU thread."
            cpu = 1

    if cpu == 0:
        # detect CPU cores
        try:
            import multiprocessing
            cpu = multiprocessing.cpu_count()
        except (ImportError, NotImplementedError):
            print "Multiprocessing not available, using single CPU thread."
            cpu = 1

    # Multiprocessing begins
    if cpu != 1:

        s=WWZPAR(args.file, args.output, args.freq_low, args.freq_high, \
                args.freq_step, args.dcon, args.time_divisions, \
                args.max_periods, args.gnuplot_compatible)

        s.readfile()
        s.maketau()
        s.makefreq()

        # the list of output filenames
        output_par_names = []
        for i in range(1,cpu+1):
            output_par_names.append('wwz.par.proc.%i' % i)

        ntau_par1 = s.ntau/cpu

        # the last ntau_par (turns out this is not necessary)
        ntau_par2 = s.ntau - (ntau_par1 * (cpu-1) )

        thread_list = []

        for i in range(1,cpu+1):
            if i != cpu:
                t = multiprocessing.Process(target=s.wwt, \
                args=(output_par_names[i-1],ntau_par1*(i-1),ntau_par1*i))
            else:
                t = multiprocessing.Process(target=s.wwt, \
                args=(output_par_names[i-1],ntau_par1*(i-1),s.ntau))
            thread_list.append(t)

        for thread in thread_list:
            thread.start()
        for thread in thread_list:
            thread.join()

        # join the outputs and delete remaining files

        if args.max_periods:
            output_max_periods = open(args.output.name + ".max_periods", "w")

        # Write the headers if asked
        if not args.no_headers:
            args.output.write("#%9s %10s %10s %10s %10s %10s\n" % \
                         ("TAU","FREQ","WWZ","AMP","COEF","NEFF"))
            if args.max_periods:
                output_max_periods.write("#%9s %10s %10s %10s %10s %10s\n" % \
                         ("TAU","FREQ","WWZ","AMP","COEF","NEFF"))

        for i in range(1,cpu+1):
            read_par = open(output_par_names[i-1],"r")
            for line in read_par:
                args.output.write(line)
            read_par.close()
            os.remove(output_par_names[i-1])

            if args.max_periods:

                read_par_max = open(output_par_names[i-1] + \
                                    ".max_periods.par","r")
                for line in read_par_max:
                    output_max_periods.write(line)
                read_par_max.close()
                os.remove(output_par_names[i-1] + ".max_periods.par")

    # for single threadding computing
    if cpu == 1:

        # Run the main class and its subroutines

        s=WWZ()

        time_data, magnitude_data = s.readfile(args.file)
        wwz_output = s.wwt(time_data, magnitude_data, args.freq_low, \
                           args.freq_high, args.freq_step, args.dcon, \
                           args.time_divisions)
        if args.gnuplot_compatible:
            send_ntau = len(wwz_output) / \
                  int(((args.freq_high - args.freq_low) / args.freq_step) + 1)

            s.writegnu(wwz_output, args.output, args.no_headers, \
                       args.max_periods, send_ntau)
        else:
            s.writefile(wwz_output, args.output, \
                        args.no_headers, args.max_periods)

    args.output.close()
    if args.max_periods == True and cpu != 1:
        output_max_periods.close()

   # Print calculated time
    if args.time:
        endtime = datetime.now()
        print "Run Time: %s seconds" % str((endtime - starttime).seconds)

max_stars_repo_name: eaydin/WWZ

max_issues_repo_name: WuShichao/WWZ

max_forks_repo_name: WuShichao/WWZ


text:
import numpy as np
from scipy.stats import chi2
from ._utils import _CheckInputs
from hyppo.ksample.base import KSampleTest, KSampleTestOutput

class SmoothCFTest(KSampleTest):
    r"""
    Smooth Characteristic Function test statistic and p-value

    The Smooth Characteristic Function test is a two-sample test that uses
    differences in the smoothed (analytic) characteristic function of
    two data distributions in order to determine how different
    the two data are :footcite:p:`chwialkowski2015fast`.

    Parameters
    ----------
    num_randfreq: int
        Used to construct random array with size ``(p, q)`` where `p` is the number of
        dimensions of the data and `q` is the random frequency at which the
        test is performed. These are the random test points at which test occurs (see notes).

    Notes
    -----
    The test statistic takes on the following form:

    .. math::

        nW_n\Sigma_n^{-1}W_n

    As seen in the above formulation, this test-statistic takes the same form as
    the Hotelling :math:`T^2` statistic. However, the components are
    defined differently in this case. Given data sets
    X and Y, define the following as :math:`Z_i`, the vector of differences:

    .. math::

        Z_i = (k(X_i, T_1) - k(Y_i, T_1), \ldots,
        k(X_i, T_J) - k(Y_i, T_J)) \in \mathbb{R}^J

    The above is the vector of differences between kernels at test points, :math:`T_j`.
    This same formulation is used in the Mean Embedding Test.
    Moving forward, :math:`W_n` can be defined:

    .. math::

        W_n = \frac{1}{n} \sum_{i = 1}^n Z_i

    This leaves :math:`\Sigma_n`, the covariance matrix as:

    .. math::

        \Sigma_n = \frac{1}{n}ZZ^T

    In the specific case of the Smooth Characteristic function test,
    the vector of differences can be defined as follows:

    .. math::

        Z_i = (f(X_i)\sin(X_iT_1) - f(Y_i)\sin(Y_iT_1),
        f(X_i)\cos(X_iT_1) - f(Y_i)\cos(Y_iT_1),\cdots) \in \mathbb{R}^{2J}

    Once :math:`S_n` is calculated, a threshold :math:`r_{\alpha}` corresponding to the
    :math:`1 - \alpha` quantile of a Chi-squared distribution w/ J degrees of freedom
    is chosen. Null is rejected if :math:`S_n` is larger than this threshold.

    References
    ----------
    .. footbibliography::
    """

    def __init__(self, num_randfreq=5):

        self.num_randfreq = num_randfreq
        KSampleTest.__init__(self)

    def statistic(self, x, y, random_state):
        r"""
        Calculates the smooth CF test statistic.

        Parameters
        ----------
        x,y : ndarray of float
            Input data matrices. ``x`` and ``y`` must have the same number of
            dimensions. That is, the shapes must be ``(n, p)`` and ``(m, p)`` where
            `n` is the number of samples and `p` and `q` are the number of
            dimensions.
        random_state: int
            Set random seed for generation of test points

        Returns
        -------
        stat : float
            The computed Smooth CF statistic.
        """
        _, p = np.shape(x)
        if random_state:
            np.random.seed(random_state)
        random_frequencies = np.random.randn(p, self.num_randfreq)

        x_smooth = np.exp(-np.linalg.norm(x, axis=1) ** 2 / 2).reshape(-1, 1)
        y_smooth = np.exp(-np.linalg.norm(y, axis=1) ** 2 / 2).reshape(-1, 1)

        x_mat = x.dot(random_frequencies)
        y_mat = y.dot(random_frequencies)

        x_smooth_cf = np.concatenate(
            (np.sin(x_mat) * x_smooth, np.cos(x_mat) * x_smooth), 1
        )
        y_smooth_cf = np.concatenate(
            (np.sin(y_mat) * y_smooth, np.cos(y_mat) * y_smooth), 1
        )

        difference = x_smooth_cf - y_smooth_cf
        return smooth_cf_distance(difference)

    def test(self, x, y, random_state=None):
        r"""
        Calculates the smooth CF test statistic and p-value.

        Parameters
        ----------
        x,y : ndarray of float
            Input data matrices. ``x`` and ``y`` must have the same number of
            dimensions. That is, the shapes must be ``(n, p)`` and ``(m, p)`` where
            `n` is the number of samples and `p` and `q` are the number of
            dimensions.
        random_state: int
            Set random seed for generation of test points

        Returns
        -------
        stat : float
            The computed Smooth CF statistic.
        pvalue : float
            The computed smooth CF p-value.

        Examples
        --------
        >>> import numpy as np
        >>> from hyppo.ksample import SmoothCFTest
        >>> np.random.seed(1234)
        >>> x = np.random.randn(500, 10)
        >>> y = np.random.randn(500, 10)
        >>> stat, pvalue = SmoothCFTest().test(x, y, random_state=1234)
        >>> '%.2f, %.3f' % (stat, pvalue)
        '4.70, 0.910'
        """
        check_input = _CheckInputs(inputs=[x, y], indep_test=None)
        x, y = check_input()

        stat = self.statistic(x, y, random_state)
        pvalue = chi2.sf(stat, 2 * self.num_randfreq)
        self.stat = stat
        self.pvalue = pvalue

        return KSampleTestOutput(stat, pvalue)

def smooth_cf_distance(difference):
    r"""
    Calculates the Smooth CF test statistic using the vector of differences.

    Parameters
    ----------
    difference : ndarray of float
        The vector of differences which indicates distance between mean embeddings.
    num_randfeatures : int
        The number of test frequencies

    Returns
    -------
    stat : float
        The computed smooth CF statistic.
    """
    num_samples, _ = np.shape(difference)
    sigma = np.cov(np.transpose(difference))

    mu = np.mean(difference, 0)

    stat = num_samples * mu.dot(np.linalg.solve(sigma, np.transpose(mu)))

    return stat

ext: py

reason: import numpy,from scipy


text:
# -*- coding: utf-8 -*-

' a module for some simple but important processes '

__author__ = 'Zitong Lu'

import nibabel as nib
import numpy as np
import os
import math
from scipy.stats import spearmanr
from scipy.stats import pearsonr
from scipy.stats import kendalltau

# get package abspath
package_root = os.path.dirname(os.path.abspath(__file__))

' a function for zeroing the value close to zero '

def limtozero(x):

    """
    zero the value close to zero

    Parameters
    ----------
    x : float

    Returns
    -------
    0
    """

    if x < 1e-15:
        x = 0

    return x

' a function for getting the affine of the fMRI-img '

def get_affine(file_name):

    """
    get the affine of the fMRI-img

    Parameters
    ----------
    file_name : string
        The filename of a sample fMRI-img in your experiment

    Returns
    -------
    affine : array
        The position information of the fMRI-image array data in a reference space.
    """

    img = nib.load(file_name)

    return img.affine

' a function for FWE-correction for fMRI RSA results '

def fwe_correct(p):

    """
    FWE correction for fMRI RSA results

    p : array
        The p-value map (3-D).
    """

    px = np.shape(p)[0]
    py = np.shape(p)[1]
    pz = np.shape(p)[2]

    n = 0

    for i in range(px):
        for j in range(py):
            for k in range(pz):

                if (math.isnan(p[i, j, k]) == False) and (p[i, j, k] != 0):
                    n = n + 1

    fwep = p*n

    print("finished FWE correct")

    return fwep

' a function for FDR-correction for fMRI RSA results '

def fdr_correct(p):

    """
    FDR correction for fMRI RSA results

    Parameters
    ----------
    p : array
        The p-value map (3-D).

    Returns
    -------
    correctp : array.
        The FDR corrected p-value map.
    """

    px = np.shape(p)[0]
    py = np.shape(p)[1]
    pz = np.shape(p)[2]

    n = 0

    for i in range(px):
        for j in range(py):
            for k in range(pz):

                if (math.isnan(p[i, j, k]) == False) and (p[i, j, k] != 0):

                    n = n + 1

    fdrp = np.full([px, py, pz], np.nan)

    pcluster = np.zeros([n], dtype=np.float)

    m = 0

    for i in range(px):
        for j in range(py):
            for k in range(pz):

                if (math.isnan(p[i, j, k]) == False) and p[i, j, k] != 0:
                    pcluster[m] = p[i, j, k]
                    m = m + 1

    index = np.argsort(pcluster)

    for l in range(n):
        pcluster[index[l]] = float(pcluster[index[l]] * n / (l + 0.5))

    """for l in range(n - 1):
        if pcluster[index[-l - 1]] < pcluster[index[-l - 2]]:
            pcluster[index[-l - 2]] = pcluster[index[-l - 1]]"""

    newpcluster = np.full([n], np.nan)

    for l in range(n):
        newpcluster[l] = pcluster[index[l]]

    m = 0

    for i in range(px):
        for j in range(py):
            for k in range(pz):

                if (math.isnan(p[i, j, k]) == False) and p[i, j, k] != 0:
                    fdrp[i, j, k] = newpcluster[m]
                    m = m + 1

    print("finished FDR correct")

    return fdrp

"""def fdr_correct(p, size=[60, 60, 60], n=64, type="sphere"):

    x = size[0]
    y = size[1]
    z = size[2]

    px = np.shape(p)[0]
    py = np.shape(p)[1]
    pz = np.shape(p)[2]

    if type == "cube":

        n = float(n*px*py*pz/(x*y*z))

        nq = 1
        ni = 1

        while nq < n:
            ni = ni + 1
            nq = ni*ni*ni

        n = nq

        print(n, ni)

        fdrp = np.full([px, py, pz], np.nan)

        for i in range(px-ni+1):
            for j in range(py-ni+1):
                for k in range(pz-ni+1):

                    pcluster = p[i:i+ni, j:j+ni, k:k+ni]
                    pcluster = np.reshape(pcluster, [n])

                    index = np.argsort(pcluster)

                    for l in range(n):
                        pcluster[index[l]] = float(pcluster[index[l]]*n/(l+1))

                    for l in range(n-1):
                        if pcluster[index[-l-1]] < pcluster[index[-l-2]]:
                            pcluster[index[-l-2]] = pcluster[index[-l-1]]

                    newpcluster = np.full([n], np.nan)

                    for l in range(n):
                        newpcluster[l] = pcluster[index[l]]

                    fdrp[i:i+ni, j:j+ni, k:k+ni] = np.reshape(newpcluster, [ni, ni, ni])

        print("finished FDR correct")

        return fdrp

    elif type == "sphere":

        ni = int(nr * px * py * pz / (x * y * z)) + 1

        c0 = [int(ni), int(ni), int(ni)]
        n0 = 0
        for i in range(int(2 * ni + 1)):
            for j in range(int(2 * ni + 1)):
                for k in range(int(2 * ni + 1)):
                    dist = np.square(i - c0[0]) + np.square(j - c0[1]) + np.square(k - c0[2])
                    if dist <= ni * ni:
                        # print(i, j)
                        n0 = n0 + 1

        n = n0
        # print(c0)
        # print(n0)
        print(n, ni)

        fdrp = np.full([px, py, pz], np.nan)

        for i in range(px - 2 * ni):
            for j in range(py - 2 * ni):
                for k in range(pz - 2 * ni):

                    nindex = 0

                    pcluster = np.zeros([n], dtype=np.float)

                    for lx in range(2 * ni + 1):
                        for ly in range(2 * ni + 1):
                            for lz in range(2 * ni + 1):
                                dist = np.square(lx - c0[0]) + np.square(ly - c0[1]) + np.square(lz - c0[2])
                                if dist <= ni * ni:
                                    pcluster[nindex] = p[i + lx, j + ly, k + lz]
                                    nindex = nindex + 1

                    index = np.argsort(pcluster)

                    for l in range(n):
                        pcluster[index[l]] = float(pcluster[index[l]] * n / (l + 1))

                    for l in range(n - 1):
                        if pcluster[index[-l - 1]] < pcluster[index[-l - 2]]:
                            pcluster[index[-l - 2]] = pcluster[index[-l - 1]]

                    newpcluster = np.full([n], np.nan)

                    for l in range(n):
                        newpcluster[l] = pcluster[index[l]]

                    nindex = 0

                    for lx in range(2 * ni + 1):
                        for ly in range(2 * ni + 1):
                            for lz in range(2 * ni + 1):
                                dist = np.square(lx - c0[0]) + np.square(ly - c0[1]) + np.square(lz - c0[2])
                                if dist <= ni * ni:
                                    fdrp[i + lx, j + ly, k + lz] = newpcluster[nindex]
                                    nindex = nindex + 1

        print("finished FDR correct")

        return fdrp"""

' a function for fMRI RSA results correction by threshold '

def correct_by_threshold(img, threshold):

    """
    correct the fMRI RSA results by threshold

    Parameters
    ----------
    img : array
        A 3-D array of the fMRI RSA results.
        The shape of img should be [nx, ny, nz]. nx, ny, nz represent the shape of the fMRI-img.
    threshold : int
        The number of voxels used in correction.
        If threshold=n, only the similarity clusters consisting more than n voxels will be visualized.

    Returns
    -------
    img : array
        A 3-D array of the fMRI RSA results after correction.
        The shape of img should be [nx, ny, nz]. nx, ny, nz represent the shape of the fMRI-img.
    """

    sx = np.shape(img)[0]
    sy = np.shape(img)[1]
    sz = np.shape(img)[2]

    nsmall = 1

    while nsmall*nsmall*nsmall < threshold:
        nsmall = nsmall + 1

    nlarge = nsmall + 2

    for i in range(sx-nlarge+1):
        for j in range(sy-nlarge+1):
            for k in range(sz-nlarge+1):

                listlarge = list(np.reshape(img[i:i+nlarge, j:j+nlarge, k:k+nlarge], [nlarge*nlarge*nlarge]))
                print(listlarge.count(0))

                if listlarge.count(0) < nlarge*nlarge*nlarge:

                    index1 = 0

                    for l in range(nlarge):
                        for m in range(nlarge):

                            if img[i + l, j + m, k] == 0:
                                index1 = index1 + 1

                            if img[i + l, j + m, k + nlarge - 1] == 0:
                                index1 = index1 + 1

                    for l in range(nlarge-1):
                        for m in range(nlarge-2):

                            if img[i + l, j, k + m] == 0:
                                index1 = index1 + 1

                            if img[i, j + l + 1, k + m] == 0:
                                index1 = index1 + 1

                            if img[i + nlarge - 1, j + l, k + m] == 0:
                                index1 = index1 + 1

                            if img[i + l + 1, j + nlarge - 1, k + m] == 0:
                                index1 = index1 + 1

                    nex = nlarge * nlarge * nlarge - nsmall * nsmall * nsmall
                    print("index1:"+str(index1))

                    if index1 == nex:
                        print("**************************")
                        unit = img[i+1:i+1+nsmall, j+1:j+1+nsmall, k+1:k+1+nsmall]
                        unit = np.reshape(unit, [nsmall*nsmall*nsmall])
                        list_internal = list(unit)
                        index2 = nsmall*nsmall*nsmall-list_internal.count(0)
                        print(index1, index2)

                        if index2 < threshold:
                            img[i+1:i+1+nsmall, j]

                            for l in range(nsmall):
                                for m in range(nsmall):
                                    for p in range(nsmall):
                                        img[i+1:i+1+nsmall, j+1:j+1+nsmall, k+1:k+1+nsmall] = np.zeros([nsmall, nsmall, nsmall])

    print("finished correction")

    return img

' a function for getting ch2.nii.gz '

def get_bg_ch2():

    """
    get ch2.nii.gz

    Returns
    -------
    path : string
        The absolute file path of 'ch2.nii.gz'
    """

    return os.path.join(package_root, 'template/ch2.nii.gz')

' a function for getting ch2bet.nii.gz '

def get_bg_ch2bet():

    """
    get ch2bet.nii.gz

    Returns
    -------
    path : string
        The absolute file path of 'ch2bet.nii.gz'
    """

    return os.path.join(package_root, 'template/ch2bet.nii.gz')

' a function for filtering the data by a ROI mask '

def datamask(fmri_data, mask_data):

    """
    filter the data by a ROI mask

    Parameters:
    fmri_data : array
        The fMRI data.
        The shape of fmri_data is [nx, ny, nz]. nx, ny, nz represent the size of the fMRI data.
    mask_data : array
        The mask data.
        The shape of mask_data is [nx, ny, nz]. nx, ny, nz represent the size of the fMRI data.

    Returns
    -------
    newfmri_data : array
        The new fMRI data.
        The shape of newfmri_data is [nx, ny, nz]. nx, ny, nz represent the size of the fMRI data.
    """

    nx, ny, nz = fmri_data.shape

    newfmri_data = np.full([nx, ny, nz], np.nan)

    for i in range(nx):
        for j in range(ny):
            for k in range(nz):

                if (mask_data[i, j, k] != 0) and (math.isnan(mask_data[i, j, k]) is False):
                    newfmri_data[i, j, k] = fmri_data[i, j, k]

    return newfmri_data

' a function for permutation test '

def permutation_test(v1, v2, iter=5000):

    """
    Conduct Permutation test

    Parameters
    ----------
    v1 : array
        Vector 1.
    v2 : array
        Vector 2.
    iter : int. Default is 5000.
        The times for iteration.

    Returns
    -------
    p : float
        The permutation test result, p-value.
    """

    # permutation test

    diff = abs(np.average(v1) - np.average(v2))
    v = np.hstack((v1, v2))
    nv = v.shape[0]
    ni = 0

    for i in range(iter):
        vshuffle = np.random.permutation(v)
        vshuffle1 = vshuffle[:int(nv/2)]
        vshuffle2 = vshuffle[int(nv/2):]
        diff_i = np.average(vshuffle1) - np.average(vshuffle2)

        if diff_i >= diff:
            ni = ni + 1

    # permunitation test p-value
    p = np.float64(ni/iter)

    return p

' a function for permutation test for correlation coefficients '

def permutation_corr(v1, v2, method="spearman", iter=5000):

    """
    Conduct Permutation test for correlation coefficients

    Parameters
    ----------
    v1 : array
        Vector 1.
    v2 : array
        Vector 2.
    iter : int. Default is 5000.
        The times for iteration.

    Returns
    -------
    p : float
        The permutation test result, p-value.
    """

    # permutation test

    if method == "spearman":

        rtest = spearmanr(v1, v2)[0]

        ni = 0

        for i in range(iter):
            v1shuffle = np.random.permutation(v1)
            v2shuffle = np.random.permutation(v2)
            rperm = spearmanr(v1shuffle, v2shuffle)[0]

            if rperm>rtest:
                ni = ni + 1

    if method == "pearson":

        rtest = spearmanr(v1, v2)[0]

        ni = 0

        for i in range(iter):
            v1shuffle = np.random.permutation(v1)
            v2shuffle = np.random.permutation(v2)
            rperm = pearsonr(v1shuffle, v2shuffle)[0]

            if rperm>rtest:
                ni = ni + 1

    if method == "kendalltau":

        rtest = spearmanr(v1, v2)[0]

        ni = 0

        for i in range(iter):
            v1shuffle = np.random.permutation(v1)
            v2shuffle = np.random.permutation(v2)
            rperm = kendalltau(v1shuffle, v2shuffle)[0]

            if rperm>rtest:
                ni = ni + 1

    p = np.float64((ni+1)/(iter+1))

    return p

max_stars_repo_name: guishuyunye-lyw/NeuroRA

max_issues_repo_name: guishuyunye-lyw/NeuroRA

max_forks_repo_name: guishuyunye-lyw/NeuroRA


text:
# This file is part of datacube-ows, part of the Open Data Cube project.
# See https://opendatacube.org for more information.
#
# Copyright (c) 2017-2021 OWS Contributors
# SPDX-License-Identifier: Apache-2.0
from __future__ import division

import numpy

# Style index functions

def scale_data(imgband_data, scale_from, scale_to):
    sc_min, sc_max = scale_from
    tc_min, tc_max = scale_to
    clipped = imgband_data.clip(sc_min, sc_max)
    normalised = (clipped - sc_min) / (sc_max - sc_min)
    scaled = normalised * (tc_max - tc_min)
    return scaled + tc_min

def scalable(undecorated):
    def decorated(*args, **kwargs):
        scale_from = kwargs.pop("scale_from", None)
        scale_to = kwargs.pop("scale_to", None)
        if scale_from is not None and scale_to is None:
            scale_to = [0, 255]
        unscaled = undecorated(*args, **kwargs)
        if scale_from:
            return scale_data(unscaled, scale_from, scale_to)
        return unscaled

    return decorated

def band_modulator(undecorated):
    def decorated(data, *args, **kwargs):
        band_mapper = kwargs.get("band_mapper", None)
        mult_band = kwargs.pop("mult_band", None)
        raw_data = undecorated(data, *args, **kwargs)
        if mult_band and band_mapper:
            mult_band = band_mapper(mult_band)
        if mult_band:
            return data[mult_band] * raw_data
        return raw_data
    return decorated

def sum_bands(data, band1, band2, band_mapper=None):
    if band_mapper:
        band1 = band_mapper(band1)
        band2 = band_mapper(band2)
    return data[band1] + data[band2]

def delta_bands(data, band1, band2, band_mapper=None):
    if band_mapper:
        band1 = band_mapper(band1)
        band2 = band_mapper(band2)
    typ1 = data[band1].dtype
    typ2 = data[band2].dtype
    if typ1.name.startswith('uint'):
        nodata = data[band1].nodata
        data[band1] = data[band1].astype('int32')
        data[band1].attrs["nodata"] = nodata
    if typ2.name.startswith('uint'):
        nodata = data[band2].nodata
        data[band2] = data[band2].astype('int32')
        data[band2].attrs["nodata"] = nodata
    # if typ1.name.startswith('uint') or typ2.name.startswith('uint'):
        # data = data.astype('int32', copy=False)
    return data[band1] - data[band2]

# N.B. Modifying scale_to would be dangerous - don't do it.
# pylint: disable=dangerous-default-value
@scalable
def norm_diff(data, band1, band2, band_mapper=None):
    # Calculate a normalised difference index.
    return delta_bands(data, band1, band2, band_mapper) / sum_bands(data, band1, band2, band_mapper)

@scalable
def constant(data, band, const, band_mapper=None):
    # Covert an xarray for a flat constant.
    # Useful for displaying mask extents as a flat colour and debugging.
    # params is assumed to be a tuple containing a constant value and a band name/alias.

    if band_mapper:
        band = band_mapper(band)
    return data[band] * 0.0 + const

@scalable
def single_band(data, band, band_mapper=None):
    # Use the raw value of a band directly as the index function.

    if band_mapper:
        band = band_mapper(band)
    return data[band]

@scalable
def band_quotient(data, band1, band2, band_mapper=None):
    if band_mapper:
        band1 = band_mapper(band1)
        band2 = band_mapper(band2)
    return data[band1] / data[band2]

@scalable
def band_quotient_sum(data, band1a, band1b, band2a, band2b, band_mapper=None):
    return band_quotient(data, band1a, band1b, band_mapper) + band_quotient(data, band2a, band2b, band_mapper)

@scalable
def sentinel2_ndci(data, b_red_edge, b_red, b_green, b_swir, band_mapper=None):
    red_delta = delta_bands(data, b_red_edge, b_red, band_mapper)
    red_sum = sum_bands(data, b_red_edge, b_red, band_mapper)
    mndwi = norm_diff(data, b_green, b_swir, band_mapper)

    return red_delta / red_sum.where(mndwi > 0.1)

def multi_date_delta(data, time_direction=-1):
    data1, data2 = (data.sel(time=dt) for dt in data.coords["time"].values)

#    data1, data2 = data.values.item(0), data.values.item(1)
    if time_direction >= 0:
        return data1 - data2
    else:
        return data2 - data1

@band_modulator
@scalable
def single_band_log(data, band, scale_factor, exponent, band_mapper=None):
    if band_mapper:
        band = band_mapper(band)
    d = data[band]
    return scale_factor * ((d ** exponent) - 1.0)

@band_modulator
@scalable
def single_band_arcsec(data, band, band_mapper=None):
    if band_mapper:
        band = band_mapper(band)
    d = data[band]
    return numpy.arccos(1 / (d + 1))

@band_modulator
@scalable
def single_band_offset_log(data, band, scale=1.0, offset=None, band_mapper=None):
    if band_mapper:
        band = band_mapper(band)
    d = data[band]
    if offset is not None:
        d = data[band] * scale + offset
        unscaled = numpy.log(d)
    else:
        unscaled = numpy.log1p(d * scale)
    return unscaled

@scalable
def radar_vegetation_index(data, band_hv, band_hh, band_mapper=None):
    if band_mapper:
        band_hv = band_mapper(band_hv)
        band_hh = band_mapper(band_hh)
    hv_sq = data[band_hv] * data[band_hv]
    hh_sq = data[band_hh] * data[band_hh]
    return (hv_sq * 4.0) / (hh_sq + hv_sq)

max_stars_repo_name: LiamOSullivan/datacube-ows

max_issues_repo_name: LiamOSullivan/datacube-ows

max_forks_repo_name: LiamOSullivan/datacube-ows

reason: import numpy


text:
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for `tf.data.Dataset.map()`."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import map_fn
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops

class MapBenchmark(benchmark_base.DatasetBenchmarkBase):
  """Benchmarks for `tf.data.Dataset.map()`."""

  def benchmark_chain_of_maps(self):

    def benchmark_helper(chain_length, fn, use_inter_op_parallelism, label,
                         benchmark_id):
      dataset = dataset_ops.Dataset.range(10000)
      for _ in range(chain_length):
        dataset = dataset_ops.MapDataset(
            dataset, fn, use_inter_op_parallelism=use_inter_op_parallelism)
      self.run_and_report_benchmark(
          dataset,
          num_elements=10000,
          extras={
              "model_name": "map.benchmark.%d" % benchmark_id,
              "parameters": "%d" % chain_length,
          },
          name="chain_length_%d%s" % (chain_length, label))

    chain_lengths = [0, 1, 2, 5, 10, 20, 50]
    for chain_length in chain_lengths:
      benchmark_helper(
          chain_length=chain_length,
          fn=lambda x: x + 1,
          use_inter_op_parallelism=True,
          label="",
          benchmark_id=1)
      benchmark_helper(
          chain_length=chain_length,
          fn=lambda x: x + 1,
          use_inter_op_parallelism=False,
          label="_single_threaded",
          benchmark_id=2)
      benchmark_helper(
          chain_length=chain_length,
          fn=lambda x: x,
          use_inter_op_parallelism=True,
          label="_short_circuit",
          benchmark_id=3)

  def benchmark_map_fan_out(self):
    fan_outs = [1, 2, 5, 10, 20, 50, 100]

    def benchmark_helper(fan_out, fn, use_inter_op_parallelism, label,
                         benchmark_id):
      dataset = dataset_ops.Dataset.from_tensors(
          tuple(0 for _ in range(fan_out))).repeat(None)
      dataset = dataset_ops.MapDataset(
          dataset, fn, use_inter_op_parallelism=use_inter_op_parallelism)
      self.run_and_report_benchmark(
          dataset,
          num_elements=10000,
          extras={
              "model_name": "map.benchmark.%d" % benchmark_id,
              "parameters": "%d" % fan_out,
          },
          name="fan_out_%d%s" % (fan_out, label))

    for fan_out in fan_outs:
      benchmark_helper(
          fan_out=fan_out,
          fn=lambda *xs: [x + 1 for x in xs],
          use_inter_op_parallelism=True,
          label="",
          benchmark_id=4)
      benchmark_helper(
          fan_out=fan_out,
          fn=lambda *xs: [x + 1 for x in xs],
          use_inter_op_parallelism=False,
          label="_single_threaded",
          benchmark_id=5)
      benchmark_helper(
          fan_out=fan_out,
          fn=lambda *xs: xs,
          use_inter_op_parallelism=True,
          label="_short_circuit",
          benchmark_id=6)

  def benchmark_sequential_control_flow(self):
    dataset = dataset_ops.Dataset.from_tensors(100000)

    def fn(x):
      i = constant_op.constant(0)

      def body(i, x):
        return math_ops.add(i, 1), x

      return control_flow_ops.while_loop(math_ops.less, body, [i, x])

    num_elements = 1
    dataset = dataset.map(fn)
    self.run_and_report_benchmark(
        dataset,
        num_elements=num_elements,
        extras={
            "model_name": "map.benchmark.8",
            "parameters": "%d" % num_elements,
        },
        name="sequential_control_flow",
        apply_default_optimizations=True)

  def benchmark_parallel_control_flow(self):
    dataset = dataset_ops.Dataset.from_tensors(
        random_ops.random_uniform([100, 10000000]))

    def fn(x):
      return map_fn.map_fn(
          lambda y: y * array_ops.transpose(y), x, parallel_iterations=10)

    num_elements = 1
    dataset = dataset.map(fn)
    self.run_and_report_benchmark(
        dataset,
        num_elements=1,
        extras={
            "model_name": "map.benchmark.9",
            "parameters": "%d" % num_elements,
        },
        name="parallel_control_flow",
        apply_default_optimizations=True)

  def _benchmark_nested_parallel_map(self, cycle_length, num_parallel_calls):
    k = 1024 * 1024
    num_map_elements = 10
    num_range_elements = 2000

    def g(_):
      return np.random.rand(50 * k).sum()

    def f(_):
      return dataset_ops.Dataset.range(num_map_elements).map(
          g, num_parallel_calls=num_parallel_calls)

    dataset = dataset_ops.Dataset.range(num_range_elements)
    dataset = dataset.interleave(
        f, cycle_length=cycle_length, num_parallel_calls=dataset_ops.AUTOTUNE)

    cycle_length_str = ("default"
                        if cycle_length is None else str(cycle_length))
    num_parallel_calls_str = ("autotune"
                              if num_parallel_calls == dataset_ops.AUTOTUNE else
                              str(num_parallel_calls))
    map_dataset_str = ("map" if num_parallel_calls is None else
                       "parallel_map_num_parallel_calls_%s" %
                       num_parallel_calls_str)

    self.run_and_report_benchmark(
        dataset,
        num_elements=num_map_elements * num_range_elements,
        extras={
            "model_name": "map.benchmark.10",
            "parameters": "%s_%s" % (cycle_length_str, num_parallel_calls_str),
        },
        name=("%s_cycle_length_%s" % (map_dataset_str, cycle_length_str)))

  def benchmark_nested_parallel_map(self):
    cycle_lengths = [None, 100]
    nums_parallel_calls = [None, 1, 10, 100, dataset_ops.AUTOTUNE]
    for cycle_length in cycle_lengths:
      for num_parallel_calls in nums_parallel_calls:
        self._benchmark_nested_parallel_map(cycle_length, num_parallel_calls)

if __name__ == "__main__":
  benchmark_base.test.main()

max_stars_repo_name: ashutom/tensorflow-upstream


text:
from dynamic_graph.sot.torque_control.tests.test_control_manager import cm
from dynamic_graph.sot.torque_control.tests.robot_data_test import initRobotData
from dynamic_graph.sot.core.matrix_util import matrixToTuple, vectorToTuple,rotate, matrixToRPY
from dynamic_graph.sot.torque_control.position_controller import *
from numpy import matrix, identity, zeros, eye, array, pi, ndarray, ones

# Instanciate the free flyer
pc = PositionController("pc_test")

q=zeros(initRobotData.nbJoints+6)
dq=zeros(initRobotData.nbJoints)
qRef=zeros(initRobotData.nbJoints)
dqRef=zeros(initRobotData.nbJoints)

# Setting the robot configuration
pc.base6d_encoders.value=q
pc.qRef.value=qRef

# Setting the robot velocities
pc.jointsVelocities.value=dq
pc.dqRef.value=dqRef

Kp=ones(initRobotData.nbJoints)
pc.Kp.value=Kp
Ki=ones(initRobotData.nbJoints)
pc.Ki.value=Ki
Kd=ones(initRobotData.nbJoints)
pc.Kd.value=Kd

# Initializing the entity.
pc.init(cm.controlDT, "control-manager-robot"); #initRobotData.robotRef)

pc.pwmDes.recompute(10)
pc.qError.recompute(10)

max_stars_repo_name: jviereck/sot-torque-control

max_issues_repo_name: jviereck/sot-torque-control

max_forks_repo_name: jviereck/sot-torque-control

reason: from numpy


text:
####
# Read in LSHTM results and perform inference on

####
import matplotlib
matplotlib.use('Agg')
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import norm
from scipy.special import expit
from sys import argv
from datetime import timedelta, datetime
import pystan
import pickle
import os, glob
from Reff_functions import *
from Reff_constants import *

iterations=5000
chains=2

### Read in md surveys

surveys = pd.DataFrame()
##Improve this to read by glob.glob and get all of them

path = "data/md/Barometer wave*.csv"
for file in glob.glob(path):
    surveys = surveys.append(pd.read_csv(file,parse_dates=['date']))

surveys = surveys.sort_values(by='date')
print("Latest Microdistancing survey is {}".format(surveys.date.values[-1]))

surveys.loc[surveys.state!='ACT','state'] = surveys.loc[surveys.state!='ACT','state'].map(states_initials).fillna(surveys.loc[surveys.state!='ACT','state'])
surveys['proportion'] = surveys['count']/surveys.respondents
surveys.date = pd.to_datetime(surveys.date)

always =surveys.loc[surveys.response=='Always'].set_index(["state",'date'])
always = always.unstack(['state'])

idx = pd.date_range('2020-03-01',pd.to_datetime("today"))

always = always.reindex(idx, fill_value=np.nan)

always.index.name = 'date'

#fill back to earlier and between weeks.
# Assume survey on day x applies for all days up to x - 6
always =always.fillna(method='bfill')

#assume values continue forward if survey hasn't completed
always = always.fillna(method='ffill')
always = always.stack(['state'])

#Zero out before first survey 20th March
always = always.reset_index().set_index('date')
always.loc[:'2020-03-20','count'] =0
always.loc[:'2020-03-20','respondents'] =0
always.loc[:'2020-03-20','proportion'] =0

always = always.reset_index().set_index(['state','date'])

survey_X = pd.pivot_table(data=always,
                          index='date',columns='state',values='proportion')
survey_counts_base =pd.pivot_table(data=always,
                          index='date',columns='state',values='count').drop(['Australia','Other'],axis=1).astype(int)

survey_respond_base = pd.pivot_table(data=always,
                          index='date',columns='state',values='respondents').drop(['Australia','Other'],axis=1).astype(int)

## Read in pystan model that is saved on disk
rho_model_gamma = """
data {
    int N; //data length num days
    int K; //Number of mobility indices
    int j; //Number of states
    matrix[N,j] Reff; //response
    matrix[N,K] Mob[j]; //Mobility indices
    matrix[N,K] Mob_std[j]; ///std of mobility
    matrix[N,j] sigma2; //Variances of R_eff from previous study
    vector[N] policy; //Indicators for post policy or not
    matrix[N,j] local; //local number of cases
    matrix[N,j] imported; //imported number of cases

    int N_v; //length of VIC days
    int j_v; //second wave states
    matrix[N_v,j_v] Reff_v; //Reff for VIC in June
    matrix[N_v,K] Mob_v[j_v]; //Mob for VIC June
    matrix[N_v,K] Mob_v_std[j_v];// std of mobility
    matrix[N_v,j_v] sigma2_v;// variance of R_eff from previous study
    vector[N_v] policy_v;// micro distancing compliance
    matrix[N_v,j_v] local_v; //local cases in VIC
    matrix[N_v,j_v] imported_v; //imported cases in VIC

    vector[N] count_md[j]; //count of always
    vector[N] respond_md[j]; // num respondants

    vector[N_v] count_md_v[j_v]; //count of always
    vector[N_v] respond_md_v[j_v]; // num respondants

    int map_to_state_index[j_v];// indices of second wave to map to first
    int total_N_pv; //total number of data in sec wave, entire state first
    vector[N_v] include_in_sec_wave[j_v]; // dates include in sec_wave
    int pos_starts[j_v];//starting positions for each state
}
parameters {
    vector[K] bet; //coefficients
    real<lower=0> R_I; //base level imports,
    real<lower=0> R_L; //base level local
    vector<lower=0>[j] R_Li; //state level estimates
    real<lower=0> sig; //state level variance
    real<lower=0> theta_md; // md weighting
    matrix<lower=0,upper=1>[N,j] prop_md; // proportion who are md'ing
    vector<lower=0,upper=1>[total_N_pv] prop_md_v;
    matrix<lower=0,upper=1>[N,j] brho; //estimate of proportion of imported cases
    matrix<lower=0,upper=1>[N,K] noise[j];
    //real<lower=0> R_temp;

    vector<lower=0,upper=1>[total_N_pv] brho_v; //estimate of proportion of imported cases
    //matrix<lower=0,upper=1>[N_v,K] noise_v[j_v];
}
transformed parameters {
    matrix<lower=0>[N,j] mu_hat;
    vector<lower=0>[total_N_pv] mu_hat_v;
    matrix<lower=0>[N,j] md; //micro distancing
    vector<lower=0>[total_N_pv] md_v;

    for (i in 1:j) {
        for (n in 1:N){
            md[n,i] = pow(1+theta_md , -1*prop_md[n,i]);

            mu_hat[n,i] = brho[n,i]*R_I + (1-brho[n,i])*2*R_Li[i]*(
            (1-policy[n]) + md[n,i]*policy[n] )*inv_logit(
            Mob[i][n,:]*(bet)); //mean estimate
        }
    }
    for (i in 1:j_v){
        int pos;
        if (i==1){
            pos=1;
        }
        else {
            //Add 1 to get to start of new group, not end of old group
            pos =pos_starts[i-1]+1;
            }
        for (n in 1:N_v){
            if (include_in_sec_wave[i][n]==1){
                md_v[pos] = pow(1+theta_md ,-1*prop_md_v[pos]);
                if (map_to_state_index[i] == 5) {
                    mu_hat_v[pos] = brho_v[pos]*R_I + (1-brho_v[pos])*(2*R_Li[
                    map_to_state_index[i]
                    ])*(
                    (1-policy_v[n]) + md_v[pos]*policy_v[n] )*inv_logit(
                    Mob_v[i][n,:]*(bet)
                    ); //mean estimate
                }
                else {
                    mu_hat_v[pos] = brho_v[pos]*R_I + (1-brho_v[pos])*2*R_Li[
                    map_to_state_index[i]
                    ]*(
                    (1-policy_v[n]) + md_v[pos]*policy_v[n] )*inv_logit(
                    Mob_v[i][n,:]*(bet)); //mean estimate
                }
                pos += 1;
            }
        }
    }

}
model {
    int pos2;
    bet ~ normal(0,1);
    theta_md ~ lognormal(0,0.5);
    //md ~ beta(7,3);

    R_L ~ gamma(1.8*1.8/0.05,1.8/0.05); //hyper-prior
    R_I ~ gamma(0.5*0.5/.2,0.5/.2);
    sig ~ exponential(200); //mean is 1/50=0.02
    R_Li ~ gamma( R_L*R_L/sig, R_L/sig); //partial pooling of state level estimates
    for (i in 1:j) {
        for (n in 1:N){
            prop_md[n,i] ~ beta(1 + count_md[i][n], 1+ respond_md[i][n] - count_md[i][n]);
            brho[n,i] ~ beta( 1+ imported[n,i], 1+ local[n,i]); //ratio imported/ (imported + local)
            //noise[i][n,:] ~ normal( Mob[i][n,:] , Mob_std[i][n,:]);
            mu_hat[n,i] ~ gamma( Reff[n,i]*Reff[n,i]/(sigma2[n,i]), Reff[n,i]/sigma2[n,i]); //Stan uses shape/inverse scale
        }
    }

    for (i in 1:j_v){
        if (i==1){
            pos2=1;
        }
        else {
            //Add 1 to get to start of new group, not end of old group
            pos2 =pos_starts[i-1]+1;
            }
        for (n in 1:N_v){
            if (include_in_sec_wave[i][n]==1){
                prop_md_v[pos2] ~ beta(1 + count_md_v[i][n], 1+ respond_md_v[i][n] - count_md_v[i][n]);
                brho_v[pos2] ~ beta( 1+ imported_v[n,i], 1+ local_v[n,i]); //ratio imported/ (imported + local)
                //noise_v[i][n,:] ~ normal( Mob_v[i][n,:] , Mob_v_std[i][n,:]);
                mu_hat_v[pos2] ~ gamma( Reff_v[n,i]*Reff_v[n,i]/(sigma2_v[n,i]), Reff_v[n,i]/sigma2_v[n,i]);
                pos2+=1;
            }
        }
    }
}
"""

sm_pol_gamma = pystan.StanModel(
    model_code = rho_model_gamma,
    model_name ='gamma_pol_state'
)
#sm_pol_gamma = pickle.load(open('model/sm_pol_gamma.pkl','rb'))

###Create dates
try:
    cprs_start_date = pd.to_datetime(argv[1])#2020-04-01')
    cprs_end_date = pd.to_datetime(argv[1])#'2020-07-22')
except:
    print("Running full validation dates")
    cprs_start_date = pd.to_datetime('2020-04-01')
    cprs_end_date = pd.to_datetime('2020-10-07')
cprs_all_dates = pd.date_range(cprs_start_date, cprs_end_date, freq='7D')
cprs_dates = cprs_all_dates[cprs_all_dates!='2020-09-09']

#if argv[1]=='2020-09-09':
#    print("This won't run due to cprs date definitions, please comment out line 215.")

for data_date in cprs_dates:
    print(data_date)
    print(data_date.strftime('%d%b%Y'))
    #data_date =  pd.to_datetime('2020-08-17')

    ## also filter Reff by 10 days!
    ## need to truncate most recent days of Reff
    #########
    ### here is where I can loop over to perform inference##
    #######

    # Reff estimates from Price et al 2020
    #df_Reff = read_in_Reff() #estimates up to 14th April

    #df_Reff = read_in_LSHTM()#read_in_Reff()
    #df_Reff = df_Reff.loc[df_Reff.date_of_analysis==data_date.strftime("%Y-%m-%d")]

    df_Reff = pd.read_csv("results/EpyReff/Reff"+
                data_date.strftime("%Y-%m-%d")+"tau_4.csv",parse_dates=['INFECTION_DATES'])
    df_Reff['date'] = df_Reff.INFECTION_DATES
    df_Reff['state'] = df_Reff.STATE
    print('data loaded')
    if data_date < pd.to_datetime('2020-06-02'):
        #no leading zero on early dates
        if data_date.day <10:
            df_state = read_in_cases(case_file_date=data_date.strftime('%d%b%Y')[1:])
        else:
            df_state = read_in_cases(case_file_date=data_date.strftime('%d%b%Y'))
    else:
        df_state = read_in_cases(case_file_date=data_date.strftime('%d%b%Y'))

    print('cases loaded')
    df_Reff = df_Reff.merge(df_state,how='left',left_on=['state','date'], right_on=['STATE','date_inferred']) #how = left to use Reff days, NNDSS missing dates
    df_Reff['rho_moving'] = df_Reff.groupby(['state'])['rho'].transform(lambda x: x.rolling(7,1).mean()) #minimum number of 1

    #some days have no cases, so need to fillna
    df_Reff['rho_moving'] = df_Reff.rho_moving.fillna(method='bfill')

    #shift counts to align with infection date not symptom date
    # dates should be complete at this point, no days skipped
    # will be some end days with NaN, but that should be fine since
    # we don't use the most recent 10 days
    df_Reff['local'] = df_Reff.local.shift(periods=-5)
    df_Reff['imported'] = df_Reff.imported.shift(periods=-5)
    df_Reff['rho_moving'] = df_Reff.rho_moving.shift(periods=-5)
    df_Reff['rho'] = df_Reff.rho.shift(periods=-5)
    df_Reff['local'] = df_Reff.local.fillna(0)
    df_Reff['imported'] = df_Reff.imported.fillna(0)
    #Add Insight traffic
    #df_ai = read_AddInsight()

    df_google = read_in_google(local=True,moving=True)

    df= df_google.merge(df_Reff[['date','state','mean','lower','upper',
                                'top','bottom','std','rho','rho_moving','local','imported']], on=['date','state'],how='inner')

    ##### Create useable dataset
    ## ACT and NT not in original estimates, need to extrapolated
    states_to_fit = sorted(['NSW','VIC','QLD','SA','WA','TAS']) #sorting keeps consistent with sort in data_by_state
    fit_post_March = True
    ban = '2020-03-20'
    start_date = '2020-03-01'
    end_date = '2020-03-31'

    ##Second wave inputs
    sec_states=sorted(['NSW'])
    sec_start_date = '2020-06-01'
    sec_end_date = '2021-01-19'

    # Setting a fixed end date to the second wave
    # if data_date > pd.to_datetime("2020-06-12"):
    #     possible_end_date = data_date - timedelta(10)#subtract 10 days to aovid right truncation
    # else:
    #     possible_end_date = pd.to_datetime("2020-06-01")
    # sec_end_date = possible_end_date.strftime('%Y-%m-%d')
    #min('2020-08-14',possible_end_date.strftime('%Y-%m-%d')) #all we have for now

    fit_mask = df.state.isin(states_to_fit)
    if fit_post_March:
        fit_mask = (fit_mask) & (df.date >= start_date)

    fit_mask = (fit_mask) & (df.date <= end_date )

    second_wave_mask = df.state.isin(sec_states)
    second_wave_mask = (second_wave_mask) & (df.date >= sec_start_date)
    second_wave_mask = (second_wave_mask) & (df.date <= sec_end_date)

    predictors = mov_values.copy()
    #predictors.extend(['driving_7days','transit_7days','walking_7days','pc'])

    #remove residential to see if it improves fit
    predictors.remove('residential_7days')

    df['post_policy'] = (df.date >= ban).astype(int)

    dfX = df.loc[fit_mask].sort_values('date')
    df2X = df.loc[second_wave_mask].sort_values('date')

    #filter out the surveys we don't have
    if df2X.shape[0]>0:
        survey_respond = survey_respond_base.loc[:df2X.date.values[-1]]
        survey_counts = survey_counts_base.loc[:df2X.date.values[-1]]
    else:
        survey_respond = survey_respond_base.loc[:dfX.date.values[-1]]
        survey_counts = survey_counts_base.loc[:dfX.date.values[-1]]

    #choose dates for each state for sec wave
    sec_date_range = {
        'NSW':pd.date_range(start=sec_start_date,end=sec_end_date).values,
        # 'VIC':pd.date_range(start=sec_start_date,end='2020-10-28').values
    }
    df2X['is_sec_wave'] =0
    for state in sec_states:
        df2X.loc[df2X.state==state,'is_sec_wave'] = df2X.loc[df2X.state==state].date.isin(
            sec_date_range[state]
            ).astype(int).values

    data_by_state= {}
    sec_data_by_state={}
    for value in ['mean','std','local','imported']:
        data_by_state[value] = pd.pivot(dfX[['state',value,'date']],
                            index='date',columns='state',values=value).sort_index(
            axis='columns')
        sec_data_by_state[value] = pd.pivot(df2X[['state',value,'date']],
                            index='date',columns='state',values=value).sort_index(
            axis='columns')
            #account for dates pre second wave
        if df2X.loc[df2X.state==sec_states[0]].shape[0]==0:
            print("making empty")
            sec_data_by_state[value] = pd.DataFrame(columns=sec_states).astype(float)

    mobility_by_state =[]
    mobility_std_by_state=[]
    count_by_state =[]
    respond_by_state=[]

    #FIRST PHASE
    for state in states_to_fit:

        mobility_by_state.append(dfX.loc[dfX.state==state, predictors].values/100)
        mobility_std_by_state.append(
            dfX.loc[dfX.state==state,[val+'_std' for val in predictors]].values/100
        )
        count_by_state.append(survey_counts.loc[start_date:end_date,state].values)
        respond_by_state.append(survey_respond.loc[start_date:end_date,state].values)

    sec_mobility_by_state =[]
    sec_mobility_std_by_state=[]
    sec_count_by_state=[]
    sec_respond_by_state=[]
    include_in_sec_wave=[]
    #SECOND PHASE
    for state in sec_states:

        sec_mobility_by_state.append(df2X.loc[df2X.state==state, predictors].values/100)
        sec_mobility_std_by_state.append(
            df2X.loc[df2X.state==state,[val+'_std' for val in predictors]].values/100
        )
        sec_count_by_state.append(survey_counts.loc[sec_start_date:sec_end_date,state].values)
        sec_respond_by_state.append(survey_respond.loc[sec_start_date:sec_end_date,state].values)
        include_in_sec_wave.append(df2X.loc[df2X.state==state,'is_sec_wave'].values)
    policy_v = [1]*df2X.loc[df2X.state==sec_states[0]].shape[0]
    policy = dfX.loc[dfX.state==states_to_fit[0],'post_policy']

    state_index = { state : i+1  for i, state in enumerate(states_to_fit)}
    ##Make state by state arrays
    input_data ={
        'N': dfX.loc[dfX.state==states_to_fit[0]].shape[0],
        'K': len(predictors),
        'j':len(states_to_fit),
        'Reff': data_by_state['mean'].values,
        'Mob': mobility_by_state,
        'Mob_std':mobility_std_by_state,
        'sigma2': data_by_state['std'].values**2,
        'policy': policy.values,
        'local':data_by_state['local'].values,
        'imported':data_by_state['imported'].values,

        'N_v': df2X.loc[df2X.state==sec_states[0]].shape[0],
        'j_v': len(sec_states),
        'Reff_v': sec_data_by_state['mean'].values,
        'Mob_v': sec_mobility_by_state,
        'Mob_v_std':sec_mobility_std_by_state,
        'sigma2_v': sec_data_by_state['std'].values**2,
        'policy_v': policy_v,
        'local_v':sec_data_by_state['local'].values,
        'imported_v':sec_data_by_state['imported'].values,

        'count_md':count_by_state,
        'respond_md':respond_by_state,
        'count_md_v':sec_count_by_state,
        'respond_md_v':sec_respond_by_state,
        'map_to_state_index': [state_index[state] for state in sec_states],
        'total_N_pv':sum( [sum(x) for x in include_in_sec_wave]),
        'include_in_sec_wave': include_in_sec_wave,
        'pos_starts': np.cumsum([sum(x) for x in include_in_sec_wave])
    }

    fit = sm_pol_gamma.sampling(
        data=input_data,
        iter=iterations,
        chains=chains,
        #control={'max_treedepth':15}
    )

    #make results dir
    results_dir ="results/soc_mob_posterior/"
    os.makedirs(results_dir,exist_ok=True)

    filename = "stan_posterior_fit" + data_date.strftime("%Y-%m-%d") + ".txt"
    with open(results_dir+filename, 'w') as f:
        print(fit.stansummary(pars=['bet','R_I','R_L','R_Li','theta_md','sig']), file=f)
    samples_mov_gamma = fit.to_dataframe(pars=['bet','R_I','R_L','R_Li','sig','brho','theta_md','brho_v'])

    # Plot ratio of imported to total cases
    # First phase
    #rho calculated at data entry
    if isinstance(df_state.index, pd.MultiIndex):
        df_state = df_state.reset_index()

    states=sorted(['NSW','QLD','VIC','TAS','SA','WA','ACT','NT'])
    fig,ax = plt.subplots(figsize=(24,9), ncols=len(states),sharey=True)
    states_to_fitd = {state: i+1 for i,state in enumerate(states_to_fit)      }

    for i, state in enumerate(states):
        if state in states_to_fit:
            dates = df_Reff.loc[(df_Reff.date>=start_date) &
                                (df_Reff.state==state)&(df_Reff.date<=end_date)].date
            rho_samples = samples_mov_gamma[['brho['+str(j+1)+','+str(states_to_fitd[state])+']' for j in range(dfX.loc[dfX.state==states_to_fit[0]].shape[0])]]
            ax[i].plot(dates, rho_samples.median(),label='fit',color='C0')
            ax[i].fill_between(dates, rho_samples.quantile(0.25),rho_samples.quantile(0.75),color='C0',alpha=0.4)

            ax[i].fill_between(dates, rho_samples.quantile(0.05),rho_samples.quantile(0.95),color='C0',alpha=0.4)
        else:
            sns.lineplot(x='date_inferred',y='rho',
                data=df_state.loc[(df_state.date_inferred>=start_date) & (df_state.STATE==state)&(df_state.date_inferred<=end_date)], ax=ax[i],color='C1',label='data')
        sns.lineplot(x='date',y='rho',
                data=df_Reff.loc[(df_Reff.date>=start_date) & (df_Reff.state==state)&(df_Reff.date<=end_date)], ax=ax[i],color='C1',label='data')
        sns.lineplot(x='date',y='rho_moving',
                data=df_Reff.loc[(df_Reff.date>=start_date) & (df_Reff.state==state)&(df_Reff.date<=end_date)], ax=ax[i],color='C2',label='moving')

        dates = dfX.loc[dfX.state==states_to_fit[0]].date

        ax[i].tick_params('x',rotation=20)
        ax[i].xaxis.set_major_locator(plt.MaxNLocator(4))
        ax[i].set_title(state)
    ax[0].set_ylabel('Proportion of imported cases')
    plt.legend()
    plt.savefig(results_dir+data_date.strftime("%Y-%m-%d")+"rho_first_phase.png",dpi = 144)

    # Second phase
    if df2X.shape[0]>0:
        fig,ax = plt.subplots(figsize=(24,9), ncols=len(sec_states),sharey=True, squeeze=False)
        states_to_fitd = {state: i+1 for i,state in enumerate(sec_states)      }
        pos = 1
        for i, state in enumerate(sec_states):
            #Google mobility only up to a certain date, so take only up to that value
            dates = df2X.loc[(df2X.state==state)&(
                df2X.is_sec_wave==1
            )].date.values
            #df_Reff.loc[(df_Reff.date>=sec_start_date) &
            #                    (df_Reff.state==state)&(df_Reff.date<=sec_end_date)].date
            rho_samples = samples_mov_gamma[
                ['brho_v['+str(j)+']'
                for j in range(pos, pos+df2X.loc[df2X.state==state].is_sec_wave.sum() ) ]
                ]
            pos = pos + df2X.loc[df2X.state==state].is_sec_wave.sum()

            ax[0,i].plot(dates, rho_samples.median(),label='fit',color='C0')
            ax[0,i].fill_between(dates, rho_samples.quantile(0.25),rho_samples.quantile(0.75),color='C0',alpha=0.4)

            ax[0,i].fill_between(dates, rho_samples.quantile(0.05),rho_samples.quantile(0.95),color='C0',alpha=0.4)

            sns.lineplot(x='date_inferred',y='rho',
                data=df_state.loc[(df_state.date_inferred>=sec_start_date) & (df_state.STATE==state)&(df_state.date_inferred<=sec_end_date)], ax=ax[0,i],color='C1',label='data')
            sns.lineplot(x='date',y='rho',
                    data=df_Reff.loc[(df_Reff.date>=sec_start_date) & (df_Reff.state==state)&(df_Reff.date<=sec_end_date)], ax=ax[0,i],color='C1',label='data')
            sns.lineplot(x='date',y='rho_moving',
                    data=df_Reff.loc[(df_Reff.date>=sec_start_date) & (df_Reff.state==state)&(df_Reff.date<=sec_end_date)], ax=ax[0,i],color='C2',label='moving')

            dates = dfX.loc[dfX.state==sec_states[0]].date

            ax[0,i].tick_params('x',rotation=20)
            ax[0,i].xaxis.set_major_locator(plt.MaxNLocator(4))
            ax[0,i].set_title(state)
        ax[0,0].set_ylabel('Proportion of imported cases')
        plt.legend()
        plt.savefig(results_dir+data_date.strftime("%Y-%m-%d")+"rho_sec_phase.png",dpi = 144)

    #plot marginal distributions
    fig,ax = plt.subplots(figsize=(12,9))
    samples_mov_gamma['R_L_prior'] = np.random.gamma(
    1.8*1.8/0.05,0.05/1.8, size=samples_mov_gamma.shape[0])

    samples_mov_gamma['R_I_prior'] = np.random.gamma(
    0.5**2/0.2, .2/0.5, size=samples_mov_gamma.shape[0])

    samples_mov_gamma['R_L_national'] = np.random.gamma(
        samples_mov_gamma.R_L.values **2/ samples_mov_gamma.sig.values,
        samples_mov_gamma.sig.values / samples_mov_gamma.R_L.values
    )
    df_R_values = pd.melt(samples_mov_gamma[[col for col in samples_mov_gamma if 'R' in col]])
    print(df_R_values.variable.unique())
    sns.violinplot(x='variable',y='value',
                data=pd.melt(samples_mov_gamma[[col for col in samples_mov_gamma if 'R' in col]]),
                ax=ax,
                cut=0)

    ax.set_yticks([1],minor=True,)
    ax.set_yticks([0,2,3],minor=False)
    ax.set_yticklabels([0,2,3],minor=False)
    ax.set_ylim((0,4))
    #state labels in alphabetical
    ax.set_xticklabels(['R_I','R_L0 mean',
    'R_L0 NSW','R_L0 QLD','R_L0 SA','R_L0 TAS','R_L0 VIC','R_L0 WA',#'R_temp',
    'R_L0 prior','R_I prior','R_L0 national'])
    ax.tick_params('x',rotation=45)
    ax.yaxis.grid(which='minor',linestyle='--',color='black',linewidth=2)
    plt.savefig(results_dir+data_date.strftime("%Y-%m-%d")+"R_priors.png",dpi = 144)

    posterior = samples_mov_gamma[['bet['+str(i)+']' for i in range(1,1+len(predictors))]
                                ]

    split=True
    md = 'power'#samples_mov_gamma.md.values

    posterior.columns = [val for val in predictors]
    long = pd.melt(posterior)

    fig,ax2 =plt.subplots(figsize=(12,9))

    ax2 = sns.violinplot(x='variable',y='value',#hue='policy',
                        data=long,
                        ax=ax2,
                        color='C0'
                    )

    ax2.plot([0]*len(predictors), linestyle='dashed',alpha=0.6, color = 'grey')
    ax2.tick_params(axis='x',rotation=90)

    #ax =plot_posterior_violin(posterior)

    #ax2.set_title('Coefficients of mobility indices')
    ax2.set_xlabel('Social mobility index')
    ax2.set_xticklabels([var[:-6] for var in predictors])
    ax2.tick_params('x',rotation=15)
    plt.savefig(
        results_dir+data_date.strftime("%Y-%m-%d")+'mobility_posteriors.png', dpi =144)

    RL_by_state = { state: samples_mov_gamma[
        'R_Li['+str(i)+']'].values for state, i in state_index.items()
    }
    ax3 =predict_plot(samples_mov_gamma,df.loc[(df.date>=start_date)&(df.date<=end_date)],gamma=True,
                    moving=True,split=split,grocery=True,ban = ban,
                    R=RL_by_state, var= True, md_arg=md,
                    rho=states_to_fit, R_I =samples_mov_gamma.R_I.values,prop=survey_X.loc[start_date:end_date])#by states....
    for ax in ax3:
        for a in ax:
            a.set_ylim((0,3))
            #a.set_xlim((start_date,end_date))
    plt.savefig(
        results_dir+data_date.strftime("%Y-%m-%d")+"total_Reff_allstates.png", dpi=144)

    if df2X.shape[0]>0:
        df['is_sec_wave'] =0
        for state in sec_states:
            df.loc[df.state==state,'is_sec_wave'] = df.loc[df.state==state].date.isin(
                sec_date_range[state]
                ).astype(int).values
        #plot only if there is second phase data
        ax4 =predict_plot(samples_mov_gamma,df.loc[(df.date>=sec_start_date)&(df.date<=sec_end_date)],gamma=True, moving=True,split=split,grocery=True,ban = ban,
                        R=RL_by_state, var= True, md_arg=md,
                        rho=sec_states, second_phase=True,
                        R_I =samples_mov_gamma.R_I.values,prop=survey_X.loc[sec_start_date:sec_end_date])#by states....
        for ax in ax4:
            for a in ax:
                a.set_ylim((0,3))
                #a.set_xlim((start_date,end_date))
        plt.savefig(
            results_dir+data_date.strftime("%Y-%m-%d")+"Reff_sec_phase.png", dpi=144)

        #remove plots from memory
        fig.clear()
        plt.close(fig)

    var_to_csv = predictors
    samples_mov_gamma[predictors] = samples_mov_gamma[['bet['+str(i)+']' for i in range(1,1+len(predictors))]]
    var_to_csv = ['R_I']+['R_L','sig']+['theta_md']+predictors + [
        'R_Li['+str(i+1)+']' for i in range(len(states_to_fit))
        ]

    samples_mov_gamma[var_to_csv].to_hdf('data/soc_mob_posterior'+data_date.strftime("%Y-%m-%d")+'.h5',key='samples')