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()