Automatically chunks a Dataset or DataArray so that each chunk contains
approximately target_chunk_size number of elements.
Parameters:
dataset_to_chunk: xarray.Dataset or xarray.DataArray
target_chunk_size: Target total number of elements per chunk
Returns:
a copy of the dataset which is chunked xarray.Dataset or xarray.DataArray
Source code in pythermogis/dask.py
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42 | def auto_chunk_dataset(
dataset_to_chunk: xr.Dataset | xr.DataArray, target_chunk_size: int
) -> xr.Dataset | xr.DataArray:
"""
Automatically chunks a Dataset or DataArray so that each chunk contains
approximately `target_chunk_size` number of elements.
Parameters:
dataset_to_chunk: xarray.Dataset or xarray.DataArray
target_chunk_size: Target total number of elements per chunk
Returns:
a copy of the dataset which is chunked xarray.Dataset or xarray.DataArray
"""
dim_sizes = dataset_to_chunk.sizes
total_size = np.prod(list(dim_sizes.values()))
if (
total_size <= target_chunk_size
): # No need to chunk, return reservoir property as is
return dataset_to_chunk
# Start with full size
chunking = {dim: dim_sizes[dim] for dim in dataset_to_chunk.dims}
# Greedy algorithm: reduce chunk size along largest dimensions,
# until target_chunk_size is reached
current_chunk_size = total_size
while current_chunk_size > target_chunk_size:
# Sort dims by current chunk size (largest first)
dim_by_size = sorted(chunking.items(), key=lambda kv: kv[1], reverse=True)
for dim, _size in dim_by_size:
if chunking[dim] > 1:
chunking[dim] = max(1, chunking[dim] // 2)
break # reduce one dimension at a time
current_chunk_size = np.prod(list(chunking.values()))
return dataset_to_chunk.copy().chunk(chunking)
|