Exposed initial values

pyxem/signals/beam_shift.py

@@ -43,7 +43,7 @@ def make_linear_plane(self, **kwargs):
         self.data = s_linear_plane.data
         self.events.data_changed.trigger(None)
 
-    def get_linear_plane(self, mask=None, fit_corners=None, constrain_magnitude=False):
+    def get_linear_plane(self, mask=None, fit_corners=None, initial_values=None, constrain_magnitude=False):
         """Fit linear planes to the beam shifts, and returns a BeamShift signal
         with the planes.
 
@@ -81,6 +81,13 @@ def get_linear_plane(self, mask=None, fit_corners=None, constrain_magnitude=Fals
         fit_corners : float, optional
             Make a mask so that the planes are fitted to the corners of the
             signal. This mush be set with a number, like 0.05 (5%) or 0.10 (10%).
+        initial_values : array of floats, optional
+            Initial guess for the plane parameters. Useful to vary if the plane fitting
+            does not give desirable results. The parameters are as follows:
+            [d/dx, d/dy, x_0, d/dx, d/dy, y_0]
+            where the first three entries are for the x-shift, being in order the
+            step in x, step in y and the initial value at (0, 0). Similarly for the
+            last three entries for the y-shift.
         constrain_magnitude : bool, optional
             Fits the linear planes so there are deflection with constant magnitude.
             In the presence of electromagnetic fields in the sample area, least squares 
@@ -128,7 +135,7 @@ def get_linear_plane(self, mask=None, fit_corners=None, constrain_magnitude=Fals
             if mask.dtype != bool:
                 raise ValueError("mask needs to be an array of bools")
         if constrain_magnitude:
-            plane_image = bst._get_linear_plane_by_minimizing_magnitude_variance(self, mask=mask)
+            plane_image = bst._get_linear_plane_by_minimizing_magnitude_variance(self, mask=mask, initial_values=initial_values)
         else:
             plane_image_x = bst._get_linear_plane_from_signal2d(s_shift_x, mask=mask)
             plane_image_y = bst._get_linear_plane_from_signal2d(s_shift_y, mask=mask)

pyxem/tests/signals/test_beam_shift.py

@@ -65,6 +65,28 @@ def test_mask(self):
         s_lp = s.get_linear_plane(mask=s_mask)
         assert s_lp.data == approx(s_orig.data, abs=1e-7)
 
+    def test_constrain_magnitude(self):
+
+        p = [0.5]*6 # Plane parameters
+        x, y = np.meshgrid(np.arange(256), np.arange(256))
+        base_plane_x = p[0] * x + p[1] * y + p[2]
+        base_plane_y = p[3] * x + p[4] * y + p[5]
+
+        base_plane = np.stack((base_plane_x,base_plane_y)).T
+        data = base_plane.copy()
+
+        shifts = np.zeros_like(data)
+        shifts[:128,128:] =(10,10)
+        shifts[:128,:128] =(-10,-10)
+        shifts[128:,128:] =(-10,10)
+        shifts[128:,:128] =(10,10)
+        data += shifts
+
+        s = BeamShift(data)
+
+        s_lp = s.get_linear_plane(constrain_magnitude=True)
+        assert s_lp.data == approx(base_plane.data, abs=1e-7)
+
     def test_lazy_input_error(self):
         s = LazyBeamShift(da.zeros((50, 40, 2)))
         with pytest.raises(ValueError):