99 Plot Raw Input Maps

.github/workflows/test_code.yaml

@@ -36,6 +36,9 @@ jobs:
         uses: astral-sh/setup-uv@v6
         with:
           version: "0.8.3"
+      - name: Install ffmpeg (so matplotlib can call it)
+        run: |
+          sudo apt-get update && sudo apt-get install -y ffmpeg
       - name: Run pytest
         run: uv run --group dev pytest
       - name: Run mypy

.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       # Run the linter
       - id: ruff-check
         args: ["--fix", "--show-fixes"]
-        pass_filenames: false
+        pass_filenames: true
       # Run the formatter
       - id: ruff-format
-        pass_filenames: false
+        pass_filenames: true

README.md

@@ -84,6 +84,22 @@ Run `uv run zebra train` to train using the datasets specified in the config.
 
 Run `uv run zebra evaluate --checkpoint PATH_TO_A_CHECKPOINT` to evaluate using a checkpoint from a training run.
 
+### Visualisations
+
+Plot raw input variables from the test dataset:
+
+```bash
+uv run zebra visualisations plot-raw-inputs --config-name <your local config>.yaml --sample-idx 0
+```
+
+Create animations of raw inputs over time:
+
+```bash
+uv run zebra visualisations animate-raw-inputs --config-name <your local config>.yaml
+```
+
+Settings (output directories, styling, animation parameters) are read from `config.evaluate.callbacks.raw_inputs` in your YAML config files. Command-line options can override config values if needed.
+
 ## Adding a new model
 
 ### Background

ice_station_zebra/callbacks/__init__.py

@@ -1,11 +1,13 @@
 from .ema_weight_averaging_callback import EMAWeightAveragingCallback
 from .metric_summary_callback import MetricSummaryCallback
 from .plotting_callback import PlottingCallback
+from .raw_inputs_callback import RawInputsCallback
 from .unconditional_checkpoint import UnconditionalCheckpoint
 
 __all__ = [
     "EMAWeightAveragingCallback",
     "MetricSummaryCallback",
     "PlottingCallback",
+    "RawInputsCallback",
     "UnconditionalCheckpoint",
 ]

ice_station_zebra/callbacks/plotting_callback.py

@@ -1,3 +1,8 @@
+# Set matplotlib backend before any plotting imports
+import matplotlib as mpl
+
+mpl.use("Agg")
+
 import dataclasses
 import logging
 from collections.abc import Mapping, Sequence
@@ -207,7 +212,7 @@ def log_static_plots(
                     )
         except InvalidArrayError as err:
             logger.warning("Static plotting skipped due to invalid arrays: %s", err)
-        except (ValueError, MemoryError, OSError):
+        except Exception:
             logger.exception("Static plotting failed")
 
     def log_video_plots(