Fix DiffractionPattern k values/vectors

freud/diffraction.pyx

@@ -46,6 +46,15 @@ cdef class DiffractionPattern(_Compute):
     as a multiplication in Fourier space. The computed diffraction pattern
     can be accessed as a square array of shape ``(output_size, output_size)``.
 
+    The :math:`\vec{k}=0` peak is always located at index
+    ``(output_size // 2, output_size // 2)`` and is normalized to have a value
+    of :math:`S(\vec{k}=0) = 1` (not :math:`N`, a common convention). The
+    remaining :math:`\vec{k}`-vectors are computed such that each peak in the
+    diffraction pattern satisfies the relationship :math:`\vec{k} \cdot
+    \vec{R} = 2\piN` for some integer :math: `N` and lattice vector of
+    the system :math: `R` (see
+    `this article <https://en.wikipedia.org/wiki/Reciprocal_lattice>`).
+
     This method is based on the implementations in the open-source
     `GIXStapose application <https://github.com/cmelab/GIXStapose>`_ and its
     predecessor, diffractometer :cite:`Jankowski2017`.
@@ -67,6 +76,7 @@ cdef class DiffractionPattern(_Compute):
     cdef unsigned int _frame_counter
     cdef double _box_matrix_scale_factor
     cdef double[:] _view_orientation
+    cdef double _k_scale_factor
     cdef cbool _k_values_cached
     cdef cbool _k_vectors_cached
 
@@ -130,7 +140,7 @@ cdef class DiffractionPattern(_Compute):
         """Zoom, shear, and scale diffraction intensities.
 
         Args:
-            img ((``grid_size//zoom, grid_size//zoom``) :class:`numpy.ndarray`):
+            img ((``grid_size, grid_size``) :class:`numpy.ndarray`):
                 Array of diffraction intensities.
             box (:class:`~.box.Box`):
                 Simulation box.
@@ -147,7 +157,8 @@ cdef class DiffractionPattern(_Compute):
         # The adjustments to roll and roll_shift ensure that the peak
         # corresponding to k=0 is located at exactly
         # (output_size//2, output_size//2), regardless of whether the grid_size
-        # and output_size are odd or even.
+        # and output_size are odd or even. This keeps the peak aligned at the
+        # center of a single pixel, which should always have the maximum value.
 
         roll = img.shape[0] / 2
         if img.shape[0] % 2 == 1:
@@ -218,8 +229,6 @@ cdef class DiffractionPattern(_Compute):
         view_orientation = freud.util._convert_array(
             view_orientation, (4,), np.double)
 
-        grid_size = int(self.grid_size / zoom)
-
         # Compute the box projection matrix
         inv_shear = self._calc_proj(view_orientation, system.box)
 
@@ -232,7 +241,7 @@ cdef class DiffractionPattern(_Compute):
         xy += 0.5
         xy %= 1
         im, _, _ = np.histogram2d(
-            xy[:, 0], xy[:, 1], bins=np.linspace(0, 1, grid_size+1))
+            xy[:, 0], xy[:, 1], bins=np.linspace(0, 1, self.grid_size+1))
 
         # Compute FFT and convolve with Gaussian
         cdef double complex[:, :] diffraction_fft
@@ -259,8 +268,7 @@ cdef class DiffractionPattern(_Compute):
         if not self._called_compute:
             # Create a 1D axis of k-vector magnitudes
             self._k_values_orig = np.fft.fftshift(np.fft.fftfreq(
-                n=self.output_size,
-                d=1/self.output_size))
+                n=self.output_size))
 
             # Create a 3D meshgrid of k-vectors with shape
             # (output_size, output_size, 3)
@@ -271,6 +279,7 @@ cdef class DiffractionPattern(_Compute):
         # lazy evaluation of k-values and k-vectors
         self._box_matrix_scale_factor = np.max(system.box.to_matrix())
         self._view_orientation = view_orientation
+        self._k_scale_factor = 2 * np.pi * self.output_size / (self._box_matrix_scale_factor * zoom)
         self._k_values_cached = False
         self._k_vectors_cached = False
 
@@ -290,16 +299,15 @@ cdef class DiffractionPattern(_Compute):
     def diffraction(self):
         """
         (``output_size``, ``output_size``) :class:`numpy.ndarray`:
-            diffraction pattern.
+            Diffraction pattern.
         """
         return np.asarray(self._diffraction) / self._frame_counter
 
     @_Compute._computed_property
     def k_values(self):
         """(``output_size``, ) :class:`numpy.ndarray`: k-values."""
         if not self._k_values_cached:
-            self._k_values = np.asarray(
-                self._k_values_orig) / self._box_matrix_scale_factor
+            self._k_values = np.asarray(self._k_values_orig) * self._k_scale_factor
             self._k_values_cached = True
         return np.asarray(self._k_values)
 
@@ -312,7 +320,7 @@ cdef class DiffractionPattern(_Compute):
         if not self._k_vectors_cached:
             self._k_vectors = rowan.rotate(
                 self._view_orientation,
-                self._k_vectors_orig) / self._box_matrix_scale_factor
+                self._k_vectors_orig) * self._k_scale_factor
             self._k_vectors_cached = True
         return np.asarray(self._k_vectors)
 

freud/plot.py

@@ -11,6 +11,7 @@
 try:
     import matplotlib.pyplot as plt
     from matplotlib.backends.backend_agg import FigureCanvasAgg
+    from matplotlib.ticker import FormatStrFormatter, MaxNLocator
 except ImportError:
     raise ImportError("matplotlib must be installed for freud.plot.")
 
@@ -522,8 +523,13 @@ def diffraction_plot(
 
     # Plot the diffraction image and color bar
     norm = matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax)
+    extent = (np.min(k_values), np.max(k_values), np.min(k_values), np.max(k_values))
     im = ax.imshow(
-        np.clip(diffraction, vmin, vmax), interpolation="nearest", cmap=cmap, norm=norm
+        np.clip(diffraction, vmin, vmax),
+        interpolation="nearest",
+        cmap=cmap,
+        norm=norm,
+        extent=extent,
     )
     ax_divider = make_axes_locatable(ax)
     cax = ax_divider.append_axes("right", size="7%", pad="10%")
@@ -539,12 +545,11 @@ def diffraction_plot(
     ticks = np.linspace(0, grid_size, num_ticks)
 
     # Set tick locations and labels
-    tick_labels = np.interp(ticks, range(grid_size), k_values)
-    tick_labels = [f"{x:.3g}" for x in tick_labels]
-    ax.xaxis.set_ticks(ticks)
-    ax.xaxis.set_ticklabels(tick_labels)
-    ax.yaxis.set_ticks(ticks)
-    ax.yaxis.set_ticklabels(tick_labels)
+    ax.xaxis.set_major_locator(MaxNLocator(nbins=6, symmetric=True, min_n_ticks=7))
+    ax.yaxis.set_major_locator(MaxNLocator(nbins=6, symmetric=True, min_n_ticks=7))
+    formatter = FormatStrFormatter("%.3g")
+    ax.xaxis.set_major_formatter(formatter)
+    ax.yaxis.set_major_formatter(formatter)
 
     # Set title, limits, aspect
     ax.set_title("Diffraction Pattern")

tests/test_diffraction_DiffractionPattern.py

@@ -3,6 +3,7 @@
 import numpy.testing as npt
 import pytest
 import rowan
+from scipy.cluster import vq
 
 import freud
 
@@ -160,11 +161,176 @@ def test_k_values_and_k_vectors(self):
         dp = freud.diffraction.DiffractionPattern()
 
         for size in [2, 5, 10]:
-            for npoints in [10, 20, 75]:
-                box, positions = freud.data.make_random_system(npoints, size)
-                dp.compute((box, positions))
-
-                output_size = dp.output_size
-                npt.assert_allclose(dp.k_values[output_size // 2], 0)
-                center_index = (output_size // 2, output_size // 2)
-                npt.assert_allclose(dp.k_vectors[center_index], [0, 0, 0])
+            box, positions = freud.data.make_random_system(size, 1)
+            zoom = 4
+            view_orientation = np.asarray([1, 0, 0, 0])
+            dp.compute((box, positions), view_orientation=view_orientation, zoom=zoom)
+
+            output_size = dp.output_size
+            npt.assert_allclose(dp.k_values[output_size // 2], 0)
+            center_index = (output_size // 2, output_size // 2)
+            npt.assert_allclose(dp.k_vectors[center_index], [0, 0, 0])
+
+            # Tests for the left and right k_value bins. Uses the math for
+            # np.fft.fftfreq and the _k_scale_factor to write an expression for
+            # the first and last bins' k-values and k-vectors.
+
+            scale_factor = 2 * np.pi * output_size / (np.max(box.to_matrix()) * zoom)
+
+            if output_size % 2 == 0:
+
+                first_k_value = -0.5 * scale_factor
+                last_k_value = (0.5 - (1 / output_size)) * scale_factor
+
+                npt.assert_allclose(dp.k_values[0], first_k_value)
+                npt.assert_allclose(dp.k_values[-1], last_k_value)
+
+                top_left_k_vector = rowan.rotate(
+                    view_orientation, [first_k_value, first_k_value, 0]
+                )
+                bottom_right_k_vector = rowan.rotate(
+                    view_orientation, [last_k_value, last_k_value, 0]
+                )
+                top_right_k_vector = rowan.rotate(
+                    view_orientation, [first_k_value, last_k_value, 0]
+                )
+                bottom_left_k_vector = rowan.rotate(
+                    view_orientation, [last_k_value, first_k_value, 0]
+                )
+
+                npt.assert_allclose(dp.k_vectors[0, 0], first_k_vector)
+                npt.assert_allclose(dp.k_vectors[-1, -1], last_k_vector)
+                npt.assert_allclose(dp.k_vectors[0, -1], top_right_k_vector)
+                npt.assert_allclose(dp.k_vectors[-1, 0], bottom_left_k_vector)
+
+                center = output_size // 2
+                top_center_k_vector = rowan.rotate(
+                    view_orientation, [0, first_k_value, 0]
+                )
+                bottom_center_k_vector = rowan.rotate(
+                    view_orientation, [0, last_k_value, 0]
+                )
+                left_center_k_vector = rowan.rotate(
+                    view_orientation, [first_k_value, 0, 0]
+                )
+                right_center_k_vector = rowan.rotate(
+                    view_orientation, [last_k_value, 0, 0]
+                )
+
+                npt.assert_allclose(dp.k_vectors[center, 0], top_center_k_vector)
+                npt.assert_allclose(dp.k_vectors[center, -1], bottom_center_k_vector)
+                npt.assert_allclose(dp.k_vectors[0, center], left_center_k_vector)
+                npt.assert_allclose(dp.k_vectors[-1, center], right_center_k_vector)
+
+            else:
+
+                first_k_value = (-0.5 + 1 / (2 * output_size)) * scale_factor
+                last_k_value = (0.5 - 1 / (2 * output_size)) * scale_factor
+
+                npt.assert_allclose(dp.k_values[0], first_k_value)
+                npt.assert_allclose(dp.k_values[-1], last_k_value)
+
+                first_k_vector = rowan.rotate(
+                    view_orientation, [first_k_value, first_k_value, 0]
+                )
+                last_k_vector = rowan.rotate(
+                    view_orientation, [last_k_value, last_k_value, 0]
+                )
+                top_right_k_vector = rowan.rotate(
+                    view_orientation, [first_k_value, last_k_value, 0]
+                )
+                bottom_left_k_vector = rowan.rotate(
+                    view_orientation, [last_k_value, first_k_value, 0]
+                )
+
+                npt.assert_allclose(dp.k_vectors[0, 0], first_k_vector)
+                npt.assert_allclose(dp.k_vectors[-1, -1], last_k_vector)
+                npt.assert_allclose(dp.k_vectors[0, -1], top_right_k_vector)
+                npt.assert_allclose(dp.k_vectors[-1, 0], bottom_left_k_vector)
+
+                center = output_size // 2
+                top_center_k_vector = rowan.rotate(
+                    view_orientation, [0, first_k_value, 0]
+                )
+                bottom_center_k_vector = rowan.rotate(
+                    view_orientation, [0, last_k_value, 0]
+                )
+                left_center_k_vector = rowan.rotate(
+                    view_orientation, [first_k_value, 0, 0]
+                )
+                right_center_k_vector = rowan.rotate(
+                    view_orientation, [last_k_value, 0, 0]
+                )
+
+                npt.assert_allclose(dp.k_vectors[center, 0], top_center_k_vector)
+                npt.assert_allclose(dp.k_vectors[center, -1], bottom_center_k_vector)
+                npt.assert_allclose(dp.k_vectors[0, center], left_center_k_vector)
+                npt.assert_allclose(dp.k_vectors[-1, center], right_center_k_vector)
+
+    def test_cubic_system(self):
+        length = 1
+        box, positions = freud.data.UnitCell.sc().generate_system(
+            num_replicas=16, scale=length, sigma_noise=0.1 * length
+        )
+        dp = freud.diffraction.DiffractionPattern()
+        dp.compute((box, positions), zoom=2)
+
+        # Make sure that the peaks are where we expect them.
+        # Identify the indices of the highest values in dp.diffraction
+        # and test that k * R == 2*pi*N for some integer N, for the corresponding peak k vectors
+        # and lattice vectors R. This will be inexact because of binning.
+        threshold = np.partition(dp.diffraction.flatten(), -50)[-50]
+        xs, ys = np.nonzero(dp.diffraction > threshold)
+        xy = np.dstack((xs, ys))[0]
+
+        centers, _ = vq.kmeans(xy.astype(float), 9)
+        centers = centers.astype(int)
+
+        for center in centers:
+            peak_k_vector = dp.k_vectors[center[0], center[1]]
+            lattice_vector = [length, length, length]
+            dot_prod = np.dot(peak_k_vector, lattice_vector)
+
+            npt.assert_allclose(dot_prod % (2 * np.pi), 0)
+
+    def test_cubic_system_parameterized(self):
+        length = 1
+        box, positions = freud.data.UnitCell.sc().generate_system(
+            num_replicas=16, scale=length, sigma_noise=0.1 * length
+        )
+        # Same as above test but with different grid_size,
+        # output_size, and zoom values.
+        for grid_size in (256, 1024):
+            for output_size in (255, 256, 1023, 1024):
+                for zoom in (1, 2.5, 4):
+                    dp = freud.diffraction.DiffractionPattern(
+                        grid_size=grid_size,
+                        output_size=output_size,
+                    )
+                    dp.compute((box, positions), zoom=zoom)
+
+                    # Locate brightest areas of diffraction pattern (intensity > threshold),
+                    # and check that the ideal diffraction peak locations, given by
+                    # k * R = 2*pi*N for some lattice vector R and integer N, are contained
+                    # within these regions. This test only checks N in range [-2, 2].
+                    threshold = 0.2
+                    xs, ys = np.nonzero(dp.diffraction > threshold)
+                    xy = np.dstack((xs, ys))[0]
+
+                    ideal_peaks = {-2: "f", -1: "f", 0: "f", 1: "f", 2: "f"}
+                    all_peaks = True
+
+                    for peak in ideal_peaks:
+                        for x, y in xy:
+                            k_vector = dp.k_vectors[x, y]
+                            lattice_vector = [1, 1, 1]
+                            dot_prod = np.dot(k_vector, lattice_vector)
+
+                            if dot_prod == (peak * 2 * np.pi):
+                                ideal_peaks[peak] = "t"
+
+                    for peak in ideal_peaks:
+                        if ideal_peaks[peak] == "f":
+                            all_peaks = False
+
+                    assert all_peaks == True