def insert_new_measurement(self, measurement, start, end, new_measurement):
if isinstance(measurement, str):
with h5py.File(measurement, 'a') as f:
f['array'][start:end] = asnumpy(new_measurement)
else:
measurement.array[start:end] = asnumpy(new_measurement)
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)
def insert_new_measurement(self, measurement, start, end, new_measurement):
for row, slic, slic_1d in zip(
*unravel_slice_2d(start, end, self.shape)):
if isinstance(measurement, str):
with h5py.File(measurement, 'a') as f:
f['array'][row, slic] = asnumpy(new_measurement[slic_1d])
else:
measurement.array[row,
slic] = asnumpy(new_measurement[slic_1d])
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)
def test_potential_build_gpu():
atoms = Atoms('CO', positions=[(2, 3, 1), (3, 2, 3)], cell=(4, 6, 4.3))
potential = Potential(atoms=atoms, sampling=.1, device='gpu')
array_potential = potential.build()
assert np.all(asnumpy(array_potential[2].array) == asnumpy(potential[2].array))
potential = Potential(atoms=atoms, sampling=.1, device='cpu')
assert np.allclose(asnumpy(array_potential[2].array), potential[2].array)
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)
def __iadd__(self, other):
if isinstance(other, self.__class__):
self.check_match_calibrations(other)
self._array += other.array
else:
self._array += asnumpy(other)
return self
def __add__(self, other):
if isinstance(other, self.__class__):
self.check_match_calibrations(other)
new_array = self.array + other.array
else:
new_array = self._array + asnumpy(other)
return self.__class__(new_array, calibrations=self.calibrations, units=self.units, name=self.name)
def __init__(self,
array: Union[np.ndarray, 'Measurement'],
calibrations: Union[Calibration,
Sequence[Union[Calibration,
None]]] = None,
units: str = '',
name: str = ''):
if issubclass(array.__class__, self.__class__):
measurement = array
array = measurement.array
calibrations = measurement.calibrations
units = measurement.array
name = measurement.name
if not isinstance(calibrations, Iterable):
calibrations = [calibrations] * len(array.shape)
if len(calibrations) != len(array.shape):
raise RuntimeError(
'The number of calibrations must equal the number of array dimensions. For undefined use None.'
)
self._array = asnumpy(array)
self._calibrations = calibrations
self._units = units
self._name = name
def __init__(self, array, calibrations, units='', name=''):
if len(calibrations) != len(array.shape):
raise RuntimeError()
self._array = asnumpy(array)
self._calibrations = calibrations
self._units = units
self._name = name
def squeeze(self):
new_meaurement = self.copy()
calibrations = [
calib
for calib, num_elem in zip(self.calibrations, self.array.shape)
if num_elem > 1
]
new_meaurement._calibrations = calibrations
new_meaurement._array = np.squeeze(asnumpy(new_meaurement.array))
return new_meaurement
def test_prism_gpu():
potential = Potential(Atoms('C', positions=[(2.5, 2.5, 2)],
cell=(5, 5, 4)))
S_builder = SMatrix(30, 80e3, 1, gpts=500, device='cpu')
S_cpu = S_builder.multislice(potential, pbar=False)
assert type(S_cpu.array) is np.ndarray
S_builder = SMatrix(30, 80e3, 1, gpts=500, device='gpu')
S_gpu = S_builder.multislice(potential, pbar=False)
assert type(S_gpu.array) is cp.ndarray
assert np.allclose(S_cpu.array, asnumpy(S_gpu.array))
def project(self):
"""
Create a 2d measurement of the projected potential.
Returns
-------
Measurement
"""
calibrations = calibrations_from_grid(self.grid.gpts,
self.grid.sampling,
names=['x', 'y'])
array = asnumpy(self.array.sum(0))
array -= array.min()
return Measurement(array, calibrations)
def squeeze(self) -> 'Measurement':
"""
Remove dimensions of length one from measurement.
Returns
-------
Measurement
"""
new_meaurement = self.copy()
calibrations = [
calib
for calib, num_elem in zip(self.calibrations, self.array.shape)
if num_elem > 1
]
new_meaurement._calibrations = calibrations
new_meaurement._array = np.squeeze(asnumpy(new_meaurement.array))
return new_meaurement
def diffraction_pattern(self) -> Measurement:
"""
:return: The intensity of the wave functions at the diffraction plane.
"""
calibrations = calibrations_from_grid(self.grid.antialiased_gpts,
self.grid.antialiased_sampling,
names=['alpha_x', 'alpha_y'],
units='mrad',
scale_factor=self.wavelength *
1000,
fourier_space=True)
calibrations = (None, ) * (len(self.array.shape) - 2) + calibrations
xp = get_array_module(self.array)
abs2 = get_device_function(xp, 'abs2')
fft2 = get_device_function(xp, 'fft2')
pattern = asnumpy(
abs2(
crop_to_center(
xp.fft.fftshift(fft2(self.array, overwrite_x=False)))))
return Measurement(pattern, calibrations)
def show(self, **kwargs):
calibrations = [
calib
for calib, num_elem in zip(self.calibrations, self.array.shape)
if num_elem > 1
]
array = np.squeeze(asnumpy(self.array))
dims = len(array.shape)
cbar_label = self._name + ' [' + self._units + ']'
if dims == 1:
return show_line(array, calibrations[0], **kwargs)
elif dims == 2:
return show_image(array,
calibrations,
cbar_label=cbar_label,
**kwargs)
else:
raise RuntimeError(
'Plotting not supported for {}d measurement, use reduction operation first'
.format(dims))
def __isub__(self, other):
if isinstance(other, self.__class__):
self.check_match_calibrations(other)
self._array -= other.array
else:
self._array -= asnumpy(other)
