Add option to have spot centers off pixel centers

diffsims/pattern/detector_functions.py

@@ -16,7 +16,10 @@
 # You should have received a copy of the GNU General Public License
 # along with diffsims.  If not, see <http://www.gnu.org/licenses/>.
 
+from typing import Tuple
+
 import numpy as np
+from numba import jit
 
 from numpy.random import default_rng
 from scipy import ndimage as ndi
@@ -31,6 +34,7 @@
     "add_shot_noise",
     "add_shot_and_point_spread",
     "constrain_to_dynamic_range",
+    "get_pattern_from_pixel_coordinates_and_intensities",
 ]
 
 
@@ -243,3 +247,114 @@ def add_detector_offset(pattern, offset):
     """
     pattern = np.add(pattern, offset)
     return constrain_to_dynamic_range(pattern)
+
+
+def get_pattern_from_pixel_coordinates_and_intensities(
+    coordinates: np.ndarray,
+    intensities: np.ndarray,
+    shape: Tuple[int, int],
+    sigma: float,
+    clip_threshold: float = 1,
+) -> np.ndarray:
+    """Generate a diffraction pattern from spot pixel-coordinates and intensities,
+    using a gaussian blur.
+    This is subpixel-precise, meaning the coordinates can be floats.
+    Values less than `clip_threshold` are rounded down to 0 to simplify computation.
+
+    Parameters
+    ----------
+    coordinates : np.ndarray
+        Coordinates of reflections, in pixels. Shape (n, 2) or (n, 3). Can be floats
+    intensities : np.ndarray
+        Intensities of each reflection. Must have same same first dimension as `coordinates`
+    shape : tuple[int, int]
+        Output shape
+    sigma : float
+        For Gaussian blur
+    intensity_scale : float
+        Scale to multiply the final diffraction pattern with
+
+    Returns
+    -------
+    np.ndarray
+        dtype int
+
+    Notes
+    -----
+    Not all values below the clipping threshold are ignored.
+    The threshold is used to estimate a radius (box) around each reflection where the pixel intensity is greater than the threshold.
+    As the radius is rounded up and as the box is square rather than circular, some values below the threshold can be included.
+
+    When using float coordinates, the intensity is spread as if the edge was not there.
+    This is in line with what should be expected from a beam on the edge of the detector, as part of the beam is simply outside the detector area.
+    However, when using integer coordinates, the total intensity is preserved for the pixels in the pattern.
+    This means that the intensity contribution from parts of the beam which would hit outside the detector are now kept in the pattern.
+    Thus, reflections wich are partially outside the detector will have higher intensities than expected, when using integer coordinates.
+    """
+    if np.issubdtype(coordinates.dtype, np.integer):
+        # Much simpler with integer coordinates
+        coordinates = coordinates.astype(int)
+        out = np.zeros(shape)
+        # coordinates are xy(z), out array indices are yx.
+        out[coordinates[:, 1], coordinates[:, 0]] = intensities
+        out = add_shot_and_point_spread(out, sigma, shot_noise=False)
+        return out
+
+    # coordinates of each pixel in the output, such that the final axis is yx coordinates
+    inds = np.transpose(np.indices(shape), (1, 2, 0))
+    return _subpixel_gaussian(
+        coordinates,
+        intensities,
+        inds,
+        shape,
+        sigma,
+        clip_threshold,
+    )
+
+
+@jit(
+    nopython=True
+)  # Not parallel, we might get a race condition with overlapping spots
+def _subpixel_gaussian(
+    coordinates: np.ndarray,
+    intensities: np.ndarray,
+    inds: np.ndarray,
+    shape: Tuple[int, int],
+    sigma: float,
+    clip_threshold: float = 1,
+) -> np.ndarray:
+    out = np.zeros(shape)
+
+    # Pre-calculate the constants
+    prefactor = 1 / (2 * np.pi * sigma**2)
+    exp_prefactor = -1 / (2 * sigma**2)
+
+    for i in range(intensities.size):
+        # Reverse since coords are xy, but indices are yx
+        coord = coordinates[i][:2][::-1]
+        intens = intensities[i]
+
+        # The gaussian is expensive to evaluate for all pixels and spots.
+        # Therefore, we limit the calculations to a box around each reflection where the intensity is above a threshold.
+        # Formula found by inverting the gaussian
+        radius = np.sqrt(np.log(clip_threshold / (prefactor * intens)) / exp_prefactor)
+
+        if np.isnan(radius):
+            continue
+        slic = (
+            slice(
+                max(0, int(np.ceil(coord[0] - radius))),
+                min(shape[0], int(np.floor(coord[0] + radius + 1))),
+            ),
+            slice(
+                max(0, int(np.ceil(coord[1] - radius))),
+                min(shape[1], int(np.floor(coord[1] + radius + 1))),
+            ),
+        )
+        # Calculate the values of the Gaussian manually
+        out[slic] += (
+            intens
+            * prefactor
+            * np.exp(exp_prefactor * np.sum((inds[slic] - coord) ** 2, axis=-1))
+        )
+    return out

diffsims/simulations/simulation2d.py

@@ -16,7 +16,7 @@
 # You should have received a copy of the GNU General Public License
 # along with diffsims.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Union, Sequence, TYPE_CHECKING, Any
+from typing import Union, Sequence, Tuple, TYPE_CHECKING, Any
 import copy
 
 import numpy as np
@@ -27,7 +27,9 @@
 from orix.vector import Vector3d
 
 from diffsims.crystallography._diffracting_vector import DiffractingVector
-from diffsims.pattern.detector_functions import add_shot_and_point_spread
+from diffsims.pattern.detector_functions import (
+    get_pattern_from_pixel_coordinates_and_intensities,
+)
 
 # to avoid circular imports
 if TYPE_CHECKING:  # pragma: no cover
@@ -343,12 +345,15 @@ def polar_flatten_simulations(self, radial_axes=None, azimuthal_axes=None):
 
     def get_diffraction_pattern(
         self,
-        shape=None,
-        sigma=10,
-        direct_beam_position=None,
-        in_plane_angle=0,
-        calibration=0.01,
-        mirrored=False,
+        shape: Tuple[int, int] = None,
+        sigma: float = 10,
+        direct_beam_position: Tuple[int, int] = None,
+        in_plane_angle: float = 0,
+        calibration: float = 0.01,
+        mirrored: bool = False,
+        fast: bool = True,
+        normalize: bool = True,
+        fast_clip_threshold: float = 1,
     ):
         """Returns the diffraction data as a numpy array with
         two-dimensional Gaussians representing each diffracted peak. Should only
@@ -368,11 +373,19 @@ def get_diffraction_pattern(
         mirrored: bool, optional
             Whether the pattern should be flipped over the x-axis,
             corresponding to the inverted orientation
-
+        fast: bool, optional
+            Whether to speed up calculations by rounding spot coordinates down to integer pixel
+        normalize: bool, optional
+            Whether to normalize the pattern to values between 0 and 1
+        fast_clip_threshold: float, optional
+            Only used when `fast` is False.
+            Pixel intensity threshold, such that pixels which would be below this value are ignored.
+            Thresholding performed before possible normalization.
+            See diffsims.pattern.detector_functions.get_pattern_from_pixel_coordinates_and_intensities for details.
         Returns
         -------
         diffraction-pattern : numpy.array
-            The simulated electron diffraction pattern, normalised.
+            The simulated electron diffraction pattern, normalized by default.
 
         Notes
         -----
@@ -381,7 +394,13 @@ def get_diffraction_pattern(
         the order of 0.5nm and the default size and sigma are used.
         """
         if direct_beam_position is None:
-            direct_beam_position = (shape[1] // 2, shape[0] // 2)
+            # Use subpixel-precise center if possible
+            if fast or np.issubdtype(
+                self.get_current_coordinates().data.dtype, np.integer
+            ):
+                direct_beam_position = (shape[1] // 2, shape[0] // 2)
+            else:
+                direct_beam_position = ((shape[1] - 1) / 2, (shape[0] - 1) / 2)
         transformed = self._get_transformed_coordinates(
             in_plane_angle,
             direct_beam_position,
@@ -395,17 +414,21 @@ def get_diffraction_pattern(
             & (transformed.data[:, 1] >= 0)
             & (transformed.data[:, 1] < shape[0])
         )
-        spot_coords = transformed.data[in_frame].astype(int)
-
+        spot_coords = transformed.data[in_frame]
+        if fast:
+            spot_coords = spot_coords.astype(int)
         spot_intens = transformed.intensity[in_frame]
-        pattern = np.zeros(shape)
+
         # checks that we have some spots
         if spot_intens.shape[0] == 0:
-            return pattern
-        else:
-            pattern[spot_coords[:, 0], spot_coords[:, 1]] = spot_intens
-            pattern = add_shot_and_point_spread(pattern.T, sigma, shot_noise=False)
-        return np.divide(pattern, np.max(pattern))
+            return np.zeros(shape)
+
+        pattern = get_pattern_from_pixel_coordinates_and_intensities(
+            spot_coords, spot_intens, shape, sigma, fast_clip_threshold
+        )
+        if normalize:
+            pattern = np.divide(pattern, np.max(pattern))
+        return pattern
 
     @property
     def num_phases(self):