"""Module for describing different types of scans."""
from abc import ABCMeta, abstractmethod
from copy import copy
from typing import Union, Sequence, Tuple, List
import h5py
import numpy as np
from matplotlib.patches import Rectangle
from abtem.base_classes import Grid, HasGridMixin
from abtem.device import asnumpy
from abtem.measure import Calibration, Measurement
from abtem.utils import subdivide_into_batches, ProgressBar
from ase import Atom
[docs]class AbstractScan(metaclass=ABCMeta):
"""Abstract class to describe scans."""
def __init__(self):
self._batches = None
def __len__(self):
return self.num_positions
@property
def num_positions(self):
return len(self.get_positions())
@property
@abstractmethod
def shape(self) -> tuple:
"""The shape the scan."""
pass
@property
@abstractmethod
def calibrations(self) -> tuple:
"""The measurement calibrations associated with the scan."""
pass
[docs] @abstractmethod
def get_positions(self):
"""Get the scan positions as numpy array."""
pass
[docs] @abstractmethod
def insert_new_measurement(self, measurement, indices, new_values):
"""
Insert new measurement values into a Measurement object or HDF5 file.
Parameters
----------
measurement : Measurement object
The measurement to which new values are inserted.
start : int
First index of slice.
end : int
Last index of slice.
new_values : ndarray
New measurement values to be inserted. Length should be (end - start).
"""
pass
def generate_positions(self, max_batch, pbar=False):
positions = self.get_positions()
self._partition_batches(max_batch)
if pbar:
pbar = ProgressBar(total=len(self))
for i in range(len(self._batches)):
indices = self.get_next_batch()
yield indices, positions[indices]
if pbar:
pbar.update(len(indices))
if pbar:
pbar.close()
def get_next_batch(self):
return self._batches.pop(0)
def _partition_batches(self, max_batch):
n = len(self)
n_batches = (n + (-n % max_batch)) // max_batch
batch_sizes = subdivide_into_batches(len(self), n_batches)
self._batches = []
start = 0
for batch_size in batch_sizes:
end = start + batch_size
indices = np.arange(start, end, dtype=int)
start += batch_size
self._batches.append(indices)
@abstractmethod
def __copy__(self):
pass
[docs] def copy(self):
"""Make a copy."""
return copy(self)
[docs]class PositionScan(AbstractScan):
"""
Position scan object.
Defines a scan based on user-provided positions.
Parameters
----------
positions : list
A list of xy scan positions [Å].
"""
def __init__(self, positions: np.ndarray):
self._positions = np.array(positions)
if (len(self._positions.shape) != 2) or (self._positions.shape[1] != 2):
raise RuntimeError('The shape of the sequence of positions must be (n, 2).')
super().__init__()
@property
def shape(self) -> tuple:
return len(self),
@property
def calibrations(self) -> tuple:
return None,
[docs] def insert_new_measurement(self, measurement, indices, new_measurement):
if isinstance(measurement, str):
with h5py.File(measurement, 'a') as f:
f['array'][indices] = asnumpy(new_measurement)
else:
measurement.array[indices] = asnumpy(new_measurement)
[docs] def get_positions(self):
return self._positions
[docs] def add_to_mpl_plot(self, ax, marker: str = 'o', color: str = 'r', **kwargs):
"""
Add a visualization of the scan positions to a matplotlib plot.
Parameters
----------
ax: matplotlib Axes
The axes of the matplotlib plot the visualization should be added to.
marker: str, optional
Style of scan position markers. Default is '-'.
color: str, optional
Color of the scan position markers. Default is 'r'.
kwargs:
Additional options for matplotlib.pyplot.plot as keyword arguments.
"""
ax.plot(*self.get_positions().T, marker=marker, linestyle='', color=color, **kwargs)
def __copy__(self):
return self.__class__(self._positions.copy())
[docs]class LineScan(AbstractScan, HasGridMixin):
"""
Line scan object.
Defines a scan along a straight line.
Parameters
----------
start : two float
Start point of the scan [Å].
end : two float
End point of the scan [Å].
gpts: int
Number of scan positions.
sampling: float
Sampling rate of scan positions [1 / Å].
endpoint: bool
If True, end is the last position. Otherwise, it is not included. Default is True.
"""
def __init__(self,
start: Union[Tuple[float, float], Atom],
end: Union[Tuple[float, float], Atom] = None,
angle: float = 0.,
gpts: int = None,
sampling: float = None,
margin: float = 0.,
endpoint: bool = True):
super().__init__()
if isinstance(start, Atom):
start = (start.x, start.y)
if isinstance(end, Atom):
end = (end.x, end.y)
# if (end is not None) & (angle is not None):
# raise ValueError('only one of "end" and "angle" may be specified')
if (gpts is None) & (sampling is None):
raise RuntimeError('grid gpts or sampling must be set')
self._grid = Grid(gpts=gpts, sampling=sampling, endpoint=endpoint, dimensions=1)
self._start = start[:2]
self._margin = margin
if end is not None:
self._set_direction_and_extent(self._start, end[:2])
else:
self.angle = angle
self.extent = 2 * self._margin
def _set_direction_and_extent(self, start: Tuple[float, float], end: Tuple[float, float]):
difference = np.array(end) - np.array(start)
extent = np.linalg.norm(difference, axis=0)
self._direction = difference / extent
extent = extent + 2 * self._margin
if extent == 0.:
raise RuntimeError('scan has no extent')
self.extent = extent
@property
def shape(self) -> Tuple[int]:
return self.gpts[0],
@property
def calibrations(self) -> Tuple[Calibration]:
return Calibration(offset=0, sampling=self.sampling[0], units='Å', name='x', endpoint=self.grid.endpoint[0]),
@property
def start(self) -> Tuple[float, float]:
"""
Start point of the scan [Å].
"""
return self._start
@start.setter
def start(self, start: Tuple[float, float]):
self._start = start
self._set_direction_and_extent(self._start, self.end)
@property
def end(self) -> Tuple[float, float]:
"""
End point of the scan [Å].
"""
return (self.start[0] + self.direction[0] * self.extent[0] - self.direction[0] * 2 * self._margin,
self.start[1] + self.direction[1] * self.extent[0] - self.direction[1] * 2 * self._margin)
@end.setter
def end(self, end: Tuple[float, float]):
self._set_direction_and_extent(self.start, end)
@property
def angle(self) -> float:
"""
End point of the scan [Å].
"""
return np.arctan2(self._direction[0], self._direction[1])
@angle.setter
def angle(self, angle: float):
self._direction = (np.cos(np.deg2rad(angle)), np.sin(np.deg2rad(angle)))
@property
def direction(self) -> Tuple[float, float]:
"""Direction of the scan line."""
return self._direction
@property
def margin(self) -> float:
return self._margin
[docs] def insert_new_measurement(self,
measurement: Measurement,
indices,
new_measurement_values: np.ndarray):
if isinstance(measurement, str):
with h5py.File(measurement, 'a') as f:
f['array'][indices] += asnumpy(new_measurement_values)
else:
measurement.array[indices] += asnumpy(new_measurement_values)
@property
def margin_start(self) -> Tuple[float, float]:
return self.start[0] - self.direction[0] * self.margin, self.start[1] - self.direction[1] * self.margin
@property
def margin_end(self) -> Tuple[float, float]:
return self.end[0] + self.direction[0] * self.margin, self.end[1] + self.direction[1] * self.margin
[docs] def get_positions(self) -> np.ndarray:
start = self.margin_start
end = self.margin_end
x = np.linspace(start[0], end[0], self.gpts[0], endpoint=self.grid.endpoint[0])
y = np.linspace(start[1], end[1], self.gpts[0], endpoint=self.grid.endpoint[0])
return np.stack((np.reshape(x, (-1,)), np.reshape(y, (-1,))), axis=1)
[docs] def add_to_mpl_plot(self, ax, color: str = 'r', width: float = 0.2, **kwargs):
"""
Add a visualization of a scan line to a matplotlib plot.
Parameters
----------
ax : matplotlib Axes
The axes of the matplotlib plot the visualization should be added to.
linestyle : str, optional
Linestyle of scan line. Default is '-'.
color : str, optional
Color of the scan line. Default is 'r'.
kwargs :
Additional options for matplotlib.pyplot.plot as keyword arguments.
"""
start = self.margin_start
end = self.margin_end
from matplotlib.lines import Line2D
class LineDataUnits(Line2D):
def __init__(self, *args, **kwargs):
_lw_data = kwargs.pop("linewidth", 1)
super().__init__(*args, **kwargs)
self._lw_data = _lw_data
def _get_lw(self):
if self.axes is not None:
ppd = 72. / self.axes.figure.dpi
trans = self.axes.transData.transform
return ((trans((1, self._lw_data)) - trans((0, 0))) * ppd)[1]
else:
return 1
def _set_lw(self, lw):
self._lw_data = lw
_linewidth = property(_get_lw, _set_lw)
line = LineDataUnits([start[0], end[0]], [start[1], end[1]], linewidth=width, color=color, **kwargs)
ax.add_line(line)
# ax.plot([start[0], end[0]], [start[1], end[1]], linestyle=linestyle, color=color, **kwargs)
def __copy__(self):
return self.__class__(start=self.start, end=self.end, gpts=self.gpts, endpoint=self.grid.endpoint[0])
[docs]class GridScan(AbstractScan, HasGridMixin):
"""
Grid scan object.
Defines a scan on a regular grid.
Parameters
----------
start : two float
Start corner of the scan [Å].
end : two float
End corner of the scan [Å].
gpts : two int
Number of scan positions in the x- and y-direction of the scan.
sampling : two float
Sampling rate of scan positions [1 / Å].
endpoint : bool
If True, end is the last position. Otherwise, it is not included. Default is False.
batch_partition : 'squares' or 'lines'
Specify how to split the scan into batches. If 'squares', the scan position batches are divided into the best
matching squares for the batch size. If 'lines', the batches are divided into lines of scan positions.
measurement_shift : two int
The insertion indices of new measurements will be shifted by this amount in x and y. This is used for
correctly inserting measurements collected from a partitioned scan.
"""
def __init__(self,
start: Sequence[float],
end: Sequence[float],
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
endpoint: bool = False,
batch_partition: str = 'squares',
measurement_shift: Sequence[int] = None):
super().__init__()
try:
self._start = np.array(start)[:2]
end = np.array(end)[:2]
assert (self._start.shape == (2,)) & (end.shape == (2,))
except:
raise ValueError('Scan start/end has incorrect shape')
if (gpts is None) & (sampling is None):
raise RuntimeError('Grid gpts or sampling must be set')
if not batch_partition.lower() in ['squares', 'lines']:
raise ValueError('batch partition must be "squares" or "lines"')
self._batch_partition = batch_partition
self._measurement_shift = measurement_shift
self._grid = Grid(extent=end - start, gpts=gpts, sampling=sampling, dimensions=2, endpoint=endpoint)
@property
def shape(self):
return self.gpts
@property
def calibrations(self) -> tuple:
return (Calibration(offset=self.start[0], sampling=self.sampling[0], units='Å', name='x',
endpoint=self.grid.endpoint[0]),
Calibration(offset=self.start[1], sampling=self.sampling[1], units='Å', name='y',
endpoint=self.grid.endpoint[1]))
@property
def start(self) -> np.ndarray:
"""Start corner of the scan [Å]."""
return self._start
@start.setter
def start(self, start: Sequence[float]):
self._start = np.array(start)
self.extent = self.end - self._start
@property
def end(self) -> np.ndarray:
"""End corner of the scan [Å]."""
return self.start + self.extent
@end.setter
def end(self, end: Sequence[float]):
self.extent = np.array(end) - self.start
[docs] def get_scan_area(self) -> float:
"""Get the area of the scan."""
height = abs(self.start[0] - self.end[0])
width = abs(self.start[1] - self.end[1])
return height * width
[docs] def get_positions(self) -> np.ndarray:
x = np.linspace(self.start[0], self.end[0], self.gpts[0], endpoint=self.grid.endpoint[0])
y = np.linspace(self.start[1], self.end[1], self.gpts[1], endpoint=self.grid.endpoint[1])
x, y = np.meshgrid(x, y, indexing='ij')
return np.stack((np.reshape(x, (-1,)),
np.reshape(y, (-1,))), axis=1)
[docs] def insert_new_measurement(self, measurement, indices: np.ndarray, new_measurement: np.ndarray):
x, y = np.unravel_index(indices, self.shape)
if self._measurement_shift is not None:
x += self._measurement_shift[0]
y += self._measurement_shift[1]
if isinstance(measurement, str):
with h5py.File(measurement, 'a') as f:
for unique in np.unique(x):
f['array'][unique, y[unique == x]] += asnumpy(new_measurement[unique == x])
else:
measurement.array[x, y] += asnumpy(new_measurement)
[docs] def partition_scan(self, partitions: Sequence[int]) -> List['GridScan']:
"""
Partition the scan into smaller grid scans
Parameters
----------
partitions : two int
The number of partitions to create in x and y.
Returns
-------
List of GridScan objects
"""
Nx = subdivide_into_batches(self.gpts[0], partitions[0])
Ny = subdivide_into_batches(self.gpts[1], partitions[1])
Sx = np.concatenate(([0], np.cumsum(Nx)))
Sy = np.concatenate(([0], np.cumsum(Ny)))
scans = []
for i, nx in enumerate(Nx):
for j, ny in enumerate(Ny):
start = [Sx[i] * self.sampling[0], Sy[j] * self.sampling[1]]
end = [start[0] + nx * self.sampling[0], start[1] + ny * self.sampling[1]]
endpoint = [False, False]
if i + 1 == partitions[0]:
endpoint[0] = self.grid.endpoint[0]
if endpoint[0]:
end[0] -= self.sampling[0]
if j + 1 == partitions[1]:
endpoint[1] = self.grid.endpoint[1]
if endpoint[1]:
end[1] -= self.sampling[1]
scan = self.__class__(start,
end,
gpts=(nx, ny),
endpoint=endpoint,
batch_partition='squares',
measurement_shift=(Sx[i], Sy[j]))
scans.append(scan)
return scans
def _partition_batches(self, max_batch: int):
if self._batch_partition == 'lines':
super()._partition_batches(max_batch)
return
if max_batch == 1:
self._batches = [[i] for i in range(len(self))]
return
max_batch_x = int(np.floor(np.sqrt(max_batch)))
max_batch_y = int(np.floor(np.sqrt(max_batch)))
Nx = subdivide_into_batches(self.gpts[0], (self.gpts[0] + (-self.gpts[0] % max_batch_x)) // max_batch_x)
Ny = subdivide_into_batches(self.gpts[1], (self.gpts[1] + (-self.gpts[1] % max_batch_y)) // max_batch_y)
self._batches = []
Sx = np.concatenate(([0], np.cumsum(Nx)))
Sy = np.concatenate(([0], np.cumsum(Ny)))
for i, nx in enumerate(Nx):
for j, ny in enumerate(Ny):
x = np.arange(Sx[i], Sx[i] + nx, dtype=int)
y = np.arange(Sy[j], Sy[j] + ny, dtype=int)
self._batches.append((y[None] + x[:, None] * self.gpts[1]).ravel())
[docs] def add_to_mpl_plot(self, ax, alpha: float = .33, facecolor: str = 'r', edgecolor: str = 'r', **kwargs):
"""
Add a visualization of the scan area to a matplotlib plot.
Parameters
----------
ax : matplotlib Axes
The axes of the matplotlib plot the visualization should be added to.
alpha : float, optional
Transparency of the scan area visualization. Default is 0.33.
facecolor : str, optional
Color of the scan area visualization.
edgecolor : str, optional
Color of the edge of the scan area visualization.
kwargs :
Additional options for matplotlib.patches.Rectangle used for scan area visualization as keyword arguments.
"""
rect = Rectangle(tuple(self.start), *self.extent, alpha=alpha, facecolor=facecolor, edgecolor=edgecolor,
**kwargs)
ax.add_patch(rect)
def __copy__(self):
return self.__class__(start=self.start,
end=self.end,
gpts=self.gpts,
endpoint=self.grid.endpoint,
batch_partition=self._batch_partition,
measurement_shift=self._measurement_shift)