"""Module for describing electrostatic potentials using the independent atom model."""
from __future__ import annotations
import warnings
from abc import ABCMeta, abstractmethod
from functools import partial, reduce
from numbers import Number
from operator import mul
from typing import TYPE_CHECKING, Optional, Sequence, Type
import dask
import dask.array as da
import numpy as np
from ase import Atoms
from ase.cell import Cell
from ase.data import chemical_symbols
from abtem.array import ArrayObject, validate_lazy
from abtem.atoms import (
best_orthogonal_cell,
cut_cell,
is_cell_orthogonal,
orthogonalize_cell,
pad_atoms,
plane_to_axes,
rotate_atoms_to_plane,
)
from abtem.core.axes import (
AxisMetadata,
FrozenPhononsAxis,
RealSpaceAxis,
ThicknessAxis,
_find_axes_type,
)
from abtem.core.backend import get_array_module, validate_device
from abtem.core.chunks import Chunks, chunk_ranges, generate_chunks, validate_chunks
from abtem.core.complex import complex_exponential
from abtem.core.energy import Accelerator, HasAcceleratorMixin, energy2sigma
from abtem.core.ensemble import Ensemble, _wrap_with_array, unpack_blockwise_args
from abtem.core.grid import Grid, HasGrid2DMixin
from abtem.core.utils import CopyMixin, EqualityMixin, get_dtype, itemset
from abtem.inelastic.phonons import (
AtomsEnsemble,
BaseFrozenPhonons,
DummyFrozenPhonons,
FrozenPhonons,
validate_seeds,
)
from abtem.integrals import (
QuadratureProjectionIntegrals,
ScatteringFactorProjectionIntegrals,
)
from abtem.measurements import Images
from abtem.slicing import (
BaseSlicedAtoms,
SlicedAtoms,
SliceIndexedAtoms,
_validate_slice_thickness,
slice_limits,
)
if TYPE_CHECKING:
from abtem.integrals import FieldIntegrator
from abtem.parametrizations import Parametrization
from abtem.waves import BaseWaves, Waves
[docs]
class BaseField(Ensemble, HasGrid2DMixin, EqualityMixin, CopyMixin, metaclass=ABCMeta):
# @property
# @abstractmethod
# def device(self) -> str:
# pass
@property
def base_shape(self):
"""Shape of the base axes of the potential."""
return (self.num_slices,) + self.gpts
@property
@abstractmethod
def num_configurations(self):
"""Number of frozen phonons in the ensemble of potentials."""
pass
@property
@abstractmethod
def base_axes_metadata(self):
pass
def _get_exit_planes_axes_metadata(self):
return ThicknessAxis(label="z", values=tuple(self.exit_thicknesses))
@property
@abstractmethod
def exit_planes(self) -> tuple[int, ...]:
"""The "exit planes" of the potential. The indices of slices where a measurement
is returned."""
pass
@property
def _exit_plane_after(self):
exit_plane_index = 0
exit_planes = self.exit_planes
if exit_planes[0] == -1:
exit_plane_index += 1
is_exit_plane = np.zeros(len(self), dtype=bool)
for i in range(len(is_exit_plane)):
if exit_plane_index < len(exit_planes) and i == exit_planes[exit_plane_index]:
is_exit_plane[i] = True
exit_plane_index += 1
return is_exit_plane
@property
def exit_thicknesses(self) -> tuple[float, ...]:
"""The "exit thicknesses" of the potential. The thicknesses in the potential
where a measurement is returned."""
thicknesses = np.cumsum(self.slice_thickness)
exit_indices = np.array(self.exit_planes, dtype=int)
exit_thicknesses = tuple(thicknesses[i] for i in exit_indices)
if self.exit_planes[0] == -1:
return (0.0,) + exit_thicknesses[1:]
else:
return exit_thicknesses
@property
def num_exit_planes(self) -> int:
"""Number of exit planes."""
return len(self.exit_planes)
@abstractmethod
def generate_slices(self, first_slice: int = 0, last_slice: Optional[int] = None):
pass
@abstractmethod
def build(
self,
first_slice: int = 0,
last_slice: Optional[int] = None,
chunks: int = 1,
lazy: Optional[bool] = None,
):
pass
def __len__(self) -> int:
return self.num_slices
@property
def num_slices(self) -> int:
"""Number of projected potential slices."""
return len(self.slice_thickness)
@property
@abstractmethod
def slice_thickness(self) -> tuple[float, ...]:
"""Slice thicknesses for each slice."""
pass
@property
def slice_limits(self) -> list[tuple[float, float]]:
"""The entrance and exit thicknesses of each slice [Å]."""
return slice_limits(self.slice_thickness)
@property
def thickness(self) -> float:
"""Thickness of the potential [Å]."""
return sum(self.slice_thickness)
def __iter__(self):
for slic in self.generate_slices():
yield slic
[docs]
def project(self) -> Images:
"""
Sum of the potential slices as an image.
Returns
-------
projected : Images
The projected potential.
"""
return self.build().project()
@property
def _default_ensemble_chunks(self) -> tuple:
return validate_chunks(self.ensemble_shape, (1,) * len(self.ensemble_shape))
[docs]
def to_images(self):
"""
Converts the potential to an ensemble of images.
Returns
-------
image_ensemble : Images
The potential slices as images.
"""
return self.build().to_images()
[docs]
def show(self, project: bool = True, **kwargs):
"""
Show the potential projection. This requires building all potential slices.
Parameters
----------
project : bool, optional
Show the projected potential (True, default) or show all potential slices.
It is recommended to index a subset of the potential slices when this
keyword set to False.
kwargs :
Additional keyword arguments for the show method of :class:`.Images`.
"""
kwargs.setdefault("interpolation", "antialiased")
if project:
return self.project().show(**kwargs)
else:
if "explode" not in kwargs.keys():
kwargs["explode"] = True
return self.to_images().show(**kwargs)
[docs]
def depth_profile(
self,
projection_axis: str = "y",
depth: Optional[float] = None,
) -> Images:
"""Create a depth profile by projecting the potential along a spatial axis.
Parameters
----------
projection_axis : str
Spatial axis to project (sum) along. ``"y"`` (default) produces an
x–z cross-section; ``"x"`` produces a y–z cross-section.
depth : float, optional
If given, project only over a finite slab of this thickness [Å],
centered on the midpoint of the projected axis. The number of grid
points is rounded to the nearest integer. If ``None``, the full
extent is projected.
Returns
-------
depth_profile : Images
2D image(s) with the remaining spatial axis horizontal and depth
(z) vertical. Any ensemble axes (e.g. frozen phonons) are preserved.
"""
return self.build().depth_profile(
projection_axis=projection_axis,
depth=depth,
)
[docs]
def show_depth_profile(
self,
projection_axis: str = "y",
depth: Optional[float] = None,
z_scale: float = 1.0,
slice_lines: bool = True,
ax=None,
cbar: bool = False,
cmap: Optional[str] = None,
vmin: Optional[float] = None,
vmax: Optional[float] = None,
power: float = 1.0,
common_color_scale: bool = False,
explode: bool | Sequence[int] = (),
figsize: Optional[tuple[int, int]] = None,
title: bool | str = True,
**kwargs,
):
"""Show a depth cross-section of the potential.
Parameters
----------
projection_axis : str
Spatial axis to project (sum) along. ``"y"`` (default) produces an
x–z cross-section; ``"x"`` produces a y–z cross-section.
depth : float, optional
If given, project only over a finite slab of this thickness [Å],
centered on the midpoint of the projected axis. The number of grid
points is rounded to the nearest integer. If ``None``, the full
extent is projected.
z_scale : float
Scaling factor for the z-axis relative to the spatial axis.
Values less than 1 compress the z-axis, making panels of thick
specimens more compact. Default is 1.0 (equal scaling).
slice_lines : bool
If True (default), draw horizontal lines at slice boundaries.
ax : matplotlib.axes.Axes, optional
If given the plot is added to the axis.
cbar : bool, optional
Add a colorbar to the plot. Default is False.
cmap : str, optional
Matplotlib colormap name.
vmin : float, optional
Minimum of the color scale.
vmax : float, optional
Maximum of the color scale.
power : float
Show image on a power scale.
common_color_scale : bool, optional
If True, all images in a grid share the same color scale.
explode : bool or sequence of int, optional
If True, create a grid of images for ensemble items.
figsize : two int, optional
Figure size as (width, height) in inches.
title : bool or str, optional
Column title for the images.
**kwargs
Additional keyword arguments passed to the show method.
Returns
-------
visualization : Visualization
"""
from abtem.visualize import Visualization
profile = self.depth_profile(
projection_axis=projection_axis,
depth=depth,
)
if figsize is None and ax is None:
spatial_extent = profile.extent[0]
z_extent = profile.extent[1]
if explode is True or (isinstance(explode, Sequence) and explode):
n_panels = (
profile.ensemble_shape[0] if profile.ensemble_shape else 1
)
else:
n_panels = 1
visual_ratio = (z_extent * z_scale) / spatial_extent
panel_width = 3.0
panel_height = panel_width * visual_ratio
if panel_height < 1.0:
panel_width = min(5.0, 1.0 / visual_ratio)
panel_height = panel_width * visual_ratio
elif panel_height > 8.0:
panel_height = 8.0
panel_width = panel_height / visual_ratio
figsize = (
panel_width * n_panels + 1.0 * n_panels + 0.5,
max(2.5, panel_height + 1.5),
)
visualization = Visualization(
measurement=profile,
ax=ax,
common_scale=common_color_scale,
figsize=figsize,
title=title,
aspect=False,
share_x=True,
share_y=True,
explode=explode,
overlay=(),
interactive=True,
value_limits=(vmin, vmax),
power=power,
cmap=cmap,
cbar=cbar,
**kwargs,
)
spatial_label = "x" if projection_axis == "y" else "y"
visualization.set_xlabel(f"{spatial_label} [Å]")
visualization.set_ylabel("z [Å]")
z_sampling = profile.sampling[1]
for idx in np.ndindex(visualization.axes.shape):
artist = visualization.artists[idx]
xlim = artist.get_xlim()
ylim = artist.get_ylim()
artist.set_extent(
(xlim[0], xlim[1], ylim[0] + z_sampling / 2, ylim[1] + z_sampling / 2)
)
visualization.adjust_coordinate_limits_to_artists()
for idx in np.ndindex(visualization.axes.shape):
visualization.axes[idx].set_aspect(z_scale)
if slice_lines:
limits = self.slice_limits
z_boundaries = sorted({z for lo, hi in limits for z in (lo, hi)})
for idx in np.ndindex(visualization.axes.shape):
for z in z_boundaries:
visualization.axes[idx].axhline(
z, color="white", linewidth=0.5, alpha=0.5
)
return visualization
[docs]
class BasePotential(BaseField, metaclass=ABCMeta):
"""Base class of all potentials. Documented in the subclasses."""
@property
def base_axes_metadata(self):
"""List of AxisMetadata for the base axes."""
return [
ThicknessAxis(
label="z", values=tuple(np.cumsum(self.slice_thickness)), units="Å"
),
RealSpaceAxis(
label="x", sampling=self.sampling[0], units="Å", endpoint=False
),
RealSpaceAxis(
label="y", sampling=self.sampling[1], units="Å", endpoint=False
),
]
[docs]
def validate_potential(
potential: Atoms | BasePotential, waves: Optional[BaseWaves] = None
) -> BasePotential:
if isinstance(potential, (Atoms, BaseFrozenPhonons)):
device = None
if waves is not None:
device = waves.device
potential = Potential(potential, device=device)
# elif not isinstance(potential, BasePotential):
# raise ValueError()
if waves is not None and potential is not None:
potential.grid.match(waves)
return potential
def _validate_exit_planes(exit_planes, num_slices):
if isinstance(exit_planes, int):
if exit_planes >= num_slices:
return (num_slices - 1,)
exit_planes = list(range(exit_planes - 1, num_slices, exit_planes))
if exit_planes[-1] != (num_slices - 1):
exit_planes.append(num_slices - 1)
exit_planes = (-1,) + tuple(exit_planes)
elif exit_planes is None:
exit_planes = (num_slices - 1,)
return exit_planes
def _require_cell_transform(cell, box, plane, origin):
if box == tuple(np.diag(cell)):
return False
if not is_cell_orthogonal(cell):
return True
if box is not None:
return True
if plane != "xy":
return True
if origin != (0.0, 0.0, 0.0):
return True
return False
class _FieldBuilder(BaseField):
def __init__(
self,
array_object: Type[FieldArray],
slice_thickness: float | tuple[float, ...],
cell: np.ndarray | Cell,
exit_planes: Optional[int | tuple[int, ...]] = None,
gpts: Optional[int | tuple[int, int]] = None,
sampling: Optional[float | tuple[float, float]] = None,
box: Optional[tuple[float, float, float]] = None,
plane: (
str | tuple[tuple[float, float, float], tuple[float, float, float]]
) = "xy",
origin: tuple[float, float, float] = (0.0, 0.0, 0.0),
periodic: bool = True,
device: Optional[str] = None,
):
self._array_object = array_object
if _require_cell_transform(cell, box=box, plane=plane, origin=origin):
if not isinstance(plane, str):
raise NotImplementedError
axes = plane_to_axes(plane)
cell = np.array(cell)[:, list(axes)]
box = tuple(best_orthogonal_cell(cell))
elif box is None:
box = tuple(np.diag(cell))
self._grid = Grid(
extent=box[:2], gpts=gpts, sampling=sampling, lock_extent=True
)
self._device = validate_device(device)
self._box = box
self._plane = plane
self._origin = origin
self._periodic = periodic
self._slice_thickness = _validate_slice_thickness(
slice_thickness, thickness=box[2]
)
self._exit_planes = _validate_exit_planes(
exit_planes, len(self._slice_thickness)
)
@property
def slice_thickness(self) -> tuple[float, ...]:
return self._slice_thickness
@property
def exit_planes(self) -> tuple[int]:
return self._exit_planes
@property
def device(self) -> str:
"""The device where the potential is created."""
return self._device
@property
def periodic(self) -> bool:
"""Specifies whether the potential is periodic."""
return self._periodic
@property
def plane(
self,
) -> str | tuple[tuple[float, float, float], tuple[float, float, float]]:
"""The plane relative to the atoms mapped to `xy` plane of the potential,
i.e. the plane is perpendicular to the propagation direction."""
return self._plane
@property
def box(self) -> tuple[float, float, float]:
"""The extent of the potential in `x`, `y` and `z`."""
return self._box
@property
def origin(self) -> tuple[float, float, float]:
"""The origin relative to the provided atoms mapped to the origin of the
potential."""
return self._origin
def __getitem__(self, item) -> PotentialArray:
return self.build(lazy=False)[item]
@staticmethod
def _wrap_build_potential(potential, first_slice, last_slice):
potential = potential.item()
array = potential.build(first_slice, last_slice, lazy=False).array
return array
def build(
self,
first_slice: int = 0,
last_slice: Optional[int] = None,
max_batch: int | str = 1,
lazy: Optional[bool] = None,
) -> FieldArray:
"""
Build the potential.
Parameters
----------
first_slice : int, optional
Index of the first slice of the generated potential.
last_slice : int, optional
Index of the last slice of the generated potential
max_batch : int or str, optional
Maximum number of slices to calculate in task. Default is 1.
lazy : bool, optional
If True, create the wave functions lazily, otherwise, calculate instantly.
If None, this defaults to the value set in the configuration file.
Returns
-------
potential_array : PotentialArray
The built potential as an array.
"""
lazy = validate_lazy(lazy)
self.grid.check_is_defined()
if last_slice is None:
last_slice = len(self)
if lazy:
blocks = self.ensemble_blocks(self._default_ensemble_chunks)
xp = get_array_module(self.device)
chunks = validate_chunks(self.ensemble_shape, self._default_ensemble_chunks)
chunks = chunks + self.base_shape
if self.ensemble_shape:
new_axis = tuple(
range(
len(self.ensemble_shape),
len(self.ensemble_shape) + len(self.base_shape),
)
)
else:
new_axis = tuple(range(0, len(self.base_shape)))
# new_axis = (0, 1, 2) # This was causing problems with FrozenPhonons
array = da.map_blocks(
self._wrap_build_potential,
blocks,
new_axis=new_axis,
first_slice=first_slice,
last_slice=last_slice,
chunks=chunks,
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
xp = get_array_module(self.device)
array = xp.zeros(
self.ensemble_shape + (last_slice - first_slice,) + self.base_shape[1:],
dtype=get_dtype(complex=False),
)
if self.ensemble_shape:
for i, _, potential_wrapped in self.generate_blocks(1):
potential = potential_wrapped.item()
for j, slic in enumerate(
potential.generate_slices(first_slice, last_slice)
):
array[i + (j,)] = slic.array[0]
else:
for j, slic in enumerate(self.generate_slices(first_slice, last_slice)):
array[j] = slic.array[0]
output_potential = self._array_object(
array,
sampling=self._valid_sampling,
slice_thickness=self.slice_thickness[first_slice:last_slice],
exit_planes=self.exit_planes,
ensemble_axes_metadata=self.ensemble_axes_metadata,
)
return output_potential
class _FieldBuilderFromAtoms(_FieldBuilder):
def __init__(
self,
atoms: Atoms | BaseFrozenPhonons,
array_object: Type[FieldArray],
gpts: Optional[int | tuple[int, int]] = None,
sampling: Optional[float | tuple[float, float]] = None,
slice_thickness: float | tuple[float, ...] = 1,
exit_planes: Optional[int | tuple[int, ...]] = None,
plane: (
str | tuple[tuple[float, float, float], tuple[float, float, float]]
) = "xy",
origin: tuple[float, float, float] = (0.0, 0.0, 0.0),
box: Optional[tuple[float, float, float]] = None,
periodic: bool = True,
integrator=None,
device: Optional[str] = None,
):
self._frozen_phonons = _validate_frozen_phonons(atoms)
self._integrator = integrator
self._sliced_atoms: Optional[BaseSlicedAtoms] = None
self._array_object = array_object
super().__init__(
array_object=array_object,
gpts=gpts,
sampling=sampling,
cell=self._frozen_phonons.cell,
slice_thickness=slice_thickness,
exit_planes=exit_planes,
device=device,
plane=plane,
origin=origin,
box=box,
periodic=periodic,
)
@property
def frozen_phonons(self) -> BaseFrozenPhonons:
"""Ensemble of atomic configurations representing frozen phonons."""
return self._frozen_phonons
@property
def num_configurations(self) -> int:
"""Size of the ensemble of atomic configurations representing frozen phonons."""
return len(self.frozen_phonons)
@property
def integrator(self):
"""The integrator determining how the projection integrals for each slice is
calculated."""
return self._integrator
def _cutoffs(self):
atoms = self.frozen_phonons.atoms
unique_numbers = np.unique(atoms.numbers)
return tuple(
self._integrator.cutoff(chemical_symbols[number])
for number in unique_numbers
)
def get_transformed_atoms(self):
"""
The atoms used in the multislice algorithm, transformed to the given plane,
origin and box.
Returns
-------
transformed_atoms : Atoms
Transformed atoms.
"""
atoms = self.frozen_phonons.atoms
if is_cell_orthogonal(atoms.cell) and self.plane != "xy":
atoms = rotate_atoms_to_plane(atoms, self.plane)
elif tuple(np.diag(atoms.cell)) != self.box:
if self.periodic:
atoms = orthogonalize_cell(
atoms,
box=self.box,
plane=self.plane,
origin=self.origin,
return_transform=False,
allow_transform=True,
)
return atoms
else:
cutoffs = self._cutoffs()
atoms = cut_cell(
atoms,
cell=self.box,
plane=self.plane,
origin=self.origin,
margin=max(cutoffs) if cutoffs else 0.0,
)
return atoms
def _prepare_atoms(self):
atoms = self.get_transformed_atoms()
if self.integrator.finite:
cutoffs = self._cutoffs()
margins = max(cutoffs) if len(cutoffs) else 0.0
else:
margins = 0.0
if self.periodic:
atoms = self.frozen_phonons.randomize(atoms)
atoms.wrap(eps=0.0)
# wrap(eps=0.0) uses strict modulo: z positions that are tiny-negative
# (floating-point artifact from ASE surface builders) become z ≈ cell_z
# instead of z = 0. The SliceIndexedAtoms bin edges are nudged down by
# 1e-12 to fix cumsum drift, so any atom in (cell_z-1e-12, cell_z) falls
# outside all bins and is silently dropped. Snap those back to 0.
cell_z = atoms.cell[2, 2]
atoms.positions[atoms.positions[:, 2] > cell_z - 1e-10, 2] = 0.0
if not self.integrator.periodic and self.integrator.finite:
atoms = pad_atoms(atoms, margins=margins)
elif self.integrator.periodic:
atoms = pad_atoms(atoms, margins=margins, directions="z")
if not self.periodic:
atoms = self.frozen_phonons.randomize(atoms)
if self.integrator.finite:
sliced_atoms = SlicedAtoms(
atoms=atoms, slice_thickness=self.slice_thickness, z_padding=margins
)
else:
sliced_atoms = SliceIndexedAtoms(
atoms=atoms, slice_thickness=self.slice_thickness
)
return sliced_atoms
def get_sliced_atoms(self) -> BaseSlicedAtoms:
"""
The atoms grouped into the slices given by the slice thicknesses.
Returns
-------
sliced_atoms : BaseSlicedAtoms
"""
if self._sliced_atoms is not None:
return self._sliced_atoms
self._sliced_atoms = self._prepare_atoms()
return self._sliced_atoms
def generate_slices(
self,
first_slice: int = 0,
last_slice: Optional[int] = None,
return_depth: float = False,
):
"""
Generate the slices for the potential.
Parameters
----------
first_slice : int, optional
Index of the first slice of the generated potential.
last_slice : int, optional
Index of the last slice of the generated potential.
return_depth : bool
If True, return the depth of each generated slice.
Yields
------
slices : generator of np.ndarray
Generator for the array of slices.
"""
if last_slice is None:
last_slice = len(self)
xp = get_array_module(self.device)
sliced_atoms = self.get_sliced_atoms()
numbers = np.unique(sliced_atoms.atoms.numbers)
exit_plane_after = self._exit_plane_after
cumulative_thickness = np.cumsum(self.slice_thickness)
for start, stop in generate_chunks(
last_slice - first_slice, chunks=1, start=first_slice
):
if len(numbers) > 1 or stop - start > 1:
array = xp.zeros(
(stop - start,) + self.base_shape[1:],
dtype=get_dtype(complex=False),
)
else:
array = None
for i, slice_idx in enumerate(range(start, stop)):
atoms = sliced_atoms.get_atoms_in_slices(slice_idx)
new_array = self._integrator.integrate_on_grid(
atoms,
a=sliced_atoms.slice_limits[slice_idx][0],
b=sliced_atoms.slice_limits[slice_idx][1],
gpts=self.gpts,
sampling=self.sampling,
device=self.device,
)
if array is not None:
array[i] += new_array
else:
array = new_array[None]
if array is None:
array = xp.zeros(
(stop - start,) + self.base_shape[1:],
dtype=get_dtype(complex=False),
)
# array -= array.min()
exit_planes = tuple(np.where(exit_plane_after[start:stop])[0])
potential_array = self._array_object(
array,
slice_thickness=self.slice_thickness[start:stop],
exit_planes=exit_planes,
extent=self.extent,
)
if return_depth:
depth = cumulative_thickness[stop - 1]
yield depth, potential_array
else:
yield potential_array
@property
def ensemble_axes_metadata(self):
return self.frozen_phonons.ensemble_axes_metadata
@property
def ensemble_shape(self) -> tuple[int, ...]:
return self.frozen_phonons.ensemble_shape
@classmethod
def _from_partitioned_args_func(cls, *args, frozen_phonons_partial, **kwargs):
args = unpack_blockwise_args(args)
frozen_phonons = frozen_phonons_partial(*args)
frozen_phonons = frozen_phonons.item()
new_potential = cls(frozen_phonons, **kwargs)
ndims = max(len(new_potential.ensemble_shape), 1)
new_potential = _wrap_with_array(new_potential, ndims)
return new_potential
def _from_partitioned_args(self, *args, **kwargs):
frozen_phonons_partial = self.frozen_phonons._from_partitioned_args()
kwargs = self._copy_kwargs(exclude=("atoms", "sampling"))
return partial(
self._from_partitioned_args_func,
frozen_phonons_partial=frozen_phonons_partial,
**kwargs,
)
def _partition_args(self, chunks: Optional[Chunks] = None, lazy: bool = True):
if chunks is None:
chunks = (1,)
return self.frozen_phonons._partition_args(chunks, lazy=lazy)
class _PotentialBuilder(_FieldBuilder, BasePotential):
pass
def _validate_frozen_phonons(atoms):
if isinstance(atoms, Atoms):
atoms = atoms.copy()
atoms.calc = None
if not hasattr(atoms, "randomize"):
if isinstance(atoms, (list, tuple)):
frozen_phonons = AtomsEnsemble(atoms)
elif isinstance(atoms, Atoms):
frozen_phonons = DummyFrozenPhonons(atoms)
else:
raise ValueError(
"Frozen phonons should be of types `FrozenPhonons`, `Atoms` or"
f"`AtomsEnsemble`, not {atoms}"
)
else:
frozen_phonons = atoms
return frozen_phonons
[docs]
class Potential(_FieldBuilderFromAtoms, BasePotential):
"""
Calculate the electrostatic potential of a set of atoms or frozen phonon
configurations. The potential is calculated with the Independent Atom Model (IAM)
using a user-defined parametrization of the atomic potentials.
Parameters
----------
atoms : ase.Atoms or abtem.FrozenPhonons
Atoms or FrozenPhonons defining the atomic configuration(s) used in the
independent atom model for calculating the electrostatic potential(s).
gpts : one or two int, optional
Number of grid points in `x` and `y` describing each slice of the potential.
Provide either "sampling" (spacing between consecutive grid points) or "gpts"
(total number of grid points).
sampling : one or two float, optional
Sampling of the potential in `x` and `y` [Å].
Provide either "sampling" or "gpts".
slice_thickness : float or sequence of float, optional
Thickness of the potential slices in the propagation direction in [Å]
(default is 1 Å).
If given as a float, the number of slices is calculated by dividing the slice
thickness into the `z`-height of supercell. The slice thickness may be given as
a sequence of values for each slice, in which case an error will be thrown if
the sum of slice thicknesses is not equal to the height of the atoms.
parametrization : 'lobato' or 'kirkland', optional
The potential parametrization describes the radial dependence of the potential
for each element. Two of the most accurate parametrizations are available
(by Lobato et al. and Kirkland; default is 'lobato').
See the citation guide for references.
projection : 'finite' or 'infinite', optional
If 'finite' the 3D potential is numerically integrated between the slice
boundaries. If 'infinite' (default), the infinite potential projection of each
atom will be assigned to a single slice.
exit_planes : int or tuple of int, optional
The `exit_planes` argument can be used to calculate thickness series.
Providing `exit_planes` as a tuple of int indicates that the tuple contains the
slice indices after which an exit plane is desired, and hence during a
multislice simulation a measurement is created. If `exit_planes` is an integer
a measurement will be collected every `exit_planes` number of slices.
plane : str or two tuples of three float, optional
The plane relative to the provided atoms mapped to `xy` plane of the potential,
i.e. provided plane is perpendicular to the propagation direction. If string,
it must be a concatenation of two of 'x', 'y' and 'z'; the default value 'xy'
indicates that potential slices are cuts along the `xy`-plane of the atoms.
The plane may also be specified with two arbitrary 3D vectors, which are mapped
to the `x` and `y` directions of the potential, respectively. The length of the
vectors has no influence. If the vectors are not perpendicular, the second
vector is rotated in the plane to become perpendicular to the first.
Providing a value of ((1., 0., 0.), (0., 1., 0.)) is equivalent to providing
'xy'.
origin : three float, optional
The origin relative to the provided atoms mapped to the origin of the potential.
This is equivalent to translating the atoms. The default is (0., 0., 0.).
box : three float, optional
The extent of the potential in `x`, `y` and `z`. If not given this is determined
from the atoms' cell. If the box size does not match an integer number of the
atoms' supercell, an affine transformation may be necessary to preserve
periodicity, determined by the `periodic` keyword.
periodic : bool, True
If a transformation of the atomic structure is required, `periodic` determines
how the atomic structure is transformed. If True, the periodicity of the Atoms
is preserved, which may require applying a small affine transformation to the
atoms. If False, the transformed potential is effectively cut out of a larger
repeated potential, which may not preserve periodicity.
integrator : ProjectionIntegrator, optional
Provide a custom integrator for the projection integrals of the potential
slicing.
device : str, optional
The device used for calculating the potential, 'cpu' or 'gpu'. The default is
determined by the user configuration file.
"""
_exclude_from_copy = ("parametrization", "projection")
[docs]
def __init__(
self,
atoms: Atoms | BaseFrozenPhonons,
gpts: int | tuple[int, int] | None = None,
sampling: float | tuple[float, float] | None = None,
slice_thickness: float | tuple[float, ...] = 1,
parametrization: str | Parametrization = "lobato",
projection: str = "infinite",
exit_planes: int | tuple[int, ...] | None = None,
plane: (
str | tuple[tuple[float, float, float], tuple[float, float, float]]
) = "xy",
origin: tuple[float, float, float] = (0.0, 0.0, 0.0),
box: tuple[float, float, float] | None = None,
periodic: bool = True,
integrator: FieldIntegrator | None = None,
device: str | None = None,
):
if integrator is None:
if projection == "finite":
integrator = QuadratureProjectionIntegrals(
parametrization=parametrization
)
elif projection == "infinite":
integrator = ScatteringFactorProjectionIntegrals(
parametrization=parametrization
)
else:
raise NotImplementedError
super().__init__(
atoms=atoms,
array_object=PotentialArray,
gpts=gpts,
sampling=sampling,
slice_thickness=slice_thickness,
exit_planes=exit_planes,
device=device,
plane=plane,
origin=origin,
box=box,
periodic=periodic,
integrator=integrator,
)
[docs]
class FieldArray(BaseField, ArrayObject):
[docs]
def __init__(
self,
array: np.ndarray | da.core.Array,
slice_thickness: float | Sequence[float],
extent: Optional[float | tuple[float, float]] = None,
sampling: Optional[float | tuple[float, float]] = None,
exit_planes: Optional[int | tuple[int, ...]] = None,
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[dict] = None,
):
# assert len(array.shape) == self._base_dims
self._slice_thickness = _validate_slice_thickness(
slice_thickness, num_slices=array.shape[-self._base_dims]
)
self._exit_planes = _validate_exit_planes(
exit_planes, len(self._slice_thickness)
)
self._grid = Grid(extent=extent, gpts=array.shape[-2:], sampling=sampling)
super().__init__(
array=array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@property
def num_configurations(self):
indices = _find_axes_type(self, FrozenPhononsAxis)
if indices:
return reduce(mul, tuple(self.array.shape[i] for i in indices))
else:
return 1
@property
def slice_thickness(self) -> tuple[float, ...]:
return self._slice_thickness
@property
def exit_planes(self) -> tuple[int, ...]:
return self._exit_planes
def build(
self,
first_slice: int = 0,
last_slice: Optional[int] = None,
chunks: int = 1,
lazy: Optional[bool] = None,
):
raise RuntimeError("potential is already built")
[docs]
def generate_slices(self, first_slice: int = 0, last_slice: Optional[int] = None):
"""
Generate the slices for the potential.
Parameters
----------
first_slice : int, optional
Index of the first slice of the generated potential.
last_slice : int, optional
Index of the last slice of the generated potential.
Yields
------
slices : generator of np.ndarray
Generator for the array of slices.
"""
if last_slice is None:
last_slice = len(self)
exit_plane_after = self._exit_plane_after
# cum_thickness = np.cumsum(self.slice_thickness)
start = first_slice
stop = first_slice + 1
for i in range(first_slice, last_slice):
s = (0,) * (len(self.array.shape) - 3) + (i,)
array = self.array[s][None]
slic = self.__class__(
array, self.slice_thickness[i : i + 1], extent=self.extent
)
exit_planes = tuple(np.where(exit_plane_after[start:stop])[0])
slic._exit_planes = exit_planes
start += 1
stop += 1
yield slic
def __getitem__(self, items):
if isinstance(items, (Number, slice)):
items = (items,)
if not len(items) <= len(self.ensemble_shape) + 1:
raise IndexError(
f"Too many indices for potential array with {len(self.ensemble_shape)}"
"ensemble axes. Only slice indices and ensemble indices are allowed."
)
ensemble_items = items[: len(self.ensemble_shape)]
slic_items = items[len(self.ensemble_shape) :]
if len(ensemble_items):
potential_array = super().__getitem__(ensemble_items)
else:
potential_array = self
if len(slic_items) == 0:
return potential_array
padded_items = (slice(None),) * len(potential_array.ensemble_shape) + slic_items
array = potential_array._array[padded_items]
slice_thickness = np.array(potential_array.slice_thickness)[slic_items]
if len(array.shape) < len(potential_array.shape):
array = array[
(slice(None),) * len(potential_array.ensemble_shape) + (None,)
]
slice_thickness = slice_thickness[None]
kwargs = potential_array._copy_kwargs(exclude=("array", "slice_thickness"))
kwargs["array"] = array
kwargs["slice_thickness"] = slice_thickness
kwargs["sampling"] = None
return potential_array.__class__(**kwargs)
[docs]
def tile(self, repetitions: tuple[int, int] | tuple[int, int, int]):
"""
Tile the potential.
Parameters
----------
repetitions: two or three int
The number of repetitions of the potential along each axis. NOTE: if three
integers are given, the last represents the number of repetitions along the
`z`-axis.
Returns
-------
PotentialArray object
The tiled potential.
"""
if len(repetitions) == 2:
repetitions = (repetitions[0], repetitions[1], 1)
assert len(repetitions) == 3
new_array = np.tile(
self.array, (repetitions[2], repetitions[0], repetitions[1])
)
if self.extent is not None:
new_extent = (
self.extent[0] * repetitions[0],
self.extent[1] * repetitions[1],
)
else:
new_extent = None
new_slice_thickness = tuple(np.tile(self.slice_thickness, repetitions[2]))
return self.__class__(
array=new_array,
slice_thickness=new_slice_thickness,
extent=new_extent,
ensemble_axes_metadata=self.ensemble_axes_metadata,
)
[docs]
def to_hyperspy(self, transpose: bool = True):
return self.to_images().to_hyperspy(transpose=transpose)
[docs]
def to_images(self):
"""Convert slices of the potential to a stack of images."""
return Images(
array=self._array,
sampling=(self.sampling[0], self.sampling[1]),
metadata=self.metadata,
ensemble_axes_metadata=self.axes_metadata[:-2],
)
[docs]
def depth_profile(
self,
projection_axis: str = "y",
depth: Optional[float] = None,
) -> Images:
"""Create a depth profile by projecting the potential along a spatial axis.
Parameters
----------
projection_axis : str
Spatial axis to project (sum) along. ``"y"`` (default) produces an
x–z cross-section; ``"x"`` produces a y–z cross-section.
depth : float, optional
If given, project only over a finite slab of this thickness [Å],
centered on the midpoint of the projected axis. The number of grid
points is rounded to the nearest integer. If ``None``, the full
extent is projected.
Returns
-------
depth_profile : Images
2D image(s) with the remaining spatial axis horizontal and depth
(z) vertical.
"""
from copy import copy
if projection_axis not in ("x", "y"):
raise ValueError("projection_axis must be 'x' or 'y'.")
array = self.array
if projection_axis == "y":
sum_axis = -1
spatial_sampling = self.sampling[0]
else:
sum_axis = -2
spatial_sampling = self.sampling[1]
if depth is not None:
proj_sampling = (
self.sampling[1] if projection_axis == "y" else self.sampling[0]
)
proj_gpts = self.gpts[1] if projection_axis == "y" else self.gpts[0]
n = max(1, min(proj_gpts, round(depth / proj_sampling)))
start = (proj_gpts - n) // 2
slices = [slice(None)] * len(array.shape)
slices[sum_axis] = slice(start, start + n)
array = array[tuple(slices)]
array = array.sum(axis=sum_axis)
xp = get_array_module(array)
if hasattr(array, "rechunk"):
array = da.moveaxis(array, -2, -1)
else:
array = xp.moveaxis(array, -2, -1)
n_z = self.num_slices
z_extent = self.thickness
z_sampling = z_extent / n_z if n_z > 0 else 1.0
metadata = copy(self.metadata)
return Images(
array,
sampling=(spatial_sampling, z_sampling),
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=metadata,
)
[docs]
def project(self) -> Images:
"""
Create a 2D array representing a projected image of the potential(s).
Returns
-------
images : Images
One or more images of the projected potential(s).
"""
array = self.array.sum(-self._base_dims)
# array -= array.min((-2, -1), keepdims=True)
ensemble_axes_metadata = (
self.ensemble_axes_metadata + self.base_axes_metadata[1:-2]
)
return Images(
array=array,
sampling=self._valid_sampling,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=self.metadata,
)
[docs]
class PotentialArray(BasePotential, FieldArray):
"""
The potential array represents slices of the electrostatic potential as an array.
All other potentials build potential arrays.
Parameters
----------
array: 3D np.ndarray
The array representing the potential slices. The first dimension is the slice
index and the last two are the spatial dimensions.
slice_thickness: float
The thicknesses of potential slices [Å]. If a float, the thickness is the same
for all slices.
If a sequence, the length must equal the length of the potential array.
extent: one or two float, optional
Lateral extent of the potential [Å].
sampling: one or two float, optional
Lateral sampling of the potential [1 / Å].
exit_planes : int or tuple of int, optional
The `exit_planes` argument can be used to calculate thickness series.
Providing `exit_planes` as a tuple of int indicates that the tuple contains the
slice indices after which an exit plane is desired, and hence during a
multislice simulation a measurement is created. If `exit_planes` is an integer a
measurement will be collected every `exit_planes` number of slices.
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 = 3
[docs]
def __init__(
self,
array: np.ndarray | da.core.Array,
slice_thickness: float | Sequence[float],
extent: Optional[float | tuple[float, float]] = None,
sampling: Optional[float | tuple[float, float]] = None,
exit_planes: Optional[int | tuple[int, ...]] = None,
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[dict] = None,
):
if metadata is None:
metadata = {}
metadata = {"label": "potential", "units": "eV / e", **metadata}
super().__init__(
array=array,
slice_thickness=slice_thickness,
extent=extent,
sampling=sampling,
exit_planes=exit_planes,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@staticmethod
def _transmission_function(array, energy):
xp = get_array_module(array)
sigma = xp.array(energy2sigma(energy), dtype=get_dtype())
array = complex_exponential(sigma * array)
return array
[docs]
def transmission_function(self, energy: float) -> TransmissionFunction:
"""
Calculate the transmission functions for each slice for a specific energy.
Parameters
----------
energy: float
Electron energy [eV].
Returns
-------
transmissionfunction : TransmissionFunction
Transmission functions for each slice.
"""
xp = get_array_module(self.array)
if self.is_lazy:
array = da.map_blocks(
self._transmission_function,
self.array,
energy=energy,
meta=xp.array((), dtype=get_dtype(complex=True)),
)
else:
array = self._transmission_function(self.array, energy=energy)
t = TransmissionFunction(
array,
slice_thickness=self.slice_thickness,
extent=self.extent,
energy=energy,
)
return t
[docs]
def transmit(self, waves: Waves, conjugate: bool = False) -> Waves:
"""
Transmit a wave function through a potential slice.
Parameters
----------
waves: Waves
Waves object to transmit.
conjugate : bool, optional
If True, use the conjugate of the transmission function. Default is False.
Returns
-------
transmission_function : TransmissionFunction
Transmission function for the wave function through the potential slice.
"""
transmission_function = self.transmission_function(waves._valid_energy)
return transmission_function.transmit(waves, conjugate=conjugate)
[docs]
class TransmissionFunction(PotentialArray, HasAcceleratorMixin):
"""Class to describe transmission functions.
Parameters
----------
array : 3D np.ndarray
The array representing the potential slices. The first dimension is the slice
index and the last two are the spatial dimensions.
slice_thickness : float
The thicknesses of potential slices [Å]. If a float, the thickness is the same
for all slices. If a sequence, the length must equal the length of the potential
array.
extent : one or two float, optional
Lateral extent of the potential [Å].
sampling : one or two float, optional
Lateral sampling of the potential [1 / Å].
energy : float
Electron energy [eV].
"""
[docs]
def __init__(
self,
array: np.ndarray,
slice_thickness: float | Sequence[float],
extent: Optional[float | tuple[float, float]] = None,
sampling: Optional[float | tuple[float, float]] = None,
energy: Optional[float] = None,
):
self._accelerator = Accelerator(energy=energy)
super().__init__(array, slice_thickness, extent, sampling)
def get_chunk(self, first_slice, last_slice) -> TransmissionFunction:
array = self.array[first_slice:last_slice]
if len(array.shape) == 2:
array = array[None]
return self.__class__(
array,
self.slice_thickness[first_slice:last_slice],
extent=self.extent,
energy=self.energy,
)
[docs]
def transmission_function(self, energy) -> TransmissionFunction:
"""
Calculate the transmission functions for each slice for a specific energy.
Parameters
----------
energy: float
Electron energy [eV].
Returns
-------
transmissionfunction : TransmissionFunction
Transmission functions for each slice.
"""
if energy != self.energy:
raise RuntimeError()
return self
[docs]
def transmit(self, waves: Waves, conjugate: bool = False) -> Waves:
"""
Transmit a wave function through a potential slice.
Parameters
----------
waves: Waves
Waves object to transmit.
conjugate : bool, optional
If True, use the conjugate of the transmission function. Default is False.
Returns
-------
transmission_function : Waves
Transmission function for the wave function through the potential slice.
"""
self.accelerator.check_match(waves)
self.grid.check_match(waves)
xp = get_array_module(self.array[0])
if conjugate:
waves._array *= xp.conjugate(self.array[0])
else:
waves._array *= self.array[0]
return waves
[docs]
class CrystalPotential(_PotentialBuilder):
"""
The crystal potential may be used to represent a potential consisting of a repeating
unit. This may allow calculations to be performed with lower computational cost by
calculating the potential unit once and repeating it.
If the repeating unit is a potential with frozen phonons, it is treated as a
pool of displaced configurations: every repetition of the unit (each lateral
tile of every `z`-repetition) draws a configuration from the pool. Draws are
balanced over the whole crystal, so reuse of a configuration is the minimum
the pool size allows -- no two tiles within a layer are identical whenever
the pool permits, and a pool of at least
``repetitions[0] * repetitions[1] * repetitions[2]`` configurations gives
every repeated unit a distinct configuration (statistically equivalent to
tiling the displaced atoms directly). If `num_frozen_phonons` is set, an
ensemble of crystal potentials is created; each member independently
rebuilds its own pool of atomic displacement snapshots (reseeded from
that member's own seed) rather than sharing one fixed pool across the
ensemble, so members are genuinely independent thermal realisations --
there is no need to size the pool for the ensemble, only for a single
crystal (see above).
Parameters
----------
potential_unit : BasePotential
The potential unit to assemble the crystal potential from.
repetitions : three int
The repetitions of the potential in `x`, `y` and `z`.
num_frozen_phonons : int, optional
Number of crystal realisations in the frozen-phonon ensemble; each
realisation independently rebuilds its own pool of atomic
displacement snapshots.
exit_planes : int or tuple of int, optional
The `exit_planes` argument can be used to calculate thickness series.
Providing `exit_planes` as a tuple of int indicates that the tuple contains the
slice indices after which an exit plane is desired, and hence during a
multislice simulation a measurement is created. If `exit_planes` is an integer
a measurement will be collected every `exit_planes` number of slices.
seeds: int or sequence of int
Seed for the random number generator (RNG), or one seed for each RNG in the
frozen phonon ensemble.
ensemble_mean : bool, optional
If True (default), the mean over the frozen-phonon ensemble is calculated.
If False, the individual configurations are returned.
"""
[docs]
def __init__(
self,
potential_unit: BasePotential,
repetitions: tuple[int, int, int],
num_frozen_phonons: int | None = None,
exit_planes: int | None = None,
seeds: int | tuple[int, ...] | None = None,
ensemble_mean: bool = True,
):
if num_frozen_phonons is None and seeds is None:
self._seeds = None
else:
if num_frozen_phonons is None and seeds:
assert isinstance(seeds, tuple)
num_frozen_phonons = len(seeds)
elif num_frozen_phonons is None and seeds is None:
num_frozen_phonons = 1
self._seeds = validate_seeds(seeds, num_frozen_phonons)
if (
(potential_unit.num_configurations == 1)
and (num_frozen_phonons is not None)
and (num_frozen_phonons > 1)
):
warnings.warn(
"'num_frozen_phonons' is greater than one, but the potential unit does"
" not have frozen phonons"
)
gpts = (
potential_unit._valid_gpts[0] * repetitions[0],
potential_unit._valid_gpts[1] * repetitions[1],
)
extent = (
potential_unit._valid_extent[0] * repetitions[0],
potential_unit._valid_extent[1] * repetitions[1],
)
box = extent + (potential_unit.thickness * repetitions[2],)
slice_thickness = potential_unit.slice_thickness * repetitions[2]
assert hasattr(potential_unit, "device")
super().__init__(
array_object=PotentialArray,
gpts=gpts,
cell=Cell(np.diag(box)),
slice_thickness=slice_thickness,
exit_planes=exit_planes,
device=potential_unit.device,
plane="xy",
origin=(0.0, 0.0, 0.0),
box=box,
periodic=True,
)
self._potential_unit = potential_unit
self._repetitions = repetitions
self._ensemble_mean = ensemble_mean
self._sliced_atoms: Optional[BaseSlicedAtoms] = None
@property
def ensemble_mean(self) -> bool:
return self._ensemble_mean
@property
def ensemble_shape(self) -> tuple[int, ...]:
if self._seeds is None:
return ()
else:
return (self.num_configurations,)
@property
def num_configurations(self):
if self._seeds is None:
return 1
else:
return len(self._seeds)
@property
def seeds(self):
return self._seeds
@property
def potential_unit(self) -> BasePotential:
return self._potential_unit
@property
def gpts(self) -> tuple[int, int] | None:
return super().gpts
@gpts.setter
def gpts(self, gpts: tuple[int, int]):
if not (
(gpts[0] % self.repetitions[0] == 0)
and (gpts[1] % self.repetitions[0] == 0)
):
raise ValueError(
"Number of grid points must be divisible by the number of potential"
"repetitions."
)
self.grid.gpts = gpts
self._potential_unit.gpts = (
gpts[0] // self._repetitions[0],
gpts[1] // self._repetitions[1],
)
@property
def sampling(self) -> tuple[float, float] | None:
return super().sampling
@sampling.setter
def sampling(self, sampling: tuple[float, float]):
self.sampling = sampling
self._potential_unit.sampling = sampling
@property
def repetitions(self) -> tuple[int, int, int]:
return self._repetitions
@property
def num_slices(self) -> int:
return self._potential_unit.num_slices * self.repetitions[2]
@property
def ensemble_axes_metadata(self) -> list[AxisMetadata]:
if self.seeds is None:
return []
else:
return [FrozenPhononsAxis(_ensemble_mean=self._ensemble_mean)]
[docs]
def get_sliced_atoms(self) -> BaseSlicedAtoms:
"""
The atoms of the full crystal grouped into the slices given by the slice
thicknesses.
The atoms are reconstructed by tiling the unit potential's transformed
(orthogonalised) atoms by the crystal repetitions. This makes
``CrystalPotential`` work with any code path that derives atomic sites
from a potential via ``get_sliced_atoms`` -- e.g. the core-loss EELS
driver's automatic site extraction -- without special-casing the
repeating-unit structure.
Notes
-----
- **Frozen phonons are not displaced.** ``get_transformed_atoms``
returns the equilibrium (mean) positions, so the returned sites are
the un-displaced atomic columns. This is deliberate: a
``CrystalPotential`` ensemble draws an independent random unit
configuration per z-repetition, so there is no single displaced
realisation to return, and atomic-column site identification (the
main consumer) wants the equilibrium column positions anyway. This
differs from ``Potential.get_sliced_atoms``, which applies the
frozen-phonon displacement of its single configuration.
- The result is cached; the tile is non-trivial for large supercells.
Returns
-------
sliced_atoms : BaseSlicedAtoms
"""
if self._sliced_atoms is not None:
return self._sliced_atoms
if not hasattr(self._potential_unit, "get_transformed_atoms"):
raise RuntimeError(
"Cannot derive atoms from a CrystalPotential whose "
f"potential_unit ({type(self._potential_unit).__name__}) does "
"not expose 'get_transformed_atoms' (e.g. a precomputed "
"PotentialArray). Pass the scattering sites explicitly instead."
)
unit_atoms = self._potential_unit.get_transformed_atoms()
tiled_atoms = unit_atoms * self._repetitions
self._sliced_atoms = SliceIndexedAtoms(
tiled_atoms, slice_thickness=self.slice_thickness
)
return self._sliced_atoms
@classmethod
def _from_partitioned_args_func(cls, *args, **kwargs):
args = unpack_blockwise_args(args)
potential, seed = args[0]
if hasattr(potential, "item"):
potential = potential.item()
if seed is not None:
num_frozen_phonons = len(seed)
else:
num_frozen_phonons = None
new = cls(
potential_unit=potential,
seeds=seed,
num_frozen_phonons=num_frozen_phonons,
**kwargs,
)
return _wrap_with_array(new)
def _from_partitioned_args(self):
kwargs = self._copy_kwargs(
exclude=("potential_unit", "seeds", "num_frozen_phonons")
)
output = partial(self._from_partitioned_args_func, **kwargs)
return output
def _partition_args(self, chunks: Optional[Chunks] = None, lazy: bool = True):
if chunks is None:
chunks = 1
chunks = validate_chunks(self.ensemble_shape, chunks)
# print(self.ensemble_shape)
if chunks == ():
old_chunks = ()
chunks = ((1,),)
else:
old_chunks = chunks
if lazy:
arrays = []
for i, (start, stop) in enumerate(chunk_ranges(chunks)[0]):
if self.seeds is not None:
seeds = self.seeds[start:stop]
else:
seeds = None
lazy_atoms = dask.delayed(self.potential_unit)
lazy_args = dask.delayed(_wrap_with_array)((lazy_atoms, seeds), ndims=1)
lazy_array = da.from_delayed(lazy_args, shape=(1,), dtype=object)
arrays.append(lazy_array)
array = da.concatenate(arrays)
if old_chunks == ():
array = array[0]
else:
potential_unit = self.potential_unit
array = np.zeros((len(chunks[0]),), dtype=object)
for i, (start, stop) in enumerate(chunk_ranges(chunks)[0]):
if self.seeds is not None:
seeds = self.seeds[start:stop]
else:
seeds = None
itemset(array, i, (potential_unit, seeds))
if old_chunks == ():
array = _wrap_with_array(array[0], ndims=0)
return (array,)
@property
def _n_lateral_tiles(self) -> int:
return self.repetitions[0] * self.repetitions[1]
def _pool_unit_for_member(self, member_seed: Optional[int]) -> BasePotential:
"""Return the unit potential to draw pool configurations from for one
ensemble member (``member_seed`` is that member's seed), or for the
single default builder (``member_seed`` is None).
Two independent adjustments are made when the unit carries frozen
phonons; a precomputed ``PotentialArray`` unit has a fixed pool and is
always returned unchanged.
1. **Enlarge to the tile count.** A frozen-phonon ``CrystalPotential``
assembles each slice as a mosaic: every lateral tile draws an
independent pool configuration. If the pool holds fewer
configurations than there are lateral tiles
(``repetitions[0] * repetitions[1]``), some tiles must reuse a
configuration -- reintroducing the artificial in-plane periodicity
the mosaic is meant to remove. The pool is transparently enlarged
to the tile count (warning).
2. **Reseed per ensemble member.** Every ensemble member is built from
the *same* ``potential_unit`` object, so without reseeding every
member would draw from an identical, fixed pool of configurations
-- differing only in how those same snapshots are arranged across
the crystal, not in which atomic displacements exist. That is a
much weaker form of independence than a frozen-phonon ensemble is
supposed to provide, and sizing the pool cannot fix it (drawing
from a bigger *shared* pool still shares it). Instead, when this
call belongs to an ensemble (``member_seed`` is not None), the pool
is quietly rebuilt with ``member_seed`` as its root seed, so each
member gets its own independent set of atomic snapshots. This adds
no cost: the pool was already rebuilt once per member.
"""
unit = self.potential_unit
n_tiles = self._n_lateral_tiles
fp = getattr(unit, "frozen_phonons", None)
if not isinstance(fp, FrozenPhonons) or fp.num_configs <= 1:
return unit
enlarge = fp.num_configs < n_tiles
reseed = member_seed is not None
if not enlarge and not reseed:
return unit
if enlarge:
warnings.warn(
f"frozen-phonon pool ({fp.num_configs}) is smaller than the "
f"number of lateral tiles ({n_tiles}); enlarging the pool to "
f"{n_tiles} so each tile draws a distinct configuration and no "
"lateral duplication occurs. Pass a unit with "
f"num_configs >= {n_tiles} to silence this."
)
new_fp = FrozenPhonons(
fp.atoms,
num_configs=n_tiles if enlarge else fp.num_configs,
sigmas=fp.sigmas,
directions=fp.directions,
ensemble_mean=fp.ensemble_mean,
seed=int(member_seed) if reseed else int(fp.seed[0]),
)
kwargs = unit._copy_kwargs(exclude=("atoms",))
return type(unit)(new_fp, **kwargs)
[docs]
def generate_slices(
self,
first_slice: int = 0,
last_slice: Optional[int] = None,
return_depth: bool = False,
):
"""
Generate the slices for the potential.
Parameters
----------
first_slice : int, optional
Index of the first slice of the generated potential.
last_slice : int, optional
Index of the last slice of the generated potential.
return_depth : bool
If True, return the depth of each generated slice.
Yields
------
slices : generator of np.ndarray
Generator for the array of slices.
"""
# if hasattr(self.potential_unit, "array")
# potentials = self.potential_unit
member_seed = None if self.seeds is None else int(self.seeds[0])
pool_unit = self._pool_unit_for_member(member_seed)
if not isinstance(pool_unit, PotentialArray):
potentials = pool_unit.build(lazy=False)
else:
potentials = pool_unit
assert isinstance(potentials, PotentialArray)
if len(potentials.shape) == 3:
potentials = potentials.expand_dims(axis=0)
rng = np.random.default_rng(member_seed)
exit_plane_after = self._exit_plane_after
cum_thickness = np.cumsum(self.slice_thickness)
start = first_slice
stop = first_slice + 1
# Lazy cache of tiled unit slices, keyed by (config_idx, slice_idx).
# Without it each (z-rep, unit-slice) pair re-tiles the same array via
# ``.tile(self.repetitions[:2])`` — for the no-frozen-phonon case
# (n_configs == 1) every z-rep produces an identical result so the
# cost scales linearly with ``repetitions[2]``. The cache turns this
# into ``n_configs * len(self.potential_unit)`` unique tile calls.
# For the SrTiO3 tutorial (reps=(4,4,25), 2 unit slices, no FP) this
# is 2 tiles instead of 50; the cache footprint is bounded by the
# tiled-unit byte size and freed when the generator is exhausted.
tiled_cache: dict[tuple[int, int], PotentialArray] = {}
unit_generators: dict[int, object] = {}
tile_xy = self.repetitions[:2]
n_configs = potentials.shape[0]
# The mosaic path (frozen-phonon pools, n_configs > 1) needs random
# per-tile access into the pool, so materialise the (small unit-cell)
# pool array once. A lazily-built PotentialArray unit carries a dask
# array here; compute it so per-tile fancy indexing works and stays on
# the target device.
xp = get_array_module(self.device)
_pool_array = potentials.array
if n_configs > 1 and hasattr(_pool_array, "compute"):
_pool_array = _pool_array.compute()
def _tiled_slice(config_idx: int, j: int) -> PotentialArray:
key = (config_idx, j)
cached = tiled_cache.get(key)
if cached is not None:
return cached
gen = unit_generators.get(config_idx)
if gen is None:
gen = potentials[config_idx].generate_slices()
unit_generators[config_idx] = gen
slic = next(gen).tile(tile_xy)
tiled_cache[key] = slic
return slic
def _mosaic_slice(config_tiles: np.ndarray, j: int) -> PotentialArray:
# Assemble sub-slice ``j`` of the lateral supercell by placing an
# *independently drawn* pool configuration at every lateral
# repetition (a mosaic), rather than replicating a single displaced
# unit across all tiles. This is what reproduces genuine lateral
# (in-plane) thermal disorder: with plain ``.tile()`` every one of
# the ``repetitions[0] * repetitions[1]`` tiles is a bit-identical
# copy, so there is no in-plane disorder at all and no diffuse
# (Kikuchi) scattering can form. ``config_tiles`` holds one pool
# index per lateral tile, shaped ``repetitions[:2]``.
sub = _pool_array[:, j] # (n_configs, uy, ux)
uy, ux = sub.shape[-2], sub.shape[-1]
mosaic = sub[xp.asarray(config_tiles)] # (rep0, rep1, uy, ux)
# Interleave to match ``PotentialArray.tile`` block layout, which
# tiles the row axis by repetitions[0] and the col axis by
# repetitions[1] (np.tile(array, (rep2, rep0, rep1))).
mosaic = mosaic.transpose(0, 2, 1, 3).reshape(
tile_xy[0] * uy, tile_xy[1] * ux
)
return potentials.__class__(
mosaic[None],
potentials.slice_thickness[j : j + 1],
extent=self.extent,
)
n_tiles = tile_xy[0] * tile_xy[1]
# Balanced global drawing: every pool configuration receives a total
# usage budget of floor/ceil(total_slots / n_configs) over the whole
# crystal (all lateral tiles x all z-repetitions), and each z-layer
# draws the ``n_tiles`` configurations with the most budget remaining
# (random tie-breaking keeps assignments uniform). Drawing each layer
# independently instead (i.e. with replacement across z) lets the
# same configuration recur in many layers even when the pool is large
# enough to avoid it, correlating slices along z and measurably
# inflating thermal-diffuse statistics above the tiled-atoms ground
# truth. With budgets, reuse is the minimum the pool size allows and
# is spread evenly: once ``n_configs >= n_tiles * repetitions[2]``
# every unit cell in the crystal receives a distinct configuration --
# statistically identical to tiling the displaced atoms directly.
# Within a layer draws remain distinct whenever the pool allows (no
# in-plane duplication), as before.
if n_configs > 1:
total_slots = n_tiles * self.repetitions[2]
base, extra = divmod(total_slots, n_configs)
budgets = np.full(n_configs, base, dtype=np.int64)
if extra:
budgets[rng.permutation(n_configs)[:extra]] += 1
def _draw_config_tiles() -> np.ndarray:
if n_configs >= n_tiles:
# The ``n_tiles`` most-underused configurations, in random
# order (permute first; the stable sort then orders by budget
# only, keeping ties shuffled).
order = rng.permutation(n_configs)
chosen = order[np.argsort(-budgets[order], kind="stable")[:n_tiles]]
else:
# Pool smaller than a single layer: in-plane repeats are
# unavoidable; cycle freshly shuffled permutations to spread
# them as evenly as possible.
n_perms = -(-n_tiles // n_configs) # ceil
chosen = np.concatenate(
[rng.permutation(n_configs) for _ in range(n_perms)]
)[:n_tiles]
np.subtract.at(budgets, chosen, 1)
return chosen.reshape(tile_xy)
for i in range(self.repetitions[2]):
# Draw an independent displaced realisation per unit cell in the x,
# y and z supercell directions. For a single-config pool this
# collapses to the cheap cached ``.tile()`` path below.
if n_configs > 1:
config_tiles = _draw_config_tiles()
else:
config_tiles = None
for j in range(len(self.potential_unit)):
if config_tiles is None:
slic = _tiled_slice(0, j)
else:
slic = _mosaic_slice(config_tiles, j)
exit_planes = tuple(np.where(exit_plane_after[start:stop])[0])
# Mutating the cached slice is safe in the standard sequential
# consumption pattern (the consumer reads ``slic.exit_planes``
# immediately upon receiving the yield and never holds a back-
# reference across iterations — see multislice.py:672). Reset
# the value on every yield so re-entering the same cached
# slice on a different z-rep still carries the right metadata.
slic._exit_planes = exit_planes
start += 1
stop += 1
if return_depth:
yield cum_thickness[stop - 1], slic
else:
yield slic
if j == last_slice:
break
