Merge pull request #150 from pc494/pyxem-coop-pr

Adds a plot method to DiffractionSimulation
pyxem · Jan 5, 2021 · a445994 · a445994
2 parents 8a5cfca + 73c3af0
commit a445994
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 1 deletion.
diff --git a/diffsims/release_info.py b/diffsims/release_info.py
@@ -1,5 +1,5 @@
 name = "diffsims"
-version = "0.4.0-rc.2"
+version = "0.4.0-rc.3.dev"
 author = "Duncan Johnstone, Phillip Crout"
 copyright = "Copyright 2017-2020, The pyXem Developers"
 credits = [

diff --git a/diffsims/sims/diffraction_simulation.py b/diffsims/sims/diffraction_simulation.py
@@ -167,6 +167,41 @@ def get_diffraction_pattern(self, size=512, sigma=10):
 
         return np.divide(pattern, np.max(pattern))
 
+    def plot(self, size_factor=1, units="real", **kwargs):
+        """A quick-plot function for a simulation of spots
+
+        Parameters
+        ----------
+        size_factor : float, optional
+            linear spot size scaling, default to 1
+        units : str, optional
+            'real' or 'pixel', only changes scalebars, falls back on 'real', the default
+        **kwargs :
+            passed to ax.scatter() method
+
+        Returns
+        -------
+        ax,sp
+
+        Notes
+        -----
+        spot size scales with the square root of the intensity.
+        """
+        _, ax = plt.subplots()
+        ax.set_aspect("equal")
+        if units == "pixel":
+            coords = self.calibrated_coordinates
+        else:
+            coords = self.coordinates
+
+        sp = ax.scatter(
+            coords[:, 0],
+            coords[:, 1],
+            s=size_factor * np.sqrt(self.intensities),
+            **kwargs
+        )
+        return ax, sp
+
 
 class ProfileSimulation:
     """Holds the result of a given kinematic simulation of a diffraction

diff --git a/diffsims/tests/sims/test_diffraction_simulation.py b/diffsims/tests/sims/test_diffraction_simulation.py
@@ -173,3 +173,13 @@ def test_assertion_free_get_diffraction_pattern(self):
             calibration=[2, 2],
         )
         z = empty_sim.get_diffraction_pattern(size=10)
+
+    @pytest.mark.parametrize("units_in",['pixel','real'])
+    def test_plot_method(self,units_in):
+        short_sim = DiffractionSimulation(
+            coordinates=np.asarray([[0.3, 1.2, 0]]),
+            intensities=np.ones(1),
+            calibration=[1, 2],
+        )
+
+        ax,sp = short_sim.plot(units=units_in)