Как обрабатывать большие наборы данных xarray/dask, чтобы минимизировать время вычислений или нехватку памяти для хранения в виде ежегодных файлов (набор данных ERA5)

В настоящее время я использую данные ERA5 для расчета переменных, связанных с ветром. Хотя я могу вычислить то, что хочу, у меня возникают проблемы с эффективной реализацией, позволяющей переносить эти тяжелые данные реальным способом. Эти данные хранятся на «суперкомпьютере», к которому я могу получить доступ и запустить свой сценарий. Доступны несколько вариантов ЦП и памяти, но большую часть времени я использую либо 4 ЦП/18 ГБ, либо 7 ЦП/32 ГБ. Данные поступают в ежемесячные файлы, которые я открываю как годовые. Общее количество лет — 74, но я решил, что лучше подходить из года в год.

Вот мой код:

# Imports
import geopandas as gpd
from osgeo import gdal
import numpy as np
import matplotlib.pyplot as plt
import glob
import pickle
import xarray as xr
import time 

# setting up my workers (e.g. 7)
from dask.distributed import Client
client = Client(n_workers=7)
client

# Declaring variables
directory = 'PATH_TO_LOCAL_DIR'

years  = np.arange(1950,2023+1)

lon0 =  129.7
lat0 = - 27.9
lonbounds = [113.0,145.0]
latbounds = [-36.0,-16.0]

d     = 300e-6
z     = 10
g     = 9.8
z0    = 1e-2
rho_s = 2650
rho_f = 1.2
kappa = 0.4
phi = 0.6
fluxconstant      = 5
thresholdconstant = 0.082
usimths = np.linspace(0.01,0.8,80,dtype=np.float32)
r = 6371.229*1e+3 
s2yr = 1/60/60/24/365

uabinwidth = 3*np.pi/180
uabine = np.arange(-180*np.pi/180,(180+1e-4)*np.pi/180,uabinwidth)
uabinm = uabine[1:]-uabinwidth/2


# Opening the data. Currently I just grabe the first year worth of data
era5  = xr.open_mfdataset(glob.glob('u10_era5-land_oper_sfc_%d*.nc'%(years[i],years[i])))
era5  = era5.reindex(latitude=list(reversed(era5.latitude)))
era5  = era5.sel(longitude=slice(lonbounds[0],lonbounds[-1]),latitude=slice(latbounds[0],latbounds[-1]))       
era5v = xr.open_mfdataset(glob.glob('v10_era5-land_oper_sfc_%d*.nc'%(years[i],years[i])))
era5v = era5v.reindex(latitude=list(reversed(era5v.latitude)))
era5v = era5v.sel(longitude=slice(lonbounds[0],lonbounds[-1]),latitude=slice(latbounds[0],latbounds[-1]))
era5  = era5.merge(era5v)

Результат работыera5 на этом этапе:

 <xarray.Dataset>
 Dimensions:    (longitude: 321, latitude: 201, time: 8759)
 Coordinates:
   * longitude  (longitude) float32 113.0 113.1 113.2 113.3 ... 144.8 144.9 145.0
   * latitude   (latitude) float32 -36.0 -35.9 -35.8 -35.7 ... -16.2 -16.1 -16.0
   * time       (time) datetime64[ns] 1950-01-01T01:00:00 ... 1950-12-31T23:00:00
 Data variables:
     u10        (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     v10        (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>

Дальше я делаю:

# Calculation 
era5 = era5.assign(ua_from=np.arctan2(-era5.v10,-era5.u10))#.compute() <--- not sure if compute in between is a good idea?

for j in range(len(usimths)):
    era5 = era5.assign(temp=(((era5.u10**2+era5.v10**2)*(kappa/np.log(z/z0))**2-usimths[j]**2)*fluxconstant*usimths[j]/g*rho_f/rho_s).where(((era5.u10**2+era5.v10**2)**0.5*(kappa/np.log(z/z0)))>usimths[j],0))#.compute()
    era5['qm_{}'.format(str(round(usimths[j],2)))] = era5.temp
    era5 = era5.drop_vars('temp')

Результат работыera5:

 Output of era5:
 <xarray.Dataset>
 Dimensions:    (longitude: 321, latitude: 201, time: 8759)
 Coordinates:
   * longitude  (longitude) float32 113.0 113.1 113.2 113.3 ... 144.8 144.9 145.0
   * latitude   (latitude) float32 -36.0 -35.9 -35.8 -35.7 ... -16.2 -16.1 -16.0
   * time       (time) datetime64[ns] 1950-01-01T01:00:00 ... 1950-12-31T23:00:00
 Data variables: (12/83)
     u10        (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     v10        (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     ua_from    (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     qm_0.01    (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     qm_0.02    (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     qm_0.03    (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     ...         ...
     qm_0.75    (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     qm_0.76    (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     qm_0.77    (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     qm_0.78    (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     qm_0.79    (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>
     qm_0.8     (time, latitude, longitude) float32 dask.array<chunksize=(52, 100, 166), meta=np.ndarray>

Окончательно:

# this part takes some time and I am still unsure if I should compute or not in between
ua_from = era5.ua_from.compute().values #<-- need these for below
bin_arrays_mean = []
bin_arrays_sum = []
for ident in range(len(uabine)-1):
    era5tempid = era5.drop_vars(["u10","v10","ua_from"])
    era5tempid = era5tempid.where((ua_from >= uabine[ident]) & (ua_from < uabine[ident+1]) & (era5tempid > 0))
    bin_arrays_mean.append(era5tempid.mean(dim = "time", skipna=True))#.compute())
    bin_arrays_sum.append(era5tempid.sum(dim = "time", skipna=True))#.compute())

# merge and rearrange for desired dataset
era5ang_mean = xr.concat(bin_arrays_mean,dim = "ang_bins").assign_coords({"ang_bins":np.arange(120, dtype=np.uint32)})
era5ang_sum = xr.concat(bin_arrays_sum,dim = "ang_bins").assign_coords({"ang_bins":np.arange(120, dtype=np.uint32)})

qm_arrays_mean = []
for var in era5ang_mean:
    qm_arrays_mean.append(era5ang_mean[var].rename("qm"))

qm_arrays_sum = []
for var in era5ang_sum:
    qm_arrays_sum.append(era5ang_sum[var].rename("qm"))

era5_bins_mean = xr.concat(qm_arrays_mean, dim = "usimths").assign_coords({"usimths":usimths})
era5_bins_sum = xr.concat(qm_arrays_sum, dim = "usimths").assign_coords({"usimths":usimths})

era5_bins_comb = xr.Dataset({"qmmean":era5_bins_mean, "qmsum":era5_bins_sum})

Выходные данныеera5_bins_comb (ПРИМЕЧАНИЕ. Выходные данные ниже (вычислены) на основе описанного выше, но с выбранной ТОЛЬКО одной долготой, что заняло около часа):

<xarray.Dataset>
Dimensions:    (latitude: 201, ang_bins: 120, usimths: 80)
Coordinates:
    longitude  float32 113.0
  * latitude   (latitude) float32 -36.0 -35.9 -35.8 -35.7 ... -16.2 -16.1 -16.0
  * ang_bins   (ang_bins) uint32 0 1 2 3 4 5 6 7 ... 113 114 115 116 117 118 119
  * usimths    (usimths) float32 0.01 0.02 0.03 0.04 0.05 ... 0.77 0.78 0.79 0.8
Data variables:
    qmmean     (usimths, ang_bins, latitude) float32 nan nan nan ... nan nan nan
    qmsum      (usimths, ang_bins, latitude) float32 0.0 0.0 0.0 ... 0.0 0.0 0.0

На последнем шаге я хочу сохранить файл, используя era5_bins_comb.to_netcdf().

Боюсь, я не слишком хорошо знаком с тем, как работают xarray и dask. Конкретно, когда и когда не использовать .compute(). Попробовав пару вещей с вычислениями между вышеперечисленными, я оказался там. Часть, которая занимает больше всего времени (например, час или больше), — это цикл for ident in range(len(uabine)-1):. Остальное работает довольно быстро (наверное, потому что до этого я ничего не вычисляю). Кроме того, последняя часть сохранения набора данных для его временного хранения и объединения с остальными 73 годами позже также кажется сложной. Похоже, у меня может не хватить памяти и/или это займет вечность. Последнее, что я тестировал, — это использовать era5_bins_comb.isel(usimth=0).compute(), который, по крайней мере, завершился быстрее, но, очевидно, потребовал бы от меня перебора всех usimth и сохранения их отдельно. В качестве альтернативы у меня также есть доступ к большему количеству процессоров и памяти, например. 14 ЦП/63 ГБ (использовалось для тестирования era5_bins_comb.isel(usimth=0).compute()) или 28 ЦП/126.

Просто такое ощущение, что я мог бы улучшить свой код, пропустить некоторые циклы и/или лучше использовать вычислительную мощность. Опять же, я мало что знаю о больших данных, dask и принципах, которым следует следовать. Какие улучшения я могу внести, чтобы получать данные, не дожидаясь вечно и не тратя впустую ресурсы?

Редактировать

Просто дополнительная информация: сегодня я снова запустил свой код с 14 процессорами и засек время цикла for ident in range(len(uabine)-1):. Один цикл занимает примерно 6 минут, что составит 12 часов (12*74=888 часов для файлов за весь год!). Кроме того, поток задач на панели задач выглядит действительно неоптимизированным. Попытка запустить вложенный цикл, перебирая все 80 переменных и все 120 uabine, занимает 20 секунд для каждого вложенного цикла, поэтому 20*120*80/60/60=53 часа, что еще хуже. Возможно, мне нужно найти другой способ избежать цикла уабина, который фильтрует мои данные в зависимости от направления ветра, откуда они дуют.