Écriture de données multi-index xarray par blocs

Sep 15 2020

J'essaie de restructurer efficacement un grand ensemble de données multidimensionnelles. Supposons que j'ai un certain nombre d'images détectées à distance au fil du temps avec un certain nombre de bandes avec des coordonnées xy pour l'emplacement des pixels, le temps pour le temps d'acquisition de l'image et la bande pour différentes données collectées.

Dans mon cas d'utilisation, supposons que les longueurs de coordonnées xarray sont à peu près x (3000), y (3000), time (10), avec des bandes (40) de données en virgule flottante. Donc 100 Go + de données.

J'ai essayé de travailler à partir de cet exemple mais j'ai du mal à le traduire dans ce cas.

Exemple de petit jeu de données

REMARQUE: les données réelles sont beaucoup plus volumineuses que cet exemple.

import numpy as np
import dask.array as da
import xarray as xr

nrows = 100
ncols = 200
row_chunks = 50
col_chunks = 50

data = da.random.random(size=(1, nrows, ncols), chunks=(1, row_chunks, col_chunks))

def create_band(data, x, y, band_name):

    return xr.DataArray(data,
                        dims=('band', 'y', 'x'),
                        coords={'band': [band_name],
                                'y': y,
                                'x': x})

def create_coords(data, left, top, celly, cellx):
    nrows = data.shape[-2]
    ncols = data.shape[-1]
    right = left + cellx*ncols
    bottom = top - celly*nrows
    x = np.linspace(left, right, ncols) + cellx/2.0
    y = np.linspace(top, bottom, nrows) - celly/2.0
    
    return x, y

x, y = create_coords(data, 1000, 2000, 30, 30)

src = []

for time in ['t1', 't2', 't3']:

    src_t = xr.concat([create_band(data, x, y, band) for band in ['blue', 'green', 'red', 'nir']], dim='band')\
                    .expand_dims(dim='time')\
                    .assign_coords({'time': [time]})
    
    src.append(src_t)

src = xr.concat(src, dim='time')

print(src)


<xarray.DataArray 'random_sample-5840d8564d778d573dd403f27c3f47a5' (time: 3, band: 4, y: 100, x: 200)>
dask.array<concatenate, shape=(3, 4, 100, 200), dtype=float64, chunksize=(1, 1, 50, 50), chunktype=numpy.ndarray>
Coordinates:
  * x        (x) float64 1.015e+03 1.045e+03 1.075e+03 ... 6.985e+03 7.015e+03
  * band     (band) object 'blue' 'green' 'red' 'nir'
  * y        (y) float64 1.985e+03 1.955e+03 1.924e+03 ... -984.7 -1.015e+03
  * time     (time) object 't1' 't2' 't3'

Restructuré - empilé et transposé

J'ai besoin de stocker la sortie de ce qui suit:

print(src.stack(sample=('y','x','time')).T)

<xarray.DataArray 'random_sample-5840d8564d778d573dd403f27c3f47a5' (sample: 60000, band: 4)>
dask.array<transpose, shape=(60000, 4), dtype=float64, chunksize=(3600, 1), chunktype=numpy.ndarray>
Coordinates:
  * band     (band) object 'blue' 'green' 'red' 'nir'
  * sample   (sample) MultiIndex
  - y        (sample) float64 1.985e+03 1.985e+03 ... -1.015e+03 -1.015e+03
  - x        (sample) float64 1.015e+03 1.015e+03 ... 7.015e+03 7.015e+03
  - time     (sample) object 't1' 't2' 't3' 't1' 't2' ... 't3' 't1' 't2' 't3'

J'espère utiliser dask et xarray pour écrire le résultat sur le disque en morceaux, accessible pour open_mfdataset . le parquet semble être une bonne option, mais je ne sais pas comment l'écrire en morceaux (src est trop gros pour être stocké en mémoire).

@dask.delayed
def stacker(data):
   return data.stack(sample=('y','x','time')).T.to_pandas() 

stacker(src).to_parquet('out_*.parquet')

def stack_write(data):
   data.stack(sample=('y','x','time')).T.to_pandas().to_parquet('out_*.parquet')
   return None

stack_write(src)

À ce stade, j'espère juste quelques bonnes idées. Merci!

Réponses

3 dcherian Sep 30 2020 at 01:26

J'ai une solution ici (https://github.com/pydata/xarray/issues/1077#issuecomment-644803374) pour écrire des ensembles de données multi-indexés dans un fichier.

Vous devrez «encoder» manuellement le jeu de données sous une forme qui peut être écrite sous forme de netCDF. Et puis "décoder" quand vous le relisez.

import numpy as np
import pandas as pd
import xarray as xr


def encode_multiindex(ds, idxname):
    encoded = ds.reset_index(idxname)
    coords = dict(zip(ds.indexes[idxname].names, ds.indexes[idxname].levels))
    for coord in coords:
        encoded[coord] = coords[coord].values
    shape = [encoded.sizes[coord] for coord in coords]
    encoded[idxname] = np.ravel_multi_index(ds.indexes[idxname].codes, shape)
    encoded[idxname].attrs["compress"] = " ".join(ds.indexes[idxname].names)
    return encoded


def decode_to_multiindex(encoded, idxname):
    names = encoded[idxname].attrs["compress"].split(" ")
    shape = [encoded.sizes[dim] for dim in names]
    indices = np.unravel_index(encoded.landpoint.values, shape)
    arrays = [encoded[dim].values[index] for dim, index in zip(names, indices)]
    mindex = pd.MultiIndex.from_arrays(arrays)

    decoded = xr.Dataset({}, {idxname: mindex})
    for varname in encoded.data_vars:
        if idxname in encoded[varname].dims:
            decoded[varname] = (idxname, encoded[varname].values)
    return decoded

1 Rivers Nov 15 2020 at 18:20

Ce n'est pas la solution, pour le moment, mais une version de votre code, modifiée pour qu'elle soit facilement reproductible si d'autres veulent essayer de résoudre ce problème:

Le problème vient de l' stackopération ( concatenated.stack(sample=('y','x','time')). À cette étape, la mémoire continue d'augmenter et le processus est killed.

L' concatenatedobjet est un "soutenu par Dask" xarray.DataArray. On pouvait donc s'attendre à ce que l' stackopération soit faite paresseusement par Dask. Alors, pourquoi le processus killeden est-il à cette étape?

2 possibilités pour ce qui se passe ici:

L' stackopération est en fait effectuée paresseusement par Dask, mais comme les données sont très énormes, même la mémoire minimale requise pour Dask est trop
L' stackopération n'est PAS soutenue par Dask


import numpy as np
import dask.array as da
import xarray as xr
from numpy.random import RandomState

nrows = 20000
ncols = 20000
row_chunks = 500
col_chunks = 500


# Create a reproducible random numpy array
prng = RandomState(1234567890)
numpy_array = prng.rand(1, nrows, ncols)

data = da.from_array(numpy_array, chunks=(1, row_chunks, col_chunks))


def create_band(data, x, y, band_name):

    return xr.DataArray(data,
                        dims=('band', 'y', 'x'),
                        coords={'band': [band_name],
                                'y': y,
                                'x': x})

def create_coords(data, left, top, celly, cellx):
    nrows = data.shape[-2]
    ncols = data.shape[-1]
    right = left + cellx*ncols
    bottom = top - celly*nrows
    x = np.linspace(left, right, ncols) + cellx/2.0
    y = np.linspace(top, bottom, nrows) - celly/2.0
    
    return x, y


x, y = create_coords(data, 1000, 2000, 30, 30)

bands = ['blue', 'green', 'red', 'nir']
times = ['t1', 't2', 't3']
bands_list = [create_band(data, x, y, band) for band in bands]

src = []

for time in times:

    src_t = xr.concat(bands_list, dim='band')\
                    .expand_dims(dim='time')\
                    .assign_coords({'time': [time]})

    src.append(src_t)


concatenated = xr.concat(src, dim='time')
print(concatenated)
# computed = concatenated.compute() # "computed" is ~35.8GB

stacked = concatenated.stack(sample=('y','x','time'))

transposed = stacked.T

On peut essayer de changer les valeurs de nrowset ncolsafin de faire varier la taille de concatenated. Et pour les performances, nous pourrions / devrions aussi varier chunks.

Remarque: j'ai même essayé ceci

concatenated.to_netcdf("concatenated.nc")
concatenated = xr.open_dataarray("concatenated.nc", chunks=10)

Ceci afin d'être sûr qu'il s'agit d'un DataArray soutenu par Dask et de pouvoir également ajuster les morceaux. J'ai essayé différentes valeurs pour chunks: mais toujours à court de mémoire.