Merge pull request #662 from openclimatefix/issue/save-batch

add option to save batches to differnt location for manager live

Author: peterdudfield

nowcasting_dataset/manager/manager_live.py

@@ -19,6 +19,7 @@
     SpaceTimeLocation,
     load_from_csv,
 )
+from nowcasting_dataset.dataset.batch import Batch
 from nowcasting_dataset.filesystem import utils as nd_fs_utils
 from nowcasting_dataset.manager.base import ManagerBase
 from nowcasting_dataset.manager.utils import callback, error_callback
@@ -251,3 +252,23 @@ def create_batches(self, use_async: Optional[bool] = True) -> None:
                             )
 
             logger.info(f"Finished creating batches for {split_name}!")
+
+    def save_batch(self, batch_idx, path: str):
+        """
+        Option to save batch to a differnt location
+
+        Args:
+            path: the path to save the file to
+            batch_idx: the batch index
+
+        """
+
+        # load batch
+        batch = Batch.load_netcdf(
+            batch_idx=batch_idx,
+            data_sources_names=self.data_sources,
+            local_netcdf_path=self.config.output_data.filepath / "live",
+        )
+
+        # save batch
+        batch.save_netcdf(batch_i=batch_idx, path=path)

tests/manager/test_manager_live.py

@@ -199,3 +199,7 @@ def test_run(test_configuration_filename):
             t0_datetime=datetime(2020, 4, 1, 15)
         )  # noqa 101
         manager.create_batches()
+
+        with tempfile.TemporaryDirectory() as local_temp_path_save:
+            manager.save_batch(batch_idx=0, path=local_temp_path_save)
+            assert os.path.exists(f"{local_temp_path_save}/gsp/000000.nc")