minor fixes

stresearch · Jun 8, 2022 · 98b0f42 · 98b0f42
1 parent 59be221
commit 98b0f42
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/create_dataset.py b/create_dataset.py
@@ -6,4 +6,9 @@
 
 if __name__=="__main__":
     NCDataConstructor.default_data(split="train",  workers  =32, prefix=cam4, train_years=3, save_location=".", cache = ".")
+    NCDataConstructor.default_data(split="test",  workers  =32, prefix=cam4, train_years=3, save_location=".", cache = ".")
+    NCDataConstructor.default_data(split="train",  workers  =32, prefix=spcam, train_years=2, save_location=".", cache = ".")
+    NCDataConstructor.default_data(split="test",  workers  =32, prefix=spcam, train_years=2, save_location=".", cache = ".")
+
+
     # NCDataConstructor.default_data(split="train")
diff --git a/gaia/data.py b/gaia/data.py
@@ -494,6 +494,7 @@ def default_data(
             ),
             outputs="PRECT,PRECC,PTEQ,PTTEND".split(","),
             flatten=split == "train",
+            shuffle = split == "train",
             subsample_factor=4,
             compute_stats=True,
             cache = os.path.join(cache,split),
@@ -689,7 +690,11 @@ def clean_up_file(self, dataset):
     def subsample_data(self, xi, yi, subsample_factor):
         size = xi.shape[0]
         new_size = size // subsample_factor
-        shuffled_index = torch.randperm(size)[:new_size]
+        if self.shuffle:
+            shuffled_index = torch.randperm(size)[:new_size]
+        else:
+            shuffled_index = torch.arange(0, size, subsample_factor)
+
         xi = xi[shuffled_index, ...]
         yi = yi[shuffled_index, ...]
         return xi, yi, shuffled_index