Skip to content

Data Loaders

N_CALENDAR_DAYS = 366 module-attribute

CombinedDataset

Bases: Dataset

n_forecast_steps = n_forecast_steps instance-attribute

n_history_steps = n_history_steps instance-attribute

climatology = climatology instance-attribute

target = next(ds for ds in datasets if ds.name == target_group_name).subset(variables=target_variables) instance-attribute

inputs = list(datasets) instance-attribute

frequency = frequencies[0] instance-attribute

dates cached property

Get list of dates that are available in all datasets.

end_date property

Return the end date of the dataset.

start_date property

Return the start date of the dataset.

climatology_for(start_date)

Return the climatology field for each forecast step following the start date.

The calendar day (month/day label) of each forecast step indexes the [366, C, H, W] climatology table, so the result has shape [n_forecast_steps, C, H, W].

Parameters:

Name Type Description Default
start_date datetime64

The start date of the sample.

 required

Returns:

Type Description
ArrayTCHW | None

The stack of calendar-day mean fields for the forecast steps, or None
ArrayTCHW | None

if no climatology table was provided.

get_forecast_steps(start_date)

Return list of consecutive forecast dates for a given start date.

get_history_steps(start_date)

Return list of consecutive history dates for a given start date.

CommonDataModule

Bases: LightningDataModule

base_path = Path(config['base_path']) instance-attribute

dataset_groups = defaultdict(list) instance-attribute

target_group_name = config['predict']['target']['group_name'] instance-attribute

batch_size = int(config['data']['split']['batch_size']) instance-attribute

predict_periods = [{str(k): None if v is None else str(v) for k, v in period.items()} for period in config['data']['split']['predict']] instance-attribute

test_periods = [{str(k): None if v is None else str(v) for k, v in period.items()} for period in config['data']['split']['test']] instance-attribute

train_periods = [{str(k): None if v is None else str(v) for k, v in period.items()} for period in config['data']['split']['train']] instance-attribute

val_periods = [{str(k): None if v is None else str(v) for k, v in period.items()} for period in config['data']['split']['validate']] instance-attribute

n_forecast_steps = int(config['predict'].get('n_forecast_steps', 1)) instance-attribute

n_history_steps = int(config['predict'].get('n_history_steps', 1)) instance-attribute

datasets cached property

Return a dictionary of dataset group names to SingleDataset objects.

hemisphere cached property

Return the hemisphere of the dataset.

input_spaces cached property

Return the data space for each input.

latitudes cached property

Return the latitudes of the dataset.

longitudes cached property

Return the longitudes of the dataset.

mask_directory cached property

Mask directory for the prediction target group.

A target group usually holds a single dataset with generated masks, but if it holds several, pick the first. Combining masks across datasets is unsupported.

output_space cached property

Return the data space of the desired output.

target_variables cached property

Return the names of the variables to predict.

target_variable_indices cached property

Return the indices of the variables to predict.

climatology cached property

Return the climatology: calendar-day means of the target variables.

The [366, C, H, W] table holds, for each calendar day (month/day label), the mean of the normalised target fields over dates sharing that calendar day within the averaging period. The averaging period is the union of the training split's date ranges, intersected with the dates available in the target dataset; it is never widened to dates outside the configured training periods. Dates that are missing from the dataset are never included in a mean.

29 February is the exception: because a training period spanning only non-leap years has no such date, it is not required to have its own data. If no date in the averaging period falls on 29 February, that slot instead copies the 28 February mean.

Raises:

Type Description
ValueError

If the training periods have no available dates at all, or a calendar day other than 29 February has no available dates in the period.

variable_names cached property

Return the variable names for each input.

assign_workers(n_workers)

Assign number of workers for data loading.

predict_dataloader()

Construct predict dataloader.

test_dataloader()

Construct test dataloader.

train_dataloader()

Construct train dataloader.

val_dataloader()

Construct validation dataloader.

SingleDataset

Bases: Dataset

A dataset containing one or more timeslices of data from a single source.

anemoi_cache = {} class-attribute

hemisphere = 'north' if any('north' in str(input_file).lower() for input_file in input_files) else 'south' instance-attribute

dataslices cached property

Get all slices of contiguous dates from the underlying Anemoi dataset.

dates cached property

Return all available dates in the dataset, removing any that are missing.

end_date cached property

Return the end date of the dataset.

frequency cached property

Return the frequency of the dataset.

Use frequency of the underlying Anemoi dataset, which is read from metadata, rather than the frequency of the dataslices which are recalculated on-the-fly by Anemoi through diff-ing the first two dates in the dataslice, which is incorrect for datasets with missing dates.

latitudes cached property

Return the latitudes of the dataset.

longitudes cached property

Return the longitudes of the dataset.

name cached property

Return the name of the dataset.

space cached property

Return the data space for this dataset.

start_date cached property

Return the start date of the dataset.

variable_names cached property

Return the variable names for this dataset.

statistics cached property

Return per-channel statistics from the underlying dataset.

Keys are 'mean', 'stdev', 'maximum', 'minimum', each with shape [C].

load_dataset(input_files) classmethod

normalise_date(np_datetime) staticmethod

Normalise a datetime to noon.

normalise_date_ranges(date_ranges) staticmethod

Sort the ranges, then merge overlapping or touching ones into single spans.

Ranges that overlap/touch must merge into one Anemoi subset, or data fetching step comes back short and fails to reshape. Two ranges merge when the later one starts on/before the day after the earlier one finishes; the merged span is the union of all ranges.

A "None" bound is open, an open-ended previous range always overlap with the later ranges, an open-started later range always overlap with previous ranges.

Each merge is logged so a user notices that their config ranges were altered and verifies.

get_tchw(dates)

Return normalised data for an arbitrary sequence of timesteps in [T, C, H, W] format.

get_tchw_slice(start_date, n_steps, *, check=True)

Return the data for consecutive timesteps in [T, C, H, W] format.

Since contiguous dates must be in a single dataslice, we simply identify which one this is and read from it.

If check is True then we check that we're not crossing the boundary between dataslices, adding a small amount of overhead.

Parameters:

Name Type Description Default
start_date datetime64

The date of the first timestep to return.

 required
n_steps int

The number of consecutive timesteps to return.

 required
check bool

Whether to check that the requested slice is valid. If False, this method may return meaningless or incorrect data if the requested slice is invalid

 True

normalise(data)

Normalise the data to [0, 1] for each channel if configured to do so.

Note that this can be applied to both ArrayCHW and ArrayTCHW.

subset(*, date_ranges=None, normalise=None, variables=None)

to_index(date)

Return the index of a given date in the dataset.

calendar_day_index(day)

Return the 0-365 calendar-day index (month/day label) for a date.