pytrec issue

evaluate

#!/usr/bin/env python3


import os
import csv
import math
import argparse
import warnings
import collections
import pytrec_eval
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import ttest_rel, pearsonr, spearmanr


_default_measures = 'recip_rank,P_1,P_3,map,ndcg'


def _get_measure_min_relevant(measure):
	for m in ("P_",):
		if measure.startswith(m) and measure[len(m):].isnumeric():
			return int(measure[len(m):])
	return None


def _evaluate_measures(qrels, scores, measures, relevance_level):
	
	# Evaluate TREC measures
	trec_measures = set(measures) - {"pearson", "spearman"}
	evaluator = pytrec_eval.RelevanceEvaluator(qrels, trec_measures, relevance_level=relevance_level)
	results = evaluator.evaluate(scores)
	
	# Evaluate non-TREC measures, if needed
	if len(measures) > len(trec_measures):
		for id1, id2_qrels in qrels.items():
			id2_scores = scores[id1]
			qrels_vector = []
			similarities_vector = []
			for id2 in set(id2_qrels.keys()) | set(id2_scores.keys()):
				id2_qrel = float(id2_qrels[id2]) if id2 in id2_qrels else 0.0
				id2_score = id2_scores[id2] if id2 in id2_scores else 0.0
				qrels_vector.append(id2_qrel)
				similarities_vector.append(id2_score)
			with warnings.catch_warnings():
				warnings.simplefilter("ignore")
				if "pearson" in measures:
					r, _ = pearsonr(qrels_vector, similarities_vector)
					results[id1]["pearson"] = r
				if "spearman" in measures:
					r, _ = spearmanr(qrels_vector, similarities_vector)
					results[id1]["spearman"] = r
	
	return results


def _evaluate_aggregate_measure(measure, values):
	if measure in {"pearson", "spearman"}:
		values = [v for v in values if not math.isnan(v)]
		return sum(values) / len(values)
	return pytrec_eval.compute_aggregated_measure(measure, values)


def _main():
	
	# Parse command line
	parser = argparse.ArgumentParser()
	parser.add_argument(
		"input", metavar="FILE", nargs="+", type=str,
		help="system results files to evaluate (first one treated as baseline)")
	parser.add_argument(
		"-r", metavar="NUM", dest="relevance_level", type=int, default=1,
		help="threshold relevance level starting from which pairs are considered similar (for binary relevance measures, default 1)")
	parser.add_argument(
		"-n", metavar="NUM", dest="n", type=int, default=20,
		help="number of top-n similar ideas returned by a system to consider for each query idea (default 20)")
	parser.add_argument(
		"-m", metavar="MEASURES", dest="measures", type=str, default=_default_measures,
		help=f"comma-separated list of evaluation measures (default: '{_default_measures}')")
	parser.add_argument(
		"-l", metavar="LABELS", dest="labels", type=str,
		help=f"comma-separated list of labels for the approaches to evaluate (default: derive label from TSV file name)")
	parser.add_argument(
		'-g', metavar="FILE", dest='gold_input', required=False, type=str,
		help='gold standard TSV file containing <id1, id2, gold_similarity> rows')
	parser.add_argument(
		'-s', dest='significance_test', action='store_true', required=False,
		help='perform a two-sided paired t-test against the baseline (first system) storing p-values')
	parser.add_argument(
		'-b', dest='boxplots', action='store_true', default=False, required=False,
		help='emits also boxplots of evaluation measure distributions along query')
	parser.add_argument(
		'-o', metavar="FILE", dest='output', required=False, type=str,
		help='output TSV file (default to <stdout>)')
	args = parser.parse_args()
	
	# Load measures
	measures = args.measures.split(",")

	# Load graded relevances from gold standard (id1 -> query, id2 -> document, similarity -> relevance)
	qrels = collections.defaultdict(dict)
	nrels = collections.defaultdict(lambda: 0)
	df_gold = pd.read_csv(args.gold_input, index_col=False, sep="\t")
	for row in df_gold.itertuples(index=False):
		if row.similarity > 0:
			qrels[str(row.id1)][str(row.id2)] = int(row.similarity)
			nrels[str(row.id1)] += 1 if row.similarity >= args.relevance_level else 0
	
	# Evaluate each system, collecting one result row for each id1 and an extra row for aggregated (over id1s) measures
	rows = []
	systems = args.labels.split(",") if args.labels is not None else [os.path.basename(input_file).split(".")[0] for input_file in args.input]
	baseline = systems[0] if len(systems) > 0 else None
	baseline_results = None
	for input_file, system in zip(args.input, systems):
		
		# Load system results
		scores = collections.defaultdict(dict)
		df = pd.read_csv(input_file, index_col=False, sep="\t")
		df = df.groupby(["id1"]).apply(lambda df_group: df_group.sort_values("similarity", ascending=False).head(args.n))
		for row in df.itertuples(index=False):
			if str(row.id1) in qrels.keys():
				scores[str(row.id1)][str(row.id2)] = np.float64(row.similarity)
		
		# Apply TREC evaluation code (note: must create evaluator each time otherwise results are inconsistent)
		results = _evaluate_measures(qrels, scores, measures, args.relevance_level)
		
		# Generates rows with evaluation measures for each id1
		for id, id_results in sorted(results.items()):
			row = [id, system]
			rows.append(row)
			for measure in measures:
				row.append(id_results[measure])
		
		# Generate extra 'all' row with aggregated evaluation measures for all id1s
		row = ["all", system]
		rows.append(row)
		for measure in measures:
			min_relevant = _get_measure_min_relevant(measure)
			values = [id_results[measure] for id, id_results in results.items() if min_relevant is None or nrels[id] >= min_relevant]
			aggregated_value = _evaluate_aggregate_measure(measure, values)
			if measure == "map_cut_3":
				print(system, measure, values)
			row.append(aggregated_value)
		
		# Perform a two-sided paired t-test against the baseline, if requested
		if args.significance_test:
			if system == baseline:
				baseline_results = results
			else:
				row = ["pvalue", system]
				rows.append(row)
				for measure in measures:
					min_relevant = _get_measure_min_relevant(measure)
					baseline_values = [id_results[measure] for id, id_results in baseline_results.items() if min_relevant is None or nrels[id] >= min_relevant]
					system_values = [id_results[measure] for id, id_results in results.items() if min_relevant is None or nrels[id] >= min_relevant]
					_, pvalue = ttest_rel(baseline_values, system_values)
					row.append(pvalue)
	
	# Generate dataframe with generated result rows
	df_results = pd.DataFrame(columns=["id", "system"] + measures, data=rows)
	df_results.set_index(["id", "system"], inplace=True)

	# Save dataframe to TSV, if requested
	if args.output is not None:
		directory = os.path.dirname(args.output)
		basename = (directory if directory != "" else ".") + "/" + os.path.basename(args.output).split(".")[0]
		df_results.sort_values(["id", "system"], inplace=False).to_csv(basename + ".tsv", index=True, sep="\t", quoting=csv.QUOTE_NONE)
		if args.boxplots:
			df_plot = df_results.reset_index().melt(['id', 'system'], var_name='measure', value_name='value')
			df_plot['measure'] = df_plot['measure'].str.replace('recip_rank','MRR')
			df_plot['measure'] = df_plot['measure'].str.replace('P_1','P@1')
			df_plot['measure'] = df_plot['measure'].str.replace('P_3','P@3')
			df_plot['measure'] = df_plot['measure'].str.replace('map','MAP')
			df_plot['measure'] = df_plot['measure'].str.replace('ndcg','NDCG')
			plt.figure(figsize=(12,5))
			sns.set(style="white", palette="Blues", color_codes=True)
			sns.set_style({'xtick.bottom': True}, {'ytick.left': True})
			f = sns.boxplot(data=df_plot, x="measure", y="value", hue="system", palette="Blues", whis=(0,100), showmeans=True, meanprops={"marker":"D", "markerfacecolor":"white", "markeredgecolor":"black", "markersize":"5"}, linewidth=.5)
			f.set(xlabel='Measure', ylabel='Value')
			plt.legend(loc="center",  bbox_to_anchor=(0.5, 1.05),  frameon=False, ncol=len(systems))
			plt.tight_layout()
			plt.savefig(basename + ".pdf")
			plt.savefig(basename + ".svg")
	
	# Display aggregated results including statistical significance (pvalue <= 0.05 -> * marker)
	def format(df):
		cells = []
		for measure in measures:
			value = df.query("id == 'all'")[measure].squeeze()
			pvalue = 1.0 if df.shape[0] == 1 else df.query("id == 'pvalue'")[measure].squeeze()
			cells.append("{:.4f}".format(value) + ("*" if pvalue <= 0.05 else " "))
		return pd.Series(cells, index=[" " + m + " " for m in measures])
	df_display = df_results.query("id in ['all', 'pvalue']").groupby(["system"]).apply(format)
	df_display = df_display.reindex(systems)
	pd.set_option("display.max_rows", 1000)
	pd.set_option("display.max_column", 100)
	pd.set_option("display.width", 120)
	print("\n" + str(df_display) + "\n")


if __name__ == "__main__": _main()

gold.tsv

scores.tsv