Clean up config and add extra validation; refactor dropout params by dfulu · Pull Request #450 · openclimatefix/ocf-data-sampler

dfulu · 2026-10-19T11:39:48Z

Pull Request

Description

This PR

Cleans up and refactors the config
Adds extra validation to the config
Refactors how we apply dropout
- Rename dropout to delay for NWPs to reflect different handling
- This includes how we configure it in the config now a dictionary instead of linked lists
- Also includes separating out the functionality to pick a dropout value from applying the dropout

Checklist:

My code follows OCF's coding style guidelines
I have performed a self-review of my own code
I have made corresponding changes to the documentation
I have added tests that prove my fix is effective or that my feature works
I have checked my code and corrected any misspellings

dfulu · 2026-10-19T11:40:31Z

 from numpy.typing import NDArray

+# A single path, a local glob, or an explicit sequence of Zarr store paths
+ZarrSource: TypeAlias = str | list[str] | tuple[str, ...]


Moved this from xr_tensorstore since we also use it in the config. It feels like it fits better in this types module

dfulu · 2026-10-19T11:44:00Z

    )

-    @field_validator("dropout_timedeltas_minutes")
-    def dropout_timedeltas_minutes_negative(cls, v: list[int]) -> list[int]:


The _MinuteProbabilities class does its own validation

dfulu · 2026-10-19T11:44:26Z

-        description="The value used to fill in dropped out data or any missing values. This is "
-        "applied after normalisation, so it is in normalised units."
+    fill_value: float | None = Field(
+        default=None,


We default to no infilling. I want to make infilling very explicit so its never glossed over. Also this isn't dropout specific, it fills any NaN values at all hence name change

Agree on making infilling explicit. Could we add a check that a finished sample has no NaNs left, so a missing fill value fails loudly rather than showing up later as a NaN loss? Also worth a line in the migration notes maybe.

Scanning through each array for NaNs can have a significant performance hit. I mean to refactor how we fill NaNs since right now the nan_to_num function we're using has a weird side effect of replacing infinites as well. So I'd rather leave it to be considered then.

I'm also not maintaining migration notes for the dev branch, and a lot more than this has changed. If we want them I'll have to build them up

dfulu · 2026-10-19T11:46:24Z


-class Base(BaseModel):
-    """Pydantic Base model where no extras can be added."""
+def _validate_grid_resolution_minutes(resolution: int) -> int:


This is an extra validator on the spacing of our temporal resolutions. We generally aren't set up to deal with time resolutions of say 11 minutes anywhere. We cn deal with the simple cases which are factors of an hour or multiple of an hour.

dfulu · 2026-10-19T11:46:58Z

+    dropout_minute_probabilities: dict[int, float] | None,
+    interval_end_minutes: int,
+) -> None:
+    """Reject choices that would leave every observation in the window untouched."""


This is extra validation. We should know if we've configured a dropout which does nothing

dfulu · 2026-10-19T11:47:29Z

+]
+
+_Embedding = tuple[
+    Annotated[str, Field(pattern=r"^[0-9]+[hy]$")],


This replaces a bunch of validation code on the t0 embeddings

dfulu · 2026-10-19T11:50:32Z

-        default=0.0,
-        description="Either a float(Chance of dropout being applied to each sample) or a list of "
-        "floats (probability that dropout of the corresponding timedelta is applied)",
+    dropout_minute_probabilities: _MinuteProbabilities | None = Field(


I've refactored the dropout configuration. Instead of being a linked-list of dropout values and probabilities it is now a dict which is configured like dict[dropout_minutes: int, probability: float]. I think this is nicer and a more secure way of configuring this than linked lists which require extra validation (like checking lengths) that this data structure does not

dfulu · 2026-10-19T11:51:32Z

    channels: list[str] = Field(
        ...,
-        description="the satellite channels that are used",
+        min_length=1,


Must include at least 1 channel

dfulu · 2026-10-19T11:53:15Z

+        description="Channels to include, in output order",
+    )
+
+    delay_minute_probabilities: _MinuteProbabilities | None = Field(


For NWP we rename dropout to delay since that better reflects the data outputs. Dropout replaces the values with NaNs which are then infilled. Delay just uses an older init-time

dfulu · 2026-10-19T11:54:19Z

    @field_validator("provider")
-    def validate_provider(cls, v: str) -> str:
-        """Validator for 'provider'."""
-        provider = v.lower()


I don't think we should be case insensitive. This is easy for us to just get right

dfulu · 2026-10-19T12:03:37Z

    )

    @field_validator("embeddings")
-    def validate_embeddings(cls, embeddings: _embedding_type) -> _embedding_type:


This validation plus more is now down in the _Embedding class and in in validate_embeddings() below

dfulu · 2026-10-19T12:04:17Z

These are just changes for the new dropout data structure

dfulu · 2026-10-19T12:06:04Z

-    if config.generation is not None:
-        for key, window_config in (
-            ("generation_input", config.generation.input),
-            ("generation_target", config.generation.target),


There should only ever be NaNs in the input generation, and those only come from dropout

dfulu · 2026-10-19T12:06:58Z



-def fill_nans(da: TArray, fill_value: float) -> TArray:
+def fill_nans(da: TArray, fill_value: float) -> None:


This function changes the inputs in-place. I want to make that more explicit

dfulu · 2026-10-19T12:09:08Z

 from ocf_data_sampler.common.types import TArray


+def choose_delay(delay_minute_probabilities: Mapping[int, float]) -> np.timedelta64 | None:


I've refactored the dropout codepaths. Now we choose a delay and apply it in different steps. This helps split out the responsibilities a bit more and we can use choose_delay() both for dropout and for choosing the NWP delay

dfulu · 2026-10-19T12:10:56Z

            f"Dropout can only be applied to DataArrays with numpy data. Got: {type(da.data)}"
        )

-    if isinstance(dropout_frac, float | int):


This is all validated at the config level. I think that's a better boundary

Agreed. Should max_delay <= max_staleness move into NWP config too, then?

I'm not sure, it would move that error message earlier, BUT if max_staleness is set to None then we calculate a max_staleness from the max forecast step minus the interval end minutes. We need to open the zarr to do that and we shouldn't be touching data when validating the config.

Therefore, we can't get rid of the check inside select_time_slice_nwp. We could add a additional validator in the config for the case where max_staleness is not None, but its not going to catch many errors much earlier. I think leaving it as None is the dominant usage pattern

dfulu · 2026-10-19T12:11:49Z

    return da.isel(init_time_utc=selected_init_time_index, step=selected_step_indices)
-
-
-def _get_nwp_dropout_available_time(


After the refactor this isn't needed. We get to reuse choose_delay() instead and have all the validation done at the config level

dfulu · 2026-10-19T13:21:07Z

    assert (expected_init_time == da_slice["init_time_utc"].values)
-
-
-def test_select_time_slice_nwp_with_weighted_dropout_list(da_nwp_like):


This validation is done at the config level now. The function expects clean inputs

dfulu · 2026-10-19T13:40:45Z

+    """NWP source configurations, accessed by name: config.nwp["ukv"]."""

-    root: dict[str, NWP]
+    root: dict[str, NWP] = Field(..., min_length=1)


Make sure NWP can't be configured as an empty dictionary. We set it to None if we don't want any NWP

dfulu · 2026-10-19T13:43:14Z

            t0,
-            dropout_timedeltas=minutes(config.satellite.dropout_timedeltas_minutes),
-            dropout_frac=config.satellite.dropout_fraction,
+            dropout_delay=choose_delay(sat_conf.dropout_minute_probabilities),


The dropout time is chosen and passed in. I think this is better separation of concerns

dfulu · 2026-10-20T08:48:08Z

        return self


-class MultiNWP(RootModel):


Turns out this class is completely unnecessary with the current version of the library. It just mimics the behaviour of a dictionary. Pydantic does actually allow you to use a dictionary as part of the config, so I've switched to that

Correct me if I'm wrong, but #451 will need a rebase onto this right?

Yeh in general I'm maintaining that all these PRs will need to rebase on the PRs that came before them. This is part_22, the PR you mention is part_23 (as per the branch names)

Clean up config and add extra validation; refactor dropout params

d8c83b6

dfulu commented Oct 19, 2026

View reviewed changes

simplify tests

1b9e26b

dfulu commented Oct 19, 2026

View reviewed changes

dfulu marked this pull request as ready for review October 19, 2026 13:37

dfulu commented Oct 19, 2026

View reviewed changes

dfulu added 2 commits October 19, 2026 13:55

Make dropout explicitly return None

47b363e

Simplify nested NWP config

f38d557

dfulu commented Oct 20, 2026

View reviewed changes

felix-e-h-p self-requested a review October 20, 2026 13:59

felix-e-h-p approved these changes Oct 21, 2026

View reviewed changes

dfulu merged commit 2f12e1c into dev_feb2026_speedups Oct 21, 2026
4 checks passed

dfulu deleted the part_22 branch October 21, 2026 16:58

dfulu mentioned this pull request Oct 21, 2026

dev - next major release - DS v2 #400

Draft



		def fill_nans(da: TArray, fill_value: float) -> TArray:
		def fill_nans(da: TArray, fill_value: float) -> None:

		from ocf_data_sampler.common.types import TArray


		def choose_delay(delay_minute_probabilities: Mapping[int, float]) -> np.timedelta64 \| None:

		return da.isel(init_time_utc=selected_init_time_index, step=selected_step_indices)


		def _get_nwp_dropout_available_time(

		assert (expected_init_time == da_slice["init_time_utc"].values)


		def test_select_time_slice_nwp_with_weighted_dropout_list(da_nwp_like):

Conversation

dfulu commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request

Description

Checklist:

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dfulu Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dfulu Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dfulu Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dfulu Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dfulu Oct 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dfulu Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

dfulu commented Oct 19, 2026 •

edited

Loading

dfulu Oct 19, 2026 •

edited

Loading

dfulu Oct 21, 2026 •

edited

Loading

dfulu Oct 19, 2026 •

edited

Loading

dfulu Oct 19, 2026 •

edited

Loading

dfulu Oct 20, 2026 •

edited

Loading

dfulu Oct 21, 2026 •

edited

Loading