diff --git a/src/weathergen/utils/plot_training.py b/src/weathergen/utils/plot_training.py index 5e5f9478e..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,7 +473,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 and int(col_split[4]) in forecast_steps ): data_cols += [col] @@ -603,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 @@ -631,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..4d958f412 100644 --- a/src/weathergen/utils/train_logger.py +++ b/src/weathergen/utils/train_logger.py @@ -194,6 +194,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] @@ -269,6 +273,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],