def test_grid_match():
grid1 = Grid(extent=10, gpts=10)
grid2 = Grid()
grid1.match(grid2)
grid1.check_match(grid2)
grid2.sampling = .2
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) == 'Grid gpts cannot be modified'
def __init__(self,
array: np.ndarray,
slice_thicknesses: Union[float, Sequence[float]],
extent: Union[float, Sequence[float]] = None,
sampling: Union[float, Sequence[float]] = None):
if (len(array.shape) != 2) & (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], )) & (len(array.shape) == 3):
raise ValueError()
self._array = array
self._slice_thicknesses = slice_thicknesses
self._grid = Grid(extent=extent,
gpts=self.array.shape[-2:],
sampling=sampling,
lock_gpts=True)
super().__init__(precalculate=False)
def __init__(self,
transitions,
atoms=None,
slice_thickness=None,
gpts: Union[int, Sequence[float]] = None,
sampling: Union[float, Sequence[float]] = None,
energy: float = None,
min_contrast=.95):
if isinstance(transitions,
(SubshellTransitions, SubshellTransitionsArrays)):
transitions = [transitions]
self._slice_thickness = slice_thickness
self._grid = Grid(gpts=gpts, sampling=sampling)
self.atoms = atoms
self._transitions = transitions
self._accelerator = Accelerator(energy=energy)
self._sliced_atoms = SlicedAtoms(
atoms, slice_thicknesses=self._slice_thickness)
self._potentials_cache = Cache(1)
def __init__(self,
start: Sequence[float],
end: Sequence[float],
gpts: int = None,
sampling: float = None,
endpoint: bool = True):
super().__init__()
try:
start = np.array(start)[:2]
end = np.array(end)[:2]
assert (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')
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,
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)
def __init__(self,
calculator,
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
# origin: Union[float, Sequence[float]] = None,
slice_thickness=.5,
core_size=.005,
storage='cpu',
precalculate=True):
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 = subdivide_into_batches(nz, num_slices)
# TODO: implement support for non-periodic extent
self._origin = (0., 0.)
extent = np.diag(orthogonalize_cell(calculator.atoms.copy()).cell)[:2]
self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling, lock_extent=True)
super().__init__(precalculate=precalculate, storage=storage)
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)
print(grid.sampling[0] == .2, type(grid.sampling[0]), type(.2))
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 grid_from_calibrations(calibrations, extent=None, gpts=None) -> Grid:
if (extent is None) and (gpts is None):
raise RuntimeError
sampling = ()
for calibration in calibrations:
sampling += (calibration.sampling, )
return Grid(extent=extent, gpts=gpts, sampling=sampling)
def _interpolate_2d(self,
new_sampling: Union[float, Tuple[float, float]] = None,
new_gpts: Union[int, Tuple[int, int]] = None,
padding: str = 'wrap',
kind: str = None) -> 'Measurement':
if kind is None:
kind = 'quintic'
if not (self.calibrations[-1].units == self.calibrations[-2].units):
raise RuntimeError(
'the units of the interpolation dimensions must match')
endpoint = tuple(
[calibration.endpoint for calibration in self.calibrations])
sampling = tuple(
[calibration.sampling for calibration in self.calibrations])
offset = tuple(
[calibration.offset for calibration in self.calibrations])
extent = (sampling[0] * (self.array.shape[0] - endpoint[0]),
sampling[1] * (self.array.shape[1] - endpoint[1]))
new_grid = Grid(extent=extent,
gpts=new_gpts,
sampling=new_sampling,
endpoint=endpoint)
array = np.pad(self.array, ((5, ) * 2, ) * 2, mode=padding)
x = self.calibrations[0].coordinates(
array.shape[0]) - 5 * self.calibrations[0].sampling
y = self.calibrations[1].coordinates(
array.shape[1]) - 5 * self.calibrations[1].sampling
interpolator = interp2d(x, y, array.T, kind=kind)
x = np.linspace(offset[0],
offset[0] + extent[0],
new_grid.gpts[0],
endpoint=endpoint[0])
y = np.linspace(offset[1],
offset[1] + extent[1],
new_grid.gpts[1],
endpoint=endpoint[1])
new_array = interpolator(x, y).T
calibrations = []
for calibration, d in zip(self.calibrations, new_grid.sampling):
calibrations.append(copy(calibration))
calibrations[-1].sampling = d
return self.__class__(new_array,
calibrations,
name=self.name,
units=self.units)
def interpolate(self,
new_sampling: Union[float, Sequence[float]],
padding: str = 'wrap',
kind: str = 'quintic') -> 'Measurement':
"""
Interpolate a 2d measurement.
Parameters
----------
new_sampling : float
Target measurement sampling. Same units as measurement calibrations.
padding : str
The padding mode as used by numpy.pad.
kind : str
The kind of spline interpolation to use. Default is 'quintic'.
Returns
-------
Measurement object
Interpolated measurement
"""
if self.dimensions != 2:
raise RuntimeError('interpolate only implemented for 2d measurements')
if not (self.calibrations[-1].units == self.calibrations[-2].units):
raise RuntimeError('the units of the interpolation dimensions must match')
endpoint = tuple([calibration.endpoint for calibration in self.calibrations])
sampling = tuple([calibration.sampling for calibration in self.calibrations])
offset = tuple([calibration.offset for calibration in self.calibrations])
extent = (sampling[0] * (self.array.shape[0] - endpoint[0]),
sampling[1] * (self.array.shape[1] - endpoint[1]))
new_grid = Grid(extent=extent, sampling=new_sampling, endpoint=endpoint)
array = np.pad(self.array, ((5,) * 2,) * 2, mode=padding)
x = self.calibrations[0].coordinates(array.shape[0]) - 5 * self.calibrations[0].sampling
y = self.calibrations[1].coordinates(array.shape[1]) - 5 * self.calibrations[1].sampling
interpolator = interp2d(x, y, array.T, kind=kind)
x = np.linspace(offset[0], offset[0] + extent[0], new_grid.gpts[0], endpoint=endpoint[0])
y = np.linspace(offset[1], offset[1] + extent[1], new_grid.gpts[1], endpoint=endpoint[1])
new_array = interpolator(x, y).T
calibrations = []
for calibration, d in zip(self.calibrations, new_grid.sampling):
calibrations.append(copy(calibration))
calibrations[-1].sampling = d
return self.__class__(new_array, calibrations, name=self.name, units=self.units)
def _polar_coordinates(self, gpts=None, extent=None, sampling=None, xp=np):
grid = Grid(gpts=gpts, extent=extent, sampling=sampling)
gpts = grid.gpts
sampling = grid.sampling
kx, ky = spatial_frequencies(gpts, sampling)
kx = kx.reshape((1, -1, 1))
ky = ky.reshape((1, 1, -1))
kx = xp.asarray(kx)
ky = xp.asarray(ky)
return polar_coordinates(xp.asarray(kx * self.wavelength),
xp.asarray(ky * self.wavelength))
def __init__(self,
calculator,
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
origin: Union[float, Sequence[float]] = None,
orthogonal_cell: Sequence[float] = None,
periodic_z: bool = True,
slice_thickness=.5,
core_size=.005,
plane='xy',
storage='cpu',
precalculate=True):
self._calculator = calculator
self._core_size = core_size
self._plane = plane
if orthogonal_cell is None:
atoms = rotate_atoms_to_plane(calculator.atoms, plane)
thickness = atoms.cell[2, 2]
nz = calculator.hamiltonian.finegd.N_c[plane_to_axes(plane)[-1]]
extent = np.diag(orthogonalize_cell(atoms.copy()).cell)[:2]
else:
if plane != 'xy':
raise NotImplementedError()
thickness = orthogonal_cell[2]
nz = calculator.hamiltonian.finegd.N_c / np.linalg.norm(calculator.atoms.cell, axis=0) * orthogonal_cell[2]
nz = int(np.ceil(np.max(nz)))
extent = orthogonal_cell[:2]
num_slices = int(np.ceil(nz / np.floor(slice_thickness / (thickness / nz))))
self._orthogonal_cell = orthogonal_cell
self._voxel_height = thickness / nz
self._slice_vertical_voxels = subdivide_into_batches(nz, num_slices)
self._origin = (0., 0., 0.)
self._periodic_z = periodic_z
self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling, lock_extent=True)
super().__init__(precalculate=precalculate, storage=storage)
def _interpolate_1d(self,
new_sampling: float = None,
new_gpts: int = None,
padding: str = 'wrap',
kind: str = None) -> 'Measurement':
if kind is None:
kind = 'quadratic'
endpoint = self.calibrations[-1].endpoint
sampling = self.calibrations[-1].sampling
offset = self.calibrations[-1].offset
extent = sampling * (self.array.shape[-1] - endpoint)
new_grid = Grid(extent=extent,
gpts=new_gpts,
sampling=new_sampling,
endpoint=endpoint)
array = np.pad(self.array, ((5, ) * 2, ), mode=padding)
x = self.calibrations[-1].coordinates(array.shape[-1]) - 5 * sampling
interpolator = interp1d(x, array, kind=kind)
x = np.linspace(offset,
offset + extent,
new_grid.gpts[0],
endpoint=endpoint)
new_array = interpolator(x)
calibrations = [
calibration.copy() for calibration in self.calibrations
]
calibrations[-1].sampling = new_grid.sampling[0]
return self.__class__(new_array,
calibrations,
name=self.name,
units=self.units)
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 __init__(self,
potential_unit: AbstractPotential,
repetitions: Tuple[int, int, int],
num_frozen_phonon_configs: int = 1):
self._potential_unit = potential_unit
self.repetitions = repetitions
self._num_frozen_phonon_configs = num_frozen_phonon_configs
if (potential_unit.num_frozen_phonon_configs
== 1) & (num_frozen_phonon_configs > 1):
warnings.warn(
'"num_frozen_phonon_configs" is greater than one, but the potential unit does not have'
'frozen phonons')
if (potential_unit.num_frozen_phonon_configs >
1) & (num_frozen_phonon_configs == 1):
warnings.warn(
'the potential unit has frozen phonons, but "num_frozen_phonon_configs" is set to 1'
)
self._cache = Cache(1)
self._changed = Event()
gpts = (self._potential_unit.gpts[0] * self.repetitions[0],
self._potential_unit.gpts[1] * self.repetitions[1])
extent = (self._potential_unit.extent[0] * self.repetitions[0],
self._potential_unit.extent[1] * self.repetitions[1])
self._grid = Grid(extent=extent,
gpts=gpts,
sampling=self._potential_unit.sampling,
lock_extent=True)
self._grid.changed.register(self._changed.notify)
self._changed.register(cache_clear_callback(self._cache))
super().__init__(precalculate=False)
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 __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)
super().__init__(precalculate=False)
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 __init__(self, Z, extent, gpts, sampling, energy):
self._Z = Z
self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling)
self._accelerator = Accelerator(energy=energy)
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',
projection: str = 'finite',
cutoff_tolerance: float = 1e-3,
device='cpu',
precalculate: bool = True,
storage=None):
self._cutoff_tolerance = cutoff_tolerance
self._parametrization = parametrization
self._slice_thickness = slice_thickness
self._storage = storage
if parametrization.lower() == 'lobato':
self._parameters = load_lobato_parameters()
self._function = lobato
self._derivative = dvdr_lobato
elif parametrization.lower() == 'kirkland':
self._parameters = load_kirkland_parameters()
self._function = kirkland
self._derivative = dvdr_kirkland
else:
raise RuntimeError(
'Parametrization {} not recognized'.format(parametrization))
if projection == 'infinite':
if parametrization.lower() != 'kirkland':
raise RuntimeError(
'Infinite projections are only implemented for the Kirkland parametrization'
)
elif (projection != 'finite'):
raise RuntimeError('Projection must be "finite" or "infinite"')
self._projection = projection
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 = {}
self._disc_indices = {}
def grid_changed_callback(*args, **kwargs):
self._integrators = {}
self._disc_indices = {}
self.grid.changed.register(grid_changed_callback)
self.changed = Event()
if storage is None:
storage = device
super().__init__(precalculate=precalculate,
device=device,
storage=storage)
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))'
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'