def __call__(self, x, out=None, scale=None):
"""
Parameters
----------
x : `numpy.ndarray`, (nx, ny, nz, 3) or list of arrays of shape [(nx,), (ny,), (nz,)]
Mesh points at which to evaluate the probe density
out : `numpy.ndarray`, (nx, ny, nz), optional
If provided then computation is performed inplace
scale : `numpy.ndarray`, (nx, ny, nz), optional
If provided then the probe density is scaled by this before being
returned.
Returns
-------
out : `numpy.ndarray`, (nx, ny, nz)
An array equal to `probe(x)*scale`
"""
if self._func is None:
raise NotImplementedError
if not (hasattr(x, "shape")):
x = to_mesh(x)
if out is None:
out = self._func(x)
else:
out[...] = self._func(x)
if scale is not None:
out *= scale
return out
def test_toMesh(shape, dx, dtype):
x = [np.linspace(0, 1, s + 1)[1:] for s in shape]
if dx:
dx = list(np.eye(len(shape), len(shape))[::-1])
else:
dx = None
var1 = to_mesh(x, dx, dtype)
var2 = to_mesh(x, dx).astype(dtype)
if len(shape) == 2:
flag = dx is None
for i in range(shape[0]):
for j in range(shape[1]):
assert abs(var1[i, j, 0 if flag else 1] - x[0][i]) < 1e-16
assert abs(var1[i, j, 1 if flag else 0] - x[1][j]) < 1e-16
assert var1.dtype == var2.dtype
assert var1.shape == var2.shape
np.testing.assert_allclose(var1, var2, 1e-16, 1e-16)
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.
As a plotting utility, if a lower dimensional mesh is provided then
the remaining coordinates are assumed to be 0 and so only the
respective 1D/2D slice of the probe is returned.
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 `ny=0` or `nz=0` then
array is of smaller dimension.
"""
if not hasattr(y, "shape"):
y = to_mesh(y)
r = self._r
if y.shape[-1] == 1 or y.ndim == 1:
y = (y * r).reshape(-1)
y[abs(y) > 1] = 1
if out is None:
out = (2 * r) * sqrt(1 - y * y)
else:
out[...] = (2 * r) * sqrt(1 - y * y)
else:
if y.shape[-1] == 3:
dy2 = []
for i in range(y.ndim - 1):
tmp = tuple(0 if j != i else 1
for j in range(y.ndim - 1)) + (2, )
dy2.append(
abs(y[tmp] -
y[..., 2].item(0)) if y.shape[-1] == 3 else 1)
eps = max(1e-16, max(dy2) * 0.5)
if out is None:
out = empty(y.shape[:3], dtype=y.dtype)
_bessFT(y.reshape(-1, 3), 1 / r**2, 2 * pi * r**2, eps,
out.reshape(-1))
else:
if out is None:
out = (2 * pi * r**2) * (abs(y * y).sum(-1) <= 1 / r**2)
else:
out[...] = (2 * pi *
r**2) * (abs(y * y).sum(-1) <= 1 / r**2)
return out
def __call__(self, x, out=None, scale=None):
"""
Parameters
----------
x : `numpy.ndarray`, (nx, ny, nz, 3) or list of arrays of shape [(nx,), (ny,), (nz,)]
Mesh points at which to evaluate the probe density.
As a plotting utility, if a lower dimensional mesh is provided then
the remaining coordinates are assumed to be 0 and so only the
respective 1D/2D slice of the probe is returned.
out : `numpy.ndarray`, (nx, ny, nz), optional
If provided then computation is performed inplace
scale : `numpy.ndarray`, (nx, ny, nz), optional
If provided then the probe density is scaled by this before being
returned.
Returns
-------
out : `numpy.ndarray`, (nx, ny, nz)
An array equal to `probe(x)*scale`. If `ny=0` or `nz=0` then array is of
smaller dimension.
"""
if not hasattr(x, "shape"):
x = to_mesh(x)
scale = ones(1, dtype=x.dtype) if scale is None else scale
if out is None:
out = empty(x.shape[:-1], dtype=scale.dtype)
if x.shape[-1] == 1 or x.ndim == 1:
x = maximum(1e-16, abs(x)).reshape(-1)
out[...] = jv(1, x) / x * scale
elif x.shape[-1] == 2:
x = maximum(1e-16, sqrt(abs(x * x).sum(-1) / self._r**2))
out[...] = jv(1, x) / x * scale
else:
d = abs(x[1, 1, 0, :2] - x[0, 0, 0, :2])
h = d.min() / 10
s = ((d[0] * x.shape[0])**2 + (d[1] * x.shape[1])**2)**0.5
fine_grid = arange(h / 2, s / self._r + h, h)
j = jv(1, fine_grid) / fine_grid
_bess(
x.reshape(-1, 3),
1 / self._r,
1 / h,
j,
scale.reshape(-1),
out.reshape(-1),
)
return out
def grid2sphere(arr, x, dx, C):
"""
Projects 3d array onto a sphere
Parameters
----------
arr : `numpy.ndarray` [`float`], (nx, ny, nz)
Input function to be projected
x : `list` [`numpy.ndarray` [`float`]], of shapes [(nx,), (ny,), (nz,)]
Vectors defining mesh of <arr>
dx : `list` [`numpy.ndarray` [`float`]], of shapes [(3,), (3,), (3,)]
Basis in which to orient sphere. Centre of sphere will be at `C*dx[2]`
and mesh of output array will be defined by the first two vectors
C : `float`
Radius of sphere
Returns
-------
out : `numpy.ndarray` [`float`], (nx, ny)
If y is the point on the line between `i*dx[0]+j*dx[1]` and `C*dx[2]`
which also lies on the sphere of radius `C` from `C*dx[2]` then:
`out[i,j] = arr(y)`
Interpolation on arr is linear.
"""
if C in (None, 0) or x[2].size == 1:
if arr.ndim == 2:
return arr
elif arr.shape[2] == 1:
return arr[:, :, 0]
y = to_mesh((x[0], x[1], array([0])), dx).reshape(-1, 3)
if C is not None: # project on line to centre
w = 1 / (1 + (y**2).sum(-1) / C**2)
y *= w[:, None]
if dx is None:
y[:, 2] = C * (1 - w)
else:
y += C * (1 - w)[:, None] * dx[2]
out = interpn(x, arr, y, method='linear', bounds_error=False, fill_value=0)
return out.reshape(x[0].size, x[1].size)
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))
