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