precision_at_k.py

import argparse


def main():
    parser = argparse.ArgumentParser(
        description='Compute precision at k.')
    parser.add_argument(
            '-k', help='k, to compute precision@k', type=int, default=1)
    parser.add_argument('qrels', help='TREC relevance file (qrels)')
    parser.add_argument('topfile', help='TREC results file')
    args = parser.parse_args()

    ground_truth = {}
    with open(args.qrels) as fin:
        for line in fin:
            fields = line.strip().split()
            qid = int(fields[0])
            aid = int(fields[2])
            label = float(fields[3])
            if label == 1:
                ground_truth.setdefault(qid, set()).add(aid)

    n_questions = 0
    accumulated_precision = 0.0
    with open(args.topfile) as fin:
        # processing first line
        line = next(fin)
        n_questions += 1
        fields = line.strip().split()
        qid = int(fields[0])
        aid = int(fields[2])
        score = float(fields[4])
        current_qid = qid
        scores = []
        scores.append(score)
        # remaining lines
        for line in fin:
            fields = line.strip().split()
            qid = int(fields[0])
            if qid != current_qid:
                n_questions += 1
                # computing precision
                rank = sorted(range(len(scores)),
                              key=scores.__getitem__, reverse=True)
                precision_at_k = len(set(rank[:args.k]) & (ground_truth[current_qid])) / args.k
                accumulated_precision += precision_at_k
                # resetting
                current_qid = qid
                scores = []
            # line processing
            aid = int(fields[2])
            score = float(fields[4])
            scores.append(score)
        # processing last line
        rank = sorted(range(len(scores)),
                      key=scores.__getitem__, reverse=True)
        precision_at_k = len(set(rank[:args.k]) & (ground_truth[current_qid])) / args.k
        accumulated_precision += precision_at_k
        # average precision
        average = accumulated_precision / n_questions
        print('p@%d: %d/%d = %.4f'
              % (args.k, accumulated_precision, n_questions, average))


if __name__ == '__main__':
    main()