[595] Changes for running a notebook script (ecmwf#598)

* Changes

* Chanegs

* work

* change

* changes

* changes

* changes

* changes

* changes

* changes

* changes

* changes

* reverse old changes

* linter

Author: tjhunter

.gitignore

@@ -79,6 +79,14 @@ instance/
 # Scrapy stuff:
 .scrapy
 
+
+# Jupyter Notebook
+*.ipynb_checkpoints
+# Use the jupytext extension instead.
+*.ipynb
+
+*.zip
+
 # Sphinx documentation
 docs/_build/
 

packages/evaluate/pyproject.toml

@@ -11,6 +11,7 @@ dependencies = [
     "panel",
     "omegaconf",
     "weathergen-common",
+    "plotly>=6.2.0",
 ]
 
 [dependency-groups]

packages/evaluate/src/weathergen/evaluate/plot_inference.py

@@ -13,9 +13,9 @@
 from collections import defaultdict
 from pathlib import Path
 
-from omegaconf import OmegaConf
-from plotter import Plotter
-from utils import (
+from omegaconf import DictConfig, OmegaConf
+
+from weathergen.evaluate.utils import (
     calc_scores_per_stream,
     metric_list_to_json,
     plot_data,
@@ -25,23 +25,8 @@
 
 _logger = logging.getLogger(__name__)
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Fast evaluation of WeatherGenerator runs."
-    )
-    parser.add_argument(
-        "--config",
-        type=str,
-        help="Path to the configuration yaml file for plotting. e.g. config/plottig_config.yaml",
-    )
-
-    args = parser.parse_args()
-
-    # configure logging
-    logging.basicConfig(level=logging.INFO)
 
-    # load configuration
-    cfg = OmegaConf.load(args.config)
+def run_main(cfg: DictConfig) -> None:
     runs = cfg.run_ids
 
     _logger.info(f"Detected {len(runs)} runs")
@@ -52,14 +37,12 @@
     out_scores_dir = Path(cfg.output_scores_dir)
     out_scores_dir.mkdir(parents=True, exist_ok=True)
 
-    results_dir = Path(cfg.results_dir)
     metrics = cfg.evaluation.metrics
 
     # to get a structure like: scores_dict[metric][stream][run_id] = plot
     scores_dict = defaultdict(lambda: defaultdict(dict))
 
     for run_id, run in runs.items():
-        plotter = Plotter(cfg, run_id)
         _logger.info(f"RUN {run_id}: Getting data...")
 
         streams = run["streams"].keys()
@@ -71,7 +54,7 @@
 
             if stream_dict.get("plotting"):
                 _logger.info(f"RUN {run_id}: Plotting stream {stream}...")
-                plots = plot_data(cfg, run_id, stream, stream_dict)
+                plot_data(cfg, run_id, stream, stream_dict)
 
             if stream_dict.get("evaluation"):
                 _logger.info(f"Retrieve or compute scores for {run_id} - {stream}...")
@@ -109,9 +92,27 @@
                         scores_dict[metric][stream][run_id] = all_metrics.sel(
                             {"metric": metric}
                         )
+    # plot summary
+    if scores_dict and cfg.summary_plots:
+        _logger.info("Started creating summary plots..")
+        plot_summary(cfg, scores_dict, print_summary=cfg.print_summary)
 
 
-# plot summary
-if scores_dict and cfg.summary_plots:
-    _logger.info("Started creating summary plots..")
-    plot_summary(cfg, scores_dict, print_summary=cfg.print_summary)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Fast evaluation of WeatherGenerator runs."
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        help="Path to the configuration yaml file for plotting. e.g. config/plottig_config.yaml",
+    )
+
+    args = parser.parse_args()
+
+    # configure logging
+    logging.basicConfig(level=logging.INFO)
+
+    # load configuration
+    cfg = OmegaConf.load(args.config)
+    run_main(cfg)

packages/evaluate/src/weathergen/evaluate/score.py

@@ -14,7 +14,8 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
-from score_utils import to_list
+
+from weathergen.evaluate.score_utils import to_list
 
 # from common.io import MockIO
 

packages/evaluate/src/weathergen/evaluate/utils.py

@@ -15,11 +15,11 @@
 import numpy as np
 import omegaconf as oc
 import xarray as xr
-from plotter import DefaultMarkerSize, LinePlots, Plotter
-from score import VerifiedData, get_score
 from tqdm import tqdm
 
 from weathergen.common.io import ZarrIO
+from weathergen.evaluate.plotter import DefaultMarkerSize, LinePlots, Plotter
+from weathergen.evaluate.score import VerifiedData, get_score
 from weathergen.evaluate.score_utils import to_list
 
 _logger = logging.getLogger(__name__)
@@ -310,9 +310,7 @@ def plot_data(cfg: str, run_id: str, stream: str, stream_dict: dict) -> list[str
     # Check if histograms should be plotted
     plot_histograms = plot_settings.get("plot_histograms", False)
     if not isinstance(plot_settings.plot_histograms, bool):
-        raise TypeError(
-            "plot_histograms must be a boolean."
-        )
+        raise TypeError("plot_histograms must be a boolean.")
 
     if plot_fsteps == "all":
         plot_fsteps = None

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
  "dask~=2025.5.1",
  "hatchling",
  "weathergen-common",
+ "numexpr>=2.11.0",
 ]
 
 [project.urls]
@@ -53,9 +54,12 @@ packages = ["src/weathergen"]
 [dependency-groups]
 # The development dependencies
 dev = [
+ "ipykernel>=6.30.0",
+ "jupytext>=1.17.2",
  "pytest~=8.3.5",
  "pytest-mock>=3.14.1",
  "ruff==0.9.7",
+ "tensorboard>=2.20.0",
 ]
 
 

scripts/actions.sh

@@ -1,13 +1,13 @@
 #!/bin/bash
 
-# Get the directory where the script is located
+# TODO: this is the root weathergenerator directory, rename the variable.
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && cd .. && pwd)"
 
 case "$1" in
   sync)
     (
       cd "$SCRIPT_DIR" || exit 1
-      uv sync
+      uv sync --all-packages
     )
     ;;
   lint)
@@ -63,8 +63,27 @@ case "$1" in
       done
     )
     ;;
+  create-jupyter-kernel)
+    (
+      cd "$SCRIPT_DIR" || exit 1
+      uv sync --all-packages
+      uv run ipython kernel install --user --env VIRTUAL_ENV $(pwd)/.venv --name=weathergen_kernel --display-name "Python (WeatherGenerator)" 
+      echo "Jupyter kernel created. You can now use it in Jupyter Notebook or JupyterLab."
+      echo "To use this kernel, select 'Python (WeatherGenerator)' from the kernel options in Jupyter Notebook or JupyterLab."
+      echo "If you want to remove the kernel later, you can run:"
+      echo "jupyter kernelspec uninstall weathergen_kernel"
+    )
+    ;;
+  jupytext-sync)
+    (
+      cd "$SCRIPT_DIR" || exit 1
+      # Run on any python or jupyter notebook files in the WeatherGenerator-private/notebooks directory
+      uv run jupytext --set-formats ipynb,py:percent --sync  ../WeatherGenerator-private/notebooks/*.ipynb ../WeatherGenerator-private/notebooks/*.py
+      echo "Jupytext sync completed."
+    )
+    ;;
   *)
-    echo "Usage: $0 {sync|lint|unit-test|integration-test|create-links}"
+    echo "Usage: $0 {sync|lint|unit-test|integration-test|create-links|create-jupyter-kernel|jupytext-sync}"
     exit 1
     ;;
 esac

src/weathergen/utils/config.py

@@ -22,8 +22,8 @@
 
 _REPO_ROOT = Path(__file__).parent.parent.parent.parent  # TODO use importlib for resources
 _DEFAULT_CONFIG_PTH = _REPO_ROOT / "config" / "default_config.yml"
-_DEFAULT_MODEL_PATH = "./models"
-_DEFAULT_RESULT_PATH = "./results"
+_DEFAULT_MODEL_PATH = _REPO_ROOT / "models"
+_DEFAULT_RESULT_PATH = _REPO_ROOT / "results"
 
 _logger = logging.getLogger(__name__)
 
@@ -133,7 +133,7 @@ def load_config(
     if from_run_id is None:
         base_config = _load_default_conf()
     else:
-        base_config = load_model_config(from_run_id, epoch, private_config["model_path"])
+        base_config = load_model_config(from_run_id, epoch, private_config.get("model_path", None))
 
     # use OmegaConf.unsafe_merge if too slow
     return OmegaConf.merge(base_config, private_config, *overwrite_configs)
@@ -283,7 +283,7 @@ def _load_private_conf(private_home: Path | None) -> DictConfig:
         )
     private_cf = OmegaConf.load(private_home)
     private_cf["model_path"] = (
-        private_cf["model_path"] if "model_path" in private_cf.keys() else "./models"
+        private_cf["model_path"] if "model_path" in private_cf.keys() else None
     )
 
     if "secrets" in private_cf:
@@ -345,8 +345,9 @@ def load_streams(streams_directory: Path) -> list[Config]:
 def set_paths(config: Config) -> Config:
     """Set the configs run_path model_path attributes to default values if not present."""
     config = config.copy()
-    config.run_path = config.get("run_path", None) or _DEFAULT_RESULT_PATH
-    config.model_path = config.get("model_path", None) or _DEFAULT_MODEL_PATH
+    # pathlib.Path are not json serializable, so we convert them to str
+    config.run_path = config.get("run_path", None) or str(_DEFAULT_RESULT_PATH)
+    config.model_path = config.get("model_path", None) or str(_DEFAULT_MODEL_PATH)
 
     return config