Fixed Field.clip()

src/empyre/fields/field.py

@@ -475,7 +475,7 @@ class Field(NDArrayOperatorsMixin):
             raise AssertionError('Can only handle 3D or 2D cases (see documentation for specifics)!')
 
     def clip(self, vmin=None, vmax=None, sigma=None, mask=None):
-        """Clip (limit) the values in an array. For vector fields, this will take the amplitude into account
+        """Clip (limit) the values in an array. For vector fields, this will take the amplitude into account.
 
         Parameters
         ----------
@@ -517,9 +517,8 @@ class Field(NDArrayOperatorsMixin):
         if vmax is None:
             vmax = np.nanmax(amp_masked)
         if self.vector:  # Vector fields need to scale components according to masked amplitude
-            mask_vec = np.logical_and(mask, amp <= vmax)  # Only vmax is important!
-            data = amp / np.where(mask_vec, 1, amp)  # TODO: needs testing!
-            # TODO: Test np.clip(field) (also ufunc?? think not...) seems to work, but what about others?
+            mask_vec = (amp <= vmax)[..., None]  # Only vmax is important for vectors! mask_vec broadcast to components!
+            data = np.where(mask_vec, self.data, vmax * self.data/amp[..., None])  # Scale outliers to vmax!
         else:  # For scalar fields, just delegate to the numpy function:
             data = np.clip(self.data, vmin, vmax)
         return Field(data, self.scale, self.vector)