memory error when building fieldset from zarr stores via xarray

[selipot]

I am using ocean parcels version 2.4.1.dev9 (because this is what I have set up on my HPC). I am trying something new which is to advect some particles in a fieldset built from zarr stores (instead of NetCDF files as previously). But I am getting a memory error when initializing the fieldset (see below) because the job is not able to allocate the necessary memory, see below. My previous code used fieldset = FieldSet.from_netcdf with all the NetCDF files as input and the option deferred_load=True. So what's the strategy when using huge zarr stores opened as xarray datasets to build the fieldset?

numpy.core._exceptions.MemoryError: Unable to allocate 341. GiB for an array with shape (720, 2, 7055, 9000) and data type float32

Code snippet that fails at fieldset initialization :

# open datasets
ds1 = xr.open_dataset(uvfiles,engine='zarr')
ds1 = ds1.assign_coords({"Latitude":ds1.Latitude,"Longitude":ds1.Longitude})

ds2 = xr.open_dataset(sshfiles,engine='zarr')
ds2 = ds2.assign_coords({"Latitude":ds2.Latitude,"Longitude":ds2.Longitude})

ds3 = xr.open_dataset(bathyfile,engine='zarr')
ds3 = ds3.assign_coords({"Latitude":ds3.Latitude,"Longitude":ds3.Longitude})

fieldset = FieldSet.from_xarray_dataset(ds1,
                                        variables1,
                                        dimensions1,
                                        indices=indices,
                                        mesh='spherical',
                                        allow_time_extrapolation=True,
                                        deferred_load=True,
                                       )

[erikvansebille]

Hi @selipot, I don't have that much experience with creating field sets from (very large) zarr stores; but it seems that the deferred_load option is not respected by from_xarray_dataset. Otherwise, Parcels wouldn't need such a large array for the data.

So I'm sorry, but I can't help you here. This is not something that we've tested or planned for. Is it an option to go back to the netcdf version of creating the FieldSet?

[selipot]

Not really. Creating zarr stores for the velocity and ssh data allowed me to reduce the model dataset from 28TB to 5.6TB so I moved to a slow storage the original NetCDF files ...