autocast.data.dataset

autocast.data.dataset#

class BatchMixin[source]#

Bases: object

A mixin class to provide Batch conversion functionality.

static to_sample(data)[source]#

Convert a dictionary of tensors to a Sample object.

Parameters:

data (dict)

Return type:

Sample

class SpatioTemporalDataset(data_path, data=None, n_steps_input=1, n_steps_output=1, stride=1, channel_idxs=None, full_trajectory_mode=False, autoencoder_mode=False, dtype=torch.float32, verbose=False, use_normalization=False, normalization_type=<class 'the_well.data.normalization.ZScoreNormalization'>, normalization_path=None, normalization_stats=None)[source]#

Bases: Dataset, BatchMixin

A class for spatio-temporal datasets.

Parameters:

  • data_path (str | None)

  • data (dict | None)

  • n_steps_input (int)

  • n_steps_output (int)

  • stride (int)

  • channel_idxs (tuple[int, ...] | None)

  • full_trajectory_mode (bool)

  • autoencoder_mode (bool)

  • dtype (dtype)

  • verbose (bool)

  • use_normalization (bool)

  • normalization_type (type[ZScoreNormalization] | None)

  • normalization_path (str | None)

  • normalization_stats (dict | DictConfig | None)

read_data(data_path)[source]#

Read data.

By default assumes HDF5 format in data_path with correct shape and fields.

Parameters:

data_path (str)

parse_data(data)[source]#

Parse data from a dictionary.

Parameters:

data (dict | None)

to_preloaded_data()[source]#

Return the in-memory payload accepted by data= without copying.

Return type:

dict[str, Any]

set_up_normalization()[source]#

Set up normalizer (None if self.use_normalization = False).

class ReactionDiffusionDataset(*args, **kwargs)[source]#

Bases: SpatioTemporalDataset

Reaction-Diffusion dataset.

class AdvectionDiffusionDataset(*args, **kwargs)[source]#

Bases: SpatioTemporalDataset

Advection-Diffusion dataset.

class BOUTDataset(*args, **kwargs)[source]#

Bases: SpatioTemporalDataset

BOUT++ dataset.

class TheWell(path=None, normalization_path='../stats.yaml', well_base_path=None, well_dataset_name=None, well_split_name=None, include_filters=None, exclude_filters=None, use_normalization=False, normalization_type=<class 'the_well.data.normalization.ZScoreNormalization'>, max_rollout_steps=100, n_steps_input=1, n_steps_output=1, min_dt_stride=1, max_dt_stride=1, flatten_tensors=True, cache_small=True, max_cache_size=1000000000.0, return_grid=True, boundary_return_type='padding', full_trajectory_mode=False, autoencoder_mode=False, name_override=None, transform=None, min_std=0.0001, storage_options=None)[source]#

Bases: SpatioTemporalDataset

A wrapper around The Well’s WellDataset to provide Batch objects.

Parameters:

  • path (None | str)

  • normalization_path (str)

  • well_base_path (None | str)

  • well_dataset_name (None | str)

  • well_split_name (Literal['train', 'valid', 'test', None])

  • include_filters (list[str] | None)

  • exclude_filters (list[str] | None)

  • use_normalization (bool)

  • normalization_type (None | Callable[[...], Any] | ZScoreNormalization)

  • n_steps_input (int)

  • n_steps_output (int)

  • min_dt_stride (int)

  • max_dt_stride (int)

  • flatten_tensors (bool)

  • cache_small (bool)

  • max_cache_size (float)

  • return_grid (bool)

  • boundary_return_type (str)

  • full_trajectory_mode (bool)

  • autoencoder_mode (bool)

  • name_override (None | str)

  • transform (None | Augmentation)

  • min_std (float)

  • storage_options (None | dict)

well_dataset: WellDataset#