Merge pull request Quasars#31 from ngergihun/masking

Masking

Author: ngergihun

pySNOM/images.py

@@ -189,36 +189,57 @@ def type_from_channelname(channelname):
 
 
 class Transformation:
+
     def transform(self, data):
         raise NotImplementedError()
 
 
-class LineLevel(Transformation):
+class MaskedTransformation(Transformation):
+
+    def calculate(self, data, mask=None):
+        raise NotImplementedError()
+    
+    def correct(self, data, correction):
+        """ Applies the calculated corrections to the data """
+        if self.datatype == DataTypes.Amplitude:
+            return data / correction
+        else:
+            return data - correction
+
+    def transform(self, data, mask=None):
+        """ Calculates and applies the corrections to the data taking into account the mask if given """
+        correction = self.calculate(data, mask=mask)
+        return self.correct(data, correction)
+
+
+class LineLevel(MaskedTransformation):
     def __init__(self, method="median", datatype=DataTypes.Phase):
         self.method = method
         self.datatype = datatype
 
-    def transform(self, data):
-        if self.method == "median":
+    def calculate(self, data, mask = None):
+        if mask is not None:
+            data = mask*data
+
+        if self.method == 'median':
             norm = np.nanmedian(data, axis=1, keepdims=True)
         elif self.method == "mean":
             norm = np.nanmean(data, axis=1, keepdims=True)
         elif self.method == "difference":
             if self.datatype == DataTypes.Amplitude:
                 norm = np.nanmedian(data[1:] / data[:-1], axis=1, keepdims=True)
+                norm = np.append(norm,1)
             else:
                 norm = np.nanmedian(data[1:] - data[:-1], axis=1, keepdims=True)
-            data = data[:-1]  # difference does not make sense for the last row
+                norm = np.append(norm,0)  # difference does not make sense for the last row
+            norm = np.reshape(norm, (norm.size,1))
         else:
             if self.datatype == DataTypes.Amplitude:
                 norm = 1
             else:
                 norm = 0
 
-        if self.datatype == DataTypes.Amplitude:
-            return data / norm
-        else:
-            return data - norm
+        return norm
 
 
 class RotatePhase(Transformation):
@@ -248,43 +269,41 @@ def transform(self, data):
             )
 
 
-class SimpleNormalize(Transformation):
+class SimpleNormalize(MaskedTransformation):
     def __init__(self, method="median", value=1.0, datatype=DataTypes.Phase):
         self.method = method
         self.value = value
         self.datatype = datatype
 
-    def transform(self, data):
+    def calculate(self, data, mask=None):
+        """ Calculates and returns the image corrections using mask (if given) without applying it to the data"""
+        if mask is not None:
+            data = mask*data
+
         match self.method:
-            case "median":
-                if self.datatype == DataTypes.Amplitude:
-                    return data / np.nanmedian(data)
-                else:
-                    return data - np.nanmedian(data)
-            case "mean":
-                if self.datatype == DataTypes.Amplitude:
-                    return data / np.nanmean(data)
-                else:
-                    return data - np.nanmean(data)
-            case "manual":
-                if self.datatype == DataTypes.Amplitude:
-                    return data / self.value
-                else:
-                    return data - self.value
-            case "min":
-                if self.datatype == DataTypes.Amplitude:
-                    return data / np.nanmin(data)
-                else:
-                    return data - np.nanmin(data)
-
-
-class BackgroundPolyFit(Transformation):
+            case 'median':
+                norm = np.nanmedian(data)
+            case 'mean':
+                norm = np.nanmean(data)
+            case 'manual':
+                norm = self.value
+            case 'min':
+                norm = np.nanmin(data)
+
+        return norm
+                
+
+class BackgroundPolyFit(MaskedTransformation):
     def __init__(self, xorder=1, yorder=1, datatype=DataTypes.Phase):
         self.xorder = int(xorder)
         self.yorder = int(yorder)
         self.datatype = datatype
 
-    def transform(self, data):
+    def calculate(self, data, mask = None):
+        """ Calculates and returns the fitted polynomial background using mask (if given) without applying it to the data"""
+        if mask is not None:
+            data = mask*data
+
         Z = copy.deepcopy(data)
         x = list(range(0, Z.shape[1]))
         y = list(range(0, Z.shape[0]))
@@ -310,18 +329,21 @@ def get_basis(x, y, max_order_x=1, max_order_y=1):
             A = np.vstack(basis).T
             c, r, rank, s = np.linalg.lstsq(A, b, rcond=None)
 
-            background = np.sum(
-                c[:, None, None]
-                * np.array(get_basis(X, Y, self.xorder, self.yorder)).reshape(
-                    len(basis), *X.shape
-                ),
-                axis=0,
-            )
+            background = np.sum(c[:, None, None] * np.array(get_basis(X, Y, self.xorder, self.yorder)).reshape(len(basis), *X.shape),axis=0)
+            
         except ValueError:
             background = np.ones(np.shape(data))
             print("X and Y order must be integer!")
 
-        if self.datatype == DataTypes["Amplitude"]:
-            return Z / background, background
-        else:
-            return Z - background, background
+        return background
+
+        
+# TODO: Helper functions to create masks or turn other types of masks into 1/Nan mask
+def mask_from_booleans(bool_mask, bad_values = False):
+    """ Turn a boolean array to an array conatining nans and ones"""
+    mshape = np.shape(bool_mask)
+    return np.where(bool_mask==bad_values,np.nan*np.ones(mshape),np.ones(mshape))
+
+def mask_from_datacondition(condition):
+    mshape = np.shape(condition)
+    return np.where(condition,np.nan*np.ones(mshape),np.ones(mshape))

pySNOM/tests/test_transform.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from pySNOM.images import LineLevel, DataTypes
+from pySNOM.images import LineLevel, BackgroundPolyFit, SimpleNormalize, DataTypes, mask_from_datacondition
 
 
 class TestLineLevel(unittest.TestCase):
@@ -46,13 +46,131 @@ def test_difference(self):
         d = np.arange(12).reshape(3, -1)[:, [0, 1, 3]]
         l = LineLevel(method="difference", datatype=DataTypes.Phase)
         out = l.transform(d)
-        np.testing.assert_almost_equal(out, [[-4.0, -3.0, -1.0], [0.0, 1.0, 3.0]])
+        np.testing.assert_almost_equal(out, [[-4., -3., -1.], [0., 1., 3.], [8., 9., 11.]])
         l = LineLevel(method="difference", datatype=DataTypes.Amplitude)
         out = l.transform(d)
         np.testing.assert_almost_equal(
-            out, [[0.0, 0.2, 0.6], [2.2222222, 2.7777778, 3.8888889]]
+            out, [[0., 0.2, 0.6], [ 2.2222222,  2.7777778,  3.8888889], [8., 9., 11.]]
         )
 
+    def test_masking_mean(self):
+        d = np.zeros([8,10])
+        d[2:6,3:7] = 1
+        mask = np.ones([8,10])
+        mask[2:6,3:7] = np.nan
+
+        l = LineLevel(method="mean", datatype=DataTypes.Phase)
+
+        out = l.transform(d)
+        np.testing.assert_almost_equal(out[5,0], -0.4)
+        out = l.transform(d,mask=mask)
+        np.testing.assert_almost_equal(out[5,0], 0.0)
+
+    def test_masking_median(self):
+        d = np.zeros([8,10])
+        d[2:6,2:9] = 1
+        mask = np.ones([8,10])
+        mask[2:6,2:9] = np.nan
+
+        l = LineLevel(method="median", datatype=DataTypes.Phase)
+        
+        out = l.transform(d)
+        np.testing.assert_almost_equal(out[5,0], -1.0)
+        out = l.transform(d,mask=mask)
+        np.testing.assert_almost_equal(out[5,0], 0.0)
+
+    def test_difference_masking(self):
+        d = np.zeros([8,10])
+        d[2:6,2:9] = 1
+        mask = np.ones([8,10])
+        mask[2:6,2:9] = np.nan
+
+        l = LineLevel(method="difference", datatype=DataTypes.Phase)
+        out = l.transform(d)
+        np.testing.assert_almost_equal(out[5,0], 1.0)
+        np.testing.assert_almost_equal(out[1,0], -1.0)
+        np.testing.assert_almost_equal(out[5,2], 2.0)
+
+class TestBackgroundPolyFit(unittest.TestCase):
+    def test_withmask(self):
+        d = np.ones([10,10])
+        d[4:8,4:8] = 10
+        mask = np.ones([10,10])
+        mask[4:8,4:8] = np.nan
+
+        t = BackgroundPolyFit(xorder=1,yorder=1,datatype=DataTypes.Phase)
+        out = t.transform(d,mask=mask)
+        np.testing.assert_almost_equal(out[0,0], 0.0)
+        np.testing.assert_almost_equal(out[9,9], 0.0)
+
+        t = BackgroundPolyFit(xorder=1,yorder=1,datatype=DataTypes.Amplitude)
+        out = t.transform(d,mask=mask)
+        np.testing.assert_almost_equal(out[0,0], 1.0)
+        np.testing.assert_almost_equal(out[9,9], 1.0)
+
+    def test_withoutmask(self):
+        d = np.ones([10,10])
+        d[4:8,4:8] = 10
+        mask = np.ones([10,10])
+        mask[4:8,4:8] = np.nan
+
+        t = BackgroundPolyFit(xorder=1,yorder=1,datatype=DataTypes.Phase)
+        out = t.transform(d)
+        np.testing.assert_almost_equal(out[0,0], -0.2975206611570238)
+        np.testing.assert_almost_equal(out[9,9], -3.439338842975202)
+
+        t = BackgroundPolyFit(xorder=1,yorder=1,datatype=DataTypes.Amplitude)
+        out = t.transform(d)
+        np.testing.assert_almost_equal(out[0,0], 0.7707006369426758)
+        np.testing.assert_almost_equal(out[9,9], 0.22525876833718098)
+
+class TestHelperFunction(unittest.TestCase):
+
+    def test_mask_from_condition(self):
+        d = np.ones([2,2])
+        d[0,0] = 0
+        out = mask_from_datacondition(d<1)
+
+        np.testing.assert_equal(out[0,0], np.nan)
+
+class TestSimpleNormalize(unittest.TestCase):
+
+    def test_median(self):
+        d = np.zeros([8,10])
+        d[2:9,1:9] = 1
+        mask = np.ones([8,10])
+        mask[2:9,1:9] = np.nan
+
+        l = SimpleNormalize(method="median", datatype=DataTypes.Phase)
+        
+        out = l.transform(d)
+        np.testing.assert_almost_equal(out[0,0], -1.0)
+        out = l.transform(d,mask=mask)
+        np.testing.assert_almost_equal(out[0,0], 0.0)
+
+    def test_mean(self):
+        d = np.zeros([8,10])
+        d[2:9,1:9] = 1
+        mask = np.ones([8,10])
+        mask[2:9,1:9] = np.nan
+
+        l = SimpleNormalize(method="mean", datatype=DataTypes.Phase)
+        
+        out = l.transform(d)
+        np.testing.assert_almost_equal(out[0,0], -0.6)
+        out = l.transform(d,mask=mask)
+        np.testing.assert_almost_equal(out[0,0], 0.0)
+
+    def test_min(self):
+        d=np.asarray([1.0, 2.0, 3.0])
+        mask = np.asarray([np.nan, 1, 1])
+
+        l = SimpleNormalize(method="min", datatype=DataTypes.Phase)
+
+        out = l.transform(d)
+        np.testing.assert_almost_equal(out, [0.0, 1.0, 2.0])
+        out = l.transform(d,mask=mask)
+        np.testing.assert_almost_equal(out, [-1.0,0.0,1.0])
 
 if __name__ == "__main__":
     unittest.main()