switch chunking to auto to improve perf on load

NVIDIA · Nov 20, 2024 · 781705f · 781705f
1 parent 12c13da
commit 781705f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/modulus/datapipes/healpix/data_modules.py b/modulus/datapipes/healpix/data_modules.py
@@ -119,7 +119,7 @@ def open_time_series_dataset_classic_on_the_fly(
         file_name = _get_file_name(directory, prefix, variable, suffix)
         logger.debug("open nc dataset %s", file_name)
 
-        ds = xr.open_dataset(file_name, chunks={"sample": batch_size}, autoclose=True)
+        ds = xr.open_dataset(file_name, autoclose=True)
 
         if "LL" in prefix:
             ds = ds.rename({"lat": "height", "lon": "width"})
@@ -212,7 +212,7 @@ def open_time_series_dataset_classic_prebuilt(
     if not ds_path.exists():
         raise FileNotFoundError(f"Dataset doesn't appear to exist at {ds_path}")
 
-    result = xr.open_zarr(ds_path, chunks={"time": batch_size})
+    result = xr.open_zarr(ds_path)
     return result
 
 
@@ -286,11 +286,11 @@ def create_time_series_dataset_classic(
         file_name = _get_file_name(src_directory, prefix, variable, suffix)
         logger.debug("open nc dataset %s", file_name)
         if "sample" in list(xr.open_dataset(file_name).dims.keys()):
-            ds = xr.open_dataset(file_name, chunks={"sample": batch_size}).rename(
+            ds = xr.open_dataset(file_name).rename(
                 {"sample": "time"}
             )
         else:
-            ds = xr.open_dataset(file_name, chunks={"time": batch_size})
+            ds = xr.open_dataset(file_name)
         if "varlev" in ds.dims:
             ds = ds.isel(varlev=0)
 

diff --git a/modulus/datapipes/healpix/timeseries_dataset.py b/modulus/datapipes/healpix/timeseries_dataset.py
@@ -19,7 +19,6 @@
 import warnings
 from dataclasses import dataclass
 from typing import Optional, Sequence, Union
-import warnings
 
 import numpy as np
 import pandas as pd