Skip to content

Sophiex/dev ssl diffusion/make write target optional - #2765

Open
sophie-xhonneux wants to merge 13 commits into
develop-ssl-diffusion-v1from
sophiex/dev-ssl-diffusion/make-write-target-optional
Open

Sophiex/dev ssl diffusion/make write target optional#2765
sophie-xhonneux wants to merge 13 commits into
develop-ssl-diffusion-v1from
sophiex/dev-ssl-diffusion/make-write-target-optional

Conversation

@sophie-xhonneux

Copy link
Copy Markdown
Contributor

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

  • tested: train+ validation space JEPA
  • tested: train+ validation temporal JEPA
  • tested: train+ validation diffusion
  • tested: train+ validation forecasting

@github-actions github-actions Bot added data Anything related to the datasets used in the project eval anything related to the model evaluation pipeline infra Issues related to infrastructure model Related to model training or definition (not generic infra) labels Aug 14, 2026
@sophie-xhonneux
sophie-xhonneux changed the base branch from develop to develop-ssl-diffusion-v1 August 14, 2026 14:24

@clessig clessig left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need the override here? Wouldn't the implementation in the base class work?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same about arg name as above

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sure, idm


# 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))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

data Anything related to the datasets used in the project eval anything related to the model evaluation pipeline infra Issues related to infrastructure model Related to model training or definition (not generic infra)

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

3 participants