Write selected channels during inference

[sbAsma]

Description (Edited)

Changes in config/default_config.yml:

Added a new filed in output to filter channels.

Changes in packages/common/src/weathergen/common/io.py:

Added a field in the zarr output for channels being written

Changes in src/weathergen/utils/validation_io.py:

Now trims targets_all and preds_all tensors to a per-stream allowlist or denylist before writing to the output zarr, via a new output_channels list derived from a deep copy of target_channels.

Changes in packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py:

Added get_channels() that reads channel names directly from the Zarr output metadata

To run with:

uv run --offline inference --from-run-id run_id --options test_config.samples_per_mini_epoch=1 validation_config.output.filter_output_channels.allow=True  validation_config.output.filter_output_channels.ERA5=["2t"]

The output zarr file has the tree:

>>> print(root.tree())

/
└── 0
    └── ERA5
        ├── 0
        │   └── source
        │       ├── coords (40320, 2) float32
        │       ├── data (40320, 71) float32
        │       ├── geoinfo (40320, 0) float32
        │       └── times (40320,) datetime64
        ├── 1
        │   ├── prediction
        │   │   ├── coords (40320, 2) float32
        │   │   ├── data (40320, 1, 1) float32
        │   │   ├── geoinfo (40320, 0) float32
        │   │   └── times (40320,) datetime64
        │   └── target
        │       ├── coords (40320, 2) float32
        │       ├── data (40320, 1) float32
        │       ├── geoinfo (40320, 0) float32
        │       └── times (40320,) datetime64
        └── 2
            ├── prediction
            │   ├── coords (40320, 2) float32
            │   ├── data (40320, 1, 1) float32
            │   ├── geoinfo (40320, 0) float32
            │   └── times (40320,) datetime64
            └── target
                ├── coords (40320, 2) float32
                ├── data (40320, 1) float32
                ├── geoinfo (40320, 0) float32
                └── times (40320,) datetime64

>>> print(dict(root["0"]["ERA5"]["1"]["target"].attrs))
{'channels': ['2t'], 'geoinfo_channels': [], 'source_interval': {'start': '2023-10-01T00:00:00.000000000', 'end': '2023-10-01T06:00:00.000000000'}}

Issue Number

Closes #1705

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

[iluise]

Thanks for the heads up! Minor suggestion to avoid retrieving the channels twice

[iluise]

I would just move it at line 152 (so before self.get_inference_stream_attr(stream, "val_target_channels")):

write_output = self.get_inference_stream_attr(stream, "write_output")
if write_output is not None:
            all_channels = [ch for ch in all_channels if ch in write_output]
else:
     all_channels = self.get_inference_stream_attr(stream, "val_target_channels")
_logger.debug(f"Channels found in config: {all_channels}")

[SavvasMel]

Why do we need to change the reader in general here? all_channels should have anyway whatever is written in the output from write_output

[SavvasMel]

So all_channels should be 2t in our example here.

[sbAsma]

@SavvasMel source is kept as it is. Only target and prediction channels are being filtered.

[SavvasMel]

We need to be careful this to be removed before merge. Or give an empty list.

[sbAsma]

I am adding another param in default_config that controls whether channels are being filtered or not.

[SavvasMel]

It seems that this respects backward compatibility but please double check

[SavvasMel]

I mean what happens if the stream does not have "write_output"?

[sbAsma]

I am working on checking the backward compatibility. To answer your question, get_inference_stream_attr returns an empty list if it doesn't find write_output

[sbAsma]

Backward compatibility checks

[sbAsma]

Hi @SavvasMel and @iluise

I had another idea how to implement this issue, and I would really appreciate your input. Among the changes, I decided to remove the filtering of the channels from the stream config, and add it in the default config instead so that it can be overwritten with --options. I also added a change requested by Christian, where the output zarr file would have metadata on the channels being written.