def test_DFT(rX, rY):
x, y = get_recip_points(len(rX), rX=rX, rY=rY)
axes = (0 if len(x) == 1 else None)
f, g = get_atoms(0, returnFunc=True)
f, g = f(_toMesh(x)), g(_toMesh(y))
ft, ift = get_DFT(x, y)
ft1, ift1 = get_DFT(X=x)
ft2, ift2 = get_DFT(Y=y)
for FT in (ft, ft1, ft2):
np.testing.assert_allclose(g, FT(f, axes=axes), 1e-5, 1e-5)
for IFT in (ift, ift1, ift2):
np.testing.assert_allclose(f, IFT(g, axes=axes), 1e-5, 1e-5)
def FT(self, y, out=None):
"""
Parameters
----------
y : `numpy.ndarray`, (nx, ny, nz, 3) or list of arrays of shape [(nx,), (ny,), (nz,)]
Mesh of Fourier coordinates at which to evaluate the probe density
out : `numpy.ndarray`, (nx, ny, nz), optional
If provided then computation is performed inplace
Returns
-------
out : `numpy.ndarray`, (nx, ny, nz)
An array equal to `FourierTransform(probe)(y)`
"""
if hasattr(y, "shape"):
y_start = y[(0, ) * (y.ndim - 1)]
y_end = y[(-1, ) * (y.ndim - 1)]
y = [
linspace(y_start[i], y_end[i], y.shape[i], endpoint=True)
for i in range(3)
]
x = from_recip(y)
ft = get_DFT(x, y)[0]
tmp = ft(self(x, out=out))
if out is None:
out = tmp
else:
out[...] = tmp
return out
def FT(self, y, out=None):
"""Returns the Fourier transform of func on the mesh `y`. Again,
if `out` is provided then computation is `inplace`. If `y` is a
list of arrays then it is converted into a mesh first. If this
function is not overridden then an approximation is made using
`func` and the `fft`.
Parameters
----------
y : numpy.ndarray, (nx, ny, nz, 3) or list of arrays of shape
[(nx,), (ny,), (nz,)]
Mesh of Fourier coordinates at which to evaluate the probe
density.
out : numpy.ndarray, (nx, ny, nz), optional
If provided then computation is performed inplace.
Returns
-------
out : numpy.ndarray, (nx, ny, nz)
An array equal to `FourierTransform(probe)(y)`.
"""
if hasattr(y, "shape"):
y_start = y[(0, ) * (y.ndim - 1)]
y_end = y[(-1, ) * (y.ndim - 1)]
y = [
linspace(y_start[i], y_end[i], y.shape[i], endpoint=True)
for i in range(3)
]
x = from_recip(y)
ft = get_DFT(x, y)[0]
tmp = ft(self(x, out=out))
if out is None:
out = tmp
else:
out[...] = tmp
return out
def get_diffraction_image(coordinates,
species,
probe,
x,
wavelength,
precession,
GPU=True,
pointwise=False,
**kwargs):
"""
Return kinematically simulated diffraction pattern
Parameters
----------
coordinates : `numpy.ndarray` [`float`], (n_atoms, 3)
List of atomic coordinates
species : `numpy.ndarray` [`int`], (n_atoms,)
List of atomic numbers
probe : `diffsims.ProbeFunction`
Function representing 3D shape of beam
x : `list` [`numpy.ndarray` [`float`] ], of shapes [(nx,), (ny,), (nz,)]
Mesh on which to compute the volume density
wavelength : `float`
Wavelength of electron beam
precession : a pair (`float`, `int`)
The float dictates the angle of precession and the int how many points are
used to discretise the integration.
dtype : (`str`, `str`)
tuple of floating/complex datatypes to cast outputs to
ZERO : `float` > 0, optional
Rounding error permitted in computation of atomic density. This value is
the smallest value rounded to 0.
GPU : `bool`, optional
Flag whether to use GPU or CPU discretisation. Default (if available) is True
pointwise : `bool`, optional
Optional parameter whether atomic intensities are computed point-wise at
the centre of a voxel or an integral over the voxel. default=False
Returns
-------
DP : `numpy.ndarray` [`dtype[0]`], (nx, ny, nz)
The two-dimensional diffraction pattern evaluated on the reciprocal grid
corresponding to the first two vectors of `x`.
"""
FTYPE = kwargs["dtype"][0]
kwargs["GPU"] = GPU
kwargs["pointwise"] = pointwise
x = [X.astype(FTYPE, copy=False) for X in x]
y = to_recip(x)
if wavelength == 0:
p = probe(x).mean(-1)
vol = get_discretisation(coordinates, species, x[:2], **kwargs)[..., 0]
ft = get_DFT(x[:-1], y[:-1])[0]
else:
p = probe(x)
vol = get_discretisation(coordinates, species, x, **kwargs)
ft = get_DFT(x, y)[0]
if precession[0] == 0:
arr = ft(vol * p)
arr = fast_abs(arr, arr).real**2
if wavelength == 0:
return normalise(arr)
else:
return normalise(grid2sphere(arr, y, None, 2 * pi / wavelength))
R = [
precess_mat(precession[0], i * 360 / precession[1])
for i in range(precession[1])
]
if wavelength == 0:
return normalise(
sum(
get_diffraction_image(coordinates.dot(r), species, probe, x,
wavelength, (0, 1), **kwargs)
for r in R))
fftshift_phase(vol) # removes need for fftshift after fft
buf = empty(vol.shape, dtype=FTYPE)
ft, buf = plan_fft(buf, overwrite=True, planner=1)
DP = None
for r in R:
probe(to_mesh(x, r.T, dtype=FTYPE), out=buf,
scale=vol) # buf = bess*vol
# Do convolution
newFT = ft()
newFT = fast_abs(newFT, buf).real
newFT *= newFT # newFT = abs(newFT) ** 2
newFT = grid2sphere(newFT.real, y, list(r), 2 * pi / wavelength)
if DP is None:
DP = newFT
else:
DP += newFT
return normalise(DP.astype(FTYPE, copy=False))
def test_fail_DFT(): get_DFT()
@pytest.mark.parametrize('shape1, shape2, n1, n2, dx', [
def test_fail_DFT():
get_DFT()
