Add Local Filesystem Support (#142)

* Add local option

* Add testing local option

* Bump version

* Update test to check for example data

Author: jacobbieker

nowcasting_dataset/dataset/datasets.py

@@ -159,15 +159,15 @@ def __init__(
         self.gcs = None
         self.s3_resource = None
 
-        assert cloud in ["gcp", "aws"]
+        assert cloud in ["gcp", "aws", "local"]
 
         if not os.path.isdir(self.tmp_path):
             os.mkdir(self.tmp_path)
 
     def per_worker_init(self, worker_id: int):
         if self.cloud == "gcp":
             self.gcs = gcsfs.GCSFileSystem()
-        else:
+        elif self.cloud == "aws":
             self.s3_resource = boto3.resource("s3")
 
     def __len__(self):
@@ -198,15 +198,18 @@ def __getitem__(self, batch_idx: int) -> example.Example:
                 local_filename=local_netcdf_filename,
                 gcs=self.gcs,
             )
-        else:
+        elif self.cloud == "aws":
             aws_download_to_local(
                 remote_filename=remote_netcdf_filename,
                 local_filename=local_netcdf_filename,
                 s3_resource=self.s3_resource,
             )
+        else:
+            local_netcdf_filename = remote_netcdf_filename
 
         netcdf_batch = xr.load_dataset(local_netcdf_filename)
-        os.remove(local_netcdf_filename)
+        if self.cloud != "local":
+            os.remove(local_netcdf_filename)
 
         batch = example.Example(
             sat_datetime_index=netcdf_batch.sat_time_coords,

setup.py

@@ -9,7 +9,7 @@
 
 setup(
     name="nowcasting_dataset",
-    version="0.1.4",
+    version="0.1.5",
     license="MIT",
     description="Nowcasting Dataset",
     author="Jack Kelly, Peter Dudfield, Jacob Bieker",

tests/data/test.nc renamed to tests/data/0.nc

@@ -24,7 +24,7 @@
 
 
 def test_subselect_date():
-    dataset = xr.open_dataset("tests/data/test.nc")
+    dataset = xr.open_dataset("tests/data/0.nc")
     x = example.Example(
         sat_data=dataset["sat_data"],
         nwp=dataset["nwp"],
@@ -44,6 +44,48 @@ def test_subselect_date():
     assert batch[NWP_DATA].shape[2] == 5
 
 
+def test_netcdf_dataset_local():
+    DATA_PATH = "tests/data"
+    TEMP_PATH = "tests/data/temp/"
+
+    train_dataset = NetCDFDataset(
+        1,
+        DATA_PATH,
+        TEMP_PATH,
+        cloud="local",
+        history_minutes=10,
+        forecast_minutes=10,
+        current_timestep_index=7,
+        required_keys=(NWP_DATA, NWP_TARGET_TIME, SATELLITE_DATA, SATELLITE_DATETIME_INDEX),
+    )
+
+    dataloader_config = dict(
+        pin_memory=True,
+        num_workers=1,
+        prefetch_factor=1,
+        worker_init_fn=worker_init_fn,
+        persistent_workers=True,
+        # Disable automatic batching because dataset
+        # returns complete batches.
+        batch_size=None,
+    )
+
+    _ = torch.utils.data.DataLoader(train_dataset, **dataloader_config)
+
+    train_dataset.per_worker_init(1)
+    t = iter(train_dataset)
+    data = next(t)
+
+    sat_data = data[SATELLITE_DATA]
+
+    # Sat is in 5min increments, so should have 2 history + current + 2 future
+    assert sat_data.shape[1] == 5
+    assert data[NWP_DATA].shape[2] == 5
+
+    # Make sure file isn't deleted!
+    assert os.path.exists("tests/data/0.nc")
+
+
 @pytest.mark.skip("CD does not have access to GCS")
 def test_get_dataloaders_gcp():
     DATA_PATH = "gs://solar-pv-nowcasting-data/prepared_ML_training_data/v5/"