abtem.core.chunks

abtem.core.chunks#

Module for determining chunk sizes of Dask arrays.

Module Contents#

Functions#

assert_chunks_match_shape

check_chunks_match_shape_length

chunk_ranges

Get the start and end indices for each chunk.

equal_sized_chunks

Split an n integer into m (almost) equal integers, such that the sum of smaller integers equals n.

estimate_potential_chunk_size

Estimate the number of potential slices that fit in the memory budget.

estimate_scan_batch_size

Estimate the maximum number of probe wavefunctions per scan batch.

fill_in_chunk_sizes

generate_chunks

Generate start and end indices for each chunks of equal sized chunks.

is_tuple_of_ints

is_tuple_of_ints_or_tuple_of_ints

is_tuple_of_ints_or_tuple_of_tuple_of_ints

is_tuple_of_tuple_of_ints

is_validated_chunks

Check if the input is are valid chunk sizes.

iterate_chunk_ranges

Iterate over the chunk ranges.

validate_chunks

Validate the chunks for a Dask array based on the shape and a maximum number of elements.

Data#

Chunks

ChunksTuple

ValidatedChunks

API#

abtem.core.chunks.Chunks#

None

abtem.core.chunks.ChunksTuple#

None

abtem.core.chunks.ValidatedChunks#

None

abtem.core.chunks.assert_chunks_match_shape(...) → None[source]#
abtem.core.chunks.check_chunks_match_shape_length(...) → None[source]#
abtem.core.chunks.chunk_ranges(...) → tuple[tuple[tuple[int, int], ...], ...][source]#

Get the start and end indices for each chunk.

Parameters:

chunks (tuple of tuple of int) – The chunk sizes of the Dask array.

Returns:

The range of indices for each chunk.

Return type:

tuple of tuple of tuple of two int

abtem.core.chunks.equal_sized_chunks(...) → tuple[int, ...][source]#

Split an n integer into m (almost) equal integers, such that the sum of smaller integers equals n.

Parameters:

  • num_items (int) – The integer to split.

  • num_chunks (int) – The number integers n will be split into.

  • chunk_size (int) – The size of each chunk.

Returns:

The split integers.

Return type:

tuple of int

abtem.core.chunks.estimate_potential_chunk_size(...) → int[source]#

Estimate the number of potential slices that fit in the memory budget.

build() places the entire slice dimension into a single dask chunk, so the full potential must fit in memory at once. This function calculates how many slices can be held simultaneously when the potential is instead built in smaller chunks via generate_chunked_slices().

On GPU the per-slice cost accounts for CuPy memory pool fragmentation — the pool may hold large contiguous blocks for live arrays (waves, probes) that prevent new allocations even when total free bytes suffice. The effective per-slice cost under fragmentation is empirically ~5× the raw slice size for scan workloads at 4096² grids.

On CPU there is no pool fragmentation and system RAM is typically abundant. The default is therefore to place the entire potential in a single chunk (no chunking), matching pre-chunking behaviour. Set potential.slice-chunk-size in the configuration to a positive integer to enable CPU chunking when memory is genuinely limited.

On GPU the budget uses the CUDA-reported free memory without calling free_all_blocks() first. Dead pool blocks represent recent memory pressure from build/propagation temporaries; leaving them gives a conservative estimate that self-adapts as the pool fills up over successive scan batches. Falls back to dask.chunk-size-gpu when CuPy is unavailable.

Parameters:

  • gpts (tuple of int) – The number of grid points (y, x).

  • device (str) – The device (‘cpu’ or ‘gpu’).

  • dtype (dtype, optional) – The dtype of the potential array. If None, uses float32.

Returns:

The estimated number of slices that fit in memory.

Return type:

int

abtem.core.chunks.estimate_scan_batch_size(...) → int[source]#

Estimate the maximum number of probe wavefunctions per scan batch.

For GPU, queries free CUDA memory at graph-construction time and allocates up to half of it for probe wavefunctions. This is intentionally generous because estimate_potential_chunk_size is called at computation time (inside generate_chunked_slices) when the probe batch is already resident in VRAM; it therefore sees the reduced free memory and sizes the potential chunk to fit in what remains. The two estimates are thus naturally coordinated without requiring explicit cross-referencing.

The raw estimate is rounded to the nearest power of two (preferring the upper power when within 25 % headroom) so that batch sizes snap to GPU-friendly values such as 8, 16, 32, 64.

For CPU, falls back to the dask.chunk-size configuration key.

Parameters:

  • gpts (tuple of int) – Spatial grid size (ny, nx) of each probe wavefunction.

  • dtype (dtype-like) – Wavefunction dtype (typically complex64).

  • device (str) – "gpu" or "cpu".

Returns:

Maximum number of probe wavefunctions per batch (≥ 1).

Return type:

int

abtem.core.chunks.fill_in_chunk_sizes(...) → abtem.core.chunks.ValidatedChunks[source]#
abtem.core.chunks.generate_chunks(...) → Generator[tuple[int, int], None, None][source]#

Generate start and end indices for each chunks of equal sized chunks.

Parameters:

  • num_items (int) – The integer to split.

  • num_chunks (int) – The number integers n will be split into.

  • chunks (int) – The size of each chunk.

  • start (int) – The starting index.

Yields:

tuple of int – The start and end indices of the current chunk.

abtem.core.chunks.is_tuple_of_ints(...) → TypeGuard[tuple[int, ...]][source]#
abtem.core.chunks.is_tuple_of_ints_or_tuple_of_ints(...) → TypeGuard[tuple[tuple[int, ...], ...]][source]#
abtem.core.chunks.is_tuple_of_ints_or_tuple_of_tuple_of_ints(...) → TypeGuard[tuple[int | tuple[int, ...], ...]][source]#
abtem.core.chunks.is_tuple_of_tuple_of_ints(...) → TypeGuard[tuple[tuple[int, ...], ...]][source]#
abtem.core.chunks.is_validated_chunks(...) → TypeGuard[abtem.core.chunks.ValidatedChunks][source]#

Check if the input is are valid chunk sizes.

Parameters:

x (int or tuple of int or tuple of tuple of int or str) – The chunk sizes of the Dask array.

Returns:

True if the input is a valid chunk size.

Return type:

TypeGuard[ValidatedChunks]

abtem.core.chunks.iterate_chunk_ranges(...)[source]#

Iterate over the chunk ranges.

Parameters:

chunks (tuple of tuple of int) – The chunk sizes of the Dask array.

Yields:

  • block_indices (tuple of int) – The indices of the current block.

  • slices (tuple of slice) – The slices indexing the current block.

abtem.core.chunks.validate_chunks(...) → abtem.core.chunks.ValidatedChunks[source]#

Validate the chunks for a Dask array based on the shape and a maximum number of elements.

Parameters:

  • shape (tuple of int) – The shape of the array.

  • chunks (int or tuple of int or str) – The chunk sizes of the Dask array. If an integer, the array will be split into equal chunks. If a tuple, the array will be split into the specified chunks. If “auto”, the chunks will be determined automatically based on the shape and the maximum number of elements.

  • max_elements (int or str) – The maximum number of elements in a chunk. If “auto”, the maximum number of elements will be determined based on the maximum number of bytes per chunk and the dtype.

  • dtype (dtype) – The dtype of the array.

  • device (str) – The device the array will be stored on.

Returns:

The chunk sizes of the Dask array.

Return type:

tuple of tuple of int