eval.py

import argparse
import h5py
import numpy as np
import csv
import glob
from datasets import DATASETS, get_fn, prepare, get_h5_item


def get_all_results(dirname):
    """Yield (path, attrs, knns) for every valid result .h5 file under dirname."""
    seen = set()
    mask = dirname + "/**/*.h5"
    print(f"Searching for results matching: {mask}")
    for fn in glob.iglob(mask, recursive=True):
        if fn in seen:
            continue
        seen.add(fn)
        with h5py.File(fn, "r") as f:
            if "knns" not in f or "dataset" not in f.attrs or "task" not in f.attrs:
                print(f"Ignoring {fn}")
                continue
            print(fn)
            yield fn, dict(f.attrs), np.array(f["knns"])


def get_recall(I, gt, k):
    """
    Compute Recall@k averaged over all queries.
    
    Handles two result formats:
    - n x k: results are expected to exclude self-loops
    - n x (k+1): results are expected to include self-loops in first position
    
    Ground truth format is assumed to be n x m where m >= k+1,
    with self-loops in the first column.
    
    Always compares the k nearest neighbors EXCLUDING self-loops.
    """
    assert I.shape[0] == gt.shape[0], "query count mismatch between results and ground truth"
    assert gt.shape[1] >= k + 1, f"Ground truth needs at least {k+1} columns (self + {k} neighbors), got {gt.shape[1]}"
    
    if I.shape[1] == k + 1:
        # Results include self-loops: use columns [1:k+1] (skip first column)
        I_to_compare = I[:, 1:k+1]
        print(f"  Results shape {I.shape}: assumed self-loops, comparing columns [1:{k+1}]")
    elif I.shape[1] == k:
        # Results exclude self-loops: use columns [:k]
        I_to_compare = I[:, :k]
        print(f"  Results shape {I.shape}: assumed no self-loops, comparing columns [0:{k}]")
    else:
        raise ValueError(f"Results shape {I.shape} is neither {I.shape[0]}x{k} nor {I.shape[0]}x{k+1}")
    
    # Ground truth: always skip first column (self-loop) and use next k columns
    gt_to_compare = gt[:, 1:k+1]  # Skip self-loop, take next k neighbors
    
    hits = np.array([
        len(np.intersect1d(I_to_compare[i], gt_to_compare[i]))
        for i in range(len(I))
    ])
    return hits.sum() / (len(I) * k)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--results",
        help="directory in which results are stored",
        default="results",
    )
    parser.add_argument("csvfile")
    args = parser.parse_args()

    gt_cache = {}  # (dataset, task) -> gt_I array
    columns = ["dataset", "task", "algo", "buildtime", "querytime", "params", "recall"]

    with open(args.csvfile, "w", newline="") as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=columns, extrasaction="ignore")
        writer.writeheader()

        for fn, attrs, knns in get_all_results(args.results):
            dataset = attrs["dataset"]
            task = attrs["task"]

            if dataset not in DATASETS or task not in DATASETS[dataset]:
                print(f"Skipping {fn}: unknown dataset={dataset!r} task={task!r}")
                continue

            prepare(dataset, task)

            cache_key = (dataset, task)
            if cache_key not in gt_cache:
                gt_fn = get_fn(dataset, task)
                print(f"Loading ground truth from {gt_fn}")
                with h5py.File(gt_fn) as gf:
                    gt_I = np.array(get_h5_item(gf, DATASETS[dataset][task]["gt_I"]))
                gt_cache[cache_key] = gt_I

            gt_I = gt_cache[cache_key]
            k = DATASETS[dataset][task]["k"]
            recall = get_recall(knns, gt_I, k)
            row = dict(attrs)
            row["recall"] = recall
            print(dataset, task, attrs.get("algo"), attrs.get("params"), "=>", recall)
            writer.writerow(row)