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 FT(self, y, out=None): """If `Y = sqrt(y[...,0]**2 + y[...,1]**2)*r` then returns an indicator function on the disc `Y < 1, y[2]==0`. Again, if `out` is provided then computation is inplace. If `y` is a list of arrays then it is converted into a mesh first. 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): """If `X = sqrt(x[...,0]**2+x[...,1]**2)/r` then returns `J_1(X)/X*scale`. If `out` is provided then this is computed inplace. If `scale` is not provided then it is assumed to be 1. If `x` is a list of arrays it is converted into a mesh first. 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 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 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 straight up # w = C - sqrt(maximum(0, C ** 2 - (y ** 2).sum(-1))) # if dx is None: # y[:, 2] = w.reshape(-1) # else: # y += w.reshape(y.shape[0], 1) * dx[2].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 __call__(self, x, out=None, scale=None): """Returns `func(x)*scale`. If `out` is provided then it is used as preallocated storage. If `scale` is not provided then it is assumed to be 1. If `x` is a list of arrays it is converted into a mesh first. 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 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))