Latent diffusion with fiddle configuration

.github/workflows/pre-commit.yml

@@ -18,6 +18,10 @@ jobs:
         with:
           enable-cache: true
 
+      # The config-diagram-* hooks require the `dot` binary from graphviz.
+      - name: Install graphviz
+        run: sudo apt-get update && sudo apt-get install -y graphviz
+
       # Run pre-commit through uv so CI uses the same project-managed environment
       # as local development, including local hooks that invoke `uv run ...`.
       - run: uv run pre-commit run --all-files

.pre-commit-config.yaml

@@ -28,11 +28,17 @@ repos:
   - id: pyproject-fmt
 - repo: local
   hooks:
-  - id: config-diagram
-    name: config diagram up to date
+  - id: config-diagram-convgru
+    name: convgru config diagram up to date
     language: system
     entry: uv run python docs/generate_base_experiment_config_diagram.py --check
-    files: ^src/mlcast/config/base\.py$
+    files: ^src/mlcast/config/(base|archetype/convgru)\.py$
+    pass_filenames: false
+  - id: config-diagram-latent-diffusion
+    name: latent diffusion config diagram up to date
+    language: system
+    entry: uv run python docs/generate_latent_diffusion_config_diagram.py --check
+    files: ^src/mlcast/config/archetype/latent_diffusion\.py$
     pass_filenames: false
 ci:
   autoupdate_schedule: monthly

AGENTS.md

@@ -23,6 +23,20 @@ Guidelines for contributors and AI agents working on this codebase.
 
 - Use `loguru` exclusively. Do not use the stdlib `logging` module.
 
+## Metric naming
+
+- Logged metric names follow TensorBoard conventions: use `/` as a hierarchy
+  separator (e.g. `val/loss`, `train/rec_loss`) so that related metrics are
+  grouped in the TensorBoard UI.
+- Use `rec_loss` for reconstruction-stage metrics and `loss` for
+  forecasting/diffusion-stage metrics to make clear which training stage a
+  metric belongs to.
+- Metric name format: `{split}/{name}` where `split` is `train`, `val`, or
+  `test`.
+- Monitoring references in config files (ModelCheckpoint, EarlyStopping,
+  lr_scheduler) must match the `val/{name}` variant of the metric they should
+  track.
+
 ## Code style
 
 - Docstrings follow NumPy style.

docs/generate_base_experiment_config_diagram.py

@@ -1,4 +1,4 @@
-"""Generate a Graphviz SVG diagram of the default training_experiment config.
+"""Generate a Graphviz SVG diagram of the included ConvGRU training config.
 
 Run without arguments to regenerate docs/config_diagram.svg:
 
@@ -15,13 +15,13 @@
 
 import fiddle.graphviz as fgv
 
-from mlcast.config import training_experiment
+from mlcast.config import convgru_training_experiment
 
 OUT = Path(__file__).parent / "config_diagram.svg"
 
 
 def main() -> None:
-    """Generate or verify the base experiment config diagram."""
+    """Generate or verify the ConvGRU training config diagram."""
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
         "--check",
@@ -30,7 +30,7 @@ def main() -> None:
     )
     args = parser.parse_args()
 
-    cfg = training_experiment.as_buildable()
+    cfg = convgru_training_experiment.as_buildable()
     g = fgv.render(cfg, max_str_length=40)
     g.format = "svg"
     new_svg = g.pipe().decode()

docs/generate_latent_diffusion_config_diagram.py

@@ -0,0 +1,52 @@
+"""Generate a Graphviz SVG diagram of the included latent diffusion training config.
+
+Run without arguments to regenerate docs/latent_diffusion_config_diagram.svg:
+
+    uv run python docs/generate_latent_diffusion_config_diagram.py
+
+Run with --check to verify the diagram is up to date:
+
+    uv run python docs/generate_latent_diffusion_config_diagram.py --check
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import fiddle.graphviz as fgv
+
+from mlcast.config import latent_diffusion_experiment
+
+OUT = Path(__file__).parent / "latent_diffusion_config_diagram.svg"
+
+
+def main() -> None:
+    """Generate or verify the latent diffusion training config diagram."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--check",
+        action="store_true",
+        help="Check that the diagram is up to date rather than regenerating it.",
+    )
+    args = parser.parse_args()
+
+    cfg = latent_diffusion_experiment.as_buildable()
+    g = fgv.render(cfg, max_str_length=40)
+    g.format = "svg"
+    new_svg = g.pipe().decode()
+
+    if args.check:
+        if not OUT.exists() or OUT.read_text() != new_svg:
+            print(
+                "docs/latent_diffusion_config_diagram.svg is out of date.\n"
+                "Run: uv run python docs/generate_latent_diffusion_config_diagram.py"
+            )
+            sys.exit(1)
+        print("docs/latent_diffusion_config_diagram.svg is up to date.")
+    else:
+        OUT.write_text(new_svg)
+        print(f"Written {OUT}")
+
+
+if __name__ == "__main__":
+    main()