Creating huge xarray datasets

[muaali]

Hi.

I have a lot of satellite images in my local directory. I want to create an xarray dataset from n satellite images (each satellite image is huge in memory). Normally, one would like to iteratively create xarray data array objects, concatenate them and then convert to xarray dataset. But this approach leads to memory issues as xr.concat forces dask arrays to load (thread: #4628)

What would be the right approach to solve this issue with xarrays and dask ?

Few things to note:

• Images are of extension .tif
• Images are not of the same size
• Each image can be up to 2 GB and I would like to scale to as many images as possible
• using xr.open_rasterio leads to dataset with max image (image with the maximum size) repeated n times where n is the number of images. I will have to resize each image to save size and then save back, which is a tedious work and does not seem an efficient way to accomplish this task.

Below is my code on high level, but it I still run into memory issues with even after chunking with dask

xdarrays_list = []


def create_xdarray_with_raster(fpath):
    rasterio_container = rasterio.open(fpath)
    Iarray = np.squeeze(rasterio_container.read()) 
    xdarray = xr.DataArray(Iarray, coords=[np.arange(Iarray.shape[0]), np.arange(Iarray.shape[1])], dims=["x", "y"])
    xdarray.attrs = rasterio_container.tags()

    return xdarray


for fpath in fpaths:
    raster_xdarray = create_xdarray_with_raster(raster_fpath).chunk(4096,4096)
    assert(dask.is_dask_collection(raster_xdarray))
    xdarrays_list.append(raster_xdarray)

xr.Dataset({"values": xr.concat(xdarrays_list, dim="band", data_vars="all", fill_value=-1)},
            coords={
                "latitude": list_of_lats, #of len same as number of images
                "longitude": list_of_longs, #of len same as number of images
            },
        )

[keewis]

that thread is about lazy concatenation without using dask, with dask it should work just fine:

In [10]: import xarray as xr
    ...: from xarray.tests import raise_if_dask_computes
    ...: 
    ...: ds = xr.tutorial.open_dataset("rasm").chunk({"time": 8, "x": 5, "y": 5})
    ...: 
    ...: with raise_if_dask_computes(max_computes=0):
    ...:     ds1 = ds.isel(time=slice(0, 18))
    ...:     ds2 = ds.isel(time=slice(18, None))
    ...: 
    ...:     concat = xr.concat([ds1, ds2], dim="time")
    ...: 
    ...: xr.testing.assert_identical(ds, concat)

In [11]:

Instead of doing this manually using a loop you should also be able to use open_mfdataset (if you need to use rasterio to load your data this requires xarray>=0.18 and a recent version of rioxarray).

[muaali]

Alright. so with xarray updated to 0.18.2 and rioxarray installed. This line works :

ds1 = xr.open_mfdataset(tif_fpaths, engine="rasterio")

But facing this issue where I only get one 1 band in ds1 where I will expect number of bands to be: num_bands=len(tif_fpaths)
Same issue takes place when I try to open_mfdataset all saved netcdf4 files from a directory ds2 = xr.open_mfdataset(nc_fpaths, combine='by_coords', engine='netcdf4'). It only returns me ds2 with one band.

Any idea why is that so ?

[muaali]

@max-sixty I would be happy to have take a look at that code. Might be helpful.

[snowman2]

But facing this issue where I only get one 1 band in ds1 where I will expect number of bands to be: num_bands=len(tif_fpaths)

Is there a new dimension/coordinate added?

[muaali]

@snowman2 All of images don't have the same size so that leads to xr.DataArray of different sizes. Could that cause this issue ?

[snowman2]

@muaali you will want to ensure they are all the same projection/resolution if you want to stack them. This can help you prepare them: https://corteva.github.io/rioxarray/stable/examples/reproject_match.html

[muaali]

@keewis so I am able to create a 21 GB of .nc with the following code. Here xr.concat does not force dask arrays to load.

xarray_list = []
for tif_fpath in tif_fpaths:
    da = xr.open_rasterio(tif_fpath)
    da = da.chunk(4096,4096) # arbitrary chunk size
    assert(dask.is_dask_collection(da))
    xarray_list.append(da)

ds = xr.Dataset(
            {"intensity": xr.concat(xarray_list, dim="band", data_vars="all", fill_value=-1)},
            coords={
                "latitude": list_of_lats, #derived from metadata
                "longitude": list_of_longs, #derived from metadata
            },
        )
ds.to_netcdf('huge_cube.nc'))

But what I don't like about this solution is I am heavily dependent on default functionalities of xr.open_rasterio. So if I would like to perform some preprocessing or create another dataset variable (apart from intensity) based on some external information, I would be back to square one. So I am just thinking how to modify following function to achieve that (and I could replicate this for others):

def create_xdarray_with_raster(fpath, load_metadata=True):
    rasterio_container = rasterio.open(fpath)
    Iarray = np.squeeze(rasterio_container.read())  # <- This right here is the issue as you pointed out.
    xdarray = xr.DataArray(Iarray, coords=[np.arange(Iarray.shape[0]), np.arange(Iarray.shape[1])], dims=["x", "y"])
    xdarray.attrs = rasterio_container.tags()

    return xdarray

[muaali]

I'm not really sure I understand you OP correctly. But looking at the lat/lon coordinates, the grids are not the same. You would have to keep that in mind when working with the data (selecting data etc).

So the way I have approached it is that for each image, there is coordinate.latitude , coordinate.longitude, coordinate.incidence_angle. The idea is to perform slicing along 3rd dimension for images that are at certain locations and within certain angles. Do you think it is a wrong approach. What could be an alternative ?

[muaali]

Did you specify chunk-sizes in the call to open_mfdataset? If so, which?

all_ds = xr.open_mfdataset(nc_fpaths, concat_dim="new_dim", combine='nested', chunks={"x": 4096, "y": 4096}, engine='netcdf4')

[kmuehlbauer]

all_ds = xr.open_mfdataset(nc_fpaths, concat_dim="new_dim", combine='nested', chunks={"x": 4096, "y": 4096}, engine='netcdf4')

Just for the sake of testing, can you drop/delete the coordinates (lat/lon/incident_angle) before exporting to netcdf?

[kmuehlbauer]

I'm not really sure I understand you OP correctly. But looking at the lat/lon coordinates, the grids are not the same. You would have to keep that in mind when working with the data (selecting data etc).

So the way I have approached it is that for each image, there is coordinate.latitude , coordinate.longitude, coordinate.incidence_angle. The idea is to perform slicing along 3rd dimension for images that are at certain locations and within certain angles. Do you think it is a wrong approach. What could be an alternative ?

Unfortunately my data sources are fixed and don't move around that much (ground weather radar). But some general notes: You might be able to slice along the third dimension, but still you have to keep in mind, that your x,y coordinates are only identical if all three of your coordinates are the same (for the different slices). Also possible, that slicing isn't possible at all, since the data is not properly ordered. I'm sorry, but this as far as I can help here.

[muaali]

Thanks alot @kmuehlbauer for your help so far. I tried deleting lats/long/angles but still memory chokes after a while.