Write to stats.yml, remove root

Author: sgreenbury

src/autosim/cli.py

@@ -257,16 +257,14 @@ def _stats_by_channel(values: torch.Tensor) -> dict[str, float]:
         }
 
     return {
-        "normalization_stats": {
-            "stats": {
-                "mean": _stats_by_channel(mean),
-                "std": _stats_by_channel(std),
-                "mean_delta": _stats_by_channel(mean_delta),
-                "std_delta": _stats_by_channel(std_delta),
-            },
-            "core_field_names": resolved_core_field_names,
-            "constant_field_names": constant_field_names or [],
-        }
+        "stats": {
+            "mean": _stats_by_channel(mean),
+            "std": _stats_by_channel(std),
+            "mean_delta": _stats_by_channel(mean_delta),
+            "std_delta": _stats_by_channel(std_delta),
+        },
+        "core_field_names": resolved_core_field_names,
+        "constant_field_names": constant_field_names or [],
     }
 
 
@@ -290,10 +288,7 @@ def _rounded_normalization_stats_payload(
         OmegaConf.to_container(OmegaConf.create(stats_payload), resolve=True),
     )
 
-    normalization_stats = rounded.get("normalization_stats")
-    if not isinstance(normalization_stats, dict):
-        return rounded
-    stats = normalization_stats.get("stats")
+    stats = rounded.get("stats")
     if not isinstance(stats, dict):
         return rounded
 
@@ -359,9 +354,7 @@ def generate_normalization_stats_yaml(
     )
 
     resolved_output_path = (
-        output_path
-        if output_path is not None
-        else dataset_dir / f"{dataset_dir.name}.yaml"
+        output_path if output_path is not None else dataset_dir / "stats.yml"
     )
     save_normalization_stats(
         stats_payload=stats_payload,
@@ -501,7 +494,7 @@ def _generate_main(cfg: Any) -> None:
     )
     save_normalization_stats(
         stats_payload=normalization_stats_payload,
-        output_path=output_dir / f"{output_dir.name}.yaml",
+        output_path=output_dir / "stats.yml",
     )
     save_example_videos(
         splits=splits,
@@ -575,10 +568,7 @@ def main() -> None:
         stats_parser.add_argument(
             "--output",
             default=None,
-            help=(
-                "Optional output YAML path (default: "
-                "<dataset_dir>/<dataset_dir_name>.yaml)."
-            ),
+            help=("Optional output YAML path (default: <dataset_dir>/stats.yml)."),
         )
         stats_parser.add_argument(
             "--field-names",

tests/test_cli.py

@@ -183,7 +183,7 @@ def test_compute_normalization_stats_includes_temporal_deltas() -> None:
     stats_payload = compute_normalization_stats(
         split_payload=split_payload, core_field_names=["U", "V"]
     )
-    stats = stats_payload["normalization_stats"]["stats"]
+    stats = stats_payload["stats"]
 
     assert stats["mean"]["U"] == pytest.approx(3.0)
     assert stats["mean"]["V"] == pytest.approx(14.0)
@@ -193,8 +193,8 @@ def test_compute_normalization_stats_includes_temporal_deltas() -> None:
     assert stats["mean_delta"]["V"] == pytest.approx(4.0)
     assert stats["std_delta"]["U"] == pytest.approx(0.0)
     assert stats["std_delta"]["V"] == pytest.approx(0.0)
-    assert stats_payload["normalization_stats"]["core_field_names"] == ["U", "V"]
-    assert stats_payload["normalization_stats"]["constant_field_names"] == []
+    assert stats_payload["core_field_names"] == ["U", "V"]
+    assert stats_payload["constant_field_names"] == []
 
 
 def test_generate_normalization_stats_yaml_from_existing_dataset(
@@ -228,11 +228,9 @@ def test_generate_normalization_stats_yaml_from_existing_dataset(
     assert output_path.exists()
     stats_cfg = OmegaConf.load(output_path)
     assert isinstance(stats_cfg, DictConfig)
-    assert stats_cfg["normalization_stats"]["core_field_names"] == ["U", "V"]
-    assert stats_cfg["normalization_stats"]["stats"]["mean"]["U"] == pytest.approx(2.5)
-    assert stats_cfg["normalization_stats"]["stats"]["mean_delta"][
-        "V"
-    ] == pytest.approx(2.0)
+    assert stats_cfg["core_field_names"] == ["U", "V"]
+    assert stats_cfg["stats"]["mean"]["U"] == pytest.approx(2.5)
+    assert stats_cfg["stats"]["mean_delta"]["V"] == pytest.approx(2.0)
 
 
 def test_cli_stats_subcommand_writes_yaml(tmp_path: Path) -> None:
@@ -269,16 +267,12 @@ def test_cli_stats_subcommand_writes_yaml(tmp_path: Path) -> None:
         cwd=repo_root,
     )
 
-    stats_path = dataset_dir / f"{dataset_dir.name}.yaml"
+    stats_path = dataset_dir / "stats.yml"
     assert stats_path.exists()
     stats_cfg = OmegaConf.load(stats_path)
     assert isinstance(stats_cfg, DictConfig)
-    assert stats_cfg["normalization_stats"]["stats"]["mean_delta"][
-        "U"
-    ] == pytest.approx(1.0)
-    assert stats_cfg["normalization_stats"]["stats"]["std_delta"]["V"] == pytest.approx(
-        0.0
-    )
+    assert stats_cfg["stats"]["mean_delta"]["U"] == pytest.approx(1.0)
+    assert stats_cfg["stats"]["std_delta"]["V"] == pytest.approx(0.0)
 
 
 def test_cli_help_outputs_usage() -> None: