Sophiex/dev ssl diffusion/make write target optional - #2765
Sophiex/dev ssl diffusion/make write target optional#2765sophie-xhonneux wants to merge 13 commits into
Conversation
clessig
left a comment
There was a problem hiding this comment.
I wouldn't merge right now. There are quite a few changes in places that also touch training. It would have to be tested properly with various configs, which is too big a distraction right now.
|
|
||
| # if False, no target datasets are built/written (skip_target_values inference); | ||
| # self.targets then only provides row counts via targets_lens | ||
| write_targets: bool = True |
There was a problem hiding this comment.
Can we put this under validation_config.output.write_targets
| return rdata | ||
|
|
||
| def get_target(self, idx: TIndex) -> ReaderData: | ||
| def get_target(self, idx: TIndex, coords_geoinfos_only: bool = False) -> ReaderData: |
There was a problem hiding this comment.
The flag is a bit misleading since coords, geoinfos, and time are provided. Maybe call is no_data or similar?
| return rdata | ||
|
|
||
| @override | ||
| def _get_coords_geoinfos_only(self, idx: int, channels_idx: list[int]) -> ReaderData: |
There was a problem hiding this comment.
Why do we need the override here? Wouldn't the implementation in the base class work?
There was a problem hiding this comment.
I think it has to do with the fact that different readers have different capability to only get the geoinfos, e.g. in anemoi it isn't possible, you can only drop them at the end of the stack frame
|
|
||
| def collect_datasources(stream_datasets: list, idx: int, type: str, rng) -> IOReaderData: | ||
| def collect_datasources( | ||
| stream_datasets: list, idx: int, type: str, rng, coords_geoinfos_only: bool = False |
There was a problem hiding this comment.
Same about arg name as above
|
|
||
| # skip_target_values: inference-only; do not read/store target values for physically | ||
| # reconstructed streams, only coords/geoinfos/datetimes (data gets zero width) | ||
| self.skip_target_values = bool(mode_cfg.get("skip_target_values", False)) |
There was a problem hiding this comment.
Why do you have a flag in io.py and a config argument? As above, the config argument should be validation_config.output.write_targets
…r a run id and for santis
Description
Don't write the target to the zarrs and drop the target channels as soon as possible. From my and Claude's understanding of anemoi zarr reading, we still need the bulk dataread to get the geoinfos.
Note this is only useful once we have the evaluation package able to read the target zarrs from predictions alone.
Issue Number
Is this PR a draft? Mark it as draft.
Checklist before asking for review