Added tests

sivborg · Nov 14, 2023 · 5bcabe8 · 5bcabe8
1 parent 0572484
commit 5bcabe8
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/pyxem/tests/signals/test_beam_shift.py b/pyxem/tests/signals/test_beam_shift.py
@@ -62,6 +62,9 @@ def test_simple(self):
         s_orig = s.deepcopy()
         s.make_linear_plane()
         np.testing.assert_almost_equal(s.data, s_orig, decimal=7)
+        s = s_orig.deepcopy()
+        s.make_linear_plane(use_joint_distance=True)
+        np.testing.assert_almost_equal(s.data, s_orig, decimal=5)
 
     def test_mask(self):
         data_x, data_y = np.meshgrid(
@@ -78,6 +81,11 @@ def test_mask(self):
         s.data[45:50, 36:41, 1] = -100000
         s.make_linear_plane(mask=s_mask)
         np.testing.assert_almost_equal(s.data, s_orig, decimal=7)
+        s = s_orig.deepcopy()
+        s.data[45:50, 36:41, 0] = 100000
+        s.data[45:50, 36:41, 1] = -100000
+        s.make_linear_plane(mask=s_mask, use_joint_distance=True)
+        np.testing.assert_almost_equal(s.data, s_orig, decimal=5)
 
     def test_lazy_input_error(self):
         s = LazyBeamShift(da.zeros((50, 40, 2)))