Coupled no-target inference not possible with time-invariant ocean forcing dataset

[jpdunc23]

In recent version of SamudrACE we use time-invariant forcings of the ocean such as the bottom geothermal heat flux, but we currently have no way to load a dataset which contains only time-invariant forcings. The current workaround is to use an ocean forcing dataset that includes a non-time-invariant variable, e.g. just using the original training dataset.

[oliverwm1]

I think the right way to resolve this is to implement the serialization of time invariant forcing fields with checkpoints.

[mcgibbon]

Should we make the Step responsible for this, and have it grab invariant forcings through DatasetInfo at initialization time?

While I think that may be the best way to implement serialization with checkpoints, it wouldn't cover multi-dataset models that have a few examples of these forcings - for those, we may still want to support loading a dataset that contains only time-invariant forcings.

When loading time-invarant-only forcing data, the configuration itself would need to specify the time axis, since the forcing data is valid for any time axis.

[mcgibbon]

I think the right way to resolve this is to implement the serialization of time invariant forcing fields with checkpoints.

To implement, I would propose something like:

• Include static data fields as a dict[str, torch.Tensor] value in DatasetInfo.
• Use a dataset_info, statics = dataset_info.retrieve_static(list_of_names) helper function to get static variables, which raises on missing statics, or drops non-retrieved static variables in the normal case (so that when we later serialize the dataset info it only has statics we actually use).
• Add a static_input_names: List[str] field to SingleModuleStepConfig, which it uses to retrieve and store static inputs.

I would suggest leaving overriding of these static inputs to later if we need it, it's definitely possible to implement but it will require dedicated API, which I suggest we shouldn't do if we want to discourage it anyways and have no specific use case in mind. If we have a use case in mind we could implement this as a follow-on task.

[oliverwm1]

Add a static_input_names: List[str] field to SingleModuleStepConfig, which it uses to retrieve and store static inputs.

Does this need to be user configurable? Could determine if it is static from the training dataset. We already have custom handling in XarrayDataset for variables that don't have a time dimension, i.e. static inputs.

[mcgibbon]

Add a static_input_names: List[str] field to SingleModuleStepConfig, which it uses to retrieve and store static inputs.

Does this need to be user configurable? Could determine if it is static from the training dataset. We already have custom handling in XarrayDataset for variables that don't have a time dimension, i.e. static inputs.

Ah true, it does need to be user-configurable in some sense (we shouldn't just grab/store every static) but arguably in_names and its presence in the dataset_info could be sufficient.

It just depends whether we want this to be automatic or not. For example if it's automatic the statics will get silently dropped when multiple datasets contain different values for the same static (outside of some almost_close epsilon we choose), and we wouldn't be able to recover our previous training behavior (producing a model that requires these in the inputs) if we wanted it for some unforeseen reason.

[mcgibbon]

@jpdunc23 let me know what you think of this solution before I go to implement it, and whether it unblocks you in the short-term or whether you also need a user-configurable dataset for static-only entries in the short-term.