update channel_mapping.py

channel_mapping.py

@@ -1,78 +1,135 @@
 import utils
+
+import time
 import os
 import numpy as np
 
 import mne
 from mne.channels import read_custom_montage
 
-def reorder_data(filename, old_idx):
+def reorder_data(filename, old_idx, fill_mode):
 	filepath = os.path.dirname(str(filename))
 	old_data = utils.read_train_data(filename)
 	new_data = np.zeros((30, old_data.shape[1]))
-	#print('old = ', old_data.shape)
+	print('old data shape: ', old_data.shape)
+	
+	if fill_mode == 'zero':
+		zero_arr = np.zeros((1, old_data.shape[1]))
+		old_data = np.concatenate((old_data, zero_arr), axis=0)
 
-	for j in range(30):
-		new_data[j, :] = old_data[old_idx[j]-1, :]
+	for i in range(30):
+		new_data[i, :] = old_data[old_idx[i]-1, :]
 
-	#print('i = ', i+1, ', ', new_data.shape)
+	print('new data shape: ', new_data.shape)
 	utils.save_data(new_data, filepath+'/mapped.csv')
 	return
 
+
+class Channel:  # (DigMontage):
+
+    def __init__(self, index, name=None, used=False, coord=None, topo_index=None, topo_location=None):
+
+		# super().__init__()
+		# self._montage = montage
+
+        self.name = name
+        self.index = index
+        self.used = used
+        self.coord = coord
+        self.topo_index = topo_index
+        self.topo_location = topo_location
+
+    def prefix(self):
+        ret = ''.join(filter(str.isalpha, self.name))
+        return ret[:len(ret) - 1] if ret[-1] == 'Z' else ret
+
+    def suffix(self):
+        return -1 if self.name[-1] == 'Z' else int(''.join(filter(str.isdigit, self.name)))
+
+
 def mapping(input_file, loc_file, fill_mode):
+	second1 = time.time()
+	
 	template_montage = read_custom_montage("./template_chanlocs.loc")
 	input_montage = read_custom_montage(loc_file)
 	#template_montage.plot()
 	#input_montage.plot()
+	
+	template = {}
+	input = {}
 
-	input_labels_dict = {}
+	# convert all channel names to uppercase
 	for i in range(30):
-		template_montage.rename_channels({template_montage.ch_names[i]:str.upper(template_montage.ch_names[i])}) # 統一大寫
+		channel = template_montage.ch_names[i]
+		template_montage.rename_channels({channel: str.upper(channel)})
+
+		channel = str.upper(channel)
+		template[channel] = Channel(index=i, name=channel)
 
 	for i in range(len(input_montage.ch_names)):
-		input_montage.rename_channels({input_montage.ch_names[i]:str.upper(input_montage.ch_names[i])}) # 統一大寫
-		input_labels_dict[input_montage.ch_names[i]] = i
+		channel = input_montage.ch_names[i]
+		input_montage.rename_channels({channel: str.upper(channel)})
+
+		channel = str.upper(channel)
+		input[channel] = Channel(index=i, coord=input_montage.get_positions()['ch_pos'][channel])
 
 
 	new_idx = [-1]*30
-	new_idx_name = ['']*30 # tmp
-	input_used = [0]*len(input_montage.ch_names)
+	new_idx_name = ['']*30  # tmp
+	missing_channels = []
+	z_row_idx = len(input_montage.ch_names)
+	
+	
+	# STAGE_1
+	
+	# match the template's channel names with the input ones
 	finish_flag = 1
-
-	alias = {'T3':'T7', 'T4':'T8', 'T5':'P7', 'T6':'P8'} # CP7,FT7 ?
+	alias = {
+		'T3': 'T7',
+		'T4': 'T8',
+		'T5': 'P7',
+		'T6': 'P8',
+		'TP7': 'T5\'',
+		'TP8': 'T6\'',
+		}
 	
-	# correct place
 	for i in range(30):
-		channel_name = template_montage.ch_names[i]
-
-		if channel_name in input_labels_dict:
-			new_idx[i] = input_labels_dict[channel_name]
-			new_idx_name[i] = channel_name # tmp
+		channel = template_montage.ch_names[i]
+		if channel not in input.keys() | alias.keys():
+			finish_flag = 0
+			continue
 
-			input_used[new_idx[i]] = 1
+		if channel not in input and channel in alias:
+			if alias[channel] in input:
+				template_montage.rename_channels({channel: alias[channel]})
+				template[alias[channel]] = template.pop(channel)
+				template[alias[channel]].name = alias[channel]
+				channel = alias[channel]
+			else:
+				finish_flag = 0
+				continue
 
-		elif channel_name in alias:
-			template_montage.rename_channels({channel_name:alias[channel_name]})
-			channel_name = template_montage.ch_names[i]
-			new_idx[i] = input_labels_dict[channel_name]
-			new_idx_name[i] = channel_name # tmp
+		new_idx[i] = input[channel].index
+		new_idx_name[i] = channel  # tmp
+		input[channel].used = True
 
-			input_used[new_idx[i]] = 1
-		else:
-			finish_flag = 0
-	
 	if finish_flag == 1:
-		print('Finish at stage 1,2 !')
-		reorder_data(input_file, new_idx) # & save data to mapped.csv
+		second2 = time.time()
+		print('Finish at stage 1 ! (',second2 - second1,'s)')
+		#print('new idx order:', new_idx)
+		#print('new_idx_name:', new_idx_name) # tmp
+
+		reorder_data(input_file, new_idx, fill_mode) # & save data to mapped.csv
 		return
 	
 	
 	
-	# store channel positions in 2-d array
-	template_pos = []
-	template_pos_idx = []
-
+	# STAGE_2
+	
+	# store channel positions in a 2-d array
+	template_topo_pos = []  
 	temporal_channels = []
-	temporal_row_prefix = ['FC','C','CP','P']
+	temporal_row_prefix = ['FC', 'C', 'CP', 'P']
 
 	cnt = 0
 	for i in range(7):
@@ -81,93 +138,136 @@ def mapping(input_file, loc_file, fill_mode):
 			if [i,j] in [[0,0],[0,2],[0,4],[6,0],[6,4]]:
 				tmp.append('')
 			else:
-				tmp.append(template_montage.ch_names[cnt])
-				template_pos_idx.append([i,j])
+				channel = template_montage.ch_names[cnt]
+				tmp.append(channel)
+
+				ver = 'front' if i<3 else 'center' if i==3 else 'back'
+				hor = 'left' if j<2 else 'center' if j==2 else 'right'
+				template[channel].topo_index = [i, j]
+				template[channel].topo_location = [ver, hor]
 
-				if i>1 and j in [0,4]:
-					temporal_channels.append(template_montage.ch_names[cnt])
+				if i > 1 and j in [0, 4]:
+					temporal_channels.append(channel)
 				cnt += 1
-		template_pos.append(tmp)
-	
-	
-	
-	# CZ
-	template_CZ_idx = 14
-	if new_idx[template_CZ_idx] == -1:
+		template_topo_pos.append(tmp)
+		
+
+	# ensure that CZ is found or imputed by another channel
+	if 'CZ' not in input and fill_mode=='adjacent':
 		min_dist = 1e5
-		nearest_channel = 'CZ'
 		for channel in input_montage.ch_names:
-			cur_x, cur_y, cur_z = input_montage.get_positions()['ch_pos'][channel]
-			if cur_x**2+cur_y**2 < min_dist and channel != 'CZ':
-				nearest_channel = channel
-				min_dist = cur_x**2+cur_y**2
-		input_labels_dict['CZ'] = input_labels_dict[nearest_channel]
+			curr_x, curr_y, curr_z = input[channel].coord.round(6)
 
+			if curr_x**2 + curr_y**2 < min_dist:
+				nearest_channel = channel
+				min_dist = curr_x**2 + curr_y**2
 
+		if input[nearest_channel].used == True:
+			missing_channels.append('CZ')
+		input[nearest_channel].used = True
+		input['CZ'] = input[nearest_channel]
 
-	finish_flag = 1
+		print("CZ's nearest neighbor:", nearest_channel)
 
-	if fill_mode == "zero":
-		z_row_idx = len(input_montage.ch_names)
 
 	for i in range(30):
 		if new_idx[i] != -1:
 			continue
 
-		channel_name = template_montage.ch_names[i]
-		channel_prefix = channel_name[:len(channel_name)-1]
-		channel_suffix = -1 if channel_name[-1]=='Z' else int(channel_name[-1])
+		channel = template_montage.ch_names[i]
 
-		# current target channel is in the middle
-		if channel_suffix == -1:
+		curr_prefix = template[channel].prefix()
+		curr_suffix = template[channel].suffix()
 
-			if fill_mode == "zero":
-				new_idx[i] = z_row_idx
+		curr_row = template[channel].topo_index[0]
+		curr_col = template[channel].topo_index[1]
+		curr_ver = template[channel].topo_location[0]
+		curr_hor = template[channel].topo_location[1]
+
+		impute_channel = ''
+
+		# if the current channel is a temporal channel
+		if channel in temporal_channels:
+			curr_prefix = temporal_row_prefix[temporal_channels.index(channel)//2]
+			curr_suffix = 7 if curr_hor=='left' else 8
+
+		if fill_mode == 'zero':
 
-			elif fill_mode == "adjacent":
-
-				if channel_prefix+str(1) in input_labels_dict: # ex: FCZ<-FC1
-					new_idx[i] = input_labels_dict[channel_prefix+str(1)]
-					new_idx_name[i] = channel_prefix+str(1) # tmp
-				elif (channel_name in ['FCZ','CPZ']): # and ('CZ' in input_labels_dict): # ex: FCZ<-CZ
-					new_idx[i] = input_labels_dict['CZ']
-					new_idx_name[i] = 'CZ' # tmp
-				elif channel_prefix+str(3) in input_labels_dict: # ex: FCZ<-FC3
-					new_idx[i] = input_labels_dict[channel_prefix+str(3)]
-					new_idx_name[i] = channel_prefix+str(3) # tmp
+			impute_channel = curr_prefix+str(1) if curr_hor=='center' else curr_prefix+str(curr_suffix-2)
+			if impute_channel not in input or input[impute_channel].used==True:
+				impute_channel = ''
+				new_idx[i] = z_row_idx
+				missing_channels.append(channel)
+				continue
+		
+		elif fill_mode == 'adjacent':
+
+			if curr_hor == 'center':  # FZ, FPZ, CZ...
+
+				if curr_prefix+str(1) in input:  # ex: FZ<-F1
+					impute_channel = curr_prefix + str(1)
+				elif curr_ver=='front' and template_topo_pos[curr_row+1][curr_col] in input: # ex: FZ<-FCZ
+					impute_channel = template_topo_pos[curr_row+1][curr_col]
+					
+				elif curr_ver=='back' and template_topo_pos[curr_row-1][curr_col] in input: # ex: PZ<-CPZ
+					impute_channel = template_topo_pos[curr_row-1][curr_col]
+					
+				elif curr_prefix+str(3) in input: # ex: FZ<-F3
+					impute_channel = curr_prefix + str(3)
+					
 				else:
-					new_idx[i] = input_labels_dict['CZ']
-					new_idx_name[i] = 'CZ' # tmp
-
-		# current target channel is in the left/right region
-		else:
-			try:
-				# if the current target channel is a temporal channel
-				potential_neighbor = temporal_row_prefix[temporal_channels.index(channel_name)//2]+str(5 if channel_suffix%2==1 else 6) # ex: FT7<-FC5
-			except:
-				potential_neighbor = channel_name[:len(channel_name)-1]+str(channel_suffix-2) # ex: FC3<-FC1, FC4<-FC2
-
-			if (potential_neighbor in input_labels_dict) and (input_used[input_labels_dict[potential_neighbor]]==0):
-				new_idx[i] = input_labels_dict[potential_neighbor]
-				new_idx_name[i] = potential_neighbor # tmp
-
-				input_used[new_idx[i]] = 1
-			else:
-				if fill_mode == "zero":
-					new_idx[i] = z_row_idx
-				elif fill_mode == "adjacent": # 先這樣暫時這樣...QQ
-					mid_channel = template_pos[template_pos_idx[i][0]][2]
-					mid_channel_idx = template_montage.ch_names.index(mid_channel)
-					new_idx[i] = new_idx[mid_channel_idx]
-					new_idx_name[i] = mid_channel # tmp
-
-					#finish_flag = 0
-
-		#if finish_flag == 1:
-		#  print('Finish at stage 3,4 !')
-		#  reorder_data(input_file, new_idx) # & save data to mapped.csv
-		#  return
-		#else:
-		#  print('Error: the channel mapping process has failed!')
-		reorder_data(input_file, new_idx)
+					impute_channel = 'CZ'
+					
+			elif curr_hor == 'left' or curr_hor == 'right':
+
+				ver_ctrl = 1 if curr_ver=='front' else 2 if curr_ver=='back' else 3 # bit0: row+1, bit1: row-1
+				ver_dir = 1 if curr_ver == 'front' else -1
+
+				# search horizontally
+				cnt = 0
+				tmp_suffix = curr_suffix
+				while tmp_suffix > 0:  # ex: F7<-F5/F3/F1
+					tmp_suffix = curr_suffix - 2*cnt
+					if curr_prefix+str(tmp_suffix) in input:
+						impute_channel = curr_prefix + str(tmp_suffix)
+						break
+
+					if cnt == 2:
+						# check row+1/row-1
+						if ver_ctrl&1 and template_topo_pos[curr_row+1][curr_col] in input:
+							impute_channel = template_topo_pos[curr_row+1][curr_col]
+							break
+						if ver_ctrl&2 and template_topo_pos[curr_row-1][curr_col] in input:
+							impute_channel = template_topo_pos[curr_row-1][curr_col]
+							break
+					cnt += 1
+
+				# search vertically
+				if impute_channel == '':
+					cnt = 0
+					tmp_row = curr_row + ver_dir
+					while tmp_row-ver_dir != 3:  # terminate if the last channel is a middle one
+						if template_topo_pos[tmp_row][curr_col] in input:
+							impute_channel = template_topo_pos[tmp_row][curr_col]
+							break
+						tmp_row += ver_dir
+
+				# if still cannot find available channel...
+				if impute_channel == '':
+					impute_channel = 'CZ'
+
+		new_idx[i] = input[impute_channel].index
+		new_idx_name[i] = impute_channel  # tmp
+		if input[impute_channel].used == True:  # this channel is shared with others
+			missing_channels.append(channel)
+		input[impute_channel].used = True
+	
+	second2 = time.time()
+	print('Finish at stage 2 ! (',second2 - second1,'s)')
+	#print('new_idx:', new_idx)
+	#print('new_idx_name:', new_idx_name) # tmp
+	print('missing_channels:', missing_channels)
+	reorder_data(input_file, new_idx, fill_mode) # & save data to mapped.csv
+	
+	return