Merge pull request #413 from pc494/bugfix-affine-transform-for-0.9

Replicates the patch from 0.8.1

pyxem/signals/electron_diffraction.py

@@ -26,10 +26,10 @@
 from pyxem.signals.diffraction_vectors import DiffractionVectors
 
 from pyxem.utils.expt_utils import _index_coords, _cart2polar, _polar2cart, \
-    radial_average, gain_normalise, remove_dead, affine_transformation, \
+    radial_average, gain_normalise, remove_dead,\
     regional_filter, subtract_background_dog, subtract_background_median, \
     subtract_reference, circular_mask, find_beam_offset_cross_correlation, \
-    peaks_as_gvectors
+    peaks_as_gvectors,convert_affine_to_transform,apply_transformation
 
 from pyxem.utils.peakfinders2D import find_peaks_zaefferer, find_peaks_stat, \
     find_peaks_dog, find_peaks_log, find_peaks_xc
@@ -255,14 +255,21 @@ def get_direct_beam_mask(self, radius):
     def apply_affine_transformation(self,
                                     D,
                                     order=3,
+                                    keep_dtype=False,
                                     inplace=True,
                                     *args, **kwargs):
         """Correct geometric distortion by applying an affine transformation.
 
         Parameters
         ----------
-        D : array
-            3x3 np.array specifying the affine transform to be applied.
+        D : array or Signal2D of arrays
+            3x3 np.array (or Signal2D thereof) specifying the affine transform
+            to be applied.
+        order : 1,2,3,4 or 5
+            The order of interpolation on the transform. Default is 3.
+        keep_dtype : bool
+            If True dtype of returned ElectronDiffraction Signal is that of
+            the input, if False, casting to higher precision may occur.
         inplace : bool
             If True (default), this signal is overwritten. Otherwise, returns a
             new signal.
@@ -277,24 +284,17 @@ def apply_affine_transformation(self,
             diffraction patterns.
 
         """
-        # Account for the transformation center not being (0,0)
-        shape = self.axes_manager.signal_shape
-        shift_x = (shape[1] - 1) / 2
-        shift_y = (shape[0] - 1) / 2
-
-        tf_shift = tf.SimilarityTransform(translation=[-shift_x, -shift_y])
-        tf_shift_inv = tf.SimilarityTransform(translation=[shift_x, shift_y])
-
-        # This defines the transform you want to perform
-        distortion = tf.AffineTransform(matrix=D)
 
-        # skimage transforms can be added like this, does matrix multiplication,
-        # hence the need for the brackets. (Note tf.warp takes the inverse)
-        transformation = (tf_shift + (distortion + tf_shift_inv)).inverse
+        shape = self.axes_manager.signal_shape
+        if type(D) == np.ndarray:
+            transformation = convert_affine_to_transform(D,shape)
+        else:
+            transformation = D.map(convert_affine_to_transform,shape=shape,inplace=False)
 
-        return self.map(affine_transformation,
+        return self.map(apply_transformation,
                         transformation=transformation,
                         order=order,
+                        keep_dtype=keep_dtype,
                         inplace=inplace,
                         *args, **kwargs)
 

pyxem/tests/test_signals/test_electron_diffraction.py

@@ -58,6 +58,23 @@ def test_apply_affine_transformation(self, diffraction_pattern):
                         [0., 0., 1.]]))
         assert isinstance(diffraction_pattern, ElectronDiffraction)
 
+    def test_apply_affine_transforms_paths(self,diffraction_pattern):
+        D=np.array([[1., 0.9, 0.],
+                    [1.1, 1., 0.],
+                    [0., 0., 1.]])
+        s = Signal2D(np.asarray([[D,D],[D,D]]))
+        static  = diffraction_pattern.apply_affine_transformation(D,inplace=False)
+        dynamic = diffraction_pattern.apply_affine_transformation(s,inplace=False)
+        assert np.allclose(static.data,dynamic.data,atol=1e-3)
+
+    def test_apply_affine_transformation_with_casting(self, diffraction_pattern):
+        diffraction_pattern.change_dtype('uint8')
+        transformed_dp = ElectronDiffraction(diffraction_pattern).apply_affine_transformation(
+            D=np.array([[1., 0., 0.],
+                        [0., 1., 0.],
+                        [0., 0., 1.2]]), order=2, keep_dtype=True, inplace=False)
+        assert transformed_dp.data.dtype == 'uint8'
+
     methods = ['average', 'nan']
 
     @pytest.mark.parametrize('method', methods)

pyxem/tests/test_utils/test_expt_utils.py

@@ -22,7 +22,7 @@
 
 from pyxem.signals.electron_diffraction import ElectronDiffraction
 from pyxem.utils.expt_utils import _index_coords, _cart2polar, _polar2cart, \
-    radial_average, gain_normalise, remove_dead, affine_transformation, \
+    radial_average, gain_normalise, remove_dead, apply_transformation, \
     regional_filter, subtract_background_dog, subtract_background_median, \
     subtract_reference, circular_mask, reference_circle, \
     find_beam_offset_cross_correlation, peaks_as_gvectors

pyxem/utils/expt_utils.py

@@ -198,18 +198,53 @@ def remove_dead(z, deadpixels, deadvalue="average", d=1):
 
     return z_bar
 
+def convert_affine_to_transform(D,shape):
+    """ Converts an affine transform on a diffraction pattern to a suitable
+    form for skimage.transform.warp()
 
-def affine_transformation(z, transformation, order, *args, **kwargs):
-    """Apply an affine transformation to a 2-dimensional array.
+    Parameters
+    ----------
+    D : np.array
+        Affine transform to be applied
+    shape : tuple
+        Shape tuple in form (y,x) for the diffraction pattern
+
+    Returns
+    -------
+    transformation : np.array
+        3x3 numpy array of the transformation to be applied.
+    
+    """    
+
+    shift_x = (shape[1] - 1) / 2
+    shift_y = (shape[0] - 1) / 2
+
+    tf_shift = tf.SimilarityTransform(translation=[-shift_x, -shift_y])
+    tf_shift_inv = tf.SimilarityTransform(translation=[shift_x, shift_y])
+
+    # This defines the transform you want to perform
+    distortion = tf.AffineTransform(matrix=D)
+
+    # skimage transforms can be added like this, does matrix multiplication,
+    # hence the need for the brackets. (Note tf.warp takes the inverse)
+    transformation = (tf_shift + (distortion + tf_shift_inv)).inverse
+
+    return transformation
+
+
+def apply_transformation(z, transformation,keep_dtype,order=1,*args, **kwargs):
+    """Apply a transformation to a 2-dimensional array.
 
     Parameters
     ----------
     z : np.array
         Array to be transformed
-    matrix : np.array
-        3x3 numpy array specifying the affine transformation to be applied.
+    transformation : np.array
+        3x3 numpy array specifying the transformation to be applied.
     order : int
         Interpolation order.
+    keep_dtype : bool
+        If True dtype of returned object is that of z
     *args :
         To be passed to skimage.warp
     **kwargs :
@@ -219,9 +254,19 @@ def affine_transformation(z, transformation, order, *args, **kwargs):
     -------
     trans : array
         Affine transformed diffraction pattern.
+
+    Notes
+    -----
+    Generally used in combination with pyxem.expt_utils.convert_affine_to_transform
     """
-    trans = tf.warp(z, transformation,
+    if keep_dtype == False:
+        trans = tf.warp(z, transformation,
                     order=order, *args, **kwargs)
+    if keep_dtype == True:
+            trans = tf.warp(z, transformation,
+                        order=order,preserve_range=True, *args, **kwargs)
+            trans = trans.astype(z.dtype)
+
     return trans
 
 

setup.py

@@ -48,11 +48,11 @@
     # adjust the tabbing
     install_requires=[
     	'scikit-image < 0.15', # See pyxem/pull/378
-        'matplotlib < 3.1.0' , # See pyxem/pull/403
-        'scikit-learn >= 0.19', # bug unknown
-    	'hyperspy >= 1.3',      # 1.2 fails, (NTU Workshop - May 2019) 
-        'transforms3d',        
-        'diffpy.structure >= 3.0.0' # First Python 3 support
+      'matplotlib < 3.1.0' , # See pyxem/pull/403
+      'scikit-learn >= 0.19', # bug unknown
+     	'hyperspy >= 1.3',      # 1.2 fails, (NTU Workshop - May 2019) 
+      'transforms3d',        
+      'diffpy.structure >= 3.0.0' # First Python 3 support
       ],                       
     package_data={
         "": ["LICENSE", "readme.rst",],