Menulis data multiindex xarray dalam potongan

Sep 15 2020

Saya mencoba merestrukturisasi kumpulan data multidimentional yang besar secara efisien. Anggaplah saya memiliki sejumlah gambar penginderaan jauh dari waktu ke waktu dengan sejumlah pita dengan koordinat xy untuk lokasi piksel, waktu untuk waktu akuisisi citra, dan pita untuk data berbeda yang dikumpulkan.

Dalam kasus penggunaan saya, mari kita asumsikan panjang coord xarray kira-kira x (3000), y (3000), waktu (10), dengan band (40) data floating point. Jadi 100 GB + data.

Saya telah mencoba bekerja dari contoh ini tetapi saya mengalami masalah dalam menerjemahkannya ke kasus ini.

Contoh set data kecil

CATATAN: data sebenarnya jauh lebih besar dari contoh ini.

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'

Direstrukturisasi - ditumpuk dan dialihkan

Saya perlu menyimpan keluaran berikut ini:

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'

Saya berharap untuk menggunakan dask dan xarray untuk menulis hasilnya ke disk dalam bentuk potongan, dapat diakses oleh open_mfdataset . parket sepertinya pilihan yang bagus, tapi saya tidak tahu bagaimana menulisnya dalam potongan (src terlalu besar untuk disimpan dalam memori).

@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)

Saat ini saya hanya mengharapkan beberapa ide bagus. Terima kasih!

Jawaban

3 dcherian Sep 30 2020 at 01:26

Saya punya solusi di sini (https://github.com/pydata/xarray/issues/1077#issuecomment-644803374) untuk menulis kumpulan data multiindeks ke file.

Anda harus secara manual "menyandikan" kumpulan data ke dalam bentuk yang dapat ditulis sebagai netCDF. Dan kemudian "memecahkan kode" saat Anda membacanya kembali.

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

Untuk saat ini, ini bukanlah solusi, tetapi versi kode Anda, yang dimodifikasi sehingga dapat dengan mudah direproduksi jika orang lain ingin mencoba memecahkan masalah ini:

Masalahnya ada pada stackoperation ( concatenated.stack(sample=('y','x','time')). Pada langkah ini, memori terus bertambah dan prosesnya berjalan killed.

The concatenatedobjek adalah "DASK didukung" xarray.DataArray. Jadi kita bisa berharap stackoperasi dilakukan dengan malas oleh Dask. Jadi, mengapa proses killedpada langkah ini?

2 kemungkinan untuk apa yang terjadi di sini:

The stackoperasi sebenarnya dilakukan malas oleh DASK, tetapi karena data yang sangat yang besar, bahkan minimum yang diperlukan memori untuk DASK terlalu banyak
The stackoperasi TIDAK DASK yang didukung


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

Seseorang dapat mencoba mengubah nilai nrowsdan ncolsuntuk memvariasikan ukuran concatenated. Dan untuk kinerja, kami juga dapat / harus bervariasi chunks.

Catatan: Saya bahkan mencoba ini

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

Ini untuk memastikan itu adalah DataArray yang didukung Dask dan untuk dapat menyesuaikan potongan juga. Saya mencoba nilai yang berbeda untuk chunks: tetapi selalu kehabisan memori.