diff --git a/config/default_config.yml b/config/default_config.yml index 53f1ffc58..1f51a398f 100644 --- a/config/default_config.yml +++ b/config/default_config.yml @@ -222,6 +222,10 @@ validation_config: normalized_samples: False, # output streams to write; default all streams: null, + # channels that will be written; default all + channels: { + "ERA5": "all" + } } # run validation before training starts (mainly for model development) diff --git a/src/weathergen/utils/validation_io.py b/src/weathergen/utils/validation_io.py index b54bf5c73..90bb568a1 100644 --- a/src/weathergen/utils/validation_io.py +++ b/src/weathergen/utils/validation_io.py @@ -20,6 +20,68 @@ _logger = logging.getLogger(__name__) +def _filter_output_channels( + filter_cfg, + stream_names: list[str], + target_channels: list[list[str]], + targets_all: list, + preds_all: list, +) -> None: + """Apply per-stream channel filtering in-place. + + Args: + filter_cfg: The ``filter_output_channels`` config mapping + ``{STREAM_NAME: [channels]}``. An empty list or ``None`` for a + stream means no filtering for that stream. + stream_names: Ordered list of stream names. + target_channels: Per-stream list of channel names (mutated in-place). + targets_all: Nested list ``[t_idx][stream_idx]`` of target arrays + (mutated in-place). + preds_all: Nested list ``[t_idx][stream_idx]`` of prediction arrays + (mutated in-place). + """ + if not filter_cfg: + return + + for stream_idx, stream_name in enumerate(stream_names): + write_vars = filter_cfg.get(stream_name) + if write_vars is None or (isinstance(write_vars, str) and write_vars.lower() == "all"): + continue + write_vars = [write_vars] if isinstance(write_vars, str) else list(write_vars) + if len(write_vars) == 0: + continue + + all_channels = target_channels[stream_idx] + write_vars_set = set(write_vars) + keep_idxs = [i for i, ch in enumerate(all_channels) if ch in write_vars_set] + + missing = write_vars_set - set(all_channels) + if missing: + _logger.warning( + f"filter_output_channels for stream {stream_name} " + f"contains unknown channels, which will be skipped: {missing}" + ) + if not keep_idxs: + _logger.warning( + f"filter_output_channels for stream {stream_name} matched no channels; " + f"skipping filter." + ) + continue + if len(keep_idxs) == len(all_channels): + continue + keep_names = [all_channels[i] for i in keep_idxs] + removed_names = [ch for ch in all_channels if ch not in keep_names] + _logger.debug( + f"Filtering output channels for stream {stream_name}: " + f"{len(all_channels)} -> {len(keep_idxs)} channels; " + f"kept: {keep_names}; removed: {removed_names}" + ) + target_channels[stream_idx] = [all_channels[i] for i in keep_idxs] + for t_idx in range(len(targets_all)): + targets_all[t_idx][stream_idx] = targets_all[t_idx][stream_idx][:, keep_idxs] + preds_all[t_idx][stream_idx] = preds_all[t_idx][stream_idx][:, :, keep_idxs] + + def write_output( cf, val_cfg, batch_size, mini_epoch, batch_idx, dn_data, batch, model_output, target_aux_out ): @@ -119,8 +181,6 @@ def write_output( for sample in batch.get_source_samples().get_samples() ] - # more prep work - # output stream names to be written, use specified ones or all if nothing specified stream_names = list(cf.streams.keys()) stream_infos = list(cf.streams.values()) @@ -135,7 +195,10 @@ def write_output( target_channels: list[list[str]] = [list(stream.val_target_channels) for stream in stream_infos] source_channels: list[list[str]] = [list(stream.val_source_channels) for stream in stream_infos] - geoinfo_channels = [[] for _ in stream_infos] # TODO obtain channels + filter_cfg = val_cfg.get("output", {}).get("channels", None) + _filter_output_channels(filter_cfg, stream_names, target_channels, targets_all, preds_all) + + geoinfo_channels = [[] for _ in stream_names] # TODO obtain channels # calculate global sample indices for this batch by offsetting by sample_start sample_start = batch_idx * batch_size