Remove n_timesteps_per_batch and _cache from DataSources.

Author: JackKelly

conftest.py

@@ -22,7 +22,7 @@
 register_xr_data_set_to_tensor()
 
 
-def pytest_addoption(parser):
+def pytest_addoption(parser):  # noqa: D103
     parser.addoption(
         "--use_cloud_data",
         action="store_true",
@@ -32,12 +32,12 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture
-def use_cloud_data(request):
+def use_cloud_data(request):  # noqa: D103
     return request.config.getoption("--use_cloud_data")
 
 
 @pytest.fixture
-def sat_filename(use_cloud_data: bool) -> Path:
+def sat_filename(use_cloud_data: bool) -> Path:  # noqa: D103
     if use_cloud_data:
         return consts.SAT_FILENAME
     else:
@@ -47,24 +47,23 @@ def sat_filename(use_cloud_data: bool) -> Path:
 
 
 @pytest.fixture
-def sat_data_source(sat_filename: Path):
+def sat_data_source(sat_filename: Path):  # noqa: D103
     return SatelliteDataSource(
         image_size_pixels=pytest.IMAGE_SIZE_PIXELS,
         zarr_path=sat_filename,
         history_minutes=0,
         forecast_minutes=5,
         channels=("HRV",),
-        n_timesteps_per_batch=2,
     )
 
 
 @pytest.fixture
-def general_data_source():
+def general_data_source():  # noqa: D103
     return MetadataDataSource(history_minutes=0, forecast_minutes=5, object_at_center="GSP")
 
 
 @pytest.fixture
-def gsp_data_source():
+def gsp_data_source():  # noqa: D103
     return GSPDataSource(
         image_size_pixels=16,
         meters_per_pixel=2000,
@@ -75,13 +74,13 @@ def gsp_data_source():
 
 
 @pytest.fixture
-def configuration():
+def configuration():  # noqa: D103
     filename = os.path.join(os.path.dirname(nowcasting_dataset.__file__), "config", "gcp.yaml")
     configuration = load_yaml_configuration(filename)
 
     return configuration
 
 
 @pytest.fixture
-def test_data_folder():
+def test_data_folder():  # noqa: D103
     return os.path.join(os.path.dirname(nowcasting_dataset.__file__), "../tests/data")

notebooks/2021-09/2021-09-07/sat_data.py

@@ -1,3 +1,4 @@
+"""Notebook"""
 from datetime import datetime
 
 from nowcasting_dataset.data_sources.satellite.satellite_data_source import SatelliteDataSource
@@ -9,7 +10,6 @@
     forecast_len=12,
     image_size_pixels=64,
     meters_per_pixel=2000,
-    n_timesteps_per_batch=32,
 )
 
 s.open()

nowcasting_dataset/data_sources/data_source.py

@@ -6,17 +6,20 @@
 from pathlib import Path
 from typing import Iterable, List, Tuple
 
+import futures
 import pandas as pd
 import xarray as xr
 
-import nowcasting_dataset.filesystem.utils as nd_fs_utils
-
 # nowcasting_dataset imports
+import nowcasting_dataset.filesystem.utils as nd_fs_utils
 import nowcasting_dataset.time as nd_time
 import nowcasting_dataset.utils as nd_utils
 from nowcasting_dataset import square
 from nowcasting_dataset.data_sources.datasource_output import DataSourceOutput
-from nowcasting_dataset.dataset.xr_utils import join_dataset_to_batch_dataset
+from nowcasting_dataset.dataset.xr_utils import (
+    join_dataset_to_batch_dataset,
+    join_list_data_array_to_batch_dataset,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -310,20 +313,12 @@ class ZarrDataSource(ImageDataSource):
     """
 
     channels: Iterable[str]
-    #: Mustn't be None, but cannot have a non-default arg in this position :)
-    n_timesteps_per_batch: int = None
     consolidated: bool = True
 
     def __post_init__(self, image_size_pixels: int, meters_per_pixel: int):
         """ Post init """
         super().__post_init__(image_size_pixels, meters_per_pixel)
         self._data = None
-        if self.n_timesteps_per_batch is None:
-            # Using hacky default for now.  The whole concept of n_timesteps_per_batch
-            # will be removed when #213 is completed.
-            # TODO: Remove n_timesteps_per_batch when #213 is completed!
-            self.n_timesteps_per_batch = 16
-            logger.warning("n_timesteps_per_batch is not set!  Using default!")
 
     @property
     def data(self):
@@ -378,7 +373,42 @@ def get_example(
         # rename 'variable' to 'channels'
         selected_data = selected_data.rename({"variable": "channels"})
 
-        return selected_data
+        return selected_data.load()
+
+    def get_batch(
+        self,
+        t0_datetimes: pd.DatetimeIndex,
+        x_locations: Iterable[Number],
+        y_locations: Iterable[Number],
+    ) -> DataSourceOutput:
+        """
+        Get batch data
+
+        Args:
+            t0_datetimes: list of timestamps for the datetime of the batches. The batch will also
+                include data for historic and future depending on `history_minutes` and
+                `future_minutes`.
+            x_locations: x center batch locations
+            y_locations: y center batch locations
+
+        Returns: Batch data
+
+        """
+        zipped = list(zip(t0_datetimes, x_locations, y_locations))
+        batch_size = len(t0_datetimes)
+
+        with futures.ThreadPoolExecutor(max_workers=batch_size) as executor:
+            future_examples = []
+            for coords in zipped:
+                t0_datetime, x_location, y_location = coords
+                future_example = executor.submit(
+                    self.get_example, t0_datetime, x_location, y_location
+                )
+                future_examples.append(future_example)
+            examples = [future_example.result() for future_example in future_examples]
+
+        output = join_list_data_array_to_batch_dataset(examples)
+        return self._dataset_to_data_source_output(output)
 
     def geospatial_border(self) -> List[Tuple[Number, Number]]:
         """
@@ -415,3 +445,6 @@ def open(self) -> None:
 
     def _open_data(self) -> xr.DataArray:
         raise NotImplementedError()
+
+    def _dataset_to_data_source_output(output: xr.Dataset) -> DataSourceOutput:
+        raise NotImplementedError()

nowcasting_dataset/data_sources/nwp/nwp_data_source.py

@@ -1,19 +1,16 @@
 """ NWP Data Source """
 import logging
-from concurrent import futures
 from dataclasses import InitVar, dataclass
-from numbers import Number
 from typing import Iterable, Optional
 
 import numpy as np
 import pandas as pd
 import xarray as xr
 
 from nowcasting_dataset import utils
+from nowcasting_dataset.consts import NWP_VARIABLE_NAMES
 from nowcasting_dataset.data_sources.data_source import ZarrDataSource
 from nowcasting_dataset.data_sources.nwp.nwp_model import NWP
-from nowcasting_dataset.dataset.xr_utils import join_list_data_array_to_batch_dataset
-from nowcasting_dataset.consts import NWP_VARIABLE_NAMES
 
 _LOG = logging.getLogger(__name__)
 
@@ -78,72 +75,12 @@ def open(self) -> None:
         data = self._open_data()
         self._data = data["UKV"].sel(variable=list(self.channels))
 
-    def get_batch(
-        self,
-        t0_datetimes: pd.DatetimeIndex,
-        x_locations: Iterable[Number],
-        y_locations: Iterable[Number],
-    ) -> NWP:
-        """
-        Get batch data
-
-        Args:
-            t0_datetimes: list of timstamps
-            x_locations: list of x locations, where the batch data is for
-            y_locations: list of y locations, where the batch data is for
-
-        Returns: batch data
-
-        """
-        # Lazily select time slices.
-        selections = []
-        for t0_dt in t0_datetimes[: self.n_timesteps_per_batch]:
-            selections.append(self._get_time_slice(t0_dt))
-
-        # Load entire time slices from disk in multiple threads.
-        data = []
-        with futures.ThreadPoolExecutor(max_workers=self.n_timesteps_per_batch) as executor:
-            data_futures = []
-            # Submit tasks.
-            for selection in selections:
-                future = executor.submit(selection.load)
-                data_futures.append(future)
-
-            # Grab tasks
-            for future in data_futures:
-                d = future.result()
-                data.append(d)
-
-        # Select squares from pre-loaded time slices.
-        examples = []
-        for i, (x_meters_center, y_meters_center) in enumerate(zip(x_locations, y_locations)):
-            selected_data = data[i % self.n_timesteps_per_batch]
-            bounding_box = self._square.bounding_box_centered_on(
-                x_meters_center=x_meters_center, y_meters_center=y_meters_center
-            )
-            selected_data = selected_data.sel(
-                x=slice(bounding_box.left, bounding_box.right),
-                y=slice(bounding_box.top, bounding_box.bottom),
-            )
-
-            # selected_sat_data is likely to have 1 too many pixels in x and y
-            # because sel(x=slice(a, b)) is [a, b], not [a, b).  So trim:
-            selected_data = selected_data.isel(
-                x=slice(0, self._square.size_pixels), y=slice(0, self._square.size_pixels)
-            )
-
-            t0_dt = t0_datetimes[i]
-            selected_data = self._post_process_example(selected_data, t0_dt)
-
-            examples.append(selected_data)
-
-        output = join_list_data_array_to_batch_dataset(examples)
-
-        return NWP(output)
-
     def _open_data(self) -> xr.DataArray:
         return open_nwp(self.zarr_path, consolidated=self.consolidated)
 
+    def _dataset_to_data_source_output(output: xr.Dataset) -> NWP:
+        NWP(output)
+
     def _get_time_slice(self, t0_dt: pd.Timestamp) -> xr.DataArray:
         """
         Select the numerical weather predictions for a single time slice.
@@ -177,20 +114,9 @@ def _get_time_slice(self, t0_dt: pd.Timestamp) -> xr.DataArray:
         selected["target_time"] = init_time + selected.step
         return selected
 
-    def _post_process_example(
-        self, selected_data: xr.DataArray, t0_dt: pd.Timestamp
-    ) -> xr.DataArray:
-        """Resamples to 5 minutely."""
-        start_dt = self._get_start_dt(t0_dt)
-        end_dt = self._get_end_dt(t0_dt)
-        selected_data = selected_data.resample({"target_time": "5T"})
-        selected_data = selected_data.interpolate()
-        selected_data = selected_data.sel(target_time=slice(start_dt, end_dt))
+    def _post_process_example(self, selected_data: xr.DataArray) -> xr.DataArray:
         selected_data = selected_data.rename({"target_time": "time"})
         selected_data = selected_data.rename({"variable": "channels"})
-
-        selected_data.data = selected_data.data.astype(np.float32)
-
         return selected_data
 
     def datetime_index(self) -> pd.DatetimeIndex: