def test_getDiscretisation(n, shape):
atoms, species = create_atoms(n, shape)
x = [np.linspace(0, 0.1, g) for g in (10, 21, 32)]
X = _toMesh(x)
f1 = get_discretisation(atoms,
species,
x,
pointwise=True,
FT=False,
**params)
FT1 = get_discretisation(atoms,
species,
x,
pointwise=True,
FT=True,
**params)
f2, FT2 = 0, 0
for i in range(len(n)):
a, b = get_atoms(n[i], True)
f2 = f2 + a(X - atoms[i].reshape(1, 1, -1))
FT2 = FT2 + b(X) * np.exp(-1j *
(X * atoms[i].reshape(1, 1, -1)).sum(-1))
# The errors here are from approximating exp
np.testing.assert_allclose(f1, f2, 1e-2)
np.testing.assert_allclose(FT1, FT2, 1e-2)
def test_CUDA(n, shape):
atoms, species = create_atoms(n, shape)
x = [np.linspace(0, 0.01, g) for g in (10, 21, 32)]
params["GPU"] = False
cpu_f = get_discretisation(atoms,
species,
x,
pointwise=True,
FT=False,
**params)
cpu_FT = get_discretisation(atoms,
species,
x,
pointwise=True,
FT=True,
**params)
params["GPU"] = True
gpu_f = get_discretisation(atoms,
species,
x,
pointwise=True,
FT=False,
**params)
gpu_FT = get_discretisation(atoms,
species,
x,
pointwise=True,
FT=True,
**params)
np.testing.assert_allclose(cpu_f, gpu_f, 1e-4)
np.testing.assert_allclose(cpu_FT, gpu_FT, 1e-4)
def test_pointwise(n, shape):
atoms, species = create_atoms(n, shape)
x = [np.linspace(0, 0.01, g) for g in (10, 21, 32)]
pw_f = get_discretisation(atoms,
species,
x,
pointwise=True,
FT=False,
**params)
pw_FT = get_discretisation(atoms,
species,
x,
pointwise=True,
FT=True,
**params)
av_f = get_discretisation(atoms,
species,
x,
pointwise=False,
FT=False,
**params)
av_FT = get_discretisation(atoms,
species,
x,
pointwise=True,
FT=True,
**params)
np.testing.assert_allclose(pw_f, av_f, 1e-2)
np.testing.assert_allclose(pw_FT, av_FT, 1e-2)
def test_getDiscretisation_str():
atoms, _ = create_atoms([14] * 3, [10] * 3)
get_discretisation(atoms,
"Si", [np.linspace(0, 1, g) for g in (10, 21, 31)],
pointwise=True,
FT=False,
**params)
def test_getDiscretisation_2d(n, shape):
atoms, species = create_atoms(n, shape)
x = [np.linspace(0, 0.01, g) for g in (10, 21, 31)]
f1 = get_discretisation(atoms,
species,
x,
pointwise=False,
FT=False,
**params).mean(-1)
FT1 = get_discretisation(atoms,
species,
x,
pointwise=False,
FT=True,
**params)[..., 0]
f2 = get_discretisation(atoms,
species,
x[:2],
pointwise=False,
FT=False,
**params).mean(-1)
FT2 = get_discretisation(atoms,
species,
x[:2],
pointwise=False,
FT=True,
**params)[..., 0]
for thing in (f1, FT1, f2, FT2):
thing /= abs(thing).max()
np.testing.assert_allclose(f1, f2, 1e-1)
np.testing.assert_allclose(FT1, FT2, 1e-1)
def test_getDiscretisation_bools(n, shape, grid):
atoms, species = create_atoms(n, shape)
x = [np.linspace(0, 0.1, g) for g in grid]
for b1 in (True, False):
for b2 in (True, False):
f = get_discretisation(atoms,
species,
x,
pointwise=b1,
FT=b2,
**params)
assert f.shape == grid
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))
