celldetective 1.3.9.post4__py3-none-any.whl → 1.4.0__py3-none-any.whl

celldetective/neighborhood.py

@@ -9,6 +9,281 @@ from celldetective.io import locate_labels, get_position_pickle, get_position_ta
 
 abs_path = os.sep.join([os.path.split(os.path.dirname(os.path.realpath(__file__)))[0], 'celldetective'])
 
+def _fill_distance_neighborhood_at_t(time_index, setA, setB, dist_map, attention_weight=None, include_dead_weight=False, symmetrize=False, compute_cum_sum=False,
+							   weights=None, closest_A=None, neigh_col="", column_labelsA=None, column_labelsB=None, statusA=None, statusB=None, distance=10):
+
+	index_A = list(setA.loc[setA[column_labelsA['time']] == time_index].index)
+	index_B = list(setB.loc[setB[column_labelsB['time']] == time_index].index)
+
+	dataA = setA.loc[setA[column_labelsA['time']] == time_index, [column_labelsA['x'], column_labelsA['y'], column_labelsA['track'],statusA]].to_numpy()
+
+	ids_A = dataA[:, 2]
+	status_A = dataA[:, 3]
+
+	dataB = setB.loc[setB[column_labelsB['time']] == time_index, [column_labelsB['x'], column_labelsB['y'], column_labelsB['track'],statusB]].to_numpy()
+	ids_B = dataB[:, 2]
+	status_B = dataB[:, 3]
+
+	for k in range(dist_map.shape[0]):
+		
+		col = dist_map[k, :]
+		col[col == 0.] = 1.0E06
+		
+		neighs_B = np.array([ids_B[i] for i in np.where((col <= distance))[0]])
+		status_neigh_B = np.array([status_B[i] for i in np.where((col <= distance))[0]])
+		dist_B = [round(col[i], 2) for i in np.where((col <= distance))[0]]
+		if len(dist_B) > 0:
+			closest_B_cell = neighs_B[np.argmin(dist_B)]
+		
+		if symmetrize and attention_weight:
+			n_neighs = float(len(neighs_B))
+			if not include_dead_weight:
+				n_neighs_alive = len(np.where(status_neigh_B == 1)[0])
+				neigh_count = n_neighs_alive
+			else:
+				neigh_count = n_neighs
+			if neigh_count > 0:
+				weight_A = 1. / neigh_count
+			else:
+				weight_A = np.nan
+			
+			if not include_dead_weight and status_A[k] == 0:
+				weight_A = 0
+		
+		neighs = []
+		setA.at[index_A[k], neigh_col] = []
+		for n in range(len(neighs_B)):
+			
+			# index in setB
+			n_index = np.where(ids_B == neighs_B[n])[0][0]
+			# Assess if neigh B is closest to A
+			if attention_weight:
+				if closest_A[n_index] == ids_A[k]:
+					closest = True
+				else:
+					closest = False
+			
+			if symmetrize:
+				# Load neighborhood previous data
+				sym_neigh = setB.loc[index_B[n_index], neigh_col]
+				if neighs_B[n] == closest_B_cell:
+					closest_b = True
+				else:
+					closest_b = False
+				if isinstance(sym_neigh, list):
+					sym_neigh.append({'id': ids_A[k], 'distance': dist_B[n], 'status': status_A[k]})
+				else:
+					sym_neigh = [{'id': ids_A[k], 'distance': dist_B[n], 'status': status_A[k]}]
+				if attention_weight:
+					sym_neigh[-1].update({'weight': weight_A, 'closest': closest_b})
+			
+			# Write the minimum info about neighborhing cell B
+			neigh_dico = {'id': neighs_B[n], 'distance': dist_B[n], 'status': status_neigh_B[n]}
+			if attention_weight:
+				neigh_dico.update({'weight': weights[n_index], 'closest': closest})
+			
+			if compute_cum_sum:
+				# Compute the integrated presence of the neighboring cell B
+				assert column_labelsB[
+						   'track'] == 'TRACK_ID', 'The set B does not seem to contain tracked data. The cumulative time will be meaningless.'
+				past_neighs = [[ll['id'] for ll in l] if len(l) > 0 else [None] for l in setA.loc[
+					(setA[column_labelsA['track']] == ids_A[k]) & (setA[column_labelsA['time']] <= time_index), neigh_col].to_numpy()]
+				past_neighs = [item for sublist in past_neighs for item in sublist]
+				
+				if attention_weight:
+					past_weights = [[ll['weight'] for ll in l] if len(l) > 0 else [None] for l in setA.loc[
+						(setA[column_labelsA['track']] == ids_A[k]) & (
+								setA[column_labelsA['time']] <= time_index), neigh_col].to_numpy()]
+					past_weights = [item for sublist in past_weights for item in sublist]
+				
+				cum_sum = len(np.where(past_neighs == neighs_B[n])[0])
+				neigh_dico.update({'cumulated_presence': cum_sum + 1})
+				
+				if attention_weight:
+					cum_sum_weighted = np.sum(
+						[w if l == neighs_B[n] else 0 for l, w in zip(past_neighs, past_weights)])
+					neigh_dico.update({'cumulated_presence_weighted': cum_sum_weighted + weights[n_index]})
+			
+			if symmetrize:
+				setB.at[index_B[n_index], neigh_col] = sym_neigh
+			
+			neighs.append(neigh_dico)
+		
+		setA.at[index_A[k], neigh_col] = neighs
+
+def _fill_contact_neighborhood_at_t(time_index, setA, setB, dist_map, intersection_map=None, attention_weight=None, include_dead_weight=False, symmetrize=False, compute_cum_sum=False,
+							   weights=None, closest_A=None, neigh_col="", column_labelsA=None, column_labelsB=None, statusA=None, statusB=None, d_filter=10):
+
+	index_A = list(setA.loc[setA[column_labelsA['time']] == time_index].index)
+	index_B = list(setB.loc[setB[column_labelsB['time']] == time_index].index)
+
+	dataA = setA.loc[setA[column_labelsA['time']] == time_index, [column_labelsA['x'], column_labelsA['y'], column_labelsA['track'], column_labelsA['mask_id'],
+	statusA]].to_numpy()
+
+	ids_A = dataA[:, 2]
+	status_A = dataA[:, 4]
+
+	dataB = setB.loc[setB[column_labelsB['time']] == time_index, [column_labelsB['x'], column_labelsB['y'], column_labelsB['track'], column_labelsB['mask_id'],
+	statusB]].to_numpy()
+	ids_B = dataB[:, 2]
+	status_B = dataB[:, 4]
+
+	for k in range(dist_map.shape[0]):
+
+		col = dist_map[k, :]
+		col_inter = intersection_map[k, :]
+		col[col == 0.] = 1.0E06
+
+		neighs_B = np.array([ids_B[i] for i in np.where((col <= d_filter))[0]])
+		status_neigh_B = np.array([status_B[i] for i in np.where((col <= d_filter))[0]])
+		dist_B = [round(col[i], 2) for i in np.where((col <= d_filter))[0]]
+		intersect_B = [round(col_inter[i], 2) for i in np.where((col <= d_filter))[0]]
+
+		if len(dist_B) > 0:
+			closest_B_cell = neighs_B[np.argmin(dist_B)]
+
+		if symmetrize and attention_weight:
+			n_neighs = float(len(neighs_B))
+			if not include_dead_weight:
+				n_neighs_alive = len(np.where(status_neigh_B == 1)[0])
+				neigh_count = n_neighs_alive
+			else:
+				neigh_count = n_neighs
+			if neigh_count > 0:
+				weight_A = 1. / neigh_count
+			else:
+				weight_A = np.nan
+
+			if not include_dead_weight and status_A[k] == 0:
+				weight_A = 0
+
+		neighs = []
+		setA.at[index_A[k], neigh_col] = []
+		for n in range(len(neighs_B)):
+
+			# index in setB
+			n_index = np.where(ids_B == neighs_B[n])[0][0]
+			# Assess if neigh B is closest to A
+			if attention_weight:
+				if closest_A[n_index] == ids_A[k]:
+					closest = True
+				else:
+					closest = False
+
+			if symmetrize:
+				# Load neighborhood previous data
+				sym_neigh = setB.loc[index_B[n_index], neigh_col]
+				if neighs_B[n] == closest_B_cell:
+					closest_b = True
+				else:
+					closest_b = False
+				if isinstance(sym_neigh, list):
+					sym_neigh.append({'id': ids_A[k], 'distance': dist_B[n], 'status': status_A[k],
+									  'intersection': intersect_B[n]})
+				else:
+					sym_neigh = [{'id': ids_A[k], 'distance': dist_B[n], 'status': status_A[k],
+								  'intersection': intersect_B[n]}]
+				if attention_weight:
+					sym_neigh[-1].update({'weight': weight_A, 'closest': closest_b})
+
+			# Write the minimum info about neighborhing cell B
+			neigh_dico = {'id': neighs_B[n], 'distance': dist_B[n], 'status': status_neigh_B[n],
+						  'intersection': intersect_B[n]}
+			if attention_weight:
+				neigh_dico.update({'weight': weights[n_index], 'closest': closest})
+
+			if compute_cum_sum:
+				# Compute the integrated presence of the neighboring cell B
+				assert column_labelsB[
+						   'track'] == 'TRACK_ID', 'The set B does not seem to contain tracked data. The cumulative time will be meaningless.'
+				past_neighs = [[ll['id'] for ll in l] if len(l) > 0 else [None] for l in setA.loc[
+					(setA[column_labelsA['track']] == ids_A[k]) & (
+							setA[column_labelsA['time']] <= time_index), neigh_col].to_numpy()]
+				past_neighs = [item for sublist in past_neighs for item in sublist]
+
+				if attention_weight:
+					past_weights = [[ll['weight'] for ll in l] if len(l) > 0 else [None] for l in
+									setA.loc[
+										(setA[column_labelsA['track']] == ids_A[k]) & (
+												setA[column_labelsA['time']] <= time_index), neigh_col].to_numpy()]
+					past_weights = [item for sublist in past_weights for item in sublist]
+
+				cum_sum = len(np.where(past_neighs == neighs_B[n])[0])
+				neigh_dico.update({'cumulated_presence': cum_sum + 1})
+
+				if attention_weight:
+					cum_sum_weighted = np.sum(
+						[w if l == neighs_B[n] else 0 for l, w in zip(past_neighs, past_weights)])
+					neigh_dico.update(
+						{'cumulated_presence_weighted': cum_sum_weighted + weights[n_index]})
+
+			if symmetrize:
+				setB.at[index_B[n_index], neigh_col] = sym_neigh
+
+			neighs.append(neigh_dico)
+
+		setA.at[index_A[k], neigh_col] = neighs
+
+
+def _compute_mask_contact_dist_map(setA, setB, labelsA, labelsB=None, distance=10, mode="self", column_labelsA=None, column_labelsB=None):
+
+	coordinates_A = setA.loc[:, [column_labelsA['x'], column_labelsA['y']]].to_numpy()
+	coordinates_B = setB.loc[:, [column_labelsB['x'], column_labelsB['y']]].to_numpy()
+	ids_A = setA.loc[:, column_labelsA["track"]].to_numpy()
+	ids_B = setB.loc[:, column_labelsB["track"]].to_numpy()
+	mask_ids_A = setA.loc[:, column_labelsA["mask_id"]].to_numpy()
+	mask_ids_B = setB.loc[:, column_labelsB["mask_id"]].to_numpy()
+
+	# compute distance matrix
+	dist_map = cdist(coordinates_A, coordinates_B, metric="euclidean")
+	intersection_map = np.zeros_like(dist_map).astype(float)
+
+	# Do the mask contact computation
+	labelsA = np.where(np.isin(labelsA, mask_ids_A), labelsA.copy(), 0.)
+
+	if labelsB is not None:
+		labelsB = np.where(np.isin(labelsB, mask_ids_B), labelsB.copy(), 0.)
+
+	contact_pairs = contact_neighborhood(labelsA, labelsB=labelsB, border=distance, connectivity=2)
+
+	# Put infinite distance to all non-contact pairs (something like this)
+	flatA = labelsA.flatten()
+	if labelsB is not None:
+		flatB = labelsB.flatten()
+
+	if len(contact_pairs) > 0:
+		mask = np.ones_like(dist_map).astype(bool)
+
+		indices_to_keep = []
+		for cp in contact_pairs:
+
+			cp = np.abs(cp)
+			mask_A, mask_B = cp
+			idx_A = np.where(mask_ids_A == int(mask_A))[0][0]
+			idx_B = np.where(mask_ids_B == int(mask_B))[0][0]
+
+			intersection = 0
+			if labelsB is not None:
+				intersection = len(flatA[(flatA == int(mask_A)) & (flatB == int(mask_B))])
+
+			indices_to_keep.append([idx_A, idx_B, intersection])
+			print(f'Ref cell #{ids_A[idx_A]} matched with neigh. cell #{ids_B[idx_B]}...')
+			print(f'Computed intersection: {intersection} px...')
+
+		if len(indices_to_keep) > 0:
+			indices_to_keep = np.array(indices_to_keep)
+			mask[indices_to_keep[:, 0], indices_to_keep[:, 1]] = False
+			if mode == 'self':
+				mask[indices_to_keep[:, 1], indices_to_keep[:, 0]] = False
+			dist_map[mask] = 1.0E06
+			intersection_map[indices_to_keep[:, 0], indices_to_keep[:, 1]] = indices_to_keep[:, 2]
+		else:
+			dist_map[:, :] = 1.0E06
+	else:
+		dist_map[:, :] = 1.0E06
+
+	return dist_map, intersection_map
+
 
 def set_live_status(setA, setB, status, not_status_option):
 	"""
@@ -202,15 +477,12 @@ def distance_cut_neighborhood(setA, setB, distance, mode='two-pop', status=None,
 			int)
 		for t in tqdm(timeline):
 
-			index_A = list(setA.loc[setA[cl[0]['time']] == t].index)
 			coordinates_A = setA.loc[setA[cl[0]['time']] == t, [cl[0]['x'], cl[0]['y']]].to_numpy()
 			ids_A = setA.loc[setA[cl[0]['time']] == t, cl[0]['track']].to_numpy()
 			status_A = setA.loc[setA[cl[0]['time']] == t, status[0]].to_numpy()
 
-			index_B = list(setB.loc[setB[cl[1]['time']] == t].index)
 			coordinates_B = setB.loc[setB[cl[1]['time']] == t, [cl[1]['x'], cl[1]['y']]].to_numpy()
 			ids_B = setB.loc[setB[cl[1]['time']] == t, cl[1]['track']].to_numpy()
-			status_B = setB.loc[setB[cl[1]['time']] == t, status[1]].to_numpy()
 
 			if len(ids_A) > 0 and len(ids_B) > 0:
 
@@ -218,96 +490,14 @@ def distance_cut_neighborhood(setA, setB, distance, mode='two-pop', status=None,
 				dist_map = cdist(coordinates_A, coordinates_B, metric="euclidean")
 
 				if attention_weight:
-					weights, closest_A = compute_attention_weight(dist_map, d, status_A, ids_A, axis=1,
-																  include_dead_weight=include_dead_weight)
+					weights, closest_A = compute_attention_weight(dist_map, d, status_A, ids_A, axis=1, include_dead_weight=include_dead_weight)
 
-				# Target centric
-				for k in range(dist_map.shape[0]):
-
-					col = dist_map[k, :]
-					col[col == 0.] = 1.0E06
-
-					neighs_B = np.array([ids_B[i] for i in np.where((col <= d))[0]])
-					status_neigh_B = np.array([status_B[i] for i in np.where((col <= d))[0]])
-					dist_B = [round(col[i], 2) for i in np.where((col <= d))[0]]
-					if len(dist_B) > 0:
-						closest_B_cell = neighs_B[np.argmin(dist_B)]
-
-					if symmetrize and attention_weight:
-						n_neighs = float(len(neighs_B))
-						if not include_dead_weight:
-							n_neighs_alive = len(np.where(status_neigh_B == 1)[0])
-							neigh_count = n_neighs_alive
-						else:
-							neigh_count = n_neighs
-						if neigh_count > 0:
-							weight_A = 1. / neigh_count
-						else:
-							weight_A = np.nan
-
-						if not include_dead_weight and status_A[k] == 0:
-							weight_A = 0
-
-					neighs = []
-					setA.at[index_A[k], neigh_col] = []
-					for n in range(len(neighs_B)):
-
-						# index in setB
-						n_index = np.where(ids_B == neighs_B[n])[0][0]
-						# Assess if neigh B is closest to A
-						if attention_weight:
-							if closest_A[n_index] == ids_A[k]:
-								closest = True
-							else:
-								closest = False
-
-						if symmetrize:
-							# Load neighborhood previous data
-							sym_neigh = setB.loc[index_B[n_index], neigh_col]
-							if neighs_B[n] == closest_B_cell:
-								closest_b = True
-							else:
-								closest_b = False
-							if isinstance(sym_neigh, list):
-								sym_neigh.append({'id': ids_A[k], 'distance': dist_B[n], 'status': status_A[k]})
-							else:
-								sym_neigh = [{'id': ids_A[k], 'distance': dist_B[n], 'status': status_A[k]}]
-							if attention_weight:
-								sym_neigh[-1].update({'weight': weight_A, 'closest': closest_b})
-
-						# Write the minimum info about neighborhing cell B
-						neigh_dico = {'id': neighs_B[n], 'distance': dist_B[n], 'status': status_neigh_B[n]}
-						if attention_weight:
-							neigh_dico.update({'weight': weights[n_index], 'closest': closest})
-
-						if compute_cum_sum:
-							# Compute the integrated presence of the neighboring cell B
-							assert cl[1][
-									   'track'] == 'TRACK_ID', 'The set B does not seem to contain tracked data. The cumulative time will be meaningless.'
-							past_neighs = [[ll['id'] for ll in l] if len(l) > 0 else [None] for l in setA.loc[
-								(setA[cl[0]['track']] == ids_A[k]) & (setA[cl[0]['time']] <= t), neigh_col].to_numpy()]
-							past_neighs = [item for sublist in past_neighs for item in sublist]
-
-							if attention_weight:
-								past_weights = [[ll['weight'] for ll in l] if len(l) > 0 else [None] for l in setA.loc[
-									(setA[cl[0]['track']] == ids_A[k]) & (
-												setA[cl[0]['time']] <= t), neigh_col].to_numpy()]
-								past_weights = [item for sublist in past_weights for item in sublist]
-
-							cum_sum = len(np.where(past_neighs == neighs_B[n])[0])
-							neigh_dico.update({'cumulated_presence': cum_sum + 1})
-
-							if attention_weight:
-								cum_sum_weighted = np.sum(
-									[w if l == neighs_B[n] else 0 for l, w in zip(past_neighs, past_weights)])
-								neigh_dico.update({'cumulated_presence_weighted': cum_sum_weighted + weights[n_index]})
-
-						if symmetrize:
-							setB.at[index_B[n_index], neigh_col] = sym_neigh
-
-						neighs.append(neigh_dico)
-
-					setA.at[index_A[k], neigh_col] = neighs
+				_fill_distance_neighborhood_at_t(t, setA, setB, dist_map,
+												attention_weight=attention_weight,
+												include_dead_weight=include_dead_weight, symmetrize=symmetrize,
+												compute_cum_sum=compute_cum_sum, weights=weights, closest_A=closest_A,
+												neigh_col=neigh_col, column_labelsA=cl[0], column_labelsB=cl[1],
+												statusA=status[0], statusB=status[1], distance=d)
 
 	return setA, setB
 
@@ -435,6 +625,7 @@ def compute_neighborhood_at_position(pos, distance, population=['targets', 'effe
 		df_B = df_B.drop(columns=unwanted)
 
 	df_A, df_B = distance_cut_neighborhood(df_A, df_B, distance, **neighborhood_kwargs)
+
 	if df_A is None or df_B is None or len(df_A)==0:
 		return None
 
@@ -442,6 +633,7 @@ def compute_neighborhood_at_position(pos, distance, population=['targets', 'effe
 
 		if neighborhood_kwargs['mode'] == 'two-pop':
 			neigh_col = f'neighborhood_2_circle_{d}_px'
+
 		elif neighborhood_kwargs['mode'] == 'self':
 			neigh_col = f'neighborhood_self_circle_{d}_px'
 
@@ -464,17 +656,27 @@ def compute_neighborhood_at_position(pos, distance, population=['targets', 'effe
 					df_A = mean_neighborhood_after_event(df_A, neigh_col, event_time_col)
 				print('Done...')
 
-	df_A.to_pickle(path_A.replace('.csv', '.pkl'))
 	if not population[0] == population[1]:
-		# Remove neighborhood column
+		# Remove neighborhood column from neighbor table, rename with actual population name
 		for td, d in zip(theta_dist, distance):
 			if neighborhood_kwargs['mode'] == 'two-pop':
 				neigh_col = f'neighborhood_2_circle_{d}_px'
+				new_neigh_col = neigh_col.replace('_2_',f'_({population[0]}-{population[1]})_')
+				df_A = df_A.rename(columns={neigh_col: new_neigh_col})
 			elif neighborhood_kwargs['mode'] == 'self':
 				neigh_col = f'neighborhood_self_circle_{d}_px'
 			df_B = df_B.drop(columns=[neigh_col])
 		df_B.to_pickle(path_B.replace('.csv', '.pkl'))
 
+	cols_to_rename = [c for c in list(df_A.columns) if c.startswith('intermediate_count_') or c.startswith('inclusive_count_') or c.startswith('exclusive_count_') or c.startswith('mean_count_')]
+	new_col_names = [c.replace('_2_',f'_({population[0]}-{population[1]})_') for c in cols_to_rename]
+	new_name_map = {}
+	for k,c in enumerate(cols_to_rename):
+		new_name_map.update({c: new_col_names[k]})
+	df_A = df_A.rename(columns=new_name_map)
+	
+	df_A.to_pickle(path_A.replace('.csv', '.pkl'))
+
 	unwanted = df_A.columns[df_A.columns.str.startswith('neighborhood_')]
 	df_A2 = df_A.drop(columns=unwanted)
 	df_A2.to_csv(path_A, index=False)
@@ -904,16 +1106,6 @@ def mask_contact_neighborhood(setA, setB, labelsA, labelsB, distance, mode='two-
 		do not count dead cells when establishing attention weight
 	"""
 
-	# Check live_status option
-	# if setA is not None:
-	# 	setA_id = extract_identity_col(setA)
-	# 	if setA_id=="TRACK_ID":
-	# 		setA = setA.loc[~setA['TRACK_ID'].isnull(),:].copy()
-	# if setB is not None:
-	# 	setB_id = extract_identity_col(setB)
-	# 	if setB_id=="TRACK_ID":
-	# 		setB = setB.loc[~setB['TRACK_ID'].isnull(),:].copy()
-
 	if setA is not None and setB is not None:
 		setA, setB, status = set_live_status(setA, setB, status, not_status_option)
 	else:
@@ -949,6 +1141,9 @@ def mask_contact_neighborhood(setA, setB, labelsA, labelsB, distance, mode='two-
 			neigh_col = f'neighborhood_2_contact_{d}_px'
 		elif mode == 'self':
 			neigh_col = f'neighborhood_self_contact_{d}_px'
+		else:
+			print("Please provide a valid mode between `two-pop` and `self`...")
+			return None
 
 		setA[neigh_col] = np.nan
 		setA[neigh_col] = setA[neigh_col].astype(object)
@@ -961,163 +1156,30 @@ def mask_contact_neighborhood(setA, setB, labelsA, labelsB, distance, mode='two-
 			int)
 		for t in tqdm(timeline):
 
-			index_A = list(setA.loc[setA[cl[0]['time']] == t].index)
-			dataA = setA.loc[setA[cl[0]['time']] == t, [cl[0]['x'], cl[0]['y'], cl[0]['track'], cl[0]['mask_id'], status[0]]].to_numpy()
-			coordinates_A = dataA[:,[0,1]]; ids_A = dataA[:,2]; mask_ids_A = dataA[:,3]; status_A = dataA[:,4];
-
-			index_B = list(setB.loc[setB[cl[1]['time']] == t].index)
-			dataB = setB.loc[setB[cl[1]['time']] == t, [cl[1]['x'], cl[1]['y'], cl[1]['track'], cl[1]['mask_id'], status[1]]].to_numpy()
-			coordinates_B = dataB[:,[0,1]]; ids_B = dataB[:,2]; mask_ids_B = dataB[:,3]; status_B = dataB[:,4]
-
-			if len(coordinates_A) > 0 and len(coordinates_B) > 0:
-
-				# compute distance matrix
-				dist_map = cdist(coordinates_A, coordinates_B, metric="euclidean")
-				intersection_map = np.zeros_like(dist_map).astype(float)
+			setA_t = setA.loc[setA[cl[0]['time']] == t, :]
+			setB_t = setB.loc[setB[cl[1]['time']] == t, :]
 
-				# Do the mask contact computation
-				lblA = labelsA[t]
-				lblA = np.where(np.isin(lblA, mask_ids_A), lblA, 0.)
-				
-				lblB = labelsB[t]
-				if lblB is not None:
-					lblB = np.where(np.isin(lblB, mask_ids_B), lblB, 0.)
-
-				contact_pairs = contact_neighborhood(lblA, labelsB=lblB, border=d, connectivity=2)
-
-				# Put infinite distance to all non-contact pairs (something like this)
-				plot_map = False
-				flatA = lblA.flatten()
-				if lblB is not None:
-					flatB = lblB.flatten()
-
-				if len(contact_pairs) > 0:
-					mask = np.ones_like(dist_map).astype(bool)
-
-					indices_to_keep = []
-					for cp in contact_pairs:
-
-						cp = np.abs(cp)
-						mask_A, mask_B = cp
-						idx_A = np.where(mask_ids_A == int(mask_A))[0][0]
-						idx_B = np.where(mask_ids_B == int(mask_B))[0][0]
-						
-						intersection = 0
-						if lblB is not None:
-							intersection = len(flatA[(flatA==int(mask_A))&(flatB==int(mask_B))])
-
-						indices_to_keep.append([idx_A,idx_B, intersection])
-						print(f'Ref cell #{ids_A[idx_A]} matched with neigh. cell #{ids_B[idx_B]}...')
-						print(f'Computed intersection: {intersection} px...')
-
-					if len(indices_to_keep) > 0:
-						indices_to_keep = np.array(indices_to_keep)
-						mask[indices_to_keep[:, 0], indices_to_keep[:, 1]] = False
-						if mode == 'self':
-							mask[indices_to_keep[:, 1], indices_to_keep[:, 0]] = False
-						dist_map[mask] = 1.0E06
-						intersection_map[indices_to_keep[:,0], indices_to_keep[:,1]] = indices_to_keep[:,2]
-						plot_map=True
-					else:
-						dist_map[:,:] = 1.0E06
-				else:
-					dist_map[:, :] = 1.0E06
+			if len(setA_t) > 0 and len(setB_t) > 0:
+				dist_map, intersection_map = _compute_mask_contact_dist_map(setA_t, setB_t, labelsA[t], labelsB[t],
+																			distance=d, mode=mode, column_labelsA=cl[0],
+																			column_labelsB=cl[1])
 
 				d_filter = 1.0E05
 				if attention_weight:
+					status_A = setA_t[status[0]].to_numpy()
+					ids_A = setA_t[cl[0]["track"]].to_numpy()
 					weights, closest_A = compute_attention_weight(dist_map, d_filter, status_A, ids_A, axis=1,
 																  include_dead_weight=include_dead_weight)
+				else:
+					weights = None
+					closest_A = None
 
-				# Target centric
-				for k in range(dist_map.shape[0]):
-
-					col = dist_map[k, :]
-					col_inter = intersection_map[k, :]
-					col[col == 0.] = 1.0E06
-
-					neighs_B = np.array([ids_B[i] for i in np.where((col <= d_filter))[0]])
-					status_neigh_B = np.array([status_B[i] for i in np.where((col <= d_filter))[0]])
-					dist_B = [round(col[i], 2) for i in np.where((col <= d_filter))[0]]
-					intersect_B = [round(col_inter[i], 2) for i in np.where((col <= d_filter))[0]]
-
-					if len(dist_B) > 0:
-						closest_B_cell = neighs_B[np.argmin(dist_B)]
-
-					if symmetrize and attention_weight:
-						n_neighs = float(len(neighs_B))
-						if not include_dead_weight:
-							n_neighs_alive = len(np.where(status_neigh_B == 1)[0])
-							neigh_count = n_neighs_alive
-						else:
-							neigh_count = n_neighs
-						if neigh_count > 0:
-							weight_A = 1. / neigh_count
-						else:
-							weight_A = np.nan
-
-						if not include_dead_weight and status_A[k] == 0:
-							weight_A = 0
-
-					neighs = []
-					setA.at[index_A[k], neigh_col] = []
-					for n in range(len(neighs_B)):
-
-						# index in setB
-						n_index = np.where(ids_B == neighs_B[n])[0][0]
-						# Assess if neigh B is closest to A
-						if attention_weight:
-							if closest_A[n_index] == ids_A[k]:
-								closest = True
-							else:
-								closest = False
-
-						if symmetrize:
-							# Load neighborhood previous data
-							sym_neigh = setB.loc[index_B[n_index], neigh_col]
-							if neighs_B[n] == closest_B_cell:
-								closest_b = True
-							else:
-								closest_b = False
-							if isinstance(sym_neigh, list):
-								sym_neigh.append({'id': ids_A[k], 'distance': dist_B[n], 'status': status_A[k], 'intersection': intersect_B[n]})
-							else:
-								sym_neigh = [{'id': ids_A[k], 'distance': dist_B[n], 'status': status_A[k], 'intersection': intersect_B[n]}]
-							if attention_weight:
-								sym_neigh[-1].update({'weight': weight_A, 'closest': closest_b})
-
-						# Write the minimum info about neighborhing cell B
-						neigh_dico = {'id': neighs_B[n], 'distance': dist_B[n], 'status': status_neigh_B[n], 'intersection': intersect_B[n]}
-						if attention_weight:
-							neigh_dico.update({'weight': weights[n_index], 'closest': closest})
-
-						if compute_cum_sum:
-							# Compute the integrated presence of the neighboring cell B
-							assert cl[1][
-									   'track'] == 'TRACK_ID', 'The set B does not seem to contain tracked data. The cumulative time will be meaningless.'
-							past_neighs = [[ll['id'] for ll in l] if len(l) > 0 else [None] for l in setA.loc[
-								(setA[cl[0]['track']] == ids_A[k]) & (setA[cl[0]['time']] <= t), neigh_col].to_numpy()]
-							past_neighs = [item for sublist in past_neighs for item in sublist]
-
-							if attention_weight:
-								past_weights = [[ll['weight'] for ll in l] if len(l) > 0 else [None] for l in setA.loc[
-									(setA[cl[0]['track']] == ids_A[k]) & (
-												setA[cl[0]['time']] <= t), neigh_col].to_numpy()]
-								past_weights = [item for sublist in past_weights for item in sublist]
-
-							cum_sum = len(np.where(past_neighs == neighs_B[n])[0])
-							neigh_dico.update({'cumulated_presence': cum_sum + 1})
-
-							if attention_weight:
-								cum_sum_weighted = np.sum(
-									[w if l == neighs_B[n] else 0 for l, w in zip(past_neighs, past_weights)])
-								neigh_dico.update({'cumulated_presence_weighted': cum_sum_weighted + weights[n_index]})
-
-						if symmetrize:
-							setB.at[index_B[n_index], neigh_col] = sym_neigh
-
-						neighs.append(neigh_dico)
-
-					setA.at[index_A[k], neigh_col] = neighs
+				_fill_contact_neighborhood_at_t(t, setA, setB, dist_map, intersection_map=intersection_map,
+												attention_weight=attention_weight,
+												include_dead_weight=include_dead_weight, symmetrize=symmetrize,
+												compute_cum_sum=compute_cum_sum, weights=weights, closest_A=closest_A,
+												neigh_col=neigh_col, column_labelsA=cl[0], column_labelsB=cl[1],
+												statusA=status[0], statusB=status[1], d_filter=d_filter)
 
 	return setA, setB
 
@@ -1278,17 +1340,28 @@ def compute_contact_neighborhood_at_position(pos, distance, population=['targets
 					df_A = mean_neighborhood_after_event(df_A, neigh_col, event_time_col, metrics=['inclusive', 'intermediate'])
 				print('Done...')
 				
-	df_A.to_pickle(path_A.replace('.csv', '.pkl'))
 	if not population[0] == population[1]:
-		# Remove neighborhood column
+		# Remove neighborhood column from neighbor table, rename with actual population name
 		for td, d in zip(theta_dist, distance):
 			if neighborhood_kwargs['mode'] == 'two-pop':
 				neigh_col = f'neighborhood_2_contact_{d}_px'
+				new_neigh_col = neigh_col.replace('_2_',f'_({population[0]}-{population[1]})_')
+				df_A = df_A.rename(columns={neigh_col: new_neigh_col})
 			elif neighborhood_kwargs['mode'] == 'self':
 				neigh_col = f'neighborhood_self_contact_{d}_px'
 			df_B = df_B.drop(columns=[neigh_col])
 		df_B.to_pickle(path_B.replace('.csv', '.pkl'))
 
+	cols_to_rename = [c for c in list(df_A.columns) if c.startswith('intermediate_count_') or c.startswith('inclusive_count_') or c.startswith('exclusive_count_') or c.startswith('mean_count_')]
+	new_col_names = [c.replace('_2_',f'_({population[0]}-{population[1]})_') for c in cols_to_rename]
+	new_name_map = {}
+	for k,c in enumerate(cols_to_rename):
+		new_name_map.update({c: new_col_names[k]})
+	df_A = df_A.rename(columns=new_name_map)
+
+	print(f'{df_A.columns=}')
+	df_A.to_pickle(path_A.replace('.csv', '.pkl'))
+
 	unwanted = df_A.columns[df_A.columns.str.startswith('neighborhood_')]
 	df_A2 = df_A.drop(columns=unwanted)
 	df_A2.to_csv(path_A, index=False)
@@ -1355,28 +1428,29 @@ def extract_neighborhood_in_pair_table(df, distance=None, reference_population="
 	"""
 
 
-	assert reference_population in ["targets", "effectors"], "Please set a valid reference population ('targets' or 'effectors')"
+	#assert reference_population in ["targets", "effectors"], "Please set a valid reference population ('targets' or 'effectors')"
 	if neighborhood_key is None:
-		assert neighbor_population in ["targets", "effectors"], "Please set a valid neighbor population ('targets' or 'effectors')"
+		#assert neighbor_population in ["targets", "effectors"], "Please set a valid neighbor population ('targets' or 'effectors')"
 		assert mode in ["circle", "contact"], "Please set a valid neighborhood computation mode ('circle' or 'contact')"
-		if reference_population==neighbor_population:
-			type = "self"
-		else:
-			type = "2"
-
+		type = '('+'-'.join([reference_population, neighbor_population])+')'
 		neigh_col = f"neighborhood_{type}_{mode}_{distance}_px"
 	else:
 		neigh_col = neighborhood_key.replace('status_','')
-		if 'self' in neigh_col:
-			neighbor_population = reference_population
+		if '_(' in neigh_col and ')_' in neigh_col:
+			neighbor_population = neigh_col.split('_(')[-1].split(')_')[0].split('-')[-1]
 		else:
-			if reference_population=="effectors":
-				neighbor_population=='targets'
+			if 'self' in neigh_col:
+				neighbor_population = reference_population
 			else:
-				neighbor_population=='effectors'
+				if reference_population=="effectors":
+					neighbor_population='targets'
+				else:
+					neighbor_population='effectors'
 
 	assert "status_"+neigh_col in list(df.columns),"The selected neighborhood does not appear in the data..."
 
+	print(df[['reference_population','neighbor_population', "status_"+neigh_col]])
+
 	if contact_only:
 		s_keep = [1]
 	else: