Spaces:
Sleeping
Sleeping
| import pickle | |
| import os | |
| from src.utils.paths import get_path | |
| from src.utils.utils import CPU_Unpickler | |
| import argparse | |
| from src.jetfinder.clustering import get_clustering_labels | |
| import optuna | |
| from src.dataset.dataset import EventDataset | |
| from src.evaluation.clustering_metrics import compute_f1_score | |
| import torch | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| # filename = get_path("/work/gkrzmanc/jetclustering/results/train/Test_betaPt_BC_2025_01_03_15_07_14/eval_0.pkl", "results") | |
| # for rinv=0.7, see /work/gkrzmanc/jetclustering/results/train/Test_betaPt_BC_rinv07_2025_01_03_15_38_58 | |
| # keeping the clustering script here for now, so that it's separated from the GPU-heavy tasks like inference (clustering may be changed frequently...) | |
| # parameters: min-cluster-size: [5, 30] | |
| # min-samples: [2, 30] | |
| # epsilon: [0.01, 0.5] | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--input", type=str, required=True) | |
| parser.add_argument("--dataset", type=int, required=False, default=11) # Which dataset to optimize on | |
| parser.add_argument("--dataset-cap", type=int, required=False, default=-1) | |
| parser.add_argument("--spatial-components-only", "-spatial-only", action="store_true") | |
| parser.add_argument("--lorentz-cos-sim", action="store_true") | |
| parser.add_argument("--cos-sim", action="store_true") | |
| parser.add_argument("--normalize", action="store_true") | |
| # --input train/ --dataset-cap 1000 --spatial-components-only | |
| args = parser.parse_args() | |
| path = get_path(args.input, "results") | |
| suffix = "" | |
| if args.spatial_components_only: | |
| suffix = "_sp_comp_only" | |
| #if args.lorentz_norm: | |
| # suffix = "_lorentz_norm" | |
| if args.lorentz_cos_sim: | |
| suffix = "_lorentz_cos_sim" | |
| if args.cos_sim: | |
| suffix = "_cos_sim" | |
| if args.normalize: | |
| suffix = "_norm" | |
| study_file = os.path.join(path, "clustering_tuning_{}{}.log".format(args.dataset, suffix)) | |
| study_exists = os.path.exists(study_file) | |
| storage = optuna.storages.JournalStorage( | |
| optuna.storages.journal.JournalFileBackend(study_file) | |
| ) | |
| if study_exists: | |
| study = optuna.load_study(storage=storage, study_name="clustering") | |
| else: | |
| study = optuna.create_study(storage=storage, study_name="clustering", direction="maximize") | |
| eval_result_file = os.path.join(path, "eval_{}.pkl".format(args.dataset)) | |
| eval_result = CPU_Unpickler(open(eval_result_file, "rb")).load() | |
| dataset_cap = args.dataset_cap | |
| def objective(trial): | |
| min_clust_size = trial.suggest_int("min_cluster_size", 2, 20) | |
| min_samples = trial.suggest_int("min_samples", 0, 10) | |
| epsilon = trial.suggest_uniform("epsilon", 0.01, 0.5) | |
| print("Starting trial with parameters:", trial.params) | |
| suffix = "{}-{}-{}".format(min_clust_size, min_samples, epsilon) | |
| if args.spatial_components_only: | |
| suffix = "sp-" + suffix | |
| #if args.lorentz_norm: | |
| # suffix = "ln-" + suffix | |
| if args.cos_sim: | |
| suffix = "cs-" + suffix | |
| if args.lorentz_cos_sim: | |
| suffix = "lcs-" + suffix | |
| if args.normalize: | |
| suffix = "norm-" + suffix | |
| clustering_file = os.path.join(path, "clustering_{}_{}.pkl".format(suffix, args.dataset)) | |
| if not os.path.exists(clustering_file): | |
| if eval_result["pred"].shape[1] == 4: | |
| coords = eval_result["pred"][:, :3] | |
| else: | |
| if args.spatial_components_only or args.cos_sim: | |
| coords = eval_result["pred"][:, 1:4] | |
| else: | |
| coords = eval_result["pred"][:, :4] | |
| event_idx = eval_result["event_idx"] | |
| if dataset_cap > 0: | |
| filt = event_idx < dataset_cap | |
| event_idx = event_idx[filt] | |
| coords = coords[filt] | |
| if args.cos_sim or args.normalize: | |
| coords = coords / torch.norm(coords, dim=1, keepdim=True) | |
| labels = get_clustering_labels(coords, event_idx, min_cluster_size=min_clust_size, | |
| min_samples=min_samples, epsilon=epsilon, bar=True, | |
| lorentz_cos_sim=args.lorentz_cos_sim, | |
| cos_sim=args.cos_sim) | |
| with open(clustering_file, "wb") as f: | |
| pickle.dump(labels, f) | |
| print("Clustering saved to", clustering_file) | |
| #else: | |
| # labels = pickle.load(open(clustering_file, "rb")) | |
| print("Dataset:", eval_result["filename"]) | |
| dataset = EventDataset.from_directory(eval_result["filename"], | |
| model_clusters_file=clustering_file, | |
| model_output_file=eval_result_file, | |
| include_model_jets_unfiltered=True, parton_level=True, aug_soft=True) | |
| score = compute_f1_score(dataset, dataset_cap=dataset_cap) | |
| print("F1 score for", suffix, ":", score) | |
| return score | |
| study.optimize(objective, n_trials=100) | |
| print(f"Best params is {study.best_params} with value {study.best_value}") | |