Warning, file /AID2E-framework/examples/evaluators/dtlz2_panda.py was not indexed
or was modified since last indexation (in which case cross-reference links may be missing, inaccurate or erroneous).
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.evaluators.dtlz2 import dtlz2_both_objectives
0009
0010
0011 def _qualified_design_vector(design_point: Dict[str, float]) -> List[float]:
0012 """Return DTLZ2 variables from canonical qualified design keys."""
0013 keys = [
0014 "DTLZ2_variables.x1",
0015 "DTLZ2_variables.x2",
0016 "DTLZ2_variables.x3",
0017 "DTLZ2_variables.x4",
0018 "DTLZ2_variables.x5",
0019 ]
0020 missing = [key for key in keys if key not in design_point]
0021 if missing:
0022 raise ValueError(
0023 "_qualified_design_vector: missing design keys "
0024 f"{missing}; got {list(design_point.keys())}"
0025 )
0026 return [float(design_point[key]) for key in keys]
0027
0028
0029 def panda_stage_simreco(context: JobContext, particle: str, eta_point: float, **kwargs) -> Dict[str, Any]:
0030 """Sim/reco PanDA stage used by result and dataset smoke tests."""
0031 x = _qualified_design_vector(context.design_point)
0032 if particle == "pi+":
0033 xyz = ((x[0] - 0.5) ** 3 + (x[1] - 0.5) ** 3) * float(eta_point)
0034 elif particle == "kaon+":
0035 xyz = ((x[0] - 0.5) ** 2 + (x[1] - 0.5) ** 2) * float(eta_point)
0036 else:
0037 xyz = 0.1 * float(eta_point)
0038
0039 result = {"xyz": float(xyz), "particle": particle, "eta_point": float(eta_point)}
0040 result["message"] = (
0041 f"simreco produced xyz={result['xyz']} for particle={particle} "
0042 f"eta_point={float(eta_point)}"
0043 )
0044 _write_panda_stage_outputs(result, kwargs)
0045 return result
0046
0047
0048 def _write_panda_stage_outputs(result: Dict[str, Any], kwargs: Dict[str, object]) -> None:
0049 output_file_name = kwargs.get("output_file_name") or kwargs.get("output_file")
0050 if not output_file_name:
0051 return
0052
0053 import json
0054 import os
0055
0056 with open(os.path.basename(str(output_file_name)), "w", encoding="utf-8") as handle:
0057 json.dump(
0058 {
0059 "xyz": result["xyz"],
0060 "message": result["message"],
0061 "particle": result["particle"],
0062 "eta_point": result["eta_point"],
0063 },
0064 handle,
0065 )
0066 handle.write("\n")
0067
0068
0069 def panda_stage_ana(
0070 context: JobContext,
0071 simreco_results: Dict[str, Dict[str, float]] | None = None,
0072 input_file_names: List[str] | None = None,
0073 **kwargs,
0074 ) -> Dict[str, float]:
0075 """PanDA analysis stage for result-based or dataset-backed sim/reco outputs."""
0076 if simreco_results is not None:
0077 xyz_sum = sum(float(result.get("xyz", 0.0)) for result in simreco_results.values())
0078 return {"xyz": float(xyz_sum), "n_inputs": len(simreco_results)}
0079
0080 xyz_sum = 0.0
0081 input_file_names = input_file_names or []
0082 for input_file_name in input_file_names:
0083 import json
0084
0085 with open(input_file_name, "r") as handle:
0086 data = json.load(handle)
0087 message = data.get("message", "")
0088 if message:
0089 print(f"ana read from {input_file_name}: {message}")
0090 xyz_sum += float(data.get("xyz", 0.0))
0091 return {"xyz": float(xyz_sum), "n_inputs": len(input_file_names)}
0092
0093
0094 def panda_stage_final(
0095 context: JobContext,
0096 ana_result: Dict[str, float] | None = None,
0097 ana_results: Dict[str, Dict[str, float]] | None = None,
0098 **kwargs,
0099 ) -> Dict[str, float]:
0100 """Local final stage that converts analysis results into objectives."""
0101 x = _qualified_design_vector(context.design_point)
0102 base = dtlz2_both_objectives(x)
0103 if ana_results is not None:
0104 xyz = sum(float(result.get("xyz not none ana-result", 0.0)) for result in ana_results.values())
0105 else:
0106 xyz = float((ana_result or {}).get("xyz", 0.0))
0107 objectives = {
0108 "f1": float(base["f1"] + xyz * 0.1),
0109 "f2": float(base["f2"] + xyz * 0.1),
0110 }
0111 context.xcom_push("objectives", objectives)
0112 return objectives
