blog/gpu-pipeline

[giscus[bot]]

blog/gpu-pipeline

How to accelerate AI/ML workflows in Earth Sciences with GPU-native Xarray and Zarr.

https://xarray.dev/blog/gpu-pipeline

[VolkaneicFox]

Hey,

First of all, really amazing directions! This is majorly helpful.

However, I have always had problems with chunking, especially along time, since with daily data, I typically have to interpolate. I understand this is very important, that is chunking, when the netcdf is global, but I am currently working with Brazilian ERA5. I found it better to batch-wise load and save the data as .zarrr. I have about 13000 daily netcdfs of Brazil, and I did this code. Stored it all as individual zarr files, and then work on that.

The total time to save to zarr took me 3mins46sec

import glob, xarray, os, tqdm, zarr

all_paths = sorted(glob.glob("/home/darkmage/paper_three/vpd_netcdfs///*.nc"))
output_dir = "/home/darkmage/paper_three/vpd_processed_batches"
os.makedirs(output_dir, exist_ok=True)

batch_size = 1000
total_files = len(all_paths)

for i in tqdm.tqdm(range(0, total_files, batch_size)):
batch_paths = all_paths[i:i + batch_size]
batch_num = i // batch_size + 1
out_file = f"{output_dir}/brazil_vpd_batch_{batch_num:02d}.nc"

if os.path.exists(out_file):
    print(f"Batch {batch_num} already processed. Skipping...")
    continue
    
print(f"Processing Batch {batch_num} ({len(batch_paths)} files)...")

brazil_ds = xarray.open_mfdataset(batch_paths, combine="by_coords", parallel=True)
brazil_ds = brazil_ds.rename({"Vapour_Pressure_Deficit_at_Maximum_Temperature": "vpd_max"})
brazil_ds = brazil_ds.load()
brazil_ds["vpd_max"] = brazil_ds["vpd_max"].interpolate_na(dim="time", method="linear")


brazil_ds.to_zarr(f"/home/darkmage/paper_three/brazil_vpd_batch{batch_num}.zarr", mode="w",consolidated=True)