File indexing completed on 2026-11-09 09:29:52
0001 """PanDA/iDDS-specific DTLZ2 evaluator stages."""
0002
0003 from __future__ import annotations
0004
0005 from typing import Any, Dict, List
0006
0007 from aid2e.utilities.workflows import JobContext
0008 from examples.dtlz2.dtlz2_evaluator import design_vector, dtlz2_both_objectives
0009
0010
0011 def panda_stage_simreco(context: JobContext, particle: str, eta_point: float, **kwargs) -> Dict[str, Any]:
0012 """Sim/reco PanDA stage used by result and dataset smoke tests."""
0013 x = design_vector(context.design_point)
0014 if particle == "pi+":
0015 xyz = ((x[0] - 0.5) ** 3 + (x[1] - 0.5) ** 3) * float(eta_point)
0016 elif particle == "kaon+":
0017 xyz = ((x[0] - 0.5) ** 2 + (x[1] - 0.5) ** 2) * float(eta_point)
0018 else:
0019 xyz = 0.1 * float(eta_point)
0020
0021 result = {"xyz": float(xyz), "particle": particle, "eta_point": float(eta_point)}
0022 result["message"] = (
0023 f"simreco produced xyz={result['xyz']} for particle={particle} "
0024 f"eta_point={float(eta_point)}"
0025 )
0026 _write_panda_stage_outputs(result, kwargs)
0027 return result
0028
0029
0030 def _write_panda_stage_outputs(result: Dict[str, Any], kwargs: Dict[str, object]) -> None:
0031 output_file_name = kwargs.get("output_file_name") or kwargs.get("output_file")
0032 if not output_file_name:
0033 return
0034
0035 import json
0036 import os
0037
0038 with open(os.path.basename(str(output_file_name)), "w", encoding="utf-8") as handle:
0039 json.dump(
0040 {
0041 "xyz": result["xyz"],
0042 "message": result["message"],
0043 "particle": result["particle"],
0044 "eta_point": result["eta_point"],
0045 },
0046 handle,
0047 )
0048 handle.write("\n")
0049
0050
0051 def panda_stage_ana(
0052 context: JobContext,
0053 simreco_results: Dict[str, Dict[str, float]] | None = None,
0054 input_file_names: List[str] | None = None,
0055 **kwargs,
0056 ) -> Dict[str, float]:
0057 """PanDA analysis stage for result-based or dataset-backed sim/reco outputs."""
0058 if simreco_results is not None:
0059 xyz_sum = sum(float(result.get("xyz", 0.0)) for result in simreco_results.values())
0060 return {"xyz": float(xyz_sum), "n_inputs": len(simreco_results)}
0061
0062 xyz_sum = 0.0
0063 input_file_names = input_file_names or []
0064 for input_file_name in input_file_names:
0065 import json
0066
0067 with open(input_file_name, "r") as handle:
0068 data = json.load(handle)
0069 message = data.get("message", "")
0070 if message:
0071 print(f"ana read from {input_file_name}: {message}")
0072 xyz_sum += float(data.get("xyz", 0.0))
0073 return {"xyz": float(xyz_sum), "n_inputs": len(input_file_names)}
0074
0075
0076 def panda_stage_final(
0077 context: JobContext,
0078 ana_result: Dict[str, float] | None = None,
0079 ana_results: Dict[str, Dict[str, float]] | None = None,
0080 **kwargs,
0081 ) -> Dict[str, float]:
0082 """Local final stage that converts analysis results into objectives."""
0083 x = design_vector(context.design_point)
0084 base = dtlz2_both_objectives(x)
0085 if ana_results is not None:
0086 xyz = sum(float(result.get("xyz not none ana-result", 0.0)) for result in ana_results.values())
0087 else:
0088 xyz = float((ana_result or {}).get("xyz", 0.0))
0089 objectives = {
0090 "f1": float(base["f1"] + xyz * 0.1),
0091 "f2": float(base["f2"] + xyz * 0.1),
0092 }
0093 context.xcom_push("objectives", objectives)
0094 return objectives
