<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from pylab import *

# =============================================================================
# Recursively prints the tree structure of the matplotlib object
def print_tree( obj, indent = "" ):
    print indent, type(obj)
    for child in obj.get_children():
        print_tree( child, indent + "   " )
    return
# =============================================================================
def matplotlib2tikz( filepath, figurewidth=None, figureheight=None, tex_relative_path_to_data=None ):
    from os import path

    global fwidth
    fwidth = figurewidth
    global fheight
    fheight = figureheight
    global rel_data_path
    rel_data_path = tex_relative_path_to_data

    global pgfplots_libs
    pgfplots_libs = []

    global output_dir
    output_dir = path.dirname(filepath)

    # open file
    file_handle = open( filepath, "w" )

    # write the contents
    file_handle.write( "\\begin{tikzpicture}\n\n" )
    handle_children( file_handle, gcf() )
    file_handle.write( "\\end{tikzpicture}" )

    # close file
    file_handle.close()

    # print message about necessary pgfplot libs to command line
    print_pgfplot_libs_message()
    return
# =============================================================================
def draw_axes( file_handle, obj ):
    # Are we dealing with an axis that hosts a colorbar?
    # Skip then.
    # TODO instead of testing here, rather blacklist the colorbar axis
    #      plots as soon as they have been found, e.g., by find_associated_colorbar()
    if extract_colorbar(obj)!=None:
        return

    # instantiation
    nsubplots = 0
    subplot_index = 0
    issubplot = False

    if isinstance( obj, matplotlib.axes.Subplot ):
        issubplot = True
        geom = obj.get_geometry()
        nsubplots = geom[0]*geom[1]
        subplot_index = geom[2]
        if subplot_index==1:
            file_handle.write( "\\begin{groupplot}[group style={group size=%.d by %.d}]\n" % (geom[1],geom[0]) )
            pgfplots_libs.append( "groupplots" )

    axis_options = []

    if fwidth or fheight:
        axis_options.append( "scale only axis" )
    if fwidth:
        axis_options.append( "width="+fwidth )
    if fheight:
        axis_options.append( "height="+fheight )

    # get axes title
    title = obj.get_title()
    if len(title) != 0 :
        axis_options.append( "title={" + title + "}" )

    # get axes titles
    xlabel = obj.get_xlabel()
    if len(xlabel) != 0 :
        axis_options.append( "xlabel={" + xlabel + "}" )
    ylabel = obj.get_ylabel()
    if len(ylabel) != 0 :
        axis_options.append( "ylabel={" + ylabel + "}" )

    # axes limits
    xlim = obj.get_xlim()
    axis_options.append(     "xmin=" + repr(xlim[0])
                         + ", xmax=" + repr(xlim[1]) )
    ylim = obj.get_ylim()
    axis_options.append(     "ymin=" + repr(ylim[0])
                         + ", ymax=" + repr(ylim[1]) )

    # find color bar
    colorbar = find_associated_colorbar( obj )
    if colorbar!=None:
        clim = colorbar.get_clim()
        axis_options.append( "colorbar=true" )
        mycolormap = matplotlibColormap2pgfplotsColormap( colorbar.get_cmap() )
        axis_options.append( "colormap/" + mycolormap )
        axis_options.append( 'point meta min=' + repr(clim[0]) )
        axis_options.append( 'point meta max=' + repr(clim[1]) )

    # actually print the thing
    if issubplot:
        file_handle.write( "\\nextgroupplot" )
    else:
        file_handle.write( "\\begin{axis}" )

    if len(axis_options)!=0:
        options = ",\n".join( axis_options )
        file_handle.write( "[\n" + options + "\n]\n" )

    # TODO Use get_lines()?
    handle_children( file_handle, obj )

    if not issubplot:
        file_handle.write( "\\end{axis}\n\n" )
    elif issubplot and nsubplots==subplot_index:
        file_handle.write( "\\end{groupplot}\n\n" )

    return
# =============================================================================
# Converts a color map as given in matplotlib to a color map as represented
# in Pgfplots.
def matplotlibColormap2pgfplotsColormap( cmap ):
    if isinstance( cmap, matplotlib.colors.LinearSegmentedColormap ):
        if cmap.is_gray():
            return 'blackwhite'
        else:
            segdata = cmap._segmentdata
            red    = segdata['red']
            green  = segdata['green']
            blue   = segdata['blue']
            print red
            print green
            print blue
            # loop over the data, stop at each spot where the linear
            # interpolations is interrupted, and set a color mark there
            # set initial color
            k_red   = 0
            k_green = 0
            k_blue  = 0
            x = 0.0
            color_changes = []
            while True:
                # find next x
                x = min( red[k_red][0], green[k_green][0], blue[k_blue][0] )

                if ( red[k_red][0]==x ):
                    red_part = red[k_red][1]
                    k_red    = k_red+1
                else:
                    red_part = linear_interpolation( x, (red[k_red-1][0],red[k_red][0]), (red[k_red-1][2],red[k_red][1]) )

                if ( green[k_green][0]==x ):
                    green_part = green[k_green][1]
                    k_green    = k_green+1
                else:
                    green_part = linear_interpolation( x, (green[k_green-1][0],green[k_green][0]), (green[k_green-1][2],green[k_green][1]) )
                    #print green_part

                if ( blue[k_blue][0]==x ):
                    blue_part = blue[k_blue][1]
                    k_blue    = k_blue+1
                else:
                    blue_part = linear_interpolation( x, (blue[k_blue-1][0],blue[k_blue][0]), (blue[k_blue-1][2],blue[k_blue][1]) )

                #print k
                color_changes.append( "rgb(%.3fcm)=(%.3f,%.3f,%.3f)" % ( x, red_part,green_part,blue_part ) )

                if x>=1.0: break

            return "{mymap}{" + "; ".join( color_changes ) + "}"
    else :
        print "Don't know how to handle color map. Using 'blackwhite'."
        return 'blackwhite'
# =============================================================================
def linear_interpolation( x, X, Y ):
    xRet = ( Y[1]*(x-X[0]) + Y[0]*(X[1]-x) ) / ( X[1]-X[0] )
    return xRet
# =============================================================================
def draw_line2d( file_handle, obj ):
    addplot_options = []

    marker = obj.get_marker()
    if marker!="None":
        print "%TODO Implement markers."

    # get line color
    color = obj.get_color()
    xcolor = mpl_color2xcolor(color)
    if xcolor:
        addplot_options.append( xcolor )

    linestyle = mpl_linestyle2pgfp_linestyle( obj.get_linestyle() )
    if linestyle:
        addplot_options.append( linestyle )

    # process options
    file_handle.write( "\\addplot " )
    if addplot_options:
        options = ", ".join( addplot_options )
        file_handle.write( "[" + options + "]\n" )

    # print the hard numerical data
    xdata, ydata = obj.get_data()
    file_handle.write(  "coordinates {\n" )
    for k in range(len(xdata)):
        file_handle.write( "(%.15g,%.15g) " % (xdata[k], ydata[k]) )
    file_handle.write( "\n};\n" )

    return
# =============================================================================
# Translates a line style of matplotlib to the corresponding style
# in Pgfplots.
def mpl_linestyle2pgfp_linestyle( ls ):
    if (ls=='-' or ls=='None'):
        return None
    elif ls==':':
        return 'dotted'
    else:
        print '%Unknown line style \"' + ls + '\".'
        return None
# =============================================================================
def draw_image( file_handle, obj ):
    from os import path

    global img_number
    try:
       img_number = img_number+1
    except NameError, e:
       img_number = 0

    filename = path.join( output_dir,
                          "img" + repr(img_number) + ".png" )

    # store the image as in a file
    img_array = obj.get_array()
    dims = img_array.shape
    if len(dims)==2: # the values are given as one real number: look at cmap
        clims = obj.get_clim()
        imsave( fname = filename,
                arr   = img_array ,
                cmap  = obj.get_cmap(),
                vmin  = clims[0],
                vmax  = clims[1] )
    elif len(dims)==3: # RGB information
        if dims[2]==4: # RGB+alpha information at each point
            import Image
            # convert to PIL image (after upside-down flip)
            im = Image.fromarray( flipud(img_array), 'RGBA' )
            im.save( filename )

    # write the corresponding information to the TikZ file
    extent = obj.get_extent()
    rel_filepath = path.join( rel_data_path,  path.basename(filename) )
    file_handle.write( ("\\addplot graphics [xmin=%.15g, xmax=%.15g, ymin=%.15g, ymax=%.15g] {" + rel_filepath + "};\n")
                        % extent )

    handle_children( file_handle, obj )

    return
# =============================================================================
def draw_polygon( file_handle, obj ):
    # TODO do nothing for polygons?!
    handle_children( file_handle, obj )
    return
# =============================================================================
# Rather poor way of telling whether an axis has a colorbar associated:
# Check the next axis environment, and see if it is de facto a color bar;
# if yes, return the color bar object.
def find_associated_colorbar( obj ):
      for child in obj.get_children():
              try:
                      cbar = child.colorbar
              except AttributeError:
                      continue
              if not cbar == None: # really necessary?
                      # if fetch was successful, cbar contains
                      # ( reference to colorbar, reference to axis containing colorbar )
                      return cbar[0]
      return None
# =============================================================================
# TODO This function contains crude logic.
def is_colorbar( obj ):
    if isinstance( obj, matplotlib.cm.ScalarMappable ):
        arr = obj.get_array()
        dims = arr.shape
        if len(dims)==1:
            return True # o rly?
        else:
            return False
    else:
        return False
# =============================================================================
def extract_colorbar( obj ):
    colorbars = findobj( obj, is_colorbar )
    if len(colorbars)==0:
        return None
    if not equivalent( colorbars ):
        print "More than one color bar found. Use first one."
    return colorbars[0]
# =============================================================================
# checks if the vectors consists of all the same objects
def equivalent( array ):
    if len(array)==0:
        return False
    else:
        for elem in array:
            if elem!=array[0]:
                return False

    return True
# =============================================================================
def draw_polycollection( file_handle, obj ):
    print "matplotlib2tikz: Don't know how to draw a PolyCollection."
    return
# =============================================================================
# Translates a matplotlib color specification into a proper LaTeX
# xcolor.
def mpl_color2xcolor( color ):
    try: # is the color a gray-scale value?
        alpha = float(color)
        if alpha==0.0:
            return 'black'
        elif alpha==0.5:
            return 'gray'
        elif alpha==0.75:
            return 'lightgray'
        elif alpha==1.0:
            return None
        else:
            return color
    except ValueError:
      pass

    if color=='r':
        return 'red'
    elif color=='g':
        return 'green'
    elif color=='b':
        return 'blue'
    elif color=='k':
        return None
    else:
        print '%Unknown color \"' + color + '\".'
        return None # default
# =============================================================================
def handle_children( file_handle, obj ):
    for child in obj.get_children():
        if ( isinstance( child, matplotlib.axes.Axes ) ):
            draw_axes( file_handle, child )
        elif ( isinstance( child, matplotlib.lines.Line2D ) ):
            draw_line2d( file_handle, child )
        elif ( isinstance( child, matplotlib.image.AxesImage ) ):
            draw_image( file_handle, child )
        elif ( isinstance( child, matplotlib.patches.Polygon ) ):
            draw_polygon( file_handle, child )
        elif ( isinstance( child, matplotlib.collections.PolyCollection ) ):
            draw_polycollection( file_handle, child )
        elif (   isinstance( child, matplotlib.axis.XAxis )
              or isinstance( child, matplotlib.axis.YAxis )
              or isinstance( child, matplotlib.spines.Spine )
              or isinstance( child, matplotlib.patches.Rectangle )
              or isinstance( child, matplotlib.text.Text )
             ):
            pass
        else:
            print "matplotlib2tikz: Don't know how to handle object ", type(child), "."
    return
# =============================================================================
def print_pgfplot_libs_message():
    # remove duplicates
    clean_pgfplots_libs = list(set(pgfplots_libs) )
    libs = ",".join( clean_pgfplots_libs )

    print "========================================================="
    print "Please add the following line to your LaTeX preamble:\n"
    print "\usepackage{pgfplots}"
    if len(clean_pgfplots_libs)!=0:
        print "\usepgfplotslibrary{" + libs + "}"
    print "========================================================="

    return
# =============================================================================
<commit_msg>initial support for x/y ticks and tick labels<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from pylab import *

# =============================================================================
# Recursively prints the tree structure of the matplotlib object
def print_tree( obj, indent = "" ):
    print indent, type(obj)
    for child in obj.get_children():
        print_tree( child, indent + "   " )
    return
# =============================================================================
def matplotlib2tikz( filepath, figurewidth=None, figureheight=None, tex_relative_path_to_data=None ):
    from os import path

    global fwidth
    fwidth = figurewidth
    global fheight
    fheight = figureheight
    global rel_data_path
    rel_data_path = tex_relative_path_to_data

    global pgfplots_libs
    pgfplots_libs = []

    global output_dir
    output_dir = path.dirname(filepath)

    # open file
    file_handle = open( filepath, "w" )

    # write the contents
    file_handle.write( "\\begin{tikzpicture}\n\n" )
    handle_children( file_handle, gcf() )
    file_handle.write( "\\end{tikzpicture}" )

    # close file
    file_handle.close()

    # print message about necessary pgfplot libs to command line
    print_pgfplot_libs_message()
    return
# =============================================================================
def draw_axes( file_handle, obj ):
    # Are we dealing with an axis that hosts a colorbar?
    # Skip then.
    # TODO instead of testing here, rather blacklist the colorbar axis
    #      plots as soon as they have been found, e.g., by find_associated_colorbar()
    if extract_colorbar(obj)!=None:
        return

    # instantiation
    nsubplots = 0
    subplot_index = 0
    issubplot = False

    if isinstance( obj, matplotlib.axes.Subplot ):
        issubplot = True
        geom = obj.get_geometry()
        nsubplots = geom[0]*geom[1]
        subplot_index = geom[2]
        if subplot_index==1:
            file_handle.write( "\\begin{groupplot}[group style={group size=%.d by %.d}]\n" % (geom[1],geom[0]) )
            pgfplots_libs.append( "groupplots" )

    axis_options = []

    if fwidth or fheight:
        axis_options.append( "scale only axis" )
    if fwidth:
        axis_options.append( "width="+fwidth )
    if fheight:
        axis_options.append( "height="+fheight )

    # get axes title
    title = obj.get_title()
    if len(title) != 0 :
        axis_options.append( "title={" + title + "}" )

    # get axes titles
    xlabel = obj.get_xlabel()
    if len(xlabel) != 0 :
        axis_options.append( "xlabel={" + xlabel + "}" )
    ylabel = obj.get_ylabel()
    if len(ylabel) != 0 :
        axis_options.append( "ylabel={" + ylabel + "}" )

    # get x ticks
    xticks = obj.get_xticks()
    xticklabels = obj.get_xticklabels()
    pgfplots_xtick = []
    pgfplots_xticklabel = []
    for i in range(0, len(xticks)):
        pgfplots_xtick.append( xticks[i] )
        if len( xticklabels[i].get_text() ) != 0:
            pgfplots_xticklabel.append( xticklabels[i] )

    if len(pgfplots_xtick)==0:
        axis_options.append( "xtick=\\empty" )
    else:
        axis_options.append( "xtick={" +  ",".join(["%s" % el for el in pgfplots_xtick]) + "}" )

    if len(pgfplots_xticklabel)!=0:
        axis_options.append( "xticklabel={" + ",".join(pgfplots_xticklabel) + "}" )

    # get y ticks
    yticks = obj.get_yticks()
    yticklabels = obj.get_yticklabels()
    pgfplots_ytick = []
    pgfplots_yticklabel = []
    for i in range(0, len(yticks)):
        pgfplots_ytick.append( yticks[i] )
        if len( yticklabels[i].get_text() ) != 0:
            pgfplots_yticklabel.append( yticklabels[i] )

    if len(pgfplots_ytick)==0:
        axis_options.append( "ytick=\\empty" )
    else:
        axis_options.append( "ytick={" +  ",".join(["%s" % el for el in pgfplots_ytick]) + "}" )

    if len(pgfplots_yticklabel)!=0:
        axis_options.append( "yticklabel={" + ",".join(pgfplots_yticklabel) + "}" )

    # axes limits
    xlim = obj.get_xlim()
    axis_options.append(     "xmin=" + repr(xlim[0])
                         + ", xmax=" + repr(xlim[1]) )
    ylim = obj.get_ylim()
    axis_options.append(     "ymin=" + repr(ylim[0])
                         + ", ymax=" + repr(ylim[1]) )

    # find color bar
    colorbar = find_associated_colorbar( obj )
    if colorbar!=None:
        clim = colorbar.get_clim()
        axis_options.append( "colorbar=true" )
        mycolormap = matplotlibColormap2pgfplotsColormap( colorbar.get_cmap() )
        axis_options.append( "colormap/" + mycolormap )
        axis_options.append( 'point meta min=' + repr(clim[0]) )
        axis_options.append( 'point meta max=' + repr(clim[1]) )

    # actually print the thing
    if issubplot:
        file_handle.write( "\\nextgroupplot" )
    else:
        file_handle.write( "\\begin{axis}" )

    if len(axis_options)!=0:
        options = ",\n".join( axis_options )
        file_handle.write( "[\n" + options + "\n]\n" )

    # TODO Use get_lines()?
    handle_children( file_handle, obj )

    if not issubplot:
        file_handle.write( "\\end{axis}\n\n" )
    elif issubplot and nsubplots==subplot_index:
        file_handle.write( "\\end{groupplot}\n\n" )

    return
# =============================================================================
# Converts a color map as given in matplotlib to a color map as represented
# in Pgfplots.
def matplotlibColormap2pgfplotsColormap( cmap ):
    if isinstance( cmap, matplotlib.colors.LinearSegmentedColormap ):
        if cmap.is_gray():
            return 'blackwhite'
        else:
            segdata = cmap._segmentdata
            red    = segdata['red']
            green  = segdata['green']
            blue   = segdata['blue']
            print red
            print green
            print blue
            # loop over the data, stop at each spot where the linear
            # interpolations is interrupted, and set a color mark there
            # set initial color
            k_red   = 0
            k_green = 0
            k_blue  = 0
            x = 0.0
            color_changes = []
            while True:
                # find next x
                x = min( red[k_red][0], green[k_green][0], blue[k_blue][0] )

                if ( red[k_red][0]==x ):
                    red_part = red[k_red][1]
                    k_red    = k_red+1
                else:
                    red_part = linear_interpolation( x, (red[k_red-1][0],red[k_red][0]), (red[k_red-1][2],red[k_red][1]) )

                if ( green[k_green][0]==x ):
                    green_part = green[k_green][1]
                    k_green    = k_green+1
                else:
                    green_part = linear_interpolation( x, (green[k_green-1][0],green[k_green][0]), (green[k_green-1][2],green[k_green][1]) )
                    #print green_part

                if ( blue[k_blue][0]==x ):
                    blue_part = blue[k_blue][1]
                    k_blue    = k_blue+1
                else:
                    blue_part = linear_interpolation( x, (blue[k_blue-1][0],blue[k_blue][0]), (blue[k_blue-1][2],blue[k_blue][1]) )

                #print k
                color_changes.append( "rgb(%.3fcm)=(%.3f,%.3f,%.3f)" % ( x, red_part,green_part,blue_part ) )

                if x>=1.0: break

            return "{mymap}{" + "; ".join( color_changes ) + "}"
    else :
        print "Don't know how to handle color map. Using 'blackwhite'."
        return 'blackwhite'
# =============================================================================
def linear_interpolation( x, X, Y ):
    xRet = ( Y[1]*(x-X[0]) + Y[0]*(X[1]-x) ) / ( X[1]-X[0] )
    return xRet
# =============================================================================
def draw_line2d( file_handle, obj ):
    addplot_options = []

    marker = obj.get_marker()
    if marker!="None":
        print "%TODO Implement markers."

    # get line color
    color = obj.get_color()
    xcolor = mpl_color2xcolor(color)
    if xcolor:
        addplot_options.append( xcolor )

    linestyle = mpl_linestyle2pgfp_linestyle( obj.get_linestyle() )
    if linestyle:
        addplot_options.append( linestyle )

    # process options
    file_handle.write( "\\addplot " )
    if addplot_options:
        options = ", ".join( addplot_options )
        file_handle.write( "[" + options + "]\n" )

    # print the hard numerical data
    xdata, ydata = obj.get_data()
    file_handle.write(  "coordinates {\n" )
    for k in range(len(xdata)):
        file_handle.write( "(%.15g,%.15g) " % (xdata[k], ydata[k]) )
    file_handle.write( "\n};\n" )

    return
# =============================================================================
# Translates a line style of matplotlib to the corresponding style
# in Pgfplots.
def mpl_linestyle2pgfp_linestyle( ls ):
    if (ls=='-' or ls=='None'):
        return None
    elif ls==':':
        return 'dotted'
    else:
        print '%Unknown line style \"' + ls + '\".'
        return None
# =============================================================================
def draw_image( file_handle, obj ):
    from os import path

    global img_number
    try:
       img_number = img_number+1
    except NameError, e:
       img_number = 0

    filename = path.join( output_dir,
                          "img" + repr(img_number) + ".png" )

    # store the image as in a file
    img_array = obj.get_array()
    dims = img_array.shape
    if len(dims)==2: # the values are given as one real number: look at cmap
        clims = obj.get_clim()
        imsave( fname = filename,
                arr   = img_array ,
                cmap  = obj.get_cmap(),
                vmin  = clims[0],
                vmax  = clims[1] )
    elif len(dims)==3: # RGB information
        if dims[2]==4: # RGB+alpha information at each point
            import Image
            # convert to PIL image (after upside-down flip)
            im = Image.fromarray( flipud(img_array), 'RGBA' )
            im.save( filename )

    # write the corresponding information to the TikZ file
    extent = obj.get_extent()
    rel_filepath = path.join( rel_data_path,  path.basename(filename) )
    file_handle.write( ("\\addplot graphics [xmin=%.15g, xmax=%.15g, ymin=%.15g, ymax=%.15g] {" + rel_filepath + "};\n")
                        % extent )

    handle_children( file_handle, obj )

    return
# =============================================================================
def draw_polygon( file_handle, obj ):
    # TODO do nothing for polygons?!
    handle_children( file_handle, obj )
    return
# =============================================================================
# Rather poor way of telling whether an axis has a colorbar associated:
# Check the next axis environment, and see if it is de facto a color bar;
# if yes, return the color bar object.
def find_associated_colorbar( obj ):
      for child in obj.get_children():
              try:
                      cbar = child.colorbar
              except AttributeError:
                      continue
              if not cbar == None: # really necessary?
                      # if fetch was successful, cbar contains
                      # ( reference to colorbar, reference to axis containing colorbar )
                      return cbar[0]
      return None
# =============================================================================
# TODO This function contains crude logic.
def is_colorbar( obj ):
    if isinstance( obj, matplotlib.cm.ScalarMappable ):
        arr = obj.get_array()
        dims = arr.shape
        if len(dims)==1:
            return True # o rly?
        else:
            return False
    else:
        return False
# =============================================================================
def extract_colorbar( obj ):
    colorbars = findobj( obj, is_colorbar )
    if len(colorbars)==0:
        return None
    if not equivalent( colorbars ):
        print "More than one color bar found. Use first one."
    return colorbars[0]
# =============================================================================
# checks if the vectors consists of all the same objects
def equivalent( array ):
    if len(array)==0:
        return False
    else:
        for elem in array:
            if elem!=array[0]:
                return False

    return True
# =============================================================================
def draw_polycollection( file_handle, obj ):
    print "matplotlib2tikz: Don't know how to draw a PolyCollection."
    return
# =============================================================================
# Translates a matplotlib color specification into a proper LaTeX
# xcolor.
def mpl_color2xcolor( color ):
    try: # is the color a gray-scale value?
        alpha = float(color)
        if alpha==0.0:
            return 'black'
        elif alpha==0.5:
            return 'gray'
        elif alpha==0.75:
            return 'lightgray'
        elif alpha==1.0:
            return None
        else:
            return color
    except ValueError:
      pass

    if color=='r':
        return 'red'
    elif color=='g':
        return 'green'
    elif color=='b':
        return 'blue'
    elif color=='k':
        return None
    else:
        print '%Unknown color \"' + color + '\".'
        return None # default
# =============================================================================
def handle_children( file_handle, obj ):
    for child in obj.get_children():
        if ( isinstance( child, matplotlib.axes.Axes ) ):
            draw_axes( file_handle, child )
        elif ( isinstance( child, matplotlib.lines.Line2D ) ):
            draw_line2d( file_handle, child )
        elif ( isinstance( child, matplotlib.image.AxesImage ) ):
            draw_image( file_handle, child )
        elif ( isinstance( child, matplotlib.patches.Polygon ) ):
            draw_polygon( file_handle, child )
        elif ( isinstance( child, matplotlib.collections.PolyCollection ) ):
            draw_polycollection( file_handle, child )
        elif (   isinstance( child, matplotlib.axis.XAxis )
              or isinstance( child, matplotlib.axis.YAxis )
              or isinstance( child, matplotlib.spines.Spine )
              or isinstance( child, matplotlib.patches.Rectangle )
              or isinstance( child, matplotlib.text.Text )
             ):
            pass
        else:
            print "matplotlib2tikz: Don't know how to handle object ", type(child), "."
    return
# =============================================================================
def print_pgfplot_libs_message():
    # remove duplicates
    clean_pgfplots_libs = list(set(pgfplots_libs) )
    libs = ",".join( clean_pgfplots_libs )

    print "========================================================="
    print "Please add the following line to your LaTeX preamble:\n"
    print "\usepackage{pgfplots}"
    if len(clean_pgfplots_libs)!=0:
        print "\usepgfplotslibrary{" + libs + "}"
    print "========================================================="

    return
# =============================================================================
<|endoftext|>


<commit_before>"""
    setup cmssw:
        Sets up the CMS software in the workspace. Workspace must exist.

    Usage:
        setup cmssw version=<cmssw_version> [alias=<>] [--init-git]\
                    [parent_folder=<>]

    Parameters:
        version: The version of CMSSW to set up (e.g. CMSSW_8_0_0)
                 Default: What is defined in setup.json

        scram_arch: The architecture for the CMSSW version
                    (e.g. slc6_amd64_gcc530).
                    Default: What is defined in setup.json

        alias:    Instead of the CMSSW version (CMSSW_X_Y_Z) the project folder
                  will be set to the value of 'alias'

        parent_folder: Parent folder of the CMSSW project.
                       Default: $HEP_PROJECT_ROOT/workspace

        init-git: Prepares the CMSSW area for merging of git branches.
                  Default:true
"""
import logging
import os
from .. import Command as C
from hepshell.interpreter import time_function

LOG = logging.getLogger(__name__)

from l1t_cli.setup import CMSSW_VERSION, WORKSPACE, SCRAM_ARCH

class Command(C):
    DEFAULTS = {
        'version': CMSSW_VERSION,
        'init-git': True,
        'parent_folder': WORKSPACE,
        'alias': '',
        'scram_arch': SCRAM_ARCH,
    }

    def __init__(self, path=__file__, doc=__doc__):
        super(Command, self).__init__(path, doc)

    @time_function('setup cmssw', LOG)
    def run(self, args, variables):
        self.__prepare(args, variables)

        self.__version = self.__variables['version']
        self.__scram_arch = self.__variables['scram_arch']
        self.__alias = self.__variables['alias']
        if self.__alias == '':
            self.__alias = self.__version
        self.__parent_folder = self.__variables['parent_folder']
        self.__init_git = self.__variables['init-git']

        if not self.__can_run():
            return False

        commands = [
            'cd {parent_folder}',
            'export SCRAM_ARCH={scram_arch}',
            'source /cvmfs/cms.cern.ch/cmsset_default.sh',
            '/cvmfs/cms.cern.ch/common/scram project CMSSW {version} -n {alias}',
        ]
        if self.__variables['init-git']:
            commands.extend([
                'cd {alias}/src',
                'eval `/cvmfs/cms.cern.ch/common/scram runtime -sh`',
                'git cms-init',
            ])

        all_in_one = ' && '.join(commands)
        all_in_one = all_in_one.format(
            parent_folder=self.__parent_folder,
            scram_arch=self.__scram_arch,
            version=self.__version,
            alias=self.__alias,
        )

        from hepshell.interpreter import call
        call(all_in_one, logger=LOG, shell=True)

        return True

    def __can_run(self):
        if not os.path.exists(self.__parent_folder):
            LOG.error(
                'Parent folder does not exist: {0}'.format(self.__parent_folder))
            return False
        project_path = os.path.join(self.__parent_folder, self.__alias)
        if os.path.exists(project_path):
            LOG.error('CMSSW is already set up: {0}'.format(project_path))
            return False
        return True
<commit_msg>fix for setup cmssw<commit_after>"""
    setup cmssw:
        Sets up the CMS software in the workspace. Workspace must exist.

    Usage:
        setup cmssw version=<cmssw_version> [alias=<>] [--init-git]\
                    [parent_folder=<>]

    Parameters:
        version: The version of CMSSW to set up (e.g. CMSSW_8_0_0)
                 Default: What is defined in setup.json

        scram_arch: The architecture for the CMSSW version
                    (e.g. slc6_amd64_gcc530).
                    Default: What is defined in setup.json

        alias:    Instead of the CMSSW version (CMSSW_X_Y_Z) the project folder
                  will be set to the value of 'alias'

        parent_folder: Parent folder of the CMSSW project.
                       Default: $HEP_PROJECT_ROOT/workspace

        init-git: Prepares the CMSSW area for merging of git branches.
                  Default:true
"""
import logging
import os
from .. import Command as C
from hepshell.interpreter import time_function

LOG = logging.getLogger(__name__)

from l1t_cli.setup import CMSSW_VERSION, WORKSPACE, SCRAM_ARCH

class Command(C):
    DEFAULTS = {
        'version': CMSSW_VERSION,
        'init-git': True,
        'parent_folder': WORKSPACE,
        'alias': CMSSW_VERSION,
        'scram_arch': SCRAM_ARCH,
    }

    def __init__(self, path=__file__, doc=__doc__):
        super(Command, self).__init__(path, doc)

    @time_function('setup cmssw', LOG)
    def run(self, args, variables):
        self.__prepare(args, variables)

        self.__version = self.__variables['version']
        self.__scram_arch = self.__variables['scram_arch']
        self.__alias = self.__variables['alias']
        if self.__alias == '':
            self.__alias = self.__version
        self.__parent_folder = self.__variables['parent_folder']
        self.__init_git = self.__variables['init-git']

        if not self.__can_run():
            return False

        commands = [
            'cd {parent_folder}',
            'export SCRAM_ARCH={scram_arch}',
            'source /cvmfs/cms.cern.ch/cmsset_default.sh',
            '/cvmfs/cms.cern.ch/common/scram project -n {alias} CMSSW {version}',
        ]
        if self.__variables['init-git']:
            commands.extend([
                'cd {alias}/src',
                'eval `/cvmfs/cms.cern.ch/common/scram runtime -sh`',
                'git cms-init',
            ])

        all_in_one = ' && '.join(commands)
        all_in_one = all_in_one.format(
            parent_folder=self.__parent_folder,
            scram_arch=self.__scram_arch,
            version=self.__version,
            alias=self.__alias,
        )

        from hepshell.interpreter import call
        call(all_in_one, logger=LOG, shell=True)

        return True

    def __can_run(self):
        if not os.path.exists(self.__parent_folder):
            LOG.error(
                'Parent folder does not exist: {0}'.format(self.__parent_folder))
            return False
        project_path = os.path.join(self.__parent_folder, self.__alias)
        if os.path.exists(project_path):
            LOG.error('CMSSW is already set up: {0}'.format(project_path))
            return False
        return True
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Setup script for MSG Data Processing and Operations.

The distribution archive is created using

    python setup.py sdist

Installation is performed using

    python setup.py install [--prefix=${LIBRARY_PATH} --exec-prefix=${BIN_PATH]

where the path arguments within the square brackets are optional.
"""

__author__ = 'Daniel Zhang (張道博)'
__copyright__ = 'Copyright (c) 2013, University of Hawaii Smart Energy Project'
__license__ = 'https://raw.github' \
              '.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid/master/BSD' \
              '-LICENSE.txt'

from distutils.core import setup

setup(name = 'MauiSmartGrid',
      version = '1.0.0',
      description = 'Data Processing and Data Operations for the Maui Smart '
                    'Grid Project.',
      long_description = 'The University of Hawaii at Manoa was tasked with '
                         'maintaining a data repository for use by analysts '
                         'for the Maui Smart Grid (http://www.mauismartgrid'
                         '.com) energy sustainability project through the '
                         'Hawaii Natural Energy Institute (http://www.hnei'
                         '.hawaii.edu). This software provides the data '
                         'processing and operational resources necessary to '
                         'accomplish this task. Source data arrives in '
                         'multiple formats including XML, tab-separated '
                         'values, and comma-separated values. Issues for this'
                         ' project are tracked at the Hawaii Smart Energy '
                         'Project YouTRACK instance ('
                         'http://smart-energy-project.myjetbrains'
                         '.com/youtrack/rest/agile).',
      author = 'Daniel Zhang (張道博)',
      author_email = 'See https://github.com/dz1111',
      url = 'https://github.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid',
      license = 'https://raw.github'
                '.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid/master/BSD'
                '-LICENSE.txt',
      platforms = 'OS X, Linux',

      package_dir = {'': 'src'},

      py_modules = [
          'filelock',
          'meco_data_autoloader',
          'meco_db_delete',
          'meco_db_insert',
          'meco_db_read',
          'meco_dupe_check',
          'meco_fk',
          'meco_mapper',
          'meco_plotting',
          'meco_pv_readings_in_nonpv_mlh_notifier',
          'meco_xml_parser',
          'msg_configer',
          'msg_db_connector',
          'msg_db_exporter',
          'msg_db_util',
          'msg_file_util',
          'msg_logger',
          'msg_noaa_weather_data_dupe_checker',
          'msg_noaa_weather_data_inserter',
          'msg_noaa_weather_data_parser',
          'msg_noaa_weather_data_util',
          'msg_notifier',
          'msg_time_util'
      ],

      scripts = [
          'src/autoloadNewMECOData.py',
          'src/insertCompressedNOAAWeatherData.py',
          'src/insertLocationRecords.py',
          'src/insertMECOEnergyData.py',
          'src/insertMECOMeterLocationHistoryData.py',
          'src/insertMeterRecords.py',
          'src/insertSingleMECOEnergyDataFile.py',
          'src/retrieveNOAAWeatherData.py'
      ]
)
<commit_msg>adding export script<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Setup script for MSG Data Processing and Operations.

The distribution archive is created using

    python setup.py sdist

Installation is performed using

    python setup.py install [--prefix=${LIBRARY_PATH} --exec-prefix=${BIN_PATH]

where the path arguments within the square brackets are optional.
"""

__author__ = 'Daniel Zhang (張道博)'
__copyright__ = 'Copyright (c) 2013, University of Hawaii Smart Energy Project'
__license__ = 'https://raw.github' \
              '.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid/master/BSD' \
              '-LICENSE.txt'

from distutils.core import setup

setup(name = 'MauiSmartGrid',
      version = '1.0.0',
      description = 'Data Processing and Data Operations for the Maui Smart '
                    'Grid Project.',
      long_description = 'The University of Hawaii at Manoa was tasked with '
                         'maintaining a data repository for use by analysts '
                         'for the Maui Smart Grid (http://www.mauismartgrid'
                         '.com) energy sustainability project through the '
                         'Hawaii Natural Energy Institute (http://www.hnei'
                         '.hawaii.edu). This software provides the data '
                         'processing and operational resources necessary to '
                         'accomplish this task. Source data arrives in '
                         'multiple formats including XML, tab-separated '
                         'values, and comma-separated values. Issues for this'
                         ' project are tracked at the Hawaii Smart Energy '
                         'Project YouTRACK instance ('
                         'http://smart-energy-project.myjetbrains'
                         '.com/youtrack/rest/agile).',
      author = 'Daniel Zhang (張道博)',
      author_email = 'See https://github.com/dz1111',
      url = 'https://github.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid',
      license = 'https://raw.github'
                '.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid/master/BSD'
                '-LICENSE.txt',
      platforms = 'OS X, Linux',

      package_dir = {'': 'src'},

      py_modules = [
          'filelock',
          'meco_data_autoloader',
          'meco_db_delete',
          'meco_db_insert',
          'meco_db_read',
          'meco_dupe_check',
          'meco_fk',
          'meco_mapper',
          'meco_plotting',
          'meco_pv_readings_in_nonpv_mlh_notifier',
          'meco_xml_parser',
          'msg_configer',
          'msg_db_connector',
          'msg_db_exporter',
          'msg_db_util',
          'msg_file_util',
          'msg_logger',
          'msg_noaa_weather_data_dupe_checker',
          'msg_noaa_weather_data_inserter',
          'msg_noaa_weather_data_parser',
          'msg_noaa_weather_data_util',
          'msg_notifier',
          'msg_time_util'
      ],

      scripts = [
          'src/autoloadNewMECOData.py',
          'src/exportDBsToCloud.py',
          'src/insertCompressedNOAAWeatherData.py',
          'src/insertLocationRecords.py',
          'src/insertMECOEnergyData.py',
          'src/insertMECOMeterLocationHistoryData.py',
          'src/insertMeterRecords.py',
          'src/insertSingleMECOEnergyDataFile.py',
          'src/retrieveNOAAWeatherData.py'
      ]
)
<|endoftext|>


<commit_before>import functools
import unittest

import numpy

import cupy
from cupy.random import _generator
from cupy.random import _bit_generator
from cupy import testing
from cupy.testing import attr
from cupy.testing import condition

def numpy_cupy_equal_continuous_distribution(significance_level, name='xp'):
    """Decorator that tests the distributions of NumPy samples and CuPy ones.

    Args:
        significance_level (float): The test fails if p-value is lower than
            this argument.
        name(str): Argument name whose value is either
            ``numpy`` or ``cupy`` module.

    Decorated test fixture is required to return samples from the same
    distribution even if ``xp`` is ``numpy`` or ``cupy``.

    """
    def decorator(impl):
        @functools.wraps(impl)
        def test_func(self, *args, **kw):
            kw[name] = cupy
            cupy_result = impl(self, *args, **kw)

            kw[name] = numpy
            numpy_result = impl(self, *args, **kw)

            assert cupy_result is not None
            assert numpy_result is not None
            d_plus, d_minus, p_value = \
                two_sample_Kolmogorov_Smirnov_test(
                    cupy.asnumpy(cupy_result), numpy_result)
            if p_value < significance_level:
                message = '''Rejected null hypothesis:
p: %f
D+ (cupy < numpy): %f
D- (cupy > numpy): %f''' % (p_value, d_plus, d_minus)
                raise AssertionError(message)
        return test_func
    return decorator

def two_sample_Kolmogorov_Smirnov_test(observed1, observed2):
    """Computes the Kolmogorov-Smirnov statistic on 2 samples

    Unlike `scipy.stats.ks_2samp`, the returned p-value is not accurate
    for large p.
    """
    assert observed1.dtype == observed2.dtype
    n1, = observed1.shape
    n2, = observed2.shape
    assert n1 >= 100 and n2 >= 100
    observed = numpy.concatenate([observed1, observed2])
    indices = numpy.argsort(observed)
    observed = observed[indices]  # sort
    ds = numpy.cumsum(numpy.where(indices < n1, -n2, n1).astype(numpy.int64))
    assert ds[-1] == 0
    check = numpy.concatenate([observed[:-1] < observed[1:], [True]])
    ds = ds[check]
    d_plus = float(ds.max()) / (n1 * n2)
    d_minus = -float(ds.min()) / (n1 * n2)
    d = max(d_plus, d_minus)
    # Approximate p = special.kolmogorov(d * numpy.sqrt(n1 * n2 / (n1 + n2)))
    p = min(1.0, 2.0 * numpy.exp(-2.0 * d**2 * n1 * n2 / (n1 + n2)))
    return d_plus, d_minus, p

class GeneratorTestCase(unittest.TestCase):

    target_method = None

    def setUp(self):
        self.__seed = testing.generate_seed()
        # TODO(ecastill) test it with other generators?
        self.rs = _generator.Generator(
            _bit_generator.Philox4x3210(seed=self.__seed))

    def _get_generator_func(self, *args, **kwargs):
        assert isinstance(self.target_method, str), (
            'generate_method must be overridden')
        f = getattr(self.rs, self.target_method)
        return lambda: f(*args, **kwargs)

    def _generate_check_repro(self, func, seed):
        # Sample a random array while checking reproducibility
        self.rs.bit_generator = _bit_generator.Philox4x3210(seed=seed)
        x = func()
        self.rs.bit_generator = _bit_generator.Philox4x3210(seed=seed)
        y = func()
        testing.assert_array_equal(
            x, y,
            'Randomly generated arrays with the same seed did not match')
        return x

    def generate(self, *args, **kwargs):
        # Pick one sample from generator.
        # Reproducibility is checked by repeating seed-and-sample cycle twice.
        func = self._get_generator_func(*args, **kwargs)
        return self._generate_check_repro(func, self.__seed)

    def generate_many(self, *args, **kwargs):
        # Pick many samples from generator.
        # Reproducibility is checked only for the first sample,
        # because it's very slow to set seed every time.
        _count = kwargs.pop('_count', None)
        assert _count is not None, '_count is required'
        func = self._get_generator_func(*args, **kwargs)

        if _count == 0:
            return []

        vals = [self._generate_check_repro(func, self.__seed)]
        for _ in range(1, _count):
            vals.append(func())
        return vals

    def check_ks(self, significance_level, cupy_len=100, numpy_len=1000):
        return functools.partial(
            self._check_ks, significance_level, cupy_len, numpy_len)

    def _check_ks(
            self, significance_level, cupy_len, numpy_len,
            *args, **kwargs):
        assert 'size' in kwargs

        # cupy
        func = self._get_generator_func(*args, **kwargs)
        vals_cupy = func()
        assert vals_cupy.size > 0
        count = 1 + (cupy_len - 1) // vals_cupy.size
        vals_cupy = [vals_cupy]
        for _ in range(1, count):
            vals_cupy.append(func())
        vals_cupy = cupy.stack(vals_cupy).ravel()

        # numpy
        kwargs['size'] = numpy_len
        dtype = kwargs.pop('dtype', None)
        numpy_rs = numpy.random.Generator(numpy.random.MT19937(self.__seed))
        vals_numpy = getattr(numpy_rs, self.target_method)(*args, **kwargs)
        if dtype is not None:
            vals_numpy = vals_numpy.astype(dtype, copy=False)

        # test
        d_plus, d_minus, p_value = \
            two_sample_Kolmogorov_Smirnov_test(
                cupy.asnumpy(vals_cupy), vals_numpy)
        if p_value < significance_level:
            message = '''Rejected null hypothesis:
p: %f
D+ (cupy < numpy): %f
D- (cupy > numpy): %f''' % (p_value, d_plus, d_minus)
            raise AssertionError(message)

def _xp_random(xp, method_name):
    # Unfortunately cupy and numpy does not share the random generators
    if xp == cupy:
        bit_gen = cupy.random.Philox4x3210()
    else:
        bit_gen = numpy.random.MT19937()
    method = getattr(xp.random.Generator(bit_gen), method_name)
    if xp == cupy:
        return method

    def f(*args, **kwargs):
        dtype = kwargs.pop('dtype', None)
        ret = method(*args, **kwargs)
        if dtype is not None:
            ret = ret.astype(dtype, copy=False)
        return ret

    return f

@testing.parameterize(
    {'a': 1.0, 'b': 3.0},
    {'a': 3.0, 'b': 3.0},
    {'a': 3.0, 'b': 1.0},
)
@testing.gpu
@testing.fix_random()
class TestBeta(GeneratorTestCase):

    target_method = 'beta'

    def test_beta(self):
        self.generate(a=self.a, b=self.b, size=(3, 2))

    @testing.for_dtypes('fd')
    @condition.repeat_with_success_at_least(10, 3)
    def test_beta_ks(self, dtype):
        self.check_ks(0.05)(
            a=self.a, b=self.b, size=2000, dtype=dtype)

@testing.gpu
@testing.fix_random()
class TestStandardExponential(GeneratorTestCase):

    target_method = 'standard_exponential'

    def test_standard_exponential(self):
        self.generate(size=(3, 2))

    @attr.slow
    @condition.repeat(10)
    def test_standard_exponential_isfinite(self):
        x = self.generate(size=10**7)
        assert cupy.isfinite(x).all()

    @testing.for_dtypes('fd')
    @condition.repeat_with_success_at_least(10, 3)
        self.check_ks(0.05)(
            size=2000, dtype=dtype)

@testing.gpu
@testing.fix_random()
class TestIntegers(GeneratorTestCase):
    # TODO(niboshi):
    #   Test soundness of distribution.
    #   Currently only reprocibility is checked.

    target_method = 'integers'

    def test_randint_1(self):
        self.generate(3)

    def test_randint_2(self):
        self.generate(3, 4, size=(3, 2))

    def test_randint_empty1(self):
        self.generate(3, 10, size=0)

    def test_randint_empty2(self):
        self.generate(3, size=(4, 0, 5))

    def test_randint_overflow(self):
        self.generate(numpy.int8(-100), numpy.int8(100))

    def test_randint_float1(self):
        self.generate(-1.2, 3.4, 5)

    def test_randint_float2(self):
        self.generate(6.7, size=(2, 3))

    def test_randint_int64_1(self):
        self.generate(2**34, 2**40, 3)

"""
@testing.gpu
class TestRandomStateThreadSafe(unittest.TestCase):

    def setUp(self):
        cupy.random.reset_states()

    def test_get_random_state_thread_safe(self):
        seed = 10
        threads = [
            threading.Thread(target=lambda: cupy.random.seed(seed)),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
        ]

        for t in threads:
            t.start()
        for t in threads:
            t.join()

        actual = cupy.random.uniform()
        cupy.random.seed(seed)
        expected = cupy.random.uniform()
        assert actual == expected

    def test_set_random_state_thread_safe(self):
        rs = cupy.random.RandomState()
        threads = [
            threading.Thread(target=lambda: cupy.random.set_random_state(rs)),
            threading.Thread(target=lambda: cupy.random.set_random_state(rs)),
        ]

        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert cupy.random.get_random_state() is rs
"""
<commit_msg>fix test<commit_after>import functools
import unittest

import numpy

import cupy
from cupy.random import _generator
from cupy.random import _bit_generator
from cupy import testing
from cupy.testing import attr
from cupy.testing import condition

def numpy_cupy_equal_continuous_distribution(significance_level, name='xp'):
    """Decorator that tests the distributions of NumPy samples and CuPy ones.

    Args:
        significance_level (float): The test fails if p-value is lower than
            this argument.
        name(str): Argument name whose value is either
            ``numpy`` or ``cupy`` module.

    Decorated test fixture is required to return samples from the same
    distribution even if ``xp`` is ``numpy`` or ``cupy``.

    """
    def decorator(impl):
        @functools.wraps(impl)
        def test_func(self, *args, **kw):
            kw[name] = cupy
            cupy_result = impl(self, *args, **kw)

            kw[name] = numpy
            numpy_result = impl(self, *args, **kw)

            assert cupy_result is not None
            assert numpy_result is not None
            d_plus, d_minus, p_value = \
                two_sample_Kolmogorov_Smirnov_test(
                    cupy.asnumpy(cupy_result), numpy_result)
            if p_value < significance_level:
                message = '''Rejected null hypothesis:
p: %f
D+ (cupy < numpy): %f
D- (cupy > numpy): %f''' % (p_value, d_plus, d_minus)
                raise AssertionError(message)
        return test_func
    return decorator

def two_sample_Kolmogorov_Smirnov_test(observed1, observed2):
    """Computes the Kolmogorov-Smirnov statistic on 2 samples

    Unlike `scipy.stats.ks_2samp`, the returned p-value is not accurate
    for large p.
    """
    assert observed1.dtype == observed2.dtype
    n1, = observed1.shape
    n2, = observed2.shape
    assert n1 >= 100 and n2 >= 100
    observed = numpy.concatenate([observed1, observed2])
    indices = numpy.argsort(observed)
    observed = observed[indices]  # sort
    ds = numpy.cumsum(numpy.where(indices < n1, -n2, n1).astype(numpy.int64))
    assert ds[-1] == 0
    check = numpy.concatenate([observed[:-1] < observed[1:], [True]])
    ds = ds[check]
    d_plus = float(ds.max()) / (n1 * n2)
    d_minus = -float(ds.min()) / (n1 * n2)
    d = max(d_plus, d_minus)
    # Approximate p = special.kolmogorov(d * numpy.sqrt(n1 * n2 / (n1 + n2)))
    p = min(1.0, 2.0 * numpy.exp(-2.0 * d**2 * n1 * n2 / (n1 + n2)))
    return d_plus, d_minus, p

class GeneratorTestCase(unittest.TestCase):

    target_method = None

    def setUp(self):
        self.__seed = testing.generate_seed()
        # TODO(ecastill) test it with other generators?
        self.rs = _generator.Generator(
            _bit_generator.Philox4x3210(seed=self.__seed))

    def _get_generator_func(self, *args, **kwargs):
        assert isinstance(self.target_method, str), (
            'generate_method must be overridden')
        f = getattr(self.rs, self.target_method)
        return lambda: f(*args, **kwargs)

    def _generate_check_repro(self, func, seed):
        # Sample a random array while checking reproducibility
        self.rs.bit_generator = _bit_generator.Philox4x3210(seed=seed)
        x = func()
        self.rs.bit_generator = _bit_generator.Philox4x3210(seed=seed)
        y = func()
        testing.assert_array_equal(
            x, y,
            'Randomly generated arrays with the same seed did not match')
        return x

    def generate(self, *args, **kwargs):
        # Pick one sample from generator.
        # Reproducibility is checked by repeating seed-and-sample cycle twice.
        func = self._get_generator_func(*args, **kwargs)
        return self._generate_check_repro(func, self.__seed)

    def generate_many(self, *args, **kwargs):
        # Pick many samples from generator.
        # Reproducibility is checked only for the first sample,
        # because it's very slow to set seed every time.
        _count = kwargs.pop('_count', None)
        assert _count is not None, '_count is required'
        func = self._get_generator_func(*args, **kwargs)

        if _count == 0:
            return []

        vals = [self._generate_check_repro(func, self.__seed)]
        for _ in range(1, _count):
            vals.append(func())
        return vals

    def check_ks(self, significance_level, cupy_len=100, numpy_len=1000):
        return functools.partial(
            self._check_ks, significance_level, cupy_len, numpy_len)

    def _check_ks(
            self, significance_level, cupy_len, numpy_len,
            *args, **kwargs):
        assert 'size' in kwargs

        # cupy
        func = self._get_generator_func(*args, **kwargs)
        vals_cupy = func()
        assert vals_cupy.size > 0
        count = 1 + (cupy_len - 1) // vals_cupy.size
        vals_cupy = [vals_cupy]
        for _ in range(1, count):
            vals_cupy.append(func())
        vals_cupy = cupy.stack(vals_cupy).ravel()

        # numpy
        kwargs['size'] = numpy_len
        dtype = kwargs.pop('dtype', None)
        numpy_rs = numpy.random.Generator(numpy.random.MT19937(self.__seed))
        vals_numpy = getattr(numpy_rs, self.target_method)(*args, **kwargs)
        if dtype is not None:
            vals_numpy = vals_numpy.astype(dtype, copy=False)

        # test
        d_plus, d_minus, p_value = \
            two_sample_Kolmogorov_Smirnov_test(
                cupy.asnumpy(vals_cupy), vals_numpy)
        if p_value < significance_level:
            message = '''Rejected null hypothesis:
p: %f
D+ (cupy < numpy): %f
D- (cupy > numpy): %f''' % (p_value, d_plus, d_minus)
            raise AssertionError(message)

def _xp_random(xp, method_name):
    # Unfortunately cupy and numpy does not share the random generators
    if xp == cupy:
        bit_gen = cupy.random.Philox4x3210()
    else:
        bit_gen = numpy.random.MT19937()
    method = getattr(xp.random.Generator(bit_gen), method_name)
    if xp == cupy:
        return method

    def f(*args, **kwargs):
        dtype = kwargs.pop('dtype', None)
        ret = method(*args, **kwargs)
        if dtype is not None:
            ret = ret.astype(dtype, copy=False)
        return ret

    return f

@testing.parameterize(
    {'a': 1.0, 'b': 3.0},
    {'a': 3.0, 'b': 3.0},
    {'a': 3.0, 'b': 1.0},
)
@testing.gpu
@testing.fix_random()
class TestBeta(GeneratorTestCase):

    target_method = 'beta'

    def test_beta(self):
        self.generate(a=self.a, b=self.b, size=(3, 2))

    @testing.for_dtypes('fd')
    @condition.repeat_with_success_at_least(10, 3)
    def test_beta_ks(self, dtype):
        self.check_ks(0.05)(
            a=self.a, b=self.b, size=2000, dtype=dtype)

@testing.gpu
@testing.fix_random()
class TestStandardExponential(GeneratorTestCase):

    target_method = 'standard_exponential'

    def test_standard_exponential(self):
        self.generate(size=(3, 2))

    @attr.slow
    @condition.repeat(10)
    def test_standard_exponential_isfinite(self):
        x = self.generate(size=10**7)
        assert cupy.isfinite(x).all()

    @testing.for_dtypes('fd')
    @condition.repeat_with_success_at_least(10, 3)
    def test_standard_exponential_ks(self, dtype):
        self.check_ks(0.05)(size=2000, dtype=dtype)

@testing.gpu
@testing.fix_random()
class TestIntegers(GeneratorTestCase):
    # TODO(niboshi):
    #   Test soundness of distribution.
    #   Currently only reprocibility is checked.

    target_method = 'integers'

    def test_randint_1(self):
        self.generate(3)

    def test_randint_2(self):
        self.generate(3, 4, size=(3, 2))

    def test_randint_empty1(self):
        self.generate(3, 10, size=0)

    def test_randint_empty2(self):
        self.generate(3, size=(4, 0, 5))

    def test_randint_overflow(self):
        self.generate(numpy.int8(-100), numpy.int8(100))

    def test_randint_float1(self):
        self.generate(-1.2, 3.4, 5)

    def test_randint_float2(self):
        self.generate(6.7, size=(2, 3))

    def test_randint_int64_1(self):
        self.generate(2**34, 2**40, 3)

"""
@testing.gpu
class TestRandomStateThreadSafe(unittest.TestCase):

    def setUp(self):
        cupy.random.reset_states()

    def test_get_random_state_thread_safe(self):
        seed = 10
        threads = [
            threading.Thread(target=lambda: cupy.random.seed(seed)),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
            threading.Thread(target=lambda: cupy.random.get_random_state()),
        ]

        for t in threads:
            t.start()
        for t in threads:
            t.join()

        actual = cupy.random.uniform()
        cupy.random.seed(seed)
        expected = cupy.random.uniform()
        assert actual == expected

    def test_set_random_state_thread_safe(self):
        rs = cupy.random.RandomState()
        threads = [
            threading.Thread(target=lambda: cupy.random.set_random_state(rs)),
            threading.Thread(target=lambda: cupy.random.set_random_state(rs)),
        ]

        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert cupy.random.get_random_state() is rs
"""
<|endoftext|>


<commit_before>import argparse
import contextlib
import cStringIO
import itertools
import json
import os
import re
import sys
import zipfile

import bs4
import requests

BASE_URL = 'http://www.mathworks.com/matlabcentral/fileexchange'
TAG_REGEX = re.compile(r'([\w\s]+)(?:\(\d+\))?')

def get_soup(*args, **kwargs):
    return bs4.BeautifulSoup(requests.get(*args, **kwargs).text)

class Project(object):
    def __init__(self, url_or_name):
        if url_or_name.startswith(BASE_URL):
            self.url = url_or_name
        else:
            self.url = '%s/%s' % (BASE_URL, url_or_name)

    def __repr__(self):
        return '<Project: %s>' % self.name

    @property
    def name(self):
        return self.url.split('/')[-1]

    # The way downloads work on MatlabCentral is that the download link
    # redirects you to the url of the actual file you're downloading.
    # Instead of following the redirect automatically, we manually get
    # the value of the Location: header; this lets us get at the name of
    # the downloaded file (and so we can tell whether it's a zip, etc.)
    def download(self, path, extract_archives=True):
        download_url = '%s?download=true' % self.url
        response = requests.head(download_url, allow_redirects=False)
        real_download_url = response.headers['Location']
        download_filename = real_download_url.split('/')[-1]
        response = requests.get(real_download_url)
        if extract_archives and download_filename.endswith('.zip'):
            archive = zipfile.ZipFile(cStringIO.StringIO(response.content))
            archive.extractall(path)
        else:
            with open(os.path.join(path, download_filename), 'w') as f:
                f.write(response.text.encode('utf-8'))

    def get_metadata(self):
        metadata = {}
        soup = get_soup(self.url)

        tags_div = soup.find('div', id='all_tags')
        raw_tags = tags_div.get_text().strip().split('\n\n')[0].split(', ')
        tags = []
        for tag in raw_tags:
            match = TAG_REGEX.match(tag)
            if match:
                tags.append(match.group(1))
        metadata['tags'] = tags

        details = soup.find('div', id='details')

        metadata['title'] = details.h1.text
        metadata['summary'] = details.find('p', id='summary').text

        author_link = details.find('p', id='author').a
        metadata['author'] = author_link.text
        metadata['author_url'] = author_link['href']

        metadata['date_submitted'] = details.find(
            'span', id='submissiondate').text.strip()
        date_span = details.find('span', id='date_updated')
        if date_span is not None:
            metadata['date_updated'] = date_span.text[len('(Updated '):-1]
        return metadata

    def get_json(self):
        data = dict(name=self.name, url=self.url)
        data.update(self.get_metadata())
        return data

def fileindex_projects(num_projects, sort):
    def helper():
        for page in itertools.count(1):
            params = dict(page=page, term='type:Function', sort=sort)
            soup = get_soup(BASE_URL, params=params)
            for title in soup.find_all('p', attrs={'class': 'file_title'}):
                yield Project(title.a['href'])
    return itertools.islice(helper(), num_projects)

def parse_args():
    parser = argparse.ArgumentParser(
        description='download files from MatlabCentral')
    parser.add_argument(
        '--num_projects', default=10, type=int,
        help='Number of projects to fetch.')
    parser.add_argument(
        '--extract_archives', type=bool, default=True,
        help='If true, automatically extract archives.')
    SORT_CHOICES = ('downloads_desc', 'downloads_asc',
                    'date_desc_updated', 'date_asc_updated',
                    'date_desc_submitted', 'date_asc_submitted',
                    'comments_desc', 'comments_asc',
                    'ratings_desc', 'ratings_asc')
    parser.add_argument(
        '--sort', choices=SORT_CHOICES, default='downloads_desc',
        help='Sorting criteria.')
    parser.add_argument(
        '--to', default='',
        help='Directory to download projects to (default current).')
    return parser.parse_args()

@contextlib.contextmanager
def status(message):
    print '%s...' % message,
    sys.stdout.flush()
    yield
    print 'done.'

def main():
    args = parse_args()
    projects = []
    for project in fileindex_projects(args.num_projects, args.sort):
        download_path = os.path.join(args.to, project.name)
        os.makedirs(download_path)
        with status('Downloading %s' % project.name):
            project.download(download_path, args.extract_archives)
        projects.append(project.get_json())
    with open(os.path.join(args.to, 'manifest.json'), 'w') as f:
        json.dump({'projects': projects}, f, indent=2, sort_keys=True)

if __name__ == '__main__':
    main()
<commit_msg>Fix crash when zip file contains non-ascii names.<commit_after>import argparse
import contextlib
import cStringIO
import errno
import itertools
import json
import os
import re
import shutil
import sys
import zipfile

import bs4
import requests

BASE_URL = 'http://www.mathworks.com/matlabcentral/fileexchange'
TAG_REGEX = re.compile(r'([\w\s]+)(?:\(\d+\))?')

def get_soup(*args, **kwargs):
    return bs4.BeautifulSoup(requests.get(*args, **kwargs).text)

# From http://stackoverflow.com/a/600612/281108
def mkdir_p(path):
    try:
        os.makedirs(path)
    except OSError as e:
        if e.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

def extractall(archive, path):
    "A version of ZipFile.extractall that doesn't choke on utf-8 filenames."
    for filename in archive.namelist():
        dest_filename = os.path.join(path, filename.decode('cp437'))
        if dest_filename.endswith('/'):
            mkdir_p(dest_filename)
            continue
        mkdir_p(os.path.dirname(dest_filename))
        with archive.open(filename) as src, open(dest_filename, 'w') as dest:
            shutil.copyfileobj(src, dest)

class Project(object):
    def __init__(self, url_or_name):
        if url_or_name.startswith(BASE_URL):
            self.url = url_or_name
        else:
            self.url = '%s/%s' % (BASE_URL, url_or_name)

    def __repr__(self):
        return '<Project: %s>' % self.name

    @property
    def name(self):
        return self.url.split('/')[-1]

    # The way downloads work on MatlabCentral is that the download link
    # redirects you to the url of the actual file you're downloading.
    # Instead of following the redirect automatically, we manually get
    # the value of the Location: header; this lets us get at the name of
    # the downloaded file (and so we can tell whether it's a zip, etc.)
    def download(self, path='.', extract_archives=True):
        download_url = '%s?download=true' % self.url
        response = requests.head(download_url, allow_redirects=False)
        real_download_url = response.headers['Location']
        download_filename = real_download_url.split('/')[-1]
        response = requests.get(real_download_url)
        if extract_archives and download_filename.endswith('.zip'):
            with zipfile.ZipFile(cStringIO.StringIO(response.content)) as f:
                extractall(f, path)
        else:
            with open(os.path.join(path, download_filename), 'w') as f:
                f.write(response.text.encode('utf-8'))

    def get_metadata(self):
        metadata = {}
        soup = get_soup(self.url)

        tags_div = soup.find('div', id='all_tags')
        raw_tags = tags_div.get_text().strip().split('\n\n')[0].split(', ')
        tags = []
        for tag in raw_tags:
            match = TAG_REGEX.match(tag)
            if match:
                tags.append(match.group(1))
        metadata['tags'] = tags

        details = soup.find('div', id='details')

        metadata['title'] = details.h1.text
        metadata['summary'] = details.find('p', id='summary').text

        author_link = details.find('p', id='author').a
        metadata['author'] = author_link.text
        metadata['author_url'] = author_link['href']

        metadata['date_submitted'] = details.find(
            'span', id='submissiondate').text.strip()
        date_span = details.find('span', id='date_updated')
        if date_span is not None:
            metadata['date_updated'] = date_span.text[len('(Updated '):-1]
        return metadata

    def get_json(self):
        data = dict(name=self.name, url=self.url)
        data.update(self.get_metadata())
        return data

def fileindex_projects(num_projects, sort):
    def helper():
        for page in itertools.count(1):
            params = dict(page=page, term='type:Function', sort=sort)
            soup = get_soup(BASE_URL, params=params)
            for title in soup.find_all('p', attrs={'class': 'file_title'}):
                yield Project(title.a['href'])
    return itertools.islice(helper(), num_projects)

def parse_args():
    parser = argparse.ArgumentParser(
        description='download files from MatlabCentral')
    parser.add_argument(
        '--num_projects', default=10, type=int,
        help='Number of projects to fetch.')
    parser.add_argument(
        '--extract_archives', type=bool, default=True,
        help='If true, automatically extract archives.')
    SORT_CHOICES = ('downloads_desc', 'downloads_asc',
                    'date_desc_updated', 'date_asc_updated',
                    'date_desc_submitted', 'date_asc_submitted',
                    'comments_desc', 'comments_asc',
                    'ratings_desc', 'ratings_asc')
    parser.add_argument(
        '--sort', choices=SORT_CHOICES, default='downloads_desc',
        help='Sorting criteria.')
    parser.add_argument(
        '--to', default='',
        help='Directory to download projects to (default current).')
    return parser.parse_args()

@contextlib.contextmanager
def status(message):
    print '%s...' % message,
    sys.stdout.flush()
    yield
    print 'done.'

def main():
    args = parse_args()
    projects = []
    for project in fileindex_projects(args.num_projects, args.sort):
        download_path = os.path.join(args.to, project.name)
        os.makedirs(download_path)
        with status('Downloading %s' % project.name):
            project.download(download_path, args.extract_archives)
        projects.append(project.get_json())
    with open(os.path.join(args.to, 'manifest.json'), 'w') as f:
        json.dump({'projects': projects}, f, indent=2, sort_keys=True)

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


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
CrossPM (Cross Package Manager) version: {version} The MIT License (MIT)

Usage:
    crosspm download [options]
    crosspm promote [options]       * Temporarily off
    crosspm pack <OUT> <SOURCE> [options]
    crosspm cache [size | age | clear [hard]]
    crosspm -h | --help
    crosspm --version

Options:
    <OUT>                           Output file.
    <SOURCE>                        Source directory path.
    -h, --help                      Show this screen.
    --version                       Show version.
    -L, --list                      Do not load packages and its dependencies. Just show what's found.
    -v LEVEL, --verbose=LEVEL       Set output verbosity: ({verb_level}) [default: ].
    -l LOGFILE, --log=LOGFILE       File name for log output. Log level is '{log_default}' if set when verbose doesn't.
    -c FILE, --config=FILE          Path to configuration file.
    -o OPTIONS, --options OPTIONS   Extra options.
    --depslock-path=FILE            Path to file with locked dependencies [./{deps_lock_default}]
    --out-format=TYPE               Output data format. Available formats:({out_format}) [default: {out_format_default}]
    --output=FILE                   Output file name (required if --out_format is not stdout)
    --out-prefix=PREFIX             Prefix for output variable name [default: ] (no prefix at all)
    --no-fails                      Ignore fails config if possible.

"""

import logging
from docopt import docopt
import os
from crosspm import config
from crosspm.helpers.archive import Archive
from crosspm.helpers.config import (
    CROSSPM_DEPENDENCY_LOCK_FILENAME,
    Config,
)
from crosspm.helpers.downloader import Downloader
# from crosspm.helpers.promoter import Promoter
from crosspm.helpers.output import Output
from crosspm.helpers.exceptions import *

# TODO: Upgrade exceptions handling
class CrossPM(object):
    _config = None
    _args = None
    _output = Output()
    _ready = False
    _throw_exceptions = True

    def __init__(self, args=None, throw_exceptions=None):
        if throw_exceptions is not None:
            self._throw_exceptions = throw_exceptions
        elif args is not None:
            self._throw_exceptions = False
        self._log = logging.getLogger('crosspm')
        self._args = docopt(__doc__.format(version=config.__version__,
                                           verb_level=Config.get_verbosity_level(),
                                           log_default=Config.get_verbosity_level(0, True),
                                           deps_lock_default=CROSSPM_DEPENDENCY_LOCK_FILENAME,
                                           out_format=self._output.get_output_types(),
                                           out_format_default='stdout',
                                           ),
                            argv=args,
                            version=config.__version__)

        if type(self._args) is str:
            if self._throw_exceptions:
                print(self._args)
                exit()
        self._ready = True

    def read_config(self):
        self._config = Config(self._args['--config'], self._args['--options'], self._args['--no-fails'], self._args['--depslock-path'])

    def run(self):
        if self._ready:
            errorcode, msg = self.do_run(self.check_common_args)
            if errorcode == 0:
                errorcode, msg = self.do_run(self.read_config)

                if errorcode == 0:
                    if self._args['download']:
                        errorcode, msg = self.do_run(self.download)
                        # self.download()

                    elif self._args['promote']:
                        errorcode, msg = self.do_run(self.promote)

                    elif self._args['pack']:
                        errorcode, msg = self.do_run(self.pack)

                    elif self._args['cache']:
                        errorcode, msg = self.do_run(self.cache)
        else:
            errorcode, msg = CROSSPM_ERRORCODE_WRONG_ARGS, self._args
        return (errorcode, msg)

    def do_run(self, func, *args, **kwargs):
        try:
            func(*args, **kwargs)
        except CrosspmExceptionWrongArgs as e:
            if self._throw_exceptions:
                print(__doc__)
                self._log.critical(e.msg)
                sys.exit(e.error_code)
            else:
                return (e.error_code, e.msg)

        except CrosspmException as e:
            if self._throw_exceptions:
                print_stdout('')
                self._log.critical(e.msg)
                sys.exit(e.error_code)
            else:
                return (e.error_code, e.msg)

        except Exception as e:
            if self._throw_exceptions:
                print_stdout('')
                self._log.exception(e)
                self._log.critical('Unknown error occurred!')
                sys.exit(CROSSPM_ERRORCODE_UNKNOWN_ERROR)
            else:
                return (CROSSPM_ERRORCODE_UNKNOWN_ERROR, 'Unknown error occurred!')
        return (0, '')

    def check_common_args(self):
        if self._args['--output']:
            output = self._args['--output'].strip().strip("'").strip('"')
            output_abs = os.path.abspath(output)
            if os.path.isdir(output_abs):
                raise CrosspmExceptionWrongArgs(
                    '"%s" is a directory - can\'t write to it'
                )
            self._args['--output'] = output

        self.set_logging_level()

    def set_logging_level(self):
        level_str = self._args['--verbose'].strip().lower()

        log = self._args['--log']
        if log:
            log = log.strip().strip("'").strip('"')
            log_abs = os.path.abspath(log)
            if os.path.isdir(log_abs):
                raise CrosspmExceptionWrongArgs(
                    '"%s" is a directory - can\'t write log to it'
                )
            else:
                log_dir = os.path.dirname(log_abs)
                if not os.path.exists(log_dir):
                    os.makedirs(log_dir)
        else:
            log_abs = None

        level = Config.get_verbosity_level(level_str or 'console')
        if level or log_abs:
            self._log.setLevel(level)
            format_str = '%(asctime)-19s [%(levelname)-9s] %(message)s'
            if level_str == 'debug':
                format_str = '%(asctime)-19s [%(levelname)-9s] %(name)-12s: %(message)s'
            formatter = logging.Formatter(format_str, datefmt="%Y-%m-%d %H:%M:%S")

            if level:
                sh = logging.StreamHandler(stream=sys.stderr)
                sh.setLevel(level)
                # sh.setFormatter(formatter)
                self._log.addHandler(sh)

            if log_abs:
                if not level_str:
                    level = Config.get_verbosity_level(0)
                fh = logging.FileHandler(filename=log_abs)
                fh.setLevel(level)
                fh.setFormatter(formatter)
                self._log.addHandler(fh)

    def download(self):
        if self._args['--out-format'] == 'stdout':
            if self._args['--output']:
                raise CrosspmExceptionWrongArgs(
                    "unwanted argument '--output' while argument '--out-format={}'".format(
                        self._args['--out-format'],
                    ))
        elif not self._args['--output']:
            raise CrosspmExceptionWrongArgs(
                "argument '--output' required when argument '--out-format={}'".format(
                    self._args['--out-format'],
                ))

        params = {
            'out_format': ['--out-format', ''],
            'output': ['--output', ''],
            'out_prefix': ['--out-prefix', ''],
            # 'depslock_path': ['--depslock-path', ''],
        }

        for k, v in params.items():
            params[k] = self._args[v[0]].strip('"').strip("'") if v[0] in self._args else v[1]

        do_load = not self._args['--list']
        if do_load:
            self._config.cache.auto_clear()
        cpm_downloader = Downloader(self._config, do_load)
        # cpm_downloader = Downloader(self._config, params.pop('depslock_path'), do_load)
        packages = cpm_downloader.download_packages()

        _not_found = any(_pkg is None for _pkg in packages.values())
        if _not_found:
            raise CrosspmException(
                CROSSPM_ERRORCODE_PACKAGE_NOT_FOUND,
                'Some package(s) not found.'
            )
        if do_load:
            self._output.write(params, packages)

    def promote(self):
        self._log.warning('This option is temporarily off.')
        # cpm_promoter = Promoter(self._config)
        # cpm_promoter.promote_packages()

    def pack(self):
        Archive.create(self._args['<OUT>'], self._args['<SOURCE>'])

    def cache(self):
        if self._args['clear']:
            self._config.cache.clear(self._args['hard'])
        elif self._args['size']:
            self._config.cache.size()
        elif self._args['age']:
            self._config.cache.age()
        else:
            self._config.cache.info()

if __name__ == '__main__':
    app = CrossPM()
    app.run()
<commit_msg>parameters parsing error fixed<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
CrossPM (Cross Package Manager) version: {version} The MIT License (MIT)

Usage:
    crosspm download [options]
    crosspm promote [options]       * Temporarily off
    crosspm pack <OUT> <SOURCE> [options]
    crosspm cache [size | age | clear [hard]]
    crosspm -h | --help
    crosspm --version

Options:
    <OUT>                           Output file.
    <SOURCE>                        Source directory path.
    -h, --help                      Show this screen.
    --version                       Show version.
    -L, --list                      Do not load packages and its dependencies. Just show what's found.
    -v LEVEL, --verbose=LEVEL       Set output verbosity: ({verb_level}) [default: ].
    -l LOGFILE, --log=LOGFILE       File name for log output. Log level is '{log_default}' if set when verbose doesn't.
    -c FILE, --config=FILE          Path to configuration file.
    -o OPTIONS, --options OPTIONS   Extra options.
    --depslock-path=FILE            Path to file with locked dependencies [./{deps_lock_default}]
    --out-format=TYPE               Output data format. Available formats:({out_format}) [default: {out_format_default}]
    --output=FILE                   Output file name (required if --out_format is not stdout)
    --out-prefix=PREFIX             Prefix for output variable name [default: ] (no prefix at all)
    --no-fails                      Ignore fails config if possible.

"""

import logging
from docopt import docopt
import os
from crosspm import config
from crosspm.helpers.archive import Archive
from crosspm.helpers.config import (
    CROSSPM_DEPENDENCY_LOCK_FILENAME,
    Config,
)
from crosspm.helpers.downloader import Downloader
# from crosspm.helpers.promoter import Promoter
from crosspm.helpers.output import Output
from crosspm.helpers.exceptions import *

# TODO: Upgrade exceptions handling
class CrossPM(object):
    _config = None
    _args = None
    _output = Output()
    _ready = False
    _throw_exceptions = True

    def __init__(self, args=None, throw_exceptions=None):
        if throw_exceptions is not None:
            self._throw_exceptions = throw_exceptions
        elif args is not None:
            self._throw_exceptions = False
        self._log = logging.getLogger('crosspm')
        self._args = docopt(__doc__.format(version=config.__version__,
                                           verb_level=Config.get_verbosity_level(),
                                           log_default=Config.get_verbosity_level(0, True),
                                           deps_lock_default=CROSSPM_DEPENDENCY_LOCK_FILENAME,
                                           out_format=self._output.get_output_types(),
                                           out_format_default='stdout',
                                           ),
                            argv=args,
                            version=config.__version__)

        if type(self._args) is str:
            if self._throw_exceptions:
                print(self._args)
                exit()
        self._ready = True

    def read_config(self):
        self._config = Config(self._args['--config'], self._args['--options'], self._args['--no-fails'], self._args['--depslock-path'])

    def run(self):
        if self._ready:
            errorcode, msg = self.do_run(self.check_common_args)
            if errorcode == 0:
                errorcode, msg = self.do_run(self.read_config)

                if errorcode == 0:
                    if self._args['download']:
                        errorcode, msg = self.do_run(self.download)
                        # self.download()

                    elif self._args['promote']:
                        errorcode, msg = self.do_run(self.promote)

                    elif self._args['pack']:
                        errorcode, msg = self.do_run(self.pack)

                    elif self._args['cache']:
                        errorcode, msg = self.do_run(self.cache)
        else:
            errorcode, msg = CROSSPM_ERRORCODE_WRONG_ARGS, self._args
        return (errorcode, msg)

    def do_run(self, func, *args, **kwargs):
        try:
            func(*args, **kwargs)
        except CrosspmExceptionWrongArgs as e:
            if self._throw_exceptions:
                print(__doc__)
                self._log.critical(e.msg)
                sys.exit(e.error_code)
            else:
                return (e.error_code, e.msg)

        except CrosspmException as e:
            if self._throw_exceptions:
                print_stdout('')
                self._log.critical(e.msg)
                sys.exit(e.error_code)
            else:
                return (e.error_code, e.msg)

        except Exception as e:
            if self._throw_exceptions:
                print_stdout('')
                self._log.exception(e)
                self._log.critical('Unknown error occurred!')
                sys.exit(CROSSPM_ERRORCODE_UNKNOWN_ERROR)
            else:
                return (CROSSPM_ERRORCODE_UNKNOWN_ERROR, 'Unknown error occurred!')
        return (0, '')

    def check_common_args(self):
        if self._args['--output']:
            output = self._args['--output'].strip().strip("'").strip('"')
            output_abs = os.path.abspath(output)
            if os.path.isdir(output_abs):
                raise CrosspmExceptionWrongArgs(
                    '"%s" is a directory - can\'t write to it'
                )
            self._args['--output'] = output

        self.set_logging_level()

    def set_logging_level(self):
        level_str = self._args['--verbose'].strip().lower()

        log = self._args['--log']
        if log:
            log = log.strip().strip("'").strip('"')
            log_abs = os.path.abspath(log)
            if os.path.isdir(log_abs):
                raise CrosspmExceptionWrongArgs(
                    '"%s" is a directory - can\'t write log to it'
                )
            else:
                log_dir = os.path.dirname(log_abs)
                if not os.path.exists(log_dir):
                    os.makedirs(log_dir)
        else:
            log_abs = None

        level = Config.get_verbosity_level(level_str or 'console')
        if level or log_abs:
            self._log.setLevel(level)
            format_str = '%(asctime)-19s [%(levelname)-9s] %(message)s'
            if level_str == 'debug':
                format_str = '%(asctime)-19s [%(levelname)-9s] %(name)-12s: %(message)s'
            formatter = logging.Formatter(format_str, datefmt="%Y-%m-%d %H:%M:%S")

            if level:
                sh = logging.StreamHandler(stream=sys.stderr)
                sh.setLevel(level)
                # sh.setFormatter(formatter)
                self._log.addHandler(sh)

            if log_abs:
                if not level_str:
                    level = Config.get_verbosity_level(0)
                fh = logging.FileHandler(filename=log_abs)
                fh.setLevel(level)
                fh.setFormatter(formatter)
                self._log.addHandler(fh)

    def download(self):
        if self._args['--out-format'] == 'stdout':
            if self._args['--output']:
                raise CrosspmExceptionWrongArgs(
                    "unwanted argument '--output' while argument '--out-format={}'".format(
                        self._args['--out-format'],
                    ))
        elif not self._args['--output']:
            raise CrosspmExceptionWrongArgs(
                "argument '--output' required when argument '--out-format={}'".format(
                    self._args['--out-format'],
                ))

        params = {
            'out_format': ['--out-format', ''],
            'output': ['--output', ''],
            'out_prefix': ['--out-prefix', ''],
            # 'depslock_path': ['--depslock-path', ''],
        }

        for k, v in params.items():
            params[k] = self._args[v[0]] if v[0] in self._args else v[1]
            if isinstance(params[k], str):
                params[k] = params[k].strip('"').strip("'")

        do_load = not self._args['--list']
        if do_load:
            self._config.cache.auto_clear()
        cpm_downloader = Downloader(self._config, do_load)
        # cpm_downloader = Downloader(self._config, params.pop('depslock_path'), do_load)
        packages = cpm_downloader.download_packages()

        _not_found = any(_pkg is None for _pkg in packages.values())
        if _not_found:
            raise CrosspmException(
                CROSSPM_ERRORCODE_PACKAGE_NOT_FOUND,
                'Some package(s) not found.'
            )
        if do_load:
            self._output.write(params, packages)

    def promote(self):
        self._log.warning('This option is temporarily off.')
        # cpm_promoter = Promoter(self._config)
        # cpm_promoter.promote_packages()

    def pack(self):
        Archive.create(self._args['<OUT>'], self._args['<SOURCE>'])

    def cache(self):
        if self._args['clear']:
            self._config.cache.clear(self._args['hard'])
        elif self._args['size']:
            self._config.cache.size()
        elif self._args['age']:
            self._config.cache.age()
        else:
            self._config.cache.info()

if __name__ == '__main__':
    app = CrossPM()
    app.run()
<|endoftext|>


<commit_before>import matplotlib.pyplot as plt
from nptdms import TdmsFile
from numpy import array
from sys import argv

from chamber.models import props

def get_tdms_obj(file_path):
	# Returns the npTDMS object coresponding to the argument file
	print('Loading TDMS file...')
	tdms_obj = TdmsFile(file_path)
	print('TDMS file loaded.')
	return tdms_obj

def get_is_mass(tdms_obj):
	# Gets the IsMass value
	is_mass = tdms_obj.object('Settings', 'IsMass').data[0]
	return is_mass

def get_m_data(tdms_obj):
	# Gets the mass data if IsMass = 1
	m_data = tdms_obj.object('Data', 'Mass').data
	return m_data

def get_p_data(tdms_obj):
	# Gets the pressure data
	p_data = tdms_obj.object('Data', 'Pressure').data
	return p_data

def get_t_data(tdms_obj):
	# Gets the temperature data of the thermocouples
	# measureing chamber temperature
	t_data = array([tdms_obj.object(
		"Data", "TC{}".format(cpl_idx)).data for cpl_idx in range(4,14)])
	return t_data

def get_t_dp_data(tdms_obj):
	# Gets the dewpoint data
	t_dp_data = tdms_obj.object('Data', 'DewPoint').data
	return t_dp_data

def get_rh_data(t_data, p_data, t_dp_data):
	# Uses props to calculate the rh data
	t_data_avg = t_data.mean(axis=0)
	rh_data = props.get_rh(p_data, t_data_avg, t_dp_data)
	return rh_data

def plot_mass(m_data):
	# Plots mass vs time
	plt.plot(m_data, label='Mass ata')
	plt.xlabel('Time (s)')
	plt.ylabel('Mass (kg)')
	plt.title('Mass (kg) vs. Time (s)')
	plt.show()

def plot_rh_t_dp(rh_data, t_dp_data):
	# Plots Relative Humidity vs time and Dew Point vs time

	# Mass vs. time plot
	fig, ax1 = plt.subplots()
	ax1.plot(t_dp_data, 'r', label='Dew Point Data')
	ax1.set_xlabel('Time (s)')
	ax1.set_ylabel('Dew Point Temperature (K)', color='r')
	ax1.tick_params('y', colors='r')

	# TH vs time plot plot
	ax2 = ax1.twinx()
	ax2.plot(rh_data, 'b', label='Relative Humidity Data')
	ax2.set_ylabel('Relative Humidity', color='b')
	ax2.tick_params('y', colors='b')

	plt.title('Dew Point Temperature (K) and Relative Humidity vs. Time (s)')
	plt.show()

def plot_t_p(t_data, p_data):
	# Plots temperature vs time and pressure vs time

	# Temperature vs time plot
	fig, ax1 = plt.subplots()
	for idx in range(10):
		ax1.plot(t_data[idx], label='Tcpl {}'.format(idx + 4))
	ax1.set_xlabel('Time (s)')
	ax1.set_ylabel('Temperature (K)', color='r')
	ax1.tick_params('y', colors='r')
	plt.legend(bbox_to_anchor=(0, -0.05), loc=2, ncol=5)

	# Pressure vs time plot
	ax2 = ax1.twinx()
	ax2.plot(p_data, 'b', label='Pressure Data')
	ax2.set_ylabel('Pressure (kPa)', color='b')
	ax2.tick_params('y', colors='b')

	plt.title('Temperature (K) and Pressure (kPa) vs. Time (s)')
	plt.show()

def make_plots(file_path):
	tdms_obj = get_tdms_obj(file_path)

	print('Gathering data...')
	is_mass = get_is_mass(tdms_obj)
	if is_mass:
		m_data = get_m_data(tdms_obj)
	p_data = get_p_data(tdms_obj)
	t_data = get_t_data(tdms_obj)
	t_dp_data = get_t_dp_data(tdms_obj)
	rh_data = get_rh_data(t_data, p_data, t_dp_data)
	print('Data gathered.')

	print('Plotting...')
	if is_mass:
		plot_mass(m_data)
	plot_rh_t_dp(rh_data, t_dp_data)
	plot_t_p(t_data, p_data)

print('Starting QC check...')
FILEPATH = argv[1]
make_plots(FILEPATH)
print('QC check complete.')
<commit_msg>Complies with pycodestyle and pydocstyle.<commit_after>"""
Quality control script.

This script plots the experimental data for a man-in-the-loop QC visual
inspection.

"""
import matplotlib.pyplot as plt
from nptdms import TdmsFile
from numpy import array
from sys import argv

from chamber.models import props

def get_tdms_obj(file_path):
    """Return the npTDMS object coresponding to the argument file."""
    print('Loading TDMS file...')
    tdms_obj = TdmsFile(file_path)
    print('TDMS file loaded.')
    return tdms_obj

def get_is_mass(tdms_obj):
    """Get the IsMass value."""
    is_mass = tdms_obj.object('Settings', 'IsMass').data[0]
    return is_mass

def get_m_data(tdms_obj):
    """Get the mass data if IsMass = 1."""
    m_data = tdms_obj.object('Data', 'Mass').data
    return m_data

def get_p_data(tdms_obj):
    """Get the pressure data."""
    p_data = tdms_obj.object('Data', 'Pressure').data
    return p_data

def get_t_data(tdms_obj):
    """Get the temperature data of the thermocouples."""
    # measureing chamber temperature
    t_data = array([tdms_obj.object(
        "Data", "TC{}".format(cpl_idx)).data for cpl_idx in range(4, 14)])
    return t_data

def get_t_dp_data(tdms_obj):
    """Get the dewpoint data."""
    t_dp_data = tdms_obj.object('Data', 'DewPoint').data
    return t_dp_data

def get_rh_data(t_data, p_data, t_dp_data):
    """Use props to calculate the rh data."""
    t_data_avg = t_data.mean(axis=0)
    rh_data = props.get_rh(p_data, t_data_avg, t_dp_data)
    return rh_data

def plot_mass(m_data):
    """Plot mass vs time."""
    plt.plot(m_data, label='Mass ata')
    plt.xlabel('Time (s)')
    plt.ylabel('Mass (kg)')
    plt.title('Mass (kg) vs. Time (s)')
    plt.show()

def plot_rh_t_dp(rh_data, t_dp_data):
    """Plot Relative Humidity vs time and Dew Point vs time."""
    # Mass vs. time plot
    _, ax1 = plt.subplots()
    ax1.plot(t_dp_data, 'r', label='Dew Point Data')
    ax1.set_xlabel('Time (s)')
    ax1.set_ylabel('Dew Point Temperature (K)', color='r')
    ax1.tick_params('y', colors='r')

    # TH vs time plot plot
    ax2 = ax1.twinx()
    ax2.plot(rh_data, 'b', label='Relative Humidity Data')
    ax2.set_ylabel('Relative Humidity', color='b')
    ax2.tick_params('y', colors='b')

    plt.title('Dew Point Temperature (K) and Relative Humidity vs. Time (s)')
    plt.show()

def plot_t_p(t_data, p_data):
    """Plot temperature vs time and pressure vs time."""
    # Temperature vs time plot
    _, ax1 = plt.subplots()
    for idx in range(10):
        ax1.plot(t_data[idx], label='Tcpl {}'.format(idx + 4))
    ax1.set_xlabel('Time (s)')
    ax1.set_ylabel('Temperature (K)', color='r')
    ax1.tick_params('y', colors='r')
    plt.legend(bbox_to_anchor=(0, -0.05), loc=2, ncol=5)

    # Pressure vs time plot
    ax2 = ax1.twinx()
    ax2.plot(p_data, 'b', label='Pressure Data')
    ax2.set_ylabel('Pressure (kPa)', color='b')
    ax2.tick_params('y', colors='b')

    plt.title('Temperature (K) and Pressure (kPa) vs. Time (s)')
    plt.show()

def make_plots(file_path):
    """Make plots for QC."""
    tdms_obj = get_tdms_obj(file_path)

    print('Gathering data...')
    is_mass = get_is_mass(tdms_obj)
    if is_mass:
        m_data = get_m_data(tdms_obj)
    p_data = get_p_data(tdms_obj)
    t_data = get_t_data(tdms_obj)
    t_dp_data = get_t_dp_data(tdms_obj)
    rh_data = get_rh_data(t_data, p_data, t_dp_data)
    print('Data gathered.')

    print('Plotting...')
    if is_mass:
        plot_mass(m_data)
    plot_rh_t_dp(rh_data, t_dp_data)
    plot_t_p(t_data, p_data)

print('Starting QC check...')
FILEPATH = argv[1]
make_plots(FILEPATH)
print('QC check complete.')
<|endoftext|>