Merge pull request #18 from jennyfothergill/master

Improvements to plot function

gixstapose/diffractometer.py

@@ -296,14 +296,25 @@ def diffract(self, rot, cutout=True):
         self.dp = dp
         return dp
 
-    def plot(self):
+    def plot(self, cmap=None, crop=None, tickspacing=0.5):
         """Plot the diffraction pattern.
 
         The plot will have units in inverse Angstrom calculated from the
         `length_scale` attribute.
         This function will also rotate the diffraction pattern according to the
         `up` attribute of the camera if `diffract_from_camera` was used.
 
+        Parameters
+        ----------
+        cmap : str, default None
+            Name of matplotlib colormap. If None is given, the default colormap
+            for matplotlib.pyplot.imshow will be used.
+        crop : float, default None
+            For small systems where zoom does not give enough precision, crop
+            can be used to zoom the plot to (-crop, crop) in 1/Angstroms.
+        tickspacing : float, default 0.5
+            Spacing between x and x tick values in 1/Angstroms.
+
         Returns
         -------
         matplotlib.figure.Figure, matplotlib.axes._subplots.AxesSubplot
@@ -321,27 +332,40 @@ def plot(self):
             """
             )
         fig, ax = plt.subplots(figsize=(8, 8))
-        extent = (self.N / 2 / self.zoom + 1) / (
-            np.max(self.box) * self.length_scale
+        extent = (
+            (self.N / self.zoom + 1)
+            * np.pi
+            / (np.max(self.box) * self.length_scale)
         )
         dp = self.dp
+        if crop is not None:
+            pts = np.linspace(-extent, extent, self.N)
+            left_idx = np.searchsorted(pts, -crop)
+            right_idx = np.searchsorted(pts, crop)
+            new_dp = dp[left_idx:right_idx, left_idx:right_idx]
+            idbig = self.circle_cutout(new_dp)
+            new_dp[np.unravel_index(idbig, new_dp.shape)] = np.log10(self.bot)
+            dp = new_dp
+            extent = (
+                (new_dp.shape[0] / self.zoom + 1)
+                * np.pi
+                / (np.max(self.box) * self.length_scale)
+            )
         if self.up_ang is not None:
             dp = rotate(
-                self.dp,
+                dp,
                 self.up_ang,
                 reshape=False,
                 cval=np.log10(self.bot),
                 order=1,
             )
-        ax.imshow(dp, extent=[-extent, extent, -extent, extent])
+        ax.imshow(dp, cmap=cmap, extent=[-extent, extent, -extent, extent])
         ax.set_xlabel(r"$q_{xy} (1/\AA)$", fontsize=20)
         ax.set_ylabel(r"$q_{z} (1/\AA)$", fontsize=20)
-        ticks = ticks = [
-            -round(extent, 2),
-            -round(extent / 2, 2),
-            0,
-            round(extent / 2, 2),
-            round(extent, 2),
+        maxtick = extent - (extent % tickspacing)
+        ticks = [
+            round(i, 1)
+            for i in np.arange(-maxtick, maxtick + tickspacing / 2, tickspacing)
         ]
         ax.set_xticks(ticks)
         ax.set_yticks(ticks)

gixstapose/tests/test_diffract.py

@@ -27,7 +27,7 @@ def test_diffract_plot_camera(self, positions_and_box, camera100):
         fig, ax = d.plot()
         assert isinstance(ax, plt.Axes)
         assert isinstance(fig, plt.Figure)
-        assert (-65, 65) == ax.get_xlim()
+        assert all(np.isclose(ax.get_xlim(), (-405, 405), atol=1))
 
     def test_diffract_plot_rot(self, positions_and_box, rot100):
         d = Diffractometer(length_scale=2.0)
@@ -36,7 +36,7 @@ def test_diffract_plot_rot(self, positions_and_box, rot100):
         fig, ax = d.plot()
         assert isinstance(ax, plt.Axes)
         assert isinstance(fig, plt.Figure)
-        assert (-32.5, 32.5) == ax.get_xlim()
+        assert all(np.isclose(ax.get_xlim(), (-202, 202), atol=1))
 
     def test_diffract_plot_raises(self, positions_and_box):
         with pytest.raises(ValueError):