Speedup datacube loading if persist=True by factor 10

darts-acquisition/src/darts_acquisition/arcticdem/datacube.py

@@ -317,21 +317,22 @@ def load_arcticdem_tile(
     procedural_download_datacube(storage, adjacent_tiles)
 
     # Load the datacube and set the spatial_ref since it is set as a coordinate within the zarr format
-    arcticdem_datacube = xr.open_zarr(storage, mask_and_scale=False).set_coords("spatial_ref")
+    chunks = None if persist else "auto"
+    arcticdem_datacube = xr.open_zarr(storage, mask_and_scale=False, chunks=chunks).set_coords("spatial_ref")
 
     # Get an AOI slice of the datacube
     arcticdem_aoi = arcticdem_datacube.odc.crop(reference_geobox.extent, apply_mask=False)
 
-    # Change dtype of the datamask to uint8 for later reproject_match
-    arcticdem_aoi["datamask"] = arcticdem_aoi.datamask.astype("uint8")
-
-    # The following code would load the data from disk
+    # The following code would load the lazy zarr data from disk into memory
     if persist:
         tick_sload = time.perf_counter()
-        arcticdem_aoi = arcticdem_aoi.compute()
+        arcticdem_aoi = arcticdem_aoi.load()
         tick_eload = time.perf_counter()
         logger.debug(f"ArcticDEM AOI loaded from disk in {tick_eload - tick_sload:.2f} seconds")
 
+    # Change dtype of the datamask to uint8 for later reproject_match
+    arcticdem_aoi["datamask"] = arcticdem_aoi.datamask.astype("uint8")
+
     logger.info(
         f"ArcticDEM tile {'loaded' if persist else 'lazy-opened'} in {time.perf_counter() - tick_fstart:.2f} seconds"
     )

darts-acquisition/src/darts_acquisition/tcvis.py

@@ -198,15 +198,16 @@ def load_tcvis(
     procedural_download_datacube(storage, reference_geobox)
 
     # Load the datacube and set the spatial_ref since it is set as a coordinate within the zarr format
-    tcvis_datacube = xr.open_zarr(storage, mask_and_scale=False).set_coords("spatial_ref")
+    chunks = None if persist else "auto"
+    tcvis_datacube = xr.open_zarr(storage, mask_and_scale=False, chunks=chunks).set_coords("spatial_ref")
 
     # Get an AOI slice of the datacube
     tcvis_aoi = tcvis_datacube.odc.crop(reference_geobox.extent, apply_mask=False)
 
-    # The following code would load the data from disk
+    # The following code would load the lazy zarr data from disk into memory
     if persist:
         tick_sload = time.perf_counter()
-        tcvis_aoi = tcvis_aoi.compute()
+        tcvis_aoi = tcvis_aoi.load()
         tick_eload = time.perf_counter()
         logger.debug(f"TCVIS AOI loaded from disk in {tick_eload - tick_sload:.2f} seconds")
 

darts-preprocessing/src/darts_preprocessing/preprocess.py

@@ -161,13 +161,21 @@ def preprocess_legacy_fast(
     ds_merged["ndvi"] = calculate_ndvi(ds_merged).ndvi
 
     # Reproject TCVIS to optical data
+    tick_sproj = time.perf_counter()
     ds_tcvis = ds_tcvis.odc.reproject(ds_optical.odc.geobox, resampling="cubic")
+    tick_eproj = time.perf_counter()
+    logger.debug(f"Reprojection of TCVIS done in {tick_eproj - tick_sproj:.2f} seconds.")
+
     ds_merged["tc_brightness"] = ds_tcvis.tc_brightness
     ds_merged["tc_greenness"] = ds_tcvis.tc_greenness
     ds_merged["tc_wetness"] = ds_tcvis.tc_wetness
 
     # Calculate TPI and slope from ArcticDEM
+    tick_sproj = time.perf_counter()
     ds_arcticdem = ds_arcticdem.odc.reproject(ds_optical.odc.geobox.buffered(tpi_outer_radius), resampling="cubic")
+    tick_eproj = time.perf_counter()
+    logger.debug(f"Reprojection of ArcticDEM done in {tick_eproj - tick_sproj:.2f} seconds.")
+
     ds_arcticdem = preprocess_legacy_arcticdem_fast(ds_arcticdem, tpi_outer_radius, tpi_inner_radius, device)
     ds_arcticdem = ds_arcticdem.odc.crop(ds_optical.odc.geobox.extent)
     ds_merged["dem"] = ds_arcticdem.dem