[na] Echo scored cost_model weights in stats Refs #484

quon_na/src/lib.rs

@@ -130,8 +130,8 @@ pub use schedule_view::{
     ScheduleViewZone,
 };
 pub use stats::{
-    CompactionConfig, EffectiveConfig, NA_STATS_KIND, NA_STATS_SCHEMA_VERSION, NaStats,
-    SearchDiagnostics, StageTimingsUs, VersionInfo, na_stats_to_json,
+    CompactionConfig, CostModelWeights, EffectiveConfig, NA_STATS_KIND, NA_STATS_SCHEMA_VERSION,
+    NaStats, SearchDiagnostics, StageTimingsUs, VersionInfo, na_stats_to_json,
 };
 pub use zoned::{
     AWARE_NODE_BUDGET, AgnosticPlacerMechanism, AwareSearchOutcome, AwareSearchParams,

quon_na/src/pipeline.rs

@@ -807,6 +807,7 @@ fn finish_pipeline(
             placement_strategy,
             compaction: compaction_config,
             objective: opts.objective,
+            cost_model: backend_info.cost_model,
         },
         stage_timings_us: StageTimingsUs {
             extract_us: None,

quon_na/src/plan.rs

@@ -26,7 +26,7 @@ use crate::pipeline::{
 };
 use crate::placement::{PlacementStrategy, place};
 use crate::schedule_entry::GraphScheduleRequest;
-use crate::stats::SearchDiagnostics;
+use crate::stats::{CostModelWeights, SearchDiagnostics};
 use crate::zoned::{
     AgnosticPlacerMechanism, PlacementCostModel, PlacerMode, schedule_zoned_with_aware_params,
 };
@@ -65,6 +65,9 @@ pub struct BackendStageInfo {
     /// `Some(ScheduleOptimality::Exact)` when z3 proved optimality;
     /// `Some(ScheduleOptimality::Heuristic)` on timeout/fallback.
     pub schedule_optimality: Option<crate::report::ScheduleOptimality>,
+    /// §9 weights passed to the zoned placer. `None` for flat AOD, which
+    /// does not score with `cost_model`.
+    pub cost_model: Option<CostModelWeights>,
 }
 
 /// Run the place → AOD movement (or zoned routing) backend on a
@@ -130,6 +133,12 @@ pub fn plan_backend<V: VertexId>(
                     )
                 }
             };
+            info.cost_model = match cost_model {
+                PlacementCostModel::Weighted { weights, .. } => {
+                    Some(CostModelWeights::from(weights))
+                }
+                PlacementCostModel::Time | PlacementCostModel::ErrorBudget { .. } => None,
+            };
             let stage_started = Instant::now();
             let zoned = schedule_zoned_with_aware_params(
                 req,
@@ -271,6 +280,8 @@ pub struct QecStageAccumulator {
     pub aware_search_status: Option<(u64, u64)>,
     pub agnostic_placer_mechanism: Option<AgnosticPlacerMechanism>,
     pub schedule_optimality: Option<crate::report::ScheduleOptimality>,
+    /// §9 weights the zoned placer scored with. Same across phases.
+    pub cost_model: Option<crate::stats::CostModelWeights>,
     /// Per-phase exact state-prep solver outcome, aggregated across CNOT
     /// phases (issue #397). `None` when the heuristic scheduler ran (default
     /// `state_prep` mode, or a per-phase empty CNOT list). `Timeout` is
@@ -305,6 +316,9 @@ impl QecStageAccumulator {
         if self.placement_strategy.is_none() {
             self.placement_strategy = backend.placement_strategy;
         }
+        if self.cost_model.is_none() {
+            self.cost_model = backend.cost_model;
+        }
         // search_diagnostics: accumulate sums.
         let acc = &mut self.search_diagnostics;
         let inc = &backend.search_diagnostics;
@@ -441,16 +455,24 @@ mod tests {
     #[test]
     fn accumulator_accumulates() {
         let mut acc = QecStageAccumulator::default();
+        let weights = CostModelWeights {
+            rydberg_stage_weight: 7.0,
+            movement_time_weight: 0.25,
+            trap_transfer_weight: 3.0,
+            idle_time_weight: 0.5,
+        };
         let backend = BackendStageInfo {
             placement_us: Some(10),
             movement_us: Some(20),
+            cost_model: Some(weights),
             ..Default::default()
         };
         acc.accumulate_phase(5, 8, &backend);
         assert_eq!(acc.schedule_from_graph_us, 5);
         assert_eq!(acc.entangling_layers_us, 8);
         assert_eq!(acc.placement_us, Some(10));
         assert_eq!(acc.movement_us, Some(20));
+        assert_eq!(acc.cost_model, Some(weights));
 
         let backend2 = BackendStageInfo {
             movement_us: Some(30),

quon_na/src/qec_schedule.rs

@@ -255,6 +255,7 @@ fn schedule_expanded(
             placement_strategy: stage_acc.placement_strategy,
             compaction: compaction_config,
             objective: opts.objective,
+            cost_model: stage_acc.cost_model,
         },
         stage_timings_us: StageTimingsUs {
             extract_us: None,

quon_na/src/stats.rs

@@ -186,6 +186,35 @@ pub struct CompactionConfig {
     pub legality_checked: bool,
 }
 
+/// §9 `cost_model` weights the zoned placer scored with.
+///
+/// Copied from the [`crate::zoned::PlacementCostModel::Weighted`] vector
+/// passed to `schedule_zoned_with_aware_params`. Flat AOD does not score
+/// with these weights, so the echo stays absent there.
+#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct CostModelWeights {
+    #[serde(default)]
+    pub rydberg_stage_weight: f64,
+    #[serde(default)]
+    pub movement_time_weight: f64,
+    #[serde(default)]
+    pub trap_transfer_weight: f64,
+    #[serde(default)]
+    pub idle_time_weight: f64,
+}
+
+impl From<backend::NeutralAtomCostModel> for CostModelWeights {
+    fn from(weights: backend::NeutralAtomCostModel) -> Self {
+        Self {
+            rydberg_stage_weight: weights.rydberg_stage_weight,
+            movement_time_weight: weights.movement_time_weight,
+            trap_transfer_weight: weights.trap_transfer_weight,
+            idle_time_weight: weights.idle_time_weight,
+        }
+    }
+}
+
 /// Effective (as-run, not merely as-requested) pipeline configuration.
 #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
@@ -202,10 +231,14 @@ pub struct EffectiveConfig {
     pub placement_strategy: Option<PlacementStrategy>,
     #[serde(default)]
     pub compaction: CompactionConfig,
-    /// Placement/routing objective (issue #309): `Time` (default) or
-    /// `ErrorBudget`. Mirrors the `--na-objective` CLI knob.
+    /// Placement/routing objective mode: `Time` (default) or `ErrorBudget`.
+    /// Mirrors the `--na-objective` CLI knob. The weights that mode scored
+    /// with are [`Self::cost_model`].
     #[serde(default)]
     pub objective: NaObjective,
+    /// §9 weights the zoned optimizer scored with. `None` for flat AOD.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub cost_model: Option<CostModelWeights>,
 }
 
 /// Tool/target identifiers so a stats file is self-describing without a

quon_na/tests/snapshots/stats_snapshots__zoned_routing_agnostic.snap

@@ -11,6 +11,12 @@ expression: redact_timings(&json)
       "legality_checked": false,
       "requested": true
     },
+    "cost_model": {
+      "idle_time_weight": 1e-6,
+      "movement_time_weight": 1.0,
+      "rydberg_stage_weight": 1.0,
+      "trap_transfer_weight": 1.0
+    },
     "objective": "time",
     "placer_mode": "routing_agnostic"
   },

quon_na/tests/snapshots/stats_snapshots__zoned_routing_aware_config_only.snap

@@ -11,6 +11,12 @@ expression: redact_timings(&json)
       "legality_checked": false,
       "requested": true
     },
+    "cost_model": {
+      "idle_time_weight": 1e-6,
+      "movement_time_weight": 1.0,
+      "rydberg_stage_weight": 1.0,
+      "trap_transfer_weight": 1.0
+    },
     "objective": "time",
     "placer_mode": "routing_aware"
   },

quon_na/tests/stats_snapshots.rs

@@ -18,8 +18,8 @@
 use std::path::PathBuf;
 
 use quon_na::{
-    NaBackendKind, NaScheduleOptions, NaStats, PlacementStrategy, PlacerMode,
-    cubic_commutation_graph, na_stats_to_json, run_from_graph,
+    CostModelWeights, NaBackendKind, NaObjective, NaScheduleOptions, NaStats, PlacementStrategy,
+    PlacerMode, cubic_commutation_graph, na_stats_to_json, run_from_graph,
 };
 use serde_json::Value;
 
@@ -188,6 +188,10 @@ fn flat_aod_stats_shape() {
         stats.config.placement_strategy,
         Some(PlacementStrategy::RowMajor)
     );
+    assert!(
+        stats.config.cost_model.is_none(),
+        "flat AOD does not score with cost_model weights"
+    );
     // Aware-search concept doesn't apply to flat AOD.
     assert_eq!(stats.search.aware_search_completed_layers, None);
     assert_eq!(stats.search.aware_search_node_expansions, None);
@@ -223,6 +227,58 @@ fn no_compact_flag_reflects_requested_not_applied() {
     assert!(stats.stage_timings_us.compaction_us.is_none());
 }
 
+/// Zoned stats echo the §9 weights the placer scored with, not only the
+/// `Time` / `ErrorBudget` mode name.
+#[test]
+fn stats_echo_scored_cost_weights_not_only_the_mode() {
+    let mut na = na_target();
+    let weights = backend::NeutralAtomCostModel {
+        rydberg_stage_weight: 7.0,
+        movement_time_weight: 0.25,
+        trap_transfer_weight: 3.0,
+        idle_time_weight: 0.5,
+    };
+    na.cost_model = weights;
+    assert!(
+        na.error_model.is_some(),
+        "error-budget mode needs the target error_model"
+    );
+
+    for objective in [NaObjective::Time, NaObjective::ErrorBudget] {
+        let artifacts = run_from_graph(
+            cubic_commutation_graph(4).expect("cubic"),
+            &na,
+            NaScheduleOptions {
+                backend: NaBackendKind::Zoned,
+                placer: PlacerMode::RoutingAgnostic,
+                objective,
+                ..Default::default()
+            },
+            None,
+        )
+        .expect("schedule");
+        let stats = artifacts
+            .stats
+            .expect("run_from_graph always populates stats");
+        assert_eq!(stats.config.objective, objective);
+        assert_eq!(
+            stats.config.cost_model,
+            Some(CostModelWeights::from(weights)),
+            "echoed weights must be the vector the zoned placer scored with"
+        );
+
+        let json = na_stats_to_json(&stats).expect("serialize NaStats");
+        let value: Value = serde_json::from_str(&json).expect("parse NaStats JSON");
+        let mode = serde_json::to_value(objective).expect("mode name");
+        assert_eq!(value["config"]["objective"], mode);
+        let echoed = &value["config"]["cost_model"];
+        assert_eq!(echoed["rydberg_stage_weight"], 7.0);
+        assert_eq!(echoed["movement_time_weight"], 0.25);
+        assert_eq!(echoed["trap_transfer_weight"], 3.0);
+        assert_eq!(echoed["idle_time_weight"], 0.5);
+    }
+}
+
 /// A stats JSON with only the fields present in schema v1 must still
 /// deserialize (the `#[serde(default)]` field-evolution policy, issue #307).
 #[test]