# Copyright (C) 2009-2022, Ecole Polytechnique Federale de Lausanne (EPFL) and
# Hospital Center and University of Lausanne (UNIL-CHUV), Switzerland, and CMP3 contributors
# All rights reserved.
#
# This software is distributed under the open-source license Modified BSD.
"""Module that defines CMTK utility functions for the EEG pipeline."""
import os
import copy
import csv
import networkx as nx
import numpy as np
import scipy.io as sio
[docs]def save_eeg_connectome_file(output_dir, output_basename, con_res, roi_labels, output_types=None):
"""Save a dictionary of connectivity matrices with corresponding keys to the metrics in the multiple formats of CMP3.
Parameters
----------
output_dir : str
Output directory for the connectome file(s)
output_basename : str
Base name for the connectome file(s) i.e.,
``sub-01_atlas-L20018_res-scale1_conndata-network_connectivity``
con_res : dict
Dictionary of connectivity metric / matrix pairs
roi_labels : list
List of parcellation roi labels extracted from the epo.pkl file generated with MNE
output_types : ['tsv', 'gpickle', 'mat', 'graphml']
List of output format in which to save the connectome files.
(Default: `None`)
"""
if output_types is None:
output_types = ['tsv']
con_methods = list(con_res.keys())
# Create a graph of n_nodes = shape of the connectivity matrix estimated by MNE
G = nx.Graph(np.ones(con_res[con_methods[0]].shape))
# Update node information
for u, d in G.nodes(data=True):
if ' ' in roi_labels[u]: # Cortical-only labels generated by MNE
label_split = roi_labels[u].split(' ')
label_name = f'ctx{label_split[1]}-{label_split[0]}'
else: # Sub-cortical and cortical labels extracted from the atlas index/label mapping file
label_name = roi_labels[u]
if "ctx" in label_name:
G.nodes[u]["dn_region"] = 'cortical'
G.nodes[u]["dn_hemisphere"] = 'left' if "-lh" in roi_labels[u] else "right"
else:
G.nodes[u]["dn_region"] = 'subcortical'
G.nodes[u]["dn_hemisphere"] = 'left' if "left" in roi_labels[u] else "right"
G.nodes[u]["dn_fsname"] = label_name
G.nodes[u]["dn_name"] = label_name
G.nodes[u]["dn_multiscaleID"] = int(u)
# TODO: Set position for the node based on the mean position of the
# ROI in voxel coordinates (segmentation volume )
# G.nodes[u]["dn_position_x"] = ...
# G.nodes[u]["dn_position_y"] = ...
# G.nodes[u]["dn_position_z"] = ...
# Update edge weights
G_out = copy.deepcopy(G)
for u, v, d in G.edges(data=True):
G_out.remove_edge(u, v)
edge = {}
for method in con_methods:
val = float(con_res[method][int(v), int(u)])
edge[method] = val
G_out.add_edge(u, v)
for key in edge:
G_out[u][v][key] = float(edge[key])
# Change w.r.t networkx2
edge_keys = []
for u, v, d in G_out.edges(data=True):
edge_keys = list(d.keys())
break
# Save the connectome file
con_basepath = os.path.join(
output_dir,
output_basename
)
# In TSV format by default to be BIDS compliant
print(f"Save {con_basepath}.tsv...")
# Write header fields
with open(f"{con_basepath}.tsv", "w") as out_file:
tsv_writer = csv.writer(out_file, delimiter="\t")
header = ["source", "target"]
header = header + [key for key in edge_keys]
tsv_writer.writerow(header)
# Write list of graph edges with all connectivity metrics (edge_keys)
with open(f"{con_basepath}.tsv", "ab") as out_file:
nx.write_edgelist(
G_out,
out_file,
comments="#",
delimiter="\t",
data=edge_keys,
encoding="utf-8",
)
# In GPickle format
if "gpickle" in output_types:
# Storing network/graph in gpickle that might be prefered by the user
print(f"Save {con_basepath}.gpickle...")
nx.write_gpickle(G_out, f"{con_basepath}.gpickle")
# In MAT format
if "mat" in output_types:
edge_struct = {}
for edge_key in edge_keys:
edge_struct[edge_key] = nx.to_numpy_matrix(G_out, weight=edge_key)
# nodes
size_nodes = len(list(G_out.nodes(data=True)))
# Get the node attributes/keys from the first node and then break.
# Change w.r.t networkx2
for u, d in G_out.nodes(data=True):
node_keys = list(d.keys())
break
node_struct = {}
for node_key in node_keys:
if node_key == "dn_position":
node_arr = np.zeros([size_nodes, 3], dtype=np.float)
else:
node_arr = np.zeros(size_nodes, dtype=np.object_)
node_n = 0
for _, node_data in G_out.nodes(data=True):
node_arr[node_n] = node_data[node_key]
node_n += 1
node_struct[node_key] = node_arr
print(f"Save {con_basepath}.mat...")
sio.savemat(
f"{con_basepath}.mat",
long_field_names=True,
mdict={"fc": edge_struct, "nodes": node_struct},
)
# In GRAPHML format
if "graphml" in output_types:
g2 = nx.Graph()
for u_gml, v_gml, d_gml in G_out.edges(data=True):
g2.add_edge(u_gml, v_gml)
for key in d_gml:
g2[u_gml][v_gml][key] = d_gml[key]
for u_gml, d_gml in G_out.nodes(data=True):
g2.add_node(u_gml)
g2.nodes[u_gml]["dn_multiscaleID"] = d_gml["dn_multiscaleID"]
g2.nodes[u_gml]["dn_fsname"] = d_gml["dn_fsname"]
g2.nodes[u_gml]["dn_hemisphere"] = d_gml["dn_hemisphere"]
g2.nodes[u_gml]["dn_name"] = d_gml["dn_name"]
g2.nodes[u_gml]["dn_region"] = d_gml["dn_region"]
print(f"Save {con_basepath}.graphml...")
nx.write_graphml(g2, f"{con_basepath}.graphml")