From dd87e30b9362fe51d0ff4ac17c95b53e9192d5a1 Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Mon, 20 Sep 2021 13:46:40 +0200
Subject: [PATCH 1/6] harden diffraction simulation object

---
 diffsims/sims/diffraction_simulation.py       | 49 +++++++++----------
 .../tests/sims/test_diffraction_simulation.py | 32 +++++++-----
 2 files changed, 44 insertions(+), 37 deletions(-)

diff --git a/diffsims/sims/diffraction_simulation.py b/diffsims/sims/diffraction_simulation.py
index cef91ea6..7a19dfff 100644
--- a/diffsims/sims/diffraction_simulation.py
+++ b/diffsims/sims/diffraction_simulation.py
@@ -44,7 +44,7 @@ class DiffractionSimulation:
 
     def __init__(
         self,
-        coordinates=None,
+        coordinates,
         indices=None,
         intensities=None,
         calibration=1.0,
@@ -54,25 +54,33 @@ def __init__(
         """Initializes the DiffractionSimulation object with data values for
         the coordinates, indices, intensities, calibration and offset.
         """
-        self._coordinates = None
-        self.coordinates = coordinates
-        self.indices = indices
-        self._intensities = None
-        self.intensities = intensities
-        self._calibration = (1.0, 1.0)
+        if indices is None:
+            indices = np.zeros((coordinates.shape[0], 3))
+        if intensities is None:
+            intensities = np.zeros((coordinates.shape[0]))
+        # check here whether shapes are all the same
+        if coordinates.shape[0] == indices.shape[0] == intensities.shape[0]:
+            self._coordinates = coordinates
+            self._indices = indices
+            self._intensities = intensities
+        else:
+            raise ValueError("Coordinate, intensity, and indices lists must be same size.")
         self.calibration = calibration
         self.offset = offset
         self.with_direct_beam = with_direct_beam
 
+    @property
+    def indices(self):
+        return self._indices[self.direct_beam_mask]
+
+    @indices.setter
+    def indices(self, indices):
+        self._indices[self.direct_beam_mask] = indices
+
     @property
     def calibrated_coordinates(self):
         """ndarray : Coordinates converted into pixel space."""
-        coordinates = np.copy(self.coordinates)
-        coordinates[:, 0] += self.offset[0]
-        coordinates[:, 1] += self.offset[1]
-        coordinates[:, 0] /= self.calibration[0]
-        coordinates[:, 1] /= self.calibration[1]
-        return coordinates
+        return (self.coordinates[:, :2] + np.array(self.offset))/np.array(self.calibration)
 
     @property
     def calibration(self):
@@ -106,24 +114,20 @@ def direct_beam_mask(self):
     @property
     def coordinates(self):
         """ndarray : The coordinates of all unmasked points."""
-        if self._coordinates is None:
-            return None
         return self._coordinates[self.direct_beam_mask]
 
     @coordinates.setter
     def coordinates(self, coordinates):
-        self._coordinates = coordinates
+        self._coordinates[self.direct_beam_mask] = coordinates
 
     @property
     def intensities(self):
         """ndarray : The intensities of all unmasked points."""
-        if self._intensities is None:
-            return None
         return self._intensities[self.direct_beam_mask]
 
     @intensities.setter
     def intensities(self, intensities):
-        self._intensities = intensities
+        self._intensities[self.direct_beam_mask] = intensities
 
     def get_as_mask(self, shape, radius=6., negative=True,
                     radius_function=None, direct_beam_position=None,
@@ -160,7 +164,7 @@ def get_as_mask(self, shape, radius=6., negative=True,
         cx, cy = shape[0]//2, shape[1]//2
         if direct_beam_position is not None:
             cx, cy = direct_beam_position
-        point_coordinates_shifted = self.calibrated_coordinates[:, :-1].copy()
+        point_coordinates_shifted = self.calibrated_coordinates.copy()
         x = point_coordinates_shifted[:, 0]
         y = point_coordinates_shifted[:, 1]
         mirrored_factor = -1 if mirrored else 1
@@ -269,11 +273,6 @@ def plot(self, size_factor=1, units="real", show_labels=False,
 
         if show_labels:
             millers = self.indices.astype(np.int16)
-            if millers.shape[0] != coords.shape[0]:
-                # ensure we don't label 0,0,0, smallest miller index is 001
-                miller_dist = np.linalg.norm(millers, axis=1)
-                condition = miller_dist > 0.1
-                millers = millers[condition]
             # only label the points inside the axes
             xlim = ax.get_xlim()
             ylim = ax.get_ylim()
diff --git a/diffsims/tests/sims/test_diffraction_simulation.py b/diffsims/tests/sims/test_diffraction_simulation.py
index cc26424f..4ef70dfb 100644
--- a/diffsims/tests/sims/test_diffraction_simulation.py
+++ b/diffsims/tests/sims/test_diffraction_simulation.py
@@ -88,12 +88,12 @@ def test_plot_profile_simulation(profile_simulation):
 class TestDiffractionSimulation:
     @pytest.fixture
     def diffraction_simulation(self):
-        return DiffractionSimulation()
+        return DiffractionSimulation(np.array([[0, 0, 0],[1, 2, 3], [3, 4, 5]]))
 
     def test_init(self, diffraction_simulation):
-        assert diffraction_simulation.coordinates is None
-        assert diffraction_simulation.indices is None
-        assert diffraction_simulation.intensities is None
+        assert np.allclose(diffraction_simulation.coordinates, np.array([[1, 2, 3], [3, 4, 5]]))
+        assert np.allclose(diffraction_simulation.indices, np.array([[0, 0, 0], [0, 0, 0]]))
+        assert np.allclose(diffraction_simulation.intensities, np.array([0, 0]))
         assert diffraction_simulation.calibration == (1.0, 1.0)
 
     @pytest.mark.parametrize(
@@ -127,9 +127,10 @@ def test_calibration(self, diffraction_simulation, calibration, expected):
         ],
     )
     def test_direct_beam_mask(
-        self, diffraction_simulation, coordinates, with_direct_beam, expected
+        self, coordinates, with_direct_beam, expected
     ):
-        diffraction_simulation.coordinates = coordinates
+        diffraction_simulation = DiffractionSimulation(coordinates,
+                                                       with_direct_beam=with_direct_beam)
         diffraction_simulation.with_direct_beam = with_direct_beam
         mask = diffraction_simulation.direct_beam_mask
         assert np.all(mask == expected)
@@ -141,22 +142,28 @@ def test_direct_beam_mask(
                 np.array([[1.0, 0.0, 0.0], [1.0, 2.0, 0.0]]),
                 1.0,
                 (0.0, 0.0),
+                np.array([[1.0, 0.0], [1.0, 2.0]]),
+            ),
+            (
                 np.array([[1.0, 0.0, 0.0], [1.0, 2.0, 0.0]]),
-            )
+                2.0,
+                (3.0, 1.0),
+                np.array([[2.0, 0.5], [2.0, 1.5]]),
+            ),
         ],
     )
     def test_calibrated_coordinates(
         self,
-        diffraction_simulation: DiffractionSimulation,
         coordinates,
         calibration,
         offset,
         expected,
     ):
-        diffraction_simulation.coordinates = coordinates
+        diffraction_simulation = DiffractionSimulation(coordinates)
         diffraction_simulation.calibration = calibration
         diffraction_simulation.offset = offset
-        assert np.allclose(diffraction_simulation.calibrated_coordinates, expected)
+        assert np.allclose(diffraction_simulation.calibrated_coordinates,
+                           expected)
 
     def test_assertion_free_get_diffraction_pattern(self):
         short_sim = DiffractionSimulation(
@@ -191,11 +198,12 @@ def test_get_as_mask(self):
     def test_plot_method(self,units_in):
         short_sim = DiffractionSimulation(
             coordinates=np.asarray([[0.3, 1.2, 0],
-                                    [2.1, 3.4, 0]]),
+                                    [-2, 3, 0],
+                                    [2.1, 3.4, 0],]),
             indices=np.array([[-2, 3, 4],
                               [ 2,-6, 1],
                               [ 0, 0, 0]]),
-            intensities=np.array([3., 5.]),
+            intensities=np.array([3., 5., 2.]),
             calibration=[1, 2],
         )
         ax, sp = short_sim.plot(units=units_in, show_labels=True)

From ae84241a4d07aa0d37dca6e6798e72ccda4b9625 Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Mon, 20 Sep 2021 14:34:31 +0200
Subject: [PATCH 2/6] added slicing into diffraction simulation

---
 diffsims/sims/diffraction_simulation.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/diffsims/sims/diffraction_simulation.py b/diffsims/sims/diffraction_simulation.py
index 7a19dfff..7b72d96c 100644
--- a/diffsims/sims/diffraction_simulation.py
+++ b/diffsims/sims/diffraction_simulation.py
@@ -20,9 +20,10 @@
 import numpy as np
 from diffsims.pattern.detector_functions import add_shot_and_point_spread
 from diffsims.utils import mask_utils
+from collections import Sequence
 
 
-class DiffractionSimulation:
+class DiffractionSimulation(Sequence):
     """Holds the result of a kinematic diffraction pattern simulation.
 
     Parameters
@@ -69,6 +70,21 @@ def __init__(
         self.offset = offset
         self.with_direct_beam = with_direct_beam
 
+    def __len__(self):
+        return self.direct_beam_mask.shape[0]
+
+    def __getitem__(self, sliced):
+        coords = self.coordinates[sliced]
+        inds = self.indices[sliced]
+        ints = self.intensities[sliced]
+        return DiffractionSimulation(coords,
+                                     indices=inds,
+                                     intensities=ints,
+                                     calibration=self.calibration,
+                                     offset=self.offset,
+                                     with_direct_beam=self.with_direct_beam
+                                    )
+
     @property
     def indices(self):
         return self._indices[self.direct_beam_mask]

From 114bdb3a334a9f7bce08f3104c92a6059e9b482d Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Thu, 30 Sep 2021 16:33:16 +0200
Subject: [PATCH 3/6] updated testing and added add/extend feature

---
 diffsims/sims/diffraction_simulation.py       | 170 ++++++++++++++----
 .../tests/sims/test_diffraction_simulation.py |  84 ++++++++-
 2 files changed, 213 insertions(+), 41 deletions(-)

diff --git a/diffsims/sims/diffraction_simulation.py b/diffsims/sims/diffraction_simulation.py
index 7b72d96c..1b46fad4 100644
--- a/diffsims/sims/diffraction_simulation.py
+++ b/diffsims/sims/diffraction_simulation.py
@@ -48,35 +48,45 @@ def __init__(
         coordinates,
         indices=None,
         intensities=None,
-        calibration=1.0,
+        calibration=None,
         offset=(0.0, 0.0),
         with_direct_beam=False,
     ):
         """Initializes the DiffractionSimulation object with data values for
         the coordinates, indices, intensities, calibration and offset.
         """
+        if coordinates.ndim == 1:
+            coordinates = coordinates[None,:]
         if indices is None:
-            indices = np.zeros((coordinates.shape[0], 3))
+            indices = np.full((coordinates.shape[0], 3), np.nan)
         if intensities is None:
-            intensities = np.zeros((coordinates.shape[0]))
+            intensities = np.full((coordinates.shape[0]), np.nan)
         # check here whether shapes are all the same
-        if coordinates.shape[0] == indices.shape[0] == intensities.shape[0]:
+        if (coordinates.shape[0] == indices.shape[0] == intensities.shape[0] and
+            coordinates.ndim == indices.ndim == 2 and intensities.ndim==1):
             self._coordinates = coordinates
             self._indices = indices
             self._intensities = intensities
         else:
-            raise ValueError("Coordinate, intensity, and indices lists must be same size.")
+            raise ValueError("Coordinate, intensity, and indices lists must be of the correct and matching shape.")
         self.calibration = calibration
         self.offset = offset
         self.with_direct_beam = with_direct_beam
 
     def __len__(self):
-        return self.direct_beam_mask.shape[0]
+        return self.coordinates.shape[0]
 
     def __getitem__(self, sliced):
         coords = self.coordinates[sliced]
         inds = self.indices[sliced]
         ints = self.intensities[sliced]
+        if coords.ndim == 1:
+            coords = coords[None, :]
+            inds = inds[None, :]
+            ints = ints[None]
+        # some valid numpy slices will create invalid shapes for diffraction simulation
+        if coords.ndim > 2 or coords.shape[1] > 3 or coords.shape[1] < 2:
+            raise ValueError(f"Invalid slice: {sliced}")
         return DiffractionSimulation(coords,
                                      indices=inds,
                                      intensities=ints,
@@ -85,6 +95,29 @@ def __getitem__(self, sliced):
                                      with_direct_beam=self.with_direct_beam
                                     )
 
+    def _combine_data(self, other):
+        coords = np.concatenate([self._coordinates, other._coordinates], axis=0)
+        inds = np.concatenate([self._indices, other._indices], axis=0)
+        ints = np.concatenate([self._intensities, other._intensities], axis=0)
+        return coords, inds, ints
+
+    def __add__(self, other):
+        coords, inds, ints = self._combine_data(other)
+        return DiffractionSimulation(coords,
+                                     indices=inds,
+                                     intensities=ints,
+                                     calibration=self.calibration,
+                                     offset=self.offset,
+                                     with_direct_beam=self.with_direct_beam
+                                    )
+
+    def extend(self, other):
+        """Add the diffraction spots from another DiffractionSimulation"""
+        coords, inds, ints = self._combine_data(other)
+        self._coordinates = coords
+        self._indices = inds
+        self._ints = ints
+
     @property
     def indices(self):
         return self._indices[self.direct_beam_mask]
@@ -96,7 +129,10 @@ def indices(self, indices):
     @property
     def calibrated_coordinates(self):
         """ndarray : Coordinates converted into pixel space."""
-        return (self.coordinates[:, :2] + np.array(self.offset))/np.array(self.calibration)
+        if self.calibration is not None:
+            return (self.coordinates[:, :2] + np.array(self.offset))/np.array(self.calibration)
+        else:
+            raise Exception("Pixel calibration is not set!")
 
     @property
     def calibration(self):
@@ -106,6 +142,9 @@ def calibration(self):
 
     @calibration.setter
     def calibration(self, calibration):
+        if calibration is None:
+            self._calibration = calibration
+            return
         if np.all(np.equal(calibration, 0)):
             raise ValueError("`calibration` cannot be zero.")
         if isinstance(calibration, float) or isinstance(calibration, int):
@@ -145,6 +184,47 @@ def intensities(self):
     def intensities(self, intensities):
         self._intensities[self.direct_beam_mask] = intensities
 
+    def _get_transformed_coordinates(self, angle, center=(0, 0),
+                                     mirrored=False, units="real"):
+        """
+        Get the coordinates of diffraction spots rotated, flipped or shifted
+        in-plane
+
+        Parameters
+        ----------
+        angle: float
+            Angle in degrees
+        center: 2-tuple of floats
+            Center of the patterns
+        mirrored: bool
+            Mirror accross the x axis
+        units: str
+            "real" or "pixel"
+
+        Returns
+        -------
+        coords_transformed: np.ndarray of shape (N, 3) or (N, 2)
+            transformed coordinates
+        """
+        if units == "real":
+            coords_transformed = self.coordinates.copy()
+        else:
+            coords_transformed = self.calibrated_coordinates.copy()
+        cx, cy = center
+        x = coords_transformed[:, 0]
+        y = coords_transformed[:, 1]
+        mirrored_factor = -1 if mirrored else 1
+        theta = mirrored_factor * np.arctan2(y, x) + np.deg2rad(angle)
+        rd = np.sqrt(x**2 + y**2)
+        coords_transformed[:, 0] = rd * np.cos(theta) + cx
+        coords_transformed[:, 1] = rd * np.sin(theta) + cy
+        return coords_transformed
+
+    def rotate_shift_coordinates(self, angle, center=(0, 0), mirrored=False):
+        """Translate, rotate or mirror the pattern"""
+        coords_new = self._get_transformed_coordinates(angle, center, mirrored, units="real")
+        self.coordinates = coords_new
+
     def get_as_mask(self, shape, radius=6., negative=True,
                     radius_function=None, direct_beam_position=None,
                     in_plane_angle=0, mirrored=False,
@@ -175,19 +255,18 @@ def get_as_mask(self, shape, radius=6., negative=True,
         mirrored: bool, optional
             Whether the pattern should be flipped over the x-axis,
             corresponding to the inverted orientation
+
+        Returns
+        -------
+        mask: numpy.ndarray
+            Boolean mask of the diffraction pattern
         """
         r = radius
-        cx, cy = shape[0]//2, shape[1]//2
+        cx, cy = shape[1]//2, shape[0]//2
         if direct_beam_position is not None:
             cx, cy = direct_beam_position
-        point_coordinates_shifted = self.calibrated_coordinates.copy()
-        x = point_coordinates_shifted[:, 0]
-        y = point_coordinates_shifted[:, 1]
-        mirrored_factor = -1 if mirrored else 1
-        theta = mirrored_factor * np.arctan2(y, x) + np.deg2rad(in_plane_angle)
-        rd = np.sqrt(x**2 + y**2)
-        point_coordinates_shifted[:, 0] = rd * np.cos(theta) + cx
-        point_coordinates_shifted[:, 1] = rd * np.sin(theta) + cy
+        point_coordinates_shifted = self._get_transformed_coordinates(
+                in_plane_angle, center=(cx, cy), mirrored=mirrored, units="pixels")
         if radius_function is not None:
             r = radius_function(self.intensities, *args, **kwargs)
         mask = mask_utils.create_mask(shape, fill=negative)
@@ -195,18 +274,27 @@ def get_as_mask(self, shape, radius=6., negative=True,
                                     fill=not negative)
         return mask
 
-    def get_diffraction_pattern(self, size=512, sigma=10):
+    def get_diffraction_pattern(self, shape=(512, 512), sigma=10,
+                                direct_beam_position=None, in_plane_angle=0,
+                                mirrored=False):
         """Returns the diffraction data as a numpy array with
         two-dimensional Gaussians representing each diffracted peak. Should only
         be used for qualitative work.
 
         Parameters
         ----------
-        size  : int
+        shape  : tuple of ints
             The size of a side length (in pixels)
-
         sigma : float
             Standard deviation of the Gaussian function to be plotted (in pixels).
+        direct_beam_position: 2-tuple of ints, optional
+            The (x,y) coordinate in pixels of the direct beam. Defaults to
+            the center of the image.
+        in_plane_angle: float, optional
+            In plane rotation of the pattern in degrees
+        mirrored: bool, optional
+            Whether the pattern should be flipped over the x-axis,
+            corresponding to the inverted orientation
 
         Returns
         -------
@@ -219,26 +307,29 @@ def get_diffraction_pattern(self, size=512, sigma=10):
         produces reasonably good patterns when the lattice parameters are on
         the order of 0.5nm and a the default size and sigma are used.
         """
-        side_length = np.min(np.multiply((size / 2), self.calibration))
-        mask_for_sides = np.all(
-            (np.abs(self.coordinates[:, 0:2]) < side_length), axis=1
-        )
-
-        spot_coords = np.add(
-            self.calibrated_coordinates[mask_for_sides], size / 2
-        ).astype(int)
-        spot_intens = self.intensities[mask_for_sides]
-        pattern = np.zeros([size, size])
+        cx, cy = shape[1]//2, shape[0]//2
+        if direct_beam_position is not None:
+            cx, cy = direct_beam_position
+        coordinates = self._get_transformed_coordinates(in_plane_angle, (cx, cy), mirrored, units="pixel")
+        in_frame = ((coordinates[:, 0] >= 0) & (coordinates[:, 0] < shape[1]) &
+                    (coordinates[:, 1] >= 0) & (coordinates[:, 1] < shape[0]))
+        spot_coords = coordinates[in_frame].astype(int)
+        spot_intens = self.intensities[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))
 
-    def plot(self, size_factor=1, units="real", show_labels=False,
+    def plot(self, size_factor=1,
+            direct_beam_position=None,
+            in_plane_angle=0,
+            mirrored=False,
+            units="real",
+            show_labels=False,
             label_offset=(0, 0),
             label_formatting={},
             ax=None,
@@ -249,6 +340,14 @@ def plot(self, size_factor=1, units="real", show_labels=False,
         ----------
         size_factor : float, optional
             linear spot size scaling, default to 1
+        direct_beam_position: 2-tuple of ints, optional
+            The (x,y) coordinate in pixels of the direct beam. Defaults to
+            the center of the image.
+        in_plane_angle: float, optional
+            In plane rotation of the pattern in degrees
+        mirrored: bool, optional
+            Whether the pattern should be flipped over the x-axis,
+            corresponding to the inverted orientation
         units : str, optional
             'real' or 'pixel', only changes scalebars, falls back on 'real', the default
         show_labels : bool, optional
@@ -272,14 +371,13 @@ def plot(self, size_factor=1, units="real", show_labels=False,
         -----
         spot size scales with the square root of the intensity.
         """
+        cx, cy = 0, 0
+        if direct_beam_position is not None:
+            cx, cy = direct_beam_position
         if ax is None:
             _, ax = plt.subplots()
             ax.set_aspect("equal")
-        if units == "pixel":
-            coords = self.calibrated_coordinates
-        else:
-            coords = self.coordinates
-
+        coords = self._get_transformed_coordinates(in_plane_angle, (cx, cy), mirrored, units=units)
         sp = ax.scatter(
             coords[:, 0],
             coords[:, 1],
diff --git a/diffsims/tests/sims/test_diffraction_simulation.py b/diffsims/tests/sims/test_diffraction_simulation.py
index 4ef70dfb..b07f621d 100644
--- a/diffsims/tests/sims/test_diffraction_simulation.py
+++ b/diffsims/tests/sims/test_diffraction_simulation.py
@@ -90,11 +90,63 @@ class TestDiffractionSimulation:
     def diffraction_simulation(self):
         return DiffractionSimulation(np.array([[0, 0, 0],[1, 2, 3], [3, 4, 5]]))
 
+    @pytest.fixture
+    def diffraction_simulation_calibrated(self):
+        return DiffractionSimulation(np.array([[0, 0, 0],[1, 2, 3], [3, 4, 5]]), calibration=0.5)
+
+    @pytest.mark.xfail(raises=ValueError)
+    def failed_initialization(self):
+        DiffractionSimulation(np.array([[0, 0, 0],[1, 2, 3], [3, 4, 5]]), indices=np.array([1, 2, 3]))
+
     def test_init(self, diffraction_simulation):
         assert np.allclose(diffraction_simulation.coordinates, np.array([[1, 2, 3], [3, 4, 5]]))
-        assert np.allclose(diffraction_simulation.indices, np.array([[0, 0, 0], [0, 0, 0]]))
-        assert np.allclose(diffraction_simulation.intensities, np.array([0, 0]))
-        assert diffraction_simulation.calibration == (1.0, 1.0)
+        assert np.isnan(diffraction_simulation.indices).all()
+        assert diffraction_simulation.indices.shape == (2, 3)
+        assert np.isnan(diffraction_simulation.intensities).all()
+        assert diffraction_simulation.intensities.shape == (2,)
+        assert diffraction_simulation.calibration is None
+        assert len(diffraction_simulation) == 2
+
+    def get_get_item(self, diffraction_simulation):
+        assert diffraction_simulation[1].coordinates.shape == (1, 3)
+        assert diffraction_simulation[1:3].coordinates.shape == (2, 3)
+
+    @pytest.mark.xfail(raises=ValueError)
+    def fail_get_item(self, diffraction_simulation):
+        diffraction_simulation[None]
+
+    def test_addition(self, diffraction_simulation):
+        sim = diffraction_simulation + diffraction_simulation
+        assert np.allclose(sim.coordinates, np.array([[1, 2, 3],
+                                                      [3, 4, 5],
+                                                      [1, 2, 3],
+                                                      [3, 4, 5],
+                                                      ]))
+
+    def test_extend(self, diffraction_simulation):
+        diffraction_simulation.extend(diffraction_simulation)
+        assert np.allclose(diffraction_simulation.coordinates, np.array([
+                                                      [1, 2, 3],
+                                                      [3, 4, 5],
+                                                      [1, 2, 3],
+                                                      [3, 4, 5],
+                                                      ]))
+
+    def test_indices_setter_getter(self, diffraction_simulation):
+        indices = np.array([[1, 2, 3],[2, 3,4],[3, 4, 5]])
+        diffraction_simulation.indices = indices[-1]
+        assert np.allclose(diffraction_simulation.indices, indices[-1])
+        diffraction_simulation.with_direct_beam = True
+        diffraction_simulation.indices = indices
+        assert np.allclose(diffraction_simulation.indices, indices)
+
+    def test_coordinates_setter(self, diffraction_simulation):
+        coords = np.array([[1, 2, 3],[2, 3,4],[3, 4, 5]])
+        diffraction_simulation.coordinates = coords[-1]
+        assert np.allclose(diffraction_simulation.coordinates, coords[-1])
+        diffraction_simulation.with_direct_beam = True
+        diffraction_simulation.coordinates = coords
+        assert np.allclose(diffraction_simulation.coordinates, coords)
 
     @pytest.mark.parametrize(
         "calibration, expected",
@@ -150,6 +202,13 @@ def test_direct_beam_mask(
                 (3.0, 1.0),
                 np.array([[2.0, 0.5], [2.0, 1.5]]),
             ),
+            pytest.param(
+                np.array([[1.0, 0.0, 0.0], [1.0, 2.0, 0.0]]),
+                None,
+                (0.0, 0.0),
+                None,
+                marks=pytest.mark.xfail(raises=Exception)
+            ),
         ],
     )
     def test_calibrated_coordinates(
@@ -165,13 +224,28 @@ def test_calibrated_coordinates(
         assert np.allclose(diffraction_simulation.calibrated_coordinates,
                            expected)
 
+    @pytest.mark.parametrize(
+        "units, expected",
+        [
+            ("real", np.array([[-2,1,3],[-4,3,5]])),
+            ("pixel", np.array([[-4,2],[-8,6]])),
+        ],
+    )
+    def test_transform_coordinates(self, diffraction_simulation_calibrated, units, expected):
+        tc = diffraction_simulation_calibrated._get_transformed_coordinates(90, units=units)
+        assert np.allclose(tc, expected)
+
+    def test_rotate_shift_coordinates(self, diffraction_simulation):
+        diffraction_simulation.rotate_shift_coordinates(90)
+        assert np.allclose(diffraction_simulation.coordinates,
+                           np.array([[-2,1,3],[-4,3,5]]))
+
     def test_assertion_free_get_diffraction_pattern(self):
         short_sim = DiffractionSimulation(
             coordinates=np.asarray([[0.3, 1.2, 0]]),
             intensities=np.ones(1),
             calibration=[1, 2],
         )
-
         z = short_sim.get_diffraction_pattern()
 
         empty_sim = DiffractionSimulation(
@@ -179,7 +253,7 @@ def test_assertion_free_get_diffraction_pattern(self):
             intensities=np.ones(1),
             calibration=[2, 2],
         )
-        z = empty_sim.get_diffraction_pattern(size=10)
+        z = empty_sim.get_diffraction_pattern(shape=(10, 20))
 
     def test_get_as_mask(self):
         short_sim = DiffractionSimulation(

From a07166872da437610e25bab9b2c7cf9291b32dcd Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Thu, 30 Sep 2021 17:04:26 +0200
Subject: [PATCH 4/6] coverage 100%

---
 diffsims/sims/diffraction_simulation.py       | 21 ++++++++-----------
 .../tests/sims/test_diffraction_simulation.py | 20 +++++++++++++-----
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/diffsims/sims/diffraction_simulation.py b/diffsims/sims/diffraction_simulation.py
index 1b46fad4..aeda3d87 100644
--- a/diffsims/sims/diffraction_simulation.py
+++ b/diffsims/sims/diffraction_simulation.py
@@ -262,11 +262,10 @@ def get_as_mask(self, shape, radius=6., negative=True,
             Boolean mask of the diffraction pattern
         """
         r = radius
-        cx, cy = shape[1]//2, shape[0]//2
-        if direct_beam_position is not None:
-            cx, cy = direct_beam_position
+        if direct_beam_position is None:
+            direct_beam_position = (shape[1]//2, shape[0]//2)
         point_coordinates_shifted = self._get_transformed_coordinates(
-                in_plane_angle, center=(cx, cy), mirrored=mirrored, units="pixels")
+                in_plane_angle, center=direct_beam_position, mirrored=mirrored, units="pixels")
         if radius_function is not None:
             r = radius_function(self.intensities, *args, **kwargs)
         mask = mask_utils.create_mask(shape, fill=negative)
@@ -307,10 +306,9 @@ def get_diffraction_pattern(self, shape=(512, 512), sigma=10,
         produces reasonably good patterns when the lattice parameters are on
         the order of 0.5nm and a the default size and sigma are used.
         """
-        cx, cy = shape[1]//2, shape[0]//2
-        if direct_beam_position is not None:
-            cx, cy = direct_beam_position
-        coordinates = self._get_transformed_coordinates(in_plane_angle, (cx, cy), mirrored, units="pixel")
+        if direct_beam_position is None:
+            direct_beam_position = (shape[1]//2, shape[0]//2)
+        coordinates = self._get_transformed_coordinates(in_plane_angle, direct_beam_position, mirrored, units="pixel")
         in_frame = ((coordinates[:, 0] >= 0) & (coordinates[:, 0] < shape[1]) &
                     (coordinates[:, 1] >= 0) & (coordinates[:, 1] < shape[0]))
         spot_coords = coordinates[in_frame].astype(int)
@@ -371,13 +369,12 @@ def plot(self, size_factor=1,
         -----
         spot size scales with the square root of the intensity.
         """
-        cx, cy = 0, 0
-        if direct_beam_position is not None:
-            cx, cy = direct_beam_position
+        if direct_beam_position is None:
+            direct_beam_position = (0, 0)
         if ax is None:
             _, ax = plt.subplots()
             ax.set_aspect("equal")
-        coords = self._get_transformed_coordinates(in_plane_angle, (cx, cy), mirrored, units=units)
+        coords = self._get_transformed_coordinates(in_plane_angle, direct_beam_position, mirrored, units=units)
         sp = ax.scatter(
             coords[:, 0],
             coords[:, 1],
diff --git a/diffsims/tests/sims/test_diffraction_simulation.py b/diffsims/tests/sims/test_diffraction_simulation.py
index b07f621d..14e04090 100644
--- a/diffsims/tests/sims/test_diffraction_simulation.py
+++ b/diffsims/tests/sims/test_diffraction_simulation.py
@@ -95,7 +95,7 @@ def diffraction_simulation_calibrated(self):
         return DiffractionSimulation(np.array([[0, 0, 0],[1, 2, 3], [3, 4, 5]]), calibration=0.5)
 
     @pytest.mark.xfail(raises=ValueError)
-    def failed_initialization(self):
+    def test_failed_initialization(self):
         DiffractionSimulation(np.array([[0, 0, 0],[1, 2, 3], [3, 4, 5]]), indices=np.array([1, 2, 3]))
 
     def test_init(self, diffraction_simulation):
@@ -107,12 +107,15 @@ def test_init(self, diffraction_simulation):
         assert diffraction_simulation.calibration is None
         assert len(diffraction_simulation) == 2
 
-    def get_get_item(self, diffraction_simulation):
+    def test_single_spot(self):
+        assert DiffractionSimulation(np.array([1, 2, 3])).coordinates.shape == (1, 3)
+
+    def test_get_item(self, diffraction_simulation):
         assert diffraction_simulation[1].coordinates.shape == (1, 3)
-        assert diffraction_simulation[1:3].coordinates.shape == (2, 3)
+        assert diffraction_simulation[0:2].coordinates.shape == (2, 3)
 
     @pytest.mark.xfail(raises=ValueError)
-    def fail_get_item(self, diffraction_simulation):
+    def test_fail_get_item(self, diffraction_simulation):
         diffraction_simulation[None]
 
     def test_addition(self, diffraction_simulation):
@@ -148,6 +151,14 @@ def test_coordinates_setter(self, diffraction_simulation):
         diffraction_simulation.coordinates = coords
         assert np.allclose(diffraction_simulation.coordinates, coords)
 
+    def test_intensities_setter(self, diffraction_simulation):
+        ints = np.array([1, 2, 3])
+        diffraction_simulation.intensities = ints[-1]
+        assert np.allclose(diffraction_simulation.intensities, ints[-1])
+        diffraction_simulation.with_direct_beam = True
+        diffraction_simulation.intensities = ints
+        assert np.allclose(diffraction_simulation.intensities, ints)
+
     @pytest.mark.parametrize(
         "calibration, expected",
         [
@@ -263,7 +274,6 @@ def test_get_as_mask(self):
         )
         mask = short_sim.get_as_mask((20, 10),
                                      radius_function=np.sqrt,
-                                     direct_beam_position=(10, 6)
                                      )
         assert mask.shape[0] == 20
         assert mask.shape[1] == 10

From 668e3770aef384f2445bf3aca829e85d409805b4 Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Thu, 30 Sep 2021 17:08:37 +0200
Subject: [PATCH 5/6] [skip ci] formatting

---
 diffsims/sims/diffraction_simulation.py       | 175 +++++++++++-------
 .../tests/sims/test_diffraction_simulation.py | 109 ++++++-----
 2 files changed, 177 insertions(+), 107 deletions(-)

diff --git a/diffsims/sims/diffraction_simulation.py b/diffsims/sims/diffraction_simulation.py
index aeda3d87..7b87624d 100644
--- a/diffsims/sims/diffraction_simulation.py
+++ b/diffsims/sims/diffraction_simulation.py
@@ -56,19 +56,24 @@ def __init__(
         the coordinates, indices, intensities, calibration and offset.
         """
         if coordinates.ndim == 1:
-            coordinates = coordinates[None,:]
+            coordinates = coordinates[None, :]
         if indices is None:
             indices = np.full((coordinates.shape[0], 3), np.nan)
         if intensities is None:
             intensities = np.full((coordinates.shape[0]), np.nan)
         # check here whether shapes are all the same
-        if (coordinates.shape[0] == indices.shape[0] == intensities.shape[0] and
-            coordinates.ndim == indices.ndim == 2 and intensities.ndim==1):
+        if (
+            coordinates.shape[0] == indices.shape[0] == intensities.shape[0]
+            and coordinates.ndim == indices.ndim == 2
+            and intensities.ndim == 1
+        ):
             self._coordinates = coordinates
             self._indices = indices
             self._intensities = intensities
         else:
-            raise ValueError("Coordinate, intensity, and indices lists must be of the correct and matching shape.")
+            raise ValueError(
+                "Coordinate, intensity, and indices lists must be of the correct and matching shape."
+            )
         self.calibration = calibration
         self.offset = offset
         self.with_direct_beam = with_direct_beam
@@ -87,13 +92,14 @@ def __getitem__(self, sliced):
         # some valid numpy slices will create invalid shapes for diffraction simulation
         if coords.ndim > 2 or coords.shape[1] > 3 or coords.shape[1] < 2:
             raise ValueError(f"Invalid slice: {sliced}")
-        return DiffractionSimulation(coords,
-                                     indices=inds,
-                                     intensities=ints,
-                                     calibration=self.calibration,
-                                     offset=self.offset,
-                                     with_direct_beam=self.with_direct_beam
-                                    )
+        return DiffractionSimulation(
+            coords,
+            indices=inds,
+            intensities=ints,
+            calibration=self.calibration,
+            offset=self.offset,
+            with_direct_beam=self.with_direct_beam,
+        )
 
     def _combine_data(self, other):
         coords = np.concatenate([self._coordinates, other._coordinates], axis=0)
@@ -103,13 +109,14 @@ def _combine_data(self, other):
 
     def __add__(self, other):
         coords, inds, ints = self._combine_data(other)
-        return DiffractionSimulation(coords,
-                                     indices=inds,
-                                     intensities=ints,
-                                     calibration=self.calibration,
-                                     offset=self.offset,
-                                     with_direct_beam=self.with_direct_beam
-                                    )
+        return DiffractionSimulation(
+            coords,
+            indices=inds,
+            intensities=ints,
+            calibration=self.calibration,
+            offset=self.offset,
+            with_direct_beam=self.with_direct_beam,
+        )
 
     def extend(self, other):
         """Add the diffraction spots from another DiffractionSimulation"""
@@ -130,7 +137,9 @@ def indices(self, indices):
     def calibrated_coordinates(self):
         """ndarray : Coordinates converted into pixel space."""
         if self.calibration is not None:
-            return (self.coordinates[:, :2] + np.array(self.offset))/np.array(self.calibration)
+            return (self.coordinates[:, :2] + np.array(self.offset)) / np.array(
+                self.calibration
+            )
         else:
             raise Exception("Pixel calibration is not set!")
 
@@ -184,8 +193,9 @@ def intensities(self):
     def intensities(self, intensities):
         self._intensities[self.direct_beam_mask] = intensities
 
-    def _get_transformed_coordinates(self, angle, center=(0, 0),
-                                     mirrored=False, units="real"):
+    def _get_transformed_coordinates(
+        self, angle, center=(0, 0), mirrored=False, units="real"
+    ):
         """
         Get the coordinates of diffraction spots rotated, flipped or shifted
         in-plane
@@ -215,20 +225,30 @@ def _get_transformed_coordinates(self, angle, center=(0, 0),
         y = coords_transformed[:, 1]
         mirrored_factor = -1 if mirrored else 1
         theta = mirrored_factor * np.arctan2(y, x) + np.deg2rad(angle)
-        rd = np.sqrt(x**2 + y**2)
+        rd = np.sqrt(x ** 2 + y ** 2)
         coords_transformed[:, 0] = rd * np.cos(theta) + cx
         coords_transformed[:, 1] = rd * np.sin(theta) + cy
         return coords_transformed
 
     def rotate_shift_coordinates(self, angle, center=(0, 0), mirrored=False):
         """Translate, rotate or mirror the pattern"""
-        coords_new = self._get_transformed_coordinates(angle, center, mirrored, units="real")
+        coords_new = self._get_transformed_coordinates(
+            angle, center, mirrored, units="real"
+        )
         self.coordinates = coords_new
 
-    def get_as_mask(self, shape, radius=6., negative=True,
-                    radius_function=None, direct_beam_position=None,
-                    in_plane_angle=0, mirrored=False,
-                    *args, **kwargs):
+    def get_as_mask(
+        self,
+        shape,
+        radius=6.0,
+        negative=True,
+        radius_function=None,
+        direct_beam_position=None,
+        in_plane_angle=0,
+        mirrored=False,
+        *args,
+        **kwargs,
+    ):
         """
         Return the diffraction pattern as a binary mask of type
         bool
@@ -263,19 +283,29 @@ def get_as_mask(self, shape, radius=6., negative=True,
         """
         r = radius
         if direct_beam_position is None:
-            direct_beam_position = (shape[1]//2, shape[0]//2)
+            direct_beam_position = (shape[1] // 2, shape[0] // 2)
         point_coordinates_shifted = self._get_transformed_coordinates(
-                in_plane_angle, center=direct_beam_position, mirrored=mirrored, units="pixels")
+            in_plane_angle,
+            center=direct_beam_position,
+            mirrored=mirrored,
+            units="pixels",
+        )
         if radius_function is not None:
             r = radius_function(self.intensities, *args, **kwargs)
         mask = mask_utils.create_mask(shape, fill=negative)
-        mask_utils.add_circles_to_mask(mask, point_coordinates_shifted, r,
-                                    fill=not negative)
+        mask_utils.add_circles_to_mask(
+            mask, point_coordinates_shifted, r, fill=not negative
+        )
         return mask
 
-    def get_diffraction_pattern(self, shape=(512, 512), sigma=10,
-                                direct_beam_position=None, in_plane_angle=0,
-                                mirrored=False):
+    def get_diffraction_pattern(
+        self,
+        shape=(512, 512),
+        sigma=10,
+        direct_beam_position=None,
+        in_plane_angle=0,
+        mirrored=False,
+    ):
         """Returns the diffraction data as a numpy array with
         two-dimensional Gaussians representing each diffracted peak. Should only
         be used for qualitative work.
@@ -307,10 +337,16 @@ def get_diffraction_pattern(self, shape=(512, 512), sigma=10,
         the order of 0.5nm and a the default size and sigma are used.
         """
         if direct_beam_position is None:
-            direct_beam_position = (shape[1]//2, shape[0]//2)
-        coordinates = self._get_transformed_coordinates(in_plane_angle, direct_beam_position, mirrored, units="pixel")
-        in_frame = ((coordinates[:, 0] >= 0) & (coordinates[:, 0] < shape[1]) &
-                    (coordinates[:, 1] >= 0) & (coordinates[:, 1] < shape[0]))
+            direct_beam_position = (shape[1] // 2, shape[0] // 2)
+        coordinates = self._get_transformed_coordinates(
+            in_plane_angle, direct_beam_position, mirrored, units="pixel"
+        )
+        in_frame = (
+            (coordinates[:, 0] >= 0)
+            & (coordinates[:, 0] < shape[1])
+            & (coordinates[:, 1] >= 0)
+            & (coordinates[:, 1] < shape[0])
+        )
         spot_coords = coordinates[in_frame].astype(int)
         spot_intens = self.intensities[in_frame]
         pattern = np.zeros(shape)
@@ -322,16 +358,19 @@ def get_diffraction_pattern(self, shape=(512, 512), sigma=10,
             pattern = add_shot_and_point_spread(pattern.T, sigma, shot_noise=False)
         return np.divide(pattern, np.max(pattern))
 
-    def plot(self, size_factor=1,
-            direct_beam_position=None,
-            in_plane_angle=0,
-            mirrored=False,
-            units="real",
-            show_labels=False,
-            label_offset=(0, 0),
-            label_formatting={},
-            ax=None,
-            **kwargs):
+    def plot(
+        self,
+        size_factor=1,
+        direct_beam_position=None,
+        in_plane_angle=0,
+        mirrored=False,
+        units="real",
+        show_labels=False,
+        label_offset=(0, 0),
+        label_formatting={},
+        ax=None,
+        **kwargs,
+    ):
         """A quick-plot function for a simulation of spots
 
         Parameters
@@ -374,12 +413,14 @@ def plot(self, size_factor=1,
         if ax is None:
             _, ax = plt.subplots()
             ax.set_aspect("equal")
-        coords = self._get_transformed_coordinates(in_plane_angle, direct_beam_position, mirrored, units=units)
+        coords = self._get_transformed_coordinates(
+            in_plane_angle, direct_beam_position, mirrored, units=units
+        )
         sp = ax.scatter(
             coords[:, 0],
             coords[:, 1],
             s=size_factor * np.sqrt(self.intensities),
-            **kwargs
+            **kwargs,
         )
 
         if show_labels:
@@ -387,31 +428,37 @@ def plot(self, size_factor=1,
             # only label the points inside the axes
             xlim = ax.get_xlim()
             ylim = ax.get_ylim()
-            condition = ((coords[:,0] > min(xlim)) &
-                         (coords[:,0] < max(xlim)) &
-                         (coords[:,1] > min(ylim)) &
-                         (coords[:,1] < max(ylim)))
+            condition = (
+                (coords[:, 0] > min(xlim))
+                & (coords[:, 0] < max(xlim))
+                & (coords[:, 1] > min(ylim))
+                & (coords[:, 1] < max(ylim))
+            )
             millers = millers[condition]
             coords = coords[condition]
             # default alignment options
-            if "ha" not in label_offset and "horizontalalignment" not in label_formatting:
-                label_formatting["ha"]="center"
+            if (
+                "ha" not in label_offset
+                and "horizontalalignment" not in label_formatting
+            ):
+                label_formatting["ha"] = "center"
             if "va" not in label_offset and "verticalalignment" not in label_formatting:
-                label_formatting["va"]="center"
+                label_formatting["va"] = "center"
             for miller, coordinate in zip(millers, coords):
                 label = "("
                 for index in miller:
-                    if index<0:
-                        label += r"$\bar{" + str(abs(index)) +r"}$"
+                    if index < 0:
+                        label += r"$\bar{" + str(abs(index)) + r"}$"
                     else:
                         label += str(abs(index))
                     label += " "
                 label = label[:-1] + ")"
-                ax.text(coordinate[0] + label_offset[0],
-                        coordinate[1] + label_offset[1],
-                        label,
-                        **label_formatting,
-                        )
+                ax.text(
+                    coordinate[0] + label_offset[0],
+                    coordinate[1] + label_offset[1],
+                    label,
+                    **label_formatting,
+                )
         return ax, sp
 
 
diff --git a/diffsims/tests/sims/test_diffraction_simulation.py b/diffsims/tests/sims/test_diffraction_simulation.py
index 14e04090..b04bbf95 100644
--- a/diffsims/tests/sims/test_diffraction_simulation.py
+++ b/diffsims/tests/sims/test_diffraction_simulation.py
@@ -88,18 +88,24 @@ def test_plot_profile_simulation(profile_simulation):
 class TestDiffractionSimulation:
     @pytest.fixture
     def diffraction_simulation(self):
-        return DiffractionSimulation(np.array([[0, 0, 0],[1, 2, 3], [3, 4, 5]]))
+        return DiffractionSimulation(np.array([[0, 0, 0], [1, 2, 3], [3, 4, 5]]))
 
     @pytest.fixture
     def diffraction_simulation_calibrated(self):
-        return DiffractionSimulation(np.array([[0, 0, 0],[1, 2, 3], [3, 4, 5]]), calibration=0.5)
+        return DiffractionSimulation(
+            np.array([[0, 0, 0], [1, 2, 3], [3, 4, 5]]), calibration=0.5
+        )
 
     @pytest.mark.xfail(raises=ValueError)
     def test_failed_initialization(self):
-        DiffractionSimulation(np.array([[0, 0, 0],[1, 2, 3], [3, 4, 5]]), indices=np.array([1, 2, 3]))
+        DiffractionSimulation(
+            np.array([[0, 0, 0], [1, 2, 3], [3, 4, 5]]), indices=np.array([1, 2, 3])
+        )
 
     def test_init(self, diffraction_simulation):
-        assert np.allclose(diffraction_simulation.coordinates, np.array([[1, 2, 3], [3, 4, 5]]))
+        assert np.allclose(
+            diffraction_simulation.coordinates, np.array([[1, 2, 3], [3, 4, 5]])
+        )
         assert np.isnan(diffraction_simulation.indices).all()
         assert diffraction_simulation.indices.shape == (2, 3)
         assert np.isnan(diffraction_simulation.intensities).all()
@@ -120,23 +126,34 @@ def test_fail_get_item(self, diffraction_simulation):
 
     def test_addition(self, diffraction_simulation):
         sim = diffraction_simulation + diffraction_simulation
-        assert np.allclose(sim.coordinates, np.array([[1, 2, 3],
-                                                      [3, 4, 5],
-                                                      [1, 2, 3],
-                                                      [3, 4, 5],
-                                                      ]))
+        assert np.allclose(
+            sim.coordinates,
+            np.array(
+                [
+                    [1, 2, 3],
+                    [3, 4, 5],
+                    [1, 2, 3],
+                    [3, 4, 5],
+                ]
+            ),
+        )
 
     def test_extend(self, diffraction_simulation):
         diffraction_simulation.extend(diffraction_simulation)
-        assert np.allclose(diffraction_simulation.coordinates, np.array([
-                                                      [1, 2, 3],
-                                                      [3, 4, 5],
-                                                      [1, 2, 3],
-                                                      [3, 4, 5],
-                                                      ]))
+        assert np.allclose(
+            diffraction_simulation.coordinates,
+            np.array(
+                [
+                    [1, 2, 3],
+                    [3, 4, 5],
+                    [1, 2, 3],
+                    [3, 4, 5],
+                ]
+            ),
+        )
 
     def test_indices_setter_getter(self, diffraction_simulation):
-        indices = np.array([[1, 2, 3],[2, 3,4],[3, 4, 5]])
+        indices = np.array([[1, 2, 3], [2, 3, 4], [3, 4, 5]])
         diffraction_simulation.indices = indices[-1]
         assert np.allclose(diffraction_simulation.indices, indices[-1])
         diffraction_simulation.with_direct_beam = True
@@ -144,7 +161,7 @@ def test_indices_setter_getter(self, diffraction_simulation):
         assert np.allclose(diffraction_simulation.indices, indices)
 
     def test_coordinates_setter(self, diffraction_simulation):
-        coords = np.array([[1, 2, 3],[2, 3,4],[3, 4, 5]])
+        coords = np.array([[1, 2, 3], [2, 3, 4], [3, 4, 5]])
         diffraction_simulation.coordinates = coords[-1]
         assert np.allclose(diffraction_simulation.coordinates, coords[-1])
         diffraction_simulation.with_direct_beam = True
@@ -189,11 +206,10 @@ def test_calibration(self, diffraction_simulation, calibration, expected):
             (np.array([[-1, 0, 0], [1, 0, 0]]), False, np.array([True, True])),
         ],
     )
-    def test_direct_beam_mask(
-        self, coordinates, with_direct_beam, expected
-    ):
-        diffraction_simulation = DiffractionSimulation(coordinates,
-                                                       with_direct_beam=with_direct_beam)
+    def test_direct_beam_mask(self, coordinates, with_direct_beam, expected):
+        diffraction_simulation = DiffractionSimulation(
+            coordinates, with_direct_beam=with_direct_beam
+        )
         diffraction_simulation.with_direct_beam = with_direct_beam
         mask = diffraction_simulation.direct_beam_mask
         assert np.all(mask == expected)
@@ -218,7 +234,7 @@ def test_direct_beam_mask(
                 None,
                 (0.0, 0.0),
                 None,
-                marks=pytest.mark.xfail(raises=Exception)
+                marks=pytest.mark.xfail(raises=Exception),
             ),
         ],
     )
@@ -232,24 +248,28 @@ def test_calibrated_coordinates(
         diffraction_simulation = DiffractionSimulation(coordinates)
         diffraction_simulation.calibration = calibration
         diffraction_simulation.offset = offset
-        assert np.allclose(diffraction_simulation.calibrated_coordinates,
-                           expected)
+        assert np.allclose(diffraction_simulation.calibrated_coordinates, expected)
 
     @pytest.mark.parametrize(
         "units, expected",
         [
-            ("real", np.array([[-2,1,3],[-4,3,5]])),
-            ("pixel", np.array([[-4,2],[-8,6]])),
+            ("real", np.array([[-2, 1, 3], [-4, 3, 5]])),
+            ("pixel", np.array([[-4, 2], [-8, 6]])),
         ],
     )
-    def test_transform_coordinates(self, diffraction_simulation_calibrated, units, expected):
-        tc = diffraction_simulation_calibrated._get_transformed_coordinates(90, units=units)
+    def test_transform_coordinates(
+        self, diffraction_simulation_calibrated, units, expected
+    ):
+        tc = diffraction_simulation_calibrated._get_transformed_coordinates(
+            90, units=units
+        )
         assert np.allclose(tc, expected)
 
     def test_rotate_shift_coordinates(self, diffraction_simulation):
         diffraction_simulation.rotate_shift_coordinates(90)
-        assert np.allclose(diffraction_simulation.coordinates,
-                           np.array([[-2,1,3],[-4,3,5]]))
+        assert np.allclose(
+            diffraction_simulation.coordinates, np.array([[-2, 1, 3], [-4, 3, 5]])
+        )
 
     def test_assertion_free_get_diffraction_pattern(self):
         short_sim = DiffractionSimulation(
@@ -272,22 +292,25 @@ def test_get_as_mask(self):
             intensities=np.ones(1),
             calibration=[1, 2],
         )
-        mask = short_sim.get_as_mask((20, 10),
-                                     radius_function=np.sqrt,
-                                     )
+        mask = short_sim.get_as_mask(
+            (20, 10),
+            radius_function=np.sqrt,
+        )
         assert mask.shape[0] == 20
         assert mask.shape[1] == 10
 
-    @pytest.mark.parametrize("units_in",['pixel','real'])
-    def test_plot_method(self,units_in):
+    @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],
-                                    [-2, 3, 0],
-                                    [2.1, 3.4, 0],]),
-            indices=np.array([[-2, 3, 4],
-                              [ 2,-6, 1],
-                              [ 0, 0, 0]]),
-            intensities=np.array([3., 5., 2.]),
+            coordinates=np.asarray(
+                [
+                    [0.3, 1.2, 0],
+                    [-2, 3, 0],
+                    [2.1, 3.4, 0],
+                ]
+            ),
+            indices=np.array([[-2, 3, 4], [2, -6, 1], [0, 0, 0]]),
+            intensities=np.array([3.0, 5.0, 2.0]),
             calibration=[1, 2],
         )
         ax, sp = short_sim.plot(units=units_in, show_labels=True)

From 29b0715c2a242618ca02f62decd505d13b79100d Mon Sep 17 00:00:00 2001
From: Niels Cautaerts <nielscautaerts@hotmail.com>
Date: Sat, 2 Oct 2021 00:34:15 +0200
Subject: [PATCH 6/6] implemented changes hokonanes

---
 diffsims/sims/diffraction_simulation.py       | 85 ++++++++-----------
 .../tests/sims/test_diffraction_simulation.py | 32 +++----
 2 files changed, 53 insertions(+), 64 deletions(-)

diff --git a/diffsims/sims/diffraction_simulation.py b/diffsims/sims/diffraction_simulation.py
index 7b87624d..642df076 100644
--- a/diffsims/sims/diffraction_simulation.py
+++ b/diffsims/sims/diffraction_simulation.py
@@ -20,10 +20,10 @@
 import numpy as np
 from diffsims.pattern.detector_functions import add_shot_and_point_spread
 from diffsims.utils import mask_utils
-from collections import Sequence
+import copy
 
 
-class DiffractionSimulation(Sequence):
+class DiffractionSimulation:
     """Holds the result of a kinematic diffraction pattern simulation.
 
     Parameters
@@ -75,13 +75,19 @@ def __init__(
                 "Coordinate, intensity, and indices lists must be of the correct and matching shape."
             )
         self.calibration = calibration
-        self.offset = offset
+        self.offset = np.array(offset)
         self.with_direct_beam = with_direct_beam
 
     def __len__(self):
         return self.coordinates.shape[0]
 
+    @property
+    def size(self):
+        return self.__len__()
+
     def __getitem__(self, sliced):
+        """Sliced is any valid numpy slice that does not change the number of
+        dimensions or number of columns"""
         coords = self.coordinates[sliced]
         inds = self.indices[sliced]
         ints = self.intensities[sliced]
@@ -101,26 +107,19 @@ def __getitem__(self, sliced):
             with_direct_beam=self.with_direct_beam,
         )
 
-    def _combine_data(self, other):
-        coords = np.concatenate([self._coordinates, other._coordinates], axis=0)
-        inds = np.concatenate([self._indices, other._indices], axis=0)
-        ints = np.concatenate([self._intensities, other._intensities], axis=0)
-        return coords, inds, ints
+    def deepcopy(self):
+        return copy.deepcopy(self)
 
     def __add__(self, other):
-        coords, inds, ints = self._combine_data(other)
-        return DiffractionSimulation(
-            coords,
-            indices=inds,
-            intensities=ints,
-            calibration=self.calibration,
-            offset=self.offset,
-            with_direct_beam=self.with_direct_beam,
-        )
+        new = self.deepcopy()
+        new.extend(other)
+        return new
 
     def extend(self, other):
         """Add the diffraction spots from another DiffractionSimulation"""
-        coords, inds, ints = self._combine_data(other)
+        coords = np.concatenate([self._coordinates, other._coordinates], axis=0)
+        inds = np.concatenate([self._indices, other._indices], axis=0)
+        ints = np.concatenate([self._intensities, other._intensities], axis=0)
         self._coordinates = coords
         self._indices = inds
         self._ints = ints
@@ -137,9 +136,7 @@ def indices(self, indices):
     def calibrated_coordinates(self):
         """ndarray : Coordinates converted into pixel space."""
         if self.calibration is not None:
-            return (self.coordinates[:, :2] + np.array(self.offset)) / np.array(
-                self.calibration
-            )
+            return (self.coordinates[:, :2] + self.offset) / self.calibration
         else:
             raise Exception("Pixel calibration is not set!")
 
@@ -152,18 +149,18 @@ def calibration(self):
     @calibration.setter
     def calibration(self, calibration):
         if calibration is None:
-            self._calibration = calibration
-            return
-        if np.all(np.equal(calibration, 0)):
+            pass
+        elif np.all(np.equal(calibration, 0)):
             raise ValueError("`calibration` cannot be zero.")
-        if isinstance(calibration, float) or isinstance(calibration, int):
-            self._calibration = (calibration, calibration)
+        elif isinstance(calibration, float) or isinstance(calibration, int):
+            calibration = np.array((calibration, calibration))
         elif len(calibration) == 2:
-            self._calibration = calibration
+            calibration = np.array(calibration)
         else:
             raise ValueError(
                 "`calibration` must be a float or length-2" "tuple of floats."
             )
+        self._calibration = calibration
 
     @property
     def direct_beam_mask(self):
@@ -196,26 +193,7 @@ def intensities(self, intensities):
     def _get_transformed_coordinates(
         self, angle, center=(0, 0), mirrored=False, units="real"
     ):
-        """
-        Get the coordinates of diffraction spots rotated, flipped or shifted
-        in-plane
-
-        Parameters
-        ----------
-        angle: float
-            Angle in degrees
-        center: 2-tuple of floats
-            Center of the patterns
-        mirrored: bool
-            Mirror accross the x axis
-        units: str
-            "real" or "pixel"
-
-        Returns
-        -------
-        coords_transformed: np.ndarray of shape (N, 3) or (N, 2)
-            transformed coordinates
-        """
+        """Translate, rotate or mirror the pattern spot coordinates"""
         if units == "real":
             coords_transformed = self.coordinates.copy()
         else:
@@ -231,7 +209,18 @@ def _get_transformed_coordinates(
         return coords_transformed
 
     def rotate_shift_coordinates(self, angle, center=(0, 0), mirrored=False):
-        """Translate, rotate or mirror the pattern"""
+        """
+        Rotate, flip or shift patterns in-plane
+
+        Parameters
+        ----------
+        angle: float
+            In plane rotation angle in degrees
+        center: 2-tuple of floats
+            Center coordinate of the patterns
+        mirrored: bool
+            Mirror across the x axis
+        """
         coords_new = self._get_transformed_coordinates(
             angle, center, mirrored, units="real"
         )
diff --git a/diffsims/tests/sims/test_diffraction_simulation.py b/diffsims/tests/sims/test_diffraction_simulation.py
index b04bbf95..d21bf7e5 100644
--- a/diffsims/tests/sims/test_diffraction_simulation.py
+++ b/diffsims/tests/sims/test_diffraction_simulation.py
@@ -22,13 +22,13 @@
 from diffsims.sims.diffraction_simulation import ProfileSimulation
 
 
-@pytest.mark.xfail(raises=ValueError)
 def test_wrong_calibration_setting():
-    DiffractionSimulation(
-        coordinates=np.asarray([[0.3, 1.2, 0]]),
-        intensities=np.ones(1),
-        calibration=[1, 2, 5],
-    )
+    with pytest.raises(ValueError, match="must be a float"):
+        DiffractionSimulation(
+            coordinates=np.asarray([[0.3, 1.2, 0]]),
+            intensities=np.ones(1),
+            calibration=[1, 2, 5],
+        )
 
 
 @pytest.fixture
@@ -96,11 +96,11 @@ def diffraction_simulation_calibrated(self):
             np.array([[0, 0, 0], [1, 2, 3], [3, 4, 5]]), calibration=0.5
         )
 
-    @pytest.mark.xfail(raises=ValueError)
     def test_failed_initialization(self):
-        DiffractionSimulation(
-            np.array([[0, 0, 0], [1, 2, 3], [3, 4, 5]]), indices=np.array([1, 2, 3])
-        )
+        with pytest.raises(ValueError, match="Coordinate"):
+            DiffractionSimulation(
+                np.array([[0, 0, 0], [1, 2, 3], [3, 4, 5]]), indices=np.array([1, 2, 3])
+            )
 
     def test_init(self, diffraction_simulation):
         assert np.allclose(
@@ -120,9 +120,9 @@ def test_get_item(self, diffraction_simulation):
         assert diffraction_simulation[1].coordinates.shape == (1, 3)
         assert diffraction_simulation[0:2].coordinates.shape == (2, 3)
 
-    @pytest.mark.xfail(raises=ValueError)
     def test_fail_get_item(self, diffraction_simulation):
-        diffraction_simulation[None]
+        with pytest.raises(ValueError, match="Invalid"):
+            _ = diffraction_simulation[None]
 
     def test_addition(self, diffraction_simulation):
         sim = diffraction_simulation + diffraction_simulation
@@ -179,16 +179,16 @@ def test_intensities_setter(self, diffraction_simulation):
     @pytest.mark.parametrize(
         "calibration, expected",
         [
-            (5.0, (5.0, 5.0)),
+            (5.0, np.array((5.0, 5.0))),
             pytest.param(0, (0, 0), marks=pytest.mark.xfail(raises=ValueError)),
             pytest.param((0, 0), (0, 0), marks=pytest.mark.xfail(raises=ValueError)),
-            ((1.5, 1.5), (1.5, 1.5)),
-            ((1.3, 1.5), (1.3, 1.5)),
+            ((1.5, 1.5), np.array((1.5, 1.5))),
+            ((1.3, 1.5), np.array((1.3, 1.5))),
         ],
     )
     def test_calibration(self, diffraction_simulation, calibration, expected):
         diffraction_simulation.calibration = calibration
-        assert diffraction_simulation.calibration == expected
+        assert np.allclose(diffraction_simulation.calibration, expected)
 
     @pytest.mark.parametrize(
         "coordinates, with_direct_beam, expected",