def test_grid_event():
grid = Grid()
grid.extent = 5
assert grid.changed._notify_count == 1
grid.gpts = 100
assert grid.changed._notify_count == 2
grid.sampling = .1
assert grid.changed._notify_count == 3
def test_locked_grid():
grid = Grid(gpts=200, lock_gpts=True)
grid.extent = 10
assert (grid.sampling[0] == .05) & (grid.sampling[1] == .05)
grid.extent = 20
assert (grid.sampling[0] == .1) & (grid.sampling[1] == .1)
with pytest.raises(RuntimeError) as e:
grid.gpts = 100
assert str(e.value) == 'gpts cannot be modified'
def test_change_grid():
grid = Grid(extent=(8, 6), gpts=10)
grid.sampling = .2
assert (grid.extent[0] == 8.) & (grid.extent[1] == 6.)
assert (grid.gpts[0] == 40) & (grid.gpts[1] == 30)
grid.gpts = 100
assert (grid.extent[0] == 8.) & (grid.extent[1] == 6.)
assert (grid.sampling[0] == .08) & (grid.sampling[1] == .06)
grid.extent = (16, 12)
assert (grid.gpts[0] == 100) & (grid.gpts[1] == 100)
assert (grid.extent[0] == 16.) & (grid.extent[1] == 12.)
assert (grid.sampling[0] == .16) & (grid.sampling[1] == .12)
grid.extent = (10, 10)
assert (grid.sampling[0] == grid.extent[0] /
grid.gpts[0]) & (grid.sampling[1] == grid.extent[1] / grid.gpts[1])
grid.sampling = .3
assert (grid.extent[0] == grid.sampling[0] *
grid.gpts[0]) & (grid.extent[1] == grid.sampling[1] * grid.gpts[1])
grid.gpts = 30
assert (grid.sampling[0] == grid.extent[0] /
grid.gpts[0]) & (grid.sampling[1] == grid.extent[1] / grid.gpts[1])
def allocate_measurement(self, grid: Grid, wavelength: float,
scan: AbstractScan) -> Measurement:
grid.check_is_defined()
calibrations = calibrations_from_grid(grid.gpts,
grid.sampling,
names=['x', 'y'],
units='Å')
array = np.zeros(scan.shape + grid.gpts, dtype=np.complex64)
measurement = Measurement(array,
calibrations=scan.calibrations +
calibrations)
if isinstance(self.save_file, str):
measurement = measurement.write(self.save_file)
return measurement
def __init__(self,
array: np.ndarray,
expansion_cutoff: float,
interpolation: int,
k: Tuple[np.ndarray, np.ndarray],
ctf: CTF = None,
extent: Union[float, Sequence[float]] = None,
sampling: Union[float, Sequence[float]] = None,
energy: float = None,
device='cpu'):
self._array = array
self._interpolation = interpolation
self._expansion_cutoff = expansion_cutoff
self._k = k
self._grid = Grid(extent=extent,
gpts=array.shape[1:],
sampling=sampling,
lock_gpts=True)
self._accelerator = Accelerator(energy=energy)
if ctf is None:
ctf = CTF(semiangle_cutoff=expansion_cutoff, rolloff=.1)
self.set_ctf(ctf)
self._device = device
def test_create_grid():
grid = Grid(extent=5, sampling=.2)
assert (grid.extent[0] == 5.) & (grid.extent[1] == 5.)
assert (grid.gpts[0] == 25) & (grid.gpts[1] == 25)
assert (grid.sampling[0] == .2) & (grid.sampling[1] == .2)
grid = Grid(sampling=.2, gpts=10)
assert (grid.extent[0] == 2.) & (grid.extent[1] == 2.)
grid = Grid(extent=(8, 6), gpts=10)
assert (grid.sampling[0] == .8) & (grid.sampling[1] == .6)
grid = Grid()
with pytest.raises(RuntimeError):
grid.check_is_defined()
def __init__(self,
calculator,
gpts=None,
sampling=None,
slice_thickness=.5,
core_size=.005,
storage='cpu'):
self._calculator = calculator
self._core_size = core_size
thickness = calculator.atoms.cell[2, 2]
nz = calculator.hamiltonian.finegd.N_c[2]
num_slices = int(
np.ceil(nz / np.floor(slice_thickness / (thickness / nz))))
self._voxel_height = thickness / nz
self._slice_vertical_voxels = split_integer(nz, num_slices)
# TODO: implement support for non-periodic extent
self._origin = (0., 0.)
extent = np.diag(
orthogonalize_cell(calculator.atoms.copy(),
strain_error=True).cell)[:2]
self._grid = Grid(extent=extent,
gpts=gpts,
sampling=sampling,
lock_extent=True)
super().__init__(storage)
def __init__(self, array, thickness, extent=None, sampling=None):
self._array = array
self._thickness = thickness
self._grid = Grid(extent=extent,
gpts=array.shape[-2:],
sampling=sampling,
lock_gpts=True)
def __init__(self,
expansion_cutoff: float,
interpolation: int = 1,
ctf: CTF = None,
extent: Union[float, Sequence[float]] = None,
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
energy: float = None,
device: str = 'cpu',
storage: str = None):
if not isinstance(interpolation, int):
raise ValueError('Interpolation factor must be int')
self._interpolation = interpolation
self._expansion_cutoff = expansion_cutoff
self._ctf = ctf
self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling)
self._accelerator = Accelerator(energy=energy)
self._device = device
if storage is None:
storage = device
self._storage = storage
def __init__(self,
semiangle_cutoff: float = np.inf,
rolloff: float = 0.1,
focal_spread: float = 0.,
angular_spread: float = 0.,
ctf_parameters: dict = None,
extent: Union[float, Sequence[float]] = None,
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
energy: float = None,
device='cpu',
**kwargs):
self._ctf = CTF(semiangle_cutoff=semiangle_cutoff,
rolloff=rolloff,
focal_spread=focal_spread,
angular_spread=angular_spread,
parameters=ctf_parameters,
energy=energy,
**kwargs)
self._accelerator = self._ctf._accelerator
self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling)
self._ctf_cache = Cache(1)
self._ctf.changed.register(cache_clear_callback(self._ctf_cache))
self._grid.changed.register(cache_clear_callback(self._ctf_cache))
self._accelerator.changed.register(
cache_clear_callback(self._ctf_cache))
self._device = device
def interpolated_grid(self) -> Grid:
"""
The grid of the interpolated scattering matrix.
"""
interpolated_gpts = tuple(n // self.interpolation for n in self.gpts)
return Grid(gpts=interpolated_gpts,
sampling=self.sampling,
lock_gpts=True)
def __init__(self,
extent: Union[float, Sequence[float]] = None,
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
energy: float = None,
device: str = 'cpu'):
self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling)
self._accelerator = Accelerator(energy=energy)
self._device = device
def allocate_measurement(self, grid: Grid, wavelength: float,
scan: AbstractScan) -> Measurement:
grid.check_is_defined()
shape = (grid.gpts[0] // 2, grid.gpts[1] // 2)
calibrations = calibrations_from_grid(grid.antialiased_gpts,
grid.antialiased_sampling,
names=['alpha_x', 'alpha_y'],
units='mrad',
scale_factor=wavelength * 1000,
fourier_space=True)
array = np.zeros(scan.shape + shape)
measurement = Measurement(array,
calibrations=scan.calibrations +
calibrations)
if isinstance(self.save_file, str):
measurement = measurement.write(self.save_file)
return measurement
def __init__(self, start, end, gpts=None, sampling=None, endpoint=False):
super().__init__()
self._start = np.array(start)
end = np.array(end)
if (self._start.shape != (2, )) | (end.shape != (2, )):
raise ValueError('Scan start/end has incorrect shape')
self._grid = Grid(extent=end - start,
gpts=gpts,
sampling=sampling,
dimensions=2,
endpoint=endpoint)
def show(self,
grid: Grid,
wavelength: float,
cbar_label: str = 'Detector regions',
**kwargs):
grid.check_is_defined()
array = np.full(grid.antialiased_gpts, -1, dtype=np.int)
for i, indices in enumerate(
self._get_regions(grid.antialiased_gpts,
grid.antialiased_sampling, wavelength)):
array.ravel()[indices] = i
calibrations = calibrations_from_grid(grid.antialiased_gpts,
grid.antialiased_sampling,
names=['alpha_x', 'alpha_y'],
units='mrad.',
scale_factor=wavelength * 1000,
fourier_space=True)
return show_image(array,
calibrations,
cbar_label=cbar_label,
discrete=True,
**kwargs)
def __init__(self,
array: np.ndarray,
extent: Union[float, Sequence[float]] = None,
sampling: Union[float, Sequence[float]] = None,
energy: float = None):
if len(array.shape) < 2:
raise RuntimeError(
'Wave function array should be have 2 dimensions or more')
self._array = array
self._grid = Grid(extent=extent,
gpts=array.shape[-2:],
sampling=sampling,
lock_gpts=True)
self._accelerator = Accelerator(energy=energy)
def __init__(self,
start: Sequence[float],
end: Sequence[float],
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
endpoint: bool = True):
super().__init__()
start = np.array(start)
end = np.array(end)
if (start.shape != (2, )) | (end.shape != (2, )):
raise ValueError('Scan start/end has incorrect shape')
self._grid = Grid(gpts=gpts,
sampling=sampling,
endpoint=endpoint,
dimensions=1)
self._start = start
self._direction, self.extent = self._direction_and_extent(start, end)
def __init__(self,
array: np.ndarray,
slice_thicknesses: Union[float, Sequence],
extent: Union[float, Sequence[float]] = None,
sampling: Union[float, Sequence[float]] = None):
if len(array.shape) != 3:
raise RuntimeError()
slice_thicknesses = np.array(slice_thicknesses)
if slice_thicknesses.shape == ():
slice_thicknesses = np.tile(slice_thicknesses, array.shape[0])
elif slice_thicknesses.shape != (array.shape[0], ):
raise ValueError()
self._array = array
self._slice_thicknesses = slice_thicknesses
self._grid = Grid(extent=extent,
gpts=self.array.shape[1:],
sampling=sampling,
lock_gpts=True)
def __init__(self, extent=None, gpts=None, sampling=None, num_slices=10):
self._num_slices = num_slices
self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling)
def test_grid_raises():
with pytest.raises(RuntimeError) as e:
Grid(extent=[5, 5, 5])
assert str(e.value) == 'grid value length of 3 != 2'
def test_grid_match():
grid1 = Grid(extent=10, gpts=10)
grid2 = Grid()
grid1.match(grid2)
grid1.check_match(grid2)
grid2.sampling = .2
with pytest.raises(RuntimeError) as e:
grid1.check_match(grid2)
assert str(e.value) == 'inconsistent grid gpts ((10, 10) != (50, 50))'
grid1.match(grid2)
grid1.check_match(grid2)
def interpolated_grid(self) -> Grid:
interpolated_gpts = tuple(n // self.interpolation for n in self.gpts)
return Grid(gpts=interpolated_gpts,
sampling=self.sampling,
lock_gpts=True)
def __init__(self,
atoms: Union[Atoms, AbstractFrozenPhonons] = None,
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
slice_thickness: float = .5,
parametrization: str = 'lobato',
cutoff_tolerance: float = 1e-3,
device='cpu',
storage=None):
self._cutoff_tolerance = cutoff_tolerance
self._parametrization = parametrization
self._slice_thickness = slice_thickness
self._storage = storage
if parametrization == 'lobato':
self._parameters = load_lobato_parameters()
self._function = lobato # lambda r: lobato(r, parameters[atomic_number])
self._derivative = dvdr_lobato # lambda r: dvdr_lobato(r, parameters[atomic_number])
elif parametrization == 'kirkland':
self._parameters = load_kirkland_parameters()
self._function = kirkland # lambda r: kirkland(r, parameters[atomic_number])
self._derivative = dvdr_kirkland # lambda r: dvdr_kirkland(r, parameters[atomic_number])
else:
raise RuntimeError(
'Parametrization {} not recognized'.format(parametrization))
if isinstance(atoms, AbstractFrozenPhonons):
self._frozen_phonons = atoms
else:
self._frozen_phonons = DummyFrozenPhonons(atoms)
atoms = next(iter(self._frozen_phonons))
if np.abs(atoms.cell[2, 2]) < 1e-12:
raise RuntimeError('Atoms cell has no thickness')
if not is_cell_orthogonal(atoms):
raise RuntimeError('Atoms are not orthogonal')
self._atoms = atoms
self._grid = Grid(extent=np.diag(atoms.cell)[:2],
gpts=gpts,
sampling=sampling,
lock_extent=True)
self._cutoffs = {}
self._integrators = {}
def grid_changed_callback(*args, **kwargs):
self._integrators = {}
self.grid.changed.register(grid_changed_callback)
self.changed = Event()
self._device = device
if storage is None:
self._storage = device
super().__init__(storage=storage)