@@ -102,10 +102,9 @@ class CategoricalSpec:
102102 FloatFactorSpec ("ITSVertexerParam.duplicateZCut[1]" , 0.5 , 2.0 , 1e-4 , 2.0 , log = True ),
103103 FloatFactorSpec ("ITSVertexerParam.finalSelectionZCut[1]" , 0.5 , 2.0 , 1e-4 , 2.0 , log = True ),
104104 FloatFactorSpec ("ITSVertexerParam.duplicateDistance2Cut[1]" , 0.5 , 2.0 , 1e-8 , 2.0 , log = True ),
105- IntFactorSpec ("ITSVertexerParam.clusterContributorsCut[1]" , 0.5 , 2.0 , 2 , 16 ),
106- CategoricalSpec ("ITSVertexerParam.seedMemberRadiusTime[1]" , (0 , 1 , 2 )),
107105 CategoricalSpec ("ITSVertexerParam.seedMemberRadiusZ[1]" , (0 , 1 , 2 , 3 , 4 )),
108106 FloatFactorSpec ("ITSVertexerParam.nSigmaCut[1]" , 0.5 , 2.0 , 1e-4 , 6.0 , log = True ),
107+ CategoricalSpec ("ITSVertexerParam.seedMemberRadiusZ[1]" , (0 , 1 , 2 , 3 , 4 )),
109108)
110109
111110UPC_BASE_DEFAULTS : dict [str , Any ] = {
@@ -120,7 +119,7 @@ class CategoricalSpec:
120119 "ITSVertexerParam.duplicateZCut[1]" : 0.1582193 ,
121120 "ITSVertexerParam.finalSelectionZCut[1]" : 0.1081465 ,
122121 "ITSVertexerParam.duplicateDistance2Cut[1]" : 0.0117033 ,
123- "ITSVertexerParam.clusterContributorsCut[1]" : 2 ,
122+ "ITSVertexerParam.clusterContributorsCut[1]" : 3 ,
124123 "ITSVertexerParam.seedMemberRadiusTime[1]" : 0 ,
125124 "ITSVertexerParam.seedMemberRadiusZ[1]" : 0 ,
126125 "ITSVertexerParam.nSigmaCut[1]" : 0.032841 ,
@@ -144,7 +143,9 @@ def parse_args() -> argparse.Namespace:
144143 parser .add_argument ("--input-line-cap-factor" , type = float , default = 1.5 , help = "Maximum allowed factor increase in total lines for input-stage trials" )
145144 parser .add_argument ("--input-fake-line-cap-factor" , type = float , default = 1.5 , help = "Maximum allowed factor increase in fake lines for input-stage trials" )
146145 parser .add_argument ("--input-fake-line-cap-add" , type = int , default = 200 , help = "Minimum additive fake-line allowance for input-stage trials" )
146+ parser .add_argument ("--failed-trial-score" , type = float , default = - 1.0 , help = "Objective value assigned to trials whose workflow fails" )
147147 parser .add_argument ("--artifact-dir" , type = Path , default = Path (".vertex_optuna" ), help = "Directory for trial workflows and summaries" )
148+ parser .add_argument ("--run-tracking" , action = "store_true" , help = "Run ITS tracking after vertexing instead of using vertexer-only trials" )
148149 parser .add_argument ("--dry-run" , action = "store_true" , help = "Patch and print one trial setup without running" )
149150 return parser .parse_args ()
150151
@@ -462,13 +463,16 @@ def patch_workflow_for_trial(
462463 workflow : dict [str , Any ],
463464 task_tag : str ,
464465 overrides : dict [str , Any ],
466+ run_tracking : bool = False ,
465467) -> dict [str , Any ]:
466468 patched = json .loads (json .dumps (workflow ))
467469 _ , patched_tasks = load_workflow_from_object (patched )
468470 prefix = f"{ task_tag } _"
469471 full_overrides = dict (overrides )
470472 full_overrides ["ITSVertexerParam.nIterations" ] = 2
471473 full_overrides ["ITSCATrackerParam.doUPCIteration" ] = 1
474+ if not run_tracking :
475+ full_overrides ["ITSCATrackerParam.nIterations" ] = 0
472476 for task in patched_tasks :
473477 name = str (task .get ("name" , "" ))
474478 if not name .startswith (prefix ):
@@ -530,8 +534,9 @@ def evaluate_trial(
530534 trial_name : str ,
531535 overrides : dict [str , Any ],
532536 selected_tasks : list [TaskInfo ],
537+ run_tracking : bool ,
533538) -> dict [str , Any ]:
534- patched = patch_workflow_for_trial (workflow , task_tag , overrides )
539+ patched = patch_workflow_for_trial (workflow , task_tag , overrides , run_tracking )
535540 workflow_copy = artifact_dir / f"{ trial_name } _workflow.json"
536541 with workflow_copy .open ("w" ) as stream :
537542 json .dump (patched , stream , indent = 2 )
@@ -573,7 +578,7 @@ def main() -> int:
573578 if args .dry_run :
574579 specs = INPUT_STAGE_SPECS if args .stage == "input" else VERTEX_STAGE_SPECS
575580 dry_overrides = {spec .key : base_params [spec .key ] for spec in specs }
576- patched = patch_workflow_for_trial (workflow , args .task_tag , dry_overrides )
581+ patched = patch_workflow_for_trial (workflow , args .task_tag , dry_overrides , args . run_tracking )
577582 workflow_copy = artifact_dir / "dry_run_workflow.json"
578583 with workflow_copy .open ("w" ) as stream :
579584 json .dump (patched , stream , indent = 2 )
@@ -594,6 +599,7 @@ def main() -> int:
594599 trial_name = "baseline" ,
595600 overrides = {},
596601 selected_tasks = selected_tasks ,
602+ run_tracking = args .run_tracking ,
597603 )
598604 print (f"Baseline: { summarize_result (baseline_metrics )} " )
599605
@@ -608,18 +614,25 @@ def main() -> int:
608614
609615 def objective (trial : optuna .Trial ) -> float :
610616 overrides = suggest_params (trial , args .stage , base_params )
611- metrics = evaluate_trial (
612- workflow = workflow ,
613- task_tag = args .task_tag ,
614- run_dir = run_dir ,
615- runner = runner ,
616- artifact_dir = artifact_dir ,
617- trial_name = f"trial_{ trial .number :04d} " ,
618- overrides = overrides ,
619- selected_tasks = selected_tasks ,
620- )
621617 for key , value in overrides .items ():
622618 trial .set_user_attr (key , value )
619+ try :
620+ metrics = evaluate_trial (
621+ workflow = workflow ,
622+ task_tag = args .task_tag ,
623+ run_dir = run_dir ,
624+ runner = runner ,
625+ artifact_dir = artifact_dir ,
626+ trial_name = f"trial_{ trial .number :04d} " ,
627+ overrides = overrides ,
628+ selected_tasks = selected_tasks ,
629+ run_tracking = args .run_tracking ,
630+ )
631+ except (subprocess .CalledProcessError , RuntimeError ) as exc :
632+ trial .set_user_attr ("workflow_failed" , True )
633+ trial .set_user_attr ("failure" , str (exc ))
634+ print (f"Trial { trial .number } : workflow failed -> { exc } " )
635+ return args .failed_trial_score
623636 for key in ("line_findable" , "true_lines" , "fake_lines" , "total_lines" , "unique_truth_ge2_lines" ,
624637 "findable" , "true_found" , "unique_true_findable" , "unique_true_all" , "total_found" ,
625638 "line_eff" , "eff" , "purity" , "fake_rate" , "duplicate_rate" , "f1" , "elapsed_ms" ):
@@ -643,7 +656,7 @@ def objective(trial: optuna.Trial) -> float:
643656 trial .set_user_attr ("guardrail_failed" , True )
644657 trial .set_user_attr ("guardrail_violations" , "," .join (input_guardrail_violations ))
645658 print (f"Trial { trial .number } : guardrail failed ({ ', ' .join (input_guardrail_violations )} ) -> { summarize_result (metrics )} " )
646- return - 1.0
659+ return args . failed_trial_score
647660
648661 violations = []
649662 baseline_eff_bins = baseline_metrics ["line_eff_bins" ] if args .stage == "input" else baseline_metrics ["eff_bins" ]
@@ -665,7 +678,7 @@ def objective(trial: optuna.Trial) -> float:
665678 trial .set_user_attr ("guardrail_failed" , True )
666679 trial .set_user_attr ("guardrail_violations" , "," .join (violations ))
667680 print (f"Trial { trial .number } : guardrail failed ({ ', ' .join (violations )} ) -> { summarize_result (metrics )} " )
668- return - 1.0
681+ return args . failed_trial_score
669682
670683 print (f"Trial { trial .number } : { summarize_result (metrics )} " )
671684 return metrics ["line_eff" ] if args .stage == "input" else metrics ["f1" ]
0 commit comments