diff --git a/diffsims/sims/diffraction_simulation.py b/diffsims/sims/diffraction_simulation.py index cef91ea6..642df076 100644 --- a/diffsims/sims/diffraction_simulation.py +++ b/diffsims/sims/diffraction_simulation.py @@ -20,6 +20,7 @@ import numpy as np from diffsims.pattern.detector_functions import add_shot_and_point_spread from diffsims.utils import mask_utils +import copy class DiffractionSimulation: @@ -44,35 +45,100 @@ class DiffractionSimulation: def __init__( self, - coordinates=None, + 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. """ - self._coordinates = None - self.coordinates = coordinates - self.indices = indices - self._intensities = None - self.intensities = intensities - self._calibration = (1.0, 1.0) + if coordinates.ndim == 1: + 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 + ): + 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." + ) 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] + 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, + calibration=self.calibration, + offset=self.offset, + with_direct_beam=self.with_direct_beam, + ) + + def deepcopy(self): + return copy.deepcopy(self) + + def __add__(self, other): + new = self.deepcopy() + new.extend(other) + return new + + def extend(self, other): + """Add the diffraction spots from another DiffractionSimulation""" + 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 + + @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 + if self.calibration is not None: + return (self.coordinates[:, :2] + self.offset) / self.calibration + else: + raise Exception("Pixel calibration is not set!") @property def calibration(self): @@ -82,16 +148,19 @@ def calibration(self): @calibration.setter def calibration(self, calibration): - if np.all(np.equal(calibration, 0)): + if calibration is None: + 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): @@ -106,29 +175,69 @@ 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, - in_plane_angle=0, mirrored=False, - *args, **kwargs): + def _get_transformed_coordinates( + self, angle, center=(0, 0), mirrored=False, units="real" + ): + """Translate, rotate or mirror the pattern spot 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): + """ + 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" + ) + self.coordinates = coords_new + + 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 @@ -155,38 +264,55 @@ 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 - if direct_beam_position is not None: - cx, cy = direct_beam_position - point_coordinates_shifted = self.calibrated_coordinates[:, :-1].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 + 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=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, 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 ------- @@ -199,36 +325,55 @@ 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 + 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" ) - - 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]) + 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, - 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 ---------- 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 @@ -252,56 +397,57 @@ def plot(self, size_factor=1, units="real", show_labels=False, ----- spot size scales with the square root of the intensity. """ + if direct_beam_position is None: + direct_beam_position = (0, 0) 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, 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: 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() - 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 cc26424f..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 @@ -88,27 +88,107 @@ 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]])) + + @pytest.fixture + def diffraction_simulation_calibrated(self): + return DiffractionSimulation( + np.array([[0, 0, 0], [1, 2, 3], [3, 4, 5]]), calibration=0.5 + ) + + def test_failed_initialization(self): + 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 diffraction_simulation.coordinates is None - assert diffraction_simulation.indices is None - assert diffraction_simulation.intensities is None - assert diffraction_simulation.calibration == (1.0, 1.0) + 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() + assert diffraction_simulation.intensities.shape == (2,) + assert diffraction_simulation.calibration is None + assert len(diffraction_simulation) == 2 + + 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[0:2].coordinates.shape == (2, 3) + + def test_fail_get_item(self, diffraction_simulation): + with pytest.raises(ValueError, match="Invalid"): + _ = 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) + + 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", [ - (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", @@ -126,10 +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, diffraction_simulation, coordinates, with_direct_beam, expected - ): - diffraction_simulation.coordinates = coordinates + 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) @@ -141,30 +221,62 @@ 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]]), + ), + 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( 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) + @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( @@ -172,7 +284,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( @@ -180,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, - direct_beam_position=(10, 6) - ) + 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.1, 3.4, 0]]), - indices=np.array([[-2, 3, 4], - [ 2,-6, 1], - [ 0, 0, 0]]), - intensities=np.array([3., 5.]), + 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)