<commit_before>#!/usr/bin/env python
"""
Stacks BOSS spectra
"""

import os
import math
import argparse

import numpy

from astropy.io import fits

import ROOT
import h5py

class Target(object):
    def __init__(self, plate, mjd, fiber):
        self.plate = int(plate)
        self.mjd = int(mjd)
        self.fiber = int(fiber)

    @classmethod
    def fromString(cls, targetString):
        targetString = targetString.strip().split()[0]
        plate, mjd, fiber = targetString.split('-')
        return cls(plate, mjd, fiber);

    def __str__(self):
        return '%d-%d-%d' % (self.plate,self.mjd,self.fiber)

def getFiducialWavelength(pixelIndex):
    return 3500.26*(10**(1e-4*pixelIndex))

def getFiducialWavelengthRatio(lambda1, lambda2=3500.26):
    return 1e4*math.log10(lambda1/lambda2)

def getFiducialPixelIndexOffset(coeff0, coeff1=1e-4):
    if coeff1 != 1e-4:
        return 0
    delta = (math.log10(3500.26)-coeff0)/coeff1
    offset = int(math.floor(delta+0.5))
    if math.fabs(delta-offset) > 0.01:
        return 0
    return -offset

def main():
    # parse command-line arguments
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-i","--input", type=str, default=None,
        help = "target list")
    parser.add_argument("-n","--ntargets", type=int, default=0,
        help = "number of targets to use, 0 for all")
    parser.add_argument("--first-target", type=int, default=0,
        help = "index of first target to use")
    parser.add_argument("--boss-root", type=str, default=None,
        help = "path to root directory containing BOSS data (ex: /data/boss)")
    parser.add_argument("--boss-version", type=str, default=None,
        help = "boss pipeline version tag (ex: v5_7_0)")
    parser.add_argument("--skim", type=str, default="dr12v1",
        help = "name of skim to use")
    parser.add_argument("--out-prefix", type=str, default="root",
        help = "output file prefix")
    parser.add_argument("--verbose", action="store_true",
        help = "more verbose output")
    parser.add_argument("--use-fits", action="store_true",
        help = "read spectra from PLATE/spPlate-PLATE-MJD.fits files, \
        otherwise use skim/plate-PLATE-MJD.root files")
    parser.add_argument("--zmin", type=float, default=2.1,
        help = "minimum quasar redshift to include")
    parser.add_argument("--zmax", type=float, default=3,
        help = "maximum quasar redshift to include")
    parser.add_argument("--nrestbins", type=float, default=1000,
        help = "number of redshift bins")
    parser.add_argument("--norm", action="store_true",
        help = "normalize spectra using mean flux value in a specified window")
    parser.add_argument("--norm-lo", type=float, default=1270,
        help = "min wavelength for normalization window")
    parser.add_argument("--norm-hi", type=float, default=1290,
        help = "max wavelength for normalization window")
    args = parser.parse_args()

    # set up paths
    boss_root = args.boss_root if args.boss_root is not None else os.getenv('BOSS_ROOT', None)
    boss_version = args.boss_version if args.boss_version is not None else os.getenv('BOSS_VERSION', None)

    if boss_root is None or boss_version is None:
        raise RuntimeError('Must speciy --boss-(root|version) or env var BOSS_(ROOT|VERSION)')

    skimPath = os.path.join(boss_root, 'skim', args.skim)
    fitsPath = os.path.join(boss_root, boss_version)

    # read target list
    targets = []
    with open(args.input,'r') as targetlist:
        for line in targetlist:
            words = line.strip().split()

            target = Target.fromString(line)

            target.z = float(words[3])
            targets.append(target)

    ntargets = len(targets)
    firstTarget = args.first_target
    endTarget = firstTarget + args.ntargets

    print 'Read %d targets (using %d:%d) from %s' % (ntargets,firstTarget,endTarget,args.input)

    # open output files ahead of time to catch potential io errors before processing data
    try:
        outfilename = args.out_prefix+'.hdf5'
        outfile = h5py.File(outfilename, 'w')
    except IOError:
        print 'Failed to open output file: ' % outfilename

    # initialize stack arrays
    arraySize = 4800
    zmin = args.zmin
    zmax = args.zmax

    restmin = 530
    restmax = 7000
    nrestbins = args.nrestbins

    fluxsum = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    counts = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    wfluxsum = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    weights = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    weightssq = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    fluxsq = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    wfluxsq = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)

    sqrtwflux = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    sqrtw = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    wfluxsq = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)

    skipcounter = 0

    # work on targets
    plateFileName = None
    for targetIndex in range(ntargets):
        if targetIndex < firstTarget or (endTarget > firstTarget and targetIndex >= endTarget):
            continue
        target = targets[targetIndex]

        # calculate z bin
        if target.z < zmin or target.z > zmax:
            continue

        # read spectrum
        flux = []
        wflux = []
        ivar = []
        numPixels = 0
        coeff0 = 0

        if args.use_fits:
            # load the spectrum file
            if plateFileName != 'spPlate-%s-%s.fits' % (target.plate, target.mjd):
                plateFileName = 'spPlate-%s-%s.fits' % (target.plate, target.mjd)
                if args.verbose:
                    print 'Opening plate file %s...' % os.path.join(skimPath,plateFileName)
                plateFile = fits.open(os.path.join(fitsPath,str(target.plate),plateFileName))
                # when does the file close?

            numPixels = plateFile[0].header['NAXIS1']
            coeff0 = plateFile[0].header['COEFF0']

            index = target.fiber-1
            andmask = plateFile[2].data[index]

            flux = plateFile[0].data[index]
            ivar = plateFile[1].data[index]

            ivar[andmask > 0] = 0

        else:
            # load the spectrum file
            if plateFileName != 'plate-%s-%s.root' % (target.plate, target.mjd):
                plateFileName = 'plate-%s-%s.root' % (target.plate, target.mjd)
                if args.verbose:
                    print 'Opening plate file %s...' % os.path.join(skimPath,plateFileName)
                plateFile = ROOT.TFile(os.path.join(skimPath,plateFileName))
                # when does the file close?
                plateTree = plateFile.Get('skim')

            combined = plateFile.Get('combined_%s' % target.fiber)
            numPixels = combined.GetN()

            xBuffer = combined.GetX()
            wavelength = numpy.frombuffer(xBuffer,count=numPixels)

            coeff0 = math.log10(wavelength[0])

            yBuffer = combined.GetY()
            flux = numpy.frombuffer(yBuffer,count=numPixels)

            yErrBuffer = combined.GetEY()
            fluxErr = numpy.frombuffer(yErrBuffer,count=numPixels)

            ivar = numpy.zeros(numPixels)
            nonzeroEntries = numpy.nonzero(fluxErr)
            ivar[nonzeroEntries] = 1/(fluxErr[nonzeroEntries]**2)

        # determine pixel offset
        offset = getFiducialPixelIndexOffset(coeff0)
        if numPixels + offset > arraySize:
            raise RuntimeError('woh! sprectrum out of range!')

        obswave = (10**coeff0)*numpy.power(10, 1e-4*(numpy.arange(0, numPixels)))
        restwave = obswave/(1+target.z)
        restindices = ((restwave - restmin)/(restmax - restmin)*nrestbins).astype(int)

        # normalize spectrum using flux window: 1280 +/- 10 Ang
        if args.norm:
            obslo = (1+target.z)*args.norm_lo
            obshi = (1+target.z)*args.norm_hi
            normlo = int(getFiducialWavelengthRatio(obslo)) - offset
            normhi = int(getFiducialWavelengthRatio(obshi)) - offset + 1
            # limit norm window to spectrum
            if normlo < 0:
                normlo = 0
            if normhi > numPixels or args.norm_hi < args.norm_lo:
                normhi = numPixels
            # calucuate mean flux in window
            normSlice = slice(normlo,normhi)
            norm = numpy.sum(ivar[normSlice]*flux[normSlice])
            normweight = numpy.sum(ivar[normSlice])
            if normweight == 0:
                skipcounter += 1
                print 'skipping %s (z=%.2f): 0 ivar in flux normalization range' % (target, target.z)
                continue
            norm /= normweight
            flux /= norm
            ivar *= norm*norm

        # add spectrum to stack
        #pixelSlice = slice(offset,offset+numPixels)
        obsindices = numpy.arange(offset,offset+numPixels)

        r = numpy.logical_or(restindices < 0, restindices >= nrestbins)
        ivar[r] = 0

        i = ivar > 0
        flux[ivar == 0] = 0
        fluxsum[obsindices,restindices] += flux
        # fluxsq[obsindices,restindices] += flux*flux
        counts[obsindices,restindices] += i

        wflux = flux*ivar
        wfluxsum[obsindices,restindices] += wflux
        weights[obsindices,restindices] += ivar
        # weightssq[obsindices,restindices] += ivar*ivar

        # isigma = numpy.sqrt(ivar)
        # sqrtw[obsindices,restindices] += isigma
        # sqrtwflux[obsindices,restindices] += isigma*flux

        # wfluxsq[obsindices,restindices] += wflux*flux

    print 'Skipped %d targets...' % skipcounter

    # divide by weights/number of entries
    fluxmean = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    c = numpy.nonzero(counts)
    fluxmean[c] = fluxsum[c]/counts[c]

    wfluxmean = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    w = numpy.nonzero(weights)
    wfluxmean[w] = wfluxsum[w]/weights[w]

    # wwdenom = weights**2 - weightssq
    # ww = numpy.nonzero(wwdenom)
    # wfluxvar = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    # wfluxvar[ww] = weights[ww]/(wwdenom[ww])*(wfluxsq[ww] - (wfluxmean[ww]**2)*weights[ww])

    # pullmean = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    # pullvar = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)

    # pullmean[w] = (sqrtwflux[w] - wfluxmean[w]*sqrtw[w])/counts[w]
    # pullvar[w] = (wfluxsq[w] - (wfluxmean[w]**2)*weights[w])/counts[w]

    # sn = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    # sn[c] = numpy.sqrt(wfluxsq[c]/counts[c])

    # save the stacked spectrum matrix
    print 'Saving stack to %s' % outfilename

    grp = outfile.create_group('hists2d')

    xaxis = 3500.26*numpy.power(10, 1e-4*(numpy.arange(0, arraySize+1)-0.5))
    dset = grp.create_dataset('xbinedges', data=xaxis)
    dset.attrs['label'] = 'Observed Wavelength ($\AA$)'

    yaxis = numpy.linspace(restmin,restmax,nrestbins+1,endpoint=True)
    dset = grp.create_dataset('ybinedges', data=yaxis)
    dset.attrs['label'] = 'Rest Wavelength ($\AA$)'

    dset = grp.create_dataset('fluxmean', data=fluxmean)
    dset.attrs['label'] = 'Normalized Flux Mean'# $(10^{-17} erg/cm^2/s/\AA)$'

    dset = grp.create_dataset('wfluxmean', data=wfluxmean)
    dset.attrs['label'] = 'Normalized Weighted Flux Mean'# $(10^{-17} erg/cm^2/s/\AA)$'

    # dset = grp.create_dataset('wfluxvar', data=wfluxvar)
    # dset.attrs['label'] = 'Normalized Weighted Flux Variance'# $(10^{-17} erg/cm^2/s/\AA)^{-2}$'

    dset = grp.create_dataset('counts', data=counts)
    dset.attrs['label'] = 'Counts'

    dset = grp.create_dataset('weights', data=weights)
    dset.attrs['label'] = 'Normalized Weights'# $(10^{-17} erg/cm^2/s/\AA)^{-2}$'

    # dset = grp.create_dataset('sn', data=sn)
    # dset.attrs['label'] = 'Signal to Noise Ratio'

    # dset = grp.create_dataset('pullmean', data=pullmean)
    # dset.attrs['label'] = 'Pull Mean'

    # dset = grp.create_dataset('pullvar', data=pullvar)
    # dset.attrs['label'] = 'Pull Variance'

    outfile.close()

if __name__ == '__main__':
    main()<commit_msg>add options to specify rest wavelength range<commit_after>#!/usr/bin/env python
"""
Stacks BOSS spectra
"""

import os
import math
import argparse

import numpy

from astropy.io import fits

import ROOT
import h5py

class Target(object):
    def __init__(self, plate, mjd, fiber):
        self.plate = int(plate)
        self.mjd = int(mjd)
        self.fiber = int(fiber)

    @classmethod
    def fromString(cls, targetString):
        targetString = targetString.strip().split()[0]
        plate, mjd, fiber = targetString.split('-')
        return cls(plate, mjd, fiber);

    def __str__(self):
        return '%d-%d-%d' % (self.plate,self.mjd,self.fiber)

def getFiducialWavelength(pixelIndex):
    return 3500.26*(10**(1e-4*pixelIndex))

def getFiducialWavelengthRatio(lambda1, lambda2=3500.26):
    return 1e4*math.log10(lambda1/lambda2)

def getFiducialPixelIndexOffset(coeff0, coeff1=1e-4):
    if coeff1 != 1e-4:
        return 0
    delta = (math.log10(3500.26)-coeff0)/coeff1
    offset = int(math.floor(delta+0.5))
    if math.fabs(delta-offset) > 0.01:
        return 0
    return -offset

def main():
    # parse command-line arguments
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-i","--input", type=str, default=None,
        help = "target list")
    parser.add_argument("-n","--ntargets", type=int, default=0,
        help = "number of targets to use, 0 for all")
    parser.add_argument("--first-target", type=int, default=0,
        help = "index of first target to use")
    parser.add_argument("--boss-root", type=str, default=None,
        help = "path to root directory containing BOSS data (ex: /data/boss)")
    parser.add_argument("--boss-version", type=str, default=None,
        help = "boss pipeline version tag (ex: v5_7_0)")
    parser.add_argument("--skim", type=str, default="dr12v1",
        help = "name of skim to use")
    parser.add_argument("--out-prefix", type=str, default="root",
        help = "output file prefix")
    parser.add_argument("--verbose", action="store_true",
        help = "more verbose output")
    parser.add_argument("--use-fits", action="store_true",
        help = "read spectra from PLATE/spPlate-PLATE-MJD.fits files, \
        otherwise use skim/plate-PLATE-MJD.root files")
    parser.add_argument("--zmin", type=float, default=2.1,
        help = "minimum quasar redshift to include")
    parser.add_argument("--zmax", type=float, default=3,
        help = "maximum quasar redshift to include")
    parser.add_argument("--restmin", type=float, default=750,
        help = "restframe wavelength minimum")
    parser.add_argument("--restmax", type=float, default=3500,
        help = "restframe wavelength maximum")
    parser.add_argument("--nrestbins", type=float, default=1000,
        help = "number of redshift bins")
    parser.add_argument("--norm", action="store_true",
        help = "normalize spectra using mean flux value in a specified window")
    parser.add_argument("--norm-lo", type=float, default=1270,
        help = "min wavelength for normalization window")
    parser.add_argument("--norm-hi", type=float, default=1290,
        help = "max wavelength for normalization window")
    args = parser.parse_args()

    # set up paths
    boss_root = args.boss_root if args.boss_root is not None else os.getenv('BOSS_ROOT', None)
    boss_version = args.boss_version if args.boss_version is not None else os.getenv('BOSS_VERSION', None)

    if boss_root is None or boss_version is None:
        raise RuntimeError('Must speciy --boss-(root|version) or env var BOSS_(ROOT|VERSION)')

    skimPath = os.path.join(boss_root, 'skim', args.skim)
    fitsPath = os.path.join(boss_root, boss_version)

    # read target list
    targets = []
    with open(args.input,'r') as targetlist:
        for line in targetlist:
            words = line.strip().split()

            target = Target.fromString(line)

            target.z = float(words[3])
            targets.append(target)

    ntargets = len(targets)
    firstTarget = args.first_target
    endTarget = firstTarget + args.ntargets

    print 'Read %d targets (using %d:%d) from %s' % (ntargets,firstTarget,endTarget,args.input)

    # open output files ahead of time to catch potential io errors before processing data
    try:
        outfilename = args.out_prefix+'.hdf5'
        outfile = h5py.File(outfilename, 'w')
    except IOError:
        print 'Failed to open output file: ' % outfilename

    # initialize stack arrays
    arraySize = 4800
    zmin = args.zmin
    zmax = args.zmax

    restmin = args.restmin
    restmax = args.restmax
    nrestbins = args.nrestbins

    fluxsum = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    counts = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    wfluxsum = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    weights = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    weightssq = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    fluxsq = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    wfluxsq = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)

    sqrtwflux = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    sqrtw = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)
    wfluxsq = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float64)

    skipcounter = 0

    # work on targets
    plateFileName = None
    for targetIndex in range(ntargets):
        if targetIndex < firstTarget or (endTarget > firstTarget and targetIndex >= endTarget):
            continue
        target = targets[targetIndex]

        # calculate z bin
        if target.z < zmin or target.z > zmax:
            continue

        # read spectrum
        flux = []
        wflux = []
        ivar = []
        numPixels = 0
        coeff0 = 0

        if args.use_fits:
            # load the spectrum file
            if plateFileName != 'spPlate-%s-%s.fits' % (target.plate, target.mjd):
                plateFileName = 'spPlate-%s-%s.fits' % (target.plate, target.mjd)
                if args.verbose:
                    print 'Opening plate file %s...' % os.path.join(skimPath,plateFileName)
                plateFile = fits.open(os.path.join(fitsPath,str(target.plate),plateFileName))
                # when does the file close?

            numPixels = plateFile[0].header['NAXIS1']
            coeff0 = plateFile[0].header['COEFF0']

            index = target.fiber-1
            andmask = plateFile[2].data[index]

            flux = plateFile[0].data[index]
            ivar = plateFile[1].data[index]

            ivar[andmask > 0] = 0

        else:
            # load the spectrum file
            if plateFileName != 'plate-%s-%s.root' % (target.plate, target.mjd):
                plateFileName = 'plate-%s-%s.root' % (target.plate, target.mjd)
                if args.verbose:
                    print 'Opening plate file %s...' % os.path.join(skimPath,plateFileName)
                plateFile = ROOT.TFile(os.path.join(skimPath,plateFileName))
                # when does the file close?
                plateTree = plateFile.Get('skim')

            combined = plateFile.Get('combined_%s' % target.fiber)
            numPixels = combined.GetN()

            xBuffer = combined.GetX()
            wavelength = numpy.frombuffer(xBuffer,count=numPixels)

            coeff0 = math.log10(wavelength[0])

            yBuffer = combined.GetY()
            flux = numpy.frombuffer(yBuffer,count=numPixels)

            yErrBuffer = combined.GetEY()
            fluxErr = numpy.frombuffer(yErrBuffer,count=numPixels)

            ivar = numpy.zeros(numPixels)
            nonzeroEntries = numpy.nonzero(fluxErr)
            ivar[nonzeroEntries] = 1/(fluxErr[nonzeroEntries]**2)

        # determine pixel offset
        offset = getFiducialPixelIndexOffset(coeff0)
        if numPixels + offset > arraySize:
            raise RuntimeError('woh! sprectrum out of range!')

        obswave = (10**coeff0)*numpy.power(10, 1e-4*(numpy.arange(0, numPixels)))
        restwave = obswave/(1+target.z)
        restindices = ((restwave - restmin)/(restmax - restmin)*nrestbins).astype(int)

        # normalize spectrum using flux window: 1280 +/- 10 Ang
        if args.norm:
            obslo = (1+target.z)*args.norm_lo
            obshi = (1+target.z)*args.norm_hi
            normlo = int(getFiducialWavelengthRatio(obslo)) - offset
            normhi = int(getFiducialWavelengthRatio(obshi)) - offset + 1
            # limit norm window to spectrum
            if normlo < 0:
                normlo = 0
            if normhi > numPixels or args.norm_hi < args.norm_lo:
                normhi = numPixels
            # calucuate mean flux in window
            normSlice = slice(normlo,normhi)
            norm = numpy.sum(ivar[normSlice]*flux[normSlice])
            normweight = numpy.sum(ivar[normSlice])
            if normweight == 0:
                skipcounter += 1
                print 'skipping %s (z=%.2f): 0 ivar in flux normalization range' % (target, target.z)
                continue
            norm /= normweight
            flux /= norm
            ivar *= norm*norm

        # add spectrum to stack
        #pixelSlice = slice(offset,offset+numPixels)
        obsindices = numpy.arange(offset,offset+numPixels)

        r = numpy.logical_or(restindices < 0, restindices >= nrestbins)
        ivar[r] = 0

        i = ivar > 0
        flux[ivar == 0] = 0
        fluxsum[obsindices,restindices] += flux
        # fluxsq[obsindices,restindices] += flux*flux
        counts[obsindices,restindices] += i

        wflux = flux*ivar
        wfluxsum[obsindices,restindices] += wflux
        weights[obsindices,restindices] += ivar
        # weightssq[obsindices,restindices] += ivar*ivar

        # isigma = numpy.sqrt(ivar)
        # sqrtw[obsindices,restindices] += isigma
        # sqrtwflux[obsindices,restindices] += isigma*flux

        # wfluxsq[obsindices,restindices] += wflux*flux

    print 'Skipped %d targets...' % skipcounter

    # divide by weights/number of entries
    fluxmean = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    c = numpy.nonzero(counts)
    fluxmean[c] = fluxsum[c]/counts[c]

    wfluxmean = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    w = numpy.nonzero(weights)
    wfluxmean[w] = wfluxsum[w]/weights[w]

    # wwdenom = weights**2 - weightssq
    # ww = numpy.nonzero(wwdenom)
    # wfluxvar = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    # wfluxvar[ww] = weights[ww]/(wwdenom[ww])*(wfluxsq[ww] - (wfluxmean[ww]**2)*weights[ww])

    # pullmean = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    # pullvar = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)

    # pullmean[w] = (sqrtwflux[w] - wfluxmean[w]*sqrtw[w])/counts[w]
    # pullvar[w] = (wfluxsq[w] - (wfluxmean[w]**2)*weights[w])/counts[w]

    # sn = numpy.zeros(shape=(arraySize,nrestbins), dtype=numpy.float32)
    # sn[c] = numpy.sqrt(wfluxsq[c]/counts[c])

    # save the stacked spectrum matrix
    print 'Saving stack to %s' % outfilename

    grp = outfile.create_group('hists2d')

    xaxis = 3500.26*numpy.power(10, 1e-4*(numpy.arange(0, arraySize+1)-0.5))
    dset = grp.create_dataset('xbinedges', data=xaxis)
    dset.attrs['label'] = 'Observed Wavelength ($\AA$)'

    yaxis = numpy.linspace(restmin,restmax,nrestbins+1,endpoint=True)
    dset = grp.create_dataset('ybinedges', data=yaxis)
    dset.attrs['label'] = 'Rest Wavelength ($\AA$)'

    dset = grp.create_dataset('fluxmean', data=fluxmean)
    dset.attrs['label'] = 'Normalized Flux Mean'# $(10^{-17} erg/cm^2/s/\AA)$'

    dset = grp.create_dataset('wfluxmean', data=wfluxmean)
    dset.attrs['label'] = 'Normalized Weighted Flux Mean'# $(10^{-17} erg/cm^2/s/\AA)$'

    # dset = grp.create_dataset('wfluxvar', data=wfluxvar)
    # dset.attrs['label'] = 'Normalized Weighted Flux Variance'# $(10^{-17} erg/cm^2/s/\AA)^{-2}$'

    dset = grp.create_dataset('counts', data=counts)
    dset.attrs['label'] = 'Counts'

    dset = grp.create_dataset('weights', data=weights)
    dset.attrs['label'] = 'Normalized Weights'# $(10^{-17} erg/cm^2/s/\AA)^{-2}$'

    # dset = grp.create_dataset('sn', data=sn)
    # dset.attrs['label'] = 'Signal to Noise Ratio'

    # dset = grp.create_dataset('pullmean', data=pullmean)
    # dset.attrs['label'] = 'Pull Mean'

    # dset = grp.create_dataset('pullvar', data=pullvar)
    # dset.attrs['label'] = 'Pull Variance'

    outfile.close()

if __name__ == '__main__':
    main()<|endoftext|>


<commit_before>#! /usr/bin/env python3
# Copyright 2022 Google LLC
#
# 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.

"""Sync GCE quota usage to Stackdriver for multiple projects.

This tool fetches global and/or regional quotas from the GCE API for
multiple projects, and sends them to Stackdriver as custom metrics, where they
can be used to set alert policies or create charts.
"""

import time

import base64
import datetime
import json
import logging
import os
import warnings

import click

from google.api_core.exceptions import GoogleAPIError
from google.api import label_pb2 as ga_label
from google.api import metric_pb2 as ga_metric
from google.cloud import monitoring_v3

import googleapiclient.discovery
import googleapiclient.errors

_BATCH_SIZE = 5
_METRIC_KIND = ga_metric.MetricDescriptor.MetricKind.GAUGE
_METRIC_TYPE_STEM = 'custom.googleapis.com/quota/'

_USAGE="usage"
_LIMIT="limit"
_UTILIZATION="utilization"

def _add_series(project_id, series, client=None):
  """Write metrics series to Stackdriver.

  Args:
    project_id: series will be written to this project id's account
    series: the time series to be written, as a list of
        monitoring_v3.types.TimeSeries instances
    client: optional monitoring_v3.MetricServiceClient will be used
        instead of obtaining a new one
  """
  client = client or monitoring_v3.MetricServiceClient()
  project_name = client.common_project_path(project_id)
  if isinstance(series, monitoring_v3.types.TimeSeries):
    series = [series]
  try:
    client.create_time_series(name=project_name, time_series=series)
  except GoogleAPIError as e:
    raise RuntimeError('Error from monitoring API: %s' % e)

def _configure_logging(verbose=True):
  """Basic logging configuration.

  Args:
    verbose: enable verbose logging
  """
  level = logging.DEBUG if verbose else logging.INFO
  logging.basicConfig(level=level)
  warnings.filterwarnings('ignore', r'.*end user credentials.*', UserWarning)

def _fetch_quotas(project, region='global', compute=None):
  """Fetch GCE per - project or per - region quotas from the API.

  Args:
    project: fetch global or regional quotas for this project id
    region: which quotas to fetch, 'global' or region name
    compute: optional instance of googleapiclient.discovery.build will be used
        instead of obtaining a new one
  """
  compute = compute or googleapiclient.discovery.build('compute', 'v1')
  try:
    if region != 'global':
      req = compute.regions().get(project=project, region=region)
    else:
      req = compute.projects().get(project=project)
    resp = req.execute()
    return resp['quotas']
  except (GoogleAPIError, googleapiclient.errors.HttpError) as e:
    logging.debug('API Error: %s', e, exc_info=True)
    raise RuntimeError('Error fetching quota (project: %s, region: %s)' %
                       (project, region))

def _get_series(metric_labels, value, metric_type, time_stamp, dt=None):
  """Create a Stackdriver monitoring time series from value and labels.

  Args:
    metric_labels: dict with labels that will be used in the time series
    value: time series value
    metric_type: which metric is this series for
    dt: datetime.datetime instance used for the series end time
  """
  series = monitoring_v3.types.TimeSeries()
  series.metric.type = metric_type
  series.resource.type = 'global'
  for label in metric_labels:
    series.metric.labels[label] = metric_labels[label]
  point = monitoring_v3.types.Point()
  point.value.double_value = value

  seconds = int(time_stamp)
  nanos = int((time_stamp - seconds) * 10 ** 9)
  interval = monitoring_v3.TimeInterval( {"end_time": {"seconds": seconds, "nanos": nanos}})
  point.interval=interval

  series.points.append(point)
  return series

def _quota_to_series_triplet(project, region, quota ):
  """Convert API quota objects to three Stackdriver monitoring time series: usage, limit and utilization

  Args:
    project: set in converted time series labels
    region: set in converted time series labels
    quota: quota object received from the GCE API
  """
  #labels = dict((k, str(v)) for k, v in quota.items())
  labels = dict()
  labels['project'] = project
  labels['region'] = region

  try:
    utilization = quota['usage'] / float(quota['limit'])
  except ZeroDivisionError:
    utilization = 0
  now=time.time()
  metric_type_prefix=_METRIC_TYPE_STEM+quota['metric'].lower()+'_'
  return [
    _get_series(labels, quota['usage'],metric_type_prefix+_USAGE,now),
    _get_series(labels, quota['limit'],metric_type_prefix+_LIMIT,now),
    _get_series(labels, utilization,metric_type_prefix+_UTILIZATION,now),
  ]

@click.command()
@click.option('--monitoring-project', required=True,
              help='monitoring project id')
@click.option('--gce-project', multiple=True,
              help='project ids (multiple), defaults to monitoring project')
@click.option('--gce-region', multiple=True,
              help='regions (multiple), defaults to "global"')
@click.option('--verbose', is_flag=True, help='Verbose output')
@click.argument('keywords', nargs=-1)
def main_cli(monitoring_project=None, gce_project=None, gce_region=None,
             verbose=False, keywords=None):
  """Fetch GCE quotas and writes them as custom metrics to Stackdriver.

  If KEYWORDS are specified as arguments, only quotas matching one of the
  keywords will be stored in Stackdriver.
  """
  try:
    _main(monitoring_project, gce_project, gce_region, verbose, keywords)
  except RuntimeError:
    logging.exception('exception raised')

def main(event, context):
  """Cloud Function entry point."""
  try:
    data = json.loads(base64.b64decode(event['data']).decode('utf-8'))
    _main(os.environ.get('GCP_PROJECT'), **data)
  # uncomment once https://issuetracker.google.com/issues/155215191 is fixed
  # except RuntimeError:
  #  raise
  except Exception:
    logging.exception('exception in cloud function entry point')

def _main(monitoring_project, gce_project=None, gce_region=None, verbose=False,
          keywords=None):
  """Module entry point used by cli and cloud function wrappers."""
  _configure_logging(verbose=verbose)
  gce_projects = gce_project or [monitoring_project]
  gce_regions = gce_region or ['global']
  keywords = set(keywords or [])
  logging.debug('monitoring project %s', monitoring_project)
  logging.debug('projects %s regions %s', gce_projects, gce_regions)
  logging.debug('keywords %s', keywords)
  quotas = []
  compute = googleapiclient.discovery.build(
      'compute', 'v1', cache_discovery=False)
  for project in gce_projects:
    logging.debug('project %s', project)
    for region in gce_regions:
      logging.debug('region %s', region)
      for quota in _fetch_quotas(project, region, compute=compute):
        if keywords and not any(k in quota['metric'] for k in keywords):
          # logging.debug('skipping %s', quota)
          continue
        logging.debug('quota %s', quota)
        quotas.append((project, region, quota))
  client, i = monitoring_v3.MetricServiceClient(), 0

  x= len(quotas)
  while i < len(quotas):
    series = sum([_quota_to_series_triplet(*q)  for q in quotas[i:i + _BATCH_SIZE]],[])
    _add_series(monitoring_project, series, client)
    i += _BATCH_SIZE
  #
  #for quota in quotas:
  #  series_triplet = _quota_to_series(project,region,quota)
  #  series_batch.update(series_triplet)
  #  if (series_batch.__len__==_BATCH_SIZE):
  #    _add_series(monitoring_project, series_batch, client)
  #    series_batch=dict()
  #_add_series(monitoring_project, series_batch, client)

if __name__ == '__main__':
  main_cli()
<commit_msg>address ludo review comments<commit_after>#! /usr/bin/env python3
# Copyright 2022 Google LLC
#
# 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.
"""Sync GCE quota usage to Stackdriver for multiple projects.

This tool fetches global and/or regional quotas from the GCE API for
multiple projects, and sends them to Stackdriver as custom metrics, where they
can be used to set alert policies or create charts.
"""

import base64
import datetime
import json
import logging
import os
import time
import warnings

import click

from google.api_core.exceptions import GoogleAPIError
from google.api import label_pb2 as ga_label
from google.api import metric_pb2 as ga_metric
from google.cloud import monitoring_v3

import googleapiclient.discovery
import googleapiclient.errors

_BATCH_SIZE = 5
_METRIC_KIND = ga_metric.MetricDescriptor.MetricKind.GAUGE
_METRIC_TYPE_STEM = 'custom.googleapis.com/quota/'

_USAGE = "usage"
_LIMIT = "limit"
_UTILIZATION = "utilization"

def _add_series(project_id, series, client=None):
  """Write metrics series to Stackdriver.

  Args:
    project_id: series will be written to this project id's account
    series: the time series to be written, as a list of
        monitoring_v3.types.TimeSeries instances
    client: optional monitoring_v3.MetricServiceClient will be used
        instead of obtaining a new one
  """
  client = client or monitoring_v3.MetricServiceClient()
  project_name = client.common_project_path(project_id)
  if isinstance(series, monitoring_v3.types.TimeSeries):
    series = [series]
  try:
    client.create_time_series(name=project_name, time_series=series)
  except GoogleAPIError as e:
    raise RuntimeError('Error from monitoring API: %s' % e)

def _configure_logging(verbose=True):
  """Basic logging configuration.

  Args:
    verbose: enable verbose logging
  """
  level = logging.DEBUG if verbose else logging.INFO
  logging.basicConfig(level=level)
  warnings.filterwarnings('ignore', r'.*end user credentials.*', UserWarning)

def _fetch_quotas(project, region='global', compute=None):
  """Fetch GCE per - project or per - region quotas from the API.

  Args:
    project: fetch global or regional quotas for this project id
    region: which quotas to fetch, 'global' or region name
    compute: optional instance of googleapiclient.discovery.build will be used
        instead of obtaining a new one
  """
  compute = compute or googleapiclient.discovery.build('compute', 'v1')
  try:
    if region != 'global':
      req = compute.regions().get(project=project, region=region)
    else:
      req = compute.projects().get(project=project)
    resp = req.execute()
    return resp['quotas']
  except (GoogleAPIError, googleapiclient.errors.HttpError) as e:
    logging.debug('API Error: %s', e, exc_info=True)
    raise RuntimeError('Error fetching quota (project: %s, region: %s)' %
                       (project, region))

def _get_series(metric_labels, value, metric_type, timestamp, dt=None):
  """Create a Stackdriver monitoring time series from value and labels.

  Args:
    metric_labels: dict with labels that will be used in the time series
    value: time series value
    metric_type: which metric is this series for
    dt: datetime.datetime instance used for the series end time
  """
  series = monitoring_v3.types.TimeSeries()
  series.metric.type = metric_type
  series.resource.type = 'global'
  for label in metric_labels:
    series.metric.labels[label] = metric_labels[label]
  point = monitoring_v3.types.Point()
  point.value.double_value = value

  seconds = int(timestamp)
  nanos = int((timestamp - seconds) * 10**9)
  interval = monitoring_v3.TimeInterval(
      {"end_time": {
          "seconds": seconds,
          "nanos": nanos
      }})
  point.interval = interval

  series.points.append(point)
  return series

def _quota_to_series_triplet(project, region, quota):
  """Convert API quota objects to three Stackdriver monitoring time series: usage, limit and utilization

  Args:
    project: set in converted time series labels
    region: set in converted time series labels
    quota: quota object received from the GCE API
  """
  #labels = dict((k, str(v)) for k, v in quota.items())
  labels = dict()
  labels['project'] = project
  labels['region'] = region

  try:
    utilization = quota['usage'] / float(quota['limit'])
  except ZeroDivisionError:
    utilization = 0
  now = time.time()
  metric_type_prefix = _METRIC_TYPE_STEM + quota['metric'].lower() + '_'
  return [
      _get_series(labels, quota['usage'], metric_type_prefix + _USAGE, now),
      _get_series(labels, quota['limit'], metric_type_prefix + _LIMIT, now),
      _get_series(labels, utilization, metric_type_prefix + _UTILIZATION, now),
  ]

@click.command()
@click.option('--monitoring-project', required=True,
              help='monitoring project id')
@click.option('--gce-project', multiple=True,
              help='project ids (multiple), defaults to monitoring project')
@click.option('--gce-region', multiple=True,
              help='regions (multiple), defaults to "global"')
@click.option('--verbose', is_flag=True, help='Verbose output')
@click.argument('keywords', nargs=-1)
def main_cli(monitoring_project=None, gce_project=None, gce_region=None,
             verbose=False, keywords=None):
  """Fetch GCE quotas and writes them as custom metrics to Stackdriver.

  If KEYWORDS are specified as arguments, only quotas matching one of the
  keywords will be stored in Stackdriver.
  """
  try:
    _main(monitoring_project, gce_project, gce_region, verbose, keywords)
  except RuntimeError:
    logging.exception('exception raised')

def main(event, context):
  """Cloud Function entry point."""
  try:
    data = json.loads(base64.b64decode(event['data']).decode('utf-8'))
    _main(os.environ.get('GCP_PROJECT'), **data)
  # uncomment once https://issuetracker.google.com/issues/155215191 is fixed
  # except RuntimeError:
  #  raise
  except Exception:
    logging.exception('exception in cloud function entry point')

def _main(monitoring_project, gce_project=None, gce_region=None, verbose=False,
          keywords=None):
  """Module entry point used by cli and cloud function wrappers."""
  _configure_logging(verbose=verbose)
  gce_projects = gce_project or [monitoring_project]
  gce_regions = gce_region or ['global']
  keywords = set(keywords or [])
  logging.debug('monitoring project %s', monitoring_project)
  logging.debug('projects %s regions %s', gce_projects, gce_regions)
  logging.debug('keywords %s', keywords)
  quotas = []
  compute = googleapiclient.discovery.build('compute', 'v1',
                                            cache_discovery=False)
  for project in gce_projects:
    logging.debug('project %s', project)
    for region in gce_regions:
      logging.debug('region %s', region)
      for quota in _fetch_quotas(project, region, compute=compute):
        if keywords and not any(k in quota['metric'] for k in keywords):
          # logging.debug('skipping %s', quota)
          continue
        logging.debug('quota %s', quota)
        quotas.append((project, region, quota))
  client, i = monitoring_v3.MetricServiceClient(), 0

  x = len(quotas)
  while i < len(quotas):
    series = sum(
        [_quota_to_series_triplet(*q) for q in quotas[i:i + _BATCH_SIZE]], [])
    _add_series(monitoring_project, series, client)
    i += _BATCH_SIZE
  #
  #for quota in quotas:
  #  series_triplet = _quota_to_series(project,region,quota)
  #  series_batch.update(series_triplet)
  #  if (series_batch.__len__==_BATCH_SIZE):
  #    _add_series(monitoring_project, series_batch, client)
  #    series_batch=dict()
  #_add_series(monitoring_project, series_batch, client)

if __name__ == '__main__':
  main_cli()
<|endoftext|>


<commit_before>from unittest import TestCase
from classes.categories import Categories

class TestCategories(TestCase):
    """This class will handle all the functions to test for the category name"""

    def setUp(self):
        """This method defines the test fixture for all test to be undertaken"""
        self.new_category = Categories()

    def test_pattern_categoryname(self):
        """Test if category name follows the pattern given"""
        catnamepattern = self.new_category.create_category("#^", "kaguna@gmail.com")
        self.assertEqual("catname_pattern", catnamepattern, "Invalid category name.")

    def test_empty_categoryname(self):
        """Test if category name is empty"""
        catname_empty = self.new_category.create_category(" ", "kaguna@gmail.com")
        self.assertEqual("null_empty_field", catname_empty, "Please fill in the category name.")

    def test_categoryname_exist(self):
        """Test if category name exists in the list"""
        self.new_category.categories = []
        self.new_category.create_category("lunch", "kaguna@gmail.com")
        catnameexist = self.new_category.create_category("lunch", "kaguna@gmail.com")
        self.assertEqual("catname_uniqueness", catnameexist, "Category name exists.")

    def test_categoryname_create(self):
        """Test if category name is different and register"""
        self.new_category.categories = []
        self.new_category.create_category("lunch", "kaguna@gmail.com")
        create_catname = self.new_category.create_category("supper", "kagunaa@gmail.com")
        self.assertEqual("success", create_catname, "Category created successfully.")

    # Below methods will test when updating the category name

    def test_pattern_update_categoryname(self):
        """Test if category name follows the pattern given"""
        category_name_pattern = self.new_category.edit_category("Luncho", "@@@", "karyorkir@gmail.com")
        self.assertEqual("categoryname_pattern", category_name_pattern, "Invalid category name.")

    def test_pattern_update_categoryname(self):
        """Test if category name is empty"""
        category_name_empty = self.new_category.edit_category("Luncho", " ", "karyorkir@gmail.com")
        self.assertEqual("null_empty_field", category_name_empty, "Invalid category name.")
<commit_msg>[Refactor] remove a test which is not necesary<commit_after>from unittest import TestCase
from classes.categories import Categories

class TestCategories(TestCase):
    """This class will handle all the functions to test for the category name"""

    def setUp(self):
        """This method defines the test fixture for all test to be undertaken"""
        self.new_category = Categories()

    def test_pattern_categoryname(self):
        """Test if category name follows the pattern given"""
        catnamepattern = self.new_category.create_category("#^", "kaguna@gmail.com")
        self.assertEqual("catname_pattern", catnamepattern, "Invalid category name.")

    def test_categoryname_exist(self):
        """Test if category name exists in the list"""
        self.new_category.categories = []
        self.new_category.create_category("lunch", "kaguna@gmail.com")
        catnameexist = self.new_category.create_category("lunch", "kaguna@gmail.com")
        self.assertEqual("catname_uniqueness", catnameexist, "Category name exists.")

    def test_categoryname_create(self):
        """Test if category name is different and register"""
        self.new_category.categories = []
        self.new_category.create_category("lunch", "kaguna@gmail.com")
        create_catname = self.new_category.create_category("supper", "kagunaa@gmail.com")
        self.assertEqual("success", create_catname, "Category created successfully.")

    # Below methods will test when updating the category name

    def test_pattern_update_categoryname(self):
        """Test if category name follows the pattern given"""
        category_name_pattern = self.new_category.edit_category("Luncho", "@@@", "karyorkir@gmail.com")
        self.assertEqual("categoryname_pattern", category_name_pattern, "Invalid category name.")

    def test_pattern_update_categoryname(self):
        """Test if category name is empty"""
        category_name_empty = self.new_category.edit_category("Luncho", " ", "karyorkir@gmail.com")
        self.assertEqual("null_empty_field", category_name_empty, "Invalid category name.")
<|endoftext|>


<commit_before>"""
This is the Configuration
"""
class BaseConfig(object):
    """ This is the main config class to be inherited """
    DEBUG = False
    TESTING = False

class DevelopmentConfig(BaseConfig):
    """ This is the main configuration class """
    DEBUG = True
    SQLALCHEMY_DATABASE_URI = 'sqlite:///C:\\User\\Renzo\\Documents\\Codes\\Python\\isproj2-fcth-crm\\dev\\crm.db'
    SECRET_KEY = 'developmentsecretkey'
    UPLOAD_FOLDER = 'app/upload'

class ProductionConfig(BaseConfig):
    """ This is the production configuration class """
    DATABASE = ''.join(['postgresql://zccisfiglvdagd:',
                        '5d5fe904f02e350ca33066c2a51a5c6cc0c5236ade95b403cea2418d0ca567c2',
                        '@ec2-54-243-255-57.compute-1.amazonaws.com:5432/dfn8ta0mogh2u9'])
    SECRET_KEY = '\xd6;\x9aq)\xea\xa5\x9c\xf5BM^N\x05X\\\xec5w\x0f\xb9\xab\xf9\xb8'
    UPLOAD_FOLDER = 'app/upload'
    SQLALCHEMY_DATABASE_URI = DATABASE
<commit_msg>Fixed SQLALCHEMY_DATABASE_URI for DevelopmentConfig<commit_after>"""
This is the Configuration
"""

class BaseConfig(object):
    """ This is the main config class to be inherited """
    DEBUG = False
    TESTING = False

class DevelopmentConfig(BaseConfig):
    """ This is the main configuration class """
    DEBUG = True
    SQLALCHEMY_DATABASE_URI = r'sqlite:///C:\Users\Renzo\Documents\Codes\Python\isproj2-fcth-crm\dev\crm.db'
    SECRET_KEY = 'developmentsecretkey'
    UPLOAD_FOLDER = 'app/upload'

class ProductionConfig(BaseConfig):
    """ This is the production configuration class """
    DATABASE = ''.join(['postgresql://zccisfiglvdagd:',
                        '5d5fe904f02e350ca33066c2a51a5c6cc0c5236ade95b403cea2418d0ca567c2',
                        '@ec2-54-243-255-57.compute-1.amazonaws.com:5432/dfn8ta0mogh2u9'])
    SECRET_KEY = '\xd6;\x9aq)\xea\xa5\x9c\xf5BM^N\x05X\\\xec5w\x0f\xb9\xab\xf9\xb8'
    UPLOAD_FOLDER = 'app/upload'
    SQLALCHEMY_DATABASE_URI = DATABASE
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Generated by Django 1.10.6 on 2017-08-08 16:28
from __future__ import unicode_literals

from django.conf import settings
from django.db import migrations, models
from django.contrib.auth import get_user_model
import django.db.models.deletion

from voty.initproc.globals import STATES

def migrate_eligible_voters_count(apps, schema_editor):
    # Let's set the eligable voters past initiation
    Initiative = apps.get_model('initproc', 'Initiative')
    User = get_user_model()
    for init in Initiative.objects.filter(state__in=[STATES.ACCEPTED, STATES.REJECTED]):
        init.eligible_voters = User.objects.filter(date_joined__lte=init.was_closed_at).count()
        init.save()

class Migration(migrations.Migration):

    dependencies = [
        ('initproc', '0017_auto_20170710_1738'),
    ]

    operations = [
        migrations.AddField(
            model_name='initiative',
            name='eligible_voters',
            field=models.IntegerField(blank=True, null=True),
            preserve_default=False,
        ),
        migrations.RunPython(migrate_eligible_voters_count),
    ]
<commit_msg>fixing upgrade path for beta<commit_after># -*- coding: utf-8 -*-
# Generated by Django 1.10.6 on 2017-08-08 16:28
from __future__ import unicode_literals

from django.conf import settings
from django.db import migrations, models
from django.contrib.auth import get_user_model
import django.db.models.deletion

from voty.initproc.globals import STATES

def migrate_eligible_voters_count(apps, schema_editor):
    # Let's set the eligable voters past initiation
    Initiative = apps.get_model('initproc', 'Initiative')
    User = get_user_model()
    for init in Initiative.objects.filter(state__in=[STATES.ACCEPTED, STATES.REJECTED]):
        if init.was_closed_at:
            init.eligible_voters = User.objects.filter(date_joined__lte=init.was_closed_at).count()
            init.save()

class Migration(migrations.Migration):

    dependencies = [
        ('initproc', '0017_auto_20170710_1738'),
    ]

    operations = [
        migrations.AddField(
            model_name='initiative',
            name='eligible_voters',
            field=models.IntegerField(blank=True, null=True),
            preserve_default=False,
        ),
        migrations.RunPython(migrate_eligible_voters_count),
    ]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding:utf-8 -*-
#---------------------------------------------------------------------------
# Name:        config.py
# Purpose:     各種設定
#
# Author:      Kosuke Akizuki
#
# Created:     2015/03/11
# Copyright:   (c) Kosuke Akizuki 2015
# Licence:     The MIT License (MIT)
#---------------------------------------------------------------------------

__author__ = "Kosuke Akizuki <slife1080@gmail.com>"
__status__ = "OK"
__version__ = "1.0"
__date__    = "11 Mar 2015"

import os

def join(*args):
    path = ''
    for e in args:
        path = os.path.join(path, e)
    path = os.path.normcase(path)
    path = os.path.normpath(path)
    return path

HOME = os.environ.get('HOME')

"""
このツールで使われる設定群です
"""

#[FORMAT]
DATE_FORMAT_PATTERN = '%Y.%m.%d_%H%M'

#[INOUT OUTPUT]
#SOURCE_FOLDER = join(HOME, 'Desktop/ダウンロード/KIRIN/00_input')
#OUTPUT_FOLDER = join(HOME, 'Desktop/ダウンロード/KIRIN/00_output')

#[TEST]
SOURCE_FOLDER = join(HOME, 'Desktop/mercre/test/src')
OUTPUT_FOLDER = join(HOME, 'Desktop/mercre/test/out')

#[MESSAGE]
COMPLETE_MESSAGE =\
'''~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .--  .-.  |-.-  .--. | |~~| --|-- |~~|
(    |   | | | | |  | | |--    |   |--
 "--  "-"  ' ' ' |""" | =--    |=  =--
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'''

#[IGNORE]
IGNORE = []
<commit_msg>[update] ignoreに使われそうなファイルを設定<commit_after>#!/usr/bin/env python
# -*- coding:utf-8 -*-
#---------------------------------------------------------------------------
# Name:        config.py
# Purpose:     各種設定
#
# Author:      Kosuke Akizuki
#
# Created:     2015/03/11
# Copyright:   (c) Kosuke Akizuki 2015
# Licence:     The MIT License (MIT)
#---------------------------------------------------------------------------

__author__ = "Kosuke Akizuki <slife1080@gmail.com>"
__status__ = "OK"
__version__ = "1.0"
__date__    = "11 Mar 2015"

import os

def join(*args):
    path = ''
    for e in args:
        path = os.path.join(path, e)
    path = os.path.normcase(path)
    path = os.path.normpath(path)
    return path

HOME = os.environ.get('HOME')

"""
このツールで使われる設定群です
"""

#[FORMAT]
DATE_FORMAT_PATTERN = '%Y.%m.%d_%H%M'

#[INOUT OUTPUT]
#SOURCE_FOLDER = join(HOME, 'Desktop/ダウンロード/KIRIN/00_input')
#OUTPUT_FOLDER = join(HOME, 'Desktop/ダウンロード/KIRIN/00_output')

#[TEST]
SOURCE_FOLDER = join(HOME, 'Desktop/mercre/test/src')
OUTPUT_FOLDER = join(HOME, 'Desktop/mercre/test/out')

#[MESSAGE]
COMPLETE_MESSAGE =\
'''~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .--  .-.  |-.-  .--. | |~~| --|-- |~~|
(    |   | | | | |  | | |--    |   |--
 "--  "-"  ' ' ' |""" | =--    |=  =--
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'''

#[IGNORE]
IGNORE = [
    '.*',
    'Thumbs.db',
    'nbproject'
]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import unittest

from nesasm.compiler import lexical, syntax

'''
Those tests are based on examples from:
http://nesdev.parodius.com/6502guid.txt
'''

class GuideTest(unittest.TestCase):

    def test_example_16_bit_subtraction_routine(self):
        ex_2 = '''
            SEC         ;clear the carry
            LDA $20     ;get the low byte of the first number
            SBC $22     ;add to it the low byte of the second
            STA $24     ;store in the low byte of the result
            LDA $21     ;get the high byte of the first number
            SBC $23     ;add to it the high byte of the second, plus carry
            STA $25     ;store in high byte of the result
            '''
        tokens = list(lexical(ex_2))
        self.assertEquals(21, len(tokens))
        ast = syntax(tokens)
        self.assertEquals(7, len(ast))

    def test_example_4_2(self):
        example_4_2 = '''
        ; Example 4-2.  Deleting an entry from an unordered list
        ;
        ; Delete the contents of $2F from a list whose starting
        ; address is in $30 and $31.  The first byte of the list
        ; is its length.
        ;

        deluel: LDY #$00     ; fetch element count
                LDA ($30),Y
                TAX          ; transfer length to X
                LDA $2F      ; item to delete
        nextel: INY          ; index to next element
                CMP ($30),Y  ; do entry and element match?
                BEQ delete   ; yes. delete element
                DEX          ; no. decrement element count
                BNE nextel   ; any more elements to compare?
                RTS          ; no. element not in list. done

        ; delete an element by moving the ones below it up one location

        delete: DEX          ; decrement element count
                BEQ deccnt   ; end of list?
                INY          ; no. move next element up
                LDA ($30),Y
                DEY
                STA ($30),Y
                INY
                JMP delete
        deccnt: LDA ($30,X)  ; update element count of list
                SBC #$01
                STA ($30,X)
                RTS
        '''
        tokens = list(lexical(example_4_2))
        self.assertEquals(96, len(tokens))

    def test_example_5_6(self):
        """
        example_5_6 = '''
        ; Example 5-6.  16-bit by 16-bit unsigned multiply
        ;
        ; Multiply $22 (low) and $23 (high) by $20 (low) and
        ; $21 (high) producing a 32-bit result in $24 (low) to $27 (high)
        ;

        mlt16:  LDA #$00     ; clear p2 and p3 of product
                STA $26
                STA $27
                LDX #$16     ; multiplier bit count = 16
        nxtbt:  LSR $21      ; shift two-byte multiplier right
                ROR $20
                BCC align    ; multiplier = 1?
                LDA $26      ; yes. fetch p2
                CLC
                ADC $22      ; and add m0 to it
                STA $26      ; store new p2
                LDA $27      ; fetch p3
                ADC $23      ; and add m1 to it
        align:  ROR A        ; rotate four-byte product right
                STA $27      ; store new p3
                ROR $26
                ROR $25
                ROR $24
                DEX          ; decrement bit count
                BNE nxtbt    ; loop until 16 bits are done
                RTS
        '''
        # TODO ROR A?
        # tokens = list(lexical(example_5_6))
        """

    def test_example_5_14(self):
        example_5_14 = '''
        ; Example 5-14.  Simple 16-bit square root.
        ;
        ; Returns the 8-bit square root in $20 of the
        ; 16-bit number in $20 (low) and $21 (high). The
        ; remainder is in location $21.

        sqrt16: LDY #$01     ; lsby of first odd number = 1
                STY $22
                DEY
                STY $23      ; msby of first odd number (sqrt = 0)
        again:  SEC
                LDA $20      ; save remainder in X register
                TAX          ; subtract odd lo from integer lo
                SBC $22
                STA $20
                LDA $21      ; subtract odd hi from integer hi
                SBC $23
                STA $21      ; is subtract result negative?
                BCC nomore   ; no. increment square root
                INY
                LDA $22      ; calculate next odd number
                ADC #$01
                STA $22
                BCC again
                INC $23
                JMP again
        nomore: STY $20      ; all done, store square root
                STX $21      ; and remainder
                RTS
        '''
        tokens = list(lexical(example_5_14))
        self.assertEquals(74, len(tokens))

        self.assertEquals('T_ENDLINE', tokens[0]['type'])
        self.assertEquals('T_ENDLINE', tokens[1]['type'])
        self.assertEquals('T_ENDLINE', tokens[2]['type'])
        self.assertEquals('T_ENDLINE', tokens[3]['type'])
        self.assertEquals('T_ENDLINE', tokens[4]['type'])
        self.assertEquals('T_ENDLINE', tokens[5]['type'])
        self.assertEquals('T_ENDLINE', tokens[6]['type'])

        self.assertEquals('T_LABEL', tokens[7]['type'])
        self.assertEquals('T_INSTRUCTION', tokens[8]['type'])
        self.assertEquals('T_HEX_NUMBER', tokens[9]['type'])
        self.assertEquals('T_ENDLINE', tokens[10]['type'])

        self.assertEquals('T_INSTRUCTION', tokens[11]['type'])
        self.assertEquals('T_ADDRESS', tokens[12]['type'])
        self.assertEquals('T_ENDLINE', tokens[13]['type'])

        self.assertEquals('T_INSTRUCTION', tokens[14]['type'])
        self.assertEquals('T_ENDLINE', tokens[15]['type'])

        self.assertEquals('T_INSTRUCTION', tokens[16]['type'])
        self.assertEquals('T_ADDRESS', tokens[17]['type'])
        self.assertEquals('T_ENDLINE', tokens[18]['type'])
<commit_msg>fixed docstring on guide_test<commit_after># -*- coding: utf-8 -*-
'''
Those tests are based on examples from:
http://nesdev.parodius.com/6502guid.txt
'''

import unittest
from nesasm.compiler import lexical, syntax

class GuideTest(unittest.TestCase):

    def test_example_16_bit_subtraction_routine(self):
        ex_2 = '''
            SEC         ;clear the carry
            LDA $20     ;get the low byte of the first number
            SBC $22     ;add to it the low byte of the second
            STA $24     ;store in the low byte of the result
            LDA $21     ;get the high byte of the first number
            SBC $23     ;add to it the high byte of the second, plus carry
            STA $25     ;store in high byte of the result
            '''
        tokens = list(lexical(ex_2))
        self.assertEquals(21, len(tokens))
        ast = syntax(tokens)
        self.assertEquals(7, len(ast))

    def test_example_4_2(self):
        example_4_2 = '''
        ; Example 4-2.  Deleting an entry from an unordered list
        ;
        ; Delete the contents of $2F from a list whose starting
        ; address is in $30 and $31.  The first byte of the list
        ; is its length.
        ;

        deluel: LDY #$00     ; fetch element count
                LDA ($30),Y
                TAX          ; transfer length to X
                LDA $2F      ; item to delete
        nextel: INY          ; index to next element
                CMP ($30),Y  ; do entry and element match?
                BEQ delete   ; yes. delete element
                DEX          ; no. decrement element count
                BNE nextel   ; any more elements to compare?
                RTS          ; no. element not in list. done

        ; delete an element by moving the ones below it up one location

        delete: DEX          ; decrement element count
                BEQ deccnt   ; end of list?
                INY          ; no. move next element up
                LDA ($30),Y
                DEY
                STA ($30),Y
                INY
                JMP delete
        deccnt: LDA ($30,X)  ; update element count of list
                SBC #$01
                STA ($30,X)
                RTS
        '''
        tokens = list(lexical(example_4_2))
        self.assertEquals(96, len(tokens))

    def test_example_5_6(self):
        """
        example_5_6 = '''
        ; Example 5-6.  16-bit by 16-bit unsigned multiply
        ;
        ; Multiply $22 (low) and $23 (high) by $20 (low) and
        ; $21 (high) producing a 32-bit result in $24 (low) to $27 (high)
        ;

        mlt16:  LDA #$00     ; clear p2 and p3 of product
                STA $26
                STA $27
                LDX #$16     ; multiplier bit count = 16
        nxtbt:  LSR $21      ; shift two-byte multiplier right
                ROR $20
                BCC align    ; multiplier = 1?
                LDA $26      ; yes. fetch p2
                CLC
                ADC $22      ; and add m0 to it
                STA $26      ; store new p2
                LDA $27      ; fetch p3
                ADC $23      ; and add m1 to it
        align:  ROR A        ; rotate four-byte product right
                STA $27      ; store new p3
                ROR $26
                ROR $25
                ROR $24
                DEX          ; decrement bit count
                BNE nxtbt    ; loop until 16 bits are done
                RTS
        '''
        # TODO ROR A?
        # tokens = list(lexical(example_5_6))
        """

    def test_example_5_14(self):
        example_5_14 = '''
        ; Example 5-14.  Simple 16-bit square root.
        ;
        ; Returns the 8-bit square root in $20 of the
        ; 16-bit number in $20 (low) and $21 (high). The
        ; remainder is in location $21.

        sqrt16: LDY #$01     ; lsby of first odd number = 1
                STY $22
                DEY
                STY $23      ; msby of first odd number (sqrt = 0)
        again:  SEC
                LDA $20      ; save remainder in X register
                TAX          ; subtract odd lo from integer lo
                SBC $22
                STA $20
                LDA $21      ; subtract odd hi from integer hi
                SBC $23
                STA $21      ; is subtract result negative?
                BCC nomore   ; no. increment square root
                INY
                LDA $22      ; calculate next odd number
                ADC #$01
                STA $22
                BCC again
                INC $23
                JMP again
        nomore: STY $20      ; all done, store square root
                STX $21      ; and remainder
                RTS
        '''
        tokens = list(lexical(example_5_14))
        self.assertEquals(74, len(tokens))

        self.assertEquals('T_ENDLINE', tokens[0]['type'])
        self.assertEquals('T_ENDLINE', tokens[1]['type'])
        self.assertEquals('T_ENDLINE', tokens[2]['type'])
        self.assertEquals('T_ENDLINE', tokens[3]['type'])
        self.assertEquals('T_ENDLINE', tokens[4]['type'])
        self.assertEquals('T_ENDLINE', tokens[5]['type'])
        self.assertEquals('T_ENDLINE', tokens[6]['type'])

        self.assertEquals('T_LABEL', tokens[7]['type'])
        self.assertEquals('T_INSTRUCTION', tokens[8]['type'])
        self.assertEquals('T_HEX_NUMBER', tokens[9]['type'])
        self.assertEquals('T_ENDLINE', tokens[10]['type'])

        self.assertEquals('T_INSTRUCTION', tokens[11]['type'])
        self.assertEquals('T_ADDRESS', tokens[12]['type'])
        self.assertEquals('T_ENDLINE', tokens[13]['type'])

        self.assertEquals('T_INSTRUCTION', tokens[14]['type'])
        self.assertEquals('T_ENDLINE', tokens[15]['type'])

        self.assertEquals('T_INSTRUCTION', tokens[16]['type'])
        self.assertEquals('T_ADDRESS', tokens[17]['type'])
        self.assertEquals('T_ENDLINE', tokens[18]['type'])
<|endoftext|>


<commit_before>""" Configuration of app.config dict
"""
import os.path
import logging
log = logging.getLogger()

# flask-sqlalchemy
SQLALCHEMY_TRACK_MODIFICATIONS = False # explicitly remove deprecated feature

# logging
#path = app.config['LOGGING_PATH']
#os.makedirs(path, exist_ok=True)
#info_fh = RotatingFileHandler(os.path.join(path, 'mixmind.log'), maxBytes=10000000, backupCount=3)
#info_fh.setFormatter(logging.Formatter('%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]'))
#info_fh.setLevel(logging.INFO)
#log = logging.getLogger()

# mix-mind
MIXMIND_RECIPES_DIR = "recipes/"
MIXMIND_INGREDIENTS_DIR = "ingredients/"

MIXMIND_DEFAULT_RECIPES = ["recipes_schubar.json", "IBA_all.json"]
MIXMIND_DEFAULT_INGREDIENTS = ["12BBplus.csv"]

def get_recipe_files():
    return get_checked_files(MIXMIND_RECIPES_DIR, MIXMIND_DEFAULT_RECIPES)

def get_ingredient_files():
    return get_checked_files(MIXMIND_INGREDIENTS_DIR, MIXMIND_DEFAULT_INGREDIENTS)

def get_checked_files(path, files):
    files = [os.path.join(path, f) for f in files]
    for f in files:
        if not os.path.isfile(f):
            log.warning("{} not found, skipping".format(f))
    return files
<commit_msg>Use abspath<commit_after>""" Configuration of app.config dict
"""
import os.path
import logging
log = logging.getLogger()

# flask-sqlalchemy
SQLALCHEMY_TRACK_MODIFICATIONS = False # explicitly remove deprecated feature

# logging
#path = app.config['LOGGING_PATH']
#os.makedirs(path, exist_ok=True)
#info_fh = RotatingFileHandler(os.path.join(path, 'mixmind.log'), maxBytes=10000000, backupCount=3)
#info_fh.setFormatter(logging.Formatter('%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]'))
#info_fh.setLevel(logging.INFO)
#log = logging.getLogger()

# mix-mind
MIXMIND_RECIPES_DIR = "recipes/"
MIXMIND_INGREDIENTS_DIR = "ingredients/"

MIXMIND_DEFAULT_RECIPES = ["recipes_schubar.json", "IBA_all.json"]
MIXMIND_DEFAULT_INGREDIENTS = ["12BBplus.csv"]

def get_recipe_files():
    return get_checked_files(MIXMIND_RECIPES_DIR, MIXMIND_DEFAULT_RECIPES)

def get_ingredient_files():
    return get_checked_files(MIXMIND_INGREDIENTS_DIR, MIXMIND_DEFAULT_INGREDIENTS)

def get_checked_files(path, files):
    files = [os.path.abspath(os.path.join(path, f)) for f in files]
    for f in files:
        if not os.path.isfile(f):
            log.warning("{} not found, skipping".format(f))
    return files
<|endoftext|>


<commit_before>#    plot_xye.py
#
#    Created by Stefano Martiniani in June 2013, stefano.martiniani@gmail.com
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import numpy as np
import copy
from itertools import cycle
import matplotlib
from matplotlib import rc
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

#####################LINE STYLE CYCLER####################
lines = ["-","--","-.",":"]
linecycler = cycle(lines)
##########################################################

#######################SET LATEX OPTIONS###################
rc('text', usetex=True)
rc('font',**{'family':'serif','serif':['Computer Modern']})
#rc('text.latex',preamble=r'\usepackage{times}')
plt.rcParams.update({'font.size': 12})
plt.rcParams['xtick.major.pad'] = 8
plt.rcParams['ytick.major.pad'] = 8
##########################################################

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="load energy intervals and compute Cv stdev",
                                     epilog="must write file name followed by a label string, otherwise raise the flag --nolabels")
    parser.add_argument("fname", nargs="+", type=str, help="filenames with heat capacity followed by label")
    parser.add_argument("--legloc", type=int, help="define location of the legend (default upper right 1):"
                        " \"upper right=1\" \"upper left=2\" \"lower left=3\" \"lower right=4\" \"right=5\" "
                        "\"center left=6\" \"center right=7\" \"lower center=8\" \"upper center=9\" \"center=10\"",default=1)
    parser.add_argument("--xlabel", type=str, help="set x-label",default="T")
    parser.add_argument("--ylabel", type=str, help="set y-label",default="Cv")
    parser.add_argument("--xtop", type=float, help="set x-axis top",default=None)
    parser.add_argument("--xbot", type=float, help="set x-axis bottom",default=None)
    parser.add_argument("--ytop", type=float, help="set y-axis top",default=None)
    parser.add_argument("--ybot", type=float, help="set y-axis bottom",default=None)
    parser.add_argument("--title", type=str, help="set title",default=None)
    parser.add_argument("--linewidth", type=float, help="set line width (default 1.8)",default=1.8)
    parser.add_argument("--colormap", type=str, help="set colormap (default Dark2)",default='Dark2')
    parser.add_argument("--ecolor", type=str, help="set error bars color (default light gray)",default='g')
    parser.add_argument("--ecap", type=float, help="set error bars cap size (default is None)",default=None)
    parser.add_argument("--output", type=str, help="set output name (default cv_plot)",default='cv_plot')
    parser.add_argument("--filetype", type=str, help="set output file format (default eps)",default='eps')
    parser.add_argument("--nolabels", action="store_true", help="turn on if do not want to insert labels",default=False)
    parser.add_argument("--show", action="store_true", help="show explorable plot preview",default=False)
    parser.add_argument("--ebar", action="store_true", help="plot error bars if in data file",default=False)
    parser.add_argument("--grid", action="store_true", help="plot grid",default=False)
    args = parser.parse_args()

    fname = args.fname
    legloc=args.legloc
    nolabels = args.nolabels
    xlabel = args.xlabel
    ylabel = args.ylabel
    show = args.show
    title = args.title
    grid = args.grid
    linew = args.linewidth
    xtop = args.xtop
    xbot = args.xbot
    ytop = args.ytop
    ybot = args.ybot
    colormap = args.colormap
    errcolor = args.ecolor
    ebar = args.ebar
    filename = args.output
    filetype = args.filetype
    ecap = args.ecap

    ####################################################################################################
    #DEAL WITH INPUT

    if nolabels is False:
        all_data = [[] for i in xrange(len(fname)/2)]
        all_labels = np.array(["no_name" for i in xrange(len(fname)/2)])
        j = 0
        for name,i in zip(fname,xrange(len(fname))):
            if ((i+1) % 2) is not 0:
                data = np.genfromtxt(name)
                dshape = np.shape(data)
                if dshape[1] > 3:
                    new_data = np.zeros([dshape[0],3])
                    for k in xrange(dshape[0]):
                        for l in xrange(3):
                            new_data[k][l] = data[k][l]
                    data = new_data
                elif dshape[1] < 3:
                    data = np.hstack((data, np.zeros((data.shape[0], 1), dtype=data.dtype))) ##adds an error of 0 if there's no error associated
                all_data[j] = data
            else:
                all_labels[j] = name
                j += 1
    else:
        all_data = [[] for i in xrange(len(fname))]
        all_labels = np.array(["no_name" for i in xrange(len(fname))])
        for name,i in zip(fname,xrange(len(fname))):
            data = np.genfromtxt(name)
            dshape = np.shape(data)
            if dshape[1] > 3:
                new_data = np.zeros([dshape[0],3])
                for k in xrange(dshape[0]):
                    for l in xrange(3):
                       new_data[k][l] = data[k][l]
                data = new_data
            elif dshape[1] < 3:
                data = np.hstack((data, np.zeros((data.shape[0], 1), dtype=data.dtype)))
            data = np.array(data)
            all_data[i] = data

    all_data = np.array(all_data)
    ######################################################################################################

    #####################
    ####### PLOT ########
    #####################

    ####SET COLOUR MAP######
    cm = plt.get_cmap(colormap)
    ########################
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_color_cycle([cm(1.*i/len(all_data)) for i in xrange(np.shape(all_data)[0])])
    for data, label, i in zip(all_data, all_labels, xrange(np.shape(all_data)[0])):
            ax.plot(data[:,0], data[:,1], next(linecycler), label = label, linewidth=linew)
            if np.shape(data)[1] is 3 and ebar is True:
                ax.errorbar(data[:,0], data[:,1], yerr=data[:,2], ecolor=errcolor, capsize=ecap )
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xlim((xbot,xtop))
    ax.set_ylim((ybot,ytop))

    if grid is True:
        ax.grid(True,linestyle='--',color='0.75')
    if title is not None:
        ax.set_title(title)
    if nolabels is False:
        ax.legend(frameon=False,loc=legloc)
    if show is True:
         plt.show()

    fig.savefig('{F}.{T}'.format(F=filename,T=filetype))
    <commit_msg>added all straight lines option to plot_cv<commit_after>#    plot_xye.py
#
#    Created by Stefano Martiniani in June 2013, stefano.martiniani@gmail.com
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import numpy as np
import copy
from itertools import cycle
import matplotlib
from matplotlib import rc
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

#######################SET LATEX OPTIONS###################
rc('text', usetex=True)
rc('font',**{'family':'serif','serif':['Computer Modern']})
#rc('text.latex',preamble=r'\usepackage{times}')
plt.rcParams.update({'font.size': 12})
plt.rcParams['xtick.major.pad'] = 8
plt.rcParams['ytick.major.pad'] = 8
##########################################################

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="load energy intervals and compute Cv stdev",
                                     epilog="must write file name followed by a label string, otherwise raise the flag --nolabels")
    parser.add_argument("fname", nargs="+", type=str, help="filenames with heat capacity followed by label")
    parser.add_argument("--legloc", type=int, help="define location of the legend (default upper right 1):"
                        " \"upper right=1\" \"upper left=2\" \"lower left=3\" \"lower right=4\" \"right=5\" "
                        "\"center left=6\" \"center right=7\" \"lower center=8\" \"upper center=9\" \"center=10\"",default=1)
    parser.add_argument("--xlabel", type=str, help="set x-label",default="T")
    parser.add_argument("--ylabel", type=str, help="set y-label",default="Cv")
    parser.add_argument("--xtop", type=float, help="set x-axis top",default=None)
    parser.add_argument("--xbot", type=float, help="set x-axis bottom",default=None)
    parser.add_argument("--ytop", type=float, help="set y-axis top",default=None)
    parser.add_argument("--ybot", type=float, help="set y-axis bottom",default=None)
    parser.add_argument("--title", type=str, help="set title",default=None)
    parser.add_argument("--linewidth", type=float, help="set line width (default 1.8)",default=1.8)
    parser.add_argument("--colormap", type=str, help="set colormap (default Dark2)",default='Dark2')
    parser.add_argument("--ecolor", type=str, help="set error bars color (default light gray)",default='g')
    parser.add_argument("--ecap", type=float, help="set error bars cap size (default is None)",default=None)
    parser.add_argument("--output", type=str, help="set output name (default cv_plot)",default='cv_plot')
    parser.add_argument("--filetype", type=str, help="set output file format (default eps)",default='eps')
    parser.add_argument("--nolabels", action="store_true", help="turn on if do not want to insert labels",default=False)
    parser.add_argument("--show", action="store_true", help="show explorable plot preview",default=False)
    parser.add_argument("--ebar", action="store_true", help="plot error bars if in data file",default=False)
    parser.add_argument("--grid", action="store_true", help="plot grid",default=False)
    parser.add_argument("--lines", action="store_true", help="plot all straight lines",default=False)
    args = parser.parse_args()

    fname = args.fname
    legloc=args.legloc
    nolabels = args.nolabels
    xlabel = args.xlabel
    ylabel = args.ylabel
    show = args.show
    title = args.title
    grid = args.grid
    linew = args.linewidth
    xtop = args.xtop
    xbot = args.xbot
    ytop = args.ytop
    ybot = args.ybot
    colormap = args.colormap
    errcolor = args.ecolor
    ebar = args.ebar
    filename = args.output
    filetype = args.filetype
    ecap = args.ecap
    slines = args.lines
    #####################LINE STYLE CYCLER####################
    if slines is False
        lines = ["-","--","-.",":"]
    else:
        lines = ["-"]
    linecycler = cycle(lines)
    ##########################################################

    ####################################################################################################
    #DEAL WITH INPUT

    if nolabels is False:
        all_data = [[] for i in xrange(len(fname)/2)]
        all_labels = np.array(["no_name" for i in xrange(len(fname)/2)])
        j = 0
        for name,i in zip(fname,xrange(len(fname))):
            if ((i+1) % 2) is not 0:
                data = np.genfromtxt(name)
                dshape = np.shape(data)
                if dshape[1] > 3:
                    new_data = np.zeros([dshape[0],3])
                    for k in xrange(dshape[0]):
                        for l in xrange(3):
                            new_data[k][l] = data[k][l]
                    data = new_data
                elif dshape[1] < 3:
                    data = np.hstack((data, np.zeros((data.shape[0], 1), dtype=data.dtype))) ##adds an error of 0 if there's no error associated
                all_data[j] = data
            else:
                all_labels[j] = name
                j += 1
    else:
        all_data = [[] for i in xrange(len(fname))]
        all_labels = np.array(["no_name" for i in xrange(len(fname))])
        for name,i in zip(fname,xrange(len(fname))):
            data = np.genfromtxt(name)
            dshape = np.shape(data)
            if dshape[1] > 3:
                new_data = np.zeros([dshape[0],3])
                for k in xrange(dshape[0]):
                    for l in xrange(3):
                       new_data[k][l] = data[k][l]
                data = new_data
            elif dshape[1] < 3:
                data = np.hstack((data, np.zeros((data.shape[0], 1), dtype=data.dtype)))
            data = np.array(data)
            all_data[i] = data

    all_data = np.array(all_data)
    ######################################################################################################

    #####################
    ####### PLOT ########
    #####################

    ####SET COLOUR MAP######
    cm = plt.get_cmap(colormap)
    ########################
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_color_cycle([cm(1.*i/len(all_data)) for i in xrange(np.shape(all_data)[0])])
    for data, label, i in zip(all_data, all_labels, xrange(np.shape(all_data)[0])):
            ax.plot(data[:,0], data[:,1], next(linecycler), label = label, linewidth=linew)
            if np.shape(data)[1] is 3 and ebar is True:
                ax.errorbar(data[:,0], data[:,1], yerr=data[:,2], ecolor=errcolor, capsize=ecap )
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xlim((xbot,xtop))
    ax.set_ylim((ybot,ytop))

    if grid is True:
        ax.grid(True,linestyle='--',color='0.75')
    if title is not None:
        ax.set_title(title)
    if nolabels is False:
        ax.legend(frameon=False,loc=legloc)
    if show is True:
         plt.show()

    fig.savefig('{F}.{T}'.format(F=filename,T=filetype))
    <|endoftext|>