add test

Author: peterdudfield

nowcasting_dataset/manager/manager_live.py

@@ -76,15 +76,15 @@ def create_files_specifying_spatial_and_temporal_locations_of_each_example(
 
         n_batches_requested = int(np.ceil(n_gsps / self.config.process.batch_size))
 
-        n_examples = n_batches_requested * self.config.process.batch_size
         logger.debug(
             f"Creating {n_batches_requested:,d} batches x {self.config.process.batch_size:,d}"
-            f" examples per batch = {n_examples:,d} examples for {split_name}."
+            f" examples per batch = {n_batches_requested:,d} examples for {split_name}."
         )
 
         locations = self.sample_spatial_and_temporal_locations_for_examples(
-            t0_datetime=datetimes_for_split[0],
+            t0_datetime=datetimes_for_split[0], n_examples=n_gsps
         )
+
         metadata = Metadata(
             batch_size=self.config.process.batch_size, space_time_locations=locations
         )
@@ -95,7 +95,7 @@ def create_files_specifying_spatial_and_temporal_locations_of_each_example(
         metadata.save_to_csv(path_for_csv)
 
     def sample_spatial_and_temporal_locations_for_examples(
-        self, t0_datetime: datetime
+        self, t0_datetime: datetime, n_examples: Optional[int] = None
     ) -> List[SpaceTimeLocation]:
         """
         Computes the geospatial and temporal locations for each training example.
@@ -122,6 +122,10 @@ def sample_spatial_and_temporal_locations_for_examples(
             t0_datetimes_utc=pd.DatetimeIndex([t0_datetime])
         )
 
+        # reduce locations to n_examples
+        if n_examples is not None:
+            locations = locations[:n_examples]
+
         # find out the number of examples in the last batch,
         # we maybe need to duplicate the last example into order to get a full batch
         n_examples_last_batch = len(locations) % self.config.process.batch_size

tests/manager/test_manager_live.py

@@ -63,6 +63,42 @@ def test_create_files_specifying_spatial_and_temporal_locations_of_each_example(
         assert len(locations_df) == batch_size
 
 
+def test_create_files_locations_of_each_example_reduced(
+    test_configuration_filename,
+):
+    """Test to create locations, for a reduced number of n_gsps"""
+
+    manager = ManagerLive()
+    manager.load_yaml_configuration(filename=test_configuration_filename)
+    manager.config.process.batch_size = 5
+    manager.config.process.split_method = SplitMethod.SAME
+    manager.initialize_data_sources()
+    t0_datetime = datetime(2021, 4, 1)
+
+    batch_size = manager.config.process.batch_size
+
+    with tempfile.TemporaryDirectory() as local_temp_path, tempfile.TemporaryDirectory() as dst_path:  # noqa 101
+
+        manager.config.output_data.filepath = Path(dst_path)
+        manager.local_temp_path = Path(local_temp_path)
+
+        manager.create_files_specifying_spatial_and_temporal_locations_of_each_example(
+            t0_datetime=t0_datetime, n_gsps=7
+        )  # noqa 101
+
+        live_file = f"{dst_path}/live/{SPATIAL_AND_TEMPORAL_LOCATIONS_OF_EACH_EXAMPLE_FILENAME}"
+
+        assert os.path.exists(live_file)
+        locations_df = pd.read_csv(live_file)
+        # we've asked for 7 examples, but the batchsize is 5, so we get 10 examples,
+        # the last 3 being filled in
+        assert len(locations_df) == batch_size * 2
+        # check last 3 are filled in from the first one
+        assert locations_df.iloc[-3].id == 1
+        assert locations_df.iloc[-2].id == 1
+        assert locations_df.iloc[-1].id == 1
+
+
 def test_batches(test_configuration_filename, sat, gsp):
     """Test that batches can be made"""