sync utility feature mode code 2026-07-03

workspace/scripts/train_utility_energy.py

@@ -16,6 +16,7 @@ if str(PROJECT_ROOT) not in sys.path:
 
 import torch  # noqa: E402
 
+from cil.chart_features import CHART_FEATURE_MODES  # noqa: E402
 from cil.metrics import macro_micro_summary, pairwise_causal_dominance_ece  # noqa: E402
 from cil.models import (  # noqa: E402
     ChartEncoder,
@@ -42,6 +43,12 @@ def main(argv: list[str] | None = None) -> int:
     parser.add_argument("--listwise-rank", type=float, default=1.0)
     parser.add_argument("--margin", type=float, default=0.5)
     parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument(
+        "--chart-feature-mode",
+        choices=CHART_FEATURE_MODES,
+        default="base",
+        help="Deployment-visible chart feature family used by ChartEncoder.",
+    )
     args = parser.parse_args(argv)
 
     random.seed(args.seed)
@@ -50,7 +57,11 @@ def main(argv: list[str] | None = None) -> int:
     out_dir.mkdir(parents=True, exist_ok=True)
     _write_provenance(out_dir, args)
 
-    charts, index = load_charts(args.dataset, max_charts=args.max_charts)
+    charts, index = load_charts(
+        args.dataset,
+        max_charts=args.max_charts,
+        chart_feature_mode=args.chart_feature_mode,
+    )
     if not charts:
         raise SystemExit("utility energy training requires at least one outcome-visible chart")
     normalizer = TangentNormalizer.fit(torch.cat([chart.tangents for chart in charts], dim=0))
@@ -112,6 +123,7 @@ def main(argv: list[str] | None = None) -> int:
             "margin": args.margin,
             "temperature": args.temperature,
         },
+        "chart_feature_mode": args.chart_feature_mode,
         "rows": rows,
         "summary": summary,
     }
@@ -125,6 +137,7 @@ def main(argv: list[str] | None = None) -> int:
             "tangent_dim": tangent_dim,
             "data_hash": index.get("content_hash"),
             "split_hash": index.get("split_hash"),
+            "chart_feature_mode": args.chart_feature_mode,
         },
         out_dir / "model.pt",
     )