def test_downsample():
f = AntialiasFilter()
sampling = (.1, .1)
gpts = (228, 229)
mask = f.get_mask(gpts, sampling, np)
n = np.sum(mask > 0.)
array = np.fft.ifft2(mask)
waves = Waves(array, sampling=sampling, energy=80e3, antialias_aperture=(2 / 3.,) * 2)
assert np.allclose(waves.downsample('valid', return_fourier_space=True).array.real, 1.)
def generate_transmission_functions(self,
energy,
first_slice=0,
last_slice=None,
max_batch=1):
"""
Generate the transmission functions one slice at a time.
Parameters
----------
energy: float
Electron energy [eV].
first_slice: int
First potential slice to generate.
last_slice: int, optional
Last potential slice generate.
max_batch: int
Maximum number of potential slices calculated in parallel.
Returns
-------
generator of PotentialArray objects
"""
antialias_filter = AntialiasFilter()
for start, end, potential_slice in self.generate_slices(
first_slice, last_slice, max_batch=max_batch):
yield start, end, potential_slice.as_transmission_function(
energy,
in_place=True,
max_batch=max_batch,
antialias_filter=antialias_filter)
def as_transmission_function(self,
energy: float,
in_place: bool = True,
max_batch: int = 1,
antialias_filter: AntialiasFilter = None):
"""
Calculate the transmission functions for a specific energy.
Parameters
----------
energy: float
Electron energy [eV].
Returns
-------
TransmissionFunction object
"""
xp = get_array_module(self.array)
complex_exponential = get_device_function(xp, 'complex_exponential')
array = self._array
if not in_place:
array = array.copy()
array = complex_exponential(energy2sigma(energy) * array)
t = TransmissionFunction(
array,
slice_thicknesses=self._slice_thicknesses.copy(),
extent=self.extent,
energy=energy)
if antialias_filter is None:
antialias_filter = AntialiasFilter()
for start, end, potential_slices in t.generate_slices(
max_batch=max_batch):
antialias_filter.bandlimit(potential_slices)
return t