import matplotlib.pyplot as plt import seaborn as sns import matplotlib.lines as mlines import pandas as pd import numpy as np from matplotlib.patches import FancyArrowPatch sns.set_theme(style="ticks", font_scale=1.2) plt.rcParams['axes.grid'] = True plt.rcParams['axes.grid.which'] = 'major' plt.rcParams['grid.linestyle'] = '--' plt.rcParams['grid.color'] = 'gray' plt.rcParams['grid.alpha'] = 0.3 plt.rcParams['xtick.minor.visible'] = False plt.rcParams['ytick.minor.visible'] = False plt.rcParams["font.family"] = "Helvetica" plt.rcParams["text.usetex"] = True plt.rcParams["font.weight"] = "bold" plt.rcParams["axes.labelweight"] = "bold" # Generation(LLama 1B) Generation(LLama 3B) Generation(LLama 7B) # 0.085s 0.217s 0.472s # llm_inference_time=[0.085, 0.217, 0.472, 0] # Will be replaced by CSV data # llm_inference_time_for_mac = [0.316, 0.717, 1.468, 0] # Will be replaced by CSV data def parse_latency_data(csv_path): df = pd.read_csv(csv_path) latency_data = {} llm_gen_times = {} # To store LLM generation times: (dataset, hardware) -> time for _, row in df.iterrows(): dataset = row['Dataset'] hardware = row['Hardware'] recall_target_str = row['Recall_target'].replace('%', '') try: recall_target = float(recall_target_str) except ValueError: print(f"Warning: Could not parse recall_target '{row['Recall_target']}'. Skipping row.") continue if (dataset, hardware) not in llm_gen_times: # Read once per (dataset, hardware) llm_time_val = pd.to_numeric(row.get('LLM_Gen_Time_1B'), errors='coerce') if not pd.isna(llm_time_val): llm_gen_times[(dataset, hardware)] = llm_time_val else: llm_gen_times[(dataset, hardware)] = np.nan # Store NaN if unparsable/missing cols_to_skip = ['Dataset', 'Hardware', 'Recall_target', 'LLM_Gen_Time_1B', 'LLM_Gen_Time_3B', 'LLM_Gen_Time_7B'] for col in df.columns: if col not in cols_to_skip: method_name = col key = (dataset, hardware, method_name) if key not in latency_data: latency_data[key] = [] try: latency_value = float(row[method_name]) latency_data[key].append((recall_target, latency_value)) except ValueError: # Handle cases where latency might be non-numeric (e.g., 'N/A' or empty) print(f"Warning: Could not parse latency for {method_name} at {dataset}/{hardware}/Recall {recall_target} ('{row[method_name]}'). Skipping this point.") latency_data[key].append((recall_target, np.nan)) # Or skip appending # Sort by recall for consistent plotting for key in latency_data: latency_data[key].sort(key=lambda x: x[0]) return latency_data, llm_gen_times def parse_storage_data(csv_path): df = pd.read_csv(csv_path) storage_data = {} # Assuming the first column is 'MetricType' (RAM/Storage) and subsequent columns are methods # And the header row is like: MetricType, Method1, Method2, ... # Transpose to make methods as rows for easier lookup might be an option, # but let's try direct parsing. # Find the row for RAM and Storage ram_row = df[df.iloc[:, 0] == 'RAM'].iloc[0] storage_row = df[df.iloc[:, 0] == 'Storage'].iloc[0] methods = df.columns[1:] # First column is the metric type label for method in methods: storage_data[method] = { 'RAM': pd.to_numeric(ram_row[method], errors='coerce'), 'Storage': pd.to_numeric(storage_row[method], errors='coerce') } return storage_data # Load data latency_csv_path = 'paper_plot/data/main_latency.csv' storage_csv_path = 'paper_plot/data/ram_storage.csv' latency_data, llm_generation_times = parse_latency_data(latency_csv_path) storage_info = parse_storage_data(storage_csv_path) # --- Determine unique Datasets and Hardware combinations to plot for --- unique_dataset_hardware_configs = sorted(list(set((d, h) for d, h, m in latency_data.keys()))) if not unique_dataset_hardware_configs: print("Error: No (Dataset, Hardware) combinations found in latency data. Check CSV paths and content.") exit() # --- Define constants for plotting --- all_method_names = sorted(list(set(m for d,h,m in latency_data.keys()))) if not all_method_names: # Fallback if latency_data is empty but storage_info might have method names all_method_names = sorted(list(storage_info.keys())) if not all_method_names: print("Error: No method names found in data. Cannot proceed with plotting.") exit() method_markers = { 'HNSW': 'o', 'IVF': 'X', 'DiskANN': 's', 'IVF-Disk': 'P', 'IVF-Recompute': '^', 'Our': '*', 'BM25': "v" # Add more if necessary, or make it dynamic } method_display_names = { 'IVF-Recompute': 'IVF-Recompute (EdgeRAG)', # 其他方法保持原名 } # Ensure all methods have a marker default_markers = ['^', 'v', '<', '>', 'H', 'h', '+', 'x', '|', '_'] next_default_marker = 0 for mn in all_method_names: if mn not in method_markers: print(f"mn: {mn}") method_markers[mn] = default_markers[next_default_marker % len(default_markers)] next_default_marker +=1 recall_levels_present = sorted(list(set(r for key in latency_data for r, l in latency_data[key]))) # Define colors for up to a few common recall levels, add more if needed base_recall_colors = { 85.0: "#1f77b4", # Blue 90.0: "#ff7f0e", # Orange 95.0: "#2ca02c", # Green # Add more if other recall % values exist } recall_colors = {} color_palette = sns.color_palette("viridis", n_colors=len(recall_levels_present)) for idx, r_level in enumerate(recall_levels_present): recall_colors[r_level] = base_recall_colors.get(r_level, color_palette[idx % len(color_palette)]) # --- Determine global x (latency) and y (storage) limits for consistent axes --- all_latency_values = [] all_storage_values = [] raw_data_size = 76 # Raw data size in GB for ds_hw_key in unique_dataset_hardware_configs: current_ds, current_hw = ds_hw_key for method_name in all_method_names: # Get storage for this method disk_storage = storage_info.get(method_name, {}).get('Storage', np.nan) if not np.isnan(disk_storage): all_storage_values.append(disk_storage) # Get latencies for this method under current_ds, current_hw latency_key = (current_ds, current_hw, method_name) if latency_key in latency_data: for recall, latency in latency_data[latency_key]: if not np.isnan(latency): all_latency_values.append(latency) # Add padding to limits min_lat = min(all_latency_values) if all_latency_values else 0.001 max_lat = max(all_latency_values) if all_latency_values else 1 min_store = min(all_storage_values) if all_storage_values else 0 max_store = max(all_storage_values) if all_storage_values else 1 # Convert storage values to proportion of raw data min_store_proportion = min_store / raw_data_size if all_storage_values else 0 max_store_proportion = max_store / raw_data_size if all_storage_values else 0.1 # Padding for log scale latency - adjust minimum to be more reasonable lat_log_min = -1 # Changed from -2 to -1 to set minimum to 10^-1 (0.1s) lat_log_max = np.log10(max_lat) if max_lat > 0 else 3 # default to 1000 s lat_padding = (lat_log_max - lat_log_min) * 0.05 global_xlim = [10**(lat_log_min - lat_padding), 10**(lat_log_max + lat_padding)] if global_xlim[0] <= 0: global_xlim[0] = 0.1 # Changed from 0.01 to 0.1 # Padding for linear scale storage proportion store_padding = (max_store_proportion - min_store_proportion) * 0.05 global_ylim = [max(0, min_store_proportion - store_padding), max_store_proportion + store_padding] if global_ylim[0] >= global_ylim[1]: # Avoid inverted or zero range global_ylim[1] = global_ylim[0] + 0.1 # After loading the data and before plotting, add this code to reorder the datasets # Find where you define all_datasets (around line 95) # Original code: all_datasets = sorted(list(set(ds for ds, _ in unique_dataset_hardware_configs))) # Replace with this to specify the exact order: all_datasets_unsorted = list(set(ds for ds, _ in unique_dataset_hardware_configs)) desired_order = ['NQ', 'TriviaQA', 'GPQA','HotpotQA'] all_datasets = [ds for ds in desired_order if ds in all_datasets_unsorted] # Add any datasets that might be in the data but not in our desired_order list all_datasets.extend([ds for ds in all_datasets_unsorted if ds not in desired_order]) # Then the rest of your code remains the same: a10_configs = [(ds, 'A10') for ds in all_datasets if (ds, 'A10') in unique_dataset_hardware_configs] mac_configs = [(ds, 'MAC') for ds in all_datasets if (ds, 'MAC') in unique_dataset_hardware_configs] # Create two figures - one for A10 and one for MAC hardware_configs = [a10_configs, mac_configs] hardware_names = ['A10', 'MAC'] for fig_idx, configs_for_this_figure in enumerate(hardware_configs): if not configs_for_this_figure: continue num_cols_this_figure = len(configs_for_this_figure) # 1 row, num_cols_this_figure columns fig, axs = plt.subplots(1, num_cols_this_figure, figsize=(7 * num_cols_this_figure, 6), sharex=True, sharey=True, squeeze=False) # fig.suptitle(f"Latency vs. Storage ({hardware_names[fig_idx]})", fontsize=18, y=0.98) for subplot_idx, (current_ds, current_hw) in enumerate(configs_for_this_figure): ax = axs[0, subplot_idx] # Accessing column in the first row ax.set_title(f"{current_ds}", fontsize=25) # No need to show hardware in title since it's in suptitle for method_name in all_method_names: marker = method_markers.get(method_name, '+') disk_storage = storage_info.get(method_name, {}).get('Storage', np.nan) latency_points_key = (current_ds, current_hw, method_name) if latency_points_key in latency_data: points_for_method = latency_data[latency_points_key] print(f"points_for_method: {points_for_method}") for recall, latency in points_for_method: # Only skip if latency is invalid (since we need log scale for x-axis) # But allow zero storage since y-axis is now linear if np.isnan(latency) or np.isnan(disk_storage) or latency <= 0: continue # Add LLM generation time from CSV current_llm_add_time = llm_generation_times.get((current_ds, current_hw)) if current_llm_add_time is not None and not np.isnan(current_llm_add_time): latency = latency + current_llm_add_time else: raise ValueError(f"No LLM generation time found for {current_ds} on {current_hw}") # Special handling for BM25 if method_name == 'BM25': # BM25 is only valid for 85% recall points (other points are 0) if recall != 85.0: continue color = 'grey' else: # Use the color for target recall color = recall_colors.get(recall, 'grey') # Convert storage to proportion disk_storage_proportion = disk_storage / raw_data_size size = 80 x_offset = -50 if current_ds == 'GPQA': x_offset = -32 # Apply a small vertical offset to IVF-Recompute points to make them more visible if method_name == 'IVF-Recompute': # Add a small vertical offset (adjust the 0.05 value as needed) disk_storage_proportion += 0.07 size = 80 if method_name == 'DiskANN': size = 50 if method_name == 'Our': size = 140 disk_storage_proportion += 0.05 # Add "Pareto Frontier" label to Our method points if recall == 95: ax.annotate('Ours', (latency, disk_storage_proportion), xytext=(x_offset, 25), # Increased leftward offset from -65 to -120 textcoords='offset points', fontsize=20, color='red', weight='bold', bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="red", alpha=0.7)) # Increase size for BM25 points if method_name == 'BM25': size = 70 size*=5 ax.scatter(latency, disk_storage_proportion, marker=marker, color=color, s=size, alpha=0.85, edgecolors='black', linewidths=0.7) ax.set_xscale("log") ax.set_yscale("linear") # CHANGED from log scale to linear scale for Y-axis # Generate appropriate powers of 10 based on your data range min_power = -1 max_power = 4 log_ticks = [10**i for i in range(min_power, max_power+1)] # Set custom tick positions ax.set_xticks(log_ticks) # Create custom bold LaTeX labels with 10^n format log_tick_labels = [fr'$\mathbf{{10^{{{i}}}}}$' for i in range(min_power, max_power+1)] ax.set_xticklabels(log_tick_labels, fontsize=24) # Apply global limits if subplot_idx == 0: ax.set_xlim(global_xlim) ax.set_ylim(global_ylim) ax.grid(True, which="major", linestyle="--", linewidth=0.6, alpha=0.7) # Remove minor grid lines completely ax.grid(False, which="minor") # Remove ticks # First set the shared parameters for both axes ax.tick_params(axis='both', which='both', length=0, labelsize=24) # Then set the padding only for the x-axis ax.tick_params(axis='x', which='both', pad=10) if subplot_idx == 0: # Y-label only for the leftmost subplot ax.set_ylabel("Proportional Size", fontsize=24) # X-label for all subplots in a 1xN layout can be okay, or just the middle/last one. # Let's put it on all for now. ax.set_xlabel("Latency (s)", fontsize=25) # Display 100%, 200%, 300% for yaxis ax.set_yticks([1, 2, 3]) ax.set_yticklabels(['100\%', '200\\%', '300\\%']) # Create a custom arrow with "Better" text inside # Create the arrow patch with a wider shaft arrow = FancyArrowPatch( (0.8, 0.8), # Start point (top-right) (0.65, 0.6), # End point (toward bottom-left) transform=ax.transAxes, arrowstyle='simple,head_width=40,head_length=35,tail_width=20', # Increased arrow dimensions facecolor='white', edgecolor='black', linewidth=3, # Thicker outline zorder=5 ) # Add the arrow to the plot ax.add_patch(arrow) # Calculate the midpoint of the arrow for text placement mid_x = (0.8 + 0.65) / 2 + 0.002 + 0.01 mid_y = (0.8 + 0.6) / 2 + 0.01 # Add the "Better" text at the midpoint of the arrow ax.text(mid_x, mid_y, 'Better', transform=ax.transAxes, ha='center', va='center', fontsize=16, # Increased font size from 12 to 16 fontweight='bold', rotation=40, # Rotate to match arrow direction zorder=6) # Ensure text is on top of arrow # Create legends (once per figure) method_legend_handles = [] for method, marker_style in method_markers.items(): if method in all_method_names: print(f"method: {method}") # Use black color for BM25 in the legend if method == 'BM25': method_legend_handles.append(mlines.Line2D([], [], color='black', marker=marker_style, linestyle='None', markersize=10, label=method)) else: if method in method_display_names: method = method_display_names[method] method_legend_handles.append(mlines.Line2D([], [], color='black', marker=marker_style, linestyle='None', markersize=10, label=method)) recall_legend_handles = [] sorted_recall_levels = sorted(recall_colors.keys()) for r_level in sorted_recall_levels: recall_legend_handles.append(mlines.Line2D([], [], color=recall_colors[r_level], marker='o', linestyle='None', markersize=20, label=f"Target Recall={r_level:.0f}\%")) # 将图例分成两行：第一行是方法，第二行是召回率 if fig_idx == 0: # 从方法列表中先排除'Our' other_methods = [m for m in all_method_names if m != 'Our'] # 按照需要的顺序创建方法列表（将'Our'放在最后） ordered_methods = other_methods + (['Our'] if 'Our' in all_method_names else []) # 按照新顺序创建方法图例句柄 method_legend_handles = [] for method in ordered_methods: if method in method_markers: marker_style = method_markers[method] # 使用显示名称映射 display_name = method_display_names.get(method, method) color = 'black' marker_size = 22 if method == 'Our': marker_size = 27 elif 'IVF-Recompute' in method or 'EdgeRAG' in method: marker_size = 17 elif 'DiskANN' in method: marker_size = 19 elif 'BM25' in method: marker_size = 20 method_legend_handles.append(mlines.Line2D([], [], color=color, marker=marker_style, linestyle='None', markersize=marker_size, label=display_name)) # 创建召回率图例（第二行）- 注意位置调整，放在方法图例下方 recall_legend = fig.legend(handles=recall_legend_handles, loc='upper center', bbox_to_anchor=(0.5, 1.05), # y坐标降低，放在第一行下方 ncol=len(recall_legend_handles), fontsize=28) # 创建方法图例（第一行） method_legend = fig.legend(handles=method_legend_handles, loc='upper center', bbox_to_anchor=(0.5, 0.91), ncol=len(method_legend_handles), fontsize=28) # 添加图例到渲染器 fig.add_artist(method_legend) fig.add_artist(recall_legend) # 调整布局，为顶部的两行图例留出更多空间 plt.tight_layout(rect=(0, 0, 1.0, 0.74)) # 顶部空间从0.9调整到0.85，给两行图例留出更多空间 save_path = f'./paper_plot/figures/main_exp_fig_{fig_idx+1}.pdf' plt.savefig(save_path, dpi=300, bbox_inches='tight') print(f"Saved figure {fig_idx+1} to {save_path}") plt.show()