def transmit(waves: Union['Waves', 'SMatrix', 'PartialSMatrix'],
potential_slice: ProjectedPotential):
"""
Transmit wave function or scattering matrix.
:param waves: Wave function or scattering matrix to propagate.
:param potential_slice: Projected potential to transmit the wave function through.
"""
xp = get_array_module(waves.array)
complex_exponential = get_device_function(xp, 'complex_exponential')
dim_padding = len(waves._array.shape) - len(potential_slice.array.shape)
slice_array = potential_slice.array.reshape((1, ) * dim_padding +
potential_slice.array.shape)
if np.iscomplexobj(slice_array):
waves._array *= copy_to_device(slice_array, xp)
else:
waves._array *= complex_exponential(
copy_to_device(waves.accelerator.sigma * slice_array, xp))
def get_mask(self, gpts, sampling, xp):
if sampling is None:
sampling = (1., 1.)
kx, ky = spatial_frequencies(gpts, sampling)
kx = copy_to_device(kx, xp)
ky = copy_to_device(ky, xp)
k = xp.sqrt(kx[:, None]**2 + ky[None]**2)
kcut = 1 / max(sampling) / 2 * self.cutoff
if self.rolloff > 0.:
array = .5 * (1 + xp.cos(np.pi *
(k - kcut + self.rolloff) / self.rolloff))
array[k > kcut] = 0.
array = xp.where(k > kcut - self.rolloff, array,
xp.ones_like(k, dtype=xp.float32))
else:
array = xp.array(k < kcut).astype(xp.float32)
return array
def build(self) -> SMatrix:
self.grid.check_is_defined()
self.accelerator.check_is_defined()
xp = get_array_module(self._device)
storage_xp = get_array_module_from_device(self._storage)
complex_exponential = get_device_function(xp, 'complex_exponential')
n_max = int(
xp.ceil(self.expansion_cutoff / 1000. /
(self.wavelength / self.extent[0] * self.interpolation)))
m_max = int(
xp.ceil(self.expansion_cutoff / 1000. /
(self.wavelength / self.extent[1] * self.interpolation)))
n = xp.arange(-n_max, n_max + 1, dtype=xp.float32)
w = xp.asarray(self.extent[0], dtype=xp.float32)
m = xp.arange(-m_max, m_max + 1, dtype=xp.float32)
h = xp.asarray(self.extent[1], dtype=xp.float32)
kx = n / w * xp.float32(self.interpolation)
ky = m / h * xp.float32(self.interpolation)
mask = kx[:, None]**2 + ky[None, :]**2 < (self.expansion_cutoff /
1000. / self.wavelength)**2
kx, ky = xp.meshgrid(kx, ky, indexing='ij')
kx = kx[mask]
ky = ky[mask]
x, y = coordinates(extent=self.extent,
gpts=self.gpts,
endpoint=self.grid.endpoint)
x = xp.asarray(x)
y = xp.asarray(y)
array = storage_xp.zeros((len(kx), ) + (self.gpts[0], self.gpts[1]),
dtype=np.complex64)
for i in range(len(kx)):
array[i] = copy_to_device(
complex_exponential(
-2 * np.pi * kx[i, None, None] * x[:, None]) *
complex_exponential(
-2 * np.pi * ky[i, None, None] * y[None, :]),
self._storage)
return SMatrix(array,
expansion_cutoff=self.expansion_cutoff,
interpolation=self.interpolation,
extent=self.extent,
energy=self.energy,
k=(kx, ky))
def build(self,
pbar: Union[bool, ProgressBar] = False) -> 'ArrayPotential':
self.grid.check_is_defined()
storage_xp = get_array_module_from_device(self._storage)
array = storage_xp.zeros(
(self.num_slices, ) + (self.gpts[0], self.gpts[1]),
dtype=np.float32)
slice_thicknesses = np.zeros(self.num_slices)
if isinstance(pbar, bool):
pbar = ProgressBar(total=len(self),
desc='Potential',
disable=not pbar)
pbar.reset()
for i, potential_slice in enumerate(self.generate_slices()):
array[i] = copy_to_device(potential_slice.array, self._storage)
slice_thicknesses[i] = potential_slice.thickness
pbar.update(1)
pbar.refresh()
return ArrayPotential(array, slice_thicknesses, self.extent)
def build(
self,
first_slice: int = 0,
last_slice: int = None,
energy: float = None,
max_batch: int = None,
pbar: Union[bool, ProgressBar] = False,
) -> 'PotentialArray':
"""
Precalcaulate the potential as a potential array.
Parameters
----------
first_slice: int
First potential slice to generate.
last_slice: int, optional
Last potential slice generate.
energy: float
Electron energy [eV]. If given, the transmission functions will be returned.
max_batch: int
Maximum number of potential slices calculated in parallel.
pbar: bool
If true, show progress bar.
Returns
-------
PotentialArray object
"""
self.grid.check_is_defined()
if last_slice is None:
last_slice = len(self)
if max_batch is None:
max_batch = self._estimate_max_batch()
storage_xp = get_array_module_from_device(self._storage)
if energy is None:
array = storage_xp.zeros(
(last_slice - first_slice, ) + (self.gpts[0], self.gpts[1]),
dtype=np.float32)
generator = self.generate_slices(max_batch=max_batch,
first_slice=first_slice,
last_slice=last_slice)
else:
array = storage_xp.zeros(
(last_slice - first_slice, ) + (self.gpts[0], self.gpts[1]),
dtype=np.complex64)
generator = self.generate_transmission_functions(
energy=energy,
max_batch=max_batch,
first_slice=first_slice,
last_slice=last_slice)
slice_thicknesses = np.zeros(last_slice - first_slice)
if isinstance(pbar, bool):
pbar = ProgressBar(total=len(self),
desc='Potential',
disable=not pbar)
close_pbar = True
else:
close_pbar = False
pbar.reset()
for start, end, potential_slice in generator:
array[start:end] = copy_to_device(potential_slice.array,
self._storage)
slice_thicknesses[start:end] = potential_slice.slice_thicknesses
pbar.update(end - start)
pbar.refresh()
if close_pbar:
pbar.close()
if energy is None:
return PotentialArray(array,
slice_thicknesses=slice_thicknesses,
extent=self.extent)
else:
return TransmissionFunction(array,
slice_thicknesses=slice_thicknesses,
extent=self.extent,
energy=energy)
def transmit(self, waves):
self.accelerator.check_match(waves)
xp = get_array_module(waves._array)
waves._array *= copy_to_device(self.array, xp)
return waves
