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