From e1b56a9377a9b0ed370e52335b4b7e1d3b829892 Mon Sep 17 00:00:00 2001 From: melidonis1 Date: Tue, 18 Aug 2026 10:38:25 +0200 Subject: [PATCH 1/3] add additional column case in plot_training --- src/weathergen/utils/plot_training.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/weathergen/utils/plot_training.py b/src/weathergen/utils/plot_training.py index 5e5f9478e..de3910f2b 100644 --- a/src/weathergen/utils/plot_training.py +++ b/src/weathergen/utils/plot_training.py @@ -469,6 +469,15 @@ def plot_loss_per_stream( ): data_cols += [col] title_col = col if title_col is None else title_col + elif len(col_split) == 6: + if ( + col_split[1].lower() == stream_name.lower() + and col_split[2].lower() == err.lower() + and col_split[3] + "." + col_split[4] == channel + and int(col_split[5]) in forecast_steps + ): + data_cols += [col] + title_col = col if title_col is None else title_col for col in data_cols: x_vals = np.array(run_data_mode[x_col]) From 71b7561f851db7290018311995daba6c851d125c Mon Sep 17 00:00:00 2001 From: melidonis1 Date: Mon, 24 Aug 2026 10:09:16 +0200 Subject: [PATCH 2/3] Introduce functions for adapting column names and channel names --- src/weathergen/utils/plot_training.py | 33 ++++++++------ src/weathergen/utils/train_logger.py | 65 +++++++++++++++++++++++---- 2 files changed, 76 insertions(+), 22 deletions(-) diff --git a/src/weathergen/utils/plot_training.py b/src/weathergen/utils/plot_training.py index de3910f2b..1b5154a1d 100644 --- a/src/weathergen/utils/plot_training.py +++ b/src/weathergen/utils/plot_training.py @@ -21,7 +21,7 @@ import weathergen.common.config as config from weathergen.train.utils import TRAIN -from weathergen.utils.train_logger import Metrics, TrainLogger +from weathergen.utils.train_logger import Metrics, TrainLogger, sanitize_channel_for_matching _logger = logging.getLogger(__name__) @@ -420,8 +420,17 @@ def plot_loss_per_stream( prop_cycle = plt.rcParams["axes.prop_cycle"] colors = prop_cycle.by_key()["color"] + ["r", "g", "b", "k", "m", "y"] + if channels is not None: + channels_sanitized = [sanitize_channel_for_matching(c) for c in channels] + else: + channels_sanitized = channels + for err in errs: - for channel in channels: + channels_to_use = channels if channels else [None] + channels_sanitized_to_use = channels_sanitized if channels_sanitized else [None] + for channel, channel_sanitized in zip( + channels_to_use, channels_sanitized_to_use, strict=True + ): for stream_name in stream_names: _fig = plt.figure(figsize=(10, 7), dpi=PLOT_DPI_VALUE) @@ -456,7 +465,7 @@ def plot_loss_per_stream( if ( col_split[1].lower() == stream_name.lower() and col_split[2].lower() == err.lower() - and col_split[3] == channel + and col_split[3] == channel_sanitized ): data_cols += [col] title_col = col if title_col is None else title_col @@ -464,20 +473,11 @@ def plot_loss_per_stream( if ( col_split[1].lower() == stream_name.lower() and col_split[2].lower() == err.lower() - and col_split[3] == channel + and col_split[3] == channel_sanitized and int(col_split[4]) in forecast_steps ): data_cols += [col] title_col = col if title_col is None else title_col - elif len(col_split) == 6: - if ( - col_split[1].lower() == stream_name.lower() - and col_split[2].lower() == err.lower() - and col_split[3] + "." + col_split[4] == channel - and int(col_split[5]) in forecast_steps - ): - data_cols += [col] - title_col = col if title_col is None else title_col for col in data_cols: x_vals = np.array(run_data_mode[x_col]) @@ -612,6 +612,11 @@ def plot_loss_per_run( if errs is None: errs = ["mse"] + if channels is not None: + channels_sanitized = [sanitize_channel_for_matching(c) for c in channels] + else: + channels_sanitized = channels + plot_dir = Path(plot_dir) modes = [modes] if type(modes) is not list else modes @@ -640,7 +645,7 @@ def plot_loss_per_run( if ( len(col_split) >= 4 and col_split[2].lower() == err.lower() - and col_split[3] in channels + and col_split[3] in channels_sanitized ): data_cols += [col] diff --git a/src/weathergen/utils/train_logger.py b/src/weathergen/utils/train_logger.py index d7501a1cd..c0dc7eb6a 100644 --- a/src/weathergen/utils/train_logger.py +++ b/src/weathergen/utils/train_logger.py @@ -19,7 +19,6 @@ import numpy as np import polars as pl import torch - import weathergen.common.config as config # from weathergen.train.trainer import cfg_keys_to_filter @@ -62,7 +61,9 @@ def __init__(self, cf, path_run: Path) -> None: self.cf = cf self.path_run = path_run - def log_metrics(self, stage: Stage, metrics: dict[str, float], step: int | None = None) -> None: + def log_metrics( + self, stage: Stage, metrics: dict[str, float], step: int | None = None + ) -> None: """ Log metrics to a file. For now, just scalar values are expected. There is no check. @@ -145,7 +146,9 @@ def read( # Load config from given model_path if provided, otherwise use path from private config if model_path: - cf = config.load_run_config(run_id=run_id, mini_epoch=mini_epoch, model_path=model_path) + cf = config.load_run_config( + run_id=run_id, mini_epoch=mini_epoch, model_path=model_path + ) else: cf = config.load_merge_configs( private_home=None, from_run_id=run_id, mini_epoch=mini_epoch @@ -158,13 +161,17 @@ def read( cols1 = [_weathergen_timestamp, "num_samples", "loss_avg_mean", "learning_rate"] cols1_patterns = ["loss_avg"] + cols_patterns - metrics_train = read_metrics(cf, run_id, "train", cols1, cols1_patterns, result_dir_base) + metrics_train = read_metrics( + cf, run_id, "train", cols1, cols1_patterns, result_dir_base + ) # define cols for validation cols2 = [_weathergen_timestamp, "num_samples"] cols2_patterns = ["loss_avg"] + cols_patterns - metrics_val = read_metrics(cf, run_id, "val", cols2, cols2_patterns, result_dir_base) + metrics_val = read_metrics( + cf, run_id, "val", cols2, cols2_patterns, result_dir_base + ) return Metrics(run_id, "train", metrics_train, metrics_val, None) @@ -194,6 +201,10 @@ def read_metrics( # TODO: this should be a config option df = read_metrics_file(metrics_path) + # Sanitize LossPhysical column names to replace dots in channel names with underscores + new_columns = {col: sanitize_loss_physical_column(col) for col in df.columns} + df = df.rename(new_columns) + if cols_patterns is not None: for col_pattern in cols_patterns: cols += [col for col in df.columns if col_pattern in col] @@ -216,10 +227,14 @@ def clean_df(df, columns: list[str] | None): # Convert timestamp column to date df = df.with_columns( - pl.from_epoch(df[_weathergen_timestamp], time_unit="ms").alias(_weathergen_timestamp) + pl.from_epoch(df[_weathergen_timestamp], time_unit="ms").alias( + _weathergen_timestamp + ) ) df = df.with_columns( - (df[_weathergen_timestamp] - df[_weathergen_timestamp].min()).alias(_weathergen_reltime) + (df[_weathergen_timestamp] - df[_weathergen_timestamp].min()).alias( + _weathergen_reltime + ) ) if columns: @@ -269,6 +284,38 @@ def _key_stddev(st_name: str) -> str: return f"stream.{st_name}.stddev_avg" +def sanitize_loss_physical_column(col: str) -> str: + """Replace dots in channel names with underscores for LossPhysical columns. + + Adapts column names like: + LossPhysical.EERIE_OCEAN.mse.avg_thetao_2.5.0 + to: + LossPhysical.EERIE_OCEAN.mse.avg_thetao_2_5.0 + + The format is: LossPhysical.{stream}.{error}.{channel}[.{forecast_step}] + """ + if not col.startswith("LossPhysical."): + return col + parts = col.split(".") + if len(parts) < 4: + return col + + if len(parts) > 4 and parts[-1].isdigit(): + channel_parts = parts[3:-1] + forecast_step = parts[-1] + sanitized_channel = "_".join(channel_parts) + return f"LossPhysical.{parts[1]}.{parts[2]}.{sanitized_channel}.{forecast_step}" + else: + channel_parts = parts[3:] + sanitized_channel = "_".join(channel_parts) + return f"LossPhysical.{parts[1]}.{parts[2]}.{sanitized_channel}" + + +def sanitize_channel_for_matching(channel: str) -> str: + """Replace dots with underscores for matching against sanitized column names.""" + return channel.replace(".", "_") + + def prepare_losses_for_logging( loss_hist: list, losses_unweighted_hist: list[dict], @@ -292,7 +339,9 @@ def prepare_losses_for_logging( for d in losses_unweighted_hist: for key, value in flatten_dict(d).items(): - value = torch.tensor(value, device="cuda") if type(value) is float else value + value = ( + torch.tensor(value, device="cuda") if type(value) is float else value + ) losses_all[key].append(ddp_average(value).item()) for d in stddev_unweighted_hist: From ae96c2f82e39657ec06954b5d123128b2388b43b Mon Sep 17 00:00:00 2001 From: melidonis1 Date: Mon, 24 Aug 2026 10:10:12 +0200 Subject: [PATCH 3/3] Linting --- src/weathergen/utils/train_logger.py | 29 ++++++++-------------------- 1 file changed, 8 insertions(+), 21 deletions(-) diff --git a/src/weathergen/utils/train_logger.py b/src/weathergen/utils/train_logger.py index c0dc7eb6a..4d958f412 100644 --- a/src/weathergen/utils/train_logger.py +++ b/src/weathergen/utils/train_logger.py @@ -19,6 +19,7 @@ import numpy as np import polars as pl import torch + import weathergen.common.config as config # from weathergen.train.trainer import cfg_keys_to_filter @@ -61,9 +62,7 @@ def __init__(self, cf, path_run: Path) -> None: self.cf = cf self.path_run = path_run - def log_metrics( - self, stage: Stage, metrics: dict[str, float], step: int | None = None - ) -> None: + def log_metrics(self, stage: Stage, metrics: dict[str, float], step: int | None = None) -> None: """ Log metrics to a file. For now, just scalar values are expected. There is no check. @@ -146,9 +145,7 @@ def read( # Load config from given model_path if provided, otherwise use path from private config if model_path: - cf = config.load_run_config( - run_id=run_id, mini_epoch=mini_epoch, model_path=model_path - ) + cf = config.load_run_config(run_id=run_id, mini_epoch=mini_epoch, model_path=model_path) else: cf = config.load_merge_configs( private_home=None, from_run_id=run_id, mini_epoch=mini_epoch @@ -161,17 +158,13 @@ def read( cols1 = [_weathergen_timestamp, "num_samples", "loss_avg_mean", "learning_rate"] cols1_patterns = ["loss_avg"] + cols_patterns - metrics_train = read_metrics( - cf, run_id, "train", cols1, cols1_patterns, result_dir_base - ) + metrics_train = read_metrics(cf, run_id, "train", cols1, cols1_patterns, result_dir_base) # define cols for validation cols2 = [_weathergen_timestamp, "num_samples"] cols2_patterns = ["loss_avg"] + cols_patterns - metrics_val = read_metrics( - cf, run_id, "val", cols2, cols2_patterns, result_dir_base - ) + metrics_val = read_metrics(cf, run_id, "val", cols2, cols2_patterns, result_dir_base) return Metrics(run_id, "train", metrics_train, metrics_val, None) @@ -227,14 +220,10 @@ def clean_df(df, columns: list[str] | None): # Convert timestamp column to date df = df.with_columns( - pl.from_epoch(df[_weathergen_timestamp], time_unit="ms").alias( - _weathergen_timestamp - ) + pl.from_epoch(df[_weathergen_timestamp], time_unit="ms").alias(_weathergen_timestamp) ) df = df.with_columns( - (df[_weathergen_timestamp] - df[_weathergen_timestamp].min()).alias( - _weathergen_reltime - ) + (df[_weathergen_timestamp] - df[_weathergen_timestamp].min()).alias(_weathergen_reltime) ) if columns: @@ -339,9 +328,7 @@ def prepare_losses_for_logging( for d in losses_unweighted_hist: for key, value in flatten_dict(d).items(): - value = ( - torch.tensor(value, device="cuda") if type(value) is float else value - ) + value = torch.tensor(value, device="cuda") if type(value) is float else value losses_all[key].append(ddp_average(value).item()) for d in stddev_unweighted_hist: