def test_freq2(shape):
x = [np.linspace(0, 1, s) if s > 1 else np.array([0]) for s in shape]
y = to_recip(from_recip(x))
z = from_recip(to_recip(x))
for i in range(len(shape)):
assert abs(x[i] - y[i] + y[i].min()).max() < 1e-6
assert abs(x[i] - z[i] + z[i].min()).max() < 1e-6
def test_freq(shape, dX, rX, dY, rY):
x, y = get_recip_points(len(shape), shape, dX, rX, dY, rY)
X, Y = from_recip(y), to_recip(x)
assert len(x) == len(shape)
assert len(y) == len(shape)
assert len(X) == len(shape)
assert len(Y) == len(shape)
for i in range(len(shape)):
assert y[i].size == x[i].size
if shape[i] is None:
if dX[i] is not None:
assert abs(x[i].item(1) - x[i].item(0)) <= dX[i] + 1e-8
if rY[i] is not None:
assert y[i].ptp() >= rY[i] - 1e-8
if rX[i] is not None:
assert x[i].ptp() >= rX[i] - 1e-8
if dY[i] is not None:
assert abs(y[i].item(1) - y[i].item(0)) <= dY[i] + 1e-8
np.testing.assert_allclose(x[i], X[i], 1e-5)
np.testing.assert_allclose(y[i], Y[i], 1e-5)
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))