Merge pull request pyxem#949 from sivborg/color_manuallim

Expose magnitude limits to DPCSignal2D methods

CHANGELOG.rst

@@ -11,6 +11,10 @@ and this project adheres to `Semantic Versioning <https://semver.org/spec/v2.0.0
 Unreleased
 ==========
 
+Added
+-----
+- Added `magnitude_limits` to `DPCSignal2D` methods (#949)
+
 Fixed
 -----
 - Remove `ipywidgets` from requirements as it is not a dependency

pyxem/signals/differential_phase_contrast.py

@@ -260,7 +260,9 @@ def to_beamshift(self):
         s_beam_shift.set_signal_type("beam_shift")
         return s_beam_shift
 
-    def get_magnitude_signal(self, autolim=True, autolim_sigma=4):
+    def get_magnitude_signal(
+        self, autolim=True, autolim_sigma=4, magnitude_limits=None
+    ):
         """Get DPC magnitude image visualized as greyscale.
 
         Converts the x and y beam shifts into a magnitude map, showing the
@@ -272,6 +274,9 @@ def get_magnitude_signal(self, autolim=True, autolim_sigma=4):
         ----------
         autolim : bool, default True
         autolim_sigma : float, default 4
+        magnitude_limits : tuple of floats, default None
+            Manually sets the value limits for the magnitude signal.
+            For this, autolim needs to be False.
 
         Returns
         -------
@@ -292,11 +297,17 @@ def get_magnitude_signal(self, autolim=True, autolim_sigma=4):
         inav02 = np.abs(self.inav[0].data) ** 2
         inav12 = np.abs(self.inav[1].data) ** 2
         magnitude = np.sqrt(inav02 + inav12)
-        magnitude_limits = None
+
         if autolim:
+            if magnitude_limits is not None:
+                raise ValueError(
+                    "If autolim==True then `magnitude_limits` must be set to None"
+                )
+
             magnitude_limits = pst._get_limits_from_array(
                 magnitude, sigma=autolim_sigma
             )
+        if magnitude_limits is not None:
             np.clip(magnitude, magnitude_limits[0], magnitude_limits[1], out=magnitude)
 
         signal = Signal2D(magnitude)
@@ -481,7 +492,9 @@ def get_phase_signal(self, rotation=None):
         pst._copy_signal2d_axes_manager_metadata(self, signal_rgb)
         return signal_rgb
 
-    def get_color_signal(self, rotation=None, autolim=True, autolim_sigma=4):
+    def get_color_signal(
+        self, rotation=None, autolim=True, autolim_sigma=4, magnitude_limits=None
+    ):
         """Get DPC image visualized using continuous color scale.
 
         Converts the x and y beam shifts into an RGB array, showing the
@@ -496,6 +509,9 @@ def get_color_signal(self, rotation=None, autolim=True, autolim_sigma=4):
             scan direction and diffraction pattern rotation.
         autolim : bool, default True
         autolim_sigma : float, default 4
+        magnitude_limits : tuple of floats, default None
+            Manually sets the value limits for the color signal.
+            For this, autolim needs to be False.
 
         Returns
         -------
@@ -528,8 +544,12 @@ def get_color_signal(self, rotation=None, autolim=True, autolim_sigma=4):
         phase = np.arctan2(inav0, inav1) % (2 * np.pi)
         magnitude = np.sqrt(np.abs(inav0) ** 2 + np.abs(inav1) ** 2)
 
-        magnitude_limits = None
         if autolim:
+            if magnitude_limits is not None:
+                raise ValueError(
+                    "If autolim==True then `magnitude_limits` must be set to None"
+                )
+
             magnitude_limits = pst._get_limits_from_array(
                 magnitude, sigma=autolim_sigma
             )
@@ -552,6 +572,7 @@ def get_color_image_with_indicator(
         only_phase=False,
         autolim=True,
         autolim_sigma=4,
+        magnitude_limits=None,
         scalebar_size=None,
         ax=None,
         ax_indicator=None,
@@ -570,6 +591,9 @@ def get_color_image_with_indicator(
             If True, will only plot the phase.
         autolim : bool, default True
         autolim_sigma : float, default 4
+        magnitude_limits : tuple of floats, default None
+            Manually sets the value limits for the color signal.
+            For this, autolim needs to be False.
         scalebar_size : int, optional
         ax : Matplotlib subplot, optional
         ax_indicator : Matplotlib subplot, optional
@@ -607,7 +631,10 @@ def get_color_image_with_indicator(
             s = self.get_phase_signal(rotation=phase_rotation)
         else:
             s = self.get_color_signal(
-                rotation=phase_rotation, autolim=autolim, autolim_sigma=autolim_sigma
+                rotation=phase_rotation,
+                autolim=autolim,
+                autolim_sigma=autolim_sigma,
+                magnitude_limits=magnitude_limits,
             )
         s.change_dtype("uint16")
         s.change_dtype("float64")

pyxem/tests/signals/test_differential_phase_contrast.py

@@ -228,6 +228,18 @@ def test_get_color_signal(self):
         s_random = DPCSignal2D(data_random)
         s_random.get_color_signal()
         s_random.get_color_signal(rotation=45)
+        s_random.get_color_signal(autolim=False, magnitude_limits=(0, 30))
+
+    def test_get_color_signal_errors(self):
+        array_x, array_y = np.meshgrid(range(64), range(64))
+        data_tilt = np.swapaxes(
+            np.dstack((array_x + array_y, np.fliplr(array_x) + array_y)), 0, 2
+        ).astype("float64")
+        data_random = data_tilt + np.random.random(size=(2, 64, 64)) * 10
+        s_random = DPCSignal2D(data_random)
+
+        with pytest.raises(ValueError):
+            s_random.get_color_signal(autolim=True, magnitude_limits=(0, 30))
 
     def test_get_color_signal_zeros(self):
         s = DPCSignal2D(np.zeros((2, 100, 100)))
@@ -241,6 +253,12 @@ def test_get_magnitude_signal_zeros(self):
         s_magnitude = s.get_magnitude_signal()
         assert (s_magnitude.data == 0).all()
 
+    def test_get_magnitude_signal_errors(self):
+        s = DPCSignal2D(np.zeros((2, 100, 100)))
+
+        with pytest.raises(ValueError):
+            s.get_color_signal(autolim=True, magnitude_limits=(0, 30))
+
     def test_get_phase_signal(self):
         s = DPCSignal2D(np.zeros((2, 100, 100)))
         s.get_phase_signal()
@@ -257,6 +275,7 @@ def test_get_color_image_with_indicator(self):
             scalebar_size=10,
         )
         s.get_color_image_with_indicator(only_phase=True)
+        s.get_color_image_with_indicator(autolim=False, magnitude_limits=(0, 0.5))
         fig, ax = subplots()
         s.get_color_image_with_indicator(ax=ax)