Adding built-in datasets class

PARCtorch/PARCv2.py

@@ -1,6 +1,6 @@
 import torch
 
-from PARCtorch.utilities.unet import UNet
+from PARCtorch.models.unet import UNet
 from PARCtorch.differentiator.finitedifference import FiniteDifference
 from differentiator.differentiator import ADRDifferentiator
 from integrator.rk4 import RK4

PARCtorch/data/dataset.py

@@ -7,7 +7,7 @@
 from tqdm import tqdm
 from the_well.data import WellDataset
 
-# TODO: Wipe out this entire file. Instead, implement the following classes:
+# TODO: Wipe out this entire file. Instead, implement the following classes under PARCtorch/datasets/dataset.py:
 # class Dataset(torch.utils.data.Dataset):
 # class RegularMeshDataset(Dataset):
 # class IrregularMeshDataset(Dataset):

PARCtorch/datasets/navier_stokes.py

@@ -0,0 +1,50 @@
+from PARCtorch.data.dataset import GenericPhysicsDataset
+from PARCtorch.utils.common import CACHE_DIR
+
+from PARCtorch.datasets.utils import download, extract_zip
+
+import os
+
+
+class NavierStokes(GenericPhysicsDataset):
+
+    """
+    Navier-Stokes data set. TODO: More details
+
+    Args:
+        split: Expected data split. Can be `train`, `test` TODO
+        data_dir: Directory to read/write data. Defaults to None, in which case,
+                  data will be stored in `CACHE_DIR/datasets/NavierStokes`.
+                  If the data does not exist in the specified directory,
+                  it will be automatically downloaded.
+        future_steps: Number of timesteps in the future the model will predict.
+                      Must be between 1 (single step prediction) and TODO (default).
+    """
+
+    url = "https://zenodo.org/records/13909869/files/NavierStokes.zip?download=1"
+
+    def __init__(
+        self, split=None, data_dir=None, future_steps=2,
+    ):
+        super().__init__()
+
+        # Set up data directory
+        if data_dir is None:
+            data_dir = os.path.join(CACHE_DIR, 'datasets', 'NavierStokes')
+        self.data_dir = data_dir
+        self.zip_dir = os.path.join(self.data_dir, 'NavierStokes.zip')
+
+        if not os.path.exists(self.data_dir):
+            os.makedirs(self.data_dir)
+
+        # Download and unzip data if it doesn't exist already.
+        self.download(force=False)
+
+        from pathlib import Path
+        filelist = Path(self.data_dir).glob('*.*')
+
+
+    def download(self, force=False):
+        if not os.path.exists(self.zip_dir):
+            download(self.url, self.zip_dir)
+            extract_zip(self.zip_dir, self.data_dir)

PARCtorch/datasets/utils.py

@@ -0,0 +1,46 @@
+import os
+import requests
+import zipfile
+from io import BytesIO
+from tqdm import tqdm
+from pathlib import Path
+
+
+def download(url, target_path, chunk_size=1024):
+    """
+    Downloads a file from the specified URL and saves it to the target_path.
+    """
+    print(f"Downloading a file from {url}")
+    if os.path.exists(target_path):
+        print(f"Found a cached data in {target_path}. Aborting the download.")
+        print(
+            "If you want to re-download the data, please remove the cached data and try again."
+        )
+        return target_path
+
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+
+    total_size = int(response.headers.get("content-length", 0))
+    progress = tqdm(total=total_size, unit="B", unit_scale=True, desc=f"Downloading to {target_path}")
+
+    with open(target_path, "wb") as f:
+        for chunk in response.iter_content(chunk_size=chunk_size):
+            if chunk:  # filter out keep-alive chunks
+                f.write(chunk)
+                progress.update(len(chunk))
+    progress.close()
+
+    print(f"Download complete. File saved to {target_path}")
+    return target_path
+
+
+def extract_zip(zip_path, target_folder):
+    """
+    Extracts the contents of a ZIP file at zip_path into the target folder.
+    """
+    Path(target_folder).mkdir(parents=True, exist_ok=True)
+    print(f"Extracting contents to {target_folder} ...")
+    with zipfile.ZipFile(zip_path, "r") as zip_ref:
+        zip_ref.extractall(target_folder)
+    print("Extraction complete.")

PARCtorch/differentiator/mappingandrecon.py

@@ -3,8 +3,8 @@
 import torch.nn as nn
 
 # Import custom utilities
-from PARCtorch.utilities.spade import SPADEGeneratorUnit
-from PARCtorch.utilities.resnet import ResNet
+from PARCtorch.models.spade import SPADEGeneratorUnit
+from PARCtorch.models.resnet import ResNet
 
 
 class MappingAndRecon(nn.Module):

PARCtorch/integrator/datadrivenintegrator.py

@@ -1,7 +1,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
-from PARCtorch.utilities.spade import SPADEGeneratorUnit
-from PARCtorch.utilities.resnet import ResNet
+from PARCtorch.models.spade import SPADEGeneratorUnit
+from PARCtorch.models.resnet import ResNet
 
 
 class DataDrivenIntegrator(nn.Module):

PARCtorch/integrator/poisson.py

@@ -3,7 +3,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from PARCtorch.utilities.resnet import ResNet
+from PARCtorch.models.resnet import ResNet
 
 
 class Poisson(nn.Module):

PARCtorch/scripts/ns_slurm.py

@@ -39,13 +39,13 @@
 # ---------------------------
 from data.normalization import compute_min_max
 from data.dataset import GenericPhysicsDataset, custom_collate_fn
-from utilities.viz import visualize_channels
+from PARCtorch.utils.viz import visualize_channels
 from PARCtorch.PARCv2 import PARCv2
 from PARCtorch.differentiator.differentiator import Differentiator
 from PARCtorch.differentiator.finitedifference import FiniteDifference
 from PARCtorch.integrator.integrator import Integrator
 from PARCtorch.integrator.heun import Heun
-from PARCtorch.utilities.unet import UNet
+from PARCtorch.models.unet import UNet
 
 
 # ---------------------------

PARCtorch/utils/common.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+
+
+CACHE_DIR = Path.home() / ".parc"