# This file contains functions that could be used in different parts of the project import numpy as np import networkx as nx import os import pickle import sys import matplotlib.pyplot as plt import csv import scipy from scipy.sparse import csr_matrix import copy import seaborn as sns import pandas as pd import random from collections import Counter import time from heapq import heappop, heappush import itertools from collections import deque BASE_DIR = os.path.dirname(os.path.abspath(__file__)) # These functions are used to handle the reservoir during computation def compute(x, W, Win, u, alpha): """ Computes a single step of the neuron activity """ W_spr = W x_spr = csr_matrix(x) Win_spr = csr_matrix(Win) u_spr = csr_matrix(u) newx_spr = np.tanh(W_spr @ x_spr + Win_spr @ u_spr) newx = newx_spr.toarray() newx = (1-alpha)*x + newx*alpha return newx def step_compute(x, W, Win, Wout, u_in, alpha, mask): """ Computes a single step of the neuron activity and returns the output y u_in should be the raw input vector WITHOUT the bias term. """ # Reshape u_in to make sure it's a column vector (Nu, 1) u_in = np.array(u_in).reshape(-1, 1) # Append bias 1 for the reservoir compute u_in_with_bias = np.vstack((u_in, [[1]])) # Calculate new reservoir state newx = compute(x, W, Win, u_in_with_bias, alpha) # Calculate output y = Wout @ [1; x[mask]] (Skip connection u_in removed) X_col = np.vstack((np.ones((1,1)), newx[mask].reshape(-1, 1))) y = Wout @ X_col return newx, y def simulate(x, W, Win, u, alpha, XX = None, XY = None, Yhat = None, building_matrices = False, mask=None): """ Sumulates the neural response to the entire input """ #time0 = time.time() # Simulates the Echo State Network over the input u (L x (1+Nu)) if building_matrices: # this makes the time scale linearly with respect to the number of samples if mask is None: raise("Need to provide a mask for this step") if XX is None or XY is None or Yhat is None: raise("You need to provide the the basic matrices to build them incremetally") retval = [] newx = copy.deepcopy(x) for u_idx, u_in in enumerate(u): #print(time.time()) newx = compute(newx, W, Win, u_in.reshape((u.shape[1],-1)), alpha) retval.append(newx) #print(time.time()) # Skip connection removed newx_reshaped = np.hstack((np.ones((1,1)), newx[mask].reshape(1,-1))) # Assuming newx_reshaped and Yhat are numpy arrays newx_reshaped_col = newx_reshaped.reshape(-1, 1) newx_reshaped_row = newx_reshaped.reshape(1, -1) Yhat_col = Yhat[:, u_idx].reshape(-1, 1) #print(time.time()) XX += np.dot(newx_reshaped_col, newx_reshaped_row) XY += np.dot(Yhat_col, newx_reshaped_row) #print(time.time()) retval = np.array(retval).reshape((-1,len(x))).transpose() #print(f"computation in {time.time() - time0} : {time.time()} - {time0}") return XX, XY, retval else: retval = [] newx = copy.deepcopy(x) for u_in in u: newx = compute(newx, W, Win, u_in.reshape((u.shape[1],-1)), alpha) retval.append(newx) retval = np.array(retval).reshape((-1,len(x))).transpose() #print(f"computation in {time.time() - time0} : {time.time()} - {time0}") return retval # These functions are used to handle the complete graph of the brain and cut it down into smaller subgraphs def create_connectivity_matrix(num_neu, graph_folder_name, sel_crit, biologically_accurate=False, showing_figure = False, make_comparison = True): # Creates the connectivity matrix ad W, and also W_in, W_out and biases. Saves them in the current directory, ready to be used by main.ipynb # get cells classes, that will be useful for the input definition csv_file = os.path.join(BASE_DIR, 'classes_by_cell_type.csv') input_cell_types = ['olfactory', 'visual', 'mechanosensory', 'hygrosensory', 'unknown_sensory', 'ocellar', 'gustatory', 'thermosensory'] output_cell_types = [ "MBON", "DAN", "LHCENT", "clock", "pars_intercerebralis", "pars_lateralis", "Kenyon_Cell", "ALON", "LOP>ME", "LOP>LO.ME", "LOP>LO", "LOP", "TuBu" ] data = [] with open(csv_file, 'r') as file: csv_reader = csv.reader(file) header = next(csv_reader) # Skip the header row id_index = header.index('pt_root_id') cell_type_index = header.index('cell_type') for row in csv_reader: cell_id = row[id_index] cell_type = row[cell_type_index] data += [(cell_id, cell_type)] data_dict = {str(cell_id): cell_type for cell_id, cell_type in data} # Load the network from a file file_path = os.path.join(BASE_DIR ,'networks_graphs', graph_folder_name) G = nx.read_graphml(os.path.join( file_path,'graph.graphml')) unique_neurons = set(G.nodes()) # Load the biases from a file with open(os.path.join( file_path,'biases.pkl'), 'rb') as file: biases = pickle.load(file) F, theta = selection_criterion(G=G, num_neu=num_neu, unique_neurons=unique_neurons, biases=biases, data_dict=data_dict, mode=sel_crit) first_elements = {int(u) for u, v in F.edges} second_elements = {int(v) for u, v in F.edges} only_first = first_elements - second_elements only_second = second_elements - first_elements both = first_elements & second_elements if showing_figure: # Create a figure with 4 subplots fig, axs = plt.subplots(2, 2, figsize=(12, 8)) # Degree distribution degrees = [F.degree(n) for n in F.nodes()] axs[0, 0].hist(degrees, bins=range(min(degrees), max(degrees) + 1), edgecolor='black') axs[0, 0].set_title('Degree Distribution') axs[0, 0].set_xlabel('Degree') axs[0, 0].set_ylabel('Frequency') # Clustering coefficient distribution clustering_coeffs = nx.clustering(F).values() axs[0, 1].hist(clustering_coeffs, bins=10, edgecolor='black') axs[0, 1].set_title('Clustering Coefficient Distribution') axs[0, 1].set_xlabel('Clustering Coefficient') axs[0, 1].set_ylabel('Frequency') # Shortest path length distribution if nx.is_connected(F): path_lengths = dict(nx.all_pairs_shortest_path_length(F)) lengths = [] for source in path_lengths: for target in path_lengths[source]: if source != target: lengths.append(path_lengths[source][target]) axs[1, 0].hist(lengths, bins=range(min(lengths), max(lengths) + 1), edgecolor='black', align='left') axs[1, 0].set_title('Shortest Path Length Distribution') axs[1, 0].set_xlabel('Path Length') axs[1, 0].set_ylabel('Frequency') else: axs[1, 0].text(0.5, 0.5, "The graph is not connected,\nso shortest path lengths cannot be computed for all pairs of nodes.", horizontalalignment='center', verticalalignment='center', transform=axs[1, 0].transAxes) # Connectivity of Unique IDs labels = ['Only as Pre-syn', 'Only as Post-syn', 'Both'] sizes = [len(only_first), len(only_second), len(both)] axs[1, 1].bar(labels, sizes, color=['blue', 'orange', 'green']) axs[1, 1].set_title('Connectivity of Unique Neurons') axs[1, 1].set_xlabel('Category') axs[1, 1].set_ylabel('Number of Unique Neurons') # Adjust the spacing between subplots plt.tight_layout() # Show the figure plt.show() # Save the figure in the /graph_stats/ folder if not os.path.exists('graph_stats'): os.makedirs('graph_stats') graph_size = len(F.nodes()) figure_path = os.path.join('graph_stats', f'graph_stats_size_{graph_size}.png') fig.savefig(figure_path) print(f"Figure saved to {figure_path}") if make_comparison: # Create a graph D with the same sparsity as G but with randomly placed edges num_nodes = len(F.nodes()) num_edges = len(F.edges()) # Generate a random graph with the same number of nodes and edges D = nx.gnm_random_graph(num_nodes, num_edges) while not nx.is_connected(D): D = nx.gnm_random_graph(num_nodes, num_edges) # Relabel the nodes of D to match the node labels of F mapping = {i: node for i, node in enumerate(F.nodes())} D = nx.relabel_nodes(D, mapping) # Create a figure with 4 subplots fig, axs = plt.subplots(2, 2, figsize=(12, 8)) # Degree distribution degrees = [D.degree(n) for n in D.nodes()] axs[0, 0].hist(degrees, bins=range(min(degrees), max(degrees) + 1), edgecolor='black') axs[0, 0].set_title('Degree Distribution') axs[0, 0].set_xlabel('Degree') axs[0, 0].set_ylabel('Frequency') # Clustering coefficient distribution clustering_coeffs = nx.clustering(D).values() axs[0, 1].hist(clustering_coeffs, bins=10, edgecolor='black') axs[0, 1].set_title('Clustering Coefficient Distribution') axs[0, 1].set_xlabel('Clustering Coefficient') axs[0, 1].set_ylabel('Frequency') # Shortest path length distribution if nx.is_connected(D): path_lengths = dict(nx.all_pairs_shortest_path_length(D)) lengths = [] for source in path_lengths: for target in path_lengths[source]: if source != target: lengths.append(path_lengths[source][target]) axs[1, 0].hist(lengths, bins=range(min(lengths), max(lengths) + 1), edgecolor='black', align='left') axs[1, 0].set_title('Shortest Path Length Distribution') axs[1, 0].set_xlabel('Path Length') axs[1, 0].set_ylabel('Frequency') else: axs[1, 0].text(0.5, 0.5, "The graph is not connected,\nso shortest path lengths cannot be computed for all pairs of nodes.", horizontalalignment='center', verticalalignment='center', transform=axs[1, 0].transAxes) first_elements_D = {int(u) for u, v in D.edges} second_elements_D = {int(v) for u, v in D.edges} only_first_D = first_elements_D - second_elements_D only_second_D = second_elements_D - first_elements_D both_D = first_elements_D & second_elements_D # Connectivity of Unique IDs labels = ['Only as Pre-syn', 'Only as Post-syn', 'Both'] sizes = [len(only_first_D ), len(only_second_D ), len(both_D )] axs[1, 1].bar(labels, sizes, color=['blue', 'orange', 'green']) axs[1, 1].set_title('Connectivity of Unique Neurons') axs[1, 1].set_xlabel('Category') axs[1, 1].set_ylabel('Number of Unique Neurons') # Adjust the spacing between subplots plt.tight_layout() # Show the figure plt.show() # Save the figure in the /graph_stats/ folder if not os.path.exists('graph_stats'): os.makedirs('graph_stats') graph_size = len(F.nodes()) figure_path = os.path.join('graph_stats', f'graph_stats_size_{graph_size}_random_perm.png') fig.savefig(figure_path) print(f"Figure saved to {figure_path}") # Convert the connectivity matrix to a sparse matrix connectivity_matrix = scipy.sparse.csr_matrix(nx.to_scipy_sparse_array(F, weight='weight')) target_radius = 1 spectral_radius, _ = scipy.sparse.linalg.eigs(connectivity_matrix, k=1, which='LM') # Rescale the connectivity matrix theta = theta.astype('float') rescaled_matrix = connectivity_matrix * target_radius/np.linalg.norm(spectral_radius) theta *= target_radius/abs(float(np.linalg.norm(spectral_radius))) print("Initial Spectral Radius:", spectral_radius) # Keep scaling the matrix down until the spectral radius is below 1 scaling_factor = .99 while True: spectral_radius, _ = scipy.sparse.linalg.eigs(rescaled_matrix, k=1, which='LM') print(f"\rSpectral Radius: {np.linalg.norm(spectral_radius)}", end='', flush=True) if np.linalg.norm(spectral_radius) < 1: break rescaled_matrix *= scaling_factor theta *= scaling_factor print(f"\n") # Create W_in array W_in = np.zeros((len(F.nodes()) , 1)) # Update W_in based on input_cell_types for i, neuron_id in enumerate(F.nodes()): if int(neuron_id) in [int(x) for x in list(data_dict.keys())] and data_dict[neuron_id] in input_cell_types: W_in[i] = 1 if not biologically_accurate or np.max(W_in) == 0: print("No input neurons found in the selected neurons, switching to first element only neurons.") for i, neuron_id in enumerate(F.nodes()): if int(neuron_id) in only_first: W_in[i] = 1 if np.max(W_in) == 0: print("no first element only neurons found") W_in = None W_out = np.zeros((len(F.nodes()), 1)) # Update W_out based on output_cell_types for i, neuron_id in enumerate(F.nodes()): if int(neuron_id) in [int(x) for x in list(data_dict.keys())] and data_dict[neuron_id] in output_cell_types: W_out[i] = 1 if not biologically_accurate or np.max(W_out) == 0: print("No output neurons found in the selected neurons, switching to second element only neurons.") for i, neuron_id in enumerate(F.nodes()): if int(neuron_id) in only_second: W_out[i] = 1 if np.max(W_out) == 0: print("no second element only neurons found") W_out = None # Save W_out to a file called W_out with open(os.path.join(BASE_DIR, 'W_out.pkl'), 'wb') as file: pickle.dump(W_out, file) # Save connectivity_matrix to a file called W with open(os.path.join(BASE_DIR, 'W.pkl'), 'wb') as file: pickle.dump(rescaled_matrix, file) # Save W_in to a file called W_in with open(os.path.join(BASE_DIR, 'W_in.pkl'), 'wb') as file: pickle.dump(W_in, file) # Save theta to a file called theta.pkl with open(os.path.join(BASE_DIR, 'bias.pkl'), 'wb') as file: pickle.dump(theta, file) # This function is similar to the one above, but it does not save the matrices in src/** and only returns the neurons IDs and the unique cell types def get_neurons_id(num_neu, mode, graph_folder_name): # get cells classes, that will be useful for the input definition csv_file = os.path.join(BASE_DIR, 'classes_by_cell_type.csv') data = [] with open(csv_file, 'r') as file: csv_reader = csv.reader(file) header = next(csv_reader) # Skip the header row id_index = header.index('pt_root_id') cell_type_index = header.index('cell_type') for row in csv_reader: cell_id = row[id_index] cell_type = row[cell_type_index] data += [(cell_id, cell_type)] data_dict = {str(cell_id): cell_type for cell_id, cell_type in data} # Load the network from a file file_path = os.path.join(BASE_DIR, 'networks_graphs', graph_folder_name) G = nx.read_graphml(os.path.join( file_path,'graph.graphml')) unique_neurons = set(G.nodes()) # Calculate the total number of edges for each neuron edge_counts = {neuron: G.degree(neuron) for neuron in unique_neurons} # Sort the neuron IDs based on the edge counts in descending order sorted_neurons = sorted(unique_neurons, key=lambda neuron: edge_counts[neuron], reverse=True) # Create a sorted list of neuron IDs neuron_id_list = list(sorted_neurons) F, _ = selection_criterion(G=G, num_neu=num_neu, unique_neurons=unique_neurons, mode=mode, biases=None, data_dict=data_dict) neuron_ids_in_F = list(F.nodes) unique_cell_types_in_F = set(data_dict[neuron_id] for neuron_id in neuron_ids_in_F if neuron_id in data_dict) return neuron_ids_in_F, unique_cell_types_in_F def selection_criterion(G, num_neu, unique_neurons, data_dict, biases = None, mode='connectivity_first', params=None): #This function selects a smaller subset of N neurons from the full NetworkX graph G, and returns the NetworkX graph F, containing N nodes if 'connectivity_first' in mode: # Calculate the total number of edges for each neuron edge_counts = {neuron: G.degree(neuron) for neuron in unique_neurons} # Sort the neuron IDs based on the edge counts in descending order sorted_neurons = sorted(unique_neurons, key=lambda neuron: edge_counts[neuron], reverse=True) if biases is not None: sorted_biases = np.array([biases[nrn] for nrn in sorted_neurons]) # Create a sorted list of neuron IDs neuron_id_list = list(sorted_neurons) num_neuron = num_neu # Create a subgraph of G with the first num_neurons neurons F = G.subgraph(neuron_id_list[:num_neuron]) # Create a copy of the subgraph F = F.copy() if '2' in mode: if not nx.is_connected(F): print("Graph not fully connected, connecting the components") components = list(nx.connected_components(F)) meta_graph = nx.Graph() for i, comp in enumerate(components): meta_graph.add_node(i) candidate_nodes = set(G.nodes()) - set(F.nodes()) best_addition = set() F = min_node_connected_subgraph(G, components) print("Pruning components") # Check the number of nodes in F while len(F.nodes()) > num_neuron: # Find the least connected node sorted_nodes = sorted(F.nodes, key=lambda node: F.degree(node)) for node in sorted_nodes: F_tentative = F.copy() F_tentative.remove_node(node) if nx.is_connected(F_tentative): F = F_tentative break print(f"\rCurrent number of nodes: {len(F.nodes())}", end='', flush=True) print("Adding components") while len(F.nodes()) < num_neuron: # Find the node in G that is most connected to nodes already in F candidate_nodes = set(G.nodes()) - set(F.nodes()) best_node = max(candidate_nodes, key=lambda node: len(set(G.neighbors(node)) & set(F.nodes()))) F.add_node(best_node) for neighbor in G.neighbors(best_node): if neighbor in F.nodes: F.add_edge(best_node, neighbor, weight=G[best_node][neighbor]['weight']) else: # Check if the graph is connected if not nx.is_connected(F): # Find the largest connected component largest_component = max(nx.connected_components(F), key=len) # Create a new graph with only the largest connected component F = F.subgraph(largest_component) # Print the number of unique neurons in the new graph print("The graph is not fully connected. Only the largest connected component is considered.") print("Number of unique neurons in the biggest element:", len(F)) if biases is not None: theta = sorted_biases[:len(F)] F = F.copy() if not '2' in mode: # Add neurons connected to F from G according to the order in sorted_neurons until num_neurons is reached for neuron_id in sorted_neurons: if len(F.nodes) >= num_neuron: break if neuron_id not in F.nodes: for neighbor in G.neighbors(neuron_id): if neighbor in F.nodes: F.add_node(neuron_id) F.add_edge(neuron_id, neighbor, weight=G[neuron_id][neighbor]['weight']) if biases is not None: theta += [biases[neuron_id]] break print("Number of unique neurons after adding neighbors:", len(F)) elif 'proportional_selection' in mode: # Find the largest connected component if G.is_directed(): largest_component = max(nx.weakly_connected_components(G), key=len) else: largest_component = max(nx.connected_components(G), key=len) # Create a new graph with only the largest connected component G = G.subgraph(largest_component) # Remove keys from data_dict that are no longer in G data_dict = {k: v for k, v in data_dict.items() if k in G.nodes()} if num_neu > len(G.nodes()): print(f"Number of neurons requested is greater than the number of neurons in the graph, setting number of neuron to {G.nodes()}") return None class_count = Counter(data_dict.values()) D = [] nodes_to_keep = set() for key_idx, key in enumerate(class_count.keys()): fraction = class_count[key] / len(data_dict) target_class_ids = [k for k, v in data_dict.items() if v == key] edge_counts = {neuron: G.degree(neuron) for neuron in target_class_ids} try: sorted_target_class_ids = sorted(target_class_ids, key=lambda neuron: edge_counts[neuron], reverse=True) except: print(f"Class {key} has no neurons in the graph") num_neuron = num_neu D_temp = G.subgraph(sorted_target_class_ids[:int(np.ceil(fraction*num_neuron))]).copy() D += [D_temp] for node in D_temp.nodes(): nodes_to_keep.add(node) F = G.subgraph(list(nodes_to_keep)).copy() is_conn = nx.is_weakly_connected(F) if F.is_directed() else nx.is_connected(F) if not is_conn: print("Graph not fully connected, quickly connecting the components with synthetic edges") components = list(nx.weakly_connected_components(F) if F.is_directed() else nx.connected_components(F)) # Sort components by size descending components.sort(key=len, reverse=True) # Connect all smaller components to the largest one with a random weak edge largest_comp = list(components[0]) for i in range(1, len(components)): u = np.random.choice(largest_comp) v = np.random.choice(list(components[i])) # Add a weak directed edge F.add_edge(u, v, weight=0.01) print("Pruning components") # Check the number of nodes in F while len(F.nodes()) > num_neuron: # Find the least connected node sorted_nodes = sorted(F.nodes, key=lambda node: F.degree(node)) for node in sorted_nodes: F_tentative = F.copy() F_tentative.remove_node(node) is_conn_tentative = nx.is_weakly_connected(F_tentative) if F_tentative.is_directed() else nx.is_connected(F_tentative) if is_conn_tentative: F = F_tentative break print(f"\rCurrent number of nodes: {len(F.nodes())}", end='', flush=True) print("Adding components") while len(F.nodes()) < num_neuron: # Find the node in G that is most connected to nodes already in F candidate_nodes = set(G.nodes()) - set(F.nodes()) best_node = max(candidate_nodes, key=lambda node: len(set(G.neighbors(node)) & set(F.nodes()))) F.add_node(best_node) for neighbor in G.neighbors(best_node): if neighbor in F.nodes: F.add_edge(best_node, neighbor, weight=G[best_node][neighbor]['weight']) node_ids = list(F.nodes()) if biases is not None: theta = np.array([biases[nrn] for nrn in node_ids]) else: raise ValueError(f"Invalid selection criterion mode: {mode}") if biases is None: theta = None return F, theta def min_node_connected_subgraph(G, components): """ Finds the minimal set of additional nodes needed to connect all given components in G. Parameters: G (networkx.Graph): The input graph. components (list of sets): Each set contains nodes forming a component. Returns: networkx.Graph: The minimal connected subgraph with the fewest extra nodes. """ # Step 1: Identify component representative nodes component_representatives = [next(iter(comp)) for comp in components] # Step 2: Build a shortest-path metric graph based on node count metric_graph = nx.Graph() shortest_paths = {} undirected_G = G.to_undirected(as_view=True) if G.is_directed() else G for u, v in itertools.combinations(component_representatives, 2): try: # First try directed path path = nx.shortest_path(G, source=u, target=v) except nx.NetworkXNoPath: try: # Try the reverse directed path path = nx.shortest_path(G, source=v, target=u) except nx.NetworkXNoPath: # Fallback to undirected path try: path = nx.shortest_path(undirected_G, source=u, target=v) except nx.NetworkXNoPath: continue # No path at all metric_graph.add_edge(u, v, weight=len(path) - 1) shortest_paths[(u, v)] = path # Step 3: Compute MST on the metric graph to ensure minimal connectivity mst = nx.minimum_spanning_tree(metric_graph, weight="weight") # Step 4: Extract corresponding paths from the original graph added_nodes = set() subgraph_edges = set() for u, v in mst.edges: path = shortest_paths[(u, v)] for i in range(len(path) - 1): n1 = path[i] n2 = path[i+1] added_nodes.update([n1, n2]) if G.has_edge(n1, n2): subgraph_edges.add((n1, n2)) elif G.has_edge(n2, n1): subgraph_edges.add((n2, n1)) # Step 5: Construct the minimal connected subgraph H = G.edge_subgraph(subgraph_edges).copy() return H def create_csr_matrix(size, target_sparsity, data_rvs): density = 1 - target_sparsity num_nonzero_elements = int(size * size * density) # Generate random row and column indices for the non-zero elements row_indices = np.random.randint(0, size, num_nonzero_elements) col_indices = np.random.randint(0, size, num_nonzero_elements) # Generate random values using the provided data_rvs function data = data_rvs(num_nonzero_elements) # Create the CSR matrix csr_matrix = scipy.sparse.csr_matrix((data, (row_indices, col_indices)), shape=(size, size)) return csr_matrix def is_fully_connected(W): # Convert the sparse matrix to a NetworkX graph graph = nx.from_scipy_sparse_array(W, create_using=nx.DiGraph if scipy.sparse.isspmatrix_csr(W) else nx.Graph) # Check if the graph is strongly connected (for directed graphs) or connected (for undirected graphs) if isinstance(graph, nx.DiGraph): return nx.is_strongly_connected(graph) else: return nx.is_connected(graph)