"""Module for describing array objects."""
from __future__ import annotations
import copy
import json
import operator
import warnings
from abc import ABCMeta, abstractmethod
from contextlib import contextmanager, nullcontext
from functools import partial
from numbers import Number
from types import ModuleType
from typing import (
TYPE_CHECKING,
Any,
Callable,
Generator,
Optional,
Self,
Sequence,
TypeVar,
Union,
)
import dask
import dask.array as da
import numpy as np
from dask.array.utils import validate_axis
from dask.diagnostics import Profiler, ProgressBar, ResourceProfiler
from tqdm.dask import TqdmCallback
from abtem._version import __version__
from abtem.core import config
from abtem.core.axes import (
AxesMetadataList,
AxisMetadata,
LinearAxis,
OrdinalAxis,
UnknownAxis,
_normalized_axis_weights,
axis_from_dict,
axis_to_dict,
)
from abtem.core.backend import (
asnumpy,
check_cupy_is_installed,
copy_to_device,
cp,
device_name_from_array_module,
get_array_module,
push_config_to_workers,
)
from abtem.core.chunks import Chunks, iterate_chunk_ranges, validate_chunks
from abtem.core.ensemble import Ensemble, _wrap_with_array, unpack_blockwise_args
from abtem.core.utils import (
CopyMixin,
EqualityMixin,
get_dtype,
interleave,
itemset,
normalize_axes,
number_to_tuple,
tuple_range,
)
from abtem.transform import TransformFromFunc
if TYPE_CHECKING:
from abtem.transform import ArrayObjectTransform
tifffile: Optional[ModuleType] = None
try:
import tifffile # type: ignore
except ImportError:
pass
hs: Optional[ModuleType] = None
try:
import hyperspy.api as hs # type: ignore
except ImportError:
pass
xr: Optional[ModuleType] = None
try:
import xarray as xr # type: ignore
except ImportError:
pass
em: Optional[ModuleType] = None
try:
import quantem as em # type: ignore
except ImportError:
pass
ArrayObjectType = TypeVar("ArrayObjectType", bound="ArrayObject")
ArrayObjectTypeAlt = TypeVar("ArrayObjectTypeAlt", bound="ArrayObject")
ArrayItemType = Union[int, slice, list, np.ndarray, None]
def _to_hyperspy_axes_metadata(
axes_metadata: list[AxisMetadata], shape: int | tuple[int, ...]
):
hyperspy_axes = []
if not isinstance(shape, (list, tuple)):
shape = (shape,)
for metadata, n in zip(axes_metadata, shape):
hyperspy_axis = {"size": n, "name": metadata.label}
if isinstance(metadata, LinearAxis):
hyperspy_axis["scale"] = metadata.sampling
hyperspy_axis["offset"] = metadata.offset
hyperspy_axis["units"] = metadata.units
elif isinstance(metadata, OrdinalAxis):
if all(isinstance(value, Number) for value in metadata.values) and (
all(
metadata.values[i] < metadata.values[i + 1]
for i in range(len(metadata.values) - 1)
)
):
hyperspy_axis["axis"] = metadata.values
hyperspy_axis["units"] = metadata.units
else:
warnings.warn(
f"Axis ({metadata.label}) not supported by hyperspy, some metadata"
f"will be lost."
)
else:
raise RuntimeError()
hyperspy_axes.append(hyperspy_axis)
return hyperspy_axes
def _extract_blockwise_multi_output(arr: np.ndarray, index: int) -> np.ndarray:
arr = arr.item()[index]
return arr
[docs]
def multi_output_blockwise(
func: Callable,
array: da.core.Array,
chunks: tuple[tuple[int, ...], ...],
array_axes: tuple[da.core.Array, ...],
new_axes: tuple[da.core.Array, ...],
out_metas: tuple,
drop_axes: tuple[tuple[int, ...], ...],
new_shapes: tuple[tuple[int, ...], ...],
**kwargs,
) -> tuple[da.core.Array, ...]:
assert len(out_metas) == len(drop_axes)
base_ndim = array.ndim
new_ndim = sum(len(axis.shape) for axis in new_axes)
out_ndim = new_ndim + base_ndim
old_chunks = array.chunks
assert isinstance(array_axes, tuple)
assert len(array_axes) == base_ndim
assert all(len(base_axis.numblocks) == 1 for base_axis in array_axes)
assert all(
base_axis.numblocks[0] == len(c) for base_axis, c in zip(array_axes, old_chunks)
)
out_symbols = tuple(range(out_ndim))
array_symbols = tuple(range(new_ndim, out_ndim))
old_axes_symbols = tuple((i,) for i in array_symbols)
old_axes_args = interleave(array_axes, old_axes_symbols)
new_axes_args: tuple[da.core.Array | tuple[int, ...], ...] = ()
current_symbol = 0
for axis in new_axes:
new_axes_args += (
axis,
tuple(range(current_symbol, current_symbol + len(axis.shape))),
)
current_symbol += len(axis.shape)
out_array = da.blockwise(
func,
out_symbols,
array,
array_symbols,
*old_axes_args,
*new_axes_args,
align_arrays=False,
meta=np.array((), dtype=object),
**kwargs,
)
outputs: tuple[da.core.Array, ...] = ()
for i, (drop_axis, out_meta, new_shape) in enumerate(
zip(drop_axes, out_metas, new_shapes)
):
if not all(len(out_array.chunks[i]) == 1 for i in drop_axis):
raise RuntimeError()
drop_chunks = []
for j, (item, ns) in enumerate(zip(chunks, new_shape)):
if j not in drop_axis:
if sum(item) != ns:
assert len(item) == 1
item = (ns,)
drop_chunks.append(item)
drop_chunks = tuple(drop_chunks)
new_output = da.map_blocks(
_extract_blockwise_multi_output,
out_array,
chunks=drop_chunks,
drop_axis=drop_axis,
index=i,
meta=out_meta,
)
outputs += (new_output,)
return outputs
# Codecs (Blosc/zlib) reject a single buffer over 2**31 - 1 bytes. Writing a
# whole array as one zarr chunk -- as this module used to always do -- hits
# that limit for any reasonably large array, especially at float64. Stay well
# under it (also a more sensible chunk size for parallel IO in general) by
# halving the largest axis of the chunk shape until it fits the budget.
_MAX_ZARR_CHUNK_BYTES = 512 * 1024**2
def _safe_zarr_chunks(
shape: tuple[int, ...],
itemsize: int,
max_bytes: Optional[int] = None,
n_fixed_trailing_axes: int = 0,
) -> tuple[int, ...]:
"""Pick a chunk shape that stays under ``max_bytes``, preferring to
shrink leading (ensemble) axes and never touching the trailing
``n_fixed_trailing_axes`` (an ArrayObject's ``_base_dims``) unless there
is no other choice.
Several of abTEM's own lazy dask operations on measurements (e.g.
``interpolate_line``'s ``da.map_blocks(..., drop_axis=...)``) require
their base/measurement axes to be a single dask chunk -- splitting them
doesn't raise, it silently produces wrong results (confirmed by direct
reproduction: a DiffractionPatterns array whose spatial axes were split
this way, then reloaded lazily, gave an all-zero momentum-resolved
spectrum while eagerly-computed data from the same file was correct).
Chunking only the ensemble axes keeps a round-tripped array safe for
every lazy operation that already assumes whole base-axis chunks,
matching e.g. ``PotentialArray``'s own ``("auto",) * n + (-1,) *
_base_dims`` convention elsewhere in this module.
"""
if max_bytes is None:
# Read fresh rather than as a default-argument value, so overriding
# the module-level budget (e.g. in tests) takes effect.
max_bytes = _MAX_ZARR_CHUNK_BYTES
chunks = list(shape)
def nbytes() -> int:
n = itemsize
for c in chunks:
n *= c
return n
n_fixed = min(n_fixed_trailing_axes, len(chunks))
splittable = list(range(len(chunks) - n_fixed))
while nbytes() > max_bytes and any(chunks[i] > 1 for i in splittable):
i = max(splittable, key=lambda j: chunks[j])
chunks[i] = max(1, chunks[i] // 2)
# Only reached if the base axes are themselves so large that shrinking
# every ensemble axis to 1 still isn't enough (e.g. a single very large
# image with little to no ensemble axis to shrink) -- fall back to
# splitting them too rather than failing outright; this is the one case
# where the lazy-op hazard described above is genuinely unavoidable.
while nbytes() > max_bytes and any(c > 1 for c in chunks):
i = max(range(len(chunks)), key=lambda j: chunks[j])
chunks[i] = max(1, chunks[i] // 2)
return tuple(chunks)
[docs]
class ComputableList(list):
"""A list with methods for conveniently computing its items."""
[docs]
def to_zarr(
self,
url: str,
compute: bool = True,
overwrite: bool = False,
progress_bar: Optional[bool] = None,
compression_level: int | None = 4,
**kwargs: Any,
):
"""Write data to a zarr file.
Parameters
----------
url : str
Location of the data. For Zarr 3 zip stores, use a .zip extension.
For directory stores, use .zarr extension or a directory path.
compute : bool
If true compute immediately; return dask.delayed.Delayed otherwise.
overwrite : bool
If given array already exists, overwrite=False will cause an error, where
overwrite=True will replace the existing data.
progress_bar : bool
Display a progress bar in the terminal or notebook during computation. The
progress bar is only displayed with a local scheduler.
compression_level : int or None
If set (0–9), applies Zstandard compression with Blosc backend at that level.
Level 0 disables compression. Default is 4, raises ValueError if > 9.
kwargs :
Keyword arguments passed to `dask.array.to_zarr`.
"""
import os
import zarr
# Validate compression level
# Based on quantem: https://github.com/electronmicroscopy/quantem
if compression_level is not None:
if not (0 <= compression_level <= 9):
raise ValueError(
f"Compression_level must be between 0 and 9 instead of {compression_level}"
)
compressors = {
"name": "blosc",
"configuration": {
"cname": "zstd",
"clevel": int(compression_level),
"shuffle": "bitshuffle",
},
}
else:
compressors = None
# Helper functions for type preservation
def encode_types(obj):
"""Recursively encode tuples for JSON serialization."""
if isinstance(obj, tuple):
return {
"_type": "tuple",
"_value": [encode_types(item) for item in obj],
}
elif isinstance(obj, list):
return [encode_types(item) for item in obj]
elif isinstance(obj, dict):
return {key: encode_types(value) for key, value in obj.items()}
# --- NumPy scalar handling ---
elif isinstance(obj, (np.integer,)):
return int(obj)
elif isinstance(obj, (np.floating,)):
return float(obj)
elif isinstance(obj, (np.bool_,)):
return bool(obj)
elif isinstance(obj, np.ndarray):
return obj.tolist()
# --- GPU (CuPy) array/scalar handling ---
elif hasattr(obj, "__cuda_array_interface__"):
return obj.tolist()
else:
return obj
# Determine if this is a zip file
is_zip = url.endswith(".zip")
arrays_to_write = []
metadata_list = []
base_dims_by_index = {}
for i, has_array in enumerate(self):
has_array = has_array.ensure_lazy()
array = has_array.copy_to_device("cpu").array
metadata_dict = has_array._metadata_to_dict()
packed_kwargs = has_array._pack_kwargs(
has_array._copy_kwargs(exclude=("array",))
)
metadata_dict["kwargs"] = packed_kwargs
metadata_dict = encode_types(metadata_dict)
arrays_to_write.append((i, array))
metadata_list.append({f"metadata{i}": metadata_dict})
base_dims_by_index[i] = has_array._base_dims
if is_zip:
# Use ZipStore for .zip files
@dask.delayed
def write_to_zipstore(
computed_arrays,
url,
metadata_list,
overwrite,
compressors=compressors,
):
if overwrite and os.path.exists(url):
os.remove(url)
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message="Duplicate name:.*zarr.json",
category=UserWarning,
)
store = zarr.storage.ZipStore(url, mode="w")
try:
root = zarr.group(store=store, overwrite=True)
for metadata_dict in metadata_list:
for key, value in metadata_dict.items():
root.attrs[key] = value
for i, computed_array in computed_arrays:
root.create_array(
name=f"array{i}",
data=computed_array,
chunks=_safe_zarr_chunks(
computed_array.shape,
computed_array.dtype.itemsize,
n_fixed_trailing_axes=base_dims_by_index[i],
),
overwrite=True,
compressors=compressors,
)
except BaseException:
# A failed write (e.g. a chunk still over a codec's
# buffer limit) can leave a .zip with valid-looking
# metadata (shape/dtype/chunks) but no actual chunk
# data -- silently readable later as all zeros
# (zarr's fill_value for a declared-but-missing
# chunk) instead of raising again. Don't leave that
# behind.
store.close()
if os.path.exists(url):
os.remove(url)
raise
else:
store.close()
return url
delayed_arrays = [
(i, dask.delayed(array.compute)()) for i, array in arrays_to_write
]
delayed_write = write_to_zipstore(
delayed_arrays, url, metadata_list, overwrite, compressors=compressors
)
else:
# Use directory store for non-.zip files
@dask.delayed
def write_to_directory(computed_arrays, url, metadata_list, overwrite):
import shutil
if overwrite and os.path.exists(url):
shutil.rmtree(url)
root = zarr.open(url, mode="w")
def _close():
store = getattr(root, "store", None)
if store is not None:
close_fn = getattr(store, "close", None)
if callable(close_fn):
close_fn()
try:
for metadata_dict in metadata_list:
for key, value in metadata_dict.items():
root.attrs[key] = value
for i, computed_array in computed_arrays:
root.create_array(
name=f"array{i}",
data=computed_array,
chunks=_safe_zarr_chunks(
computed_array.shape,
computed_array.dtype.itemsize,
n_fixed_trailing_axes=base_dims_by_index[i],
),
overwrite=True,
)
except BaseException:
# See the matching comment in write_to_zipstore: don't
# leave a partially-written store with valid-looking
# metadata but missing/incomplete chunk data.
_close()
if os.path.exists(url):
shutil.rmtree(url)
raise
else:
_close()
return url
delayed_arrays = [
(i, dask.delayed(array.compute)()) for i, array in arrays_to_write
]
delayed_write = write_to_directory(
delayed_arrays, url, metadata_list, overwrite
)
if not compute:
# The arrays are computed later, inside a task on the client's
# workers, which must then hold the configuration of this call.
_push_config_to_active_client()
return delayed_write
_push_config_to_active_client()
with _compute_context(
progress_bar, profiler=False, resource_profiler=False
) as (_, profiler, resource_profiler):
output = dask.compute(delayed_write, **kwargs)[0]
profilers = tuple(p for p in (profiler, resource_profiler) if p is not None)
if profilers:
return output, profilers
else:
return output
[docs]
def compute(self, **kwargs) -> list[ArrayObject] | tuple[list[ArrayObject], tuple]:
"""Turn a list of lazy ArrayObjects object into the in-memory equivalents.
kwargs :
Keyword arguments passed to `ArrayObject.compute`.
"""
output, profilers = _compute(self, **kwargs)
if profilers:
return output, profilers
return output
def _get_progress_bar(
progress_bar: Optional[bool] = None,
) -> Union[ProgressBar, TqdmCallback, nullcontext]:
if progress_bar is None:
progress_bar = config.get("diagnostics.progress_bar")
progress_bar_obj: Union[ProgressBar, TqdmCallback, nullcontext]
if progress_bar:
if progress_bar == "tqdm":
progress_bar_obj = TqdmCallback(desc="tasks", delay=0.5)
else:
progress_bar_obj = ProgressBar()
else:
progress_bar_obj = nullcontext()
return progress_bar_obj
@contextmanager
def _compute_context(
progress_bar: Optional[bool] = None,
profiler: int = False,
resource_profiler: int = False,
) -> Generator[tuple[Any, Any, Any], None, None]:
progress_bar_ctx = _get_progress_bar(progress_bar)
profiler_ctx: Union[Profiler, nullcontext]
resource_profiler_ctx: Union[ResourceProfiler, nullcontext]
if profiler:
profiler_ctx = Profiler()
else:
profiler_ctx = nullcontext()
if resource_profiler:
resource_profiler_ctx = ResourceProfiler()
else:
resource_profiler_ctx = nullcontext()
with (
progress_bar_ctx as progress_bar_ctx1,
profiler_ctx as profiler_ctx1,
resource_profiler_ctx as resource_profiler_ctx1,
):
yield progress_bar_ctx1, profiler_ctx1, resource_profiler_ctx1
def _push_config_to_active_client():
"""Mirror the configuration onto an active distributed client's workers.
Any distributed client (a CPU LocalCluster, a SLURMCluster, ...) executes
tasks in worker processes that never saw the client's ``abtem.config.set``.
Deduplicated inside push_config_to_workers, so calling this on every
dispatch is cheap.
"""
try:
from distributed import get_client
client = get_client()
except (ImportError, ValueError):
return
push_config_to_workers(client)
def _compute(
array_objects: list[ArrayObjectType],
progress_bar: Optional[bool] = None,
profiler: bool = False,
resource_profiler: bool = False,
**kwargs,
) -> tuple[list[ArrayObjectType], tuple]:
is_gpu = config.get("device") == "gpu" or any(
hasattr(obj, "device") and obj.device == "gpu" for obj in array_objects
)
_push_config_to_active_client()
if is_gpu:
check_cupy_is_installed()
if "scheduler" not in kwargs:
kwargs["scheduler"] = "synchronous"
if "num_workers" not in kwargs:
kwargs["num_workers"] = cp.cuda.runtime.getDeviceCount()
if "threads_per_worker" not in kwargs:
kwargs["threads_per_worker"] = cp.cuda.runtime.getDeviceCount()
with _compute_context(
progress_bar, profiler=profiler, resource_profiler=resource_profiler
) as (_, profiler, resource_profiler):
arrays = dask.compute([wrapper.array for wrapper in array_objects], **kwargs)[0]
for array, wrapper in zip(arrays, array_objects):
wrapper._array = array
profilers = tuple(p for p in (profiler, resource_profiler) if p is not None)
return array_objects, profilers
[docs]
def validate_lazy(lazy: Optional[bool]) -> bool:
"""
Validate the lazy argument.
Parameters
----------
lazy : bool or None
If None, return the default value from the configuration. If a boolean, return
the value.
Returns
-------
lazy : bool
The validated value of lazy.
"""
if lazy is None:
return config.get("dask.lazy")
elif not isinstance(lazy, bool):
raise ValueError("lazy must be a boolean")
return lazy
def _validate_array_items(
items: ArrayItemType | tuple[ArrayItemType, ...],
shape: tuple[int, ...],
keepdims: bool = False,
) -> tuple[ArrayItemType, ...]:
if isinstance(items, (int, slice, type(None), list, np.ndarray)):
items = (items,)
elif not isinstance(items, tuple):
raise NotImplementedError(
(
"Indices must be integers or slices or a tuple of integers or"
f" slices or None, not {type(items).__name__}."
)
)
if keepdims:
items = tuple(
slice(item, item + 1) if isinstance(item, int) else item for item in items
)
if any(isinstance(item, (type(...),)) for item in items):
raise NotImplementedError
if len(tuple(item for item in items if item is not None)) > len(shape):
raise RuntimeError("too many indices for array")
return items
[docs]
class ArrayObject(Ensemble, EqualityMixin, CopyMixin, metaclass=ABCMeta):
"""A base class for simulation objects described by an array and associated
metadata.
Parameters
----------
array : ndarray
Array representing the array object.
ensemble_axes_metadata : list of AxesMetadata
Axis metadata for each ensemble axis. The axis metadata must be compatible with
the shape of the array.
metadata : dict
A dictionary defining wave function metadata. All items will be added to the
metadata of measurements derived from the waves.
"""
_base_dims: int
[docs]
def __init__(
self,
array: np.ndarray | da.core.Array,
ensemble_axes_metadata: list[AxisMetadata] | None = None,
metadata: dict | None = None,
**kwargs,
):
if ensemble_axes_metadata is None:
ensemble_axes_metadata = []
if metadata is None:
metadata = {}
self._array = array
self._ensemble_axes_metadata = ensemble_axes_metadata
self._metadata = metadata
if len(array.shape) < self._base_dims:
raise RuntimeError(
f"{self.__class__.__name__} must be {self._base_dims}D or greater, not "
f"{len(array.shape)}D"
)
self._check_axes_metadata()
super().__init__(**kwargs)
@property
def base_dims(self) -> int:
"""Number of base dimensions."""
return self._base_dims
@property
def ensemble_dims(self) -> int:
"""Number of ensemble dimensions."""
return len(self.shape) - self.base_dims
@property
def base_axes_metadata(self) -> list[AxisMetadata]:
return [UnknownAxis() for _ in range(self._base_dims)]
@property
def shape(self) -> tuple[int, ...]:
"""Shape of the underlying array."""
return self.array.shape
@property
def base_shape(self) -> tuple[int, ...]:
"""Shape of the base axes of the underlying array."""
return self.shape[self.ensemble_dims :]
@property
def ensemble_shape(self) -> tuple[int, ...]:
"""Shape of the ensemble axes of the underlying array."""
return self.shape[: self.ensemble_dims]
@property
def ensemble_axes_metadata(self) -> list[AxisMetadata]:
"""List of AxisMetadata of the ensemble axes."""
return self._ensemble_axes_metadata
@property
def axes_metadata(self) -> AxesMetadataList:
"""List of AxisMetadata."""
return AxesMetadataList(
self.ensemble_axes_metadata + self.base_axes_metadata, self.shape
)
def _check_axes_metadata(self) -> None:
if len(self.shape) != len(self.axes_metadata):
raise RuntimeError(
f"number of array dimensions ({len(self.shape)}) does not match number"
f"of axis metadata items ({len(self.axes_metadata)})"
)
for n, axis in zip(self.shape, self.axes_metadata):
if isinstance(axis, OrdinalAxis) and len(axis) != n:
raise RuntimeError(
f"number of values for ordinal axis ({len(axis)}), does not match"
f"size of dimension ({n})"
)
def _is_base_axis(self, axis: int | tuple[int, ...]) -> bool:
axis = number_to_tuple(axis)
base_axes = tuple(range(len(self.ensemble_shape), len(self.shape)))
return len(set(axis).intersection(base_axes)) > 0
[docs]
def apply_func(self, func: Callable, **kwargs) -> Self:
"""Apply a function to the array object. The function must take an array as its
first argument, only the array is modified, the metadata is not changed. The
function is applied lazily if the array object is lazy.
Parameters
----------
func : callable
Function to apply to the array object.
kwargs :
Additional keyword arguments passed to the function.
Returns
-------
array_object : ArrayObject or subclass of ArrayObject
The array object with the function applied.
"""
transform = TransformFromFunc(func, func_kwargs=kwargs)
return transform.apply(self)
def get_from_metadata(self, name: str, broadcastable: bool = False):
axes_metadata_index = None
data = None
for i, (n, axis) in enumerate(zip(self.shape, self.ensemble_axes_metadata)):
if axis.label == name:
data = axis.coordinates(n)
axes_metadata_index = i
if axes_metadata_index is not None and broadcastable:
return np.array(data)[
(
*((None,) * axes_metadata_index),
slice(None),
*((None,) * (len(self.ensemble_shape) - 1 - axes_metadata_index)),
)
]
elif axes_metadata_index is not None:
if name in self.metadata.keys():
raise RuntimeError(
f"Could not resolve metadata for {name}, found in both ensemble"
"axes metadata and metadata"
)
return data
else:
try:
return self.metadata[name]
except KeyError as exc:
raise RuntimeError(f"Could not resolve metadata for {name}") from exc
[docs]
def rechunk(self, chunks: Chunks, **kwargs) -> ArrayObject:
"""Rechunk dask array.
chunks : int or tuple or str
How to rechunk the array. See `dask.array.rechunk`.
kwargs :
Additional keyword arguments passes to `dask.array.rechunk`.
"""
if not self.is_lazy:
raise RuntimeError("cannot rechunk non-lazy array object")
if isinstance(chunks, tuple) and len(chunks) < len(self.shape):
chunks = chunks + ("auto",) * max((self.ensemble_dims - len(chunks), 0))
chunks = chunks + (-1,) * max((len(self.shape) - len(chunks), 0))
array = self._lazy_array.rechunk(chunks=chunks, **kwargs)
# by keyword: not every subclass takes the array as its first argument
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
return self.__class__(**kwargs)
@property
def metadata(self) -> dict:
"""Metadata stored as a dictionary."""
return self._metadata
def __len__(self) -> int:
return len(self.array)
@property
def array(self) -> np.ndarray | da.core.Array:
"""Underlying array describing the array object."""
return self._array
@array.setter
def array(self, array: np.ndarray | da.core.Array):
"""Set underlying array describing the array object."""
if not array.shape == self.shape:
raise ValueError("Shape of array must match shape of object.")
if not array.dtype == self.dtype:
raise ValueError("Datatype of array must match datatype of object.")
if self.is_lazy != isinstance(array, da.core.Array):
raise ValueError("Type of array must match type of object.")
self._array = array
@property
def _lazy_array(self) -> da.core.Array:
"""Underlying lazy array describing the array object."""
if not self.is_lazy:
raise RuntimeError("array object is not lazy")
assert isinstance(self.array, da.core.Array)
return self.array
@property
def _eager_array(self) -> np.ndarray:
"""Underlying eager array describing the array object."""
if self.is_lazy:
raise RuntimeError("array object is lazy")
assert not isinstance(self.array, da.core.Array)
return self.array
@property
def dtype(self) -> np.dtype:
"""Datatype of array."""
return self._array.dtype
@property
def device(self) -> str:
"""The device where the array is stored."""
return device_name_from_array_module(get_array_module(self.array))
@property
def is_lazy(self) -> bool:
"""True if array is lazy."""
return isinstance(self.array, da.core.Array)
@classmethod
def _to_delayed_func(cls, array, **kwargs):
kwargs["array"] = array
return cls(**kwargs)
@property
def is_complex(self) -> bool:
"""True if array is complex."""
return np.iscomplexobj(self.array)
def _check_is_compatible(self, other: ArrayObject):
if not isinstance(other, self.__class__):
raise RuntimeError(
f"incompatible types ({self.__class__} != {other.__class__})"
)
[docs]
def mean(
self,
axis: Optional[int | tuple[int, ...]] = None,
keepdims: bool = False,
split_every: int = 2,
) -> Self:
"""Mean of array object over one or more axes. Only ensemble axes can be
reduced.
Parameters
----------
axis : int or tuple of ints, optional
Axis or axes along which a means are calculated. The default is to compute
the mean of the flattened array. If this is a tuple of ints, the mean is
calculated over multiple axes. The indicated axes must be ensemble axes.
keepdims : bool, optional
If True, the reduced axes are left in the result as dimensions with size
one. Default is False.
split_every : int
Only used for lazy arrays. See `dask.array.reductions`.
Returns
-------
reduced_array : ArrayObject or subclass of ArrayObject
The reduced array object.
"""
self._warn_if_weighted_axes(axis, "mean")
return self._reduction(
"mean", axes=axis, keepdims=keepdims, split_every=split_every
)
[docs]
def sum(
self,
axis: Optional[int | tuple[int, ...]] = None,
keepdims: bool = False,
split_every: int = 2,
) -> ArrayObject:
"""Sum of array object over one or more axes. Only ensemble axes can be reduced.
Parameters
----------
axis : int or tuple of ints, optional
Axis or axes along which a sums are performed. The default is to compute the
mean of the flattened array. If this is a tuple of ints, the sum is
performed over multiple axes. The indicated axes must be ensemble axes.
keepdims : bool, optional
If True, the reduced axes are left in the result as dimensions with size
one. Default is False.
split_every : int
Only used for lazy arrays. See `dask.array.reductions`.
Returns
-------
reduced_array : ArrayObject or subclass of ArrayObject
The reduced array object.
"""
self._warn_if_weighted_axes(axis, "sum")
return self._reduction(
"sum", axes=axis, keepdims=keepdims, split_every=split_every
)
[docs]
def std(
self,
axis: Optional[int | tuple[int, ...]] = None,
keepdims: bool = False,
split_every: int = 2,
) -> ArrayObject:
"""Standard deviation of array object over one or more axes. Only ensemble axes
can be reduced.
Parameters
----------
axis : int or tuple of ints, optional
Axis or axes along which a standard deviations are calculated. The default
is to compute the mean of the flattened array. If this is a tuple of ints,
the standard deviations are calculated over multiple axes.
The indicated axes must be ensemble axes.
keepdims : bool, optional
If True, the reduced axes are left in the result as dimensions with size
one. Default is False.
split_every : int
Only used for lazy arrays. See `dask.array.reductions`.
Returns
-------
reduced_array : ArrayObject or subclass of ArrayObject
The reduced array object.
"""
return self._reduction(
"std", axes=axis, keepdims=keepdims, split_every=split_every
)
[docs]
def min(
self,
axis: Optional[int | tuple[int, ...]] = None,
keepdims: bool = False,
split_every: int = 2,
) -> ArrayObject:
"""Minmimum of array object over one or more axes. Only ensemble axes can be
reduced.
Parameters
----------
axis : int or tuple of ints, optional
Axis or axes along which a minima are calculated. The default is to compute
the mean of the flattened array. If this is a tuple of ints, the minima are
calculated over multiple axes. The indicated axes must be ensemble axes.
keepdims : bool, optional
If True, the reduced axes are left in the result as dimensions with size
one. Default is False.
split_every : int
Only used for lazy arrays. See `dask.array.reductions`.
Returns
-------
reduced_array : ArrayObject or subclass of ArrayObject
The reduced array object.
"""
return self._reduction(
"min", axes=axis, keepdims=keepdims, split_every=split_every
)
[docs]
def max(
self,
axis: Optional[int | tuple[int, ...]] = None,
keepdims: bool = False,
split_every: int = 2,
) -> ArrayObject:
"""Maximum of array object over one or more axes. Only ensemble axes can be
reduced.
Parameters
----------
axis : int or tuple of ints, optional
Axis or axes along which a maxima are calculated. The default is to compute
the mean of the flattened array. If this is a tuple of ints, the maxima are
calculated over multiple axes. The indicated axes must be ensemble axes.
keepdims : bool, optional
If True, the reduced axes are left in the result as dimensions with size
one. Default is False.
split_every : int
Only used for lazy arrays. See `dask.array.reductions`.
Returns
-------
reduced_array : ArrayObject or subclass of ArrayObject
The reduced array object.
"""
return self._reduction(
"max", axes=axis, keepdims=keepdims, split_every=split_every
)
def _weighted_axes(self, axes: tuple[int, ...]) -> dict[int, np.ndarray]:
"""Normalized probability weights of the given (non-negative) axes, for the
axes that carry non-uniform weights."""
weights = {}
for axis in axes:
axis_weights = _normalized_axis_weights(
self.axes_metadata[axis], self.shape[axis]
)
if axis_weights is not None:
weights[axis] = axis_weights
return weights
def _warn_if_weighted_axes(
self, axes: Optional[int | tuple[int, ...]], reduction_func: str
) -> None:
if axes is None:
# A reduction over all axes includes every ensemble axis.
axes = tuple(range(len(self.ensemble_shape)))
else:
axes = tuple(
axis if axis >= 0 else len(self.shape) + axis
for axis in number_to_tuple(axes)
)
if any(axis >= len(self.ensemble_shape) for axis in axes):
return # _reduction raises for base axes
if self._weighted_axes(axes):
warnings.warn(
f"`{reduction_func}` over an ensemble axis carrying distribution "
"(probability) weights ignores the weights. The ensemble members "
"are unweighted and the distribution weights are applied when the "
"ensemble is reduced: use `reduce_ensemble(axis=...)` for the "
"probability-weighted mean.",
UserWarning,
stacklevel=3,
)
def _weighted_ensemble_mean(
self, axes: int | tuple[int, ...], split_every: int = 2
) -> Self:
"""Probability-weighted mean over ensemble axes.
Computes ``Σ_i p_i A_i / Σ_i p_i`` over each axis, where ``p_i`` are the
weights carried by the axis metadata (see :class:`OrdinalAxis`); axes
without (or with equal) weights are averaged with a plain mean. Over
several axes the weights are the outer product of the per-axis weights.
Works for NumPy, CuPy and lazy Dask arrays (the result stays lazy) and
preserves the array dtype.
"""
axes = tuple(
axis if axis >= 0 else len(self.shape) + axis
for axis in number_to_tuple(axes)
)
if self._is_base_axis(axes):
raise RuntimeError("base axes cannot be reduced")
weights = self._weighted_axes(axes)
if not weights:
# Equal weights (e.g. frozen phonons): the plain mean is exact, and
# keeping it leaves such results bitwise unchanged.
return self._reduction("mean", axes=axes, split_every=split_every)
# Integer arrays (e.g. counts) take the configured float precision rather
# than truncating the weights.
array_object = self
if self.array.dtype.kind not in "fc":
array_object = self.__class__(
**{
**self._copy_kwargs(exclude=("array",)),
"array": self.array.astype(get_dtype(complex=False)),
}
)
# Average the equal-weight axes first: it shrinks the array before the
# weighted pass, and the outer-product weights factorize over the axes.
unweighted = tuple(axis for axis in axes if axis not in weights)
if unweighted:
array_object = array_object._reduction(
"mean", axes=unweighted, split_every=split_every
)
weights = {
axis - sum(other < axis for other in unweighted): axis_weights
for axis, axis_weights in weights.items()
}
# A single multiply by the outer product of the per-axis weights, a small
# host array spanning only the weighted axes. The real dtype matches the
# array's precision (complex64 -> float32, ...), so the weights never
# promote the result.
ndim = len(array_object.shape)
combined = np.ones((1,) * ndim)
for axis, axis_weights in weights.items():
shape = [1] * ndim
shape[axis] = len(axis_weights)
combined = combined * axis_weights.reshape(shape)
xp = get_array_module(array_object.array)
real_dtype = np.finfo(array_object.array.dtype).dtype
array = array_object.array * xp.asarray(combined, dtype=real_dtype)
weighted = array_object.__class__(
**{**array_object._copy_kwargs(exclude=("array",)), "array": array}
)
return weighted._reduction(
"sum", axes=tuple(weights), split_every=split_every
)
def _reduction(
self,
reduction_func: str,
axes: Optional[int | tuple[int, ...]] = None,
keepdims: bool = False,
split_every: int = 2,
**kwargs,
) -> Self:
xp = get_array_module(self.array)
if axes is None:
if self.is_lazy:
return getattr(da, reduction_func)(self.array)
else:
return getattr(xp, reduction_func)(self.array)
axes = number_to_tuple(axes)
axes = tuple(axis if axis >= 0 else len(self.shape) + axis for axis in axes)
if self._is_base_axis(axes):
raise RuntimeError("base axes cannot be reduced")
ensemble_axes_metadata = copy.deepcopy(self.ensemble_axes_metadata)
if not keepdims:
ensemble_axes = tuple(range(len(self.ensemble_shape)))
ensemble_axes_metadata = [
axis_metadata
for axis_metadata, axis in zip(ensemble_axes_metadata, ensemble_axes)
if axis not in axes
]
default_kwargs = self._copy_kwargs(exclude=("array",))
kwargs = {**default_kwargs, **kwargs}
if self.is_lazy:
kwargs["array"] = getattr(da, reduction_func)(
self.array, axes, split_every=split_every, keepdims=keepdims
)
else:
kwargs["array"] = getattr(xp, reduction_func)(
self.array, axes, keepdims=keepdims
)
kwargs["ensemble_axes_metadata"] = ensemble_axes_metadata
return self.__class__(**kwargs)
def _arithmetic(
self, other: Self | np.ndarray | da.core.Array | Number, func: str
) -> Self:
other_array: np.ndarray | da.core.Array | Number
if isinstance(other, self.__class__):
self._check_is_compatible(other)
other_array = other.array
elif isinstance(other, (np.ndarray, da.core.Array, Number)) or (
cp is not None and isinstance(other, cp.ndarray)
):
other_array = other
elif func.startswith("__r"):
# The left operand's own method has already declined, so Python
# raises its usual TypeError naming both types.
return NotImplemented
else:
raise NotImplementedError(
f"arithmetic operation not implemented for {type(other).__name__}"
)
self_array = self.array
if (
cp is not None
and isinstance(other_array, (np.ndarray, cp.ndarray, da.core.Array))
and {get_array_module(self_array), get_array_module(other_array)}
== {np, cp}
):
# CuPy refuses a host array as an operand, a 0-d one included,
# though it takes a NumPy scalar. The side on the host, measurement
# or operand, is moved to the device, so both operand orders give
# the same result there. A lazy host operand is copied to the device
# block by block when the result is computed.
if get_array_module(self_array) is cp:
other_array = copy_to_device(other_array, "gpu")
else:
self_array = copy_to_device(self_array, "gpu")
name = func.strip("_")
if isinstance(self_array, da.core.Array) or isinstance(
other_array, da.core.Array
):
# CuPy raises on a dask operand instead of deferring to it, so the
# operation is carried out by the dask array's own method, which
# accepts NumPy and CuPy operands alike.
if isinstance(self_array, da.core.Array):
array = getattr(self_array, func)(other_array)
else: # "__add__" is other.__radd__(self), "__rsub__" other.__sub__(self)
if func.startswith("__r"):
reflected = "__" + func[3:]
else:
reflected = "__r" + func[2:]
array = getattr(other_array, reflected)(self_array)
elif hasattr(operator, name):
# Through the operator module rather than a direct call of the method
# named by func, so that an operand that returns NotImplemented hands
# over to the other operand's reflected method, and a pair neither
# side supports raises TypeError.
array = getattr(operator, name)(self_array, other_array)
else: # a reflected operation: "__rtruediv__" is other / self
array = getattr(operator, name[1:])(other_array, self_array)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
return self.__class__(**kwargs)
def _in_place_arithmetic(
self, other: Self | np.ndarray | da.core.Array | Number, func: str
) -> Self:
if (
self.is_lazy
or (hasattr(other, "is_lazy") and other.is_lazy)
or isinstance(other, da.core.Array)
):
raise RuntimeError(
"inplace arithmetic operation not implemented for lazy measurement or"
" arrays"
)
return self._arithmetic(other, func)
def __mul__(self, other: Self) -> Self:
return self._arithmetic(other, "__mul__")
def __imul__(self, other: Self) -> Self:
return self._in_place_arithmetic(other, "__imul__")
def __truediv__(self, other: Self) -> Self:
return self._arithmetic(other, "__truediv__")
def __itruediv__(self, other: Self) -> Self:
return self._in_place_arithmetic(other, "__itruediv__")
def __rtruediv__(self, other: Self) -> Self:
return self._arithmetic(other, "__rtruediv__")
def __sub__(self, other: Self) -> Self:
return self._arithmetic(other, "__sub__")
def __rsub__(self, other: Self) -> Self:
return self._arithmetic(other, "__rsub__")
def __isub__(self, other: Self) -> Self:
return self._in_place_arithmetic(other, "__isub__")
def __add__(self, other: Self) -> Self:
return self._arithmetic(other, "__add__")
def __iadd__(self, other: Self) -> Self:
return self._in_place_arithmetic(other, "__iadd__")
def __pow__(self, other: Self) -> Self:
return self._arithmetic(other, "__pow__")
def __radd__(self, other: Self) -> Self:
return self._arithmetic(other, "__radd__")
def __rmul__(self, other: Self) -> Self:
return self._arithmetic(other, "__rmul__")
def __rpow__(self, other: Self) -> Self:
return self._arithmetic(other, "__rpow__")
# NumPy scalars and arrays, dask arrays and CuPy arrays on the left of an
# operator defer to the reflected methods above, instead of coercing this
# object to an array through __len__ and __getitem__, because their priority
# is lower. __array_ufunc__ = None would also defer them, but dask treats any
# object with shape, dtype and __array_ufunc__ as array-like and then reads
# its ndim, which an ArrayObject lacks. The same deferral applies to NumPy's
# in-place operators: `ndarray += m` leaves the ndarray unchanged and binds
# the name to the new object returned by `m.__radd__`.
__array_priority__ = 1000
def _get_ensemble_axes_metadata_items(self, items):
expanded_axes_metadatas = [
axis_metadata.copy() for axis_metadata in self.ensemble_axes_metadata
]
for i, item in enumerate(items):
if item is None:
expanded_axes_metadatas.insert(i, UnknownAxis())
metadata = {}
axes_metadata = []
last_indexed = 0
for item, expanded_axes_metadata in zip(items, expanded_axes_metadatas):
last_indexed += 1
if isinstance(item, Number):
metadata = {
**metadata,
**expanded_axes_metadata.item_metadata(item, self.metadata),
}
else:
try:
axes_metadata += [expanded_axes_metadata[item].copy()]
except TypeError:
axes_metadata += [expanded_axes_metadata.copy()]
axes_metadata += expanded_axes_metadatas[last_indexed:]
return axes_metadata, metadata
[docs]
def get_items(
self,
items: ArrayItemType | tuple[ArrayItemType, ...],
keepdims: bool = False,
) -> dict:
"""Index the array and the corresponding axes metadata. Only ensemble axes can
be indexed.
Parameters
----------
items : int or tuple of int or slice
The array is indexed according to this.
keepdims : bool, optional
If True, all ensemble axes are left in the result as dimensions with size
one. Default is False.
Returns
-------
indexed_array : ArrayObject or subclass of ArrayObject
The indexed array object.
"""
items = _validate_array_items(
items, shape=self.ensemble_shape, keepdims=keepdims
)
ensemble_axes_metadata, metadata = self._get_ensemble_axes_metadata_items(items)
kwargs = self._copy_kwargs(
exclude=("array", "ensemble_axes_metadata", "metadata")
)
kwargs["array"] = self._array[items]
kwargs["ensemble_axes_metadata"] = ensemble_axes_metadata
kwargs["metadata"] = {**self.metadata, **metadata}
return kwargs
def __getitem__(self: ArrayObjectType, items) -> ArrayObjectType:
return self.__class__(**self.get_items(items))
[docs]
def expand_dims(
self,
axis: Optional[int | tuple[int, ...]] = None,
axis_metadata: Optional[list[AxisMetadata]] = None,
) -> Self:
"""Expand the shape of the array object.
Parameters
----------
axis : int or tuple of ints
Position in the expanded axes where the new axis (or axes) is placed.
axis_metadata : AxisMetadata or List of AxisMetadata, optional
The axis metadata describing the expanded axes. Default is UnknownAxis.
Returns
-------
expanded : ArrayObject or subclass of ArrayObject
View of array object with the number of dimensions increased.
"""
if axis is None:
axis = (0,)
axis = number_to_tuple(axis)
if axis_metadata is None:
axis_metadata = [UnknownAxis()] * len(axis)
axis = normalize_axes(axis, self.shape)
if any(a >= (len(self.ensemble_shape) + len(axis)) for a in axis):
raise RuntimeError()
ensemble_axes_metadata = copy.deepcopy(self.ensemble_axes_metadata)
for a, am in zip(axis, axis_metadata):
ensemble_axes_metadata.insert(a, am)
kwargs = self._copy_kwargs(exclude=("array", "ensemble_axes_metadata"))
kwargs["array"] = _expand_dims(self.array, axis=axis)
kwargs["ensemble_axes_metadata"] = ensemble_axes_metadata
return self.__class__(**kwargs)
[docs]
def squeeze(self, axis: Optional[tuple[int, ...]] = None) -> Self:
"""Remove axes of length one from array object.
Parameters
----------
axis : int or tuple of ints, optional
Selects a subset of the entries of length one in the shape.
Returns
-------
squeezed : ArrayObject or subclass of ArrayObject
The input array object, but with all or a subset of the dimensions of length
1 removed.
"""
if len(self.array.shape) < len(self.base_shape):
return self
if axis is None:
axis = tuple(range(len(self.shape)))
else:
axis = normalize_axes(axis, self.shape)
# Not `self.shape[: -len(self.base_shape)]`: Python has no negative
# zero, so for a base-less object (base_shape == ()) that slice is
# `[:0]`, always empty, rather than `[:len(self.shape)]`.
n_ensemble_dims = len(self.shape) - len(self.base_shape)
shape = self.shape[:n_ensemble_dims]
squeezed = tuple(
np.where([(n == 1) and (i in axis) for i, n in enumerate(shape)])[0]
)
xp = get_array_module(self.array)
kwargs = self._copy_kwargs(exclude=("array", "ensemble_axes_metadata"))
kwargs["array"] = xp.squeeze(self.array, axis=squeezed)
kwargs["ensemble_axes_metadata"] = [
element
for i, element in enumerate(self.ensemble_axes_metadata)
if i not in squeezed
]
return self.__class__(**kwargs)
[docs]
def ensure_lazy(self, chunks: Chunks = "auto") -> Self:
"""Creates an equivalent lazy version of the array object.
Parameters
----------
chunks : int or tuple or str
How to chunk the array. See `dask.array.from_array`.
Returns
-------
lazy_array_object : ArrayObject or subclass of ArrayObject
Lazy version of the array object.
"""
if self.is_lazy:
return self
if chunks == "auto":
chunks = ("auto",) * len(self.ensemble_shape) + (-1,) * len(self.base_shape)
array = da.from_array(self.array, chunks=chunks)
# by keyword: not every subclass takes the array as its first argument
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
return self.__class__(**kwargs)
def lazy(self, chunks: str = "auto") -> Self:
return self.ensure_lazy(chunks)
[docs]
def compute(
self,
progress_bar: bool | None = None,
profiler: bool = False,
resource_profiler: bool = False,
**kwargs,
) -> Self | tuple[Self, tuple]:
"""Turn a lazy abTEM object into its in-memory equivalent.
Parameters
----------
progress_bar : bool
Display a progress bar in the terminal or notebook during computation. The
progress bar is only displayed with a local scheduler.
profiler : bool
Return Profiler class used to profile Dask's execution at the task level.
Only execution with a local scheduler is profiled.
resource_profiler : bool
Return ResourceProfiler class used to profile Dask’s execution at the
resource level.
kwargs :
Additional keyword arguments passed to `dask.compute`.
"""
if not self.is_lazy:
return self
output, profilers = _compute(
[self],
progress_bar=progress_bar,
profiler=profiler,
resource_profiler=resource_profiler,
**kwargs,
)
output_value = output[0]
if profilers:
return output_value, profilers
return output_value
[docs]
def copy_to_device(self, device: str) -> Self:
"""Copy array to specified device.
Parameters
----------
device : str
Returns
-------
object_on_device : ArrayObject
"""
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = copy_to_device(self.array, device)
return self.__class__(**kwargs)
[docs]
def to_cpu(self) -> Self:
"""Move the array to the host memory from an arbitrary source array."""
return self.copy_to_device("cpu")
[docs]
def to_gpu(self, device: str = "gpu") -> Self:
"""Move the array from the host memory to a gpu."""
return self.copy_to_device(device)
[docs]
def to_zarr(
self, url: str, compute: bool = True, overwrite: bool = False, **kwargs
):
"""Write data to a zarr file.
Parameters
----------
url : str
Location of the data, typically a path to a local file. A URL can also
include a protocol specifier like s3:// for remote data.
compute : bool
If true compute immediately; return dask.delayed.Delayed otherwise.
overwrite : bool
If given array already exists, overwrite=False will cause an error, where
overwrite=True will replace the existing data.
kwargs :
Keyword arguments passed to `dask.array.to_zarr`.
"""
return ComputableList([self]).to_zarr(
url=url, compute=compute, overwrite=overwrite, **kwargs
)
@classmethod
def _pack_kwargs(cls, kwargs):
attrs = {}
for key, value in kwargs.items():
if key == "ensemble_axes_metadata":
attrs[key] = [axis_to_dict(axis) for axis in value]
else:
attrs[key] = value
return attrs
@classmethod
def _unpack_kwargs(cls, attrs):
kwargs = dict()
kwargs["ensemble_axes_metadata"] = []
for key, value in attrs.items():
if key == "ensemble_axes_metadata":
kwargs["ensemble_axes_metadata"] = [axis_from_dict(d) for d in value]
elif key == "type":
pass
else:
kwargs[key] = value
return kwargs
def _metadata_to_dict(self):
metadata = copy.copy(self.metadata)
metadata["axes"] = {
f"axis_{i}": axis_to_dict(axis) for i, axis in enumerate(self.axes_metadata)
}
metadata["data_origin"] = f"abTEM_v{__version__}"
metadata["type"] = self.__class__.__name__
return metadata
def _metadata_to_json(self):
metadata = copy.copy(self.metadata)
metadata["axes"] = {
f"axis_{i}": axis_to_dict(axis) for i, axis in enumerate(self.axes_metadata)
}
metadata["data_origin"] = f"abTEM_v{__version__}"
metadata["type"] = self.__class__.__name__
return json.dumps(metadata)
[docs]
def to_tiff(self, filename: str, **kwargs):
"""Write data to a tiff file.
Parameters
----------
filename : str
The filename of the file to write.
kwargs :
Keyword arguments passed to `tifffile.imwrite`.
"""
if tifffile is None:
raise RuntimeError(
"This functionality of abTEM requires tifffile, see https://github.com/cgohlke/tifffile."
)
array = self.array
if self.is_lazy:
warnings.warn("Lazy arrays are computed in memory before writing to tiff.")
_push_config_to_active_client()
array = self._lazy_array.compute()
return tifffile.imwrite(
filename, array, description=self._metadata_to_json(), **kwargs
)
[docs]
@classmethod
def from_zarr(cls, url: str, chunks: Chunks = "auto") -> Self:
"""Read wave functions from a Zarr file.
url : str
Location of the data, typically a path to a local file. A URL can also
include a protocol specifier like s3:// for remote data.
chunks : tuple of ints or tuples of ints
Passed to dask.array.from_array(), allows setting the chunks on
initialisation, if the chunking scheme in the on-disc dataset is not optimal
for the calculations to follow.
"""
return from_zarr(url, chunks=chunks)
@property
def _has_base_chunks(self) -> bool:
if not isinstance(self.array, da.core.Array):
return False
# Not `chunks[-len(self.base_shape):]`: the same -0 problem, mirrored --
# for a base-less object that slice is `[0:]`, everything, rather than
# `[len(chunks):]`, nothing.
n_ensemble_dims = len(self.array.chunks) - len(self.base_shape)
base_chunks = self.array.chunks[n_ensemble_dims:]
return any(len(c) > 1 for c in base_chunks)
[docs]
def no_base_chunks(self):
"""Rechunk to remove chunks across the base dimensions."""
if not self._has_base_chunks:
return self
# See the comment in _has_base_chunks: -0 is 0, not "the end".
n_ensemble_dims = len(self.array.chunks) - len(self.base_shape)
chunks = self.array.chunks[:n_ensemble_dims] + (-1,) * len(self.base_shape)
return self.rechunk(chunks)
@staticmethod
def _apply_transform(
array: np.ndarray,
*args: tuple[np.ndarray, ...],
transform_partial: Callable,
array_object_partial: Callable,
base_ndims: int,
out_ndim: int,
) -> np.ndarray:
axes = unpack_blockwise_args(args)
array_axes = axes[: len(array.shape)]
# Not `array_axes[:-base_ndims]`: -0 is 0, not "the end" -- see the
# comment in ArrayObject.squeeze for a base-less array_object.
n_ensemble_axes = len(array_axes) - base_ndims
ensemble_axes = array_axes[:n_ensemble_axes]
transform_axes = axes[len(array.shape) :]
array_object = array_object_partial((array, list(ensemble_axes))).item()
transform = transform_partial(*transform_axes).item()
out_arrays = transform._calculate_new_array(array_object)
if not isinstance(out_arrays, tuple):
out_arrays = (out_arrays,)
packing = np.zeros((1,) * out_ndim, dtype=object)
itemset(packing, 0, out_arrays)
return packing
[docs]
def to_hyperspy(self, transpose: bool = True):
"""Convert ArrayObject to a Hyperspy signal.
Parameters
----------
transpose : bool, optional
If True, transpose the base axes of the array before converting to a
Hyperspy signal. Default is True.
Returns
-------
signal : Hyperspy signal
The converted Hyperspy signal.
Raises
------
ImportError
If Hyperspy is not installed.
RuntimeError
If the number of base dimensions is not 1 or 2.
Notes
-----
This method requires Hyperspy to be installed. You can find more information
about Hyperspy at https://hyperspy.org.
"""
if hs is None:
raise ImportError(
"This functionality of *ab*TEM requires Hyperspy, see https://hyperspy.org."
)
if self._base_dims == 1:
signal_type = hs.signals.Signal1D
elif self._base_dims == 2:
signal_type = hs.signals.Signal2D
else:
raise RuntimeError()
axes_base = _to_hyperspy_axes_metadata(
self.base_axes_metadata,
self.base_shape,
)
ensemble_axes_metadata = _to_hyperspy_axes_metadata(
self.ensemble_axes_metadata,
self.ensemble_shape,
)
xp = get_array_module(self.device)
axes_base_indices = tuple_range(
offset=len(self.ensemble_shape), length=len(self.base_shape)
)
ensemble_axes = tuple_range(
offset=0,
length=len(self.ensemble_shape),
)
source = ensemble_axes + axes_base_indices
destination = ensemble_axes + axes_base_indices[::-1]
if transpose:
if self.is_lazy:
array = da.moveaxis(self.array, source=source, destination=destination)
else:
array = xp.moveaxis(self.array, source=source, destination=destination)
else:
array = self.array
s = signal_type(array, axes=ensemble_axes_metadata[::-1] + axes_base[::-1])
if self.is_lazy:
s = s.as_lazy()
return s
[docs]
def to_data_array(self):
"""Convert ArrayObject to a xarray DataArray. Requires xarray to be installed.
Returns
-------
xarray.DataArray
The converted xarray DataArray.
Raises
------
ImportError
If xarray is not installed.
"""
if xr is None:
raise ImportError(
"This functionality of *ab*TEM requires xarray, see https://xarray.dev/."
)
coords = {}
dims = []
for n, axis in zip(self.shape, self.axes_metadata):
x = np.array(axis.coordinates(n))
if isinstance(x, np.ndarray) and len(x.shape) == 2:
x = [f"{i}" for i in x]
elif len(x.shape) == 1:
pass
else:
raise ValueError("The shape of the coordinates is not supported.")
dims.append(axis.label)
dr = xr.DataArray(
x, name=axis.label, dims=(axis.label,), attrs={"units": axis.units}
)
coords[axis.label] = dr
attrs = copy.deepcopy(self.metadata)
if "label" in attrs:
attrs["long_name"] = attrs["label"]
return xr.DataArray(self.array, dims=dims, coords=coords, attrs=attrs)
[docs]
def to_quantem(self):
"""
Convert ArrayObject to quantem Dataset object.
The units of the axes are converted to quantem's spelling (for example "A" and
"A^-1"), and an axis without units gets "pixels". The signal units are the
"units" entry of the metadata, passed on unchanged; if there is none they are
"arb. unit". The metadata is deep-copied into the metadata of the Dataset, so
changing one does not change the other.
A quantem Dataset holds NumPy arrays only: a lazy object is computed (the
object itself stays lazy) and an array on a GPU is copied to the host.
Returns
-------
dataset : quantem.core.datastructures.Dataset
"""
if em is None:
raise ImportError(
"This functionality of *ab*TEM requires quantem, see https://github.com/electronmicroscopy/quantem."
)
def _normalize_unit(unit):
if unit is None:
return "pixels"
UNIT_MAP = {
"Ã…": "A",
"Ångström": "A",
"Angstrom": "A",
"1/Ã…": "A^-1",
"Ã…^-1": "A^-1",
"1/A": "A^-1",
}
return UNIT_MAP.get(unit, unit)
sampling = []
origin = []
units = []
for axis in self.axes_metadata:
sampling.append(getattr(axis, "sampling", 1.0))
units.append(_normalize_unit(getattr(axis, "units", None)))
origin.append(0.0)
# The signal units are the metadata units, as abTEM's own measurements set them
# ("arb. unit" by default). quantem only displays them, and reads the same
# symbols as abTEM ("Ã…", "%", "electrons"), so they are not normalized like
# the axis units. An empty string is a dimensionless signal and stays empty.
signal_units = self.metadata.get("units")
if signal_units is None:
signal_units = "arb. unit"
# compute() replaces the array of the object it is called on; computing a
# shallow copy leaves the caller lazy.
exported = copy.copy(self).compute() if self.is_lazy else self
# A quantem Dataset holds NumPy arrays only: it rejects CuPy arrays.
dataset = em.core.datastructures.Dataset.from_array(
array=asnumpy(exported.array),
name=type(self).__name__,
origin=tuple(origin),
sampling=tuple(sampling),
units=tuple(units),
signal_units=signal_units,
)
# Dataset.from_array takes no metadata; Dataset.metadata is the dict it
# holds, empty for a new dataset. Deep-copied, as in to_data_array, so that
# nested values are not shared with this object.
dataset.metadata.update(copy.deepcopy(self.metadata))
return dataset
@classmethod
def _stack(
cls,
array_objects: tuple[ArrayObject, ...],
axis_metadata: AxisMetadata,
axis: int,
) -> Self:
xp = get_array_module(array_objects[0].device)
if any(array.is_lazy for array in array_objects):
array = da.stack(
[array_object.array for array_object in array_objects], axis=axis
)
else:
array = xp.stack(
[array_object.array for array_object in array_objects], axis=axis
)
kwargs = array_objects[0]._copy_kwargs(exclude=("array",))
kwargs["array"] = array
ensemble_axes_metadata = [
axis_metadata.copy() for axis_metadata in kwargs["ensemble_axes_metadata"]
]
ensemble_axes_metadata.insert(axis, axis_metadata)
kwargs["ensemble_axes_metadata"] = ensemble_axes_metadata
return cls(**kwargs)
def _partition_ensemble_axes_metadata(
self, chunks: Optional[Chunks] = None, lazy: bool = True
):
if len(self.ensemble_shape) == 0:
ensemble_axes_metadata = _wrap_with_array([], 0)
else:
chunks = self._validate_ensemble_chunks(chunks)
chunk_shape = tuple(len(c) for c in chunks)
ensemble_axes_metadata = np.zeros(chunk_shape, dtype=object)
for index, slic in iterate_chunk_ranges(chunks):
new_ensemble_axes_metadata = [
axis[slic[i]] if hasattr(axis, "__getitem__") else axis.copy()
for i, axis in enumerate(self.ensemble_axes_metadata)
]
itemset(ensemble_axes_metadata, index, new_ensemble_axes_metadata)
if lazy:
ensemble_axes_metadata = da.from_array(ensemble_axes_metadata, chunks=1)
return ensemble_axes_metadata
@property
def _default_ensemble_chunks(self):
if self.is_lazy:
return self._lazy_array.chunks[: self.ensemble_dims]
else:
return -1
def _partition_args(self, chunks: Optional[Chunks] = None, lazy: bool = True):
if chunks is None and self.is_lazy:
# See the comment in ArrayObject.squeeze: -0 is 0, not "the end".
n_ensemble_dims = len(self._lazy_array.chunks) - len(self.base_shape)
chunks = self._lazy_array.chunks[:n_ensemble_dims]
elif chunks is None:
chunks = (1,) * len(self.ensemble_shape)
ensemble_chunks = self._validate_ensemble_chunks(chunks)
if lazy:
xp = get_array_module(self.array)
array = self.ensure_lazy()._lazy_array
if ensemble_chunks != array.chunks:
array = array.rechunk(
ensemble_chunks + array.chunks[len(ensemble_chunks) :]
)
ensemble_axes_metadata = self._partition_ensemble_axes_metadata(
chunks=ensemble_chunks
)
def _combine_args(*args):
combined = args[0], args[1].item()
return _wrap_with_array(combined, 1)
ndims = len(self.ensemble_shape)
blocks = da.blockwise(
_combine_args,
tuple_range(ndims),
array,
tuple_range(len(array.shape)),
ensemble_axes_metadata,
tuple_range(ndims),
align_arrays=False,
concatenate=True,
dtype=object,
meta=np.array((), dtype=object),
)
else:
# compute() works in place: materialise a copy, so that partitioning
# (e.g. of a lazy potential for an eager multislice) leaves the
# caller's object lazy
array = self.copy().compute().array if self.is_lazy else self.array
if len(self.ensemble_shape) == 0:
blocks = np.zeros((), dtype=object)
else:
chunk_shape = tuple(len(c) for c in chunks)
blocks = np.zeros(chunk_shape, dtype=object)
ensemble_axes_metadata = self._partition_ensemble_axes_metadata(
chunks, lazy=False
)
for block_indices, chunk_range in iterate_chunk_ranges(chunks):
itemset(
blocks,
block_indices,
(array[chunk_range], ensemble_axes_metadata[block_indices]),
)
return (blocks,)
@classmethod
def _from_partitioned_args_func(cls, *args, **kwargs):
args = unpack_blockwise_args(args)
array, ensemble_axes_metadata = args[0]
assert isinstance(ensemble_axes_metadata, list)
new_array_object = cls(
array=array, ensemble_axes_metadata=ensemble_axes_metadata, **kwargs
)
ndims = new_array_object.ensemble_dims
return _wrap_with_array(new_array_object, ndims)
def _from_partitioned_args(self):
return partial(
self._from_partitioned_args_func,
**self._copy_kwargs(exclude=("array", "ensemble_axes_metadata")),
)
def _expand_dims(
array: np.ndarray | da.core.Array, axis: int | tuple | list
) -> np.ndarray:
if isinstance(axis, int):
axis = (axis,)
out_ndim = len(axis) + array.ndim
axis = validate_axis(axis, out_ndim)
assert not isinstance(axis, int)
shape_it = iter(array.shape)
shape = [1 if ax in axis else next(shape_it) for ax in range(out_ndim)]
return array.reshape(shape)
[docs]
def from_zarr(url: str, chunks: Optional[Chunks] = None):
"""Read abTEM data from zarr (supports legacy and canonical formats).
Parameters
----------
url : str
Location of the data. A URL can include a protocol specifier like s3:// for
remote data. For Zarr 3 zip stores, this should be a path to a .zip file.
chunks : tuple of ints or tuples of ints
Passed to dask.array.from_array(), allows setting the chunks on initialisation,
if the chunking scheme in the on-disc dataset is not optimal for the
calculations to follow.
Returns
-------
imported : ArrayObject
"""
import zarr
# --- Helper: decode tuples ---
def decode_types(obj):
if isinstance(obj, dict):
if obj.get("_type") == "tuple":
return tuple(decode_types(v) for v in obj["_value"])
return {k: decode_types(v) for k, v in obj.items()}
elif isinstance(obj, list):
return [decode_types(v) for v in obj]
else:
return obj
# --- Open store ---
is_zip = url.endswith(".zip")
if is_zip:
store = zarr.storage.ZipStore(url, mode="r")
root = zarr.open(store=store, mode="r")
else:
root = zarr.open(url, mode="r")
# --- Detect format ---
if "metadata0" in root.attrs:
return _from_zarr_canonical(root, chunks, decode_types)
elif "kwargs0" in root.attrs:
return _from_zarr_legacy(root, chunks, decode_types)
else:
raise ValueError("Zarr file does not contain recognizable abTEM metadata.")
def _from_zarr_canonical(root, chunks, decode_types):
import dask.array as da
import abtem
imported = []
i = 0
while True:
key = f"metadata{i}"
if key not in root.attrs:
break
metadata = decode_types(root.attrs[key]).copy()
metadata.pop("data_origin", None)
cls_name = metadata.pop("type")
cls = getattr(abtem, cls_name)
kwargs = metadata.pop("kwargs", None)
axes_dict = metadata.pop("axes")
# Preserve axis order
axes_list = [
axis_from_dict(axes_dict[k])
for k in sorted(axes_dict, key=lambda x: int(x.split("_")[1]))
]
zarr_array = root[f"array{i}"]
if chunks == "auto":
# Respect cls._base_dims the same way the legacy loader below
# already does: dask's own "auto" heuristic doesn't know which
# trailing axes are an ArrayObject's base (measurement) axes, and
# several of abTEM's own lazy operations (e.g. interpolate_line)
# assume those are never split across chunks. chunks=None (the
# default) already avoids this by mirroring whatever to_zarr
# actually wrote, which itself never splits base axes -- this
# branch only matters if a caller explicitly opts into "auto".
num_ensemble_axes = zarr_array.ndim - cls._base_dims
array_chunks = ("auto",) * num_ensemble_axes + (-1,) * cls._base_dims
elif chunks is None:
array_chunks = zarr_array.chunks
else:
array_chunks = chunks
array = da.from_array(zarr_array, chunks=array_chunks)
try:
obj = cls.from_array_and_metadata(
array,
axes_metadata=axes_list,
metadata=metadata,
)
except Exception:
if kwargs is None:
raise
kwargs = cls._unpack_kwargs(kwargs)
kwargs["array"] = array
kwargs["metadata"] = metadata
obj = cls(**kwargs)
imported.append(obj)
i += 1
return imported[0] if len(imported) == 1 else imported
def _from_zarr_legacy(root, chunks, decode_types):
import dask.array as da
from dask import config
import abtem
imported = []
i = 0
types = []
# Collect legacy types
while True:
key = f"type{i}"
if key not in root.attrs:
break
types.append(root.attrs[key])
i += 1
for i, cls_name in enumerate(types):
# Legacy files (written by abTEM <= 1.0.9) may hold any ArrayObject
# subclass, e.g. Waves or PotentialArray, not just measurements; all
# of them are exported from the top-level abtem namespace.
cls = getattr(abtem, cls_name)
packed_kwargs = decode_types(root.attrs[f"kwargs{i}"])
kwargs = cls._unpack_kwargs(packed_kwargs)
num_ensemble_axes = len(kwargs["ensemble_axes_metadata"])
zarr_array = root[f"array{i}"]
if chunks == "auto":
array_chunks = ("auto",) * num_ensemble_axes + (-1,) * cls._base_dims
elif chunks is None:
array_chunks = zarr_array.chunks
else:
array_chunks = chunks
array = da.from_array(zarr_array, chunks=array_chunks)
with config.set({"warnings.overspecified-grid": False}):
imported.append(cls(array, **kwargs))
return imported[0] if len(imported) == 1 else imported
[docs]
def stack(
arrays: Sequence[ArrayObjectType],
axis_metadata: Optional[AxisMetadata | Sequence[str] | dict] = None,
axis: int = 0,
) -> ArrayObjectType:
"""Stack multiple array objects (e.g. Waves and BaseMeasurement) along a new
ensemble axis.
Parameters
----------
arrays : sequence of array objects
Each abTEM array object must have the same type and shape.
axis_metadata : AxisMetadata
The axis metadata describing the new axis.
axis : int
The ensemble axis in the resulting array object along which the input arrays are
stacked.
Returns
-------
array_object : ArrayObject
The stacked array object of the same type as the input.
"""
# if not all(isinstance(array, ArrayObject) for array in arrays):
# raise ValueError("arrays must be a sequence of array objects.")
assert axis <= len(arrays[0].ensemble_shape)
assert axis >= 0
axis_metadata = validate_axis_metadata(axis_metadata)
return arrays[0]._stack(arrays, axis_metadata, axis)
[docs]
def concatenate(arrays: Sequence[ArrayObject], axis: int = 0) -> ArrayObject:
"""Join a sequence of abTEM array classes along an existing axis.
Parameters
----------
arrays : sequence of array objects
Each abTEM array object must have the same type and shape, except in the
dimension corresponding to axis. The axis metadata along the concatenated
axis must be compatible for concatenation.
axis : int, optional
The axis along which the arrays will be joined. Default is 0.
Returns
-------
array_object : ArrayObject
The concatenated array object of the same type as the input.
"""
xp = get_array_module(arrays[0].array)
if any(has_array.is_lazy for has_array in arrays):
array = da.concatenate([has_array.array for has_array in arrays], axis=axis)
else:
array = xp.concatenate([has_array.array for has_array in arrays], axis=axis)
cls = arrays[0].__class__
concatenated_axes_metadata = arrays[0].axes_metadata[axis]
for has_array in arrays[1:]:
concatenated_axes_metadata = concatenated_axes_metadata.concatenate(
has_array.axes_metadata[axis]
)
axes_metadata = copy.deepcopy(arrays[0].axes_metadata)
axes_metadata[axis] = concatenated_axes_metadata
return cls.from_array_and_metadata(
array=array, axes_metadata=axes_metadata, metadata=arrays[0].metadata
)
[docs]
def swapaxes(array_object, axis1, axis2):
xp = get_array_module(array_object.array)
if array_object.is_lazy:
array = da.swapaxes(array_object.array, axis1, axis2)
else:
array = xp.swapaxes(array_object.array, axis1, axis2)
cls = array_object.__class__
axes_metadata = copy.copy(array_object.axes_metadata)
axes_metadata[axis2], axes_metadata[axis1] = (
axes_metadata[axis1],
axes_metadata[axis2],
)
return cls.from_array_and_metadata(
array=array, axes_metadata=axes_metadata, metadata=array_object.metadata
)
[docs]
def moveaxis(
array_object: ArrayObject,
source: tuple[int, ...],
destination: tuple[int, ...],
) -> ArrayObject:
xp = get_array_module(array_object.array)
if array_object.is_lazy:
array = da.moveaxis(array_object.array, source, destination)
else:
array = xp.moveaxis(array_object.array, source, destination)
axes_metadata = copy.copy(array_object.axes_metadata)
for s, d in zip(reversed(source), reversed(destination)):
element = axes_metadata.pop(s)
axes_metadata.insert(d, element)
return array_object.__class__.from_array_and_metadata(
array=array, axes_metadata=axes_metadata, metadata=array_object.metadata
)
