From 81d6a21094c4943bf8a17aff9ddc4f29d46ee8fd Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Thu, 10 Sep 2026 22:05:35 +0200
Subject: [PATCH 01/13] Allow private config to override output directories and
 filename

---
 .../common/src/weathergen/common/config.py    | 31 +++++++++++++++----
 .../common/src/weathergen/common/logger.py    |  6 ++--
 src/weathergen/model/model_interface.py       |  2 +-
 src/weathergen/run_train.py                   |  6 ++--
 4 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 1f506fffb7..fe5e023394 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -208,7 +208,9 @@ def save(config: Config, mini_epoch: int | None):
         f.write(json_str)
 
 
-def load_run_config(run_id: str, mini_epoch: int | None, model_path: str | None) -> Config:
+def load_run_config(
+    run_id: str, mini_epoch: int | None, model_path: str | None, *, config: Config | None = None
+) -> Config:
     """
     Load a configuration file from a given run_id and mini_epoch.
     If run_id is a full path, loads it from the full path.
@@ -216,7 +218,8 @@ def load_run_config(run_id: str, mini_epoch: int | None, model_path: str | None)
     Args:
         run_id: Run ID of the pretrained WeatherGenerator model
         mini_epoch: Mini_epoch of the checkpoint to load. -1 indicates last checkpoint available.
-        model_path: Path to the model directory. If None, uses the model_path from private config.
+        model_path: Parent model directory containing run-id subdirectories.
+        config: Active configuration for resolving path_model when model_path is None.
 
     Returns:
         Configuration object loaded from the specified run and mini_epoch.
@@ -228,7 +231,7 @@ def load_run_config(run_id: str, mini_epoch: int | None, model_path: str | None)
     else:
         # Load model config here. In case model_path is not provided, get it from private conf
         if model_path is None:
-            path = get_path_model(run_id=run_id)
+            path = get_path_model(config, run_id=run_id)
         else:
             path = Path(model_path) / run_id
 
@@ -451,7 +454,7 @@ def load_merge_configs(
     if from_run_id is None:
         base_config = _load_base_conf(base)
     else:
-        base_config = load_run_config(from_run_id, mini_epoch, None)
+        base_config = load_run_config(from_run_id, mini_epoch, None, config=private_config)
         from_run_id = get_run_id_from_config(base_config)
     with open_dict(base_config):
         base_config.from_run_id = from_run_id
@@ -698,9 +701,23 @@ def load_streams(streams_directory: Path) -> Config:
     return OmegaConf.create(streams)
 
 
+def _get_output_directory(config: Config, key: str, folder: str, run_id: str) -> Path:
+    """Use an exact directory override, or the existing shared per-run location."""
+    if config.get(key) is not None:
+        return Path(config[key])
+    working_dir = config.get("path_shared_working_dir")
+    root = Path(working_dir) if working_dir is not None else _get_shared_wg_path()
+    return root / folder / run_id
+
+
+def get_path_logs(config: Config) -> Path:
+    """Get the application log directory."""
+    return _get_output_directory(config, "path_logs", "logs", get_run_id_from_config(config))
+
+
 def get_path_run(config: Config) -> Path:
     """Get the current runs results_path for storing run results and logs."""
-    return _get_shared_wg_path() / "results" / get_run_id_from_config(config)
+    return _get_output_directory(config, "path_results", "results", get_run_id_from_config(config))
 
 
 def get_path_model(config: Config | None = None, run_id: str | None = None) -> Path:
@@ -710,7 +727,8 @@ def get_path_model(config: Config | None = None, run_id: str | None = None) -> P
     else:
         msg = f"Missing run_id and cannot infer it from config: {config}"
         raise ValueError(msg)
-    return _get_shared_wg_path() / "models" / run_id
+    config = config if config is not None else _load_private_conf()
+    return _get_output_directory(config, "path_model", "models", run_id)
 
 
 def get_path_results(config: Config, mini_epoch: int) -> Path:
@@ -718,6 +736,7 @@ def get_path_results(config: Config, mini_epoch: int) -> Path:
     ext = StoreType(config.zarr_store).value  # validate extension
     base_path = get_path_run(config)
     fname = f"validation_chkpt{mini_epoch:05d}_rank{config.rank:04d}.{ext}"
+    fname = config.get("output_name") or fname
 
     return base_path / fname
 
diff --git a/packages/common/src/weathergen/common/logger.py b/packages/common/src/weathergen/common/logger.py
index a32d44ba32..b5312a3eda 100644
--- a/packages/common/src/weathergen/common/logger.py
+++ b/packages/common/src/weathergen/common/logger.py
@@ -96,7 +96,7 @@ def format(self, record, *args, **kwargs):
 
 
 @cache
-def init_loggers(run_id=None, logging_config=None):
+def init_loggers(run_id=None, logging_config=None, log_path=None):
     """
     Initialize the logger for the package and set output streams/files.
 
@@ -123,7 +123,7 @@ def init_loggers(run_id=None, logging_config=None):
     # output_dir = f"./output/{timestamp}-{run_id}"
     output_dir = ""
     if run_id is not None:
-        output_dir = f"./logs/{run_id}"
+        output_dir = str(log_path) if log_path is not None else f"./logs/{run_id}"
 
     # load the structure for logging config
     if logging_config is None:
@@ -142,7 +142,7 @@ def init_loggers(run_id=None, logging_config=None):
                 filename = f"{output_dir}/{v}"
                 ofile = pathlib.Path(filename)
                 # make sure the path is independent of path where job is launched
-                if not ofile.is_absolute():
+                if not ofile.is_absolute() and log_path is None:
                     work_dir = pathlib.Path(_load_private_conf().get("path_shared_working_dir"))
                     ofile = work_dir / ofile
                 pathlib.Path(ofile.parent).mkdir(parents=True, exist_ok=True)
diff --git a/src/weathergen/model/model_interface.py b/src/weathergen/model/model_interface.py
index c0b475b156..7617eaf684 100644
--- a/src/weathergen/model/model_interface.py
+++ b/src/weathergen/model/model_interface.py
@@ -181,7 +181,7 @@ def load_model(cf, model, device, run_id: str, mini_epoch=-1):
         mini_epoch : The mini_epoch to load. Default (-1) is the latest mini_epoch
     """
 
-    path_run = get_path_model(run_id=run_id)
+    path_run = get_path_model(cf, run_id=run_id)
     mini_epoch_id = (
         f"chkpt{mini_epoch:05d}" if mini_epoch != -1 and mini_epoch is not None else "latest"
     )
diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py
index 7995b5864f..91f3f96076 100644
--- a/src/weathergen/run_train.py
+++ b/src/weathergen/run_train.py
@@ -100,7 +100,7 @@ def run_inference(args):
     devices = Trainer.init_torch()
     cf = Trainer.init_ddp(cf)
 
-    init_loggers(cf.general.run_id)
+    init_loggers(cf.general.run_id, log_path=config.get_path_logs(cf))
 
     logger.info(f"DDP initialization: rank={cf.rank}, world_size={cf.world_size}")
 
@@ -139,7 +139,7 @@ def run_continue(args):
     devices = Trainer.init_torch(multiprocessing_method=mp_method)
     cf = Trainer.init_ddp(cf)
 
-    init_loggers(cf.general.run_id)
+    init_loggers(cf.general.run_id, log_path=config.get_path_logs(cf))
 
     # track history of run to ensure traceability of results
     cf.general.run_history += [(args.from_run_id, cf.general.istep)]
@@ -176,7 +176,7 @@ def run_train(args):
 
     # this line should probably come after the processes have been sorted out else we get lots
     # of duplication due to multiple process in the multiGPU case
-    init_loggers(cf.general.run_id)
+    init_loggers(cf.general.run_id, log_path=config.get_path_logs(cf))
 
     logger.info(f"DDP initialization: rank={cf.rank}, world_size={cf.world_size}")
 

From a42cf9b051e5bffa20e7aaa90413c4bbd1861755 Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Fri, 11 Sep 2026 07:24:55 +0200
Subject: [PATCH 02/13] Avoid reassigning typed config parameter when loading a
 run

---
 packages/common/src/weathergen/common/config.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index fe5e023394..eae89c4d91 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -258,9 +258,9 @@ def load_run_config(
     with fname.open() as f:
         json_str = f.read()
 
-    config = OmegaConf.create(json.loads(json_str))
+    loaded_config = OmegaConf.create(json.loads(json_str))
 
-    return _apply_fixes(config)
+    return _apply_fixes(loaded_config)
 
 
 def _get_model_config_file_write_name(run_id: str, mini_epoch: int | None):

From 944b8e67950b03d1a5d2319e7b3f49348ec6762e Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Fri, 11 Sep 2026 09:34:06 +0200
Subject: [PATCH 03/13] Rename exact checkpoint directory override to
 full_model_path

---
 packages/common/src/weathergen/common/config.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index eae89c4d91..f47ceb6bbe 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -219,7 +219,8 @@ def load_run_config(
         run_id: Run ID of the pretrained WeatherGenerator model
         mini_epoch: Mini_epoch of the checkpoint to load. -1 indicates last checkpoint available.
         model_path: Parent model directory containing run-id subdirectories.
-        config: Active configuration for resolving path_model when model_path is None.
+        config: Active configuration whose full_model_path specifies the exact checkpoint
+            directory (without appending run_id), used when model_path is None.
 
     Returns:
         Configuration object loaded from the specified run and mini_epoch.
@@ -721,14 +722,14 @@ def get_path_run(config: Config) -> Path:
 
 
 def get_path_model(config: Config | None = None, run_id: str | None = None) -> Path:
-    """Get the current runs model_path for storing model checkpoints."""
+    """Get full_model_path if set, otherwise the shared per-run checkpoint directory."""
     if config or run_id:
         run_id = run_id if run_id else get_run_id_from_config(config)
     else:
         msg = f"Missing run_id and cannot infer it from config: {config}"
         raise ValueError(msg)
     config = config if config is not None else _load_private_conf()
-    return _get_output_directory(config, "path_model", "models", run_id)
+    return _get_output_directory(config, "full_model_path", "models", run_id)
 
 
 def get_path_results(config: Config, mini_epoch: int) -> Path:

From 5ebe17c10c634e24f0ebf389f0ecf3bcdb9307d3 Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Fri, 11 Sep 2026 09:39:30 +0200
Subject: [PATCH 04/13] Clarify private configuration argument and variable
 names

---
 packages/common/src/weathergen/common/config.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index f47ceb6bbe..9e6a4f1859 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -209,7 +209,11 @@ def save(config: Config, mini_epoch: int | None):
 
 
 def load_run_config(
-    run_id: str, mini_epoch: int | None, model_path: str | None, *, config: Config | None = None
+    run_id: str,
+    mini_epoch: int | None,
+    model_path: str | None,
+    *,
+    private_config: Config | None = None,
 ) -> Config:
     """
     Load a configuration file from a given run_id and mini_epoch.
@@ -219,7 +223,7 @@ def load_run_config(
         run_id: Run ID of the pretrained WeatherGenerator model
         mini_epoch: Mini_epoch of the checkpoint to load. -1 indicates last checkpoint available.
         model_path: Parent model directory containing run-id subdirectories.
-        config: Active configuration whose full_model_path specifies the exact checkpoint
+        private_config: Configuration whose full_model_path specifies the exact checkpoint
             directory (without appending run_id), used when model_path is None.
 
     Returns:
@@ -232,7 +236,7 @@ def load_run_config(
     else:
         # Load model config here. In case model_path is not provided, get it from private conf
         if model_path is None:
-            path = get_path_model(config, run_id=run_id)
+            path = get_path_model(private_config, run_id=run_id)
         else:
             path = Path(model_path) / run_id
 
@@ -455,7 +459,7 @@ def load_merge_configs(
     if from_run_id is None:
         base_config = _load_base_conf(base)
     else:
-        base_config = load_run_config(from_run_id, mini_epoch, None, config=private_config)
+        base_config = load_run_config(from_run_id, mini_epoch, None, private_config=private_config)
         from_run_id = get_run_id_from_config(base_config)
     with open_dict(base_config):
         base_config.from_run_id = from_run_id
@@ -728,8 +732,8 @@ def get_path_model(config: Config | None = None, run_id: str | None = None) -> P
     else:
         msg = f"Missing run_id and cannot infer it from config: {config}"
         raise ValueError(msg)
-    config = config if config is not None else _load_private_conf()
-    return _get_output_directory(config, "full_model_path", "models", run_id)
+    private_config = config if config is not None else _load_private_conf()
+    return _get_output_directory(private_config, "full_model_path", "models", run_id)
 
 
 def get_path_results(config: Config, mini_epoch: int) -> Path:

From 194500768d4401c74b82a26ce59ae6f0c8e77a8b Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Fri, 11 Sep 2026 10:03:33 +0200
Subject: [PATCH 05/13] Renamed _get_output_directory to _get_path_output

Signed-off-by: evenmn <evenmn@fys.uio.no>
---
 packages/common/src/weathergen/common/config.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 9e6a4f1859..b739559ef2 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -706,7 +706,7 @@ def load_streams(streams_directory: Path) -> Config:
     return OmegaConf.create(streams)
 
 
-def _get_output_directory(config: Config, key: str, folder: str, run_id: str) -> Path:
+def _get_path_output(config: Config, key: str, folder: str, run_id: str) -> Path:
     """Use an exact directory override, or the existing shared per-run location."""
     if config.get(key) is not None:
         return Path(config[key])
@@ -717,12 +717,12 @@ def _get_output_directory(config: Config, key: str, folder: str, run_id: str) ->
 
 def get_path_logs(config: Config) -> Path:
     """Get the application log directory."""
-    return _get_output_directory(config, "path_logs", "logs", get_run_id_from_config(config))
+    return _get_path_output(config, "path_logs", "logs", get_run_id_from_config(config))
 
 
 def get_path_run(config: Config) -> Path:
     """Get the current runs results_path for storing run results and logs."""
-    return _get_output_directory(config, "path_results", "results", get_run_id_from_config(config))
+    return _get_path_output(config, "path_results", "results", get_run_id_from_config(config))
 
 
 def get_path_model(config: Config | None = None, run_id: str | None = None) -> Path:
@@ -733,7 +733,7 @@ def get_path_model(config: Config | None = None, run_id: str | None = None) -> P
         msg = f"Missing run_id and cannot infer it from config: {config}"
         raise ValueError(msg)
     private_config = config if config is not None else _load_private_conf()
-    return _get_output_directory(private_config, "full_model_path", "models", run_id)
+    return _get_path_output(private_config, "full_model_path", "models", run_id)
 
 
 def get_path_results(config: Config, mini_epoch: int) -> Path:

From 27c6a08790869edf5a59b9f0aaab7e8e21d809ae Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Fri, 11 Sep 2026 10:21:39 +0200
Subject: [PATCH 06/13] Resolve log directories in config before initializing
 loggers

---
 .../common/src/weathergen/common/logger.py    | 30 +++++--------------
 src/weathergen/run_train.py                   |  6 ++--
 2 files changed, 11 insertions(+), 25 deletions(-)

diff --git a/packages/common/src/weathergen/common/logger.py b/packages/common/src/weathergen/common/logger.py
index b5312a3eda..f543b93c41 100644
--- a/packages/common/src/weathergen/common/logger.py
+++ b/packages/common/src/weathergen/common/logger.py
@@ -14,8 +14,6 @@
 import pathlib
 from functools import cache
 
-from weathergen.common.config import _load_private_conf
-
 LOGGING_CONFIG = """
 {
     "version": 1,
@@ -96,10 +94,13 @@ def format(self, record, *args, **kwargs):
 
 
 @cache
-def init_loggers(run_id=None, logging_config=None, log_path=None):
+def init_loggers(log_path=None, logging_config=None):
     """
     Initialize the logger for the package and set output streams/files.
 
+    log_path is the directory resolved by config.get_path_logs(cf).
+    When omitted, the default logging configuration writes only to the console.
+
     WARNING: this function resets all the logging handlers.
 
     This function follows a singleton pattern, it will only operate once per process
@@ -115,21 +116,11 @@ def init_loggers(run_id=None, logging_config=None, log_path=None):
                 not supported
     """
 
-    # Get current time
-    # Shelved until decided how to change logging directory structure
-    # now = datetime.now()
-    # timestamp = now.strftime("%Y-%m-%d-%H%M")
-
-    # output_dir = f"./output/{timestamp}-{run_id}"
-    output_dir = ""
-    if run_id is not None:
-        output_dir = str(log_path) if log_path is not None else f"./logs/{run_id}"
-
     # load the structure for logging config
     if logging_config is None:
         logging_config = json.loads(LOGGING_CONFIG)
 
-        if run_id is None:
+        if log_path is None:
             del logging_config["handlers"]["logfile"]
             del logging_config["handlers"]["errorfile"]
             del logging_config["root"]["handlers"][2:]
@@ -139,12 +130,7 @@ def init_loggers(run_id=None, logging_config=None, log_path=None):
             if k == "formatter":
                 handler[k] = v
             elif k == "filename":
-                filename = f"{output_dir}/{v}"
-                ofile = pathlib.Path(filename)
-                # make sure the path is independent of path where job is launched
-                if not ofile.is_absolute() and log_path is None:
-                    work_dir = pathlib.Path(_load_private_conf().get("path_shared_working_dir"))
-                    ofile = work_dir / ofile
+                ofile = pathlib.Path(log_path or ".") / v
                 pathlib.Path(ofile.parent).mkdir(parents=True, exist_ok=True)
                 handler[k] = ofile
             else:
@@ -153,7 +139,7 @@ def init_loggers(run_id=None, logging_config=None, log_path=None):
     # make sure the parent directory exists
     logging.config.dictConfig(logging_config)
 
-    if output_dir:
-        logging.info(f"Logging set up. Logs are in {output_dir}")
+    if log_path is not None:
+        logging.info(f"Logging set up. Logs are in {log_path}")
     else:
         logging.info("Logging set up. No log files created.")
diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py
index 91f3f96076..5af219f7cf 100644
--- a/src/weathergen/run_train.py
+++ b/src/weathergen/run_train.py
@@ -100,7 +100,7 @@ def run_inference(args):
     devices = Trainer.init_torch()
     cf = Trainer.init_ddp(cf)
 
-    init_loggers(cf.general.run_id, log_path=config.get_path_logs(cf))
+    init_loggers(log_path=config.get_path_logs(cf))
 
     logger.info(f"DDP initialization: rank={cf.rank}, world_size={cf.world_size}")
 
@@ -139,7 +139,7 @@ def run_continue(args):
     devices = Trainer.init_torch(multiprocessing_method=mp_method)
     cf = Trainer.init_ddp(cf)
 
-    init_loggers(cf.general.run_id, log_path=config.get_path_logs(cf))
+    init_loggers(log_path=config.get_path_logs(cf))
 
     # track history of run to ensure traceability of results
     cf.general.run_history += [(args.from_run_id, cf.general.istep)]
@@ -176,7 +176,7 @@ def run_train(args):
 
     # this line should probably come after the processes have been sorted out else we get lots
     # of duplication due to multiple process in the multiGPU case
-    init_loggers(cf.general.run_id, log_path=config.get_path_logs(cf))
+    init_loggers(log_path=config.get_path_logs(cf))
 
     logger.info(f"DDP initialization: rank={cf.rank}, world_size={cf.world_size}")
 

From 07c1027e2d7ab2495f200702a8eb781a740265d6 Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Fri, 11 Sep 2026 10:29:56 +0200
Subject: [PATCH 07/13] Reverted 'loaded_config' to 'config' like it was before

Signed-off-by: evenmn <evenmn@fys.uio.no>
---
 packages/common/src/weathergen/common/config.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index b739559ef2..cc6135a980 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -263,9 +263,9 @@ def load_run_config(
     with fname.open() as f:
         json_str = f.read()
 
-    loaded_config = OmegaConf.create(json.loads(json_str))
+    config = OmegaConf.create(json.loads(json_str))
 
-    return _apply_fixes(loaded_config)
+    return _apply_fixes(config)
 
 
 def _get_model_config_file_write_name(run_id: str, mini_epoch: int | None):

From 2333e85d61b31edef5e18269fb348a0d5fb3b90c Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Sat, 12 Sep 2026 16:49:46 +0200
Subject: [PATCH 08/13] Add template formatting to output_name

---
 packages/common/src/weathergen/common/config.py | 17 ++++++++++++++---
 src/weathergen/utils/validation_io.py           |  2 +-
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index cc6135a980..fae3098b4d 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -736,12 +736,23 @@ def get_path_model(config: Config | None = None, run_id: str | None = None) -> P
     return _get_path_output(private_config, "full_model_path", "models", run_id)
 
 
-def get_path_results(config: Config, mini_epoch: int) -> Path:
+def get_path_results(config: Config, mini_epoch: int, step: int | None = None) -> Path:
     """Get the path to validation results for a specific mini_epoch and rank."""
     ext = StoreType(config.zarr_store).value  # validate extension
     base_path = get_path_run(config)
-    fname = f"validation_chkpt{mini_epoch:05d}_rank{config.rank:04d}.{ext}"
-    fname = config.get("output_name") or fname
+    default_name = f"validation_chkpt{mini_epoch:05d}_rank{config.rank:04d}.{ext}"
+    fname_template = config.get("output_name")
+    if fname_template is None:
+        fname = default_name
+    else:
+        try:
+            fname = fname_template.format(
+                epoch=mini_epoch,
+                step=step,
+                rank=config.rank,
+            )
+        except (IndexError, KeyError, ValueError):
+            fname = default_name
 
     return base_path / fname
 
diff --git a/src/weathergen/utils/validation_io.py b/src/weathergen/utils/validation_io.py
index 6f989b2bb2..9216ef635e 100644
--- a/src/weathergen/utils/validation_io.py
+++ b/src/weathergen/utils/validation_io.py
@@ -187,7 +187,7 @@ def write_output(
         forecast_offset=forecast_offset,
     )
 
-    store_path = config.get_path_results(cf, mini_epoch)
+    store_path = config.get_path_results(cf, mini_epoch, batch_idx)
 
     with zarrio_writer(store_path) as zio:
         for subset in data.items():

From 42c4460ce390a69db018c2aac7245b5afaba5068 Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Sat, 12 Sep 2026 17:41:09 +0200
Subject: [PATCH 09/13] Use get_path_results as run results directory and
 simplify path helpers

---
 .../common/src/weathergen/common/config.py    | 27 ++++++++++---------
 .../weathergen/evaluate/io/wegen_reader.py    |  4 +--
 src/weathergen/train/trainer.py               |  4 +--
 src/weathergen/utils/train_logger.py          |  8 +++---
 4 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index fae3098b4d..b318a39c88 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -212,7 +212,6 @@ def load_run_config(
     run_id: str,
     mini_epoch: int | None,
     model_path: str | None,
-    *,
     private_config: Config | None = None,
 ) -> Config:
     """
@@ -706,11 +705,11 @@ def load_streams(streams_directory: Path) -> Config:
     return OmegaConf.create(streams)
 
 
-def _get_path_output(config: Config, key: str, folder: str, run_id: str) -> Path:
+def _get_path_output(path_config: Config, key: str, folder: str, run_id: str) -> Path:
     """Use an exact directory override, or the existing shared per-run location."""
-    if config.get(key) is not None:
-        return Path(config[key])
-    working_dir = config.get("path_shared_working_dir")
+    if path_config.get(key) is not None:
+        return Path(path_config[key])
+    working_dir = path_config.get("path_shared_working_dir")
     root = Path(working_dir) if working_dir is not None else _get_shared_wg_path()
     return root / folder / run_id
 
@@ -720,11 +719,6 @@ def get_path_logs(config: Config) -> Path:
     return _get_path_output(config, "path_logs", "logs", get_run_id_from_config(config))
 
 
-def get_path_run(config: Config) -> Path:
-    """Get the current runs results_path for storing run results and logs."""
-    return _get_path_output(config, "path_results", "results", get_run_id_from_config(config))
-
-
 def get_path_model(config: Config | None = None, run_id: str | None = None) -> Path:
     """Get full_model_path if set, otherwise the shared per-run checkpoint directory."""
     if config or run_id:
@@ -736,10 +730,17 @@ def get_path_model(config: Config | None = None, run_id: str | None = None) -> P
     return _get_path_output(private_config, "full_model_path", "models", run_id)
 
 
-def get_path_results(config: Config, mini_epoch: int, step: int | None = None) -> Path:
-    """Get the path to validation results for a specific mini_epoch and rank."""
+def get_path_results(
+    config: Config,
+    mini_epoch: int | None = None,
+    step: int | None = None,
+) -> Path:
+    """Get the path for run results. Returns the results directory when mini_epoch is None."""
+    base_path = _get_path_output(config, "path_results", "results", get_run_id_from_config(config))
+    if mini_epoch is None:
+        return base_path
+
     ext = StoreType(config.zarr_store).value  # validate extension
-    base_path = get_path_run(config)
     default_name = f"validation_chkpt{mini_epoch:05d}_rank{config.rank:04d}.{ext}"
     fname_template = config.get("output_name")
     if fname_template is None:
diff --git a/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py b/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py
index 8e4ef8ca4e..2d1d8cdf15 100644
--- a/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py
+++ b/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py
@@ -21,7 +21,7 @@
 
 # Local application / package
 from weathergen.common.config import (
-    get_path_run,
+    get_path_results,
     load_merge_configs,
     load_run_config,
 )
@@ -52,7 +52,7 @@ def __init__(self, eval_cfg: dict, run_id: str, private_paths: dict | None = Non
         self.inference_cfg = self.get_inference_config()
 
         if not self.results_base_dir:
-            self.results_base_dir = get_path_run(self.inference_cfg)
+            self.results_base_dir = get_path_results(self.inference_cfg)
             _logger.info(f"Results directory obtained from private config: {self.results_base_dir}")
         else:
             _logger.info(f"Results directory parsed: {self.results_base_dir}")
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 276da0bd67..13c60edc55 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -156,10 +156,10 @@ def init(self, cf: Config, devices):
 
         # create output directory
         if is_root():
-            config.get_path_run(cf).mkdir(exist_ok=True, parents=True)
+            config.get_path_results(cf).mkdir(exist_ok=True, parents=True)
             config.get_path_model(cf).mkdir(exist_ok=True, parents=True)
 
-        self.train_logger = TrainLogger(cf, config.get_path_run(self.cf))
+        self.train_logger = TrainLogger(cf, config.get_path_results(self.cf))
 
         # Initialize collapse monitor for SSL training
         collapse_config = cf.train_logging.get("collapse_monitoring", {})
diff --git a/src/weathergen/utils/train_logger.py b/src/weathergen/utils/train_logger.py
index d7501a1cd5..8d391e1bd5 100644
--- a/src/weathergen/utils/train_logger.py
+++ b/src/weathergen/utils/train_logger.py
@@ -58,9 +58,9 @@ def by_mode(self, s: str) -> pl.DataFrame:
 
 class TrainLogger:
     #######################################
-    def __init__(self, cf, path_run: Path) -> None:
+    def __init__(self, cf, path_results: Path) -> None:
         self.cf = cf
-        self.path_run = path_run
+        self.path_results = path_results
 
     def log_metrics(self, stage: Stage, metrics: dict[str, float], step: int | None = None) -> None:
         """
@@ -86,7 +86,7 @@ def log_metrics(self, stage: Stage, metrics: dict[str, float], step: int | None
         # but we can probably do better and rely for example on the logging module.
 
         metrics_path = get_train_metrics_path(
-            base_path=config.get_path_run(self.cf), run_id=self.cf.general.run_id
+            base_path=self.path_results, run_id=self.cf.general.run_id
         )
         with open(metrics_path, "ab") as f:
             s = json.dumps(clean_metrics) + "\n"
@@ -152,7 +152,7 @@ def read(
             )
         run_id = cf.general.run_id
 
-        result_dir_base = config.get_path_run(cf)
+        result_dir_base = config.get_path_results(cf)
 
         # define cols for training
         cols1 = [_weathergen_timestamp, "num_samples", "loss_avg_mean", "learning_rate"]

From 2401193befb937bca2cbe362cfa2a2ac0f7b6925 Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Sat, 12 Sep 2026 23:39:59 +0200
Subject: [PATCH 10/13] Clarify get_path_model argument name

---
 packages/common/src/weathergen/common/config.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index b318a39c88..929ca64b1d 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -719,14 +719,14 @@ def get_path_logs(config: Config) -> Path:
     return _get_path_output(config, "path_logs", "logs", get_run_id_from_config(config))
 
 
-def get_path_model(config: Config | None = None, run_id: str | None = None) -> Path:
+def get_path_model(path_config: Config | None = None, run_id: str | None = None) -> Path:
     """Get full_model_path if set, otherwise the shared per-run checkpoint directory."""
-    if config or run_id:
-        run_id = run_id if run_id else get_run_id_from_config(config)
+    if path_config or run_id:
+        run_id = run_id if run_id else get_run_id_from_config(path_config)
     else:
-        msg = f"Missing run_id and cannot infer it from config: {config}"
+        msg = f"Missing run_id and cannot infer it from config: {path_config}"
         raise ValueError(msg)
-    private_config = config if config is not None else _load_private_conf()
+    private_config = path_config if path_config is not None else _load_private_conf()
     return _get_path_output(private_config, "full_model_path", "models", run_id)
 
 

From 03e692b965fe670d13d4b0ceed753ce845b0b16f Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Sat, 12 Sep 2026 23:43:52 +0200
Subject: [PATCH 11/13] Rename get_path_results argument for clarity

---
 packages/common/src/weathergen/common/config.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 929ca64b1d..98e222da3a 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -731,18 +731,20 @@ def get_path_model(path_config: Config | None = None, run_id: str | None = None)
 
 
 def get_path_results(
-    config: Config,
+    run_config: Config,
     mini_epoch: int | None = None,
     step: int | None = None,
 ) -> Path:
     """Get the path for run results. Returns the results directory when mini_epoch is None."""
-    base_path = _get_path_output(config, "path_results", "results", get_run_id_from_config(config))
+    base_path = _get_path_output(
+        run_config, "path_results", "results", get_run_id_from_config(run_config)
+    )
     if mini_epoch is None:
         return base_path
 
-    ext = StoreType(config.zarr_store).value  # validate extension
-    default_name = f"validation_chkpt{mini_epoch:05d}_rank{config.rank:04d}.{ext}"
-    fname_template = config.get("output_name")
+    ext = StoreType(run_config.zarr_store).value  # validate extension
+    default_name = f"validation_chkpt{mini_epoch:05d}_rank{run_config.rank:04d}.{ext}"
+    fname_template = run_config.get("output_name")
     if fname_template is None:
         fname = default_name
     else:
@@ -750,7 +752,7 @@ def get_path_results(
             fname = fname_template.format(
                 epoch=mini_epoch,
                 step=step,
-                rank=config.rank,
+                rank=run_config.rank,
             )
         except (IndexError, KeyError, ValueError):
             fname = default_name

From 35b42284a3e7f3cbc6cdb99d66047402e76d1dca Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Sat, 12 Sep 2026 23:48:46 +0200
Subject: [PATCH 12/13] Align get_path_model arg name with model config usage

---
 packages/common/src/weathergen/common/config.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 98e222da3a..c1e3e0db02 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -719,14 +719,14 @@ def get_path_logs(config: Config) -> Path:
     return _get_path_output(config, "path_logs", "logs", get_run_id_from_config(config))
 
 
-def get_path_model(path_config: Config | None = None, run_id: str | None = None) -> Path:
+def get_path_model(model_config: Config | None = None, run_id: str | None = None) -> Path:
     """Get full_model_path if set, otherwise the shared per-run checkpoint directory."""
-    if path_config or run_id:
-        run_id = run_id if run_id else get_run_id_from_config(path_config)
+    if model_config or run_id:
+        run_id = run_id if run_id else get_run_id_from_config(model_config)
     else:
-        msg = f"Missing run_id and cannot infer it from config: {path_config}"
+        msg = f"Missing run_id and cannot infer it from config: {model_config}"
         raise ValueError(msg)
-    private_config = path_config if path_config is not None else _load_private_conf()
+    private_config = model_config if model_config is not None else _load_private_conf()
     return _get_path_output(private_config, "full_model_path", "models", run_id)
 
 

From d5a8d1a47a644a988f6ce57aae71fce8f14dbeaa Mon Sep 17 00:00:00 2001
From: evenmn <evenmn@fys.uio.no>
Date: Sat, 12 Sep 2026 23:51:18 +0200
Subject: [PATCH 13/13] Use model_config name in get_path_results

---
 packages/common/src/weathergen/common/config.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index c1e3e0db02..ccc543cb13 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -731,20 +731,20 @@ def get_path_model(model_config: Config | None = None, run_id: str | None = None
 
 
 def get_path_results(
-    run_config: Config,
+    model_config: Config,
     mini_epoch: int | None = None,
     step: int | None = None,
 ) -> Path:
     """Get the path for run results. Returns the results directory when mini_epoch is None."""
     base_path = _get_path_output(
-        run_config, "path_results", "results", get_run_id_from_config(run_config)
+        model_config, "path_results", "results", get_run_id_from_config(model_config)
     )
     if mini_epoch is None:
         return base_path
 
-    ext = StoreType(run_config.zarr_store).value  # validate extension
-    default_name = f"validation_chkpt{mini_epoch:05d}_rank{run_config.rank:04d}.{ext}"
-    fname_template = run_config.get("output_name")
+    ext = StoreType(model_config.zarr_store).value  # validate extension
+    default_name = f"validation_chkpt{mini_epoch:05d}_rank{model_config.rank:04d}.{ext}"
+    fname_template = model_config.get("output_name")
     if fname_template is None:
         fname = default_name
     else:
@@ -752,7 +752,7 @@ def get_path_results(
             fname = fname_template.format(
                 epoch=mini_epoch,
                 step=step,
-                rank=run_config.rank,
+                rank=model_config.rank,
             )
         except (IndexError, KeyError, ValueError):
             fname = default_name