examples/drich/drich_eval.py

0001 #!/usr/bin/env python3
0002 """
0003 Example: running one stage of an AID2E detector-optimization workflow.
0004 This stage runner is invoked by the workflow commands in workflow.yml.
0005 """
0006
0007 import argparse
0008 import contextlib
0009 import json
0010 from pathlib import Path
0011
0012 from aid2e.utilities.epic_utils import EpicLayerConfig
0013 from aid2e.utilities.workflows import JobContext, StackExecutionEngine, WorkflowSharedContext
0014 from drich_utils import (
0015     apply_overlap_policy,
0016     build_analysis_arguments,
0017     build_reco_arguments,
0018     build_sim_arguments,
0019     compute_drich_objectives,
0020     load_drich_config,
0021     make_paths,
0022     sim_reco_files,
0023     stage_scan_work,
0024 )
0025
0026
0027 # Stage runner
0028 def run_stack_job(layers, context):
0029     # This worker process reconstructs the minimum AID2E context needed by
0030     # StackExecutionEngine after DAGExecutor has already prepared geometry.
0031     workflow_context = WorkflowSharedContext(
0032         workflow_id=context["workflow_id"],
0033         parameters={"prepared_geometry_dir": str(context["prepared_geometry_dir"])},
0034     )
0035     job_context = JobContext(
0036         task_id=f"{context['stage']}:{context['run_name']}",
0037         job_id=context["run_name"],
0038         stage_id=context["stage"],
0039         workflow_id=context["workflow_id"],
0040         execution_dir=str(context["output_root"]),
0041         problem_config=context["problem"],
0042         workflow_context=workflow_context,
0043     )
0044
0045     try:
0046         return StackExecutionEngine(
0047             engine_id=context["run_name"],
0048             stack_type="epic",
0049             layers=layers,
0050         ).execute(job_context)
0051     except Exception as error:
0052         payload = job_context.xcom_pull(job_context.task_id) or {}
0053         error.returncode = payload.get("returncode", 1)
0054         raise
0055
0056
0057 def run_stack(context, layer_names):
0058     eval_config = context["eval_config"]
0059     layer_names = set(layer_names)
0060
0061     # geo layer: checkOverlaps failures can become configured penalty metrics.
0062     if "geo" in layer_names:
0063         overlap_log = context["overlap_log"]
0064         penalty_file = context["penalty_file"]
0065         penalty_file.unlink(missing_ok=True)
0066         stack_layers = [
0067             EpicLayerConfig(
0068                 layer="geo",
0069                 inputs=[context["detector_xml_ref"]],
0070                 outputs=[str(overlap_log)],
0071             )
0072         ]
0073         try:
0074             run_stack_job(stack_layers, context)
0075             return apply_overlap_policy(overlap_log, penalty_file, eval_config["failure_policy"]) or {"ok": 1.0}
0076         except Exception as error:
0077             return apply_overlap_policy(overlap_log, penalty_file, eval_config["failure_policy"], error=error) or {
0078                 "ok": 1.0
0079             }
0080
0081     npart = eval_config["npart"]
0082     particles = eval_config["particles"]
0083     trial_tag = context["trial_tag"]
0084     trial_xml = context["detector_xml_ref"]
0085     point, stage_particles = stage_scan_work(layer_names, context["job_index"], eval_config)
0086
0087     stack_layers = []
0088     for particle in stage_particles:
0089         sim_file, reco_file = sim_reco_files(context["sim_dir"], npart, trial_tag, point, particle)
0090         if "sim" in layer_names:
0091             stack_layers.append(
0092                 EpicLayerConfig(
0093                     layer="sim",
0094                     inputs=[],
0095                     outputs=[str(sim_file)],
0096                     arguments=["--compactFile", trial_xml, *build_sim_arguments(npart, point, particle)],
0097                     rule="{{command}} {{arguments}} {{inputs}} {{outputs}}",
0098                 )
0099             )
0100         if "rec" in layer_names:
0101             stack_layers.append(
0102                 EpicLayerConfig(
0103                     layer="rec",
0104                     inputs=[str(sim_file)],
0105                     outputs=[str(reco_file)],
0106                     arguments=build_reco_arguments(trial_xml),
0107                 )
0108             )
0109
0110     if "ana" in layer_names:
0111         analysis_binary = Path(context["config_path"]).parent / eval_config["analysis_executable"]
0112         analysis_output = (
0113             context["results_dir"]
0114             / f"recon_scan_{npart}_{trial_tag}_p_{point['p']}_eta_{point['eta_min']}_{point['eta_max']}.txt"
0115         )
0116         stack_layers.append(
0117             EpicLayerConfig(
0118                 layer="ana",
0119                 inputs=[
0120                     str(sim_reco_files(context["sim_dir"], npart, trial_tag, point, particle)[1])
0121                     for particle in particles
0122                 ],
0123                 outputs=[str(analysis_output)],
0124                 arguments=build_analysis_arguments(point, eval_config),
0125                 command=str(analysis_binary),
0126                 rule="{{command}} {{inputs}} {{outputs}} {{arguments}}",
0127             )
0128         )
0129
0130     run_stack_job(stack_layers, context)
0131     return {"ok": 1.0}
0132
0133
0134 def evaluate_design_point(trial_index, output_dir, config_path, stage, job_index, prepared_geometry_dir=None):
0135     config_path, cfg, eval_config = load_drich_config(config_path)
0136     paths = make_paths(output_dir)
0137     for path in (paths.log_dir, paths.results_dir):
0138         path.mkdir(parents=True, exist_ok=True)
0139
0140     trial_tag = str(trial_index)
0141     penalty_file = paths.log_dir / f"penalty_{trial_tag}.json"
0142     workflow_id = cfg.workflows.workflows[0].name
0143
0144     if stage == "retrieve_results":
0145         # retrieve_results is the objective-merge stage.
0146         if penalty_file.exists():
0147             return compute_drich_objectives(
0148                 paths.results_dir,
0149                 trial_tag=trial_tag,
0150                 eval_config=eval_config,
0151                 penalty=True,
0152             )
0153         return compute_drich_objectives(
0154             paths.results_dir,
0155             trial_tag,
0156             eval_config,
0157         )
0158
0159     layer_names = stage.split("_")
0160     if not all(layer_name in {"geo", "sim", "rec", "ana"} for layer_name in layer_names):
0161         raise ValueError(f"Unsupported stage: {stage}")
0162     if stage != "geo" and penalty_file.exists():
0163         return {"ok": 1.0}
0164     if prepared_geometry_dir is None:
0165         raise ValueError("--prepared-geometry-dir is required for stack stages")
0166
0167     for path in (paths.log_dir, paths.sim_dir):
0168         path.mkdir(parents=True, exist_ok=True)
0169
0170     run_name = f"{stage}_{trial_tag}_{job_index}"
0171     problem = cfg.problem
0172     with (
0173         (paths.log_dir / f"{run_name}.log").open("a") as log,
0174         contextlib.redirect_stdout(log),
0175         contextlib.redirect_stderr(log),
0176     ):
0177         # Activation sets EIC_SHELL/EPIC_INSTALL values used by EpicStack.
0178         problem.environment_config.activate()
0179
0180     context = {
0181         "output_root": paths.output_root,
0182         "config_path": config_path,
0183         "workflow_id": workflow_id,
0184         "stage": stage,
0185         "job_index": job_index,
0186         "problem": problem,
0187         "eval_config": eval_config,
0188         "trial_tag": trial_tag,
0189         "run_name": run_name,
0190         "results_dir": paths.results_dir,
0191         "sim_dir": paths.sim_dir,
0192         "prepared_geometry_dir": Path(prepared_geometry_dir),
0193         "detector_xml_ref": "{{geometry_dir}}/${DETECTOR_CONFIG}.xml",
0194         "overlap_log": paths.log_dir / f"overlap_log_{trial_tag}.txt",
0195         "penalty_file": penalty_file,
0196     }
0197     return run_stack(context, layer_names)
0198
0199
0200 # Main
0201
0202 def main(argv=None):
0203     parser = argparse.ArgumentParser(description="Run one detector worker stage")
0204     parser.add_argument("--trial-index", type=int, required=True)
0205     parser.add_argument("--output-dir", required=True)
0206     parser.add_argument("--config-path", required=True)
0207     parser.add_argument("--stage", required=True)
0208     parser.add_argument("--job-index", type=int, default=0)
0209     parser.add_argument("--prepared-geometry-dir")
0210     args = parser.parse_args(argv)
0211
0212     paths = make_paths(args.output_dir)
0213     metrics = evaluate_design_point(
0214         args.trial_index,
0215         args.output_dir,
0216         args.config_path,
0217         args.stage,
0218         args.job_index,
0219         args.prepared_geometry_dir,
0220     )
0221
0222     if args.stage == "retrieve_results":
0223         out_path = paths.results_dir / f"out-{args.trial_index}.json"
0224         out_path.parent.mkdir(parents=True, exist_ok=True)
0225         out_path.write_text(json.dumps(metrics, indent=2))
0226     return 0
0227
0228
0229 if __name__ == "__main__":
0230     raise SystemExit(main())