plot_score.py

from __future__ import print_function
from IMP import ArgumentParser


__doc__ = "Plot score distributions of good-scoring models."


def parse_args():
    parser = ArgumentParser(
                description="Plot score distributions of good-scoring models")
    parser.add_argument("-show", action="store_true",
                        help="Interactively show the plot")
    parser.add_argument("score_file",
                        help="Score file generated by select_good")
    parser.add_argument("column",
                        help="PMI stat file column to plot (or 'all' to "
                             "plot all columns)")
    return parser.parse_args()


def plot_column(plt, scores, column, show):
    print(column, "|| Mean:", scores[column].mean(),
          "| SD:", scores[column].std())
    plt.hist(scores[column], bins=100)
    plt.xlabel(column)
    plt.ylabel('P')
    plt.title(column)
    if show:
        plt.show()
    else:
        plt.savefig(column+".png", dpi=300)
    return plt


def main():
    args = parse_args()

    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import pandas

    scores = pandas.read_csv(args.score_file, sep=' ')

    if args.column == "all":
        print("Plotting all columns")
        exclude = frozenset(('Model_index', 'Replica_id', 'Frame_id', 'Run_id'))
        for k in scores.columns:
            if k not in exclude:
                p = plot_column(plt, scores, k, show=False)
                p.clf()
    elif args.column in scores.columns.values:
        plot_column(plt, scores, args.column, show=args.show)
    else:
        raise KeyError(args.column
              + " is not a valid score parameter. Use 'all' or one of: "
              + ", ".join(scores.columns.values))


if __name__ == '__main__':
    main()