Updated channel name recognition (images), new phase correction (spectra)

Author: ngergihun

pySNOM/images.py

@@ -1,5 +1,6 @@
 import numpy as np
 import copy
+import re
 from enum import Enum
 from pySNOM.defaults import Defaults
 from gwyfile.objects import GwyDataField
@@ -121,21 +122,26 @@ def data(self,value):
             raise ValueError("The provided data object does not contain 'data' attribute")
 
 def type_from_channelname(channelname):
-    if channelname[0] == 'O':
+    channel_strings = ['M(.?)A', 'M(.?)P', 'O(.?)A', 'O(.?)P', 'Z C', 'Z raw']
+    for pattern in channel_strings:
+        if re.search(pattern, channelname) is not None:
+            channel_name = re.search(pattern, channelname)[0]
+
+    if channel_name[0] == 'O':
         channeltype = ChannelTypes["Optical"]
-    elif 'M' in channelname:
+    elif 'M' in channel_name:
         channeltype = ChannelTypes["Mechanical"]
     else:
         channeltype = ChannelTypes["None"]
 
-    if 'Z' in channelname:
+    if 'Z' in channel_name:
         order = 0
     else:
-        order = int(channelname[1])
+        order = int(channel_name[1])
 
-    if channelname[2] == 'A':
+    if channel_name[2] == 'A':
         datatype = DataTypes["Amplitude"]
-    elif 'Z' in channelname:
+    elif 'Z' in channel_name:
         datatype = DataTypes["Topography"]
     else:
         datatype = DataTypes["Phase"]
@@ -224,6 +230,11 @@ def transform(self, data):
                     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):
 

pySNOM/spectra.py

@@ -145,6 +145,34 @@ def transform(self, spectrum, wnaxis):
 
         return np.angle(spectrum*np.exp(angles*complex(1j)))
 
+class ShiftPhaseToZero(Transformation):
+    """ 
+    Calculates and applies the phase rotation needed to get a flat, levelled phase spectrum. 
+    Two reference frequencies have to be provided.
+    """
+    def __init__(self, wavenumber1=1000.0, wavenumber2=2200.0):
+        self.wn1 = wavenumber1
+        self.wn2 = wavenumber2
+
+    def transform(self, spectrum, wnaxis):
+        if not np.iscomplex(spectrum).any():
+            spectrum = np.exp(spectrum*complex(1j))
+
+        wn1idx = np.argmin(abs(wnaxis - self.wn1))
+        wn2idx = np.argmin(abs(wnaxis - self.wn2))
+
+        theta1 = np.angle(spectrum[wn1idx])
+        wn1 = wnaxis[wn1idx]
+        spectrum = spectrum*np.exp(-theta1*complex(1j))
+
+        theta2 = np.angle(spectrum[wn2idx])
+        wn2 = wnaxis[wn2idx]
+
+        angles = (wnaxis-wn1) * theta2 / (wn2-wn1)
+        spectrum = np.angle(spectrum*np.exp(-1*angles*complex(1j)))
+
+        return spectrum
+    
 class NormalizeSpectrum(Transformation):
 
     def __init__(self, datatype=DataTypes.Phase, dounwrap=False):