"Big new design" for nowcasting_dataset

[JackKelly]

Detailed Description

Over the last few weeks, we've come up with a bunch of ideas for how to simplify nowcasting_dataset.

This issue exists to keep track of all of those "new design" issues and to discuss how these individual issues hang together into a single coherent design; and to discuss how to implement this new design in a sequence of easy-to-digest chunks :)

The plan:

In terms of sequencing this work, I'm now thinking I'll do something along the lines of this:

First, do some stand-alone "preparatory" work. Specifically:

• Simplify the calculation of available datetimes across all DataSources #204
• Tweak the way DataSources are represented in the Configuration model #217
• Get intersection of _Periods_ of available data across all DataSources, instead of using datetimes. #223
• Change each DataSource subclass' constructor arguments to align with the new YAML config field names #270
• Assert there's no overlap between train, test and validation datetimes at end of split() function #299

Implement and write tests for some of the functions in the draft design in the GitHub comment below. (for now, these functions won't actually be called in the code):

• sample_spatial_and_temporal_locations_for_examples() (done in PR Implement DataSourceList.sample_spatial_and_temporal_locations_for_examples() #278)
• DataSource.prepare_batches()

Then, the big one (where the code size will hopefully shrink a lot!):

• Re-write prepare_ml_data.py so it looks like the sketch above. Use click to pass in command-line params.
• implement the remaining utility functions
• Change DataSourceList into Manager; and maintain DataSources in a dict instead of a list? #298
• delete the unnecessary code,
• Remove PyTorch from the code #86
• Detect & remove unused functions #170
• Remove n_timesteps_per_batch?
• Stop using the first entry in DataSourceList as the one which defines the geospatial location of each example.
• Search through the code for TODOs associated with issues associated with this big design change.
• Allow user to configure the frequency of the t0 datetimes in the config yaml #277
• Document the new architecture. Ideally with a diagram 🙂
• Save history_length and forecast_length to disk for each modality, and use these in dataloader #293
• Turn off temporal interpolation of numerical weather predictions (NWPs) #135

Remove:

• All of dataset/datamodule.py
• dataset.datasets.NowcastingDataset
• dataset.datasets.worker_init_fn()
• All the batch_to_dataset and dataset_to_batch stuff
• All the to_numpy and from_numpy stuff (assuming we can go straight from xr.Dataset to torch.Tensor)

Related issues

• Machine-readable schema & validator for xarray.Dataset #211
• Can we simplify the code by always keeping the data in one data type (e.g. xr.DataArray) per modality? #209
• Use independent processes for each "modality" #202
• Example --> Pydantic #166
• Remove time_30 and time_5 #230
• Experiment with loading entire batches at once #212
• Compute datetime features on-the-fly in nowcasting_dataloader. Remove datetime features from the on-disk batches. #208
• Simplify the calculation of available datetimes across all DataSources #204
• Implement a thin "data loading" layer to help ML training #97
• refactor DataSources to have different functions for selecting timesteps, selecting geo locations, and post-processing #48
• Use NamedTensors to name dimensions #25
• Do we still need all the representations of sequence lengths in DataSource? #219
• Detect & remove unused functions #170

A bit more context

I think I've made our lives far harder than they need to be by trying to support two different use-cases:

1. Loading training data on-the-fly from multiple Zarr files during ML training; and
2. Pre-preparing batches.

I think we can make life way easier by dropping support for use-case 1 :)

Here's the broad proposal that it'd be great to discuss:

We drop support for loading data directly from Zarr on-the-fly during ML training (which we haven't done for months, and - now that we're using large NWP images - it would be far too slow). nowcasting_dataset becomes laser-focused on pre-preparing batches (just as we use it now).

This allows us to completely rip out PyTorch from nowcasting_dataset (#86); and enables each "modality" to stays in a single data type throughout nowcasting_dataset (#209). e.g. Satellite data stays in an xr.Dataset. Each modality would be processed concurrently in different processes; and would be output into different directories (e.g. train/satellite/ and train/nwp/) (#202).

Inspired by and making use of @peterdudfield's Pydantic PR (#195), we'd have a formal schema for the data structures in nowcasting_dataset (#211).

The ultimate aim is to simplify the code (I'm lazy!), whilst keeping all the useful functionality, and making the code easier to extend & maintain 🙂

Of course, we'll still use a pytorch dataloader to load the pre-prepared batches off disk into an ML model. But that's fine; and should work in very a similar (maybe identical?) fashion to how it works now 🙂

I certainly can't claim to have thought this all through properly! And everything's up for discussion, of course!

[JackKelly]

In terms of the order in which to do this work... here's some initial thoughts (very much up for discussion!)

1. Merge Issue/166 batch pydantic #195 into main DONE
2. Simplify the calculation of available datetimes across all DataSources #204 [newly added]
3. Tweak the way DataSources are represented in the Configuration model #217 [newly added]
4. Use independent processes for each "modality" #202
5. Remove PyTorch from the code #86
6. Can we simplify the code by always keeping the data in one data type (e.g. xr.DataArray) per modality? #209
7. Machine-readable schema & validator for xarray.Dataset #211

The thinking being that we can't keep each modality in a single datatype until we've removed pytorch (because pytorch requires us to move data from xarray to numpy to tensors); which we can't do until we've removed the pytorch dataloader; and the best way to remove the pytorch dataloader is to implement #202 :) Does that sound roughly right?!

[peterdudfield]

Great work @JackKelly putting this all together. I think this will be super useful at making a really solid base for our ML learning problems we are trying to solve.

I'm slightly biased by 1. , but yea I would do 1. first. #195
2. #202 then is a good next logic step. I would be tempted to split this into two tasks.
A. Change current save to netcdf to save several files. B. Change core code sot hat it can run in different processes. Doing A, first will allows us to remake datasets more easily, and just copy the ones that are the same from previous versions.
3. 4. Perhaps this can be done at the same time. I feel like doing #209 first gives us a quicker win, and makes things simpler.
5. Then #211 will be a nice bonus on the end.

I'll just share my thoughts on how #209 can be done, that allows data to end up as numpy / tensor values

[jacobbieker]

Yeah, great work! I also think the outline here for moving forward makes sense. For #86, I think that can mostly come as we refactor the code, we remove PyTorch as we see it, I think I agree with @peterdudfield that we can probably do 3 and 4 at the same time.

With the focus on removing the on-the-fly loading, would we be moving the PyTorch dataloader that does currently load the prepared batches to somewhere else, like nowcasting_utils? Or would we be leaving it in here? I generally would think keeping it in nowcasting_dataset keeps it tied to how the data is being created easier, but then we would still have some dependency on PyTorch in this repo.

[peterdudfield]

Yeah, great work! I also think the outline here for moving forward makes sense. For #86, I think that can mostly come as we refactor the code, we remove PyTorch as we see it, I think I agree with @peterdudfield that we can probably do 3 and 4 at the same time.

With the focus on removing the on-the-fly loading, would we be moving the PyTorch dataloader that does currently load the prepared batches to somewhere else, like nowcasting_utils? Or would we be leaving it in here? I generally would think keeping it in nowcasting_dataset keeps it tied to how the data is being created easier, but then we would still have some dependency on PyTorch in this repo.

I like the idea of leaving it here too, some how keeps it rapped up nicely. I wonder if we could do it as an optional requirement. That means people dont' have to install it, but they do need to install it to make the data loader. And it keeps things all here

[jacobbieker]

Yeah, I think an optional requirement would be the way to go with it

[JackKelly]

Here's a draft design sketch for the end-goal, once all the issues listed above are done:

To recap: One of the main aims is to significantly simplify nowcasting_dataset (maybe reducing the total lines of code to something like half the current size!) Some specific design features of this new design:

• Data for each modality flows through nowcasting_dataset completely independently of other modalities. As such, the data can stay in one data type throughout the code (e.g. xr.Dataset), and the data never needs to leave its DataSource.
• Each modality will be saved to a different directory, concurrently

This makes it easy to:

• run each data source concurrently.
• run multiple processes per data source (by splitting positions_of_each_batch into, say, quarters and sending each quarter to a different process).
• start processing at any batch number (e.g. if processing crashed part-way-through and needs to re-start).

Once this is implemented, I think we may be able to literally delete several large chunks of code:

• All of dataset/datamodule.py
• dataset.datasets.NowcastingDataset
• dataset.datasets.worker_init_fn()
• All the batch_to_dataset and dataset_to_batch stuff
• All the to_numpy and from_numpy stuff (assuming we can go straight from xr.Dataset to torch.Tensor)

prepare_ml_data.py rough sketch

# Instantiate the `DataSource` objects requested by the user at the command line and/or config YAML file.
# The first data_source will be used to define the geospatial location of each example.
# We maybe want to make a few small changes to the Configuration model:
# - include a list of DataSources.
# - maybe split the config up so each data source has its own section, so that entire
#   section can be passed into the DataSource object.  e.g. there would be an `NWP` section, 
#   with fields zarr_path, image_size, and channels.
data_sources = get_data_sources(config)

# If necessary, pre-compute three `positions_of_each_example.csv` files:
# One file for `train/`, one for `test/`, and one for `/validation`. 
# This CSV file will have four columns: example_number, t0_datetime, center_x, center_y.
# (We might want to use a binary file rather than CSV.  
# But CSV is nice because it's super-easy to inspect manually).
# (These `positions_of_each_example.csv` files could be re-used when we prepare different versions of
# the pre-prepared batches, to minimise variation between different datasets).
filenames_of_positions_of_each_example: dict[SplitName, Pathy] = get_filenames_of_positions_of_each_example(
    config.output_data.filepath)

if not exists(filenames_of_positions_of_each_example[SplitName.TRAIN]):
    compute_and_save_positions_of_each_example_of_each_split(
        data_sources[0], 
        split_method=config.process.split_method, 
        dst_path=config.output_data.filepath)

# Generate the batches.  Something like this:
for split_name, filename_of_positions_of_each_example in filenames_of_positions_of_each_example.items():
    dst_path = config.output_data.filepath / split_name
    positions_of_each_example = pd.read_csv(filename_of_positions_of_each_example)
    positions_for_each_batch = group_positions_of_each_example_into_batches(
        positions_of_each_example, config.process.batch_size)

    with futures.ProcessPoolExecutor() as executor:
        for data_source in data_sources:
            executor.submit(data_source.prepare_batches, positions_for_each_batch, dst_path, overwrite=False)

validate_ml_data()  # ?

compute_and_save_positions_of_each_example_of_each_split() rough sketch

def compute_and_save_positions_of_each_example_of_each_split(
    data_source: DataSource, # The DataSource that defines the location of each example.
    split_method: str,
    n_examples_per_split: dict[SplitName, int],
    dst_path: Pathy
) -> None:
    # Get intersection of all available t0_datetimes.  Current done by NowcastingDataModule._get_datetimes():
    # github.com/openclimatefix/nowcasting_dataset/blob/main/nowcasting_dataset/dataset/datamodule.py#L364
    t0_datetimes_for_all_data_sources = [data_source.get_t0_datetimes() for data_source in data_sources]
    intersection_of_t0_datetimes = nd_time.intersection_of_datetimeindexes(t0_datetimes_for_all_data_sources)
    
    # Split t0_datetimes into train, test and validation sets (being careful to ensure each group is
    # at least `total_seq_duration` apart).  Currently done by NowcastingDataModule._split_data():
    # github.com/openclimatefix/nowcasting_dataset/blob/main/nowcasting_dataset/dataset/datamodule.py#L315
    t0_datetimes_per_split: dict[SplitName, pd.DatetimeIndex] = split_datetimes(
        intersection_of_t0_datetimes, method=split_method)
    
    for split_name, t0_datetimes_for_split in t0_datetimes_per_split.items():
        n_examples = n_examples_per_split[split_name]
        positions = compute_positions_of_each_example(t0_datetimes_for_split, data_source, n_examples)
        filename = dst_path / split_name / 'positions_of_each_example.csv'
        positions.to_csv(filename)


def compute_positions_of_each_example(
    t0_datetimes_for_split: pd.DatetimeIndex, 
    data_source: DataSource, 
    n_examples: int
) -> pd.DataFrame:
    
    positions = pd.DataFrame(index=range(n_examples), columns=['t0_datetime', 'x_center', 'y_center'])
    sampled_t0_datetimes = np.random.choice(t0_datetimes_for_split, size=n_examples)
    for example_i, t0_datetime in enumerate(sampled_t0_datetimes):
        x_center, y_center = data_source.get_geo_location_for_t0_datetime(t0_datetime)
        positions.iloc[i] = dict(t0_datetime=t0_datetime, x_center=x_center, y_center=y_center)
    
    return positions
    

DataSource.prepare_batches() rough sketch

def prepare_batches(
    self, 
    positions_of_each_batch: pd.DataFrame, 
    dst_path: Pathy,
    overwrite: bool = False,
    temp_path: Optional[Pathy] = None,
    upload_every_n_batches: Optional[int] = None
) -> None:
    """
    Args:
        positions_of_each_batch: DataFrame with MultiIndex of the batch_index and the the example_index. 
            The columns are: t0_datetime, x_center, y_center
        dst_path: The path to save to.  Can start with '/' (for local filesystem) or 's3://' or 'gs://'.
        overwrite: If True then overwrite existing batches on disk.
    """
    
    self.open()  # Need to open Zarr stores in each process.
    
    dst_path = dst_path / self.name  # Where `name` is a string like 'NWP' or 'satellite'.
    if not overwrite:
        indicies_of_existing_batches = self._find_existing_batches(dst_path)
        _LOG.info(f"Found {len(indicies_of_existing_batches)} existing batches in {dst_path}.")
        indicies_of_existing_batches = positions_of_each_batch.drop(indicies_of_existing_batches)
    
    for positions_for_batch in positions_for_each_batch:
        batch = self.get_batch(positions_for_batch)
        # TODO: Logic to save upload_every_n_batches to temp_path, then upload.  But only if on the cloud!
        self.save(batch, dst_path)

Comments more than welcome!

[JackKelly]

(UPDATE: Plan moved to top of this thread!)

Does that sound OK, @peterdudfield ?

[peterdudfield]

Think that sounds good. I would move the test writing from 2. to 1.

[JackKelly]

I've got the afternoon free for coding so I'm going to make a start on the plan in the comment above... Unless that will clash with your work, @peterdudfield ?

[peterdudfield]

go for it