Skip to content

Data Loaders

CalendarDayClimatology dataclass

Pixel-wise statistics for each calendar day, in the dataset's own units.

Dates are grouped by month/day label ("MM-DD") rather than ordinal day-of-year, so 29 February keeps its own slot and every other calendar day stays aligned between leap and non-leap years.

Attributes:

Name Type Description
mean ArrayTCHW

[366, C, H, W] mean field for each calendar day.
std ArrayTCHW

[366, C, H, W] population standard deviation for each calendar day.
n_dates ndarray

[366] number of dates falling on each calendar day.

DAY_INDEX = {str(day)[5:]: index for index, day in enumerate(np.arange(np.datetime64('2000-01-01'), np.datetime64('2001-01-01'), dtype='datetime64[D]'))} class-attribute

HALF_WINDOW = 7 class-attribute

mean instance-attribute

std instance-attribute

n_dates instance-attribute

day_index(day) classmethod

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

from_dataset(dataset, dates) classmethod

Compute smoothed pixel-wise mean and standard deviation for each calendar day.

The data is read once, reducing it to a count, sum and sum of squares of the finite values on each calendar day. These are then smoothed around the 366-day calendar with triangular weights, falling linearly from 1 at the centre to 1 / (HALF_WINDOW + 1) at each end, so each calendar day draws on every date within HALF_WINDOW calendar days of it, wrapping from 31 December to 1 January. Every year is treated as having a 29 February slot: that slot draws on 28 February and 1 March from every year, and in non-leap years those two days are treated as two days apart rather than one.

Sums are accumulated in float64. Non-finite values are excluded pixel by pixel, so a pixel is NaN only if it has no finite values within the window; the caller decides whether that is an error.

Parameters:

Name Type Description Default
dataset SingleDataset

The dataset to read fields from.

 required
dates Sequence[datetime64]

The dates to average over.

 required

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.

variable_list()

Return a pretty-formatted list of input variables in the combined dataset.

Returns:

Type Description
list[str]

A list of pretty-formatted strings like:
list[str]

[ "Combined dataset has 3 input variables:", " 1) dataset1/variableA", " 2) dataset1/variableB", " 3) dataset2/variableC",
list[str]

]

CommonDataModule

Bases: LightningDataModule

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

dataset_groups = defaultdict(list) instance-attribute

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

predict_periods = self._normalise(config['data']['split']['predict']) instance-attribute

test_periods = self._normalise(config['data']['split']['test']) instance-attribute

train_periods = self._normalise(config['data']['split']['train']) instance-attribute

val_periods = self._normalise(config['data']['split']['validate']) instance-attribute

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

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

climatology cached property

Return the calendar-day statistics of the target variables.

For each calendar day (month/day label), the [366, C, H, W] mean and standard deviation tables hold statistics of the normalised target fields over dates sharing that calendar day within the date range under consideration. This range is the overlap of the training periods with the dates available in the target dataset. NaN pixels resulting from missing data are excluded.

Each calendar day's statistics are smoothed over a weighted window of CalendarDayClimatology.HALF_WINDOW days either side, so 29th February draws on its neighbours in every year (see CalendarDayClimatology.from_dataset).

Returns:

Type Description
CalendarDayClimatology | None

A CalendarDayClimatology object holding the mean, standard deviation,
CalendarDayClimatology | None

and number of dates contributing to each calendar day or None if the
CalendarDayClimatology | None

climatology cannot be built.

datasets cached property

Return a filtered dictionary of dataset group names to SingleDataset objects.

Only include requested variables for each dataset group. If no variables are requested for a dataset group, ignore it.

datasets_unfiltered cached property

Return an unfiltered 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_group_name cached property

Return the name of the target variable group.

target_variables cached property

Return the names of the variables to predict, in on-disk order.

target_variable_indices cached property

Return the indices of the target variables within their dataset group.

These indices are used to select target channels from the target SingleDataset, so they must be resolved against the actual variable order in that dataset. This is determined by the underlying data's storage order rather than the arbitrary order the variables are listed in variables.input.

training_dataset cached property

Return the dataset used for training.

assign_workers(n_workers)

Assign number of workers for data loading.

predict_dataloader()

Construct predict dataloader.

test_dataloader()

Construct test dataloader.

train_dataloader(*, shuffle=True)

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 = Hemisphere.NORTH if any('north' in str(input_file).lower() for input_file in input_files) else Hemisphere.SOUTH instance-attribute

name = name instance-attribute

group = name if group is None else group 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.

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.

variable_names_on_disk cached property

Return all variable names stored on disk, ignoring any variable selection.

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.