Getting the credits correct (hopefully)

diffsims/__init__.py

@@ -31,20 +31,6 @@
 
 from .sims.diffraction_simulation import DiffractionSimulation
 
-from .utils.atomic_diffraction_generator_support.probe_utils import (
-    ProbeFunction,
-    BesselProbe,
-)
-from .utils.atomic_diffraction_generator_support.fourier_transform import (
-    to_recip,
-    from_recip,
-    get_recip_points,
-    get_DFT,
-)
-from .utils.atomic_diffraction_generator_support.discretise_utils import (
-    get_discretisation,
-)
-
 from . import release_info
 
 __version__ = release_info.version

diffsims/generators/diffraction_generator.py

@@ -36,9 +36,7 @@
     get_vectorized_list_for_atomic_scattering_factors,
     is_lattice_hexagonal,
 )
-from diffsims.utils.atomic_diffraction_generator_support.fourier_transform import (
-    from_recip,
-)
+from diffsims.utils.fourier_transform import from_recip
 
 
 class DiffractionGenerator(object):
@@ -432,7 +430,7 @@ def calculate_ed_data(
         ]
 
         if mode == "kinematic":
-            from diffsims.utils import atomic_diffraction_generator_utils as simlib
+            from diffsims.sims import kinematic_simulation as simlib
         else:
             raise NotImplementedError(
                 "<mode> = %s is not currently supported" % repr(mode)

diffsims/generators/rotation_list_generators.py

@@ -44,7 +44,8 @@
     "triclinic": [180, 360, 0],
 }
 
-def get_list_from_orix(grid,rounding=2):
+
+def get_list_from_orix(grid, rounding=2):
     """
     Converts an orix sample to a rotation list
 
@@ -64,12 +65,13 @@ def get_list_from_orix(grid,rounding=2):
     rotation_list = z.data.tolist()
     i = 0
     while i < len(rotation_list):
-        rotation_list[i] = tuple(np.round(rotation_list[i],decimals=rounding))
+        rotation_list[i] = tuple(np.round(rotation_list[i], decimals=rounding))
         i += 1
 
     return rotation_list
 
-def get_fundamental_zone_grid(resolution=2, point_group=None,space_group=None):
+
+def get_fundamental_zone_grid(resolution=2, point_group=None, space_group=None):
     """
     Generates an equispaced grid of rotations within a fundamental zone.
 
@@ -88,10 +90,11 @@ def get_fundamental_zone_grid(resolution=2, point_group=None,space_group=None):
         Grid of rotations lying within the specified fundamental zone
     """
 
-    orix_grid = get_sample_fundamental(resolution=resolution,space_group=space_group)
-    rotation_list = get_list_from_orix(orix_grid,rounding=2)
+    orix_grid = get_sample_fundamental(resolution=resolution, space_group=space_group)
+    rotation_list = get_list_from_orix(orix_grid, rounding=2)
     return rotation_list
 
+
 def get_local_grid(resolution=2, center=None, grid_width=10):
     """
     Generates a grid of rotations about a given rotation
@@ -112,10 +115,13 @@ def get_local_grid(resolution=2, center=None, grid_width=10):
     -------
     rotation_list : list of tuples
     """
-    orix_grid =  get_sample_local(resolution=resolution, center=center, grid_width=grid_width)
-    rotation_list = get_list_from_orix(orix_grid,rounding=2)
+    orix_grid = get_sample_local(
+        resolution=resolution, center=center, grid_width=grid_width
+    )
+    rotation_list = get_list_from_orix(orix_grid, rounding=2)
     return rotation_list
 
+
 def get_grid_around_beam_direction(beam_rotation, resolution, angular_range=(0, 360)):
     """
     Creates a rotation list of rotations for which the rotation is about given beam direction

diffsims/libraries/structure_library.py

@@ -105,7 +105,9 @@ def from_crystal_systems(
         NotImplementedError:
             "This function has been removed in version 0.3.0, in favour of creation from orientation lists"
         """
-        raise NotImplementedError("This function has been removed in version 0.3.0, in favour of creation from orientation lists")
+        raise NotImplementedError(
+            "This function has been removed in version 0.3.0, in favour of creation from orientation lists"
+        )
 
     def get_library_size(self, to_print=False):
         """

diffsims/sims/diffraction_simulation.py

@@ -148,16 +148,17 @@ def get_diffraction_pattern(self, size=512, sigma=10):
         the order of 0.5nm and a the default size and sigma are used.
         """
         side_length = np.min(np.multiply((size / 2), self.calibration))
-        mask_for_sides = np.all((np.abs(self.coordinates[:, 0:2]) < side_length), axis=1)
-
+        mask_for_sides = np.all(
+            (np.abs(self.coordinates[:, 0:2]) < side_length), axis=1
+        )
 
         spot_coords = np.add(
             self.calibrated_coordinates[mask_for_sides], size / 2
         ).astype(int)
         spot_intens = self.intensities[mask_for_sides]
         pattern = np.zeros([size, size])
-        #checks that we have some spots
-        if spot_intens.shape[0]==0:
+        # checks that we have some spots
+        if spot_intens.shape[0] == 0:
             return pattern
         else:
             pattern[spot_coords[:, 0], spot_coords[:, 1]] = spot_intens

diffsims/sims/kinematic_simulation.py

@@ -0,0 +1,224 @@
+# -*- coding: utf-8 -*-
+# Copyright 2017-2019 The diffsims developers
+#
+# This file is part of diffsims.
+#
+# diffsims is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# diffsims is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with diffsims.  If not, see <http://www.gnu.org/licenses/>.
+
+"""
+Created on 1 Nov 2019
+
+Back end for computing diffraction patterns with a kinematic model.
+
+@author: Rob Tovey
+"""
+from diffsims.utils.discretise_utils import get_discretisation
+from numpy import array, pi, sin, cos, empty, maximum, sqrt
+from scipy.interpolate import interpn
+from diffsims.utils.fourier_transform import (
+    get_DFT,
+    to_recip,
+    fftshift_phase,
+    plan_fft,
+    fast_abs,
+)
+from diffsims.utils.generic_utils import to_mesh
+
+
+def normalise(arr):
+    return arr / arr.max()
+
+
+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, **kwargs).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 precess_mat(alpha, theta):
+    """
+    Generates rotation matrices for precession curves.
+
+    Parameters
+    ----------
+    alpha : `float`
+        Angle (in degrees) of precession tilt
+    theta : `float`
+        Angle (in degrees) along precession curve
+
+    Returns
+    -------
+    R : `numpy.ndarray` [`float`], (3, 3)
+        Rotation matrix associated to the tilt of `alpha` away from the vertical
+        axis and a rotation of `theta` about the vertical axis.
+    """
+    if alpha == 0:
+        return array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+    alpha, theta = alpha * pi / 180, theta * pi / 180
+    R_a = array([[1, 0, 0], [0, cos(alpha), -sin(alpha)], [0, sin(alpha), cos(alpha)]])
+    R_t = array([[cos(theta), -sin(theta), 0], [sin(theta), cos(theta), 0], [0, 0, 1]])
+    R = R_t.T.dot(R_a.dot(R_t))
+
+    return R
+
+
+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)