Merged #98 into suitable sub-branch

diffsims/generators/diffraction_generator.py

@@ -287,3 +287,129 @@ def calculate_profile_data(
         y = y / max(y) * 100
 
         return ProfileSimulation(x, y, hkls)
+
+
+class AtomicDiffractionGenerator:
+    '''
+    Computes electron diffraction patterns for an atomic lattice.
+
+    Parameters
+    ----------
+    accelerating_voltage : float, 'inf'
+        The accelerating voltage of the microscope in kV
+    detector : list of 1D float-type arrays
+        List of mesh vectors defining the (flat) detector size and sensor positions
+    reciprocal_mesh : bool, optional
+        If True then `detector` is assumed to be a reciprocal grid, else
+        (default) it is assumed to be a real grid.
+
+    '''
+
+    def __init__(self, accelerating_voltage, detector,
+                 reciprocal_mesh=False, debye_waller_factors=None):
+        self.wavelength = get_electron_wavelength(accelerating_voltage)
+        # Always store a 'real' mesh
+        self.detector = detector if not reciprocal_mesh else from_recip(detector)
+
+        if debye_waller_factors:
+            raise NotImplementedError('Not implemented for this simulator')
+        self.debye_waller_factors = debye_waller_factors or {}
+
+    def calculate_ed_data(self, structure, probe, slice_thickness,
+                          probe_centre=None, z_range=200, precessed=False, dtype='float64',
+                          ZERO=1e-14, mode='kinematic', **kwargs):
+        '''
+        Calculates single electron diffraction image for particular atomic
+        structure and probe.
+
+        Parameters
+        ----------
+        coordinates : ndarray of floats, shape [n_atoms, 3]
+            List of atomic coordinates, i.e. atom i is centred at <coordinates>[i]
+        species : ndarray of integers, shape [n_atoms]
+            List of atomic numbers, i.e. atom i has atomic number <species>[i]
+        probe : instance of probeFunction
+            Function representing 3D shape of beam
+        slice_thickness : float
+            Discretisation thickness in the z-axis
+        probe_centre : ndarray (or iterable), shape [3] or [2]
+            Translation vector for the probe. Either of the same dimension of the
+            space or the dimension of the detector. default=None focusses the
+            probe at [0,0,0]
+        zrange : float
+            z-thickness to discretise. Only required if sample is not thick enough to
+            fully resolve the Ewald-sphere. Default value is 200.
+        precessed : bool, float, or (float, int)
+            Dictates whether beam precession is simulated. If False or the float is
+            0 then no precession is computed. If <precessed> = (alpha, n) then the
+            precession arc of tilt alpha (in degrees) is discretised into n
+            projections. If n is not provided then default of 30 is used.
+        dtype : str or numpy.dtype
+            Defines the precision to use whilst computing diffraction image.
+        ZERO : float > 0
+            Rounding error permitted in computation of atomic density. This value is
+            the smallest value rounded to 0. Default is 1e-14.
+        mode : str
+            Only <mode>='kinematic' is currently supported.
+        kwargs : dictionary
+            Extra key-word arguments to pass to child simulator
+            For kinematic:
+                GPU : bool
+                    Flag to use GPU if available, default is True
+                pointwise: bool
+                    Flag to evaluate charge pointwise on voxels rather than average,
+                    default is False
+
+
+        Returns
+        -------
+        ndarray
+            Diffraction data to be interpreted as a discretisation on the original
+            detector mesh.
+
+        '''
+
+        species = structure.element
+        coordinates = structure.xyz_cartn.reshape(species.size, -1)
+        dim = coordinates.shape[1]
+        assert dim == 3
+
+        if probe_centre is None:
+            probe_centre = np.zeros(dim)
+        elif len(probe_centre) < dim:
+            probe_centre = np.array(list(probe_centre) + [0])
+        probe_centre = np.array(probe_centre)
+        coordinates = coordinates - probe_centre[None]
+
+        if not precessed:
+            precessed = (float(0), 1)
+        elif np.isscalar(precessed):
+            precessed = (float(precessed), 30)
+
+        dtype = np.dtype(dtype)
+        dtype = round(dtype.itemsize / (1 if dtype.kind == 'f' else 2))
+        dtype = 'f' + str(dtype), 'c' + str(2 * dtype)
+
+        assert ZERO > 0
+
+        # Filter list of atoms
+        for d in range(dim - 1):
+            ind = coordinates[:, d] >= self.detector[d].min() - 20
+            coordinates, species = coordinates[ind, :], species[ind]
+            ind = coordinates[:, d] <= self.detector[d].max() + 20
+            coordinates, species = coordinates[ind, :], species[ind]
+
+        # Add z-coordinate
+        z_range = max(z_range, coordinates[:, -1].ptp())  # enforce minimal resolution in reciprocal space
+        x = [self.detector[0], self.detector[1],
+             np.arange(coordinates[:, -1].min() - 20, coordinates[:, -1].min() + z_range + 20, slice_thickness)]
+
+        if mode == 'kinematic':
+            from diffsims.sims import kinematic_simulation as simlib
+        else:
+            raise NotImplementedError('<mode> = %s is not currently supported' % repr(mode))
+
+        kwargs['dtype'] = dtype
+        kwargs['ZERO'] = ZERO
+        return simlib.get_diffraction_image(coordinates, species, probe, x,
+                                            self.wavelength, precessed, **kwargs)

diffsims/generators/library_generator.py

@@ -136,6 +136,74 @@ def get_diffraction_library(
 
         return diffraction_library
 
+
+def _generate_lookup_table(recip_latt,
+                           reciprocal_radius: float,
+                           unique: bool = True):
+    """Generate a look-up table with all combinations of indices,
+    including their reciprocal distances and the angle between
+    them.
+
+    Parameters
+    ----------
+    recip_latt : :class:`diffpy.structure.lattice.Lattice`
+        Reciprocal lattice
+    reciprocal_radius : float
+        The maximum g-vector magnitude to be included in the library.
+    unique : bool
+        Return a unique list of phase measurements
+
+    Returns
+    -------
+    indices : np.array
+        Nx2x3 numpy array containing the miller indices for
+        reflection1, reflection2
+    measurements : np.array
+        Nx3 numpy array containing len1, len2, angle
+
+    """
+    miller_indices, coordinates, distances = get_points_in_sphere(
+        recip_latt,
+        reciprocal_radius)
+
+    # Create pair_indices for selecting all point pair combinations
+    num_indices = len(miller_indices)
+    pair_a_indices, pair_b_indices = np.mgrid[:num_indices, :num_indices]
+
+    # Only select one of the permutations and don't pair an index with
+    # itself (select above diagonal)
+    upper_indices = np.triu_indices(num_indices, 1)
+    pair_a_indices = pair_a_indices[upper_indices].ravel()
+    pair_b_indices = pair_b_indices[upper_indices].ravel()
+
+    # Mask off origin (0, 0, 0)
+    origin_index = num_indices // 2
+    pair_a_indices = pair_a_indices[pair_a_indices != origin_index]
+    pair_b_indices = pair_b_indices[pair_b_indices != origin_index]
+
+    pair_indices = np.vstack([pair_a_indices, pair_b_indices])
+
+    # Create library entries
+    angles = get_angle_cartesian_vec(coordinates[pair_a_indices], coordinates[pair_b_indices])
+    pair_distances = distances[pair_indices.T]
+    # Ensure longest vector is first
+    len_sort = np.fliplr(pair_distances.argsort(axis=1))
+    # phase_index_pairs is a list of [hkl1, hkl2]
+    phase_index_pairs = np.take_along_axis(miller_indices[pair_indices.T], len_sort[:, :, np.newaxis], axis=1)
+    # phase_measurements is a list of [len1, len2, angle]
+    phase_measurements = np.column_stack((np.take_along_axis(pair_distances, len_sort, axis=1), angles))
+
+    if unique:
+        # Only keep unique triplets
+        measurements, measurement_indices = np.unique(phase_measurements, axis=0, return_index=True)
+        indices = phase_index_pairs[measurement_indices]
+    else:
+        measurements = phase_measurements
+        indices = phase_index_pairs
+
+    return measurements, indices
+
+
 class VectorLibraryGenerator:
     """Computes a library of diffraction vectors and pairwise inter-vector
     angles for a specified StructureLibrary.
@@ -177,13 +245,13 @@ def get_vector_library(self, reciprocal_radius):
             # Get reciprocal lattice points within reciprocal_radius
             recip_latt = structure.lattice.reciprocal()
 
-            measurements, indices = _generate_lookup_table(
-                recip_latt=recip_latt, reciprocal_radius=reciprocal_radius, unique=True
-            )
+            measurements, indices = _generate_lookup_table(recip_latt=recip_latt,
+                                                           reciprocal_radius=reciprocal_radius,
+                                                           unique=True)
 
             vector_library[phase_name] = {
-                "indices": indices,
-                "measurements": measurements,
+                'indices': indices,
+                'measurements': measurements
             }
 
         # Pass attributes to diffraction library from structure library.