generate_training_data.py

# model input: <query, dist_start, dist_1st, dist_10th, 1st_to_start, 10th_to_start>
# model output: log2(predicted steps)


import matplotlib as mpl
mpl.use('Agg')  # noqa
import matplotlib.pyplot as plt
import numpy as np
import argparse
import os
import pickle
from tqdm import tqdm
import time
import json
from scipy import stats


from benchmark.datasets import DATASETS
from benchmark.algorithms.definitions import get_definitions
from benchmark.plotting.metrics import all_metrics as metrics
from benchmark.plotting.metrics import get_all_recall_values, get_count_at_certain_recall
from benchmark.plotting.utils import (get_plot_label, compute_metrics,
        create_linestyles, create_pointset)
from benchmark.results import (store_results, load_all_results, load_all_results_without_read, 
        get_result_filename, get_unique_algorithms)
from benchmark.dataset_io import knn_result_read
import benchmark.streaming.compute_gt
from benchmark.streaming.load_runbook import load_runbook
from benchmark.utils import read_gt_fromdir

from meta_analysis import read_float_arg_from_filename, read_int_arg_from_filename


def clopper_pearson_interval(successes, n, confidence=0.95):
    """
    计算二项比例的Clopper-Pearson精确置信区间
    
    参数:
    successes: 成功次数
    n: 总试验次数
    confidence: 置信水平，默认为0.95
    
    返回:
    (lower_bound, upper_bound): 置信区间的下界和上界
    """
    if n == 0:
        return 0.0, 0.0
    
    alpha = 1 - confidence
    lower_bound = stats.beta.ppf(alpha/2, successes, n - successes + 1)
    upper_bound = stats.beta.ppf(1 - alpha/2, successes + 1, n - successes)
    
    # 处理边界情况
    if successes == 0:
        lower_bound = 0.0
    if successes == n:
        upper_bound = 1.0
    
    return lower_bound, upper_bound


def calculate_gt_collected_confidence_intervals(gt_collected_queries, confidence=0.95):
    """
    为每个gt_collected矩阵元素计算双侧置信区间
    
    参数:
    gt_collected_queries: 形状为(nq, count+1, count)的数组，包含每个查询的结果
    confidence: 置信水平，默认为0.95
    
    返回:
    (mean, lower_ci, upper_ci): 平均值和置信区间
    """
    # 计算均值
    mean = np.mean(gt_collected_queries, axis=0)
    
    # 获取矩阵维度
    rows, cols = mean.shape
    
    # 初始化置信区间矩阵
    lower_ci = np.zeros_like(mean)
    upper_ci = np.zeros_like(mean)
    
    # 获取查询数量
    nq = gt_collected_queries.shape[0]
    
    # 对每个元素计算置信区间
    for i in range(rows):
        for j in range(cols):
            # 获取所有查询在该位置的值
            values = gt_collected_queries[:, i, j]
            
            # 计算成功次数（因为值是0或1）
            successes = np.sum(values)
            
            # 计算Clopper-Pearson置信区间
            lower_bound, upper_bound = clopper_pearson_interval(successes, nq, confidence)
            lower_ci[i, j] = lower_bound
            upper_ci[i, j] = upper_bound
    
    return mean, lower_ci, upper_ci


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dataset',
        metavar="DATASET",
        required=True)
    parser.add_argument(
        '--count',
        default=-1,
        type=int)
    parser.add_argument(
        '--effective_count',
        default=-1,
        type=int)
    parser.add_argument(
        '--definitions',
        metavar='FILE',
        help='load algorithm definitions from FILE',
        default='algos-2021.yaml')
    parser.add_argument(
        '--neurips23track',
        choices=['filter', 'ood', 'sparse', 'streaming', 'none'],
        default='none'
    )
    parser.add_argument(
        '--runbook_path',
        metavar='FILE',
        help='paths to runbooks',
    )
    parser.add_argument(
        '--results_base_path',
        type=str,
        default='results'
    )
    parser.add_argument(
        '--private-query',
        help='Use the private queries and ground truth',
        action='store_true')
    parser.add_argument(
        '--training_data_path',
        type=str,
        required=True)
    parser.add_argument(
        '--mode',
        type=str,
        required=True)
    parser.add_argument(
        '--filtered',
        help='Use filtered queries.',
        action='store_true'
    )
    parser.add_argument(
        '--label_file', 
        type=str, 
        default=None,
        help='Path to the label file.'
    )
    parser.add_argument(
        '--filter_label_file', 
        type=str, 
        default=None,
        help='Path to the filter file.'
    )
    parser.add_argument(
        '--num_queries',
        type=int,
        default=0,
        help='Number of queries to run.'
    )
    parser.add_argument(
        '--intermediate_cmps',
        type=int,
        default=None,
        help='Intermediate cmps to use in LAET mode.'
    )
    parser.add_argument(
        '--dump_gt_collected_confidence_intervals',
        help='Dump gt collected confidence intervals.',
        action='store_true'
    )
    args = parser.parse_args()

    assert args.mode.startswith('train'), "generate_training_data.py only supports train mode"

    dataset = DATASETS[args.dataset]()
    dim = dataset.d

    if args.count == -1:
        args.count = dataset.default_count()
    if args.effective_count == -1:
        args.effective_count = args.count

    count = int(args.count)
    effective_count = int(args.effective_count)

    Q = dataset.get_training_queries().astype(np.float32)
    if args.num_queries > 0:
        Q = Q[:args.num_queries]
    nq = Q.shape[0]
    print(fr"Got {nq} queries")

    max_pts, runbook = load_runbook(args.dataset, dataset.nb, args.runbook_path)
    results = load_all_results(args.dataset, count, neurips23track=args.neurips23track, runbook_path=args.runbook_path, \
                               filtered=args.filtered, label_file=args.label_file, filter_label_file=args.filter_label_file, base_path=args.results_base_path)

    for i, (fileroot, filename, properties, run) in enumerate(results):
        if not filename.startswith(args.mode):
            continue

        # Skip if intermediate cmps does not match in LAET mode
        if args.intermediate_cmps is not None:
            if read_int_arg_from_filename(filename, "intermediate", 0) != args.intermediate_cmps:
                continue

        print(f"from {fileroot}/{filename}")
        if fileroot.split("/")[-1].endswith('.txt'):
            label_file, filter_label_file = fileroot.split("/")[-2], fileroot.split("/")[-1]
        else:
            label_file, filter_label_file = "", ""
        gt_dir = benchmark.streaming.compute_gt.gt_dir(dataset, args.runbook_path, label_file, filter_label_file)

        nums_active_nodes = [0]
        for step, entry in enumerate(runbook):
            if entry['operation'] == 'search':
                nums_active_nodes.append(nums_active_nodes[-1])
            elif entry['operation'] == 'insert':
                nums_active_nodes.append(nums_active_nodes[-1] + (entry['end'] - entry['start']))
            elif entry['operation'] == 'delete':
                nums_active_nodes.append(nums_active_nodes[-1] - (entry['end'] - entry['start']))
            else:
                raise Exception(f'Undefined runbook operation {entry["operation"]}')

        for i in range(0, properties['num_searches']):
            start_time = time.time()

            search_step_id = properties['step_' + str(i)]
            step_suffix = str(search_step_id)
            N = nums_active_nodes[search_step_id]
            step_latency = properties['latency_step_' + step_suffix] / 10**6
            print(f"step {search_step_id}, N = {N}, latency = {step_latency} s")

            neighbors = np.array(run['neighbors_step' +  step_suffix])
            total_cmps = np.array(run['total_cmps_step' + step_suffix])
            total_latency = np.array(run['total_latency_step' + step_suffix])
            cmps = np.array(run['cmps_step' + step_suffix])
            lats = np.array(run['lats_step' + step_suffix])

            dists_start = np.array(run['dists_start_step' + step_suffix])
            dists_1st = np.array(run['dists_1st_step' + step_suffix])

            if args.mode.startswith("train_opt"): # train_opt
                dist_1st_hops = np.array(run['dist_1st_hops_step' + step_suffix])
                dist_1st_cmps = np.array(run['dist_1st_cmps_step' + step_suffix])
            elif args.mode.startswith("train_darth"): # train_darth, train_darth_opt
                dists_kth = np.array(run['dists_kth_step' + step_suffix])
            else: # train
                dists_10th = np.array(run['dists_10th_step' + step_suffix])

            if args.mode.startswith("train_opt"): # train_opt
                dists_visited = np.array(run['dists_visited_step' + step_suffix])
                cmps_visited = np.array(run['cmps_visited_step' + step_suffix])
                hops_visited = np.array(run['hops_visited_step' + step_suffix])
            elif args.mode.startswith("train_darth"): # train_darth, train_darth_opt
                cmps_visited = np.array(run['cmps_visited_step' + step_suffix])
                hops_visited = np.array(run['hops_visited_step' + step_suffix])
                inserts_visited = np.array(run['inserts_visited_step' + step_suffix])

            if args.mode.startswith('train_opt'):
                traversal_window_stats = np.array(run['traversal_window_stats_step' + step_suffix])

            if args.mode.startswith('train_darth'):
                result_set_stats = np.array(run['result_set_stats_step' + step_suffix])

            groundtruths_all, groundtruth_distances_all = read_gt_fromdir(gt_dir, step_suffix, count, train=True)
            groundtruths, groundtruth_distances = groundtruths_all[:, :effective_count], groundtruth_distances_all[:, :effective_count]

            mean_recall = 0
            
            training_data_path = args.training_data_path
            if not os.path.exists(training_data_path):
                os.makedirs(training_data_path)

            input_data_path = os.path.join(training_data_path, f"input_data_step{step_suffix}.npy")
            output_data_path = os.path.join(training_data_path, f"output_data_step{step_suffix}.npy")
            gt_cmps_path = os.path.join(training_data_path, f"gt_cmps_step{step_suffix}.npy")
            latency_path = os.path.join(training_data_path, f"latency_step{step_suffix}.json")

            # allocate input_data and output_data
            if args.mode.startswith('train_opt') or args.mode.startswith('train_darth'):
                rows = 0
                for total_cmps_per_query in total_cmps:
                    rows += total_cmps_per_query - 1
            else:
                rows = nq

            if args.mode.startswith('train_opt'):
                cols = 12
            elif args.mode.startswith('train_darth'):
                cols = 12
            else:
                cols = dim + 6

            input_data = np.zeros((rows, cols), dtype=np.float32)
            output_data = np.zeros((rows, 1), dtype=np.float32)
            gt_cmps = np.zeros((nq, effective_count), dtype=np.int32)
            curr_row = 0

            # 保存所有查询的gt_collected用于置信区间计算
            gt_collected_queries = np.zeros((nq, count+1, count), dtype=np.float32)
            gt_cmps_all = np.zeros((nq, count), dtype=np.int32)

            for (query_id, (
                neighbors_per_query, groundtruths_per_query, groundtruth_distances_per_query, 
                cmps_per_query, lats_per_query, 
                total_cmps_per_query, total_latency_per_query,
                dist_start
            )) in tqdm(enumerate(zip(
                neighbors, groundtruths, groundtruth_distances,
                cmps, lats, 
                total_cmps, total_latency,
                dists_start
            )), total=nq, desc="Generating training data"):

                recall_per_query = 0
                true_cmps_per_query = []
                true_cmps_per_query_sorted = []
                for rank, groundtruth in enumerate(groundtruths_per_query):
                    if groundtruth in neighbors_per_query:
                        rank_in_neighbors = np.where(neighbors_per_query == groundtruth)[0][0]
                        recall_per_query += 1
                        true_cmps_per_query.append(cmps_per_query[rank_in_neighbors])
                        true_cmps_per_query_sorted.append(cmps_per_query[rank_in_neighbors])
                    else:
                        true_cmps_per_query.append(total_cmps_per_query)
                mean_recall += recall_per_query
                true_cmps_per_query_sorted.sort()

                assert(len(true_cmps_per_query) == effective_count)
                np.copyto(gt_cmps[query_id], np.array(true_cmps_per_query))

                if args.mode.startswith('train_opt'):
                    assert(len(hops_visited[query_id]) == total_cmps_per_query - 1)

                    # model input: <curr_hops, curr_cmps, dist_1st, dist_start> + traversal_window_stats
                    # model output: probability that Top1 has been collected
                    input_data_rows = np.concatenate((
                        np.ones((total_cmps_per_query - 1, 1)) * query_id,
                        np.array(hops_visited[query_id]).reshape(-1, 1),
                        np.array(cmps_visited[query_id]).reshape(-1, 1),
                        np.array(dists_1st[query_id]).reshape(-1, 1),
                        np.ones((total_cmps_per_query - 1, 1)) * dist_start,
                        np.array(traversal_window_stats[query_id]).reshape(-1, 7),
                    ), axis=1)

                    output_rows = np.array(cmps_visited[query_id]) >= true_cmps_per_query[-1]
                    output_rows = output_rows.reshape(-1, 1)

                    # copy to input_data and output_data
                    np.copyto(input_data[curr_row : curr_row + total_cmps_per_query - 1], input_data_rows)
                    np.copyto(output_data[curr_row : curr_row + total_cmps_per_query - 1], output_rows)

                    curr_row += total_cmps_per_query - 1

                    groundtruths_all_per_query = groundtruths_all[query_id]
                    # 优化：使用向量化操作替代循环
                    neighbors_array = np.array(neighbors_per_query)
                    cmps_array = np.array(cmps_per_query)
                    
                    # 创建一个映射数组，初始值为total_cmps_per_query
                    true_cmps_all_per_query = np.full(len(groundtruths_all_per_query), total_cmps_per_query, dtype=np.int32)
                    
                    # 使用向量化操作找出在neighbors中的groundtruth
                    groundtruths_array = np.array(groundtruths_all_per_query)
                    
                    # 创建neighbors到索引的映射
                    neighbor_to_index = {neighbor: i for i, neighbor in enumerate(neighbors_array)}
                    
                    # 向量化查找和赋值
                    for i, groundtruth in enumerate(groundtruths_array):
                        if groundtruth in neighbor_to_index:
                            true_cmps_all_per_query[i] = cmps_array[neighbor_to_index[groundtruth]]
                    
                    # 使用部分向量化的优化方法
                    true_cmps_all_per_query = np.array(true_cmps_all_per_query)
                    gt_cmps_all[query_id] = true_cmps_all_per_query
                    
                    # 预先计算每个 collected 值对应的最大阈值
                    thresholds = np.full(count + 1, -np.inf, dtype=np.float32)
                    for c in range(1, min(count + 1, len(true_cmps_all_per_query) + 1)):
                        thresholds[c] = np.max(true_cmps_all_per_query[:c])
                    
                    # 逐步填充 gt_collected 数组
                    gt_collected_query = np.zeros((count+1, count), dtype=np.float32)
                    for collected in range(0, count + 1):
                        if collected == 0:
                            # 当 collected 为 0 时，所有元素都是 0
                            gt_collected_query[collected] = 0
                        else:
                            # 当 collected > 0 时
                            # 前 collected 个元素为 1
                            gt_collected_query[collected][:collected] = 1
                            
                            # 后面的元素根据条件判断
                            if collected < count and collected <= len(true_cmps_all_per_query):
                                cmps_threshold = thresholds[collected]
                                
                                # 向量化比较剩余元素
                                remaining_indices = np.arange(collected, count)
                                # 修复：直接使用remaining_indices作为comparison_indices，不再减1
                                comparison_indices = np.clip(remaining_indices, 0, len(true_cmps_all_per_query) - 1)
                                comparison_values = true_cmps_all_per_query[comparison_indices]
                                
                                # 执行比较并更新结果
                                mask = comparison_values <= cmps_threshold
                                gt_collected_query[collected][remaining_indices] = mask.astype(np.float32)
                    
                    # 保存当前查询的gt_collected
                    gt_collected_queries[query_id] = gt_collected_query

                elif args.mode.startswith('train_darth'):
                    assert(len(hops_visited[query_id]) == total_cmps_per_query - 1)

                    # model input: <curr_hops, curr_cmps, curr_inserts, dist_start, dist_1st, dist_kth> + result_set_stats
                    # model output: curr recall
                    input_data_rows = np.concatenate((
                        np.ones((total_cmps_per_query - 1, 1)) * query_id,
                        np.array(hops_visited[query_id]).reshape(-1, 1),
                        np.array(cmps_visited[query_id]).reshape(-1, 1),
                        np.array(inserts_visited[query_id]).reshape(-1, 1),
                        np.ones((total_cmps_per_query - 1, 1)) * dist_start,
                        np.array(dists_1st[query_id]).reshape(-1, 1),
                        np.array(dists_kth[query_id]).reshape(-1, 1),
                        np.array(result_set_stats[query_id]).reshape(-1, 5),
                    ), axis=1)

                    recalls = np.searchsorted(np.array(true_cmps_per_query_sorted), np.array(cmps_visited[query_id]), side='right')
                    output_rows = recalls / effective_count
                    output_rows = output_rows.reshape(-1, 1)

                    # copy to input_data and output_data
                    np.copyto(input_data[curr_row : curr_row + total_cmps_per_query - 1], input_data_rows)
                    np.copyto(output_data[curr_row : curr_row + total_cmps_per_query - 1], output_rows)

                    curr_row += total_cmps_per_query - 1

                else:
                    # model input: query + <dist_start, dist_1st, dist_10th, 1st_to_start, 10th_to_start>
                    # model output: log2(predicted steps)
                    if abs(dist_start) < 1e-6:
                        dist_start = 1e-6
                    input_data_row = [query_id] + Q[query_id].tolist() + [
                        dist_start, dists_1st[query_id], dists_10th[query_id],
                        dists_1st[query_id] / dist_start, dists_10th[query_id] / dist_start,
                    ]

                    if recall_per_query > 0:
                        output_data_row = [np.log2(true_cmps_per_query_sorted[-1])]
                    else:
                        output_data_row = [np.log2(total_cmps_per_query)]

                    # copy to input_data and output_data
                    np.copyto(input_data[curr_row], np.array(input_data_row))
                    np.copyto(output_data[curr_row], np.array(output_data_row))

                    curr_row += 1

            mean_recall = mean_recall / (effective_count * nq)
            print(f"recall: {mean_recall:.4f}")

            print(f"input shape: {input_data.shape}")
            print(f"output shape: {output_data.shape}")
            print(f"gt cmps shape: {gt_cmps.shape}")

            print(f"input examples: {input_data[:10]}")
            print(f"output examples: {output_data[:10]}")
            print(f"gt cmps examples: {gt_cmps[:10]}")

            np.save(input_data_path, input_data)
            np.save(output_data_path, output_data)
            np.save(gt_cmps_path, gt_cmps)
            print(f"saved training data to {input_data_path} and {output_data_path}")
            print(f"saved gt cmps to {gt_cmps_path}")

            if args.mode.startswith('train_opt'):
                gt_collected_avg = np.mean(gt_collected_queries, axis=0)    # (K+1, K)
                assert gt_collected_avg.shape == (count+1, count)
                gt_collected_path = os.path.join(training_data_path, f"gt_collected_step{step_suffix}.npy")
                np.save(gt_collected_path, gt_collected_avg)
                print(f"saved gt collected to {gt_collected_path}")

                gt_cmps_all_sorted = np.sort(gt_cmps_all, axis=0)           # (nq, K)
                gt_cmps_all_sorted = gt_cmps_all_sorted.T                   # (K, nq)
                gt_cmps_all_sorted_percentiles = np.percentile(gt_cmps_all_sorted, np.arange(1, 101), axis=1).T                                 # (K, 100)
                gt_cmps_all_sorted_percentiles = np.concatenate([np.zeros((1, 100), dtype=np.float32), gt_cmps_all_sorted_percentiles], axis=0)   # (K+1, 100)
                assert gt_cmps_all_sorted_percentiles.shape == (count+1, 100)
                gt_cmps_all_path = os.path.join(training_data_path, f"gt_cmps_all_step{step_suffix}.npy")
                np.save(gt_cmps_all_path, gt_cmps_all_sorted_percentiles)
                print(f"saved gt cmps all to {gt_cmps_all_path}")
                
                # 计算并保存双侧95%置信区间
                if args.dump_gt_collected_confidence_intervals:
                    mean, lower_ci, upper_ci = calculate_gt_collected_confidence_intervals(gt_collected_queries, confidence=0.95)
                    lower_bound_gt_collected_path = os.path.join(training_data_path, f"lower_bound_gt_collected_step{step_suffix}.npy")
                    upper_bound_gt_collected_path = os.path.join(training_data_path, f"upper_bound_gt_collected_step{step_suffix}.npy")
                    np.save(lower_bound_gt_collected_path, lower_ci)
                    np.save(upper_bound_gt_collected_path, upper_ci)
                    print(f"saved lower bound gt collected (95% CI) to {lower_bound_gt_collected_path}")
                    print(f"saved upper bound gt collected (95% CI) to {upper_bound_gt_collected_path}")

            generate_latency = time.time() - start_time
            print(f"run_latency: {step_latency} s, generate_latency: {generate_latency} s, total: {step_latency + generate_latency} s")
            with open(latency_path, "w") as f:
                json.dump({
                    "run_latency": step_latency,
                    "generate_latency": generate_latency,
                }, f, indent=4)

        print("")