check columns names

Author: JackKelly

nowcasting_dataset/data_sources/data_source.py

@@ -167,7 +167,7 @@ def create_batches(
         assert batch_size > 0
         assert len(spatial_and_temporal_locations_of_each_example) % batch_size == 0
         assert upload_every_n_batches >= 0
-        assert spatial_and_temporal_locations_of_each_example.columns.to_tuple() == (
+        assert spatial_and_temporal_locations_of_each_example.columns.to_list() == list(
             SPATIAL_AND_TEMPORAL_LOCATIONS_COLUMN_NAMES
         )
 

nowcasting_dataset/manager.py

@@ -342,8 +342,8 @@ def create_batches(self, overwrite_batches: bool) -> None:
         for split_name in splits_which_need_more_batches:
             filename = self._filename_of_locations_csv_file(split_name.value)
             logger.info(f"Loading {filename}.")
-            locations_for_each_example = pd.read_csv(filename)
-            assert locations_for_each_example.columns.to_tuple() == (
+            locations_for_each_example = pd.read_csv(filename, index_col=0)
+            assert locations_for_each_example.columns.to_list() == list(
                 SPATIAL_AND_TEMPORAL_LOCATIONS_COLUMN_NAMES
             )
             # Converting to datetimes is much faster using `pd.to_datetime()` than
@@ -399,7 +399,12 @@ def create_batches(self, overwrite_batches: bool) -> None:
                     future_create_batches_jobs.append(future)
 
                 # Wait for all futures to finish:
-                for future in future_create_batches_jobs:
+                for future, data_source_name in zip(
+                    future_create_batches_jobs, self.data_sources.keys()
+                ):
                     # Call exception() to propagate any exceptions raised by the worker process into
                     # the main process, and to wait for the worker to finish.
-                    future.exception()
+                    exception = future.exception()
+                    if exception is not None:
+                        logger.exception(f"Worker process {data_source_name} raised exception!")
+                        raise exception

scripts/prepare_ml_data.py

@@ -68,6 +68,7 @@ def main(config_filename: str, data_source: list[str], overwrite_batches: bool):
     manager.create_batches(overwrite_batches)
     # TODO: Issue #316: save_yaml_configuration(config)
     # TODO: Issue #317: Validate ML data.
+    logger.info("Done!")
 
 
 if __name__ == "__main__":