diff --git a/diffsims/pattern/detector_functions.py b/diffsims/pattern/detector_functions.py
index 58d5342a..133cd11a 100644
--- a/diffsims/pattern/detector_functions.py
+++ b/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 .
+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
diff --git a/diffsims/simulations/simulation2d.py b/diffsims/simulations/simulation2d.py
index 1ca57d1c..a7ca9cb8 100644
--- a/diffsims/simulations/simulation2d.py
+++ b/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 .
-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):
diff --git a/diffsims/tests/patterns/test_detector_functions.py b/diffsims/tests/patterns/test_detector_functions.py
index bed8e858..694421eb 100644
--- a/diffsims/tests/patterns/test_detector_functions.py
+++ b/diffsims/tests/patterns/test_detector_functions.py
@@ -26,6 +26,7 @@
add_dead_pixels,
add_linear_detector_gain,
add_detector_offset,
+ get_pattern_from_pixel_coordinates_and_intensities,
)
@@ -154,3 +155,152 @@ def test_offset_array(self, pattern):
g1 = add_detector_offset(pattern, pattern)
g2 = add_detector_offset(g1, -pattern)
assert np.allclose(pattern, g2)
+
+
+class TestGetPatternFromPixelCoordinatesAndIntensities:
+ def test_2d_vs_3d_coordinates(self):
+ c2 = np.asarray(
+ [
+ [10, 10],
+ [20, 30],
+ [15, 20],
+ ]
+ )
+ c3 = np.asarray(
+ [
+ [c2[0, 0], c2[0, 1], 92],
+ [c2[1, 0], c2[1, 1], 0],
+ [c2[2, 0], c2[2, 1], -192],
+ ]
+ )
+ shape = (50, 50)
+ intensities = np.ones(3) * 100
+ sigma = 1
+ p2 = get_pattern_from_pixel_coordinates_and_intensities(
+ c2,
+ intensities,
+ shape,
+ sigma,
+ )
+ p3 = get_pattern_from_pixel_coordinates_and_intensities(
+ c3,
+ intensities,
+ shape,
+ sigma,
+ )
+ assert np.array_equal(p2, p3)
+
+ def test_integer_vs_float_coordinates(self):
+ # With cutoff at 1, rounding the patterns to int should yield the same result
+ ci = np.asarray(
+ [
+ [10, 10],
+ [20, 30],
+ [15, 20],
+ ]
+ ).astype(int)
+ cf = ci.astype(float)
+ intensities = np.ones(3) * 100
+ shape = (50, 50)
+ sigma = 1
+ pi = get_pattern_from_pixel_coordinates_and_intensities(
+ ci,
+ intensities,
+ shape,
+ sigma,
+ ).astype(int)
+ pf = get_pattern_from_pixel_coordinates_and_intensities(
+ cf,
+ intensities,
+ shape,
+ sigma,
+ ).astype(int)
+ assert np.allclose(pi, pf)
+
+ def test_low_intensity(self):
+ # No pixel intensities above the cutoff
+ ci = np.asarray(
+ [
+ [10, 10],
+ [20, 30],
+ [15, 20],
+ ]
+ ).astype(float)
+ p = get_pattern_from_pixel_coordinates_and_intensities(
+ ci,
+ np.ones(3),
+ (50, 50),
+ 1,
+ )
+ assert np.sum(p) == 0.0
+
+ def test_total_intensity_preservation(self):
+ # total intensity always preserved for integer coords
+ ci = np.asarray(
+ [
+ [10, 10],
+ ]
+ ).astype(int)
+ i = 100
+ p = get_pattern_from_pixel_coordinates_and_intensities(
+ ci,
+ np.array([i]),
+ (50, 50),
+ 3,
+ )
+ assert np.allclose(np.sum(p), i)
+
+ def test_total_intensity_preservation_with_threshold(self):
+ # With large intensity, total is preserved for floats as well
+ # (since the values below the cutoff don't contribute)
+
+ # Large intensity and default cutoff
+ ci = np.asarray(
+ [
+ [10, 10],
+ ]
+ ).astype(float)
+ i = 1000
+ p = get_pattern_from_pixel_coordinates_and_intensities(
+ ci,
+ np.array([i]),
+ (50, 50),
+ 1,
+ )
+ assert np.sum(p) / i > 0.999 # 0.1% error
+ # Smaller intensity and cutoff
+ i = 20
+ p = get_pattern_from_pixel_coordinates_and_intensities(
+ ci,
+ np.array([i]),
+ (50, 50),
+ 1,
+ clip_threshold=0.01,
+ )
+ assert np.sum(p) / i > 0.999 # 0.1% error
+
+ def test_spot_in_corner(self):
+ ci = np.asarray(
+ [
+ [0.3, 0.1],
+ ]
+ )
+ i = 100
+ sigma = 3
+ p = get_pattern_from_pixel_coordinates_and_intensities(
+ ci,
+ np.array([i]),
+ (50, 50),
+ sigma,
+ )
+ # Intensity spread out of the pattern is lost
+ assert np.sum(p) < i
+
+ p = get_pattern_from_pixel_coordinates_and_intensities(
+ ci.astype(int),
+ np.array([i]),
+ (50, 50),
+ sigma,
+ )
+ # Integer coordinates, total intensity is preserved even at edge
+ assert np.allclose(np.sum(p), i)
diff --git a/diffsims/tests/simulations/test_simulations2d.py b/diffsims/tests/simulations/test_simulations2d.py
index 5fd1769c..a4afc750 100644
--- a/diffsims/tests/simulations/test_simulations2d.py
+++ b/diffsims/tests/simulations/test_simulations2d.py
@@ -113,6 +113,55 @@ def test_deepcopy(self, single_simulation):
copied = single_simulation.deepcopy()
assert copied is not single_simulation
+ def test_get_diffraction_pattern_fast_vs_slow(self, al_phase):
+ gen = SimulationGenerator()
+ rot = Rotation.identity()
+ xyz = np.asarray([[0, 0, 0]])
+ i = np.array([30])
+ integer_coords = DiffractingVector(
+ phase=al_phase, xyz=xyz.astype(int), intensity=i
+ )
+ float_coords = DiffractingVector(
+ phase=al_phase, xyz=xyz.astype(float), intensity=i
+ )
+ int_sim = Simulation2D(
+ phases=al_phase,
+ simulation_generator=gen,
+ coordinates=integer_coords,
+ rotations=rot,
+ )
+ float_sim = Simulation2D(
+ phases=al_phase,
+ simulation_generator=gen,
+ coordinates=float_coords,
+ rotations=rot,
+ )
+ sim_kwargs = dict(
+ shape=(11, 11), # This shape has a center directly at a pixel
+ sigma=1,
+ calibration=1,
+ normalize=False,
+ fast_clip_threshold=0.01,
+ )
+ # Check that fast/slow are the same when coordinates are dtype int
+ fast = int_sim.get_diffraction_pattern(fast=True, **sim_kwargs)
+ slow = int_sim.get_diffraction_pattern(fast=False, **sim_kwargs)
+ # Should be exactly equal
+ assert np.array_equal(fast, slow)
+
+ # Check that float coords are the same too, when the center is in a pixel
+ float_fast = float_sim.get_diffraction_pattern(fast=True, **sim_kwargs)
+ float_slow = float_sim.get_diffraction_pattern(fast=False, **sim_kwargs)
+ assert np.allclose(float_fast, fast)
+ assert np.all((float_slow - float_fast) < sim_kwargs["fast_clip_threshold"])
+
+ # Change the size, now the center is between pixels.
+ # Slow should reflect this
+ sim_kwargs["shape"] = (10, 10)
+ float_fast = float_sim.get_diffraction_pattern(fast=True, **sim_kwargs)
+ float_slow = float_sim.get_diffraction_pattern(fast=False, **sim_kwargs)
+ assert not np.allclose(float_fast, float_slow)
+
class TestSimulationInitFailures:
def test_different_size(self, al_phase):