diff --git a/.gitignore b/.gitignore index 107baa93..9935c271 100644 --- a/.gitignore +++ b/.gitignore @@ -108,4 +108,31 @@ ENV/ .vscode # lbfgs solver stuff -pyuoi/lbfgs/_lowlevel.c +_lowlevel.c + +*~ +out/ +*.prof +dataDale/ +*.png +*.npy +data/ +model/ +\#* +*.h5 +.cursor +*.npz +old/ +result/ +.bash_history +.viminfo +.ipython/ +.jupyter/ +.local/ +.emacs.d/ +.lesshst +*.aux +*.toc +*.out +.codex +*.dvi diff --git a/causal_net/Readme b/causal_net/Readme new file mode 100644 index 00000000..0cc31984 --- /dev/null +++ b/causal_net/Readme @@ -0,0 +1,80 @@ +Detailed manual at +https://docs.google.com/document/d/1Ou98jsFpTR7s3cTf6AQwF92tqd2guwA_cGuvQrm8RSE/edit?usp=sharing + + +IMG=nersc/causal-net:v3 # Mar 28 +IMG=nersc/causal-net:v4 # May 13 + + export OMP_NUM_THREADS=2 + salloc -q interactive -C cpu --image=$IMG -t 4:00:00 -A m2043 -N 1 + + +Directory structure to stor real and simulated data +/global/cfs/cdirs/m2043/causal_inference/dale-gen-sim-may27/input_uoi/ + +basePath=/global/homes/b/balewski/prjs/bioDataVault2025/causalNet_tmp/ + +exp_raw/ raw data from experiment +gen_dale/ daleMatrix+samples for generated network +input_uoi/ data formatted for UoI +model_uoi/ UoI fit output +postproc/ any final analysis + + +A------ dale matrix generator + simulator +runs on PM w/o image, just load python + + + += = = = = OLD = = = = = = = +Deplyting software - for now as root +Must be repeated for each image start + + +cd pyuoi +pip install -e .[dev] + +cd ../toys/ + +mkdir 2025_UoI-VAR +cd 2025_UoI-VAR +git clone -b uoi-var https://github.com/BouchardLab/pyuoi.git + +# build podman image using podman build -f ubu24-casual-net.dockerfile -t balewski/casual-net:p1 . + + +# start image + +IMG=balewski/casual-net:p1 +WORK_DIR=/2025_UoI-VAR +DATA_VAULT=/shared_volumes/dataVault2025 + +eval podman run -it \ + --volume /shared_volumes/$WORK_DIR:$WORK_DIR \ + --volume $DATA_VAULT:/dataVault2025 \ + --workdir $WORK_DIR \ + $IMG /bin/bash + +# inside image +cd ../pyuoi +cd /2025_UoI-VAR +pip install -e .[dev] + +python < 0: + for i, cycle in enumerate(cycle_edges): + color = colors[i % len(colors)] + cycle_edge_list = [(u, v) for u, v in cycle if G.has_edge(u, v)] + if cycle_edge_list: + nx.draw_networkx_edges(G, pos, edgelist=cycle_edge_list, edge_color=color, width=3, alpha=0.8, style='dashed') + + # Highlight nodes in this cycle + cycle_nodes = set() + for u, v in cycle_edge_list: + cycle_nodes.add(u) + cycle_nodes.add(v) + nx.draw_networkx_nodes(G, pos, nodelist=list(cycle_nodes), node_color=color, node_size=200, alpha=0.8) + + # Add cycle number label to the first edge of the cycle + if cycle_edge_list: + u, v = cycle_edge_list[0] + edge_x = (pos[u][0] + pos[v][0]) / 2 + edge_y = (pos[u][1] + pos[v][1]) / 2 + self.plt.text(edge_x, edge_y, f'{i+1}', fontsize=12, fontweight='bold', + ha='center', va='center', + bbox=dict(boxstyle='circle,pad=0.3', facecolor='white', edgecolor=color, linewidth=1)) + + # Add title and info + self.plt.title(f'Graph Cycles Visualization\n' + f'Nodes: {graph_conf["n_nodes"]}, Edges: {graph_conf["n_edges"]}, ' + f'β₁ = {betti_results["betti_numbers"].get(1, 0)} cycles', + fontsize=14, fontweight='bold') + + # Add legend if there are cycles + if cycle_edges and len(cycle_edges) > 0: + legend_elements = [] + for i in range(min(len(cycle_edges), len(colors))): + color = colors[i % len(colors)] + legend_elements.append(self.plt.Line2D([0], [0], color=color, lw=1, label=f'Cycle {i+1}')) + self.plt.legend(handles=legend_elements, loc='upper right') + + self.plt.axis('off') + self.plt.tight_layout() + +#...!...!.................. + def plot_voids_3d(self, grfD, grfMD, figId=3): + """Plot 3D voids as tetrahedral projections""" + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(12,8)) + ax = self.plt.gca() + + # Extract data from structures + edges = grfD['edges'] + void_tetrahedra = grfD.get('void_tetrahedra', []) + graph_conf = grfMD['graph_conf'] + betti_results = grfMD['betti_results'] + n_nodes = graph_conf['n_nodes'] + + # Create NetworkX graph + G = nx.Graph() + G.add_nodes_from(range(n_nodes)) + G.add_edges_from(edges) + + # Generate layout + pos = nx.spring_layout(G, seed=42, k=1, iterations=50) + + # Draw all nodes (smaller, no labels) + nx.draw_networkx_nodes(G, pos, node_color='lightgray', node_size=150, alpha=0.7) + + # Draw all edges in light gray + nx.draw_networkx_edges(G, pos, edge_color='lightgray', alpha=0.3, width=1) + + # Color palette for 3D voids + void3d_colors = ['red', 'blue', 'green', 'orange', 'purple', 'brown', 'pink', 'gray', 'olive', 'navy'] + + # Plot 3D voids as tetrahedral wireframes (projected to 2D) + if void_tetrahedra and len(void_tetrahedra) > 0: + for i, tetrahedra_group in enumerate(void_tetrahedra): + color = void3d_colors[i % len(void3d_colors)] + + # Plot each tetrahedron in the group + for tetrahedron in tetrahedra_group: + if len(tetrahedron) == 4: + # Get positions of tetrahedron vertices + tetra_pos = [pos[node] for node in tetrahedron if node in pos] + if len(tetra_pos) == 4: + # Draw all 6 edges of the tetrahedron + tetra_edges = [] + for j in range(4): + for k in range(j+1, 4): + tetra_edges.append((tetrahedron[j], tetrahedron[k])) + + # Highlight tetrahedral edges with thick dash-dot lines + for edge in tetra_edges: + if G.has_edge(edge[0], edge[1]): + nx.draw_networkx_edges(G, pos, edgelist=[edge], + edge_color=color, width=4, alpha=0.9, + style=(0, (3, 1, 1, 1))) # dash-dot pattern + + # Draw tetrahedral faces as transparent polygons + # Each tetrahedron has 4 triangular faces + faces = [ + [tetrahedron[0], tetrahedron[1], tetrahedron[2]], + [tetrahedron[0], tetrahedron[1], tetrahedron[3]], + [tetrahedron[0], tetrahedron[2], tetrahedron[3]], + [tetrahedron[1], tetrahedron[2], tetrahedron[3]] + ] + + for face in faces: + face_pos = [pos[node] for node in face if node in pos] + if len(face_pos) == 3: + face_x = [p[0] for p in face_pos] + [face_pos[0][0]] + face_y = [p[1] for p in face_pos] + [face_pos[0][1]] + self.plt.fill(face_x, face_y, color=color, alpha=0.1, + edgecolor=color, linewidth=0.5, linestyle='--') + + # Add 3D void number label at centroid of first tetrahedron + if tetrahedra_group: + first_tetrahedron = tetrahedra_group[0] + if len(first_tetrahedron) == 4: + centroid_x = sum(pos[node][0] for node in first_tetrahedron if node in pos) / 4 + centroid_y = sum(pos[node][1] for node in first_tetrahedron if node in pos) / 4 + self.plt.text(centroid_x, centroid_y, f'3D{i+1}', fontsize=16, fontweight='bold', + ha='center', va='center', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', edgecolor=color, linewidth=1)) + + # Add title and info + self.plt.title(f'Graph 3D Voids Visualization\n' + f'Nodes: {graph_conf["n_nodes"]}, Edges: {graph_conf["n_edges"]}, ' + f'β₃ = {betti_results["betti_numbers"].get(3, 0)} voids', + fontsize=14, fontweight='bold') + + # Add legend if there are 3D voids + if void_tetrahedra and len(void_tetrahedra) > 0: + legend_elements = [] + for i in range(min(len(void_tetrahedra), len(void3d_colors))): + color = void3d_colors[i % len(void3d_colors)] + legend_elements.append(self.plt.Line2D([0], [0], color=color, lw=4, + linestyle=(0, (3, 1, 1, 1)), label=f'3D Void {i+1}')) + self.plt.legend(handles=legend_elements, loc='upper right') + + self.plt.axis('off') + self.plt.tight_layout() + +#...!...!.................. + def plot_voids_2d(self, grfD, grfMD, figId=2): + """Plot 2D voids (cavities) as triangular faces""" + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(12,8)) + ax = self.plt.gca() + + # Extract data from structures + edges = grfD['edges'] + void_triangles = grfD.get('void_triangles', []) + graph_conf = grfMD['graph_conf'] + betti_results = grfMD['betti_results'] + n_nodes = graph_conf['n_nodes'] + + # Create NetworkX graph + G = nx.Graph() + G.add_nodes_from(range(n_nodes)) + G.add_edges_from(edges) + + # Generate layout + pos = nx.spring_layout(G, seed=42, k=1, iterations=50) + + # Draw all nodes (smaller, no labels) + nx.draw_networkx_nodes(G, pos, node_color='lightgray', node_size=150, alpha=0.7) + + # Draw all edges in light gray + nx.draw_networkx_edges(G, pos, edge_color='lightgray', alpha=0.3, width=1) + + # Color palette for voids + void_colors = ['magenta', 'cyan', 'yellow', 'lime', 'orange', 'purple', 'pink', 'brown', 'olive', 'navy'] + + # Plot 2D voids as filled triangles + if void_triangles and len(void_triangles) > 0: + for i, triangle_group in enumerate(void_triangles): + color = void_colors[i % len(void_colors)] + + # Plot each triangle in the group + for triangle in triangle_group: + if len(triangle) == 3: + # Get positions of triangle vertices + triangle_pos = [pos[node] for node in triangle if node in pos] + if len(triangle_pos) == 3: + # Create triangle coordinates + triangle_x = [p[0] for p in triangle_pos] + [triangle_pos[0][0]] + triangle_y = [p[1] for p in triangle_pos] + [triangle_pos[0][1]] + + # Fill triangle with transparency + self.plt.fill(triangle_x, triangle_y, color=color, alpha=0.3, edgecolor=color, linewidth=1) + + # Highlight triangle edges + for j in range(3): + u, v = triangle[j], triangle[(j+1) % 3] + if G.has_edge(u, v): + nx.draw_networkx_edges(G, pos, edgelist=[(u, v)], + edge_color=color, width=3, alpha=0.8, style='dotted') + + # Add void number label at centroid of first triangle + if triangle_group: + first_triangle = triangle_group[0] + if len(first_triangle) == 3: + centroid_x = sum(pos[node][0] for node in first_triangle if node in pos) / 3 + centroid_y = sum(pos[node][1] for node in first_triangle if node in pos) / 3 + self.plt.text(centroid_x, centroid_y, f'V{i+1}', fontsize=14, fontweight='bold', + ha='center', va='center', + bbox=dict(boxstyle='square,pad=0.3', facecolor='white', edgecolor=color, linewidth=1)) + + # Add title and info + self.plt.title(f'Graph 2D Voids Visualization\n' + f'Nodes: {graph_conf["n_nodes"]}, Edges: {graph_conf["n_edges"]}, ' + f'β₂ = {betti_results["betti_numbers"].get(2, 0)} voids', + fontsize=14, fontweight='bold') + + # Add legend if there are voids + if void_triangles and len(void_triangles) > 0: + legend_elements = [] + for i in range(min(len(void_triangles), len(void_colors))): + color = void_colors[i % len(void_colors)] + legend_elements.append(patches.Patch(color=color, alpha=0.3, label=f'Void {i+1}')) + self.plt.legend(handles=legend_elements, loc='upper right') + + self.plt.axis('off') + self.plt.tight_layout() + +#...!...!.................. + def plot_voids_3d(self, grfD, grfMD, figId=3): + """Plot 3D voids as tetrahedral projections""" + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(12,8)) + ax = self.plt.gca() + + # Extract data from structures + edges = grfD['edges'] + void_tetrahedra = grfD.get('void_tetrahedra', []) + graph_conf = grfMD['graph_conf'] + betti_results = grfMD['betti_results'] + n_nodes = graph_conf['n_nodes'] + + # Create NetworkX graph + G = nx.Graph() + G.add_nodes_from(range(n_nodes)) + G.add_edges_from(edges) + + # Generate layout + pos = nx.spring_layout(G, seed=42, k=1, iterations=50) + + # Draw all nodes (smaller, no labels) + nx.draw_networkx_nodes(G, pos, node_color='lightgray', node_size=150, alpha=0.7) + + # Draw all edges in light gray + nx.draw_networkx_edges(G, pos, edge_color='lightgray', alpha=0.3, width=1) + + # Color palette for 3D voids + void3d_colors = ['red', 'blue', 'green', 'orange', 'purple', 'brown', 'pink', 'gray', 'olive', 'navy'] + + # Plot 3D voids as tetrahedral wireframes (projected to 2D) + if void_tetrahedra and len(void_tetrahedra) > 0: + for i, tetrahedra_group in enumerate(void_tetrahedra): + color = void3d_colors[i % len(void3d_colors)] + + # Plot each tetrahedron in the group + for tetrahedron in tetrahedra_group: + if len(tetrahedron) == 4: + # Get positions of tetrahedron vertices + tetra_pos = [pos[node] for node in tetrahedron if node in pos] + if len(tetra_pos) == 4: + # Draw all 6 edges of the tetrahedron + tetra_edges = [] + for j in range(4): + for k in range(j+1, 4): + tetra_edges.append((tetrahedron[j], tetrahedron[k])) + + # Highlight tetrahedral edges with thick dash-dot lines + for edge in tetra_edges: + if G.has_edge(edge[0], edge[1]): + nx.draw_networkx_edges(G, pos, edgelist=[edge], + edge_color=color, width=4, alpha=0.9, + style=(0, (3, 1, 1, 1))) # dash-dot pattern + + # Draw tetrahedral faces as transparent polygons + # Each tetrahedron has 4 triangular faces + faces = [ + [tetrahedron[0], tetrahedron[1], tetrahedron[2]], + [tetrahedron[0], tetrahedron[1], tetrahedron[3]], + [tetrahedron[0], tetrahedron[2], tetrahedron[3]], + [tetrahedron[1], tetrahedron[2], tetrahedron[3]] + ] + + for face in faces: + face_pos = [pos[node] for node in face if node in pos] + if len(face_pos) == 3: + face_x = [p[0] for p in face_pos] + [face_pos[0][0]] + face_y = [p[1] for p in face_pos] + [face_pos[0][1]] + self.plt.fill(face_x, face_y, color=color, alpha=0.1, + edgecolor=color, linewidth=0.5, linestyle='--') + + # Add 3D void number label at centroid of first tetrahedron + if tetrahedra_group: + first_tetrahedron = tetrahedra_group[0] + if len(first_tetrahedron) == 4: + centroid_x = sum(pos[node][0] for node in first_tetrahedron if node in pos) / 4 + centroid_y = sum(pos[node][1] for node in first_tetrahedron if node in pos) / 4 + self.plt.text(centroid_x, centroid_y, f'3D{i+1}', fontsize=16, fontweight='bold', + ha='center', va='center', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', edgecolor=color, linewidth=1)) + + # Add title and info + self.plt.title(f'Graph 3D Voids Visualization\n' + f'Nodes: {graph_conf["n_nodes"]}, Edges: {graph_conf["n_edges"]}, ' + f'β₃ = {betti_results["betti_numbers"].get(3, 0)} voids', + fontsize=14, fontweight='bold') + + # Add legend if there are 3D voids + if void_tetrahedra and len(void_tetrahedra) > 0: + legend_elements = [] + for i in range(min(len(void_tetrahedra), len(void3d_colors))): + color = void3d_colors[i % len(void3d_colors)] + legend_elements.append(self.plt.Line2D([0], [0], color=color, lw=2, + linestyle=(0, (3, 1, 1, 1)), label=f'3D Void {i+1}')) + self.plt.legend(handles=legend_elements, loc='upper right') + + self.plt.axis('off') + self.plt.tight_layout() diff --git a/causal_net/betti_numbers/betti_gen_ana2.py b/causal_net/betti_numbers/betti_gen_ana2.py new file mode 100755 index 00000000..f01e236e --- /dev/null +++ b/causal_net/betti_numbers/betti_gen_ana2.py @@ -0,0 +1,517 @@ +#!/usr/bin/env python3 + +import numpy as np +import argparse +import time +from typing import Dict, Tuple, List +from dataclasses import dataclass +import json +import networkx as nx +from PlotterBetti import Plotter + +try: + import gudhi +except ImportError as exc: + gudhi = None + _GUDHI_IMPORT_ERROR = exc +else: + _GUDHI_IMPORT_ERROR = None + +@dataclass +class BettiResults: + """Container for Betti number computation results""" + n_nodes: int + edge_probability: float + n_edges: int + betti_numbers: Dict[int, int] + computation_time: float + n_simplices: Dict[int, int] + cycle_edges: List[List[Tuple[int, int]]] = None # List of cycles, each cycle is a list of edges + +class RandomGraphGenerator: + """CPU-based random graph generator""" + + def __init__(self, n_nodes: int, edge_prob: float, seed: int = None): + self.n_nodes = n_nodes + self.edge_prob = edge_prob + if seed is not None: + np.random.seed(seed) + + def generate(self) -> Tuple[np.ndarray, int]: + """ + Generate random graph using CPU + Returns: (edges_array, n_edges) + """ + # Generate upper triangular part only (no self-loops) + n_possible = self.n_nodes * (self.n_nodes - 1) // 2 + + # Generate random values for all possible edges + random_vals = np.random.random(n_possible) + edge_mask = random_vals < self.edge_prob + + # Convert to edge list + edges = self._mask_to_edges(edge_mask) + n_edges = len(edges) + + return edges, n_edges + + def _mask_to_edges(self, edge_mask: np.ndarray) -> np.ndarray: + """Convert boolean mask to edge list""" + edge_indices = np.where(edge_mask)[0] + n_edges = len(edge_indices) + + if n_edges == 0: + return np.empty((0, 2), dtype=np.int32) + + edges = np.empty((n_edges, 2), dtype=np.int32) + + # Convert linear indices to (i,j) pairs + # Using inverse of k = i*n - i*(i+1)/2 + j - i - 1 + for idx, k in enumerate(edge_indices): + k = int(k) + # Solve for i: i = floor((2*n - 1 - sqrt((2*n-1)^2 - 8*k))/2) + i = int((2*self.n_nodes - 1 - np.sqrt((2*self.n_nodes-1)**2 - 8*k)) / 2) + j = k - i*self.n_nodes + i*(i+1)//2 + i + 1 + edges[idx, 0] = i + edges[idx, 1] = j + + return edges + +class BettiComputer: + """Compute Betti numbers using GUDHI""" + + def __init__(self, max_dimension: int = 3): + if gudhi is None: + raise ImportError( + "GUDHI is required for Betti number computation. " + "Install it in this environment before running this script." + ) from _GUDHI_IMPORT_ERROR + self.max_dimension = max_dimension + + def compute(self, edges: np.ndarray, n_nodes: int) -> Tuple[Dict[int, int], Dict[int, int], List[List[Tuple[int, int]]], List[List[List[int]]], List[List[List[int]]]]: + """ + Compute Betti numbers up to max_dimension + Returns: (betti_numbers, simplex_counts, cycle_edges, void_triangles, void_tetrahedra) + """ + # Build simplex tree + st = gudhi.SimplexTree() + + # Add vertices + for v in range(n_nodes): + st.insert([v]) + + # Add edges + for edge in edges: + st.insert(edge.tolist()) + + # CRITICAL FIX: Only expand if we actually want higher dimensional features + # For dense graphs, expansion creates too many simplices + if self.max_dimension > 1: + # Use flag complex for clique detection, but limit dimension + # This is more memory efficient than full expansion + print(f" Building flag complex up to dimension {self.max_dimension}...") + st.expansion(self.max_dimension + 1) + + # Count simplices by dimension + simplex_counts = {dim: 0 for dim in range(self.max_dimension + 2)} + for simplex, _ in st.get_filtration(): + dim = len(simplex) - 1 + if dim <= self.max_dimension + 1: + simplex_counts[dim] += 1 + + print(f" Computing homology...") + # Compute persistence and Betti numbers + st.compute_persistence() + + # CRITICAL FIX: Get persistence intervals to compute Betti numbers (compatible with this GUDHI version) + betti_numbers = {} + + # Method 1: Direct computation for graphs (most reliable for β₀ and β₁) + if self.max_dimension >= 0: + # β₀ = number of connected components + # Can be computed from 0-dimensional persistence + intervals0 = st.persistence_intervals_in_dimension(0) + betti_numbers[0] = int(np.sum(np.isinf(intervals0[:, 1]))) if intervals0.size > 0 else 0 + if betti_numbers[0] == 0: # Fallback + betti_numbers[0] = 1 # Assume connected if edges exist + + if self.max_dimension >= 1: + intervals1 = st.persistence_intervals_in_dimension(1) + betti_numbers[1] = int(np.sum(np.isinf(intervals1[:, 1]))) if intervals1.size > 0 else 0 + + # For higher dimensions, use GUDHI's computation + for dim in range(2, self.max_dimension + 1): + intervals_d = st.persistence_intervals_in_dimension(dim) + betti_numbers[dim] = int(np.sum(np.isinf(intervals_d[:, 1]))) if intervals_d.size > 0 else 0 + + # Sanity check and warning + # Extract cycles from 1D persistence and voids from 2D and 3D persistence + cycle_edges = [] + void_triangles = [] + void_tetrahedra = [] + if self.max_dimension >= 1 and betti_numbers[1] > 0: + # Get persistence pairs for dimension 1 + intervals1 = st.persistence_intervals_in_dimension(1) + persistent_cycles = [] + + if intervals1.size > 0: + # Count infinite intervals (persistent cycles) + infinite_mask = np.isinf(intervals1[:, 1]) + n_persistent = int(np.sum(infinite_mask)) + + # For visualization, create representative cycles from the graph structure + # Using a systematic approach to find independent cycles + G = nx.Graph() + G.add_edges_from(edges) + + # Find a spanning tree + if G.number_of_nodes() > 0 and nx.is_connected(G): + spanning_tree = nx.minimum_spanning_tree(G) + tree_edges = set(spanning_tree.edges()) + + # Non-tree edges create fundamental cycles + cycle_count = 0 + for edge in edges: + if cycle_count >= n_persistent: + break + u, v = edge[0], edge[1] + if (u, v) not in tree_edges and (v, u) not in tree_edges: + # Find path in spanning tree between u and v + try: + tree_path = nx.shortest_path(spanning_tree, u, v) + # Create cycle: tree_path + direct edge back + cycle_edge_list = [] + for i in range(len(tree_path) - 1): + cycle_edge_list.append((tree_path[i], tree_path[i+1])) + cycle_edge_list.append((v, u)) # Close the cycle + cycle_edges.append(cycle_edge_list) + cycle_count += 1 + except: + continue + else: + # For disconnected graphs, use simple cycle detection + all_cycles = [] + for component in nx.connected_components(G): + subG = G.subgraph(component) + if subG.number_of_edges() > subG.number_of_nodes() - 1: + # This component has cycles + try: + cycles = nx.minimum_cycle_basis(subG) + all_cycles.extend(cycles) + except: + continue + + # Convert to edge format + for i, cycle in enumerate(all_cycles): + if i >= n_persistent: + break + cycle_edge_list = [] + for j in range(len(cycle)): + u, v = cycle[j], cycle[(j+1) % len(cycle)] + cycle_edge_list.append((u, v)) + cycle_edges.append(cycle_edge_list) + + # Extract 2D voids from persistence + if self.max_dimension >= 2 and betti_numbers[2] > 0: + intervals2 = st.persistence_intervals_in_dimension(2) + if intervals2.size > 0: + # Count infinite intervals (persistent 2D voids) + infinite_mask = np.isinf(intervals2[:, 1]) + n_voids = int(np.sum(infinite_mask)) + + # Find triangles that bound the voids + # Get all 2-simplices (triangles) from the filtration + triangles_in_complex = [] + for simplex, _ in st.get_filtration(): + if len(simplex) == 3: # Triangle + triangles_in_complex.append(sorted(simplex)) + + # For visualization, select representative triangles for each void + # This is a simplified approach - in practice, void boundaries are more complex + if triangles_in_complex and n_voids > 0: + # Group triangles by connectivity to identify void boundaries + triangle_groups = [] + used_triangles = set() + + for i, triangle in enumerate(triangles_in_complex): + if i in used_triangles or len(triangle_groups) >= n_voids: + continue + + # Start a new group with this triangle + current_group = [triangle] + used_triangles.add(i) + + # Find connected triangles (sharing an edge) + for j, other_triangle in enumerate(triangles_in_complex): + if j in used_triangles: + continue + + # Check if triangles share an edge + shared_edges = 0 + for k in range(3): + edge1 = (triangle[k], triangle[(k+1)%3]) + for l in range(3): + edge2 = (other_triangle[l], other_triangle[(l+1)%3]) + if (edge1[0] == edge2[0] and edge1[1] == edge2[1]) or \ + (edge1[0] == edge2[1] and edge1[1] == edge2[0]): + shared_edges += 1 + + if shared_edges > 0 and len(current_group) < 6: # Limit size for visualization + current_group.append(other_triangle) + used_triangles.add(j) + + triangle_groups.append(current_group) + + void_triangles = triangle_groups[:n_voids] + + # Extract 3D voids from persistence + if self.max_dimension >= 3 and betti_numbers[3] > 0: + intervals3 = st.persistence_intervals_in_dimension(3) + if intervals3.size > 0: + # Count infinite intervals (persistent 3D voids) + infinite_mask = np.isinf(intervals3[:, 1]) + n_3d_voids = int(np.sum(infinite_mask)) + + # Find tetrahedra that bound the 3D voids + # Get all 3-simplices (tetrahedra) from the filtration + tetrahedra_in_complex = [] + for simplex, _ in st.get_filtration(): + if len(simplex) == 4: # Tetrahedron + tetrahedra_in_complex.append(sorted(simplex)) + + # For visualization, select representative tetrahedra for each 3D void + if tetrahedra_in_complex and n_3d_voids > 0: + # Group tetrahedra by connectivity to identify 3D void boundaries + tetrahedra_groups = [] + used_tetrahedra = set() + + for i, tetrahedron in enumerate(tetrahedra_in_complex): + if i in used_tetrahedra or len(tetrahedra_groups) >= n_3d_voids: + continue + + # Start a new group with this tetrahedron + current_group = [tetrahedron] + used_tetrahedra.add(i) + + # Find connected tetrahedra (sharing a triangular face) + for j, other_tetrahedron in enumerate(tetrahedra_in_complex): + if j in used_tetrahedra: + continue + + # Check if tetrahedra share a triangular face + shared_faces = 0 + for k in range(4): + # Get triangular face of first tetrahedron + face1 = sorted([tetrahedron[l] for l in range(4) if l != k]) + for m in range(4): + # Get triangular face of second tetrahedron + face2 = sorted([other_tetrahedron[l] for l in range(4) if l != m]) + if face1 == face2: + shared_faces += 1 + + if shared_faces > 0 and len(current_group) < 4: # Limit size for visualization + current_group.append(other_tetrahedron) + used_tetrahedra.add(j) + + tetrahedra_groups.append(current_group) + + void_tetrahedra = tetrahedra_groups[:n_3d_voids] + + if self.max_dimension >= 2 and simplex_counts.get(3, 0) > 100000: + print(f" ⚠️ Warning: {simplex_counts[3]:,} tetrahedra found - this may affect accuracy") + + return betti_numbers, simplex_counts, cycle_edges, void_triangles, void_tetrahedra + + def compute_simple(self, edges: np.ndarray, n_nodes: int) -> Tuple[Dict[int, int], Dict[int, int]]: + """ + Alternative: Compute only β₀ and β₁ without expansion (fast and accurate) + """ + st = gudhi.SimplexTree() + + # Add all simplices + for v in range(n_nodes): + st.insert([v]) + for edge in edges: + st.insert(edge.tolist()) + + # No expansion - just work with graph + st.compute_persistence() + + # Count components using persistence intervals (compatible API) + intervals0 = st.persistence_intervals_in_dimension(0) + betti_0 = int(np.sum(np.isinf(intervals0[:, 1]))) if intervals0.size > 0 else 0 + if betti_0 == 0: + betti_0 = 1 + + # Use Euler characteristic for cycles + betti_1 = len(edges) - n_nodes + betti_0 + + return {0: betti_0, 1: betti_1}, {0: n_nodes, 1: len(edges)} + +def print_summary(grfD: dict, grfMD: dict, verbose: bool = True): + """Print formatted summary of results""" + + print("\n" + "="*60) + print("RANDOM GRAPH BETTI NUMBER COMPUTATION SUMMARY") + print("="*60) + + # Extract data from the structured dictionaries + graph_conf = grfMD['graph_conf'] + betti_results = grfMD['betti_results'] + + # Graph properties + print(f"\n📊 Graph Properties:") + print(f" Nodes: {graph_conf['n_nodes']}") + print(f" Edge probability: {graph_conf['edge_probability']:.1%}") + print(f" Actual edges: {graph_conf['n_edges']:,}") + print(f" Expected edges: {int(graph_conf['n_nodes']*(graph_conf['n_nodes']-1)/2 * graph_conf['edge_probability']):,}") + print(f" Average degree: {2*graph_conf['n_edges']/graph_conf['n_nodes']:.1f}") + + # Simplex counts + if verbose and betti_results.get('n_simplices'): + print(f"\n🔺 Simplex Counts:") + simplex_names = {0: "vertices", 1: "edges", 2: "triangles", + 3: "tetrahedra", 4: "4-simplices"} + for dim, count in betti_results['n_simplices'].items(): + if count > 0: # Only show non-zero counts + name = simplex_names.get(dim, f"{dim}-simplices") + print(f" {name:12}: {count:,}") + + # Betti numbers + print(f"\n🎯 Betti Numbers:") + interpretations = { + 0: "connected components", + 1: "independent cycles", + 2: "voids/cavities", + 3: "3D voids" + } + + betti_numbers = betti_results['betti_numbers'] + for dim in range(4): + value = betti_numbers.get(dim, 0) + interpretation = interpretations.get(dim, f"{dim}-dimensional holes") + symbol = f"β_{dim}" + print(f" {symbol} = {value:4d} ({interpretation})") + + + # Performance + print(f"\n⏱️ Computation Time: {betti_results['computation_time']:.3f} seconds") + + +def save_results(results: BettiResults, filename: str): + """Save results to JSON file""" + data = { + 'n_nodes': results.n_nodes, + 'edge_probability': results.edge_probability, + 'n_edges': results.n_edges, + 'betti_numbers': results.betti_numbers, + 'computation_time': results.computation_time, + 'n_simplices': results.n_simplices + } + + with open(filename, 'w') as f: + json.dump(data, f, indent=2) + print(f"Results saved to {filename}") + +def main(): + parser = argparse.ArgumentParser( + description='Compute Betti numbers for random graphs' + ) + parser.add_argument('-n', '--nodes', type=int, default=100, + help='Number of nodes in the graph (default: 500)') + parser.add_argument('-g', '--probability', type=float, default=0.15, + help='Edge probability (default: 0.15)') + parser.add_argument('-d', '--max-dim', type=int, default=3, + help='Maximum dimension for Betti numbers (default: 3)') + parser.add_argument('-s', '--seed', type=int, default=None, + help='Random seed for reproducibility') + parser.add_argument('-v', '--verb', type=int, default=1, + help='Verbosity level') + parser.add_argument('-o', '--output', type=str, default=None, + help='Output JSON file for results') + + parser.add_argument('--outPath', type=str, default='out/', + help='Output path for plots') + parser.add_argument('-X', '--noXterm', action='store_true', + help='Disable X terminal for plotting') + parser.add_argument('-p',"--showPlots", type=str,nargs='+', default=" ", help="Plot types to show: a=cycles, b=voids2D, c=voids3D") + + args = parser.parse_args() + args.showPlots=''.join(args.showPlots) + print(f"\n🚀 Starting Betti number computation...") + print(f" Configuration: {args.nodes} nodes, {args.probability:.1%} edge probability") + + start_time = time.time() + + # Generate random graph + generator = RandomGraphGenerator(args.nodes, args.probability, args.seed) + g_edges, n_edges = generator.generate() + + # Compute Betti numbers + computer = BettiComputer(max_dimension=args.max_dim) + + + betti_numbers, simplex_counts, cycle_edges, void_triangles, void_tetrahedra = computer.compute(g_edges, args.nodes) + + computation_time = time.time() - start_time + + # Prepare graph data structures following the pattern from sim_dale_poissonV5.py + grfD = { + 'edges': g_edges, + 'cycle_edges': cycle_edges if ('a' in args.showPlots or args.verb > 1) else [], + 'void_triangles': void_triangles if ('b' in args.showPlots or args.verb > 1) else [], + 'void_tetrahedra': void_tetrahedra if ('c' in args.showPlots or args.verb > 1) else [] + } + + grfMD = { + 'graph_conf': { + 'n_nodes': args.nodes, + 'edge_probability': args.probability, + 'n_edges': n_edges, + 'seed': args.seed, + 'max_dimension': args.max_dim, + }, + 'betti_results': { + 'betti_numbers': betti_numbers, + 'n_simplices': simplex_counts if args.verb > 1 else {}, + 'computation_time': computation_time + }, + 'short_name': f"betti_{args.nodes}n_{int(args.probability*100)}p" + } + + # Store results (keep for backwards compatibility) + results = BettiResults( + n_nodes=args.nodes, + edge_probability=args.probability, + n_edges=n_edges, + betti_numbers=betti_numbers, + computation_time=computation_time, + n_simplices=simplex_counts if args.verb > 1 else {}, + cycle_edges=cycle_edges if args.verb > 1 else [] + ) + + # Print summary + print_summary(grfD, grfMD, verbose=args.verb > 0) + + # Save results if requested + if args.output: + save_results(results, args.output) + + # Setup plotter if needed + args.prjName = grfMD['short_name'] + plot = Plotter(args) + + if 'a' in args.showPlots: + plot.plot_cycles(grfD, grfMD, figId=1) + + if 'b' in args.showPlots: + plot.plot_voids_2d(grfD, grfMD, figId=2) + + if 'c' in args.showPlots: + plot.plot_voids_3d(grfD, grfMD, figId=3) + + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/betti_numbers/toolbox b/causal_net/betti_numbers/toolbox new file mode 120000 index 00000000..c9203e9b --- /dev/null +++ b/causal_net/betti_numbers/toolbox @@ -0,0 +1 @@ +../toolbox/ \ No newline at end of file diff --git a/causal_net/container/laptop_python.src b/causal_net/container/laptop_python.src new file mode 100644 index 00000000..2ba3f648 --- /dev/null +++ b/causal_net/container/laptop_python.src @@ -0,0 +1,38 @@ +# load python (see ubu24-python.dockerfile: matplotlib, scipy, jupyter, imageio-ffmpeg, …) +# podman build --network=host -f ubu24-python.dockerfile -t balewski/ubu24-python:p1 + +IMG=ubu24-python:p3b + +echo launch image $IMG +echo you are launching container ... remember to exit + + +JNB_PORT='' +BASE_DIR=/2026_UoI-VAR # here git has home +WORK_DIR=$BASE_DIR/causal_net +DATA_VAULT=/shared_volumes/dataVault2026 + +echo "The number of arguments is: $#" +# encoded variables: jnb +for var in "$@"; do + echo "The length of argument '$var' is: ${#var}" + if [[ "jnb" == $var ]]; then + JNB_PORT="-p 8837:8837" + echo added $JNB_PORT + echo " cd notebooks; jupyter notebook --ip 0.0.0.0 --no-browser --allow-root --port 8837" + + fi + # ... more ... +done + +podman run -it --rm \ + --volume /shared_volumes/$BASE_DIR:$BASE_DIR \ + --volume /shared_volumes/$WORK_DIR:$WORK_DIR \ + --volume $DATA_VAULT:/dataVault2026 \ + --workdir $WORK_DIR $JNB_PORT \ + --user $(id -u):$(id -g) \ + -e DISPLAY=host.containers.internal:0 \ + $IMG /bin/bash + +# -e LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1 \ +# xhost + 127.0.0.1 diff --git a/causal_net/container/laptop_uoi_podman.source b/causal_net/container/laptop_uoi_podman.source new file mode 100644 index 00000000..8e4e014e --- /dev/null +++ b/causal_net/container/laptop_uoi_podman.source @@ -0,0 +1,46 @@ +#!/bin/bash + +IMG=balewski/causal-net:p2f + +echo "Launching image $IMG" +echo "You are launching Podman image... remember to exit." + +# Initialize variables +JNB_PORT='' # No spaces +BASE_DIR=/2025_UoI-VAR # here git has home +WORK_DIR=$BASE_DIR/causal_net +DATA_VAULT=/shared_volumes/bioDataVault2025 + +# Default port for Jupyter Notebook +PORT=8846 + +# Parse arguments +echo "The number of arguments is: $#" +for var in "$@"; do + echo "The length of argument '$var' is: ${#var}" + if [[ "jnb" == $var ]]; then + JNB_PORT="-p${PORT}:${PORT}" # Remove space before ${PORT} + echo "Added port mapping: $JNB_PORT" + echo " cd notebooks; jupyter notebook --ip 0.0.0.0 --no-browser --allow-root --port $PORT" + fi +done + +# Display the port mapping being added (for debugging) +echo "Final Jupyter port mapping: xx${JNB_PORT}xxx" + +# Check final command for debugging +CMD="podman run -it \ + --volume /shared_volumes/$BASE_DIR:$BASE_DIR \ + --volume /shared_volumes/$WORK_DIR:$WORK_DIR \ + --volume $DATA_VAULT:/dataVault2025 \ + -e My_dataVault=/dataVault2025/causalNet_tmp \ + --workdir $WORK_DIR \ + --user $(id -u):$(id -g) \ + $JNB_PORT \ + $IMG /bin/bash" + +# Print the command to be executed +echo "Running command: $CMD" + +# Execute the command +eval $CMD diff --git a/causal_net/container/pm_uoi-var.source b/causal_net/container/pm_uoi-var.source new file mode 100644 index 00000000..9388c74d --- /dev/null +++ b/causal_net/container/pm_uoi-var.source @@ -0,0 +1,44 @@ +BAd speed - drop it - uses instruction from stable/fit_uoiVar.py + +#!/bin/bash +# salloc -q interactive -C cpu -t 4:00:00 -A m2043 -N 1 + +IMG=balewski/casual-net:p1m # no MPI +IMG=balewski/causal-net:p2a # with MPI, from Neil +IMG=balewski/causal-net:m1b # with MPI, from Neil + UOI code in + +CFSH=/global/cfs/cdirs/mpccc/balewski/ +M2043=/global/cfs/cdirs/m2043/causal_inference + +echo launch image $IMG +echo you are launching Podman-HPC image ... remeber to exit + +JNB_PORT=' ' +BASE_DIR=/2025_UoI-VAR # here git has home +WORK_DIR=$BASE_DIR/causal_net +DATA_VAULT=${CFSH}/bioDataVault2025 +DATA_DIR=/data_tmp + +echo "The number of arguments is: $#" +# encoded variables: jnb +for var in "$@"; do + echo "The length of argument '$var' is: ${#var}" + if [[ "jnb" == $var ]]; then + JNB_PORT=" -p 8833:8833 " + echo added $JNB_PORT + echo " cd notebooks; jupyter notebook --ip 0.0.0.0 --no-browser --allow-root --port 8833 " + fi + # ... more ... +done + +podman-hpc run -it \ + --volume $CFSH/$BASE_DIR:$BASE_DIR \ + --volume $CFSH/$WORK_DIR:$WORK_DIR \ + --volume $M2043:/m2043 \ + --volume ${DATA_VAULT}:/dataVault2025 \ + --volume ${DATA_VAULT}/$DATA_DIR:/causalNet_tmp \ + -e DISPLAY --net=host -v $HOME:$HOME -e HOME \ + -e HDF5_USE_FILE_LOCKING='FALSE' \ + --workdir $WORK_DIR \ + $IMG /bin/bash + diff --git a/causal_net/container/restart_podman.source b/causal_net/container/restart_podman.source new file mode 100644 index 00000000..b7f9f532 --- /dev/null +++ b/causal_net/container/restart_podman.source @@ -0,0 +1,4 @@ +#reset : do it manually +pkill -9 qemu +podman machine stop +podman machine start diff --git a/causal_net/container/ubu22-cuda-mpi-causal-net.dockerfile b/causal_net/container/ubu22-cuda-mpi-causal-net.dockerfile new file mode 100644 index 00000000..9b6595b3 --- /dev/null +++ b/causal_net/container/ubu22-cuda-mpi-causal-net.dockerfile @@ -0,0 +1,80 @@ +# From Neil Metha, Mar 13 --> +# podman-hpc build -f ubu22-cuda-mpi-causal-net.dockerfile -t balewski/causal-net:m1a . +# on PM use 'podman-hpc' instead of 'podman' and all should work +# additionaly do 1 time: podman-hpc migrate balewski/causal-net:m1b + +FROM nvcr.io/nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 +WORKDIR /opt +ENV DEBIAN_FRONTEND noninteractive + +RUN \ + apt-get update && \ + apt-get install --yes \ + build-essential autoconf cmake flex bison zlib1g-dev \ + fftw-dev fftw3 apbs libicu-dev libbz2-dev libgmp-dev \ + bc libblas-dev liblapack-dev git libtool swig uuid-dev \ + libfftw3-dev automake lsb-core libxc-dev libgsl-dev \ + unzip libhdf5-serial-dev ffmpeg libcurl4-openssl-dev \ + libedit-dev libyaml-cpp-dev make libquadmath0 gfortran \ + python3-yaml automake pkg-config libc6-dev libzmq3-dev \ + libjansson-dev liblz4-dev libarchive-dev python3-pip \ + libsqlite3-dev lua5.1 liblua5.1-dev lua-posix jq opam \ + python3-dev python3-cffi python3-ply python3-sphinx \ + aspell aspell-en valgrind libyaml-cpp-dev wget vim \ + make libzmq3-dev python3-yaml time valgrind libeigen3-dev \ + ocaml ocamlbuild ocaml-findlib indent libnum-ocaml libnum-ocaml-dev \ + fig2dev texinfo ghostscript \ + mlocate python3-jsonschema python-is-python3 &&\ + apt-get clean all + + +WORKDIR /opt +ARG mpich=4.2.2 +ARG mpich_prefix=mpich-$mpich +RUN \ + wget https://www.mpich.org/static/downloads/$mpich/$mpich_prefix.tar.gz && \ + tar xvzf $mpich_prefix.tar.gz && \ + cd $mpich_prefix && \ + ./configure FFLAGS=-fallow-argument-mismatch FCFLAGS=-fallow-argument-mismatch \ + --prefix=/opt/mpich/install && \ + make -j 16 && \ + make install && \ + make clean && \ + cd .. && \ + rm -rf $mpich_prefix.tar.gz +ENV PATH=$PATH:/opt/mpich/install/bin +ENV PATH=$PATH:/opt/mpich/install/include +ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/mpich/install/lib +RUN /sbin/ldconfig + +ENV PATH=$PATH:/usr/local/cuda/lib64/stubs +ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64/stubs +ENV PATH=$PATH:/usr/local/cuda-11.8/targets/x86_64-linux/lib/stubs +ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.8/targets/x86_64-linux/lib/stubs + +RUN ln -s /usr/local/cuda-11.8/targets/x86_64-linux/lib/stubs/libnvidia-ml.so /usr/local/cuda-11.8/targets/x86_64-linux/lib/stubs/libnvidia-ml.so.1 + +#Install HWLOC +WORKDIR /opt +RUN git clone -b v2.11 https://github.com/open-mpi/hwloc.git hwloc && \ + cd hwloc && \ + ./autogen.sh && \ + ./configure && \ + make -j 16 && \ + make install + + +RUN pip install setuptools numpy +RUN python -m pip install mpi4py -i https://pypi.anaconda.org/mpi4py/simple +RUN pip install matplotlib pytest flake8 cython sphinx-gallery sphinx-rtd-theme +RUN pip install h5py scikit-learn +RUN pip install pandas tqdm + +# Clone the repository and install in editable mode +WORKDIR /opt +RUN git clone -b uoi-var https://github.com/BouchardLab/pyuoi.git +RUN cd /opt/pyuoi && \ + pip install -e .[dev] + +# Add /opt/pyuoi/examples to PYTHONPATH +ENV PYTHONPATH=$PYTHONPATH:/opt/pyuoi/examples diff --git a/causal_net/container/ubu24-causal-net.dockerfile b/causal_net/container/ubu24-causal-net.dockerfile new file mode 100644 index 00000000..e860694e --- /dev/null +++ b/causal_net/container/ubu24-causal-net.dockerfile @@ -0,0 +1,48 @@ +FROM ubuntu:24.04 + +# podman build -f ubu24-causal-net.dockerfile -t balewski/causal-net:p2b . +# on PM use 'podman-hpc' instead of 'podman' and all should work +# additionaly do 1 time: podman-hpc migrate balewski/casual-net:p1k + +# Set non-interactive mode for apt-get +ARG DEBIAN_FRONTEND=noninteractive +ENV TZ=America/Los_Angeles + +# Update the OS and install required packages +RUN echo "1a-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA OS update" && \ + apt-get update && \ + apt-get install -y locales autoconf automake gcc g++ make vim wget ssh openssh-server sudo git emacs aptitude build-essential xterm python3-pip python3-tk python3-scipy python3-dev iputils-ping net-tools screen feh hdf5-tools python3-bitstring plocate graphviz tzdata x11-apps python3-venv dnsutils iputils-ping && \ + apt-get clean + + +# Create a virtual environment for Python packages to avoid the externally managed environment issue +RUN python3 -m venv /opt/venv + +# Activate the virtual environment +ENV PATH="/opt/venv/bin:$PATH" + +# Install ML libraries +RUN echo "2c-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA math libs" && \ + /opt/venv/bin/pip install scikit-learn pandas seaborn[stats] networkx[default] tqdm + +# Install additional Python libraries +RUN echo "2d-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA python libs" && \ + /opt/venv/bin/pip install --upgrade pip && \ + /opt/venv/bin/pip install matplotlib h5py scipy jupyter notebook bitstring lmfit pytest + + +#RUN apt-get install -y xterm python3-pip x11-apps + +# Clone the repository and install in editable mode +RUN git clone -b uoi-var https://github.com/BouchardLab/pyuoi.git /opt/pyuoi \ + && cd /opt/pyuoi \ + && pip install -e .[dev] + +# Add /opt/pyuoi/examples to PYTHONPATH +ENV PYTHONPATH="/opt/pyuoi/examples:${PYTHONPATH}" + +# Final cleanup +RUN apt-get clean + +# Set the default command to bash +CMD ["/bin/bash"] diff --git a/causal_net/container/ubu24-python.dockerfile b/causal_net/container/ubu24-python.dockerfile new file mode 100644 index 00000000..ba38dd7f --- /dev/null +++ b/causal_net/container/ubu24-python.dockerfile @@ -0,0 +1,47 @@ +FROM ubuntu:24.04 + +# podman build --network=host -f ubu24-python.dockerfile -t ubu24-python:p3b --platform linux/arm64 +# --platform linux/amd64 works w/o LD_PRELOAD but generates WARNING: image platform (linux/amd64) does not match the expected platform (linux/arm64) +# --no-cache tells Podman not to use any cached layers +# for omp_get_num_threads: # -e LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1 \ +# on PM use 'podman-hpc' instead of 'podman' and all should work +# additionaly do 1 time: podman-hpc migrate balewski/ubuXX-qiskit-qml:p1 + + +# Set non-interactive mode for apt-get +ARG DEBIAN_FRONTEND=noninteractive +ARG TARGETARCH +ENV TZ=America/Los_Angeles + +# Update the OS and install required packages +RUN echo "1a-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA OS update" && \ + apt-get update && \ + apt-get install -y locales autoconf automake gcc g++ make vim wget ssh openssh-server sudo git emacs aptitude build-essential xterm python3-pip python3-tk python3-scipy python3-dev iputils-ping net-tools screen feh hdf5-tools python3-bitstring plocate graphviz tzdata x11-apps python3-venv dnsutils iputils-ping && \ + apt-get clean + + +# Create a virtual environment for Python packages to avoid the externally managed environment issue +RUN python3 -m venv /opt/venv + +# Activate the virtual environment +ENV PATH="/opt/venv/bin:$PATH" + +# Install additional Python libraries +RUN echo "2d-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA python libs" && \ + /opt/venv/bin/pip install --upgrade pip && \ + /opt/venv/bin/pip install matplotlib h5py scipy jupyter notebook bitstring lmfit pytest scikit-learn pytz networkx[default] imageio-ffmpeg && \ + if [ "${TARGETARCH}" = "arm64" ]; then \ + echo "Skipping gudhi on linux/arm64: no compatible wheel is available from PyPI."; \ + else \ + /opt/venv/bin/pip install gudhi; \ + fi + + +# Final cleanup +RUN apt-get clean + +# Set the default command to bash +CMD ["/bin/bash"] + + +# check the latest version: pip index versions qiskit diff --git a/causal_net/nonStation_ver3a_states_simu/PlotterDaleMatrix.py b/causal_net/nonStation_ver3a_states_simu/PlotterDaleMatrix.py new file mode 100644 index 00000000..fb2b1daf --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/PlotterDaleMatrix.py @@ -0,0 +1,313 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for simulated Poisson process visualization. + +This module provides specialized plotting capabilities for analyzing +simulated Dale's principle neural networks. The Plotter class extends +PlotterBackbone to create visualizations including: +- Dale connectivity matrix plots with excitatory/inhibitory separation +- Eigenvalue analysis and network stability visualization +- Poisson process statistics and firing rate distributions +- Network dynamics and temporal evolution plots + +Designed specifically for validating and analyzing simulated neural +networks that follow Dale's principle with Poisson spiking dynamics. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec +import matplotlib.colors as colors +import numpy as np +import matplotlib.pyplot as plt +from scipy.optimize import curve_fit +from Util_pseudospectra import compute_pseudospectrum +from UtilDalePoisson import get_offdiag_triplets + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def Dale_matrix_and_eigen(self,A,md,trueD,figId=3): + + figId=self.smart_append(figId) + nrow,ncol=1,2 + fig=self.plt.figure(figId,facecolor='white', figsize=(12,5.)) + + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel if R_sel is not None else '' + + vmin = A.min() + vmax = A.max() + normMap = colors.TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax) + + #.... Position 1: Dale matrix (natural order) ...... + ax = self.plt.subplot(nrow,ncol,1) + im=ax.imshow(A, aspect=1., origin='lower', cmap='bwr', norm=normMap, interpolation='nearest') + ax.set( xlabel='presyn. neuron index (output)', ylabel='postsyn. neuron index (input)') + ax.set_aspect(1.0) + ax.grid() + cbar = fig.colorbar(im, ax=ax, extend="both", shrink=0.7) + + tit='True Dale, N%d%s, %s'%(A.shape[0], R_tag, md['short_name']) + ax.set(title=tit) + ax.axhline(numExc-0.5,color='k',ls='--', label='E/I boundary') + ax.axvline(numExc-0.5,color='k',ls='--') + ax.plot([0,numNeur],[0,numNeur],'--',lw=0.8,color='magenta') + + #..... Position 2: Eigenvalues...... + ax = self.plt.subplot(nrow,ncol,2) + Eigen=np.linalg.eigvals(A) + real_parts = np.real(Eigen) + imag_parts = np.imag(Eigen) + ax.scatter(real_parts, imag_parts, color='blue', marker='o') + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + tit3='Eigenvalues, N%d%s, %s'%(A.shape[0], R_tag, md['short_name']) + ax.set_title(tit3) + ax.axhline(0, color='black', lw=0.5) + ax.axvline(0, color='black', lw=0.5) + ax.axvline(0,color='red', linestyle='--') + ax.set_aspect('equal') + if R_sel is not None: + theta = np.linspace(0, 2*np.pi, 200) + ax.plot(R_sel*np.cos(theta), R_sel*np.sin(theta), color='magenta', linestyle='--', lw=1.5, label='R=%.3f'%R_sel) + ax.legend() + ax.grid(True) + + def Dale_matrix_pseudospectra(self, A, md, trueD, figId=5): # p=e + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(8, 6)) + + npts = 80 + minY = -0.2 + epsMin = 0.03 + + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + ax = self.plt.subplot(1, 1, 1) + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel if R_sel is not None else '' + title = 'Pseudospectra, true N%d%s, %s' % (A.shape[0], R_tag, md['short_name']) + + levels = np.logspace(-2.5, -0.5, 10) + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + ax.text(0.05, 0.85, 'green: ε<%.3f'%epsMin, transform=ax.transAxes) + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) + + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(minY,1.1) + ax.set_xlim(-1.1,1.1) + ax.set_aspect(1.) + ax.axhline(1,color='m',linestyle='--') + ax.axvline(1,color='m',linestyle='--') + ax.axvline(-1,color='m',linestyle='--') + +#...!...!.................. + def histo_weights_rates(self,trueD,spikeD,md,figId=3): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,3.5)) + data_name=md['short_name'] + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.2f' % R_sel + + A=trueD['A_true'] + single_rates = spikeD['single_rates'] + + # output: np.column_stack([i_indices, j_indices, values]) + EposT=get_offdiag_triplets(A,isPos=True) + EnegT=get_offdiag_triplets(A,isPos=False) + print('True0 num edges pos=%d neg=%d'%(EposT.shape[0],EnegT.shape[0])) + + A_diag = np.diag(A) + wMin = min(np.min(EnegT[:,2]), np.min(A_diag)) + wMax = max(np.max(EposT[:,2]), np.max(A_diag)) + + def count_elements(E, Nn=numNeur): + i_indices = E[:, 0].astype(int) + counts = np.bincount(i_indices, minlength=Nn) + return counts + + edgeCount=count_elements(EnegT) + count_elements(EposT) + + #.... weights histogram + ax = self.plt.subplot(nrow,ncol,1) + binX= np.linspace(wMin, wMax, 100) + ax.hist(EnegT[:,2], bins=binX, color='blue', alpha=0.7, edgecolor=None,label='neg:%d'%EnegT.shape[0]) + ax.hist(EposT[:,2], bins=binX, color='red', alpha=0.7, edgecolor=None,label='pos:%d'%EposT.shape[0]) + ax.hist(A_diag, bins=binX, color='cyan', alpha=0.65, edgecolor=None, + label='diag:%d'%A_diag.shape[0]) + + ax.legend(loc='upper left') + tit='True Dale, N%d%s, %s'%(A.shape[0], R_tag, md['short_name']) + ax.set(title=tit, xlabel='True weight value',ylabel='num edges') + ax.axvline(0,color='k',ls='--') + ax.grid(True, alpha=0.3) + + #.... : histogram of rates + ax = self.plt.subplot(nrow,ncol,3) + ax.hist(single_rates, bins=20) + ax.set_xlabel('Firing rate (Hz)') + ax.set_ylabel('num neurons') + ax.grid(True, alpha=0.3) + ax.set_title('Single rates spectrum') + ax.set_xlim(0,) + median_val = np.median(single_rates) + ax.axvline(median_val, color='r', linestyle='--', linewidth=1.5) + y_max = ax.get_ylim()[1] + median_text = f"median: {median_val:.2f} (Hz)\n N={single_rates.shape[0]}" + ax.text( x=median_val * 1.1, y=y_max * 0.7, s=median_text, color='red') + + # Natural neuron indexing: first numExc are excitatory, rest inhibitory + inh_mask = np.zeros(numNeur, dtype=bool) + inh_mask[numExc:] = True + exc_mask = ~inh_mask + neurXlab = 'neuron index' + + x_vals = np.arange(numNeur) + #.... edge count + ax = self.plt.subplot(nrow,ncol,2) + ax.fill_between(x_vals, edgeCount, step='mid', color='salmon', alpha=0.7) + ax.set_xlabel(neurXlab) + ax.set_ylabel('num true edges') + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + probLo, probHi = dmd['edge_prob'] + rho_title = f'outgoing edges, true, prob=[{probLo:.2f}, {probHi:.2f}]' + ax.set_title(rho_title) + + + #.... firing rates ..... + ax = self.plt.subplot(nrow,ncol,4) + chanW=0.9 + + ax.bar(x_vals[exc_mask], single_rates[exc_mask], width=chanW, color='red', align='center', alpha=0.7, label='Excitatory') + ax.bar(x_vals[inh_mask], single_rates[inh_mask], width=chanW, color='blue', align='center', alpha=0.7, label='Inhibitory') + + ax.set_xlabel(neurXlab) + ax.set_ylabel('Firing rate (Hz)') + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + state_tag = md.get('sel_state', None) + if state_tag is None: + ax.set_title('Single Neurons') + else: + ax.set_title(f'Single Neurons, state={state_tag}') + ax.legend() + + +#............................ +#............................ +#............................ + def rates_study(self,trueD,spikeD,md,figId=4): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,4)) + data_name=md['short_name'] + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel if R_sel is not None else '' + + B_idle = trueD['B_true'] + single_rates = spikeD['single_rates'] + + # Natural indexing: first numExc are excitatory, rest inhibitory + exc_mask = np.zeros(numNeur, dtype=bool) + exc_mask[:numExc] = True + inh_mask = ~exc_mask + + # 1) Scatter: x=B_idle, y=log(single_rates) + ax = self.plt.subplot(nrow,ncol,1) + ax.scatter(B_idle[exc_mask], single_rates[exc_mask], s=14, alpha=0.6, marker='^', facecolors='none', edgecolors='red', label='Excitatory') + ax.scatter(B_idle[inh_mask], single_rates[inh_mask], s=14, alpha=0.6, marker='^', facecolors='none', edgecolors='blue', label='Inhibitory') + ax.set_xlabel('true B_idle ') + ax.set_ylabel('single_rates (Hz)') + ax.set_yscale('log') + ax.grid(True, alpha=0.3) + tit='True Dale, N%d%s, %s'%(numNeur, R_tag, data_name) + ax.set_title(tit) + + # reference line: y = exp(B), clipped to central 80% of B range + bmin = float(np.min(B_idle)) + bmax = float(np.max(B_idle)) + bLo = bmin + 0.1 * (bmax - bmin) + bHi = bmax - 0.1 * (bmax - bmin) + bx = np.linspace(bLo, bHi, 100) + by = np.exp(bx) + ax.plot(bx, by, linestyle='--', color='black', linewidth=0.8, label='y=exp(x)') + ax.legend() + + # 2) Histogram: true B_idle (all neurons) + ax2 = self.plt.subplot(nrow,ncol,2) + b_bins = min(60, max(20, int(np.sqrt(numNeur) * 3))) + ax2.hist(B_idle, bins=b_bins, color='dimgray', alpha=0.8, edgecolor=None) + ax2.set_xlabel('true B_idle') + ax2.set_ylabel('num neurons') + ax2.grid(True, alpha=0.3) + state_tag = md.get('sel_state', None) + if state_tag is None: + ax2.set_title('true B_idle') + else: + ax2.set_title(f'true B_idle, state={state_tag}') + + # 3) Histogram: single_rates for excitatory + ax3 = self.plt.subplot(nrow,ncol,3) + exc_vals = single_rates[exc_mask] + inh_vals = single_rates[inh_mask] + # compute common bins and range (start at 0) + exc_max = np.max(exc_vals) if exc_vals.size > 0 else 0.0 + inh_max = np.max(inh_vals) if inh_vals.size > 0 else 0.0 + x_max = max(exc_max, inh_max) + if x_max <= 0: x_max = 1.0 + #common_bins = np.linspace(0, x_max, 21) + common_bins = np.linspace(0, x_max, int(2*x_max)) + ax3.hist(exc_vals, bins=common_bins, color='red', alpha=0.7, edgecolor=None) + ax3.set_xlabel('single_rates (Hz)') + ax3.set_ylabel('num neurons') + ax3.grid(True, alpha=0.3) + ax3.set_title('Rates: Excitatory (N=%d)' % (numExc)) + + # 4) Histogram: single_rates for inhibitory + ax4 = self.plt.subplot(nrow,ncol,4) + ax4.hist(inh_vals, bins=common_bins, color='blue', alpha=0.7, edgecolor=None) + ax4.set_xlabel('single_rates (Hz)') + ax4.set_ylabel('num neurons') + ax4.grid(True, alpha=0.3) + ax4.set_title('Rates: Inhibitory (N=%d)' % (numNeur-numExc)) + # unify x-range starting at 0 for both histograms + ax3.set_xlim(0, x_max) + ax4.set_xlim(0, x_max) diff --git a/causal_net/nonStation_ver3a_states_simu/PlotterLassoFitEval.py b/causal_net/nonStation_ver3a_states_simu/PlotterLassoFitEval.py new file mode 100644 index 00000000..2034101b --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/PlotterLassoFitEval.py @@ -0,0 +1,946 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for LASSO fit evaluation and visualization. + +This module provides comprehensive plotting capabilities for analyzing +LASSO Poisson model fitting results. The Plotter class extends PlotterBackbone +to create specialized visualizations including: +- Connectivity matrix heatmaps with frequency-sorted neurons +- Weight distribution histograms and statistical summaries +- Network structure plots showing excitatory/inhibitory connections +- Reconstruction quality plots comparing fitted vs observed rates +- Comparative analysis plots for simulation validation + +The plots support both simulated and experimental data with automatic +adaptation based on available metadata. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec +import matplotlib.colors as colors +from matplotlib.colors import TwoSlopeNorm +from scipy.stats import gennorm +from matplotlib.colors import LogNorm + + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def summary_fitLasso(self,fitD, md,byFreq=False, figId=1): # p=a + figId=self.smart_append(figId) + nrow,ncol=2,5 + fig=self.plt.figure(figId,facecolor='white', figsize=(16,6)) + + #pprint(md) + fmd=md['fit_lasso'] + esmd=md['edge_selector'] + Nn=fmd['num_neurons'] + + # Create diagonal mask + diagM = np.eye(Nn, dtype=bool) + #!off_diagM = ~diagM + + fitType=md['fit_type'] + eselType=esmd['selector_type'] + isExp = md.get('data_type') == 'bioExp' # Automatically detect experimental data + # Unpack arrays from bigD + A_fit = fitD['A_'+fitType] + B = fitD['B_'+fitType] + Freq=fitD['single_rates'] + + #...... Training curves + ax = self.plt.subplot(nrow,ncol,1) + plot_trainingCurves(ax,fitD,md) + + #.... : histogram of rates + ax = self.plt.subplot(nrow,ncol,1+ncol) + ax.hist(Freq, bins=20) + x_vals = np.arange(Nn) + frLab='Firing rate (Hz)' + ax.set(ylabel='num neurons',xlabel=frLab,title='Single rates, %d neurons'%Nn) + ax.grid(True, alpha=0.4) + ax.set_xlim(0,) + # Compute median + median_val = np.median(Freq) + txtM= f'median rate: {median_val:.2f} Hz' + ax.text( median_val,ax.get_ylim()[1]*0.8,txtM, + color='red', ha='left', va='bottom', fontsize=10) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + + # ... A off-diagonal + xLab='edge value' + frLab='log10( Firing rate/ Hz )' + Aedg=A_fit[~diagM] + i_indices, j_indices = np.where(~diagM) + Freq_expanded = Freq[i_indices] + # Filter out 0 values + valid_mask =np.abs(Aedg)>1e-10 + Aedg_clean = Aedg[valid_mask] + Freq_expanded_clean = Freq_expanded[valid_mask] + # Count non-zero edges (non-NaN values) + n_edges = len(Aedg_clean) + + ax = self.plt.subplot(nrow,ncol,2) + ax.hist(Aedg_clean, bins=100,color='g',alpha=0.8) + ax.set_yscale('log') + ax.grid(True, alpha=0.4) + ax.set(ylabel='edges',xlabel=xLab,title='A off-diagonal') + txt='edge sel meth: %s \n acc frac=%.2f' %( eselType,n_edges/Nn/(Nn-1)) + ax.text(0.05, 0.75, txt, transform=ax.transAxes, fontsize=10) + # Plot 2D histogram + ax = self.plt.subplot(nrow,ncol,2+ncol) + h = ax.hist2d(Aedg_clean, np.log10(Freq_expanded_clean), bins=30, cmap='viridis', norm=LogNorm()) + cbar = fig.colorbar(h[3],ax=ax) + ax.set(ylabel=frLab,xlabel=xLab,title='accept %d of %d edges '%(n_edges,Nn*(Nn-1))) + ax.grid(True, alpha=0.4) + + # ... A sparsity vs epoch ..... + ax = self.plt.subplot(nrow,ncol,3) + epochs = fitD.get('losses_epochs', None) + sp = np.asarray(fitD['sparsity_epoch']) + if epochs is None: + epochs = np.arange(1, len(sp)+1, dtype=np.int32) + else: + epochs = np.asarray(epochs) + ax.plot(epochs, sp, color='tab:purple', linewidth=1.5) + add_delay_markers(ax, fmd) + ax.set_ylim(0.0, 1.02) + ax.set(ylabel='fraction', xlabel='epoch', title='A sparsity (off-diag)') + ax.grid(True, alpha=0.4) + + # ... spectral radius vs epoch ..... + ax = self.plt.subplot(nrow,ncol,4) + rho_ep = fitD.get('spectral_radius_epoch', None) + if rho_ep is not None: + rho_ep = np.asarray(rho_ep) + if epochs is None: + rho_epochs = np.arange(1, len(rho_ep)+1, dtype=np.int32) + else: + rho_epochs = np.asarray(epochs) + ax.plot(rho_epochs, rho_ep, color='tab:red', linewidth=1.5) + add_delay_markers(ax, fmd) + if 'rho_max' in fmd: + ax.axhline(float(fmd['rho_max']), color='k', linestyle='--', linewidth=1.0) + ax.set(ylabel='radius', xlabel='epoch', title='Spectral radius(A)') + else: + rho_now = float(np.max(np.abs(np.linalg.eigvals(A_fit)))) + ax.plot([0, 1], [rho_now, rho_now], color='tab:red', linewidth=1.5) + ax.set(ylabel='radius', xlabel='epoch', title='Spectral radius(A, const)') + ax.grid(True, alpha=0.4) + + # ... B-term ..... + xLab='B-term value' + ax = self.plt.subplot(nrow,ncol,5) + ax.hist(B, bins=30,color='darkviolet') + ax.grid(True, alpha=0.4) + ax.set(ylabel='neurons',xlabel=xLab,title='B-term') + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + # Plot 2D histogram + ax = self.plt.subplot(nrow,ncol,4+ncol) + h = ax.hist2d(B,np.log10(Freq), bins=30, cmap='Grays',vmax=2.1) + cbar = fig.colorbar(h[3],ax=ax) + ax.set(ylabel=frLab,xlabel=xLab,title='num neurons') + ax.grid(True, alpha=0.4) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + short_name = md.get('short_name', '') + if short_name: + fig.suptitle(f"Summary LASSO: {short_name}", fontsize=14) + + #...!...!.................. + def residuals(self, evalD,md, figId=1): + #pprint(md) + fitType=md['fit_type'] + fmd=md['fit_'+fitType] + + figId=self.smart_append(figId) + nrow,ncol=2,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(12,6)) + + # Unpack arrays from bigD + for j,etype in enumerate(['neg','pos']): + ax = self.plt.subplot(nrow,ncol,1+j) + + title = 'edges: %s' % etype + valT,valF=plot_correl_offdiag(fig,ax,evalD[etype]) + ax.set_title(title) + + ax = self.plt.subplot(nrow,ncol,1+j+ncol) + plot_1D_residuals(ax, valT,valF,lab='TP off-diag '+etype,col='green') + + if etype=='pos': + txt=md["short_name"] + else: + txt='fit: '+fitType + ax.text(0.05, 0.2, txt, transform=ax.transAxes, fontsize=8) + + # ... diagonal + title = 'fit (diagonal)' + ax = self.plt.subplot(nrow,ncol,3) + V=evalD['diag'] + dCol='salmon' + ax.scatter(V[:,1], V[:,0], alpha=0.6, color=dCol,marker='.',s=5) + ax.set(title=title,xlabel='truth',ylabel='fitted') + add_x45_lins(ax, only45=True) + ax.grid(True, alpha=0.5) + + ax = self.plt.subplot(nrow,ncol,3+ncol) + plot_1D_residuals(ax,V[:,1],V[:,0],lab='diag',col=dCol) + + + # ... B-term + title = 'fit (B-term)' + ax = self.plt.subplot(nrow,ncol,4) + dCol='darkviolet' + V=evalD['bterm'] + ax.scatter(V[:,1], V[:,0], alpha=0.6, color=dCol,marker='.',s=5) + ax.set(title=title,xlabel='truth',ylabel='fitted') + add_x45_lins(ax, only45=True) + ax.grid(True, alpha=0.5) + + ax = self.plt.subplot(nrow,ncol,4+ncol) + plot_1D_residuals(ax,V[:,1],V[:,0],lab='diag',col=dCol) + + short_name = md.get('short_name', '') + if short_name: + fig.suptitle(f"Residuals: {short_name}", fontsize=14) + + +#...!...!.................. + def A_histos(self, fitD, md, spikeD, figId=1, k=6): + #pprint(md); aa67 + fitType=md['fit_type'] + + fmd=md['fit_'+fitType] + + #1isExp = md.get('data_type') == 'bioExp' # Automatically detect experimental data + figId=self.smart_append(figId) + nrow,ncol=k,2 # Add 1 extra row for the neuron stats plot + fig=self.plt.figure(figId,facecolor='white', figsize=(16,12)) + + # Unpack arrays from bigD + A_fit = fitD['A_'+fitType].copy() + Freq=fitD['single_rates'] + + # Multiply each row of A_fit by the corresponding element of F + A_fit=A_fit * np.sqrt(Freq[:,None]) + + # convert 0's to Nan + A_fit[ A_fit==0.]=np.nan + + num_neurons = A_fit.shape[0] + wzoomMx=0.05 + # Remove diagonal elements by setting them to NaN + A_fit_no_diag = A_fit.copy() + + short_name = md['short_name'] + + np.fill_diagonal(A_fit_no_diag, np.nan) + A_flat = A_fit_no_diag.flatten() + xLab=r' weights * $\sqrt{ rate}$' + + #------- top left plot w/ 1D histo fo weights + ax = self.plt.subplot(nrow,ncol,1) + ax.hist(A_flat, bins=200, color='green', alpha=0.7, edgecolor=None) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + ax.grid() + method_name = md.get('edge_selection_method', 'unknown') + ax.set(xlabel=xLab,ylabel='count',title=f'All non-diag weights, Fit {fitType}, Method: {method_name}') + ax.set_yscale('log') + + # Left column: 2D histogram of A-matrix (mutiple rows) + ax = self.plt.subplot2grid((nrow, ncol), (1, 0), rowspan=5) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + # Create 2D histogram: x-axis is value, y-axis is row index + row_indices = np.repeat(np.arange(num_neurons), num_neurons) + + # Remove NaN values (diagonal elements) + valid_mask = ~np.isnan(A_flat) + A_flat = A_flat[valid_mask] + row_indices = row_indices[valid_mask] + + # Use ax.hist2d() directly with log scale + H, xedges, yedges, im = ax.hist2d(A_flat, row_indices, bins=[50, num_neurons], cmap='Greys', vmax=6) + + cbar = fig.colorbar(im, ax=ax, + orientation='horizontal', # Make it horizontal + location='bottom', # Place it below + shrink=0.7, # Make it 70% width + pad=0.1) # Add some padding from plot + cbar.set_label('num edges') + + ax.set_xlabel(xLab) + ax.set_ylabel('freq-sorted neuron index') + + pprint(md['fit_lasso']) + pprint(md) + lasso_name = md['provenance']['output_lasso_file'] + title_text = f'{lasso_name}, Fit {fitType}, epochs={fmd["num_epochs"]}, sampl/k={fmd["num_samples_used"]/1000}' + + ax.set_title(title_text) + + # Compute row indices and collect data for global binning + single_rates = spikeD['single_rates'] + idxOff = (num_neurons // k) // 2 + 2 + row_indices = [min(i * (num_neurons // k) + idxOff, num_neurons - 1) for i in range(k)] + selected_rows_data = [A_fit_no_diag[idx, :][~np.isnan(A_fit_no_diag[idx, :])] for idx in row_indices] + all_valid_data = np.concatenate([data for data in selected_rows_data if len(data) > 0]) + global_bin_edges = np.linspace(np.min(all_valid_data), np.max(all_valid_data), 51) if len(all_valid_data) > 0 else None + + # Add frequency annotations and plot histograms in single loop + xlim_offset = 0.02 * (ax.get_xlim()[1] - ax.get_xlim()[0]) + for i in range(k): + rowIdx = row_indices[i] + ax.axhline(rowIdx, color='yellow', linewidth=2, alpha=0.8) + ax.text(ax.get_xlim()[0] + xlim_offset, rowIdx, f'{single_rates[rowIdx]:.1f} Hz', verticalalignment='center', horizontalalignment='left', fontsize=12) + ax_hist = self.plt.subplot(nrow, ncol, 2 + i * ncol) + reversed_rowIdx = row_indices[k - 1 - i] + #xLab1 = xLab if (i == k - 1) else None + plot_row_histogram(ax_hist, A_fit_no_diag, reversed_rowIdx, single_rates, global_bin_edges, xLab, isExp=True) + + + + + if short_name: + fig.suptitle(f"Freq-sorted weights: {short_name}", fontsize=14) + +#...!...!.................. + def edges_fitLasso(self, fitD, md, minW=0,figId=1): + """Compare binary edge mask to ground truth.""" + fitType = md['fit_type'] + E_hat = fitD.get('E_mask') + if E_hat is None: + E_hat = fitD.get(f'E_{fitType}') + + if minW>0: + A_hat = fitD.get(f'A_{fitType}') + if A_hat is None: + A_hat = fitD.get('A_avr') + E_hat = (np.abs(A_hat) > minW) + + E_true = fitD['E_true'] + E_hat = np.array(E_hat, dtype=bool) + E_true = np.array(E_true, dtype=bool) + assert E_hat.shape == E_true.shape + + N = E_true.shape[0] + diag = np.eye(N, dtype=bool) + E_hat = E_hat.copy(); E_hat[diag] = False + E_true = E_true.copy(); E_true[diag] = False + + TP = E_hat & E_true + FP = E_hat & (~E_true) + FN = (~E_hat) & E_true + TN = (~E_hat) & (~E_true) + + counts = { + 'TP': int(TP.sum()), + 'FP': int(FP.sum()), + 'FN': int(FN.sum()), + 'TN': int(TN.sum()), + } + denom_pos = max(E_true.sum(), 1) + denom_pred = max(E_hat.sum(), 1) + recall = counts['TP'] / denom_pos + precision = counts['TP'] / denom_pred + f1 = 2*precision*recall / max(precision+recall, 1e-12) + acc = (counts['TP'] + counts['TN']) / max(E_true.size - N, 1) # exclude diag + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(16,3.5)) + gs = gridspec.GridSpec(1,4, width_ratios=[1,1,1,1.1]) + + ax = self.plt.subplot(gs[0]) + A_true = fitD['A_true'] + if A_true is not None and A_true.ndim == 3: + A_true = A_true[0] + if A_true is not None: + vmin = A_true.min() + vmax = A_true.max() + normMap = colors.TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax) + im = ax.imshow(A_true, aspect=1., origin='lower', cmap='bwr', norm=normMap, interpolation='nearest') + dmd = md.get('dale_conf', {}) + numExc = dmd.get('num_excite') + numNeur = dmd.get('num_neurons', N) + R_sel = md.get('sel_spect_radius') + R_tag = f", R={R_sel:.3f}" if R_sel is not None else '' + if 'E_true' in md: + E_true = md['E_true'] + n_diag = int(np.trace(E_true)) + n_off = int(np.sum(E_true) - n_diag) + else: + n_diag = int(np.count_nonzero(np.diag(A_true))) + n_off = int(np.count_nonzero(A_true) - n_diag) + tit = f"True Dale, N{A_true.shape[0]}{R_tag}, nEdges={n_diag}+{n_off}" + ax.set(title=tit) + if numExc is not None: + ax.axhline(numExc-0.5, color='k', ls='--') + ax.axvline(numExc-0.5, color='k', ls='--') + ax.plot([0,numNeur],[0,numNeur],'--',lw=0.8,color='magenta') + else: + im = ax.imshow(E_true, cmap='Greys', vmin=0, vmax=1, origin='lower') + ax.set(title=f'True edges (n={int(E_true.sum())})') + ax.plot([0,N],[0,N],'--',lw=0.8,color='magenta') + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + ax.set_aspect(1.0) + ax.grid() + ax.set_xlim(-0.5,N+0.5); ax.set_ylim(-0.5,N+0.5) + ax.set( xlabel='presyn. neuron index (output)', ylabel='postsyn. neuron index (input)') + ax.xaxis.labelpad = 8 + + ax = self.plt.subplot(gs[1]) + im = ax.imshow(E_hat, cmap='Greys', vmin=0, vmax=1, origin='lower') + ax.set(title=f'Predicted edges (n={int(E_hat.sum())}), minW={minW}') + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + dmd = md.get('dale_conf', {}) + numExc = dmd.get('num_excite') + if numExc is not None: + ax.axhline(numExc-0.5, color='k', ls='--') + ax.axvline(numExc-0.5, color='k', ls='--') + ax.plot([0,N],[0,N],'--',lw=0.8,color='magenta') + ax.set_aspect(1.0) + ax.grid() + ax.set_xlim(-0.5,N+0.5); ax.set_ylim(-0.5,N+0.5) + ax.set( xlabel='presyn. neuron index (output)', ylabel='postsyn. neuron index (input)') + ax.xaxis.labelpad = 8 + ax.set( xlabel='presyn. neuron index', ylabel='postsyn. neuron index') + + conf_map = np.zeros_like(E_true, dtype=int) + conf_map[TP] = 3 + conf_map[FP] = 2 + conf_map[FN] = 1 + cmap_conf = colors.ListedColormap(['white', 'red','blue' ,'green']) + bounds = [-0.5,0.5,1.5,2.5,3.5] + norm = colors.BoundaryNorm(bounds, cmap_conf.N) + ax = self.plt.subplot(gs[2]) + im = ax.imshow(conf_map, cmap=cmap_conf, norm=norm, origin='lower') + ax.set(title='A confusion map') + cbar = self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + cbar.set_ticks([0,1,2,3]) + cbar.set_ticklabels(['TN','FN','FP','TP']) + ax.plot([0,N],[0,N],'--k',lw=0.8) + ax.set_aspect(1.0) + ax.grid() + ax.set_xlim(-0.5,N+0.5); ax.set_ylim(-0.5,N+0.5) + ax.set( xlabel='presyn. neuron index', ylabel='postsyn. neuron index') + + ax = self.plt.subplot(gs[3]) + bar_names = ['TP','FP','FN'] + vals = [counts[k] for k in bar_names] + ax.bar(bar_names, vals, color=['green','royalblue','red','grey']) + ax.set_ylabel('count') + y_pos = 0.25 * max(1, max(vals)) + for name, val in zip(bar_names, vals): + ax.text(name, y_pos, f"{val}", ha='center', va='center', fontsize=10) + txt=f'precision={precision:.3f}\nrecall={recall:.3f}\nf1={f1:.3f}\nacc={acc:.3f}' + ax.text(0.55, 0.70, txt, transform=ax.transAxes) + + ax.set_title('A edges classification') + ax.grid(axis='y', alpha=0.4) + + fig.subplots_adjust(bottom=0.18) + short_name = md.get('short_name', '') + tag = f", short_name={short_name}" if short_name else "" + fig.suptitle(f"Edge detection vs truth, minW={minW}{tag}", fontsize=14) + return + +#...!...!.................. + def summary_network(self, fitD, edgeD, md, procFrac=0.8,figId=5): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(16,3.5)) + + fitType=md['fit_type'] + + # Unpack arrays from bigD + A_fit = fitD['A_'+fitType] + Neu_sum=edgeD['edge_sum'] + Neu_edg=edgeD['edge_vals'] + + # .... A-matrix .... + ax = self.plt.subplot(nrow,ncol,1) + plot_A2D(fig,ax,A_fit) + + # .... A-eigen .... + ax = self.plt.subplot(nrow,ncol,2) + ax.set_title(f'{fitType}: {md["short_name"]}') + eigF=np.linalg.eigvals(A_fit) + reF = np.real(eigF) + imF = np.imag(eigF) + + ax.scatter(reF, imF, color='red', marker='o', facecolors='none',label='fit',s=20) + ax.set_ylim(-0.1,) + #1ax.set_xlim(right=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.grid(True, alpha=0.3) + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + ax.legend() + + # .... pos vs. neg count .... + ax = self.plt.subplot(nrow,ncol,3) + # Extract positive and negative counts + n_pos = Neu_sum[:, 1] + n_neg = Neu_sum[:, 2] + # Create 2D histogram + h = ax.hist2d(n_pos, n_neg, bins=20, cmap='Greys', cmin=1,cmax=5) + # Add colorbar + self.plt.colorbar(h[3], ax=ax, label='Neurons') + ax.set_xlim(0,) + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + + # Add diagonal line (y=x) + lims = [ + np.min([ax.get_xlim(), ax.get_ylim()]), + np.max([ax.get_xlim(), ax.get_ylim()]), + ] + ax.plot(lims, lims, 'b--', alpha=0.5, linewidth=1.5, label='y=x') + ax.set(title='Edge type correlation, %d neurons'%(A_fit.shape[0]),xlabel='num pos',ylabel='num neg') + + # .... edge std vs. value .... + ax = self.plt.subplot(nrow,ncol,4) + ax.grid(True, alpha=0.3) + ax.set(title='Edge value accuracy',xlabel='edge val',ylabel='edge std') + + # Extract average weights and standard deviations + avg_weights = Neu_edg[:, 2] + std_devs = Neu_edg[:, 3] + + # Compute percentile range to contain procFrac of data by magnitude + abs_weights = np.abs(avg_weights) + percentile_cutoff = procFrac * 100 + threshold = np.percentile(abs_weights, percentile_cutoff) + + # Filter data within range + mask = abs_weights <= threshold + avg_weights_accepted = avg_weights[mask] + std_devs_accepted = std_devs[mask] + + short_name = md.get('short_name', '') + if short_name: + fig.suptitle(f"Summary network: {short_name}", fontsize=14) + + # Count accepted and rejected + n_total = len(avg_weights) + n_accepted = len(avg_weights_accepted) + n_rejected = n_total - n_accepted + + # Create 2D histogram + h = ax.hist2d(avg_weights_accepted, std_devs_accepted, bins=30, cmap='Greys', cmin=1) + + # Add colorbar + self.plt.colorbar(h[3], ax=ax, label='edges') + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + # Add text box with statistics + stats_text = (f'display fract = {procFrac:.2f}\n' + f'Accepted: {n_accepted:,} ({n_accepted/n_total*100:.1f}%)\n' + f'Rejected: {n_rejected:,} ({n_rejected/n_total*100:.1f}%)\n' + f'abs(x) cutoff: ±{threshold:.3f}') + + ax.text(0.02, 0.90, stats_text, transform=ax.transAxes, + fontsize=10, verticalalignment='top') + + + +#............................ +#............................ +#............................ + +def plot_trainingCurves(ax,fitD,md,title='aa3'): + fitType=md['fit_type'] + fmd=md['fit_'+fitType] + + epoch0=10 # skip intial loss values + lossTot = fitD['losses_total'][epoch0:] + epochsT=fitD['losses_epochs'][epoch0:] + + # Plot total loss on left y-axis + ax.plot(epochsT, lossTot, label='total '+fitType, color='blue', linestyle='-') + + titl='%dk samp'%(fmd["num_samples_used"]/1000) + + if fitType=='lasso': # Create second y-axis for L1 loss + lossL1= lossTot - fitD['losses_wo_L1'][epoch0:] + ax2 = ax.twinx() + ax2.plot(epochsT, lossL1, label='L1 loss', color='red', linestyle='--') + ax2.set_ylabel('L1 loss', color='red') + ax2.tick_params(axis='y', labelcolor='red') + # Move y-axis ticks and labels inside + ax2.tick_params(axis='y', direction='in', pad=-60, labelsize=10) + ax2.yaxis.set_label_position('right') + + # Format second y-axis ticks in scientific notation + from matplotlib.ticker import FuncFormatter + ax2.yaxis.set_major_formatter(FuncFormatter(lambda x, p: f'{x:.1e}')) + # Combine legends from both axes + lines1, labels1 = ax.get_legend_handles_labels() + lines2, labels2 = ax2.get_legend_handles_labels() + ax.legend(lines1 + lines2, labels1 + labels2, title=titl) + else: + ax.legend( title=titl) + + title = 'Loss: '+md["short_name"] + + add_delay_markers(ax, fmd) + ax.set(xlabel='Epoch', title=title) + ax.set_ylabel('Loss', color='blue') + + # Color the y-axis labels to match the lines + ax.tick_params(axis='y', labelcolor='blue') + ax.grid(True, alpha=0.3) + + +def add_delay_markers(ax, fmd): + """Add dashed black vertical lines at delayed-constraint boundaries.""" + delay = fmd.get('L1_prune_epoch', None) + if delay is None: + return + delay = int(delay) + if delay < 0: + return + # Training logic applies pruning starting at 0-based epoch>=delay, + # which corresponds to displayed epoch delay+1. + ax.axvline(delay + 1, color='k', linestyle='--', linewidth=1.0) + +def plot_A2D(fig,ax,A,title='aa',byFreq=True,trueD=None): + Am=A.copy() + + + # If byFreq=False, reorder matrix to natural neuron indexing + if not byFreq and trueD is not None and 'neur_revFreqIdx' in trueD: + neur_revFreqIdx = trueD['neur_revFreqIdx'] # freq_sorted_position → natural_index + # Reorder both rows and columns from frequency-sorted to natural order + Am = Am[np.ix_(neur_revFreqIdx, neur_revFreqIdx)] + mask = mask[np.ix_(neur_revFreqIdx, neur_revFreqIdx)] + Am[~mask]=0 # Re-apply mask after reordering + + vmin = Am.min()-0.3 + vmax = Am.max()+0.3 + + norm = TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax) + im1=ax.imshow(Am, cmap='bwr', norm=norm, origin='lower') # 'RdBu_r' + + #masked_values = Am[mask] + #vmin = masked_values.min() + #vmax = masked_values.max() + + nval=np.sum(Am!=0.) + title='%s n=%d'%(title,nval) + # Simplified axis labels based on display order + if byFreq: + ylabel_text = 'From neuron (freq-sorted)' + xlabel_text = 'To neuron (freq-sorted)' + else: + ylabel_text = 'From neuron (natural idx)' + xlabel_text = 'To neuron (natural idx)' + ax.set(title=title, ylabel=ylabel_text, xlabel=xlabel_text) + fig.colorbar(im1, ax=ax, shrink=0.7) + ax.grid(True, alpha=0.5) + #add_x45_lins(ax, only45=True) + + +def add_x45_lins(ax, only45=False): + # Add dashed lines through (0,0) + lims = [ + np.min([ax.get_xlim(), ax.get_ylim()]), # min of both axes + np.max([ax.get_xlim(), ax.get_ylim()]), # max of both axes + ] + # 45 degree line y=x + ax.plot(lims, lims, '--', color='k', linewidth=0.8) + if only45: return + + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + + +def plot_correl_offdiag(fig,ax,tripV): + TP,FP,FN=tripV + #print('ss',TP.shape) + ax.scatter(TP[:,3], TP[:,2], alpha=0.6, color='green',label='TP: %d'%TP.shape[0],marker='.',s=5) + n=FN.shape[0] + ax.scatter(FN[:,2], [0]*n, alpha=0.6, color='red',s=5,label='FN: %d'%FN.shape[0]) + + n=FP.shape[0] + ax.scatter([0]*n,FP[:,2], alpha=0.6, color='blue',s=5,label='FP: %d'%FP.shape[0]) + + ax.set(aspect=1. ,xlabel='true weight',ylabel='fitted') + ax.grid(True, alpha=0.5) + + add_x45_lins(ax) + if np.any(TP): # Only plot if there are TP points + x_cg = np.mean(TP[:,3]) + y_cg = np.mean(TP[:,2]) + ax.scatter(x_cg, y_cg, marker='+', s=200, color='k', linewidths=2) #, label='TP avr') + ax.legend() + + return TP[:,3], TP[:,2] + + + +#...!...!.................. +def XXXplot_1d_weight_histo_with_stats(ax, group_values, start_row, end_row, group_idx, color, ggd_rangeL, isExp=False): + """Plot 1D histogram with statistics for a single group""" + + ggd_range=(ggd_rangeL[0]+ggd_rangeL[1])/2. # TMP + # Create histogram for this group with log scale + counts, bins, _ = ax.hist(group_values, bins=100, alpha=0.7, color=color) + + ax.axvline(x=ggd_range, color='black', linestyle='--', linewidth=1, alpha=0.8) + ax.axvline(x=-ggd_range, color='black', linestyle='--', linewidth=1, alpha=0.8) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + # Compute mean and RMSE in the range ±ggd_range, excluding values too close to zero + mask_range = (np.abs(group_values) < ggd_range) & (np.abs(group_values) > ggd_range/5) + values_in_range = group_values[mask_range] + + if len(values_in_range) > 0: + mean_val = np.mean(values_in_range) + rmse_val = np.sqrt(np.mean(values_in_range**2)) + + # Fit Generalized Gaussian Distribution (GGD) + try: + # Fit GGD without any constraints to allow all parameters to be estimated freely + beta, loc, scale = gennorm.fit(values_in_range) + + # Format with scientific notation for very small values + if abs(loc) < 1e-3: + loc_str = f'{loc:.2e}' + else: + loc_str = f'{loc:.4f}' + + if abs(scale) < 1e-3: + scale_str = f'{scale:.2e}' + else: + scale_str = f'{scale:.4f}' + + ggd_text = f'β={beta:.4f}\nx₀={loc_str}\nμ={scale_str}' + + # Plot GGD fit curve + ggd_plot_range=2*ggd_range + x_fit = np.linspace(-ggd_plot_range, ggd_plot_range, 100) + # Scale the PDF to match histogram counts + ggd_pdf = gennorm.pdf(x_fit, beta, loc, scale) + + # Scale to match histogram by finding the maximum bin count in range + bin_centers = (bins[:-1] + bins[1:]) / 2 + mask_fit_bins = (bin_centers >= -ggd_plot_range) & (bin_centers <= ggd_plot_range) + if np.any(mask_fit_bins): + max_count_in_range = np.max(counts[mask_fit_bins]) + max_pdf = np.max(ggd_pdf) + if max_pdf > 0: + scaled_pdf = ggd_pdf * max_count_in_range / max_pdf + ax.plot(x_fit, scaled_pdf, '--', color='magenta', linewidth=2, + label='GGD fit', zorder=8) + + except Exception as e: + print(f"GGD fit failed for group {group_idx+1}: {e}") + beta, loc, scale = np.nan, np.nan, np.nan + ggd_text = f'β=failed\nx₀=failed\nμ=failed' + + # Add text with statistics + ax.text(0.25, 0.85, f'mean={mean_val:.4f}\nRMSE={rmse_val:.4f}\n{ggd_text}', + transform=ax.transAxes, fontsize=10, + bbox=dict(boxstyle='round', facecolor='white', alpha=0.8), + verticalalignment='top') + + # Add circle for mean with horizontal error bar for RMSE at half height + max_count = np.max(counts) + half_height = max_count / 200 + + # Draw horizontal error bar for RMSE + ax.errorbar(mean_val, half_height, xerr=rmse_val, fmt='o', + color='black', capsize=8, capthick=1, linewidth=1, zorder=9) + + ax.set_xlabel('off-diagonal weight value') + ax.set_ylabel('count') + ax.set_title(f'Fitted weights, group {group_idx+1}: neurons {start_row}-{end_row-1} (n={len(group_values)})') + ax.grid(True, alpha=0.3) + + # Limit y-range from 0.5 to max count + max_count = np.max(counts) + ax.set_ylim(0.5, max_count * 1.1) # Add 10% margin above max + + +#...!...!.................. +def add_k_block_lines(ax, num_neurons, k): + """Add horizontal lines to mark K block boundaries""" + rows_per_group = num_neurons // k + for i in range(1, k): # Don't draw line at the very top or bottom + y_line = i * rows_per_group - 0.5 # Position between blocks + ax.axhline(y=y_line, color='black', linestyle='--', linewidth=0.5, alpha=1.0) + +#...!...!.................. +def plot_neuron_stats(ax, A_flat_narrow, row_indices_narrow, num_neurons): + """Plot mean and RMSE for each neuron index""" + + means = [] + rmses = [] + neuron_indices = [] + + # Compute statistics for each neuron + for neuron_idx in range(num_neurons): + # Get values for this specific neuron + mask = row_indices_narrow == neuron_idx + values = A_flat_narrow[mask] + + if len(values) > 0: + mean_val = np.mean(values) + rmse_val = np.sqrt(np.mean(values**2)) + + means.append(mean_val) + rmses.append(rmse_val) + neuron_indices.append(neuron_idx) + + # Convert to numpy arrays + means = np.array(means) + rmses = np.array(rmses) + neuron_indices = np.array(neuron_indices) + + # Plot mean values with RMSE as error bars + ax.errorbar(neuron_indices, means, yerr=rmses, fmt='o', + color='blue', capsize=3, capthick=1, linewidth=1, markersize=3) + + ax.set_xlabel('Neuron freq-sorted index') + ax.set_ylabel('Mean ± RMSE') + ax.set_title('Per-neuron zero-values stats') + ax.grid(True, alpha=0.3) + ax.axhline(y=0, color='black', linestyle='-', linewidth=0.5, alpha=0.5) + + +#...!...!.................. +def plot_1D_residuals(ax, valT,valF,lab,col,first=True): + fac=100 + res=fac*(valT-valF) + + # Plot histogram of residuals + n, bins, patches = ax.hist(res, bins=30, color=col, edgecolor='none', alpha=0.7, density=True ) + + # Compute statistics: mean, standard deviation, and RMSE + mean_val = np.mean(res) + std_val = np.std(res) + n_entries = len(valT) + + # Determine half of the maximum bin height to position the error bar + half_height = np.max(n) / 2.0 + + # Draw horizontal error bar for RMSE at the computed mean and half-height + ax.errorbar(mean_val, half_height, xerr=std_val, fmt='o', + color='black', capsize=8, capthick=1, linewidth=1, zorder=9) + + # Add text to the figure displaying the mean, standard deviation, and RMSE + textstr = f"{lab}\nN = {n_entries}\nMean = {mean_val:.2f}\nStd = {std_val:.2f}" + + if first: + x0=0.95 + else: + x0=0.4 + ax.text(x0, 0.95, textstr, transform=ax.transAxes, fontsize=8, + verticalalignment='top', horizontalalignment='right') + + # Add a black vertical line at x = 0 + ax.axvline(0, color='black', linestyle='--', linewidth=1) + ax.set(xlabel='residuals x %d'%fac, ylabel='density/bin') + + +#...!...!.................. +def plot_both_eigen(ax, eigT,eigF): + reT = np.real(eigT) + imT = np.imag(eigT) + + ax.scatter(reT, imT, color='blue', marker='o',label='True',s=10) + + reF = np.real(eigF) + imF = np.imag(eigF) + + ax.scatter(reF, imF, color='red', marker='o', facecolors='none',label='fit',s=20) + ax.set_ylim(-0.1,) + ax.set_xlim(right=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.grid(True, alpha=0.3) + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + ax.legend() + + +#...!...!.................. +def plot_row_histogram(ax, A_fit_no_diag, rowIdx, single_rates=None, global_bin_edges=None, xLab=None, isExp=False): + """Plot histogram of a specific row from the 2D matrix with statistics""" + + # Extract row data (exclude diagonal element which is NaN) + row_data = A_fit_no_diag[rowIdx, :] + valid_data = row_data[~np.isnan(row_data)] + + if len(valid_data) == 0: + ax.text(0.5, 0.5, 'No valid data', transform=ax.transAxes, ha='center', va='center') + return + + # Plot histogram with consistent binning + if global_bin_edges is not None: + counts, bin_edges, _ = ax.hist(valid_data, bins=global_bin_edges, color='chocolate') + else: + neve_happens + n_bins = min(60, len(valid_data)//3) # Adaptive bin count fallback + counts, bin_edges, _ = ax.hist(valid_data, bins=n_bins, alpha=0.7, color='skyblue') + + # Calculate median value + xMedian = np.median(valid_data) + + # Define slice range around median + xDel = 0.1 + slice_mask = (valid_data >= (xMedian - xDel)) & (valid_data <= (xMedian + xDel)) + sliced_data = valid_data[slice_mask] + + # Compute standard deviation of sliced data + if len(sliced_data) > 1: + std_val = np.std(sliced_data) + x0 = np.mean(sliced_data) + else: + std_val = 0.0 + x0 = xMedian + + # Draw x=0 reference line (always shown) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + # Draw median line (always shown) + ax.axvline(xMedian, linestyle='-', color='k', linewidth=2, label='Median') + + # Draw dashed lines for the ±0.1 range (only for simulation data) + if not isExp: + ax.axvline(xMedian - xDel, linestyle='--', color='k', alpha=0.7) + ax.axvline(xMedian + xDel, linestyle='--', color='k', alpha=0.7) + else: + ax.set_yscale('log') + + # Add merged text with statistics and row/frequency info inside the plot + freq_text = f', {single_rates[rowIdx]:.1f}Hz' if single_rates is not None else '' + info_text = f'Row {rowIdx}{freq_text}\nmedian={xMedian:.3f}, std={std_val:.3f}' + + ax.text(0.98, 0.95, info_text, transform=ax.transAxes, fontsize=8, + verticalalignment='top', horizontalalignment='right', + bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.8)) + + # Set consistent x-axis limits when using global binning + if global_bin_edges is not None: + ax.set_xlim(global_bin_edges[0], global_bin_edges[-1]) + + if xLab!=None: ax.set_xlabel(xLab) + + ax.set_ylabel('Count') # Restore y-axis label + ax.grid(True, alpha=0.3) diff --git a/causal_net/nonStation_ver3a_states_simu/PlotterPrismEM.py b/causal_net/nonStation_ver3a_states_simu/PlotterPrismEM.py new file mode 100644 index 00000000..870db856 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/PlotterPrismEM.py @@ -0,0 +1,1159 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for prism EM evaluation. +""" + +from toolbox.PlotterBackbone import PlotterBackbone +import numpy as np +from matplotlib.ticker import MaxNLocator +import matplotlib.colors as colors + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def _rebin_1d(self, x, merge): + merge = int(max(1, merge)) + if merge <= 1: + return np.asarray(x) + x = np.asarray(x) + n = x.shape[0] // merge + if n <= 0: + return x + return x[: n * merge].reshape(n, merge).mean(axis=1) + + def _rebin_2d(self, x, merge): + merge = int(max(1, merge)) + if merge <= 1: + return np.asarray(x) + x = np.asarray(x) + n = x.shape[0] // merge + if n <= 0: + return x + return x[: n * merge].reshape(n, merge, x.shape[1]).mean(axis=1) + + def _rebin_mode_1d(self, x, merge): + merge = int(max(1, merge)) + x = np.asarray(x) + if merge <= 1: + return x.astype(np.int64, copy=False) + n = x.shape[0] // merge + if n <= 0: + return x.astype(np.int64, copy=False) + x = x[: n * merge].reshape(n, merge) + out = np.zeros((n,), dtype=np.int64) + for i in range(n): + vals = x[i].astype(np.int64, copy=False) + vals = vals[vals >= 0] + if vals.size == 0: + out[i] = -1 + else: + out[i] = np.bincount(vals).argmax() + return out + + def summary_prismEM(self, fitD, md, figId=1): + """EM convergence overview: 2 rows × 3 columns. + + Row 1: E-step NLL vs EM iter | M-step NLL+L1 vs M-epoch | spectral radius vs M-epoch + Row 2: nz off-diag edges vs M-epoch | mean occupancy | A off-diag weight histogram + """ + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 6)) + fig.subplots_adjust(hspace=0.45, wspace=0.35) + + trainMD = md["train"] + short_name = md["short_name"] + + e_nll = np.asarray(fitD["e_nll_em"]) + m_nll = np.asarray(fitD["m_nll_epoch"]) + m_l1 = np.asarray(fitD["m_l1_epoch"]) + m_loss = np.asarray(fitD["m_loss_epoch"]) + rho = np.asarray(fitD["rho_epoch"]) + nz = np.asarray(fitD["nz_edges_epoch"]) + lr = np.asarray(fitD["learning_rates"]) + + n_em = trainMD["num_em_iters"] + m_per_em = trainMD["m_epochs"] + n_m_total = len(m_nll) + m_epochs = np.arange(1, n_m_total + 1) + em_iters = np.arange(1, len(e_nll) + 1) + + prune_em = int(trainMD.get("delay_em_iter_4_Aprune", 0)) + rho_start_em = int(trainMD.get("delay_em_iter_4_ArhoMax", 0)) + lr_drop_em = int(trainMD.get("delay_em_iter_4_lrDecay", int(n_em * 0.7))) + dyn_weight_em = int(trainMD.get("delay_em_iter_4_dynWeight", int(n_em * 0.8))) + show_dyn_weight_marker = not bool(trainMD.get("noFreqWeight", False)) + prune_m_epoch = prune_em * m_per_em + rho_start_m_epoch = rho_start_em * m_per_em + lr_drop_m_epoch = lr_drop_em * m_per_em + dyn_weight_m_epoch = dyn_weight_em * m_per_em + + def draw_threshold_marker(ax, x_pos, x_max, txt, color,yFac=0.02): + if not (0 < x_pos <= x_max): + return + ax.axvline(x_pos, color=color, ls='--', lw=0.9, alpha=0.95) + y0, y1 = ax.get_ylim() + x0, x1 = ax.get_xlim() + x_off = 0.01 * max(1e-9, x1 - x0) + y_txt = y1 - yFac * (y1 - y0) + ax.text( + x_pos + x_off, y_txt, txt, rotation=90, color=color, fontsize=7, + ha='left', va='top', + bbox=dict(facecolor='white', alpha=0.55, edgecolor='none', pad=0.2) + ) + + jSkipEM = 0 + jSkipM = int(jSkipEM * m_per_em) + # ── Row 1, Col 1: E-step NLL vs EM iteration ──────────────── + ax = self.plt.subplot(2, 3, 1) + ax.plot(em_iters[jSkipEM:], e_nll[jSkipEM:], 'o-', color='tab:blue', markersize=3, + linewidth=1.2) + ax.set(title="E-step weighted NLL", xlabel="EM iteration", + ylabel="NLL / bin") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_em, len(e_nll), "start Aprune", "k") + draw_threshold_marker(ax, rho_start_em, len(e_nll), "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_em, len(e_nll), "start_lrDrop", "tab:gray", yFac=0.6) + if show_dyn_weight_marker: + draw_threshold_marker(ax, dyn_weight_em, len(e_nll), "start dynWeight", "tab:pink") + txt = (f"pgd_iter={trainMD['pgd_iter']}\n" + f"lr_E={trainMD['lr_estep']}\n" + f"λ₂={trainMD['lambda2']}") + ax.text(0.97, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=8, + bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')) + + # ── Row 1, Col 2: M-step NLL + L1 vs M-epoch ──────────────── + ax = self.plt.subplot(2, 3, 2) + ax.plot(m_epochs[jSkipM:], m_nll[jSkipM:], color='tab:blue', linewidth=1, label='NLL') + ax.set_ylabel('NLL', color='tab:blue') + ax.tick_params(axis='y', labelcolor='tab:blue') + ax.set(title="M-step loss", xlabel="M-epoch (global)") + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + ax2.plot(m_epochs[jSkipM:], m_l1[jSkipM:] , color='tab:red', linewidth=1, + linestyle='--', label='L1') + ax2.set_ylabel('L1', color='tab:red') + ax2.tick_params(axis='y', labelcolor='tab:red') + + lines1, lab1 = ax.get_legend_handles_labels() + lines2, lab2 = ax2.get_legend_handles_labels() + ax.legend(lines1 + lines2, lab1 + lab2, fontsize=7, loc='upper right') + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + if show_dyn_weight_marker: + draw_threshold_marker(ax, dyn_weight_m_epoch, n_m_total, "start dynWeight", "tab:pink") + + txt = (f"lr_M={trainMD['lr_mstep']}\n" + f"L1 λ3={trainMD['lambda3']}\n" + f"batch={trainMD['batch_size']}") + ax.text(0.03, 0.03, txt, transform=ax.transAxes, + va="bottom", ha="left", fontsize=8, + bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')) + + # ── Row 1, Col 3: spectral radius vs M-epoch ──────────────── + ax = self.plt.subplot(2, 3, 3) + rho_max = trainMD["rho_max"] + ax.plot(m_epochs[jSkipM:], rho[jSkipM:] , color='tab:green', linewidth=1, + label='ρ(A)') + ax.axhline(rho_max, color='red', ls='--', lw=1, + label=f'ρ_max={rho_max}') + ax.set(title="Spectral radius ρ(A)", xlabel="M-epoch (global)", + ylabel="ρ") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + if show_dyn_weight_marker: + draw_threshold_marker(ax, dyn_weight_m_epoch, n_m_total, "start dynWeight", "tab:pink") + ax.legend(fontsize=8) + + # ── Row 2, Col 1: non-zero off-diag edges vs M-epoch ──────── + ax = self.plt.subplot(2, 3, 4) + ax.plot(m_epochs[jSkipM:], nz[jSkipM:], color='tab:purple', linewidth=1) + ax.set(title=f"Non-zero off-diag edges (|A|>{trainMD['minW']})", + xlabel="M-epoch (global)", ylabel="count") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + if show_dyn_weight_marker: + draw_threshold_marker(ax, dyn_weight_m_epoch, n_m_total, "start dynWeight", "tab:pink") + + # learning rate on twin axis + ax2 = ax.twinx() + ax2.plot(m_epochs[jSkipM:], lr[jSkipM:], color='tab:orange', linewidth=0.8, + linestyle='--', alpha=0.6) + ax2.set_ylabel('learning rate', color='tab:orange') + ax2.tick_params(axis='y', labelcolor='tab:orange') + + # ── Row 2, Col 2: mean occupancy ───────────────────────────── + ax = self.plt.subplot(2, 3, 5) + c_hat = np.asarray(fitD["c_hat"]) + occ = c_hat.mean(axis=0) + M = len(occ) + ax.bar(np.arange(M), occ, color='tab:blue', alpha=0.7) + ax.set(title="Mean occupancy", xlabel="state", ylabel="mean c") + ax.set_ylim(0.0, 1.0) + ax.set_xticks(np.arange(M)) + ax.grid(True, alpha=0.3, axis='y') + + t0s, t1s = trainMD["time_range_sec"] + txt = (f"N={trainMD['num_neurons']} M={M}\n" + f"T=[{t0s:.0f},{t1s:.0f}]s\n" + f"bins={trainMD['num_time_bins']}") + ax.text(0.97, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=8, + bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')) + + # ── Row 2, Col 3: A off-diagonal weight histogram ──────────── + ax = self.plt.subplot(2, 3, 6) + A_fit = np.asarray(fitD["A_hat"]) + Nn = A_fit.shape[0] + diag_mask = np.eye(Nn, dtype=bool) + A_off = A_fit[~diag_mask] + valid = np.abs(A_off) > 1e-10 + A_off_nz = A_off[valid] + n_edges = int(A_off_nz.size) + ax.hist(A_off_nz, bins=100, color='g', alpha=0.8) + ax.set_yscale('log') + ax.set(title=f"A off-diagonal, {n_edges} edges", + xlabel="edge value", ylabel="edges") + ax.grid(True, alpha=0.3) + + fig.suptitle(f"Prism EM: {short_name}, " + f"K_EM={n_em}×K_M={m_per_em}", + fontsize=12) + fig.tight_layout(rect=[0, 0, 1, 0.95]) + + def state_init_prismEM(self, fitD, md, figId=2, time_reb=20): + """Two-panel plot: initialization vs truth state trajectories.""" + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(13.5, 8.0)) + gs = fig.add_gridspec(3, 3, height_ratios=[1.0, 1.0, 0.85], hspace=0.75, wspace=0.35) + + trainMD = md["train"] + t0_bin, t1_bin = trainMD["time_range_bins"] + dt = float(trainMD["time_step_sec"]) + + S_init = np.asarray(fitD["S_init"]) + C_init = np.asarray(fitD["c_init"]) + S_true = np.asarray(md["S_true"])[t0_bin : t1_bin + 1] + C_true = np.asarray(md["C_true"])[t0_bin : t1_bin + 1] + + assert S_init.shape[0] == S_true.shape[0] == C_init.shape[0] == C_true.shape[0], \ + "S_init/S_true/C_init/C_true must have matching lengths" + n_cmp = S_init.shape[0] + + # Use original EM time bins (dt), not coarse rate-bin or plotting rebinning. + S_init_cl = 1.0 - np.max(C_init, axis=1) + t = (t0_bin + np.arange(n_cmp)) * dt + n_states = C_init.shape[1] + + ax = fig.add_subplot(gs[0, :]) + ax.step(t, S_init, where="post", color="k", linewidth=1.0, label="S_init") + ax.plot(t, S_init, linestyle="none", marker=".", markersize=1.5, color="k", alpha=0.6) + lo = np.clip(S_init - S_init_cl, -1.0, float(n_states - 1)) + hi = np.clip(S_init + S_init_cl, -1.0, float(n_states - 1)) + ax.fill_between(t, lo, hi, color="gray", alpha=0.3, label="S_init_CL") + ax.set(title="Initialization", xlabel="time (s)", ylabel="state") + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + for m in range(C_init.shape[1]): + ax2.step(t, C_init[:, m], where="post", linewidth=0.8, alpha=0.85, label=f"C_init[{m}]") + ax2.plot(t, C_init[:, m], linestyle="none", marker=".", markersize=1.4, alpha=0.55) + ax2.set_ylabel("C_init") + ax2.set_ylim(0.0, 1.0) + + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + ax = fig.add_subplot(gs[1, :]) + ax.step(t, S_true, where="post", color="k", linewidth=1.0, label="S_true") + ax.plot(t, S_true, linestyle="none", marker=".", markersize=1.5, color="k", alpha=0.6) + ax.set(title="Truth", xlabel="time (s)", ylabel="state") + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + for m in range(C_true.shape[1]): + ax2.step(t, C_true[:, m], where="post", linewidth=0.8, alpha=0.85, label=f"C_true[{m}]") + ax2.plot(t, C_true[:, m], linestyle="none", marker=".", markersize=1.4, alpha=0.55) + ax2.set_ylabel("C_true") + + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + # Bottom row, left: histogram of classification frequency with thresholds. + ax = fig.add_subplot(gs[2, 0]) + init_state_md = md["init_state"] + + freq_h1d = np.asarray(fitD["freq_h1d"], dtype=np.float64).ravel() + bins = min(40, max(10, int(np.sqrt(max(1, freq_h1d.size))))) + ax.hist(freq_h1d, bins=bins, color='tab:blue', alpha=0.75) + rate_bin_sec = float(init_state_md["rate_bin_sec"]) + ax.set(title=f"Rate (time bin {rate_bin_sec:g} s)", xlabel="rate", ylabel="count") + ax.grid(True, alpha=0.3) + + thres = [float(x) for x in init_state_md["rate_thres"]] + for thr in thres: + ax.axvline(float(thr), color='k', linestyle='--', linewidth=1.0, alpha=0.9) + + # Label state regions 0,1,... inside histogram intervals. + xlo, xhi = ax.get_xlim() + ylo, yhi = ax.get_ylim() + edges = [xlo] + thres + [xhi] + for i in range(max(0, len(edges) - 1)): + xc = 0.5 * (edges[i] + edges[i + 1]) + ax.text( + xc, ylo + 0.88 * (yhi - ylo), f"{i}", + ha="center", va="center", fontsize=10, color="k", + bbox=dict(facecolor="white", alpha=0.6, edgecolor="none"), + ) + + # Bottom row, middle: B_init correlation vs log(single_rates). + ax_mid = fig.add_subplot(gs[2, 1]) + b_init = np.asarray(fitD["B_init"], dtype=np.float64).ravel() + single_rates = np.asarray(fitD["single_rates"], dtype=np.float64).ravel() + assert b_init.ndim == 1 and b_init.size > 0, "Expected B_init to be 1D and non-empty" + assert b_init.size == single_rates.size, "B_init and single_rates must have matching lengths" + + n = b_init.size + assert n > 1, "Need at least 2 points for correlation plot" + x = np.log(np.clip(single_rates, 1e-12, None)) + y = b_init + ax_mid.scatter(x, y, s=10, alpha=0.75, color='tab:green', edgecolors='none') + corr = float(np.corrcoef(x, y)[0, 1]) + lo = float(min(np.min(x), np.min(y))) + hi = float(max(np.max(x), np.max(y))) + ax_mid.plot([lo, hi], [lo, hi], color='k', linestyle='--', linewidth=0.8, alpha=0.6) + ax_mid.set(title=f"B_init r={corr:.3f}", + xlabel="log(single_rates)", ylabel="B_init") + ax_mid.set_aspect("equal", adjustable="box") + ax_mid.grid(True, alpha=0.3) + + # Bottom row, right: reserved cell for future diagnostics. + ax_right = fig.add_subplot(gs[2, 2]) + ax_right.axis("off") + + t0s, t1s = trainMD["time_range_sec"] + fig.suptitle( + f"Prism EM Init: {md['short_name']} T=[{t0s:.1f}, {t1s:.1f}] s", + fontsize=12, + ) + + def _draw_A_matrix(self, ax, A, title, num_exc=None, norm_map=None): + A = np.asarray(A) + if norm_map is None: + vmin = float(np.min(A)) + vmax = float(np.max(A)) + if np.isclose(vmin, vmax): + vmax = vmin + 1e-9 + norm_map = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + im = ax.imshow(A, aspect=1.0, origin='lower', cmap='bwr', + norm=norm_map, interpolation='nearest') + N = A.shape[0] + if num_exc is not None: + ax.axhline(num_exc - 0.5, color='k', ls='--', lw=0.8) + ax.axvline(num_exc - 0.5, color='k', ls='--', lw=0.8) + ax.plot([0, N], [0, N], '--', lw=0.8, color='magenta') + ax.set_xlim(-0.5, N + 0.5) + ax.set_ylim(-0.5, N + 0.5) + ax.grid(True, alpha=0.25) + ax.set_title(title) + ax.set_xlabel('presyn. neuron index (output)') + ax.set_ylabel('postsyn. neuron index (input)') + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + def _get_A_true_info(self, md): + A_true = np.asarray(md["A_true"]) + if A_true.ndim == 3: + A_true = A_true[0] + assert A_true.ndim == 2 and A_true.shape[0] == A_true.shape[1], "A_true must be square 2D" + N = A_true.shape[0] + num_exc = md["dale_conf"]["num_excite"] + off_mask = ~np.eye(N, dtype=bool) + n_diag = int(np.eye(N, dtype=bool).sum()) + n_off = int(np.count_nonzero(A_true[off_mask])) + vmin = float(np.min(A_true)) + vmax = float(np.max(A_true)) + if np.isclose(vmin, vmax): + vmax = vmin + 1e-9 + shared_norm = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + return A_true, N, num_exc, n_diag, n_off, shared_norm + + def _scatter_true_vs_cmp(self, ax, x_true, y_cmp, title, color, split_by_true_sign=False): + x_true = np.asarray(x_true, dtype=np.float64).ravel() + y_cmp = np.asarray(y_cmp, dtype=np.float64).ravel() + n = min(x_true.size, y_cmp.size) + x = x_true[:n] + y = y_cmp[:n] + ax.scatter(x, y, alpha=0.7, color=color, marker='.', s=9) + + if n > 0: + lo = float(min(np.min(x), np.min(y))) + hi = float(max(np.max(x), np.max(y))) + pad = 0.05 * max(1e-9, hi - lo) + lo -= pad + hi += pad + ax.plot([lo, hi], [lo, hi], color='k', linestyle='--', linewidth=1.0, alpha=0.6) + ax.axhline(0.0, color='k', linestyle='--', linewidth=0.8, alpha=0.6) + ax.axvline(0.0, color='k', linestyle='--', linewidth=0.8, alpha=0.6) + ax.set_xlim(lo, hi) + ax.set_ylim(lo, hi) + + if split_by_true_sign: + m_neg = x < 0.0 + m_pos = x > 0.0 + if np.any(m_neg): + xm = float(np.mean(x[m_neg])) + ym = float(np.mean(y[m_neg])) + ax.plot([xm], [ym], marker='+', markersize=14, markeredgewidth=2.5, color='k') + if np.any(m_pos): + xm = float(np.mean(x[m_pos])) + ym = float(np.mean(y[m_pos])) + ax.plot([xm], [ym], marker='+', markersize=14, markeredgewidth=2.5, color='k') + else: + xm = float(np.mean(x)) + ym = float(np.mean(y)) + ax.plot([xm], [ym], marker='+', markersize=14, markeredgewidth=2.5, color='k') + + ax.set_aspect('equal', adjustable='box') + ax.grid(True, alpha=0.4) + ax.set_title(title) + ax.set_xlabel('true weight') + ax.set_ylabel('fitted') + + def compare_A_vs_truth(self, A_cmp, md, cmp_label='A_init', figId=3): + """Reusable canvas to compare an A-matrix candidate against A_true.""" + A_true, N, num_exc, n_diag, n_off, shared_norm = self._get_A_true_info(md) + A_cmp = np.asarray(A_cmp) + assert A_cmp.ndim == 2 and A_cmp.shape == A_true.shape, "A_compare must match A_true shape" + + off_mask = ~np.eye(N, dtype=bool) + neg_mask = off_mask & (A_true < 0) + pos_mask = off_mask & (A_true > 0) + diag_mask = np.eye(N, dtype=bool) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(14, 8)) + gs = fig.add_gridspec(2, 3, hspace=0.55, wspace=0.45) + + ax = fig.add_subplot(gs[0, 0]) + self._draw_A_matrix( + ax, A_true, f"True Dale, N{N}, nEdges={n_diag}+{n_off}", + num_exc=num_exc, norm_map=shared_norm + ) + + ax = fig.add_subplot(gs[0, 1]) + self._draw_A_matrix( + ax, A_cmp, f"{cmp_label}, nEdges={n_diag}+{n_off}", + num_exc=num_exc, norm_map=shared_norm + ) + + ax = fig.add_subplot(gs[0, 2]) + A_off = A_cmp[off_mask] + A_off_nz = A_off[np.abs(A_off) > 1e-12] + ax.hist(A_off_nz, bins=120, color='saddlebrown', alpha=0.85) + ax.set_yscale('log') + ax.grid(True, alpha=0.35) + ax.set_title(f"{cmp_label} off-diagonal") + ax.set_xlabel("edge value") + ax.set_ylabel("edges") + + initA = md["init_A"] + txt = ( + "A-init diagnostics:\n" + f" cond(YpYp) = {initA['cond_YpYp']:.2e}\n" + f" rho(A_ols) = {initA['rho_A_init']:.3f}\n" + f" R2 = {initA['R2_1step']:.3f}\n" + f" ||A||_F = {initA['fro_A_init']:.3f}\n" + f" bins_used = {initA['num_bins_used']}/{initA['num_bins_total']}" + ) + ax.text( + 0.05, 0.70, txt, transform=ax.transAxes, + va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"), + family="monospace", + ) + + ax = fig.add_subplot(gs[1, 0]) + self._scatter_true_vs_cmp( + ax, + A_true[neg_mask], + A_cmp[neg_mask], + f"{cmp_label} :neg TP", + color='tab:green', + ) + + ax = fig.add_subplot(gs[1, 1]) + self._scatter_true_vs_cmp( + ax, + A_true[pos_mask], + A_cmp[pos_mask], + f"{cmp_label} :pos TP", + color='tab:green', + ) + + ax = fig.add_subplot(gs[1, 2]) + self._scatter_true_vs_cmp( + ax, + A_true[diag_mask], + A_cmp[diag_mask], + f"{cmp_label} :diag", + color='salmon', + split_by_true_sign=True, + ) + + fig.suptitle(f"A_true vs. {cmp_label} comparison: {md['short_name']}", fontsize=13) + fig.tight_layout(rect=[0, 0, 1, 0.95]) + + def matrix_init_prismEM(self, fitD, md, figId=3, est_key="A_init", est_label="A_init"): + """Wrapper for current EM use case; reusable for A_hat later.""" + self.compare_A_vs_truth(fitD[est_key], md, cmp_label=est_label, figId=figId) + + def edge_recovery_prismEM(self, fitD, md, minW=0.02, figId=4, est_key="A_hat", est_label="A_hat"): + """One-row, four-panel edge recovery: A_true | A_est edges | confusion | stats.""" + A_true, N, num_exc, n_diag, n_off, shared_norm = self._get_A_true_info(md) + A_est = np.asarray(fitD[est_key]) + assert A_est.ndim == 2 and A_est.shape == A_true.shape, "A_est must match A_true shape" + + E_true = np.asarray(md["E_true"]).astype(bool) + if E_true.ndim == 3: + E_true = E_true[0] + assert E_true.shape == A_true.shape, "E_true shape must match A_true" + + off_diag = ~np.eye(N, dtype=bool) + E_t = E_true & off_diag + E_hat = (np.abs(A_est) > float(minW)) & off_diag + + TP = E_t & E_hat + FP = (~E_t) & E_hat + FN = E_t & (~E_hat) + TN = (~E_t) & (~E_hat) + + tp = int(TP.sum()); fp = int(FP.sum()); fn = int(FN.sum()); tn = int(TN.sum()) + precision = tp / max(1, tp + fp) + recall = tp / max(1, tp + fn) + f1 = 2 * precision * recall / max(1e-12, precision + recall) + acc = (tp + tn) / max(1, tp + tn + fp + fn) + + conf_map = np.zeros((N, N), dtype=np.int8) + conf_map[FN] = 1 + conf_map[FP] = 2 + conf_map[TP] = 3 + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 3.8)) + kw = dict(origin='lower', interpolation='nearest') + + ax = self.plt.subplot(1, 4, 1) + self._draw_A_matrix( + ax, A_true, f"True Dale, N{N}, nEdges={n_diag}+{n_off}", + num_exc=num_exc, norm_map=shared_norm + ) + + ax = self.plt.subplot(1, 4, 2) + ax.imshow(E_hat.astype(float), cmap='Greys', vmin=0, vmax=1, **kw) + ax.set(title=f"{est_label} edges (n={int(E_hat.sum())}, minW={float(minW):g})", + xlabel='presyn. neuron (output)', ylabel='postsyn. neuron (input)') + ax.plot([0, N-1], [0, N-1], '--', lw=0.8, color='magenta') + + cmap_conf = colors.ListedColormap(['white', 'magenta', 'red', 'green']) + norm_conf = colors.BoundaryNorm([-0.5, 0.5, 1.5, 2.5, 3.5], cmap_conf.N) + ax = self.plt.subplot(1, 4, 3) + im = ax.imshow(conf_map, cmap=cmap_conf, norm=norm_conf, **kw) + cbar = self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + cbar.set_ticks([0, 1, 2, 3]) + cbar.set_ticklabels(['TN', 'FN', 'FP', 'TP']) + ax.set(title='Confusion map', + xlabel='presyn. neuron', ylabel='postsyn. neuron') + ax.plot([0, N-1], [0, N-1], '--', lw=0.8, color='k') + + ax = self.plt.subplot(1, 4, 4) + vals = [tp, fp, fn] + ax.bar(['TP', 'FP', 'FN'], vals, color=['green', 'red', 'magenta']) + ax.set(title='stats', ylabel='count') + ax.grid(axis='y', alpha=0.4) + y_pos = 0.25 * max(1, max(vals)) + for name, val in zip(['TP', 'FP', 'FN'], vals): + ax.text(name, y_pos, str(val), ha='center', va='center', fontsize=10) + txt = (f"precision={precision:.3f}\nrecall={recall:.3f}\n" + f"f1={f1:.3f}\nacc={acc:.3f}") + ax.text(0.55, 0.60, txt, transform=ax.transAxes, fontsize=9) + + fig.suptitle( + f"A-matrix edge recovery, minW={float(minW):g} (off-diag only): {md['short_name']}", + fontsize=12) + fig.tight_layout() + + def _scatter_tp_true_vs_est(self, ax, a_true_tp, a_est_tp, title, color): + x_true = np.asarray(a_true_tp, dtype=np.float64).ravel() + y_init = np.asarray(a_est_tp, dtype=np.float64).ravel() + n = min(x_true.size, y_init.size) + if n <= 0: + ax.set(title=f"{title}\n(no TP points)", xlabel="A_init", ylabel="A_true") + ax.grid(True, alpha=0.35) + return + x_true = x_true[:n] + y_init = y_init[:n] + ax.scatter(y_init, x_true, s=10, marker='.', alpha=0.70, color=color, zorder=4) + self._add_x45_lins(ax, only45=True) + ax.axvline(0.0, linestyle='--', color='0.35', linewidth=0.9, alpha=0.7, zorder=1) + ax.axhline(0.0, linestyle='--', color='0.35', linewidth=0.9, alpha=0.7, zorder=1) + + lo = float(min(np.min(x_true), np.min(y_init))) + hi = float(max(np.max(x_true), np.max(y_init))) + span = max(1e-6, hi - lo) + pad = 0.05 * span + ax.set_xlim(lo - pad, hi + pad) + ax.set_ylim(lo - pad, hi + pad) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.35) + ax.set(title=f"{title} (n={n})", xlabel="A_init", ylabel="A_true") + + def _hist_tp_residuals(self, ax, a_true_tp, a_est_tp, title, color): + x = np.asarray(a_true_tp, dtype=np.float64).ravel() + y = np.asarray(a_est_tp, dtype=np.float64).ravel() + n = min(x.size, y.size) + if n <= 0: + ax.set(title=f"{title}\n(no TP points)", xlabel="A_init - A_true", ylabel="count") + ax.grid(True, alpha=0.35) + return + resid = y[:n] - x[:n] + bins = min(80, max(15, int(np.sqrt(n)))) + ax.hist(resid, bins=bins, color=color, alpha=0.8) + ax.axvline(0.0, color='k', linestyle='--', linewidth=1.0, alpha=0.8) + ax.grid(True, alpha=0.35) + ax.set(title=f"{title} (n={n})", xlabel="A_init - A_true", ylabel="count") + + mae = float(np.mean(np.abs(resid))) + rmse = float(np.sqrt(np.mean(resid ** 2))) + ax.text( + 0.04, 0.96, + f"MAE={mae:.3f}\nRMSE={rmse:.3f}", + transform=ax.transAxes, va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none"), + ) + + def _hist_all_values(self, ax, vals, title, color, tp_vals=None): + v = np.asarray(vals, dtype=np.float64).ravel() + if v.size == 0: + ax.set(title=f"{title}\n(no values)", xlabel="A_init", ylabel="count") + ax.grid(True, alpha=0.35) + return + bins = min(120, max(25, int(np.sqrt(v.size)))) + counts, edges, _ = ax.hist(v, bins=bins, color=color, alpha=0.8, label="all") + if tp_vals is not None: + tp = np.asarray(tp_vals, dtype=np.float64).ravel() + if tp.size > 0: + ax.hist( + tp, bins=edges, histtype='step', color='k', linewidth=1.8, + label=f"TP (n={tp.size})" + ) + ax.axvline(0.0, color='k', linestyle='--', linewidth=1.0, alpha=0.9) + ax.grid(True, alpha=0.35) + ax.set(title=f"{title} (n={v.size})", xlabel="A_init", ylabel="count") + ax.text( + 0.04, 0.96, + f"mean={float(np.mean(v)):.3f}\nstd={float(np.std(v)):.3f}", + transform=ax.transAxes, va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none"), + ) + if tp_vals is not None: + ax.legend(loc="upper right", fontsize=8, framealpha=0.8) + + def _hist2d_values_vs_row(self, ax, A, row_lo, row_hi, title, cmap="bwr"): + A = np.asarray(A, dtype=np.float64) + n_rows = int(A.shape[0]) + n_cols = int(A.shape[1]) + r0 = max(0, int(row_lo)) + r1 = min(n_rows, int(row_hi)) + if r1 <= r0: + ax.set(title=f"{title}\n(empty row range)", xlabel="A_init", ylabel="row neuron index") + ax.grid(True, alpha=0.35) + return + + sub = A[r0:r1, :] + xx = sub.ravel() + yy = np.repeat(np.arange(r0, r1, dtype=np.float64), n_cols) + xbins = min(140, max(40, int(np.sqrt(xx.size)))) + ybins = max(12, min(80, r1 - r0)) + + x_lo = float(np.min(xx)) + x_hi = float(np.max(xx)) + if np.isclose(x_lo, x_hi): + x_hi = x_lo + 1e-9 + x_edges = np.linspace(x_lo, x_hi, xbins + 1) + y_edges = np.linspace(r0 - 0.5, r1 - 0.5, ybins + 1) + + H, _, _ = np.histogram2d(xx, yy, bins=[x_edges, y_edges]) + x_cent = 0.5 * (x_edges[:-1] + x_edges[1:]) + sign_x = np.sign(x_cent) + # Normalize by neuron count (N columns) per user request. + H_signed = (H * sign_x[:, None]) / float(n_cols) + + vmin = float(np.min(H_signed)) + vmax = float(np.max(H_signed)) + if vmin >= 0.0: + vmin = -1e-9 + if vmax <= 0.0: + vmax = 1e-9 + norm = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + + im = ax.pcolormesh(x_edges, y_edges, H_signed.T, shading='auto', cmap=cmap, norm=norm) + ax.axvline(0.0, color='k', linestyle='--', linewidth=1.0, alpha=0.8) + ax.set(title=title, xlabel="A_init", ylabel="row neuron index") + ax.set_ylim(r0 - 0.5, r1 - 0.5) + ax.grid(False) + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + def _hist2d_values_vs_row_mask(self, ax, A, mask, title, cmap="bwr"): + A = np.asarray(A, dtype=np.float64) + M = np.asarray(mask, dtype=bool) + assert A.shape == M.shape, "A/mask shape mismatch" + n_rows = int(A.shape[0]) + n_cols = int(A.shape[1]) + + rr, cc = np.nonzero(M) + if rr.size <= 0: + ax.set(title=f"{title}\n(no values)", xlabel="A_init", ylabel="row neuron index") + ax.grid(True, alpha=0.35) + return + + xx = A[rr, cc] + yy = rr.astype(np.float64, copy=False) + xbins = min(140, max(40, int(np.sqrt(xx.size)))) + ybins = max(12, min(80, n_rows)) + + x_lo = float(np.min(xx)) + x_hi = float(np.max(xx)) + if np.isclose(x_lo, x_hi): + x_hi = x_lo + 1e-9 + x_edges = np.linspace(x_lo, x_hi, xbins + 1) + y_edges = np.linspace(-0.5, n_rows - 0.5, ybins + 1) + + H, _, _ = np.histogram2d(xx, yy, bins=[x_edges, y_edges]) + x_cent = 0.5 * (x_edges[:-1] + x_edges[1:]) + sign_x = np.sign(x_cent) + H_signed = (H * sign_x[:, None]) / float(max(1, n_cols)) + + vmin = float(np.min(H_signed)) + vmax = float(np.max(H_signed)) + if vmin >= 0.0: + vmin = -1e-9 + if vmax <= 0.0: + vmax = 1e-9 + norm = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + + im = ax.pcolormesh(x_edges, y_edges, H_signed.T, shading='auto', cmap=cmap, norm=norm) + ax.axvline(0.0, color='k', linestyle='--', linewidth=1.0, alpha=0.8) + ax.set(title=f"{title} (n={xx.size})", xlabel="A_init", ylabel="row neuron index") + ax.set_ylim(-0.5, n_rows - 0.5) + ax.grid(False) + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + def initA_quality_prismEM(self, fitD, md, minW=0.02, figId=7): + """Three-row diagnostics for A_init vs A_true with diag/exc/inh categories.""" + A_true, N, num_exc, n_diag, n_off, shared_norm = self._get_A_true_info(md) + A_init = np.asarray(fitD["A_init"], dtype=np.float64) + assert A_init.ndim == 2 and A_init.shape == A_true.shape, "A_init must match A_true" + + E_true = np.asarray(md["E_true"]).astype(bool) + if E_true.ndim == 3: + E_true = E_true[0] + assert E_true.shape == A_true.shape, "E_true shape must match A_true" + + diag_mask = np.eye(N, dtype=bool) + off_diag = ~diag_mask + E_t = E_true & off_diag + E_hat = (np.abs(A_init) > float(minW)) & off_diag + TP = E_t & E_hat + FP = (~E_t) & E_hat + FN = E_t & (~E_hat) + + # Exclusive categories: diagonal | excitatory off-diagonal rows | inhibitory off-diagonal rows. + presyn_is_exc = (np.arange(N, dtype=np.int64)[:, None] < int(num_exc)) + cat_exc = off_diag & presyn_is_exc + cat_inh = off_diag & (~presyn_is_exc) + cat_diag = diag_mask + tp_exc = TP & cat_exc + tp_inh = TP & cat_inh + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(14.0, 8.0)) + gs = fig.add_gridspec(3, 4, hspace=0.42, wspace=0.38) + + # Top-left: A_true matrix in the same style as -p c. + ax = fig.add_subplot(gs[0, 0]) + self._draw_A_matrix( + ax, A_true, f"True Dale, N{N}, nEdges={n_diag}+{n_off}", + num_exc=num_exc, norm_map=shared_norm + ) + + ax = fig.add_subplot(gs[0, 1]) + self._scatter_tp_true_vs_est( + ax, A_true[tp_exc], A_init[tp_exc], + "excitatory TP scatter", color="red" + ) + x0, x1 = ax.get_xlim() + ax.set_xlim(min(x0, 0.0), max(x1, 0.0)) + + ax = fig.add_subplot(gs[0, 2]) + self._scatter_tp_true_vs_est( + ax, A_true[tp_inh], A_init[tp_inh], + "inhibitory TP scatter", color="tab:blue" + ) + x0, x1 = ax.get_xlim() + ax.set_xlim(min(x0, 0.0), max(x1, 0.0)) + + ax = fig.add_subplot(gs[0, 3]) + self._scatter_tp_true_vs_est( + ax, A_true[cat_diag], A_init[cat_diag], + "diagonal scatter", color="magenta" + ) + x0, x1 = ax.get_xlim() + ax.set_xlim(min(x0, 0.0), max(x1, 0.0)) + + # Middle-left: A_init matrix in the same style as A_true. + ax = fig.add_subplot(gs[1, 0]) + self._draw_A_matrix( + ax, A_init, "A_init", num_exc=num_exc, norm_map=shared_norm + ) + + ax = fig.add_subplot(gs[1, 1]) + self._hist_all_values( + ax, A_init[cat_exc], + "A_init: excit off-diag", color="red", + tp_vals=A_init[tp_exc] + ) + + ax = fig.add_subplot(gs[1, 2]) + self._hist_all_values( + ax, A_init[cat_inh], + "A_init: inhib off-diag", color="tab:blue", + tp_vals=A_init[tp_inh] + ) + + ax = fig.add_subplot(gs[1, 3]) + self._hist_all_values( + ax, A_init[cat_diag], + "A_init: all diag", color="magenta", + tp_vals=None + ) + + # Third row: A_init value vs presyn neuron index. + ax = fig.add_subplot(gs[2, 0]) + self._hist2d_values_vs_row( + ax, A_init, 0, N, "A_init value vs row idx (all)", cmap="bwr" + ) + + ax = fig.add_subplot(gs[2, 1]) + self._hist2d_values_vs_row_mask( + ax, A_init, cat_exc, "excit off-diag", cmap="bwr" + ) + + ax = fig.add_subplot(gs[2, 2]) + self._hist2d_values_vs_row_mask( + ax, A_init, cat_inh, "inhib off-diag", cmap="bwr" + ) + + ax = fig.add_subplot(gs[2, 3]) + self._hist2d_values_vs_row_mask( + ax, A_init, cat_diag, "diagonal", cmap="bwr" + ) + + tp = int(TP.sum()) + fp = int(FP.sum()) + fn = int(FN.sum()) + tp_e = int(tp_exc.sum()) + tp_i = int(tp_inh.sum()) + n_d = int(cat_diag.sum()) + fig.suptitle( + f"A_init TP quality vs A_true, minW={float(minW):g}: {md['short_name']}\n" + f"off-diag TP={tp} (exc={tp_e}, inh={tp_i}), FP={fp}, FN={fn}; diag n={n_d}", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.93]) + + def _add_x45_lins(self, ax, only45=False): + lims = [ + np.min([ax.get_xlim(), ax.get_ylim()]), + np.max([ax.get_xlim(), ax.get_ylim()]), + ] + ax.plot(lims, lims, '--', color='k', linewidth=0.8) + if only45: + return + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + + def _corrcoef_safe(self, x, y): + x = np.asarray(x) + y = np.asarray(y) + if x.size == 0 or y.size == 0: + return np.nan + if np.std(x) == 0 or np.std(y) == 0: + return np.nan + return float(np.corrcoef(x, y)[0, 1]) + + def _find_bimodal_divider(self, xV): + x = np.asarray(xV, dtype=float).ravel() + assert x.size >= 2, f"Need >=2 samples for bimodal split, got {x.size}" + assert np.isfinite(x).all(), "xV contains non-finite values" + assert np.std(x) > 0.0, "xV must have non-zero variance for bimodal split" + + q1, q3 = np.quantile(x, [0.25, 0.75]) + c1, c2 = float(q1), float(q3) + if c1 == c2: + c1 = float(np.min(x)) + c2 = float(np.max(x)) + assert c1 != c2, "Failed to initialize two distinct mode centers" + + for _ in range(32): + d1 = np.abs(x - c1) + d2 = np.abs(x - c2) + left = d1 <= d2 + n_left = int(left.sum()) + n_right = int((~left).sum()) + assert n_left > 0 and n_right > 0, "Bimodal split produced empty cluster" + c1_new = float(np.mean(x[left])) + c2_new = float(np.mean(x[~left])) + if abs(c1_new - c1) < 1e-10 and abs(c2_new - c2) < 1e-10: + c1, c2 = c1_new, c2_new + break + c1, c2 = c1_new, c2_new + + if c1 > c2: + c1, c2 = c2, c1 + divide = 0.5 * (c1 + c2) + assert np.isfinite(divide), "Computed non-finite bimodal divider" + return float(divide) + + def _plot_corr_A_regions(self, ax, xV, yV, minW, title, xlab, ylab, s=6, alpha=0.5, color=None): + x = np.asarray(xV) + y = np.asarray(yV) + ax.scatter(x, y, s=s, alpha=alpha, color=color) + self._add_x45_lins(ax, only45=True) + ax.axvline(-minW, color="red", linestyle="--", linewidth=1) + ax.axvline(minW, color="red", linestyle="--", linewidth=1) + ax.set(title=title, xlabel=xlab, ylabel=ylab) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.4) + + mask_left = x < -minW + mask_mid = (x >= -minW) & (x <= minW) + mask_right = x > minW + r_left = self._corrcoef_safe(x[mask_left], y[mask_left]) + r_mid = self._corrcoef_safe(x[mask_mid], y[mask_mid]) + r_right = self._corrcoef_safe(x[mask_right], y[mask_right]) + n_left = int(mask_left.sum()) + n_mid = int(mask_mid.sum()) + n_right = int(mask_right.sum()) + + ax.text(0.20, 0.05, f"rL={r_left:.3f}\nnL={n_left}", transform=ax.transAxes) + ax.text(0.50, 0.40, f"rM={r_mid:.3f}\nnM={n_mid}", transform=ax.transAxes) + ax.text(0.65, 0.60, f"rR={r_right:.3f}\nnR={n_right}", transform=ax.transAxes) + + def _plot_corr_B_divisor(self, ax, xV, yV, title, xlab, ylab, s=6, alpha=0.5, color=None): + x = np.asarray(xV) + y = np.asarray(yV) + assert x.shape == y.shape, f"x and y shape mismatch: {x.shape} vs {y.shape}" + divide = self._find_bimodal_divider(x) + + ax.scatter(x, y, s=s, alpha=alpha, color=color) + self._add_x45_lins(ax, only45=True) + ax.axvline(divide, color="red", linestyle="--", linewidth=1) + ax.set(title=title, xlabel=xlab, ylabel=ylab) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.4) + + mask_left = x < divide + mask_right = x >= divide + r_left = self._corrcoef_safe(x[mask_left], y[mask_left]) + r_right = self._corrcoef_safe(x[mask_right], y[mask_right]) + n_left = int(mask_left.sum()) + n_right = int(mask_right.sum()) + ax.text(0.05, 0.45, f"rL={r_left:.3f}\nnL={n_left}", transform=ax.transAxes, va="top") + ax.text(0.65, 0.55, f"rR={r_right:.3f}\nnR={n_right}", transform=ax.transAxes, va="top") + + def eval_ABcorr_prismEM(self, fitD, md, figId=6): + A_hat = np.asarray(fitD["A_hat"]) + B_hat = np.asarray(fitD["B_hat"]) + A_true = np.asarray(md["A_true"]) + B_true = np.asarray(md["B_true"]) + minW = float(self.args.minW) + + assert A_true.ndim == 2, f"A_true must be 2D, got shape={A_true.shape}" + assert A_hat.ndim == 2, f"A_hat must be 2D, got shape={A_hat.shape}" + + if B_hat.ndim == 1: + B_hat = B_hat[None, :] + if B_true.ndim == 1: + B_true = B_true[None, :] + n_states = min(B_hat.shape[0], B_true.shape[0]) + assert n_states >= 1, "Need at least one B-state for correlation plot" + + figId = self.smart_append(figId) + ncol = 1 + n_states + fig_w = max(10.0, 3.0 * ncol) + fig = self.plt.figure(figId, facecolor='white', figsize=(fig_w, 3.8)) + + ax = self.plt.subplot(1, ncol, 1) + self._plot_corr_A_regions( + ax, A_true.ravel(), A_hat.ravel(), minW, + "A fit, non-zero ,", "A_true", "A_hat", s=6, alpha=0.4, color="green" + ) + + for m in range(n_states): + ax = self.plt.subplot(1, ncol, 2 + m) + self._plot_corr_B_divisor( + ax, B_true[m], B_hat[m], + f"B fit, state {m}", "B_true", "B_hat", + s=8, alpha=0.5, color="blue" + ) + + fig.suptitle(f"2D correlations, minW={minW:g}: {md['short_name']}", fontsize=12) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.92]) + + def state_seq_prismEM(self, fitD, md, figId=5, time_range_sec=None): + """4-row state-sequence canvas in the style of prism_Estep_eval -p b.""" + trainMD = md["train"] + dt = float(trainMD["time_step_sec"]) + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + + t0_req, t1_req = [float(x) for x in time_range_sec] + if t1_req < t0_req: + t0_req, t1_req = t1_req, t0_req + b0 = max(t0_bin, int(np.floor(t0_req / dt))) + b1 = min(t1_bin, int(np.floor(t1_req / dt))) + if b1 <= b0: + raise ValueError("Requested --time_range_sec leaves no bins in fitted window") + + i0 = b0 - t0_bin + i1 = b1 - t0_bin + p0 = i0 + p1 = i1 + + S_hat = np.asarray(fitD["S_hat"])[i0:i1 + 1] + C_hat = np.asarray(fitD["c_hat"])[i0:i1 + 1] + S_hat_CL = np.asarray(fitD["S_hat_CL"])[i0:i1 + 1] + + S_true = np.asarray(md["S_true"])[b0:b1 + 1] + C_true = np.asarray(md["C_true"])[b0:b1 + 1] + + nll_t = np.asarray(md["eval_f"]["loss_nll_time"])[p0:p1] + l2_t = np.asarray(md["eval_f"]["loss_l2_time"])[p0:p1] + t_bins = np.arange(b0, b1 + 1, dtype=np.float64) * dt + t_pairs = np.arange(b0, b1, dtype=np.float64) * dt + x0, x1 = float(t_bins[0]), float(t_bins[-1]) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 9)) + # Add an explicit spacer row between diagnostics (top row) and Fit row + # so top-row x labels do not collide with Fit legends. + gs = fig.add_gridspec( + 5, 3, + height_ratios=[0.95, 0.18, 1.0, 1.0, 0.55], + hspace=0.75, + wspace=0.35, + ) + + # Row 2: Fit + ax = fig.add_subplot(gs[2, :]) + ax.plot(t_bins, S_hat, color="k", linewidth=2.5, label="S_hat") + ax.plot(t_bins, S_true, color="lime", linestyle="--", linewidth=1.5, label="S_true") + lo = np.clip(S_hat - S_hat_CL, 0.0, float(C_hat.shape[1] - 1)) + hi = np.clip(S_hat + S_hat_CL, 0.0, float(C_hat.shape[1] - 1)) + ax.fill_between(t_bins, lo, hi, color="gray", alpha=0.3, label="S_hat_CL") + ax.set(title=f"Fit (acc={float(md['eval_f']['acc']):.3f})", xlabel="time (s)", ylabel="state") + ax.set_xlim(x0, x1) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax2 = ax.twinx() + for m in range(C_hat.shape[1]): + ax2.plot(t_bins, C_hat[:, m], linewidth=0.8, alpha=0.85, label=f"C_hat[{m}]") + ax2.set_ylabel("C_hat") + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + # Row 3: Truth + ax = fig.add_subplot(gs[3, :]) + ax.plot(t_bins, S_true, color="lime", linestyle="--",linewidth=1.5, label="S_true") + ax.set(title="Truth", xlabel="time (s)", ylabel="state") + ax.set_xlim(x0, x1) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax2 = ax.twinx() + for m in range(C_true.shape[1]): + ax2.plot(t_bins, C_true[:, m], linewidth=0.8, alpha=0.85, label=f"C_true[{m}]") + ax2.set_ylabel("C_true") + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + # Row 1: diagnostics panels (occupancy, state accuracy, confidence) + srec = md["states_recovery_eval"] + acc_avg = float(srec["avg_acc"]) + state_acc_cl = np.asarray(srec["state_acc_cl"], dtype=np.float64) + acc_ps = state_acc_cl[:, 0] + cl_ps = state_acc_cl[:, 1] + + ax = fig.add_subplot(gs[0, 0]) + occ = C_hat.mean(axis=0) + x = np.arange(C_hat.shape[1], dtype=np.int64) + ax.bar(x, occ, color="tab:blue", alpha=0.65) + ax.set(title="Mean occupancy", xlabel="state", ylabel="mean c") + ax.set_ylim(0.0, 1.0) + ax.set_xticks(x) + ax.grid(True, alpha=0.3) + ax.set_box_aspect(0.9) + + ax = fig.add_subplot(gs[0, 1]) + x = np.arange(acc_ps.size, dtype=np.int64) + ax.bar(x, acc_ps, color="tab:green", alpha=0.7) + ax.errorbar(x, acc_ps, yerr=cl_ps, fmt='o', color='k', capsize=5, markersize=4) + ax.axhline(1.0, color="green", linestyle="--", linewidth=1) + ymin = min(0.5, float(np.nanmin(acc_ps - cl_ps)) - 0.05) + ax.set_ylim(ymin, None) + ax.set(title=f"State accuracy, avr={acc_avg:.3f}", xlabel="state", ylabel="accuracy") + ax.set_xticks(x) + ax.grid(True, alpha=0.3, axis="y") + ylo, yhi = ax.get_ylim() + y_txt = ylo + 0.80 * (yhi - ylo) + for i, v in enumerate(acc_ps): + ax.text(i + 0.3, y_txt, f"{v:.2f}", ha="center", va="center", fontsize=9) + ax.set_box_aspect(0.9) + + ax = fig.add_subplot(gs[0, 2]) + cl = np.asarray(S_hat_CL, dtype=np.float64) + ax.hist(cl, bins=30, color="tab:purple", alpha=0.7) + ax.axvline(np.mean(cl), color="k", linestyle="--", linewidth=1.0) + ax.set(title=f"Confidence (acc={acc_avg:.3f})", xlabel="S_hat_CL", ylabel="count") + ax.grid(True, alpha=0.3) + ax.set_box_aspect(0.9) + + # Row 4: Loss(time) + ax = fig.add_subplot(gs[4, :]) + ax.plot(t_pairs, nll_t, color="tab:blue", linewidth=0.9, label="nll") + ax.plot(t_pairs, l2_t, color="tab:orange", linewidth=0.9, label="l2") + ax.set_yscale("log") + ax.set(title="Loss (time)", xlabel="time (s)", ylabel="value") + ax.set_xlim(x0, x1) + ax.grid(True, alpha=0.35) + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), ncol=2, fontsize=8) + ax2 = ax.twinx() + ax2.plot(t_bins, S_true, color="k", linewidth=0.8, alpha=0.6, label="S_true") + ax2.set_ylabel("state") + ax2.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), fontsize=8) + + fig.suptitle( + f"Dataset: {md['short_name']} | λ₂={trainMD['lambda2']}, lr={trainMD['lr_estep']}, " + f"pgd_iter={trainMD['pgd_iter']}, decode_dwell={srec['decode_dwell_sec']}s", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.965]) diff --git a/causal_net/nonStation_ver3a_states_simu/PlotterPrismEstep.py b/causal_net/nonStation_ver3a_states_simu/PlotterPrismEstep.py new file mode 100644 index 00000000..ff412c39 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/PlotterPrismEstep.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for prism E-step evaluation. +""" + +from toolbox.PlotterBackbone import PlotterBackbone +import numpy as np +from matplotlib.ticker import MaxNLocator + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def _rebin_1d(self, x, merge): + merge = int(max(1, merge)) + if merge <= 1: + return np.asarray(x) + x = np.asarray(x) + n = x.shape[0] // merge + if n <= 0: + return x + return x[: n * merge].reshape(n, merge).mean(axis=1) + + def _rebin_2d(self, x, merge): + merge = int(max(1, merge)) + if merge <= 1: + return np.asarray(x) + x = np.asarray(x) + n = x.shape[0] // merge + if n <= 0: + return x + return x[: n * merge].reshape(n, merge, x.shape[1]).mean(axis=1) + + def _rebin_mode_1d(self, x, merge): + merge = int(max(1, merge)) + x = np.asarray(x) + if merge <= 1: + return x.astype(np.int64, copy=False) + n = x.shape[0] // merge + if n <= 0: + return x.astype(np.int64, copy=False) + x = x[: n * merge].reshape(n, merge) + out = np.zeros((n,), dtype=np.int64) + for i in range(n): + vals = x[i].astype(np.int64, copy=False) + out[i] = np.bincount(vals).argmax() + return out + + def plot_loss_time(self, ax, fitD, md, trainMD, time_bin_merge=20): + loss_nll_time = np.asarray(fitD["loss_nll_time"]) + loss_l2_time = np.asarray(fitD["loss_l2_time"]) + + merge = int(max(1, time_bin_merge)) + if merge > 1 and len(loss_nll_time) >= merge: + n = len(loss_nll_time) // merge + loss_nll_time = loss_nll_time[: n * merge].reshape(n, merge).mean(axis=1) + loss_l2_time = loss_l2_time[: n * merge].reshape(n, merge).mean(axis=1) + + dt = trainMD["time_step_sec"] + time_range_bins = trainMD["time_range_bins"] + t0_bin = time_range_bins[0] + t = (t0_bin + np.arange(len(loss_nll_time)) * merge) * float(dt) + + ax.plot(t, loss_nll_time, color="tab:blue", linewidth=0.8, label="nll") + ax.plot(t, loss_l2_time, color="tab:orange", linewidth=0.8, label="l2") + ax.set(title="Loss (time)", xlabel="time (s)", ylabel="value") + ax.grid(True, alpha=0.4) + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), ncol=2, fontsize=8) + ax.set_yscale('log') + + s_true = np.asarray(md["S_true"]) + s_true = s_true[t0_bin + 1 : t0_bin + 1 + len(fitD["loss_nll_time"])] + s_true_rb = self._rebin_mode_1d(s_true, merge) + + ax2 = ax.twinx() + ax2.plot(t, s_true_rb, color="k", linewidth=0.8, alpha=0.6, label="S_true") + ax2.set_ylabel("state") + ax2.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), fontsize=8) + + def plot_ll_gap(self, ax, md, trainMD, time_bin_merge=20): + ll_gap = np.asarray(md["eval_Estep"]["ll_gap"]) + merge = int(max(1, time_bin_merge)) + if merge > 1 and len(ll_gap) >= merge: + n = len(ll_gap) // merge + ll_gap = ll_gap[: n * merge].reshape(n, merge).mean(axis=1) + + dt = trainMD["time_step_sec"] + t0_bin = trainMD["time_range_bins"][0] + t = (t0_bin + np.arange(len(ll_gap)) * merge) * float(dt) + + ax.plot(t, ll_gap, color="tab:blue", linewidth=0.8, label="LL gap") + ax.set(title="LL gap (best - 2nd)", xlabel="time (s)", ylabel="value") + ax.grid(True, alpha=0.4) + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + + def summary_prismEstep(self, fitD, md, figId=1, time_bin_merge=20): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(10, 5.5)) + gs = fig.add_gridspec(2, 3, height_ratios=[1.0, 1.0], hspace=0.6, wspace=0.25) + + c_hat = fitD["c_hat"] + S_hat = fitD["S_hat"] + S_hat_CL = fitD["S_hat_CL"] + + trainMD = md["train"] + decE = md["decode_eval"] + acc = md["eval_Estep"]["acc"] + + # ---- State occupancy (mean c_hat) + hyperparams text + ax = fig.add_subplot(gs[0, 0]) + c_hat = np.asarray(c_hat) + occ = c_hat.mean(axis=0) + ax.bar(np.arange(len(occ)), occ, color="tab:blue", alpha=0.7) + ax.set(title="Mean occupancy", xlabel="state", ylabel="mean c") + ax.set_ylim(0.0, 1.0) + ax.grid(True, alpha=0.3) + txt = ( + f"lr={trainMD['lr']}\n" + f"lambda2={trainMD['lambda2']}\n" + f"pgd_iter={trainMD['pgd_iter']}\n" + f"chunk={trainMD['chunk_size']}" + ) + ax.text(0.98, 0.95, txt, transform=ax.transAxes, va="top", ha="right", fontsize=9) + + # ---- Per-state accuracy bar chart with CL error bars + ax = fig.add_subplot(gs[0, 1]) + acc_ps = np.asarray(decE.get("acc_per_state", [])) + cl_ps = np.asarray(decE.get("avg_cl_per_state", [])) + M = len(acc_ps) + x = np.arange(M) + ax.bar(x, acc_ps, color="tab:green", alpha=0.7) + if len(cl_ps) == M: + ax.errorbar(x, acc_ps, yerr=cl_ps, fmt='o', color='k', + capsize=5, markersize=4) + ax.axhline(1.0, color="green", linestyle="--", linewidth=1) + ymin = min(0.5, float(np.nanmin(acc_ps - cl_ps)) - 0.05) if len(cl_ps) == M else 0.5 + ax.set_ylim(ymin, None) + ax.set(title=f"State accuracy, avr={acc:.3f}", xlabel="state", ylabel="accuracy") + ax.set_xticks(x) + ax.grid(True, alpha=0.3, axis="y") + ylo, yhi = ax.get_ylim() + y_txt = ylo + 0.80 * (yhi - ylo) + for i, v in enumerate(acc_ps): + ax.text(i + 0.3, y_txt, f"{v:.2f}", ha="center", va="center", fontsize=9) + + # ---- Confidence histogram + ax = fig.add_subplot(gs[0, 2]) + cl = np.asarray(S_hat_CL) + ax.hist(cl, bins=30, color="tab:purple", alpha=0.7) + ax.set(title=f"Confidence (acc={acc:.3f})", xlabel="S_hat_CL", ylabel="count") + ax.grid(True, alpha=0.3) + ax.axvline(np.mean(cl), color="k", linestyle="--", linewidth=1) + + # ---- Loss vs time + ax = fig.add_subplot(gs[1, :]) + self.plot_loss_time(ax, fitD, md, trainMD, time_bin_merge=time_bin_merge) + + t0s, t1s = trainMD["time_range_sec"] + fig.suptitle(f"Prism E-step: {md['short_name']}, T=[{t0s:.0f}, {t1s:.0f}] sec", fontsize=12) + + def state_seq_prismEstep(self, fitD, md, figId=2, time_reb=20): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 8.5)) + gs = fig.add_gridspec(4, 1, height_ratios=[1.0, 1.0, 0.5, 0.5], hspace=0.9) + + S_hat = fitD["S_hat"] + S_hat_CL = fitD["S_hat_CL"] + C_hat = fitD["c_hat"] + S_true = md["S_true"] + C_true = md["C_true"] + + trainMD = md["train"] + time_range_bins = trainMD["time_range_bins"] + t0_bin, t1_bin = time_range_bins + + S_true = np.asarray(S_true)[t0_bin : t1_bin + 1] + C_true = np.asarray(C_true)[t0_bin : t1_bin + 1] + + S_hat = np.asarray(S_hat) + S_hat_CL = np.asarray(S_hat_CL) + C_hat = np.asarray(C_hat) + + S_hat_rb = self._rebin_mode_1d(S_hat, time_reb) + C_hat_rb = self._rebin_2d(C_hat, time_reb) + S_true_rb = self._rebin_mode_1d(S_true, time_reb) + C_true_rb = self._rebin_2d(C_true, time_reb) + CL_hat_rb = self._rebin_1d(S_hat_CL, time_reb) + + dt = trainMD["time_step_sec"] + t = (t0_bin + np.arange(len(S_hat_rb)) * max(1, int(time_reb))) * float(dt) + + ax = fig.add_subplot(gs[0, 0]) + ax.plot(t, S_hat_rb, color="k", linewidth=1.0, label="S_hat") + n_states = C_hat_rb.shape[1] + lo = np.clip(S_hat_rb - CL_hat_rb, 0.0, float(n_states - 1)) + hi = np.clip(S_hat_rb + CL_hat_rb, 0.0, float(n_states - 1)) + ax.fill_between(t, lo, hi, color="gray", alpha=0.3, label="S_hat_CL") + acc = md["eval_Estep"]["acc"] + ax.set(title=f"Fit (acc={acc:.3f})", xlabel="time (s)", ylabel="state") + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + for m in range(C_hat_rb.shape[1]): + ax2.plot(t, C_hat_rb[:, m], linewidth=0.8, alpha=0.8, label=f"C_hat[{m}]") + ax2.set_ylabel("C_hat") + + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + ax = fig.add_subplot(gs[1, 0]) + ax.plot(t, S_true_rb, color="k", linewidth=1.0, label="S_true") + ax.set(title="Truth ", xlabel="time (s)", ylabel="state") + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + for m in range(C_true_rb.shape[1]): + ax2.plot(t, C_true_rb[:, m], linewidth=0.8, alpha=0.8, label=f"C_true[{m}]") + ax2.set_ylabel("C_true") + + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + ax = fig.add_subplot(gs[2, 0]) + self.plot_loss_time(ax, fitD, md, trainMD, time_bin_merge=time_reb) + + ax = fig.add_subplot(gs[3, 0]) + self.plot_ll_gap(ax, md, trainMD, time_bin_merge=time_reb) + + fig.suptitle( + f"Dataset: {md['short_name']} | $\\lambda_2$={trainMD['lambda2']}, " + f"lr={trainMD['lr']}, pgd_iter={trainMD['pgd_iter']}, decode_dwell={md['decode_eval']['decode_dwell_sec']}s", + fontsize=12, + ) diff --git a/causal_net/nonStation_ver3a_states_simu/PlotterPrismMstep.py b/causal_net/nonStation_ver3a_states_simu/PlotterPrismMstep.py new file mode 100644 index 00000000..54581d21 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/PlotterPrismMstep.py @@ -0,0 +1,567 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for prism M-step evaluation. +""" + +from toolbox.PlotterBackbone import PlotterBackbone +import numpy as np +import matplotlib.colors as colors + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def _rebin_1d(self, x, merge): + merge = int(max(1, merge)) + if merge <= 1: + return np.asarray(x) + x = np.asarray(x) + n = x.shape[0] // merge + if n <= 0: + return x + return x[: n * merge].reshape(n, merge).mean(axis=1) + + def _rebin_2d(self, x, merge): + merge = int(max(1, merge)) + if merge <= 1: + return np.asarray(x) + x = np.asarray(x) + n = x.shape[0] // merge + if n <= 0: + return x + return x[: n * merge].reshape(n, merge, x.shape[1]).mean(axis=1) + + def summary_prismMstep(self, fitD, md, figId=1): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 5.)) + + trainMD = md["train"] + short_name = md["short_name"] + total = np.asarray(fitD["loss_epoch"]) + l1 = np.asarray(fitD["loss_l1_epoch"]) + sparsity = np.asarray(fitD["sparsity_epoch"]) + rho = np.asarray(fitD["rho_epoch"]) + A_fit = np.asarray(fitD["A_hat"]) + Freq = np.asarray(fitD["single_rates"]) + loss_state = np.asarray(fitD["loss_state_total_epoch"]) + + assert A_fit.ndim == 2, f"A_hat must be 2D, got shape={A_fit.shape}" + assert Freq.ndim == 1, f"single_rates must be 1D, got shape={Freq.shape}" + assert loss_state.ndim == 2, f"loss_state_total_epoch must be 2D, got shape={loss_state.shape}" + assert loss_state.shape[0] == total.shape[0], "loss_state_total_epoch first dim must equal num epochs" + + n_epochs = len(total) + epoch0 = 10 if n_epochs > 10 else 0 + epochs = np.arange(1 + epoch0, n_epochs + 1, dtype=np.int32) + + # ---- Row1 Col1: Loss + L1 (dual y-axis) + ax = self.plt.subplot(2, 4, 1) + ax.plot(epochs, total[epoch0:], label='total', color='blue', linestyle='-') + ax.set_xlabel('Epoch') + ax.set_ylabel('Loss', color='blue') + ax.tick_params(axis='y', labelcolor='blue') + ax.grid(True, alpha=0.3) + add_delay_markers(ax, trainMD) + + ax2 = ax.twinx() + ax2.plot(epochs, l1[epoch0:], label='L1 loss', color='red', linestyle='--') + ax2.set_ylabel('L1 loss', color='red') + ax2.tick_params(axis='y', labelcolor='red') + + lines1, labels1 = ax.get_legend_handles_labels() + lines2, labels2 = ax2.get_legend_handles_labels() + n_samp = int(trainMD["num_samples_used"]) + ax.legend(lines1 + lines2, labels1 + labels2, title=f"{n_samp//1000}k samp") + ax.set_title(f"Loss: {short_name}") + + # ---- Row1 Col2: A off-diagonal histogram + ax = self.plt.subplot(2, 4, 2) + Nn = A_fit.shape[0] + diag_mask = np.eye(Nn, dtype=bool) + A_off = A_fit[~diag_mask] + valid_mask = np.abs(A_off) > 1e-10 + A_off_clean = A_off[valid_mask] + n_edges = int(A_off_clean.size) + ax.hist(A_off_clean, bins=100, color='g', alpha=0.8) + ax.set_yscale('log') + ax.grid(True, alpha=0.4) + ax.set(ylabel='edges', xlabel='edge value', title='A off-diagonal') + acc_frac = n_edges / float(Nn * (Nn - 1)) + txt = f"edge sel meth: None\nacc frac={acc_frac:.2f}" + ax.text(0.05, 0.75, txt, transform=ax.transAxes, fontsize=10) + + # ---- Row1 Col3: A sparsity vs epoch + ax = self.plt.subplot(2, 4, 3) + ax.plot(np.arange(1, len(sparsity) + 1), sparsity, color='tab:purple', linewidth=1.5) + add_delay_markers(ax, trainMD) + ax.set_ylim(0.0, 1.02) + ax.set(ylabel='fraction', xlabel='epoch', title='A sparsity (off-diag)') + ax.grid(True, alpha=0.4) + + # ---- Row1 Col4: Spectral radius vs epoch + ax = self.plt.subplot(2, 4, 4) + ax.plot(np.arange(1, len(rho) + 1), rho, color='tab:red', linewidth=1.5) + add_delay_markers(ax, trainMD) + ax.axhline(float(trainMD["rho_max"]), color='k', linestyle='--', linewidth=1.0) + ax.set(ylabel='radius', xlabel='epoch', title='Spectral radius(A)') + ax.grid(True, alpha=0.4) + + # ---- Row2 Col3: per-state total loss vs epoch + ax = self.plt.subplot(2, 4, 5) + epoch_all = np.arange(1, loss_state.shape[0] + 1, dtype=np.int32) + for m in range(loss_state.shape[1]): + ax.plot(epoch_all, loss_state[:, m], linewidth=1.2, label=f'state {m}') + add_delay_markers(ax, trainMD) + ax.set(ylabel='loss', xlabel='epoch', title='Per-state total loss') + ax.grid(True, alpha=0.4) + ax.legend(fontsize=8, ncol=2) + + + # ---- Row2 Col1: histogram of rates + ax = self.plt.subplot(2, 4, 6) + ax.hist(Freq, bins=20) + ax.set(ylabel='num neurons', xlabel='Firing rate (Hz)', title=f'Single rates, {Nn} neurons') + ax.grid(True, alpha=0.4) + ax.set_xlim(0,) + median_val = float(np.median(Freq)) + ax.text(median_val, ax.get_ylim()[1] * 0.8, f'median rate: {median_val:.2f} Hz', + color='red', ha='left', va='bottom', fontsize=10) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + + # ---- Row2 Col2: 2D histogram edge value vs log10(rate) + ax = self.plt.subplot(2, 4, 7) + i_indices, j_indices = np.where(~diag_mask) + Freq_expanded = Freq[i_indices] + Freq_expanded_clean = Freq_expanded[valid_mask] + h = ax.hist2d(A_off_clean, np.log10(Freq_expanded_clean), bins=30, cmap='viridis', norm=colors.LogNorm()) + fig.colorbar(h[3], ax=ax) + ax.set(ylabel='log10( Firing rate/ Hz )', xlabel='edge value', + title=f'accept {n_edges} of {Nn*(Nn-1)} edges') + ax.grid(True, alpha=0.4) + + # ---- Row2 Col4: intentionally empty + ax = self.plt.subplot(2, 4, 8) + ax.axis("off") + + def eval_Amatrix_prismMstep(self, fitD, md, figId=2): + """One-row, four-panel A-matrix evaluation: E_true | A_hat edges | confusion map | stats bar.""" + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 3.8)) + + A_hat = np.asarray(fitD["A_hat"]) # (N, N) + if A_hat.ndim == 3: + A_hat = A_hat[0] + N = A_hat.shape[0] + + edge = md.get("edge_eval_main", None) + if edge is None: + raise ValueError("Missing md['edge_eval_main']; compute edge metrics in prism_Mstep_eval.py") + + minW = float(edge["minW"]) + E_hat = np.asarray(edge["E_hat"]).astype(bool) + E_t = np.asarray(edge["E_true_offdiag"]).astype(bool) + conf_map = np.asarray(edge["conf_map"]) + + tp = int(edge["tp"]); fp = int(edge["fp"]) + fn = int(edge["fn"]); tn = int(edge["tn"]) + precision = float(edge["precision"]) + recall = float(edge["recall"]) + f1 = float(edge["f1"]) + acc = float(edge["acc"]) + + kw = dict(origin='lower', interpolation='nearest') + + # ---- A_true (signed weights, bwr, E/I boundary) + A_true = np.asarray(md["A_true"]) + dmd = md.get("dale_conf", {}) + numExc = int(dmd.get("num_excite", N)) + vmin_t, vmax_t = A_true.min(), A_true.max() + norm_t = colors.TwoSlopeNorm(vmin=vmin_t, vcenter=0.0, vmax=vmax_t) + n_exc_edges = int(E_t[:, :numExc].sum()) + n_inh_edges = int(E_t[:, numExc:].sum()) + + ax = self.plt.subplot(1, 4, 1) + im = ax.imshow(A_true, cmap='bwr', norm=norm_t, aspect=1., **kw) + self.plt.colorbar(im, ax=ax, extend='both', shrink=0.7) + ax.axhline(numExc - 0.5, color='k', ls='--') + ax.axvline(numExc - 0.5, color='k', ls='--') + ax.plot([0, N], [0, N], '--', lw=0.8, color='magenta') + ax.set(title=f"True Dale, N{N}, nEdges={n_exc_edges}+{n_inh_edges}", + xlabel='presyn. neuron index (output)', + ylabel='postsyn. neuron index (input)') + + # ---- A_hat edges + ax = self.plt.subplot(1, 4, 2) + ax.imshow(E_hat.astype(float), cmap='Greys', vmin=0, vmax=1, **kw) + ax.set(title=f"A_hat edges (n={int(E_hat.sum())}, minW={minW:g})", + xlabel='presyn. neuron (output)', ylabel='postsyn. neuron (input)') + ax.plot([0, N-1], [0, N-1], '--', lw=0.8, color='magenta') + + # ---- Confusion map + cmap_conf = colors.ListedColormap(['white', 'magenta', 'red', 'green']) + norm_conf = colors.BoundaryNorm([-0.5, 0.5, 1.5, 2.5, 3.5], cmap_conf.N) + ax = self.plt.subplot(1, 4, 3) + im = ax.imshow(conf_map, cmap=cmap_conf, norm=norm_conf, **kw) + cbar = self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + cbar.set_ticks([0, 1, 2, 3]) + cbar.set_ticklabels(['TN', 'FN', 'FP', 'TP']) + ax.set(title='Confusion map', + xlabel='presyn. neuron', ylabel='postsyn. neuron') + ax.plot([0, N-1], [0, N-1], '--', lw=0.8, color='k') + + # ---- Stats bar + ax = self.plt.subplot(1, 4, 4) + vals = [tp, fp, fn] + bars = ax.bar(['TP', 'FP', 'FN'], vals, color=['green', 'red', 'magenta']) + ax.set(title='stats', ylabel='count') + ax.grid(axis='y', alpha=0.4) + y_pos = 0.25 * max(1, max(vals)) + for name, val in zip(['TP', 'FP', 'FN'], vals): + ax.text(name, y_pos, str(val), ha='center', va='center', fontsize=10) + txt = (f"precision={precision:.3f}\nrecall={recall:.3f}\n" + f"f1={f1:.3f}\nacc={acc:.3f}") + ax.text(0.55, 0.60, txt, transform=ax.transAxes, fontsize=9) + + fig.suptitle( + f"A-matrix edge recovery, minW={minW:g} (off-diag only): {md['short_name']}", + fontsize=12) + fig.tight_layout() + + def state_seq_prismMstep(self, estepD, md, figId=2, time_reb=20): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 6)) + + assert estepD is not None, "Missing E-step fit (S_hat/C_hat) for state sequence plot." + + S_hat = estepD["S_hat"] + S_hat_CL = estepD["S_hat_CL"] + C_hat = estepD["c_hat"] + S_true = md["S_true"] + C_true = md["C_true"] + + trainMD = md["estep_train"] + time_range_bins = trainMD["time_range_bins"] + t0_bin, t1_bin = time_range_bins + + S_true = np.asarray(S_true) + C_true = np.asarray(C_true) + S_true = S_true[t0_bin : t1_bin + 1] + C_true = C_true[t0_bin : t1_bin + 1] + + S_hat = np.asarray(S_hat) + S_hat_CL = np.asarray(S_hat_CL) + C_hat = np.asarray(C_hat) + + # Rebin + S_hat_rb = self._rebin_1d(S_hat, time_reb) + C_hat_rb = self._rebin_2d(C_hat, time_reb) + S_true_rb = self._rebin_1d(S_true, time_reb) + C_true_rb = self._rebin_2d(C_true, time_reb) + CL_hat_rb = self._rebin_1d(S_hat_CL, time_reb) + + dt = trainMD["time_step_sec"] + t = (t0_bin + np.arange(len(S_hat_rb)) * max(1, int(time_reb))) * float(dt) + + ax = self.plt.subplot(2, 1, 1) + ax.plot(t, S_hat_rb, color="k", linewidth=1.0, label="S_hat") + ax.fill_between(t, S_hat_rb - CL_hat_rb, S_hat_rb + CL_hat_rb, color="gray", alpha=0.3, label="S_hat_CL") + for m in range(C_hat_rb.shape[1]): + ax.plot(t, C_hat_rb[:, m], linewidth=0.8, alpha=0.8, label=f"C_hat[{m}]") + ax.set(title="Fit: S_hat and C_hat", xlabel="time (s)", ylabel="state") + ax.grid(True, alpha=0.3) + ax.legend(ncol=4, fontsize=8) + + ax = self.plt.subplot(2, 1, 2) + ax.plot(t, S_true_rb, color="k", linewidth=1.0, label="S_true") + for m in range(C_true_rb.shape[1]): + ax.plot(t, C_true_rb[:, m], linewidth=0.8, alpha=0.8, label=f"C_true[{m}]") + ax.set(title="Truth: S_true and C_true", xlabel="time (s)", ylabel="state") + ax.grid(True, alpha=0.3) + ax.legend(ncol=4, fontsize=8) + + fig.suptitle(f"State sequence: {md['short_name']}", fontsize=12) + + def state_corr_prismMstep(self, fitD, md, figId=2): + self.state_corr_truth_prismMstep(fitD, md, type="fit", figId=figId) + + def state_Bcorr_prismMstep(self, fitD, md, figId=5): + figId = self.smart_append(figId) + + B_hat = np.asarray(fitD["B_hat"]) + B_true = np.asarray(md["B_true"]) + if B_hat.ndim == 1: + B_hat = B_hat[None, :] + if B_true.ndim == 1: + B_true = B_true[None, :] + + n_states = min(B_hat.shape[0], B_true.shape[0]) + if n_states < 2: + fig = self.plt.figure(figId, facecolor='white', figsize=(8, 3)) + ax = self.plt.subplot(1, 1, 1) + ax.axis("off") + ax.text(0.05, 0.6, "Need at least 2 states for pairwise B-correlation", fontsize=11) + fig.suptitle(f"State B correlations: {md['short_name']}", fontsize=12) + return + + pairs = [(i, j) for i in range(n_states) for j in range(i + 1, n_states)] + ncol = len(pairs) + fig_w = max(10.0, 3.2 * ncol) + fig = self.plt.figure(figId, facecolor='white', figsize=(fig_w, 6.2)) + + for col, (i, j) in enumerate(pairs): + ax = self.plt.subplot(2, ncol, 1 + col) + plot_corr_B_divisor( + ax, B_hat[i], B_hat[j], + f"B_hat corr, s{i}-s{j}", f"B_hat s{i}", f"B_hat s{j}", + s=8, alpha=0.6, color="orange" + ) + + ax = self.plt.subplot(2, ncol, ncol + 1 + col) + plot_corr_B_divisor( + ax, B_true[i], B_true[j], + f"B_true corr, s{i}-s{j}", f"B_true s{i}", f"B_true s{j}", + s=8, alpha=0.6, color="seagreen" + ) + + fig.suptitle(f"State B correlations: {md['short_name']}", fontsize=12) + + def eval_ABcorr_prismMstep(self, fitD, md, figId=2): + A_hat = fitD["A_hat"] + B_hat = fitD["B_hat"] + A_true = md["A_true"] + B_true = md["B_true"] + figId = self.smart_append(figId) + minW = float(self.args.minW) + + A_hat = np.asarray(A_hat) + B_hat = np.asarray(B_hat) + A_true = np.asarray(A_true) + B_true = np.asarray(B_true) + + assert A_true.ndim == 2, f"A_true must be 2D, got shape={A_true.shape}" + assert A_hat.ndim == 2, f"A_hat must be 2D, got shape={A_hat.shape}" + + if B_hat.ndim == 1: + B_hat = B_hat[None, :] + if B_true.ndim == 1: + B_true = B_true[None, :] + n_states = min(B_hat.shape[0], B_true.shape[0]) + assert n_states >= 1, "Need at least one B-state for correlation plot" + + ncol = 1 + n_states + fig_w = max(10.0, 3.0 * ncol) + fig = self.plt.figure(figId, facecolor='white', figsize=(fig_w, 3.8)) + + ax = self.plt.subplot(1, ncol, 1) + plot_corr_A_regions(ax, A_true.ravel(), A_hat.ravel(), minW, "A corr,", "A_true", "A_hat", s=6, alpha=0.4, color="green") + + for m in range(n_states): + ax = self.plt.subplot(1, ncol, 2 + m) + plot_corr_B_divisor(ax, B_true[m], B_hat[m], f"B corr, state {m}", "B_true", "B_hat", s=8, alpha=0.5, color="blue") + + fig.suptitle(f"2D correlations, minW={minW:g}: {md['short_name']}", fontsize=12) + + def edge_state_prismMstep(self, fitD, md, figId=6): + A_hat = fitD["A_hat"] + figId = self.smart_append(figId) + edge_states = md.get("edge_eval_states", None) + if edge_states is None: + raise ValueError("Missing md['edge_eval_states']; compute edge metrics in prism_Mstep_eval.py") + minW = float(edge_states[0]["minW"]) + + A_hat = np.asarray(A_hat) + if A_hat.ndim == 2: + A_hat = A_hat[None, :, :] + n_states = min(A_hat.shape[0], len(edge_states)) + N = A_hat.shape[1] + fig_h = max(2.8, 2.8 * n_states) + fig = self.plt.figure(figId, facecolor='white', figsize=(16, fig_h)) + + for m in range(n_states): + edge = edge_states[m] + E_hat = np.asarray(edge["E_hat"]).astype(bool) + E_t = np.asarray(edge["E_true_offdiag"]).astype(bool) + conf_map = np.asarray(edge["conf_map"]) + + tp = int(edge["tp"]) + fp = int(edge["fp"]) + fn = int(edge["fn"]) + precision = float(edge["precision"]) + recall = float(edge["recall"]) + f1 = float(edge["f1"]) + acc = float(edge["acc"]) + + ax = self.plt.subplot(n_states, 4, m * 4 + 1) + im = ax.imshow(E_t, cmap='Greys', vmin=0, vmax=1, origin='lower') + ax.set(title=f"E_true, state {m} (n={int(E_t.sum())})") + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + ax.plot([0, N], [0, N], '--', lw=0.8, color='magenta') + ax.set_aspect(1.0) + ax.grid(True, alpha=0.4) + ax.set_xlim(-0.5, N + 0.5) + ax.set_ylim(-0.5, N + 0.5) + ax.set(xlabel='presyn. neuron index (output)', ylabel='postsyn. neuron index (input)') + + ax = self.plt.subplot(n_states, 4, m * 4 + 2) + im = ax.imshow(E_hat, cmap='Greys', vmin=0, vmax=1, origin='lower') + ax.set(title=f"A_hat edges, state {m} (n={int(E_hat.sum())})") + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + ax.plot([0, N], [0, N], '--', lw=0.8, color='magenta') + ax.set_aspect(1.0) + ax.grid(True, alpha=0.4) + ax.set_xlim(-0.5, N + 0.5) + ax.set_ylim(-0.5, N + 0.5) + ax.set(xlabel='presyn. neuron index (output)', ylabel='postsyn. neuron index (input)') + + cmap_conf = colors.ListedColormap(['white', 'magenta', 'red', 'green']) + norm_conf = colors.BoundaryNorm([-0.5, 0.5, 1.5, 2.5, 3.5], cmap_conf.N) + ax = self.plt.subplot(n_states, 4, m * 4 + 3) + im = ax.imshow(conf_map, cmap=cmap_conf, norm=norm_conf, origin='lower') + ax.set(title='Confusion map') + cbar = self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + cbar.set_ticks([0, 1, 2, 3]) + cbar.set_ticklabels(['TN', 'FN', 'FP', 'TP']) + ax.plot([0, N], [0, N], '--', lw=0.8, color='k') + ax.set_aspect(1.0) + ax.grid(True, alpha=0.4) + ax.set_xlim(-0.5, N + 0.5) + ax.set_ylim(-0.5, N + 0.5) + ax.set(xlabel='presyn. neuron index', ylabel='postsyn. neuron index') + + ax = self.plt.subplot(n_states, 4, m * 4 + 4) + vals = [tp, fp, fn] + ax.bar(['TP', 'FP', 'FN'], vals, color=['green', 'red', 'magenta']) + ax.set_ylabel('count') + y_pos = 0.25 * max(1, max(vals)) + for name, val in zip(['TP', 'FP', 'FN'], vals): + ax.text(name, y_pos, f"{val}", ha='center', va='center', fontsize=10) + txt = f'precision={precision:.3f}\nrecall={recall:.3f}\nf1={f1:.3f}\nacc={acc:.3f}' + ax.text(0.6, 0.60, txt, transform=ax.transAxes) + ax.set_title('stats') + ax.grid(axis='y', alpha=0.4) + + fig.subplots_adjust(bottom=0.12) + fig.suptitle( + f"Edge detection vs truth, minW={minW} (off diagonal only): {md['short_name']}", + fontsize=14, + ) + + +def add_x45_lins(ax, only45=False): + lims = [ + np.min([ax.get_xlim(), ax.get_ylim()]), + np.max([ax.get_xlim(), ax.get_ylim()]), + ] + ax.plot(lims, lims, '--', color='k', linewidth=0.8) + if only45: + return + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + + +def add_delay_markers(ax, train_md): + delay = train_md.get("L1_prune_epoch", None) + if delay is None: + return + delay = int(delay) + if delay < 0: + return + # Training applies pruning at 0-based epoch>=delay; displayed epoch is delay+1. + ax.axvline(delay + 1, color='k', linestyle='--', linewidth=1.0) + + +def corrcoef_safe(x, y): + x = np.asarray(x) + y = np.asarray(y) + if x.size == 0 or y.size == 0: + return np.nan + if np.std(x) == 0 or np.std(y) == 0: + return np.nan + return float(np.corrcoef(x, y)[0, 1]) + + +def plot_corr_A_regions(ax, xV, yV, minW, title, xlab, ylab, s=6, alpha=0.5, color=None): + x = np.asarray(xV) + y = np.asarray(yV) + ax.scatter(x, y, s=s, alpha=alpha, color=color) + add_x45_lins(ax, only45=True) + ax.axvline(-minW, color="red", linestyle="--", linewidth=1) + ax.axvline(minW, color="red", linestyle="--", linewidth=1) + ax.set(title=title, xlabel=xlab, ylabel=ylab) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.4) + + mask_left = x < -minW + mask_mid = (x >= -minW) & (x <= minW) + mask_right = x > minW + r_left = corrcoef_safe(x[mask_left], y[mask_left]) + r_mid = corrcoef_safe(x[mask_mid], y[mask_mid]) + r_right = corrcoef_safe(x[mask_right], y[mask_right]) + n_left = int(mask_left.sum()) + n_mid = int(mask_mid.sum()) + n_right = int(mask_right.sum()) + + posx = [0.2, 0.50, 0.65] + posy = [0.05, 0.40, 0.60] + ax.text(posx[0], posy[0], f"rL={r_left:.3f}\nnL={n_left}", transform=ax.transAxes) + ax.text(posx[1], posy[1], f"rM={r_mid:.3f}\nnM={n_mid}", transform=ax.transAxes) + ax.text(posx[2], posy[2], f"rR={r_right:.3f}\nnR={n_right}", transform=ax.transAxes) + + +def find_bimodal_divider(xV): + """Return divider between two 1D modes using strict 2-cluster k-means.""" + x = np.asarray(xV, dtype=float).ravel() + assert x.size >= 2, f"Need >=2 samples for bimodal split, got {x.size}" + assert np.isfinite(x).all(), "xV contains non-finite values" + assert np.std(x) > 0.0, "xV must have non-zero variance for bimodal split" + + q1, q3 = np.quantile(x, [0.25, 0.75]) + c1, c2 = float(q1), float(q3) + if c1 == c2: + c1 = float(np.min(x)) + c2 = float(np.max(x)) + assert c1 != c2, "Failed to initialize two distinct mode centers" + + for _ in range(32): + d1 = np.abs(x - c1) + d2 = np.abs(x - c2) + left = d1 <= d2 + n_left = int(left.sum()) + n_right = int((~left).sum()) + assert n_left > 0 and n_right > 0, "Bimodal split produced empty cluster" + c1_new = float(np.mean(x[left])) + c2_new = float(np.mean(x[~left])) + if abs(c1_new - c1) < 1e-10 and abs(c2_new - c2) < 1e-10: + c1, c2 = c1_new, c2_new + break + c1, c2 = c1_new, c2_new + + if c1 > c2: + c1, c2 = c2, c1 + divide = 0.5 * (c1 + c2) + assert np.isfinite(divide), "Computed non-finite bimodal divider" + return float(divide) + + +def plot_corr_B_divisor(ax, xV, yV, title, xlab, ylab, s=6, alpha=0.5, color=None): + x = np.asarray(xV) + y = np.asarray(yV) + assert x.shape == y.shape, f"x and y shape mismatch: {x.shape} vs {y.shape}" + divide = find_bimodal_divider(x) + + ax.scatter(x, y, s=s, alpha=alpha, color=color) + add_x45_lins(ax, only45=True) + ax.axvline(divide, color="red", linestyle="--", linewidth=1) + ax.set(title=title, xlabel=xlab, ylabel=ylab) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.4) + + mask_left = x < divide + mask_right = x >= divide + r_left = corrcoef_safe(x[mask_left], y[mask_left]) + r_right = corrcoef_safe(x[mask_right], y[mask_right]) + n_left = int(mask_left.sum()) + n_right = int(mask_right.sum()) + ax.text(0.05, 0.45, f"rL={r_left:.3f}\nnL={n_left}", transform=ax.transAxes, va="top") + ax.text(0.65, 0.55, f"rR={r_right:.3f}\nnR={n_right}", transform=ax.transAxes, va="top") diff --git a/causal_net/nonStation_ver3a_states_simu/PlotterSpikesTrain.py b/causal_net/nonStation_ver3a_states_simu/PlotterSpikesTrain.py new file mode 100644 index 00000000..a2ba787b --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/PlotterSpikesTrain.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for biological experiment data visualization. + +This module provides specialized plotting capabilities for experimental +neural data analysis. The Plotter class extends PlotterBackbone to create +visualizations tailored for biological neural recordings including: +- Time series plots of neural activity patterns +- Statistical summaries and distribution analysis +- Data quality assessment plots +- Comparative analysis between experimental conditions + +Designed specifically for processing and visualizing data from biological +neural experiments, with automatic adaptation to experimental metadata. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +import numpy as np +import matplotlib.gridspec as gridspec +from matplotlib.ticker import MaxNLocator + +#...!...!.................... +def summary_column(md): + return 'fix-me 32678' + #print(sorted(md)) + rs=md['rate_summary'] + ds=md['data_selector'] + #pprint(md) + txt='dataset: '+md['short_name'] + txt += '\nnum acc neurons: %d' % md['num_neurons'] + txt += '\ndrop neur: %d <%.1f Hz, %d >%.1f Hz' % (ds['drop_neur_by_freq_range'][0],ds['freq_range'][0],ds['drop_neur_by_freq_range'][1],ds['freq_range'][1]) + txt+='\ndata type: %s\ntime_step=%.2f sec'%(md['data_type'],md['time_step_sec']) + txt += '\nMedian rate: %.1f Hz' % rs['median_spike_rate'] + txt += '\nrate range [%.1f, %.1f Hz]' % (rs['min_spike_rate'], rs['max_spike_rate']) + txt += '\nAvg Rate: %.1f±%.1f Hz' % (rs['avg_spike_rate'], rs['std_spike_rate']) + txt += '\nAvg Fano: %.2f±%.2f' % (rs['avg_fano_factor'], rs['std_fano_factor']) + + return txt + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def freq_histo(self,spikeD,md,figId=1): + + tit='dataset '+md['short_name'] + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,4)) + + # Create gridspec with 7:3 width ratio + gs = gridspec.GridSpec(1, 2, width_ratios=[7, 3]) + + dataYield, dataRates = spikeD['spikes'], spikeD['single_rates'] + + # Compute median + median_val = np.median(dataRates) + txtM= f'median : {median_val:.2f} Hz' + + #.... freq per channel + ax = self.plt.subplot(gs[0, 0]) + chanV=np.arange(dataRates.shape[0]) + ax.bar(chanV, dataRates , width=0.8, color='g', align='center', alpha=0.7) + + ax.axhline(median_val, color='red', linestyle='--', linewidth=1) + if md['data_type']=='simDale': + ax.set_yscale('log') + ax.grid() + ax.set(xlabel='input neuron index', ylabel='avr frequency (Hz)',title=tit) + ax.text( ax.get_xlim()[1]*0.1,median_val,txtM, + color='red', ha='center', va='bottom', fontsize=10) + + txt=summary_column(md) + ax.text(0.05, 0.3, txt, transform=ax.transAxes, fontsize=12,color='b') + + #........ freq histo ..... + ax = self.plt.subplot(gs[0, 1]) + ax.hist(dataRates,bins=20) + #ax.set_yscale('log') + ax.set(ylabel='num channels', xlabel='avr frequency (Hz)',title=tit) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + # Add median text annotation + ax.text(median_val+1, ax.get_ylim()[1]*0.5,txtM, + color='red', ha='left', va='bottom', fontsize=10,rotation=90) + ax.text(0.25, 0.3, txt, transform=ax.transAxes, fontsize=11,color='m') + + ax.grid() + +#...!...!.................. + def freq_vs_time(self,rebD,md,figId=2,S_true=None,S_oracle=None): + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,11)) + if S_oracle is None and isinstance(rebD, dict) and 'S_oracle' in rebD: + S_oracle = rebD['S_oracle'] + + tit='dataset '+md['short_name'] + time_step=rebD['time_step2'] + R_sel = md['sel_spect_radius'] + state_tag = md['sel_state'] + + # clip time data for display + tL,tR=md['plot']['time_rangeLR'] + itL,itR=(md['plot']['time_rangeLR']/time_step).astype(int) + ntime_tot = rebD['rate2D'].shape[0] + itL = max(0, min(itL, ntime_tot - 1)) + itR = max(itL + 1, min(itR, ntime_tot)) + print('iTL,R', itL,itR) + + # unpack data and clip the time range + rate2D=rebD['rate2D'][itL:itR] + timeV=rebD['timeV'][itL:itR] + + _,nchan=rate2D.shape + Tbin = time_step + xL = float(timeV[0]) + xR = float(timeV[-1] + Tbin) + if isinstance(rebD, dict) and 'pop_rate_hz' in rebD: + pop_rate_hz = np.asarray(rebD['pop_rate_hz'][itL:itR], dtype=np.float64) + else: + pop_rate_hz = np.sum(rate2D, axis=1) + medRateDisp = float(np.median(rate2D)) + cntAboveMed = np.sum(rate2D > medRateDisp, axis=1) + + tit0='dataset: %s state=%d R=%.3f nchan=%d Tbin=%.2f sec'%(md['short_name'], state_tag, R_sel, nchan, time_step) + + # Layout: top-count trace, heatmap, synchronicity trace, target-state, oracle-state. + gs = fig.add_gridspec(5, 1, height_ratios=[0.14, 0.52, 0.14, 0.10, 0.10], hspace=0.36) + + # ..... top plot + ax = fig.add_subplot(gs[0, 0]) + ax.bar(timeV+Tbin*.5, cntAboveMed, width=Tbin, color='teal', align='center', alpha=0.7) + ax.set(ylabel='neurons > median', title=f'neurons above displayed-time median rate ({medRateDisp:.2f} Hz)') + ax.set_xlim(xL,xR) + ax.tick_params(axis='x', labelbottom=False) + ax.grid() + + # ....... main heatmap + ax = fig.add_subplot(gs[1, 0]) + #cmap='tab20c', 'Oranges' + # - - - Get the colormap and modify it --- + original_cmap = self.plt.get_cmap('tab20c') + custom_cmap = original_cmap.copy() + custom_cmap.set_under('white') + + cax = ax.imshow(rate2D.T, + aspect='auto', + cmap=custom_cmap, # Use our modified colormap + origin='lower', + extent=[xL, xR, 0, nchan], + interpolation='nearest', + vmin=1, # Set the lower bound of the colormap + vmax=rate2D.max()) # Optional: ensure the upper bound is set + + ax.set_ylabel('Neuron index') + ax.set_title(tit0 ) + ax.tick_params(axis='x', labelbottom=False) + cbar = fig.colorbar(cax, ax=ax, orientation='horizontal', pad=0.02, label=f'Instantanous freq (Hz) per neuron over Tbin={time_step:.1f} sec', shrink=0.55) + cbar.ax.tick_params(labelsize=10) + ax.grid() + + # ..... synchronicity (yellow) below heatmap + ax = fig.add_subplot(gs[2, 0]) + ax.bar(timeV+Tbin*.5, pop_rate_hz, width=Tbin, color='gold', align='center', alpha=0.8) + ax.set(ylabel='synchronicity (Hz)', title=f'Network synchronicity from {nchan} neurons, Tbin={Tbin:.2f} sec') + ax.set_xlim(xL,xR) + ax.tick_params(axis='x', labelbottom=True) + ax.grid() + + # ..... target state + ax = fig.add_subplot(gs[3, 0]) + ax.set_xlim(xL, xR) + if S_true is not None: + assert S_true.ndim == 1, f"S_true must be 1D, got shape={S_true.shape}" + dt0 = float(md['time_step_sec']) + t_state = np.arange(S_true.shape[0], dtype=float) * dt0 + sel = (t_state >= xL) & (t_state < xR + dt0) + if np.any(sel): + t_sel = t_state[sel] + s_sel = S_true[sel] + ax.step(t_sel, s_sel, where='post', color='k', linewidth=1.0) + smin = int(np.min(s_sel)) + smax = int(np.max(s_sel)) + ax.set_ylim(smin - 0.5, smax + 0.5) + ax.set_yticks(np.arange(smin, smax + 1, 1)) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.set_ylabel('state S') + ax.set_title(f"target state — {md['short_name']}") + ax.grid() + else: + ax.set_axis_off() + ax.tick_params(axis='x', labelbottom=False) + + # ..... oracle state + ax = fig.add_subplot(gs[4, 0]) + ax.set_xlim(xL, xR) + if S_oracle is not None: + assert S_oracle.ndim == 1, f"S_oracle must be 1D, got shape={S_oracle.shape}" + score_txt = f", avr_score={md['oracle_score']:.3f}" + dt0 = float(md['time_step_sec']) + t_state = np.arange(S_oracle.shape[0], dtype=float) * dt0 + sel = (t_state >= xL) & (t_state < xR + dt0) + if np.any(sel): + t_sel = t_state[sel] + s_sel = S_oracle[sel] + ax.step(t_sel, s_sel, where='post', color='g', linewidth=1.0) + smin = int(np.min(s_sel)) + smax = int(np.max(s_sel)) + ax.set_ylim(smin - 0.5, smax + 0.5) + ax.set_yticks(np.arange(smin, smax + 1, 1)) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.set_ylabel('oracle S') + ax.set_title(f"oracle state{score_txt}") + ax.grid() + else: + ax.set_axis_off() + ax.tick_params(axis='x', labelbottom=True) + ax.set_xlabel('Time (s)') diff --git a/causal_net/nonStation_ver3a_states_simu/PoissonGLModel.py b/causal_net/nonStation_ver3a_states_simu/PoissonGLModel.py new file mode 100644 index 00000000..ce8b7377 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/PoissonGLModel.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +""" +PyTorch implementation of Poisson Generalized Linear Model for neural spike data. + +This module defines the PoissonGLModel class, a neural network that models +multivariate Poisson spike counts using the GLM framework: + rate_t = exp(A @ Y_prev + B) * dt + +The model supports two initialization modes: +- Stage A: Standard parameterization with full A and B matrices +- Stage B: Efficient parameterization using only trainable weights with masks + +Key features: +- Recurrent connectivity matrix A (neuron-to-neuron weights) +- Bias term B (baseline firing rates) +- Forward pass computing Poisson rates from previous spike counts +- Custom negative log-likelihood loss function for Poisson distributions +- Support for sparse parameterization and noise injection +""" + +# PoissonGML refers to a Poisson Generalized Linear Model (GLM) +import torch +import torch.nn as nn + + +class PoissonGLModel(nn.Module): + """ + Generalized Linear Model for multivariate Poisson spike counts. + rate_t = exp(A @ Y_prev + B) * dt + + Constructor: + PoissonGLM(n_neurons, eta_clip) + - Stage A: random initialization of A and B + PoissonGLM(n_neurons, eta_clip, A_init, B_init, trainable_mask) + - Stage B: efficient parameterization with only trainable weights in a single tensor + """ + def __init__(self, n_neurons, eta_clip, A_init=None, B_init=None, trainable_mask=None, noise_scale=0.0): + super().__init__() + self.n_neurons = n_neurons + self.eta_clip = float(eta_clip) + + if A_init is None: # Stage A: standard parameterization + self.A = nn.Parameter(torch.randn(n_neurons, n_neurons) * 0.1) + self.B = nn.Parameter(torch.randn(n_neurons) * 0.1) + self.is_stage_b = False + else: # Stage B: efficient parameterization + self.is_stage_b = True + + # Store fixed components as buffers + A_t = torch.tensor(A_init, dtype=torch.float32) + B_t = torch.tensor(B_init, dtype=torch.float32) + mask = torch.tensor(trainable_mask, dtype=torch.bool) + + self.register_buffer('A_fixed', A_t) # Full A matrix for initialization + self.register_buffer('B_fixed', B_t) # Fixed B values (will be overridden) + self.register_buffer('trainable_mask', mask) + + # Find trainable positions + trainable_A_indices = torch.nonzero(mask, as_tuple=True) + self.register_buffer('trainable_A_row_idx', trainable_A_indices[0]) + self.register_buffer('trainable_A_col_idx', trainable_A_indices[1]) + + # Count trainable parameters + n_trainable_A = mask.sum().item() + n_trainable_B = n_neurons # All B elements are trainable + total_trainable = n_trainable_A + n_trainable_B + + # Single parameter tensor for all trainable weights + self.C = nn.Parameter(torch.zeros(total_trainable)) + + # Initialize C with current values + optional noise + with torch.no_grad(): + # First part: trainable A elements + trainable_A_values = A_t[trainable_A_indices] + self.C[:n_trainable_A] = trainable_A_values + + # Second part: B elements + self.C[n_trainable_A:] = B_t + # Add random noise if requested + if noise_scale > 0: + noise = torch.randn_like(self.C) * noise_scale + self.C += noise + #print(f"Added random noise with scale {noise_scale} to Stage B initialization") + + self.n_trainable_A = n_trainable_A + + @property + def A(self): + if not self.is_stage_b: + return self._parameters['A'] + else: + # Reconstruct A matrix from C and fixed values + A_reconstructed = self.A_fixed.clone() + A_reconstructed[self.trainable_A_row_idx, self.trainable_A_col_idx] = self.C[:self.n_trainable_A] + return A_reconstructed + + @property + def B(self): + if not self.is_stage_b: + return self._parameters['B'] + else: + # B elements are stored in the second part of C + return self.C[self.n_trainable_A:] + + def forward(self, Y_prev, dt=0.01): + linear = torch.addmm(self.B.unsqueeze(0), Y_prev, self.A.t()) + linear = torch.clamp(linear, max=self.eta_clip) + return torch.exp(linear) * dt + + +def poisson_nll_loss(spikes, targets, firing_rates): + eps = 1e-8 + firing_rates_safe = torch.maximum(firing_rates, torch.tensor(0.1, device=firing_rates.device)) + weights = 1.0 / firing_rates_safe + weights /= torch.mean(weights) + weights = weights.unsqueeze(0) + loss = -weights * targets * torch.log(spikes + eps) + weights * spikes + return loss.mean() diff --git a/causal_net/nonStation_ver3a_states_simu/Readme b/causal_net/nonStation_ver3a_states_simu/Readme new file mode 100644 index 00000000..3740c53d --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/Readme @@ -0,0 +1,20 @@ + salloc -q shared_interactive -C gpu -t 4:00:00 -N 1 -A m2043 + +python3 - << 'PY' +import torch +print("torch:", torch.__version__) +print("cuda available:", torch.cuda.is_available()) +print("torch cuda version:", torch.version.cuda) +if torch.cuda.is_available(): + print("gpu:", torch.cuda.get_device_name(0)) +PY + + += = = = Parallel execution = = = = = +If you want, next I can add a toggle to choose serial vs segmented E-step (--estep_sharded) for A/B testing accuracy/speed. + + +time OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 time torchrun --standalone --nnodes=1 --nproc_per_node=4 \ + nonStationary_ver3/prism_EM_train.py \ + --num_states 3 --dataName --num_em_iters 4 --m_epochs 2 --time_range_sec 0 10 --prescale_m_step_4_ArhoMax 50 --pgd_iter 5 + diff --git a/causal_net/nonStation_ver3a_states_simu/UtilDalePoisson.py b/causal_net/nonStation_ver3a_states_simu/UtilDalePoisson.py new file mode 100644 index 00000000..de4ab1a0 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/UtilDalePoisson.py @@ -0,0 +1,240 @@ +""" +Utility functions for Dale Poisson simulation data processing + +""" + +import numpy as np +import os +import time +from pprint import pprint + +def compute_consecutive_coincidence_rate(Y,time_evol): + """ + Compute the frequency of coincidences for 2 consecutive time bins for 2 different channels. + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + + Returns: + float: conc_rate_all - the overall coincidence rate + """ + num_steps, num_neurons = Y.shape + + # We need at least 2 time steps for consecutive bins + if num_steps < 2: + return 0.0 + + # Get consecutive time slices using vectorized operations + Y_t = Y[:-1, :] # Y[0:T-1, :] - current time bins + Y_t_plus_1 = Y[1:, :] # Y[1:T, :] - next time bins + + # Create boolean masks for non-zero values + mask_t = (Y_t != 0) # Shape: (T-1, N) + mask_t_plus_1 = (Y_t_plus_1 != 0) # Shape: (T-1, N) + + # Use broadcasting to compute all channel pairs at once + # mask_t[:, :, None] has shape (T-1, N, 1) + # mask_t_plus_1[:, None, :] has shape (T-1, 1, N) + # Broadcasting gives shape (T-1, N, N) for all pairs + coincidences = mask_t[:, :, None] & mask_t_plus_1[:, None, :] + + # Remove diagonal (same channel pairs) using boolean indexing + diagonal_mask = np.eye(num_neurons, dtype=bool) + coincidences[:, diagonal_mask] = False + + # Count total coincidences across all time steps + coincidence_count = np.sum(coincidences) + + conc_rate = coincidence_count / time_evol/num_neurons + return conc_rate # Hz, per neuron + +def estimate_rates(Y, dt, num_excite, max_samples, varTwindow=5, mxNn=5, verb=1, spect_radius=None): + """ + Evaluates spike statistics and estimates firing rates and Fano factors. + + Computes statistics over non-overlapping time windows of length varTwindow: + - Fano Factor: Var[spike count] / Mean[spike count] per neuron (dimensionless) + - Rate variance: Var[rate] per neuron (Hz^2) + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + dt (float): Time bin size in seconds. + num_excite (int): Number of excitatory neurons. + max_samples (int): The maximum number of time samples to use for calculation. + varTwindow (float, optional): Time window length in seconds for variance computation. Defaults to 5. + mxNn (int, optional): Max number of neurons to show in detailed stats. Defaults to 5. + + Returns: + tuple: A tuple containing (stats_dict, rates_dict, neur_freq_index): + - stats_dict (dict): Contains population-level spike statistics. + - rates_dict (dict): Contains per-neuron firing rates, Fano factors, and rate variance. + - neur_freq_index (np.ndarray): Neuron indices sorted by firing rate (low to high). + """ + # 1. Clip data to max_samples_for_rates + if verb > 0: + print("\n=== Estimating Rates & Stats ===") + if Y.shape[0] > max_samples: + if verb > 0: + print("Using %d samples (out of %d) for rate computation" % (max_samples, Y.shape[0])) + Y = Y[:max_samples] + + # Part 1: Compute all basic statistics once + num_steps_sim, Nn_sim = Y.shape + num_inhib = Nn_sim - num_excite + time_evol = num_steps_sim * dt + if verb > 0: + print('steps num_steps=%d, time_evol=%.1f sec, Nn=%d (%d Excit, %d Inhib)' % (num_steps_sim, time_evol, Nn_sim, num_excite, num_inhib)) + + # Compute windowed spike counts and statistics over non-overlapping time windows of length varTwindow (sec), per neuron + window_size = max(1, int(round(varTwindow / dt))) + num_windows = num_steps_sim // window_size + if verb > 0: + print('variance computation: num_windows=%d, window_size=%d' % (num_windows, window_size)) + if num_windows <= 1: + raise ValueError("varTwindow=%s is too large for data (num_windows=%d). Need at least 2 windows for variance." % (varTwindow, num_windows)) + Y_trim = Y[:num_windows * window_size] + Y_win = Y_trim.reshape(num_windows, window_size, Nn_sim) + + # Sum spike counts over each window: shape (num_windows, n_neurons) + window_spike_counts = np.sum(Y_win, axis=1) + + # Fano Factor: Var[spike count] / Mean[spike count] over windows + mean_window_spike_counts = np.mean(window_spike_counts, axis=0) + var_window_spike_counts = np.var(window_spike_counts, axis=0) + fano_factor = np.divide(var_window_spike_counts, mean_window_spike_counts, + out=np.zeros_like(var_window_spike_counts), + where=mean_window_spike_counts != 0) + + # Convert window spike counts to rates (Hz) for rate variance + window_rates = window_spike_counts / (window_size * dt) # shape: (num_windows, n_neurons), Hz + single_rates_var = np.var(window_rates, axis=0) # variance of rate across windows, per neuron (Hz^2) + + # Compute raw arrays (full time series) + spike_counts = np.sum(Y, axis=0) + spike_rates = spike_counts / time_evol + mean_counts_per_bin = np.mean(Y, axis=0) + + # Compute consecutive coincidence rate + start_time = time.time() + conc_rate_per_neuron = compute_consecutive_coincidence_rate(Y,time_evol) + elapsed_time = time.time() - start_time + if verb > 0: + print('Coincidence rate %.2g Hz, Y.shape=%s elaT %.3f sec' % (conc_rate_per_neuron, Y.shape, elapsed_time)) + + # Compute all population statistics once + med_rate_all = np.median(spike_rates) + avg_rate_all = float(np.mean(spike_rates)) + std_rate_all = float(np.std(spike_rates)) + avg_fano_all = float(np.mean(fano_factor)) + std_fano_all = float(np.std(fano_factor)) + avg_rate_excit = float(np.mean(spike_rates[:num_excite])) + std_rate_excit = float(np.std(spike_rates[:num_excite])) + avg_fano_excit = float(np.mean(fano_factor[:num_excite])) + std_fano_excit = float(np.std(fano_factor[:num_excite])) + avg_rate_inhib = float(np.mean(spike_rates[num_excite:])) + std_rate_inhib = float(np.std(spike_rates[num_excite:])) + avg_fano_inhib = float(np.mean(fano_factor[num_excite:])) + std_fano_inhib = float(np.std(fano_factor[num_excite:])) + + # Build stats dictionary with computed values + stats_dict = { + 'num_steps': num_steps_sim, + 'time_evol_sec': time_evol, + 'time_step_sec': dt, + 'num_neurons': Nn_sim, + 'num_excitatory': num_excite, + 'num_inhibitory': num_inhib, + 'avg_spike_rate_all': avg_rate_all, + 'std_spike_rate_all': std_rate_all, + 'avg_fano_factor_all': avg_fano_all, + 'std_fano_factor_all': std_fano_all, + 'avg_spike_rate_excit': avg_rate_excit, + 'std_spike_rate_excit': std_rate_excit, + 'avg_fano_factor_excit': avg_fano_excit, + 'std_fano_factor_excit': std_fano_excit, + 'avg_spike_rate_inhib': avg_rate_inhib, + 'std_spike_rate_inhib': std_rate_inhib, + 'avg_fano_factor_inhib': avg_fano_inhib, + 'std_fano_factor_inhib': std_fano_inhib, + 'conc_rate_per_neuron': float(conc_rate_per_neuron), + 'median_spike_rate_all': med_rate_all + } + + # Printing detailed stats for individual neurons + mxE = min(mxNn, num_excite) + mxI = min(mxNn, num_inhib) + + if verb > 0: + print('\n--- Stats for first %d Excitatory Neurons ---' % mxE) + np.set_printoptions(precision=2) + print('Total Spike Counts: %s' % spike_counts[:mxE]) + print('Mean Firing Rate (Hz): %s' % spike_rates[:mxE]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[:mxE])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[:mxE])) + print('Fano Factor (Var/Mean): %s' % fano_factor[:mxE]) + + print('\n--- Stats for first %d Inhibitory Neurons ---' % mxI) + np.set_printoptions(precision=2) + inhib_slice = slice(num_excite, num_excite + mxI) + print('Total Spike Counts: %s' % spike_counts[inhib_slice]) + print('Mean Firing Rate (Hz): %s' % spike_rates[inhib_slice]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[inhib_slice])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[inhib_slice])) + print('Fano Factor (Var/Mean): %s' % fano_factor[inhib_slice]) + + # Print population summary using dictionary values (always printed) + R_tag = ', R=%.3f' % spect_radius if spect_radius is not None else '' + print('\n--- Population Summary Statistics%s ---' % R_tag) + print('All (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (Nn_sim, stats_dict['avg_spike_rate_all'], stats_dict['std_spike_rate_all'], stats_dict['avg_fano_factor_all'], stats_dict['std_fano_factor_all'])) + print('Excit (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_excite, stats_dict['avg_spike_rate_excit'], stats_dict['std_spike_rate_excit'], stats_dict['avg_fano_factor_excit'], stats_dict['std_fano_factor_excit'])) + if num_inhib > 0: + print('Inhib (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_inhib, stats_dict['avg_spike_rate_inhib'], stats_dict['std_spike_rate_inhib'], stats_dict['avg_fano_factor_inhib'], stats_dict['std_fano_factor_inhib'])) + + # Print key summary values using dictionary + summary_keys = ['conc_rate_per_neuron', 'median_spike_rate_all'] + for key in summary_keys: + if key == 'conc_rate_per_neuron': + print('Coincidence rate per neuron %.2g Hz' % stats_dict[key]) + elif key == 'median_spike_rate_all': + print('Median rate %.2f Hz\n' % stats_dict[key]) + + # Part 2: Compute frequency sorting + neur_freq_index = np.argsort(spike_rates) + + rates_dict = { + 'single_rates': spike_rates, + 'neur_freqIdx':neur_freq_index, + 'sigle_rates_var': single_rates_var, + 'single_fano_fact': fano_factor + } + + return stats_dict, rates_dict, neur_freq_index + +def get_offdiag_triplets(A, isPos=True): + """ + Returns positive/negative, off-diagonal elements of 2D array A as array of [i, j, value] triplets. + + Parameters: + A : 2D numpy array + isPos : bool, if True only return positive values (>0), if False return all negative values + + Returns: + Array of shape (n, 3) where n is the number of valid off-diagonal elements + Each row is [row_index, col_index, value] + """ + # Create mask for off-diagonal elements + offdiag_mask = ~np.eye(A.shape[0], A.shape[1], dtype=bool) + + # Add value condition based on isPos + if isPos: + value_mask = A > 0 + else: + value_mask = A < 0 + + # Combine masks + mask = value_mask & offdiag_mask + + i_indices, j_indices = np.where(mask) + values = A[i_indices, j_indices] + + return np.column_stack([i_indices, j_indices, values]) diff --git a/causal_net/nonStation_ver3a_states_simu/UtilTorch.py b/causal_net/nonStation_ver3a_states_simu/UtilTorch.py new file mode 100644 index 00000000..dc1dc13b --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/UtilTorch.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +""" +PyTorch utilities for GPU setup, data preprocessing and model training. + +This module provides essential utilities for PyTorch-based neural network training, +specifically optimized for Poisson GLM fitting with distributed GPU support. +Main functionality includes: +- GPU availability checking and device configuration +- Data preprocessing with time decorrelation and shuffling options +- Custom dataset classes for paired neural data (Y_prev, Y_curr) +- Distributed training utilities with Poisson loss optimization +- Learning rate scheduling and training loop management + +Designed to work with multi-GPU setups using DistributedDataParallel +for efficient training of large-scale neural connectivity models. +""" + +import torch +import numpy as np +import time +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader +from PoissonGLModel import poisson_nll_loss + +def offdiag_soft_threshold_(A, lr, lam): + """In-place off-diagonal soft-thresholding for a single square A matrix.""" + if lam <= 0.0: + return + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError(f"offdiag_soft_threshold_ expects square 2D A, got shape={tuple(A.shape)}") + with torch.no_grad(): + N = A.shape[0] + off_diag = ~torch.eye(N, dtype=torch.bool, device=A.device) + t = lr * lam + A_off = A[off_diag] + A[off_diag] = A_off.sign() * (A_off.abs() - t).clamp(min=0.0) + +def enforce_spectral_radius_(A, rho_max, eps=1e-12): + """Scale A in-place only if spectral radius exceeds rho_max (GPU-friendly).""" + if rho_max is None: + return + if rho_max <= 0.0: + raise ValueError(f"rho_max must be > 0, got {rho_max}") + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError(f"enforce_spectral_radius_ expects square 2D A, got shape={tuple(A.shape)}") + with torch.no_grad(): + # Keep computation on current device; no CPU transfer. + rho = torch.linalg.eigvals(A).abs().max() + if torch.isfinite(rho) and (rho > rho_max): + A.mul_(float(rho_max) / float(rho + eps)) + +def check_gpu_availability(): + """Check GPU availability and set up device configuration.""" + if not torch.cuda.is_available(): + raise RuntimeError("This script requires a CUDA-enabled GPU environment.") + device = torch.device('cuda:0') + torch.cuda.set_device(0) + print(f"Using device: {torch.cuda.get_device_name(0)}") + return device + +def preprocess_data(Y, args): + import random + Nt, Nn = Y.shape + is_main = (getattr(args, "rank", 0) == 0) + + # Apply time decorrelation if requested + if args.desyncTime: + if is_main: print("\n=== Applying Time Decorrelation, it shifts time for each neuron ===") + seed = int(time.time() * 1000) % 1000000 + np.random.seed(seed) + shift_amounts = np.random.randint(1, Nt//4, size=Nn, dtype=np.int32) + if is_main: print('Generated random shifts with seed=%d'%(seed),shift_amounts[:10],'...',flush=True) + Y_shifted = np.zeros_like(Y) + for neuron_idx in range(Nn): + shift_amount = int(shift_amounts[neuron_idx]) + Y_shifted[:, neuron_idx] = np.roll(Y[:, neuron_idx], shift_amount) + if is_main: print(f"Applied time shifts, destroys temporal correlations between neurons") + Y = Y_shifted + + # Create consecutive pairs + max_pairs = Nt - 1 + num_samples = args.num_samples + if num_samples is None or num_samples > max_pairs: + num_samples = max_pairs + XY = np.stack([Y[:num_samples], Y[1:num_samples + 1]], axis=1) + + # Apply random data dropping if requested + if args.dropDataFrac > 0: + n_pairs = XY.shape[0] + drop_seed = (int(time.time() * 1000) + np.random.randint(0, 1000)) % (2**32) + np.random.seed(drop_seed) + random.seed(drop_seed) + n_keep = int(n_pairs * (1.0 - args.dropDataFrac)) + keep_indices = np.random.choice(n_pairs, size=n_keep, replace=False) + keep_indices = np.sort(keep_indices) + XY = XY[keep_indices] + if is_main: print(f"Dropped {args.dropDataFrac:.1%} of data, keeping {XY.shape[0]} samples (seed={drop_seed})") + + return XY + + +def train_Poisson_model(model, device, train_loader, n_epochs, lr, L1_alpha=0.0, use_scheduler=False, + firing_rates=None, train_sampler=None, print_every=20, apply_prox=False, + lr_end_factor=0.03, minW=1e-6, rho_max=0.99, prescale_m_step_4_ArhoMax=10, L1_prune_epoch=0): + if L1_prune_epoch < 0: + raise ValueError(f"L1_prune_epoch must be >= 0, got {L1_prune_epoch}") + if prescale_m_step_4_ArhoMax < 1: + raise ValueError(f"prescale_m_step_4_ArhoMax must be >= 1, got {prescale_m_step_4_ArhoMax}") + use_fused = (isinstance(device, torch.device) and device.type=='cuda' and torch.cuda.is_available()) + assert use_fused + optimizer = optim.Adam(model.parameters(), lr=lr, fused=True) + mdl = model.module if hasattr(model, 'module') else model + scheduler = optim.lr_scheduler.LinearLR( + optimizer, start_factor=1.0, end_factor=lr_end_factor, total_iters=n_epochs + ) if use_scheduler else None + + diag_mask = torch.eye(mdl.n_neurons, device=device).bool() + off_diag_mask = ~diag_mask + L1_weight_matrix = torch.ones(mdl.n_neurons, mdl.n_neurons, device=device) + L1_weight_matrix[diag_mask] = 0.0 + + firing_rates_tensor = torch.tensor(firing_rates, dtype=torch.float32, device=device) if firing_rates is not None else None + + train_losses_w_L1, train_losses_wo_L1, learning_rates = [], [], [] + sparsity_epoch, nz_offdiag_epoch, spectral_radius_epoch = [], [], [] + train_epochs = [] + start_time = time.time() + + for epoch in range(n_epochs): + if train_sampler is not None: train_sampler.set_epoch(epoch) + model.train() + train_loss_w_L1 = 0 + train_loss_wo_L1 = 0 + apply_rho = True + apply_prune = (epoch >= L1_prune_epoch) + for batch_idx, (Y_prev, Y_curr) in enumerate(train_loader): + Y_prev, Y_curr = Y_prev.float().to(device, non_blocking=True), Y_curr.float().to(device, non_blocking=True) + optimizer.zero_grad(set_to_none=True) + spikes = model(Y_prev) + base_loss = poisson_nll_loss(spikes, Y_curr, firing_rates_tensor) + loss_with_L1 = base_loss.clone() + if L1_alpha > 0: + loss_with_L1 += L1_alpha * torch.mean(torch.abs(mdl.A) * L1_weight_matrix) + loss_with_L1.backward() + optimizer.step() + if apply_prune and apply_prox and L1_alpha > 0: + offdiag_soft_threshold_(mdl.A, optimizer.param_groups[0]['lr'], L1_alpha) + if apply_rho and (batch_idx % prescale_m_step_4_ArhoMax == 0): + enforce_spectral_radius_(mdl.A, rho_max) + train_loss_w_L1 += loss_with_L1.item() + train_loss_wo_L1 += base_loss.item() + + # record train losses each epoch + train_losses_w_L1.append(train_loss_w_L1 / len(train_loader)) + train_losses_wo_L1.append(train_loss_wo_L1 / len(train_loader)) + train_epochs.append(epoch + 1) + learning_rates.append(optimizer.param_groups[0]['lr']) + with torch.no_grad(): + A_off = mdl.A[off_diag_mask] + nz_off = int((A_off.abs() > minW).sum().item()) + n_off = int(off_diag_mask.sum().item()) + sparsity = 1.0 - nz_off / max(1, n_off) + rho = float(torch.linalg.eigvals(mdl.A).abs().max().item()) + sparsity_epoch.append(sparsity) + nz_offdiag_epoch.append(nz_off) + spectral_radius_epoch.append(rho) + + if scheduler: + scheduler.step() + + if (epoch + 1) % print_every == 0: + print(f"Epoch {epoch+1}/{n_epochs}: Loss_Tot={train_losses_w_L1[-1]:.5g}, only_L1={(train_losses_w_L1[-1]-train_losses_wo_L1[-1]):.4g}, A_sparsity={sparsity:.3f}, nz_offdiag={nz_off}, rho(A)={rho:.4f}, Elapsed={(time.time() - start_time):.1f}s") + + + + return train_losses_w_L1, train_losses_wo_L1, learning_rates, train_epochs, sparsity_epoch, nz_offdiag_epoch, spectral_radius_epoch + + +class NumpyPairDataset(Dataset): + def __init__(self, X_np, Y_np): + self.X = X_np + self.Y = Y_np + self.n = X_np.shape[0] + def __len__(self): + return self.n + def __getitem__(self, idx): + return torch.from_numpy(self.X[idx]).to(dtype=torch.float32), torch.from_numpy(self.Y[idx]).to(dtype=torch.float32) + +def make_loader(X, Yt, args, is_dist=False, shuffle=True): + if is_dist: + raise ValueError("Distributed loading was removed from UtilTorch.make_loader; use is_dist=False.") + dataset = NumpyPairDataset(X, Yt) + return DataLoader(dataset, batch_size=args.batch_size, shuffle=shuffle, drop_last=shuffle, pin_memory=True, pin_memory_device='cuda', num_workers=8, + persistent_workers=True, prefetch_factor=8) diff --git a/causal_net/nonStation_ver3a_states_simu/Util_PrismEM.py b/causal_net/nonStation_ver3a_states_simu/Util_PrismEM.py new file mode 100644 index 00000000..16b1fa08 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/Util_PrismEM.py @@ -0,0 +1,309 @@ +#!/usr/bin/env python3 +""" +Utilities for prism EM initialization. +""" + +import numpy as np +import time + + +def _normalize_init_opt(opt): + s = str(opt).strip().lower() + if s not in ("data", "rand"): + raise ValueError(f"Unsupported init option: {opt}") + return s + + +def _rate_thresholds_min_var(rate_t, n_state): + """Find up to (n_state-1) thresholds minimizing within-group variance.""" + x = np.asarray(rate_t, dtype=np.float64).ravel() + if x.size == 0 or n_state <= 1: + return [] + + uniq, counts = np.unique(x, return_counts=True) + u = uniq.size + if u <= 1: + return [] + + k = min(int(n_state), u) + w = counts.astype(np.float64) + wx = w * uniq + wx2 = wx * uniq + + c_w = np.concatenate(([0.0], np.cumsum(w))) + c_s = np.concatenate(([0.0], np.cumsum(wx))) + c_q = np.concatenate(([0.0], np.cumsum(wx2))) + + # dp[c, j]: minimal weighted SSE for first (j+1) unique values into c groups. + dp = np.full((k + 1, u), np.inf, dtype=np.float64) + ptr = np.full((k + 1, u), -1, dtype=np.int32) + + j_all = np.arange(u, dtype=np.int32) + 1 + w0 = c_w[j_all] - c_w[0] + s0 = c_s[j_all] - c_s[0] + q0 = c_q[j_all] - c_q[0] + dp[1, :] = q0 - (s0 * s0) / np.maximum(w0, 1e-12) + ptr[1, :] = 0 + + for c in range(2, k + 1): + for j in range(c - 1, u): + i = np.arange(c - 1, j + 1, dtype=np.int32) + jj = j + 1 + + ww = c_w[jj] - c_w[i] + ss = c_s[jj] - c_s[i] + qq = c_q[jj] - c_q[i] + seg = qq - (ss * ss) / np.maximum(ww, 1e-12) + cand = dp[c - 1, i - 1] + seg + + ib = int(np.argmin(cand)) + dp[c, j] = float(cand[ib]) + ptr[c, j] = int(i[ib]) + + starts = [] + j = u - 1 + for c in range(k, 1, -1): + i = int(ptr[c, j]) + starts.append(i) + j = i - 1 + starts.reverse() + + thres = [] + for i in starts: + l = uniq[i - 1] + r = uniq[i] + thres.append(float(0.5 * (l + r))) + return thres + + +def init_states_vs_time(spikes, dt, args): + """Prepare initial c_hat probabilities and state labels from spikes. + + Returns: + c_init: float32 array (T, M) + S_init: int64 array (T,) + init_meta: None or {"rate_thres": [..], "rate_bin_sec": float, "rate_bin_bins": int} + freq_h1d: float32 array used for state classification + """ + spikes = np.asarray(spikes) + if spikes.ndim != 2: + raise ValueError("spikes must be 2D array (T, N)") + + T_full = spikes.shape[0] + n_state = int(args.num_states) + opt = _normalize_init_opt(args.init_states) + + if opt == "rand": + c_init = np.full((T_full, n_state), 1.0 / n_state, dtype=np.float32) + s_init = np.full((T_full,), -1, dtype=np.int64) + freq_h1d = np.zeros((0,), dtype=np.float32) + return c_init, s_init, None, freq_h1d + if opt != "data": + raise ValueError(f"Unsupported --init_states option: {opt}") + + dwell_sec = float(args.decode_dwell_sec) + dwell_bins = max(1, int(np.ceil(dwell_sec / float(dt)))) + + # Aggregate spikes over coarse dwell bins, then compute mean firing rate. + T = T_full + n_blk = (T + dwell_bins - 1) // dwell_bins + rate_blk = np.zeros((n_blk,), dtype=np.float64) + blk_sizes = np.zeros((n_blk,), dtype=np.int64) + for b in range(n_blk): + i0 = b * dwell_bins + i1 = min(T, i0 + dwell_bins) + blk = spikes[i0:i1] + nb = i1 - i0 + blk_sizes[b] = nb + rate_per_neuron = blk.sum(axis=0).astype(np.float64) / (float(nb) * float(dt)) + rate_blk[b] = rate_per_neuron.mean() + + thres = _rate_thresholds_min_var(rate_blk, n_state) + th = np.asarray(thres, dtype=np.float64) + s_blk = np.searchsorted(th, rate_blk, side="right").astype(np.int64) + s_blk = np.clip(s_blk, 0, max(0, n_state - 1)) + + c_blk, s_blk_out = init_probs_from_states(s_blk, n_state) + + # Map coarse-bin initialization back to original binning. + c_init = np.zeros((T_full, n_state), dtype=np.float32) + s_init = np.zeros((T_full,), dtype=np.int64) + for b in range(n_blk): + i0 = b * dwell_bins + i1 = i0 + int(blk_sizes[b]) + c_init[i0:i1] = c_blk[b] + s_init[i0:i1] = s_blk_out[b] + + freq_h1d = rate_blk.astype(np.float32, copy=False) + init_meta = { + "rate_thres": [float(x) for x in thres], + "rate_bin_sec": float(dwell_bins * float(dt)), + "rate_bin_bins": int(dwell_bins), + } + return c_init, s_init, init_meta, freq_h1d + + +def init_selfspikingB(spikes, dt, args): + """Initialize common per-state B from observed mean firing rates. + + For each neuron n: + rate_hz[n] = mean_t spikes[t, n] / dt + B[:, n] = log(max(rate_hz[n], eps)) + """ + spikes = np.asarray(spikes) + if spikes.ndim != 2: + raise ValueError("spikes must be 2D array (T, N)") + + n_state = int(args.num_states) + eps_rate_hz = 1e-6 + + rate_hz = spikes.mean(axis=0).astype(np.float64) / float(dt) + rate_hz_clip = np.maximum(rate_hz, eps_rate_hz) + b_vec = np.log(rate_hz_clip).astype(np.float32) + b_init = np.tile(b_vec[None, :], (n_state, 1)) + + init_meta = { + "method": "rate", + "common_across_states": True, + "rate_floor_hz": float(eps_rate_hz), + "formula": "B=log(max(rate_hz,eps))", + } + return b_init, init_meta + + +def init_B_from_spikes(spikes, dt, args): + """Select B initialization mode based on args.init_B.""" + opt = _normalize_init_opt(args.init_B) + if opt == "rand": + return None, None + if opt != "data": + raise ValueError(f"Unsupported --init_B option: {opt}") + return init_selfspikingB(spikes, dt, args) + + +def init_edgesA(spikes, Tmax=50000, verbose=True): + """OLS/covariance initialization of A from spike data Y (T x N).""" + t0 = time.perf_counter() + Y = np.asarray(spikes, dtype=np.float64) + if Y.ndim != 2: + raise ValueError("spikes must be 2D array (T, N)") + T_full, _ = Y.shape + T_use = min(int(Tmax), int(T_full)) + if T_use < 2: + raise ValueError("Need at least 2 time bins to initialize A") + Y = Y[:T_use] + + YpYp = Y[:-1].T @ Y[:-1] # Gram matrix + YYp = Y[1:].T @ Y[:-1] # one-step cross-correlation + + A_ols = YYp @ np.linalg.pinv(YpYp) + + kappa = float(np.linalg.cond(YpYp)) + rho = float(np.max(np.abs(np.linalg.eigvals(A_ols)))) + Y_pred = Y[:-1] @ A_ols.T + var_y = float(np.var(Y[1:])) + if var_y <= 0.0: + R2 = float("nan") + else: + R2 = float(1.0 - np.var(Y[1:] - Y_pred) / var_y) + frob = float(np.linalg.norm(A_ols, ord="fro")) + elapsed_sec = float(time.perf_counter() - t0) + + if verbose: + print("A-init OLS diagnostics:") + print(f" cond(YpYp) = {kappa:.2e}") + print(f" rho(A_ols) = {rho:.3f}") + print(f" R2 = {R2:.3f}") + print(f" ||A||_F = {frob:.3f}") + print(f" bins_used = {T_use}/{T_full}") + print(f" elapsed_s = {elapsed_sec:.3f}") + + meta = { + "method": "cov_ols", + "Tmax": int(Tmax), + "num_bins_used": int(T_use), + "num_bins_total": int(T_full), + "cond_YpYp": kappa, + "rho_A_init": rho, + "R2_1step": R2, + "fro_A_init": frob, + "elapsed_init_edgesA_sec": elapsed_sec, + } + return A_ols.astype(np.float32), meta + + +def init_A_from_spikes(spikes, args): + """Select A initialization mode based on args.init_A.""" + + opt = _normalize_init_opt(args.init_A) + if opt == "rand": + return None, None + if opt != "data": + raise ValueError(f"Unsupported --init_A option: {opt}") + A_init, meta = init_edgesA(spikes, verbose=False) + + if 1: # rescale A + offDiagFact = 3 + meta["offDiag_A_init_fact"] = offDiagFact + A_init = A_init * offDiagFact + + if 0: # thresholded off-diagonal renormalization + off_diag = ~np.eye(A_init.shape[0], A_init.shape[1], dtype=bool) + above_thr = (np.abs(A_init) > float(args.minW)) + below_thr = (np.abs(A_init) < float(args.minW)) + strong_off_diag = off_diag & above_thr + weak_off_diag = off_diag & below_thr + A_init[strong_off_diag] *= offDiagFact + A_init[weak_off_diag] = 0.0 + + + rho_max = float(args.rho_max) + if 0: # global spectral radius rescaling + rho_before = float(np.max(np.abs(np.linalg.eigvals(A_init)))) + if rho_before > rho_max: + A_init = A_init * (rho_max / rho_before) + rho_after = float(np.max(np.abs(np.linalg.eigvals(A_init)))) + meta["rho_A_init_raw"] = rho_before + meta["rho_A_init"] = rho_after + meta["rho_max_target"] = rho_max + + if args.verb > 0: + print("A-init OLS diagnostics:") + print(f" cond(YpYp) = {meta['cond_YpYp']:.1f}") + print(f" rho(A_ols) = {meta['rho_A_init']:.3f}") + print(f" R2 = {meta['R2_1step']:.3f}") + print(f" ||A||_F = {meta['fro_A_init']:.3f}") + print(f" bins_used = {meta['num_bins_used']}/{meta['num_bins_total']}") + print(f" elapsed_s = {meta['elapsed_init_edgesA_sec']:.3f}") + + return A_init, meta + + +def init_probs_from_states(S_rate, M): + """Build c_init from S_rate with transition smoothing.""" + S_rate = np.asarray(S_rate, dtype=np.int64) + T = S_rate.size + if T == 0: + return np.zeros((0, M), dtype=np.float32), np.zeros((0,), dtype=np.int64) + if M <= 1: + return np.ones((T, 1), dtype=np.float32), S_rate.copy() + + # Base: dominant state at 0.8, remaining 0.2 spread over all others. + other = 0.2 / float(M - 1) + c_init = np.full((T, M), other, dtype=np.float32) + c_init[np.arange(T), S_rate] = 0.8 + S_out = S_rate.copy() + + # Transition smoothing: (t-1) gets 2/3 old + 1/3 new, t gets 1/3 old + 2/3 new. + tr_idx = np.where(S_rate[1:] != S_rate[:-1])[0] + 1 + for t in tr_idx: + a = int(S_rate[t - 1]) + b = int(S_rate[t]) + c_init[t - 1, :] = 0.0 + c_init[t, :] = 0.0 + c_init[t - 1, a] = 2.0 / 3.0 + c_init[t - 1, b] = 1.0 / 3.0 + c_init[t, a] = 1.0 / 3.0 + c_init[t, b] = 2.0 / 3.0 + + return c_init, S_out diff --git a/causal_net/nonStation_ver3a_states_simu/Util_pseudospectra.py b/causal_net/nonStation_ver3a_states_simu/Util_pseudospectra.py new file mode 100755 index 00000000..23fc3e1f --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/Util_pseudospectra.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +import numpy as np +import time + +def create_example_matrices(size=100): + A_base = np.diag(np.linspace(-5, -0.5, size)) + np.diag(np.ones(size - 2) * 4, k=2) + sparsity = 0.15 + random_connections = np.random.randn(size, size) + mask = np.random.rand(size, size) > sparsity + random_connections[mask] = 0 + A_true = A_base + random_connections * 0.1 + + noise = np.random.randn(size, size) * 0.10 + noise_mask = np.random.rand(size, size) > sparsity + noise[noise_mask] = 0 + A_fit = A_true + noise + + return A_true, A_fit + +def compute_pseudospectrum(A, npts, minY): + print(f'computing pseudospectrum A:{A.shape} .... ') + eigs = np.linalg.eigvals(A) + + # Calculate bounds + real_min, real_max = np.real(eigs).min(), np.real(eigs).max() + imag_min, imag_max = np.imag(eigs).min(), np.imag(eigs).max() + + real_pad = (real_max - real_min) * 0.2 + imag_pad = (imag_max - imag_min) * 0.2 + + bbox = [ + real_min - real_pad, + min(real_max + real_pad, 1.0), + imag_min, + imag_max + imag_pad + ] + + x_coords = np.linspace(bbox[0], bbox[1], npts) + y_coords = np.linspace(minY, bbox[3], npts) + #y_coords = np.linspace(bbox[2], bbox[3], npts) + X, Y = np.meshgrid(x_coords, y_coords) + Z = X + 1j * Y + + sigma_grid = np.zeros_like(Z, dtype=float) + I = np.eye(A.shape[0]) + + for i in range(npts): + for j in range(npts): + z = Z[i, j] + s = np.linalg.svd(z * I - A, compute_uv=False) + sigma_grid[i, j] = s[-1] + + return X, Y, sigma_grid, eigs + +def plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin): + levels = np.logspace(-2.5, -0.5, 10) # Use 10 contour levels + + # Plot the normal contour lines + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + + # Create a mask for the area greater than epsMin + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) # Masking areas greater than epsMin + + # Fill the area below epsMin + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + # Scatter plot for eigenvalues + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + # Vertical line at Re=0 + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + # Set titles and labels + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + # Additional formatting + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(bottom=minY) # Clip at minY + + +def main(size=100, minY=-0.9, epsMin=0.032): + A_true, A_fit = create_example_matrices(size) + + fig, axes = plt.subplots(1, 2, figsize=(16, 7)) + + matrices_to_plot = [ + ('True Matrix ($A_{true}$)', A_true), + ('Distorted Matrix ($A_{fit}$) ', A_fit) + ] + + print("Starting pseudospectrum calculations...", minY) + start_time = time.time() + + for i, (title, A) in enumerate(matrices_to_plot): + ax = axes[i] + print("Processing '{}'...".format(title)) + + npts = 80 # Grid resolution + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin) # Pass epsMin + + total_time = time.time() - start_time + print(f"Calculation finished in {total_time:.2f} seconds.") + + fig.suptitle('Pseudospectrum Comparison (Clipped Imaginary Part)', fontsize=16) + plt.tight_layout(rect=[0, 0.03, 1, 0.95]) + + outFile = 'out/pseudospectra_comparison.png' + plt.savefig(outFile) + print(f"Plot saved as '{outFile}'") + plt.close(fig) + +# --- Main execution --- +if __name__ == '__main__': + + import matplotlib.pyplot as plt + + main(size=50, minY=-0.5, epsMin=0.024) diff --git a/causal_net/nonStation_ver3a_states_simu/batchFitPrism.slr b/causal_net/nonStation_ver3a_states_simu/batchFitPrism.slr new file mode 100755 index 00000000..f6d34263 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/batchFitPrism.slr @@ -0,0 +1,119 @@ +#!/bin/bash +#SBATCH -N 1 -C gpu -A m2043 +#SBATCH --time=30:00 -q debug +#SBATCH --time 2:28:00 -q regular +#SBATCH --output=out/%j.out +#SBATCH --licenses=scratch +# - - - E N D O F SLURM C O M M A N D S +set -u ; # exit if you try to use an uninitialized variable + +basePath=/pscratch/sd/b/balewski/2026_causalNet_tmp3/ + +#dataName=daleN100_23b4f4_787a5b ; num_states=3 # N100 13/21/28 Hz +dataName=daleN100_07a07b_0fa911 ; num_states=3 # N100 13/14/15 Hz +dataName=daleN200_8cbdb0_9efd1f ; num_states=3 # N200 10/19/27 Hz +dataName=daleN200_1c12df_eca436 ; num_states=4 # N200 10/15/20/25 Hz + +miscArgs="" +#miscArgs="--lr_estep 0.02 --lr_mstep 6e-4 --lr_end_factor 0.1 --L1_alpha 0.02" # try + +# spare --noFreqWeight + +# Optional overrides passed to this script +num_em_iters=5 +m_epochs=100 +time_range_sec="0 80" + +while [[ $# -gt 0 ]]; do + case "$1" in + --num_em_iters) + if [[ $# -lt 2 ]]; then + echo "ERROR: --num_em_iters requires a value" >&2 + exit 1 + fi + num_em_iters="$2" + shift 2 + ;; + --num_em_iters=*) + num_em_iters="${1#*=}" + shift + ;; + --m_epochs) + if [[ $# -lt 2 ]]; then + echo "ERROR: --m_epochs requires a value" >&2 + exit 1 + fi + m_epochs="$2" + shift 2 + ;; + --m_epochs=*) + m_epochs="${1#*=}" + shift + ;; + --time_range_sec) + if [[ $# -lt 3 ]]; then + echo "ERROR: --time_range_sec requires 2 values: " >&2 + exit 1 + fi + time_range_sec="$2 $3" + shift 3 + ;; + --time_range_sec=*) + range_val="${1#*=}" + if [[ "$range_val" == *,* ]]; then + time_range_sec="${range_val/,/ }" + else + time_range_sec="$range_val" + fi + shift + ;; + -h|--help) + echo "Usage: $0 [--num_em_iters N] [--m_epochs N] [--time_range_sec t0 t1]" + exit 0 + ;; + *) + echo "ERROR: Unknown argument: $1" >&2 + exit 1 + ;; + esac +done + +if ! [[ "$num_em_iters" =~ ^[0-9]+$ ]] || [ "$num_em_iters" -lt 1 ]; then + echo "ERROR: --num_em_iters must be a positive integer, got: $num_em_iters" >&2 + exit 1 +fi +if ! [[ "$m_epochs" =~ ^[0-9]+$ ]] || [ "$m_epochs" -lt 1 ]; then + echo "ERROR: --m_epochs must be a positive integer, got: $m_epochs" >&2 + exit 1 +fi +read -r t0_sec t1_sec extra <<< "$time_range_sec" +if [ -n "${extra:-}" ] || [ -z "${t0_sec:-}" ] || [ -z "${t1_sec:-}" ]; then + echo "ERROR: --time_range_sec must contain exactly 2 values, got: $time_range_sec" >&2 + exit 1 +fi +num_re='^-?[0-9]+([.][0-9]+)?$' +if ! [[ "$t0_sec" =~ $num_re ]] || ! [[ "$t1_sec" =~ $num_re ]]; then + echo "ERROR: --time_range_sec values must be numeric, got: $time_range_sec" >&2 + exit 1 +fi + + +N=${SLURM_NNODES} +jobId=${SLURM_JOBID} + +echo S: job dataName=$dataName time_range_sec=$time_range_sec jobId=$jobId +echo S: num_em_iters=$num_em_iters m_epochs=$m_epochs +module load pytorch + +last_four=${jobId: -4} # Extract last 4 characters +fitName=${dataName}_nsn${last_four} + +time ./fitPrismEM.sh --basePath $basePath --dataName $dataName --num_states $num_states $miscArgs --num_em_iters $num_em_iters --m_epochs $m_epochs --time_range_sec $time_range_sec --fitName $fitName +echo +echo 'S:fit done' +echo +./prism_EM_eval.py --basePath $basePath --dataName $fitName -p a b c d e f g -X + +echo +echo 'S:eval done' +echo diff --git a/causal_net/nonStation_ver3a_states_simu/docs/build3eval.sh b/causal_net/nonStation_ver3a_states_simu/docs/build3eval.sh new file mode 100755 index 00000000..2ed4a417 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/docs/build3eval.sh @@ -0,0 +1,32 @@ +#!/bin/bash + +set -e + +doc_name=evalDaleStability_ver3 +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$script_dir" + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $script_dir/$pdf_file" + +# Open the PDF based on the specific environment. +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" || echo "Could not open PDF automatically. PDF is ready at $script_dir/$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" || echo "Could not open PDF automatically. PDF is ready at $script_dir/$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + if [ -n "${DISPLAY:-}" ] && command -v gv >/dev/null 2>&1; then + gv "$pdf_file" || echo "Could not open PDF automatically. PDF is ready at $script_dir/$pdf_file" + else + echo "No graphical display found. PDF is ready at $script_dir/$pdf_file" + fi +else + echo "Environment not matched. PDF is ready at $script_dir/$pdf_file" +fi diff --git a/causal_net/nonStation_ver3a_states_simu/docs/build3fit.sh b/causal_net/nonStation_ver3a_states_simu/docs/build3fit.sh new file mode 100755 index 00000000..cfccdb8b --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/docs/build3fit.sh @@ -0,0 +1,32 @@ +#!/bin/bash + +set -e + +doc_name=fitDale_states_ver3 +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$script_dir" + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $script_dir/$pdf_file" + +# Open the PDF based on the specific environment. +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" || echo "Could not open PDF automatically. PDF is ready at $script_dir/$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" || echo "Could not open PDF automatically. PDF is ready at $script_dir/$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + if [ -n "${DISPLAY:-}" ] && command -v gv >/dev/null 2>&1; then + gv "$pdf_file" || echo "Could not open PDF automatically. PDF is ready at $script_dir/$pdf_file" + else + echo "No graphical display found. PDF is ready at $script_dir/$pdf_file" + fi +else + echo "Environment not matched. PDF is ready at $script_dir/$pdf_file" +fi diff --git a/causal_net/nonStation_ver3a_states_simu/docs/build3gen.sh b/causal_net/nonStation_ver3a_states_simu/docs/build3gen.sh new file mode 100755 index 00000000..51c12bf2 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/docs/build3gen.sh @@ -0,0 +1,32 @@ +#!/bin/bash + +set -e + +doc_name=genDale_topo_states_ver3 +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$script_dir" + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $script_dir/$pdf_file" + +# Open the PDF based on the specific environment. +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" || echo "Could not open PDF automatically. PDF is ready at $script_dir/$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" || echo "Could not open PDF automatically. PDF is ready at $script_dir/$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + if [ -n "${DISPLAY:-}" ] && command -v gv >/dev/null 2>&1; then + gv "$pdf_file" || echo "Could not open PDF automatically. PDF is ready at $script_dir/$pdf_file" + else + echo "No graphical display found. PDF is ready at $script_dir/$pdf_file" + fi +else + echo "Environment not matched. PDF is ready at $script_dir/$pdf_file" +fi diff --git a/causal_net/nonStation_ver3a_states_simu/docs/evalDaleStability_ver3.pdf b/causal_net/nonStation_ver3a_states_simu/docs/evalDaleStability_ver3.pdf new file mode 100644 index 00000000..3bc0f4f9 Binary files /dev/null and b/causal_net/nonStation_ver3a_states_simu/docs/evalDaleStability_ver3.pdf differ diff --git a/causal_net/nonStation_ver3a_states_simu/docs/evalDaleStability_ver3.tex b/causal_net/nonStation_ver3a_states_simu/docs/evalDaleStability_ver3.tex new file mode 100644 index 00000000..d7b4bb20 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/docs/evalDaleStability_ver3.tex @@ -0,0 +1,431 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} +\usepackage{enumitem} + +\geometry{margin=0.9in} + +\title{Evaluating Recovery of a Dale-Constrained Connectivity Matrix\\ +from Synthetic Spike Data:\\ +Study Design and Comparison Metrics} +\author{Jan Balewski, NERSC} +\date{\today} + +\newcommand{\Atrue}{A_{\mathrm{true}}} +\newcommand{\Ahat}{\hat{A}} +\newcommand{\Etrue}{E_{\mathrm{true}}} +\newcommand{\minW}{\mathrm{minW}} +\newcommand{\neff}{n_{\mathrm{eff}}} +\newcommand{\Trec}{T_{\mathrm{rec}}} + +\begin{document} + +\maketitle + +\begin{abstract} +We define a protocol for quantifying how accurately the shared connectivity +matrix $\Atrue$ of a Dale-law-constrained recurrent Poisson GLM is recovered +by PRISM-EM from non-stationary spike trains, as a function of recording +duration. Ground truth is known by construction. The evaluation is split into +three independent concerns---\emph{support} (which edges exist), \emph{sign} +(excitatory vs.\ inhibitory), and \emph{magnitude} (edge weight)---of which +this document treats support and magnitude in detail. We fix the pruning +threshold $\minW$ and the regularization strength for the present study, apply +$\minW$ to the raw fitted weights \emph{before} any spectral rescaling, and +compare magnitudes only after normalizing both matrices to unit spectral +radius. We motivate the expected $\propto 1/\sqrt{\neff}$ scaling of weight +estimation error with effective sample size and use it as the quantitative +null hypothesis for the duration sweep. +\end{abstract} + +%\tableofcontents +%\newpage + +% =============================================================== +\section{Background: Data and Fit in Brief} +\label{sec:background} +% =============================================================== + +\subsection{Generator} +A network of $N$ neurons ($N_E$ excitatory, $N_I = N - N_E$ inhibitory) is +simulated as a lag-1 Poisson vector autoregression (VAR(1)). A single shared +connectivity matrix $\Atrue \in \mathbb{R}^{N\times N}$ governs all dynamics; +$M$ latent states differ only in a per-state bias (log-baseline) vector $B_m$. +At each time bin $t$ of width $\Delta t$, +\begin{equation} + \eta_t = \Atrue\, Y_{t-1} + B_{\mathrm{eff},t},\qquad + \lambda_t = \exp\!\bigl(\min(\eta_t,\eta_{\mathrm{clip}})\bigr)\Delta t,\qquad + Y_t \sim \mathrm{Poisson}(\lambda_t), +\end{equation} +where $B_{\mathrm{eff},t}=\sum_m c_{t,m}B_m$ mixes the state biases along a +smooth simplex trajectory $c_t$. Column convention: $(\Atrue)_{ij}$ is the +weight \emph{from} presynaptic neuron $j$ \emph{to} postsynaptic neuron $i$. +Dale's law is enforced at the source: every outgoing weight of an excitatory +neuron is positive, every outgoing weight of an inhibitory neuron is negative, +and self-loops are negative. The binary support $\Etrue\in\{0,1\}^{N\times N}$ +is generated by per-neuron i.i.d.\ Bernoulli connectivity, and the spectral +radius is fixed at $\rho(\Atrue)=R<1$. Ground-truth arrays +$(\Atrue, \{B_m\}, \Etrue)$ and the $(x,y)$ coordinates of all neurons are +saved. + +\subsection{Fitter (PRISM-EM)} +PRISM-EM recovers $\Ahat$ and $\{\hat B_m\}$ by block coordinate descent: an +E-step infers the simplex coefficients $c_{1:T}$ and Viterbi-decodes them to +one-hot state assignments; an M-step updates $\Ahat,\hat{\mathbf B}$ by Adam on +the Poisson negative log-likelihood with off-diagonal $\ell_1$ shrinkage +(strength $\lambda_3$) and a spectral-radius projection to $\rho_{\max}$. The +fitter does \emph{not} impose Dale's law, so recovered rows may contain mixed +signs; this is exploited later only as a sign-consistency diagnostic and is out +of scope here. Because $\Atrue$ is shared across all states, every time bin +contributes to estimating $\Ahat$, and the arbitrary permutation of recovered +state labels does \emph{not} affect $\Ahat$; the connectivity comparison +therefore requires no label alignment. + +% =============================================================== +\section{Part I: Study Design and Conventions} +\label{sec:design} +% =============================================================== + +\subsection{Operating Regime} +We target $N\in[50,500]$ neurons with a median single-neuron firing rate of +$2$--$3$\,Hz (range $[1,5]$\,Hz), spikes binned at $\Delta t = 10$\,ms, and +recording durations from $10$ to $60$ minutes. The latent dynamics use +$M\in\{2,3\}$ states with mean dwell time $\sim 1$\,s. The central scientific +question is: \emph{how much recovery quality is lost by using a $10$-minute +record instead of a $60$-minute record?} + +\subsection{Master-Record Subsampling} +To remove generator variability from the duration comparison, a single +\emph{master} record of $120$ minutes is generated per configuration +$(N,\text{seed})$. Shorter records are obtained as contiguous temporal +prefixes, +\begin{equation} + Y^{(D)} = Y_{1:T_D},\qquad T_D = \lfloor D/\Delta t\rfloor, +\end{equation} +for durations $D\in\{10,15,20,30,45,60\}$ minutes, with the full $120$-minute +record retained as an asymptotic ceiling. Contiguous prefixes (rather than +random bin subsets) preserve the temporal lag structure and the state-dwell +statistics that the fitter depends on. Each metric is reported relative to its +$120$-minute value, so results read as ``fraction of asymptotic recovery +retained.'' At least $5$ independent seeds are run per $(N,D)$ pair; sampling +variability in this sparse, low-rate regime is large and single-seed results +are not interpretable. + +\subsection{Parameters Held Constant (Present Study)} +For this phase the following are fixed across all durations and seeds: +\begin{itemize}[nosep] + \item Pruning threshold $\minW = 0.02$ (applied to raw $\Ahat$; see below). + \item $\ell_1$ strength $\lambda_3$ and all other fit hyperparameters + (learning rates, EM/M-step iteration counts, $\rho_{\max}$, delay + schedules). + \item Network structure $(N, N_E, R, \Etrue, \Atrue, \{B_m\})$ within a seed. +\end{itemize} +Holding $\lambda_3$ fixed means we measure \emph{fixed-regularizer} recovery. +Because the optimal $\ell_1$ penalty scales with the noise level +($\propto 1/\sqrt{\neff}$, Section~\ref{sec:neff}), a fixed $\lambda_3$ +over-shrinks short records; the measured duration effect therefore includes a +regularization component that we acknowledge explicitly and revisit in later +work. + +\subsection{Threshold Then Rescale: Order of Operations} +\label{sec:order} +Support and magnitude use different normalizations of $\Ahat$, and the order +matters. +\begin{enumerate}[nosep] + \item \textbf{Support} is evaluated on the \emph{raw} fitted matrix. An edge + $(i,j)$ is declared present iff $|\Ahat_{ij}| > \minW$, $i\neq j$. The + threshold is absolute and is applied at the fitter's native scale + ($\rho \approx \rho_{\max}$). + \item \textbf{Magnitude} is evaluated after rescaling both matrices to unit + spectral radius, + \begin{equation} + \Ahat^{\star} = \Ahat / \rho(\Ahat),\qquad + \Atrue^{\star} = \Atrue / \rho(\Atrue), + \end{equation} + which removes the global scale mismatch between the fitter's + $\rho_{\max}$ constraint and the generator's $R$. This rescaling is a + single scalar per matrix; it preserves all signs, the zero pattern, and + all relative orderings. +\end{enumerate} +The two normalizations are kept in separate pipelines. Applying the absolute +$\minW$ after a $\rho$-rescaling would make its effective cut size depend on +$N$ (since unit-$\rho$ entries scale roughly as the inverse square root of the +number of active edges), introducing an $N$-dependent artifact into the support +results; thresholding the raw weights avoids this. All metrics, both support +and magnitude, exclude the diagonal. + +\subsection{Effective Sample Size and the $1/\sqrt{\neff}$ Law} +\label{sec:neff} +The recoverability of an edge $j\!\to\!i$ is governed not by raw duration but +by how often the presynaptic neuron $j$ is observed firing. The effective +sample size for that edge is approximately the number of bins in which $j$ is +active, +\begin{equation} + \neff(j) \;\approx\; f_j \,\Trec, +\end{equation} +with $f_j$ the firing rate of neuron $j$ and $\Trec$ the record duration. At +the median rate $f_j \approx 2.5$\,Hz this gives $\neff \approx 1{,}500$ +presynaptic spikes at $10$ minutes and $\approx 9{,}000$ at $60$ minutes; for +the slowest neurons ($1$\,Hz) the $10$-minute count drops to $\approx 600$. +The slowest presynaptic neurons therefore bound the worst-case edge recovery. + +For maximum-likelihood estimation of GLM coefficients, the standard error of +each weight scales as +\begin{equation} + \mathrm{SE}(\Ahat_{ij}) \;\propto\; \frac{1}{\sqrt{\neff(j)}}. +\end{equation} +Reducing the record from $60$ to $10$ minutes is a $6\times$ reduction in data, +predicting +\begin{equation} + \frac{\mathrm{SE}_{10}}{\mathrm{SE}_{60}} + \;\approx\; \sqrt{6} \;\approx\; 2.45. +\end{equation} +This is the quantitative \emph{null hypothesis} for the duration sweep: +magnitude-estimation error should grow by $\approx\!2.45\times$ from $60$ to +$10$ minutes if sampling noise alone is responsible. Larger growth indicates an +additional failure mode (e.g.\ EM converging to a worse optimum, or support +errors corrupting the surviving magnitudes); smaller growth indicates a +regularization floor where $\lambda_3$ dominates sampling noise and duration +matters little. \textbf{Important caveat:} support recovery does \emph{not} +follow this smooth law. Weak edges sit near a detection threshold; as $\neff$ +falls, edges near that threshold drop out abruptly, so recall (especially of +weak edges) degrades in a step-like fashion while precision can collapse from +proliferating noise-floor false positives. + +% =============================================================== +\section{Part II: Comparing Support} +\label{sec:support} +% =============================================================== + +The predicted support is the off-diagonal binary mask +$\hat E_{ij} = \mathbf{1}\{|\Ahat_{ij}| > \minW\}$; ground truth is $\Etrue$. +On the off-diagonal entries define the confusion counts +\begin{equation} + \mathrm{TP},\;\mathrm{FP},\;\mathrm{FN},\;\mathrm{TN} +\end{equation} +(edge present in both / only in $\hat E$ / only in $\Etrue$ / absent in both). +Because the network is sparse ($\sim 5$--$20\%$ of entries are edges), the +overwhelming majority of entries are TN, and any metric that rewards TN will +read optimistically. We therefore favor edge-focused metrics and always report +the raw counts $(\mathrm{TP},\mathrm{FP},\mathrm{FN})$ so that a change in any +summary can be attributed to missed edges (FN) versus hallucinated edges (FP), +which are distinct failure modes. + +\subsection{Candidate Scalar Metrics} +All are evaluated at the fixed $\minW = 0.02$. + +\paragraph{Jaccard index (primary).} +\begin{equation} + J = \frac{\mathrm{TP}}{\mathrm{TP}+\mathrm{FP}+\mathrm{FN}}. +\end{equation} +Ignores TN entirely; answers ``of all edges either matrix claims, what +fraction do both agree on.'' Most conservative and most edge-focused of the +candidates. + +\paragraph{$F_1$ score.} +\begin{equation} + F_1 = \frac{2\,\mathrm{TP}}{2\,\mathrm{TP}+\mathrm{FP}+\mathrm{FN}}. +\end{equation} +Also ignores TN. $F_1$ and $J$ are monotone transforms of one another, +$J = F_1/(2-F_1)$, so they rank fits identically; $J$ is simply harsher on the +same scale. Retained for comparison only. + +\paragraph{Precision and Recall.} +\begin{equation} + \mathrm{Prec}=\frac{\mathrm{TP}}{\mathrm{TP}+\mathrm{FP}},\qquad + \mathrm{Rec}=\frac{\mathrm{TP}}{\mathrm{TP}+\mathrm{FN}}. +\end{equation} +Reported as the decomposition underlying $J$ and $F_1$: recall tracks loss of +real edges as $\neff$ falls, precision tracks the proliferation of noise-floor +false positives. + +\paragraph{Matthews correlation coefficient (secondary).} +\begin{equation} + \mathrm{MCC} = + \frac{\mathrm{TP}\cdot\mathrm{TN} - \mathrm{FP}\cdot\mathrm{FN}} + {\sqrt{(\mathrm{TP}+\mathrm{FP})(\mathrm{TP}+\mathrm{FN}) + (\mathrm{TN}+\mathrm{FP})(\mathrm{TN}+\mathrm{FN})}}. +\end{equation} +Uses all four cells, including TN, so it reads higher than $J$ in the sparse +regime; retained for comparability with the broader literature. Reports the +overall correlation between the two binary matrices. + +We will run all of the above initially and discard the least discriminating +once empirical sensitivity is known. + +\subsection{Block-Resolved Support (E/I structure)} +\label{sec:blocks} +The neuron ordering induces four meaningful blocks, indexed by presynaptic type +(\emph{column} block, by the from-$j$ convention) and postsynaptic type +(\emph{row} block): +\begin{equation} + \text{E}\!\to\!\text{E},\quad + \text{E}\!\to\!\text{I},\quad + \text{I}\!\to\!\text{E},\quad + \text{I}\!\to\!\text{I}. +\end{equation} +Each scalar metric of Section~\ref{sec:support} is recomputed within each block. +Because inhibitory neurons are a minority ($\sim 20\%$) and typically carry +different effective drive, we expect $\text{I}\!\to\!\cdot$ edges to recover +more poorly; the block table makes this circuit-meaningful rather than +arbitrary. + +\subsection{Spatially Coarse-Grained Support} +\label{sec:grid} +Using the known $(x,y)$ coordinates, overlay a coarse $G\times G$ grid on the +neuron positions. Let $g(i)\in\{1,\dots,G^2\}$ be the grid cell of neuron $i$, +and partition neurons additionally by type $\tau(i)\in\{\mathrm{E,I}\}$. Define +a \emph{lumped} (super-node) connectivity between ordered cell-type pairs by +counting only edges that leave one cell and enter another. For source cell $a$ +of type $\sigma$ and target cell $b$ of type (any, or fixed), the lumped true +and recovered edge counts are +\begin{align} + \mathcal{N}^{\mathrm{true}}_{(a,\sigma)\to b} + &= \sum_{\substack{j:\,g(j)=a,\,\tau(j)=\sigma \\ i:\,g(i)=b,\; i\neq j}} + (\Etrue)_{ij}, \\ + \mathcal{N}^{\mathrm{hat}}_{(a,\sigma)\to b} + &= \sum_{\substack{j:\,g(j)=a,\,\tau(j)=\sigma \\ i:\,g(i)=b,\; i\neq j}} + \mathbf{1}\{|\Ahat_{ij}|>\minW\}. +\end{align} +Restricting the source sum to a single type $\sigma$ keeps excitatory and +inhibitory super-edges separate, as intended. Comparison at the super-node +level can use either a count-agreement measure, +\begin{equation} + \mathrm{err}_{(a,\sigma)\to b} + = \frac{\bigl|\mathcal{N}^{\mathrm{hat}}_{(a,\sigma)\to b} + - \mathcal{N}^{\mathrm{true}}_{(a,\sigma)\to b}\bigr|} + {\mathcal{N}^{\mathrm{true}}_{(a,\sigma)\to b}+1}, +\end{equation} +or a fine-grained Jaccard restricted to the edge set between the two cells. +Rendering $\mathrm{err}_{(a,\sigma)\to b}$ as a spatial heatmap reveals whether +recovery quality varies across the tissue. \textbf{Interpretive caution:} in +the current generator the binary support is spatially i.i.d., so $\Etrue$ has no +intrinsic spatial structure. Any spatial pattern in recovery is therefore an +\emph{inference artifact}, and is expected to be explained almost entirely by +the spatial distribution of presynaptic firing rates rather than by position +per se (Section~\ref{sec:strat}). The grid view is thus a descriptive lens; the +mechanistic axis is firing rate. + +\subsection{Stratified Recovery} +\label{sec:strat} +The mechanistic drivers of edge recoverability are presynaptic firing rate and +true edge magnitude. We therefore stratify recall: +\begin{itemize}[nosep] + \item by presynaptic rate $f_j$ (binned), + \item by true magnitude $|Atrue_{ij}|$ (binned), + \item jointly, as a two-dimensional recall heatmap over + $(f_j,,|Atrue_{ij}|)$. +\end{itemize} +The joint heatmap is the single most explanatory support diagnostic: it shows +exactly which edges become unrecoverable as duration shrinks, and it should +account for any apparent spatial or block-wise pattern. + +% =============================================================== +\section{Part III: Comparing Magnitude} +\label{sec:magnitude} +% =============================================================== + +Magnitude comparison is restricted to the true support +($Etrue_{ij}=1$, $i\neq j$) and uses the unit-$\rho$ matrices +$Ahat^{\star},Atrue^{\star}$ of Section~\ref{sec:order}. False-positive edges +are the concern of the support metrics and are excluded here, except where +noted (the global-vs-support Frobenius gap below). Because the true edges span +both signs (positive E, negative I), a metric computed across all signs is +partly rewarded for merely separating E from I---the easy part. To isolate +magnitude fidelity from sign, the primary forms of each metric are computed +\emph{within} the same-sign edge sets (within-E and within-I), with the global +(both-sign) form reported as a secondary, sign-inflated reference. + +\subsection{Rank Fidelity: Spearman Correlation} +Over the chosen edge set, the Spearman correlation $\rho_S$ between +$\{Atrue^{\star}_{ij}\}$ and $\{Ahat^{\star}_{ij}\}$ measures preservation of +the \emph{ordering} of edge strengths: +\begin{equation} + \rho_S = \mathrm{corr}\bigl(\mathrm{rank}(Atrue^{\star}),\; + \mathrm{rank}(Ahat^{\star})\bigr). +\end{equation} +Computed within-E, within-I, and globally. Spearman is invariant to any +monotone rescaling, so it is unaffected by the unit-$\rho$ normalization---a +useful internal consistency check (a rescaling bug moves Pearson but leaves +Spearman fixed). It weights all ranks equally and is therefore sensitive to the +fidelity of \emph{weak} edges. + +\subsection{Magnitude Error: Relative Frobenius Norm} +Two complementary versions are reported. +\begin{align} + \mathrm{RFE}_{\mathrm{supp}} + &= \frac{\bigl\|(Ahat^{\star}-Atrue^{\star})\odot \Etrue\bigr\|_F} + {\bigl\|Atrue^{\star}\odot \Etrue\bigr\|_F}, \\[4pt] + \mathrm{RFE}_{\mathrm{glob}} + &= \frac{\bigl\|Ahat^{\star}-Atrue^{\star}\bigr\|_F} + {\bigl\|Atrue^{\star}\bigr\|_F}, +\end{align} +both over off-diagonal entries. $\mathrm{RFE}_{\mathrm{supp}}$ measures +magnitude error on real edges; $\mathrm{RFE}_{\mathrm{glob}}$ additionally +includes the mass placed on false-positive entries. The gap +$\mathrm{RFE}_{\mathrm{glob}}-\mathrm{RFE}_{\mathrm{supp}}$ is a single-number +summary of spurious recovered mass and is expected to grow as duration shrinks. +The Frobenius norm is dominated by the largest-magnitude edges and is thus +sensitive to the \emph{strong} edges---complementary to Spearman. + +\subsection{Systematic Shrinkage: Regression Slope} +Fit, over the chosen edge set, the linear relation +\begin{equation} + Ahat^{\star}_{ij} = \beta\,Atrue^{\star}_{ij} + \gamma, +\end{equation} +reporting the slope $\beta$ and intercept $\gamma$ within-E and within-I, plus a +zero-intercept variant ($\gamma\equiv 0$). The slope $\beta$ is the shrinkage +diagnostic. After unit-$\rho$ rescaling the spectral-constraint mismatch is +removed, so a residual $\beta\neq 1$ is attributable to $\ell_1$ shrinkage and +finite-sample estimation bias; $\ell_1$ is expected to push $\beta<1$, more so +for the noisier inhibitory edges. A nonzero $\gamma$ on the support-restricted +fit flags a magnitude offset, which is physically unexpected (no edge implies +no systematic shift). + +\subsection{The Underlying Scatter} +The three scalars above are summaries of one object: the scatter of +$Ahat^{\star}_{ij}$ against $Atrue^{\star}_{ij}$ over the true support, +colored by presynaptic rate $f_j$. This plot reveals shrinkage (slope), +heteroscedasticity (weak/slow edges noisier), and outliers (slow-presynaptic +edges) at a glance, and should be inspected before interpreting the scalars. + +\subsection{Sensitivity Complementarity (Summary)} +The three magnitude metrics are deliberately complementary in what they +emphasize: +\begin{center} +\begin{tabular}{lll} +\toprule +Metric & Emphasizes & Scale dependence \\ +\midrule +Spearman $\rho_S$ & ordering of all edges (incl.\ weak) & invariant \\ +Relative Frobenius & large-magnitude (strong) edges & requires unit-$\rho$ \\ +Regression slope $\beta$ & systematic shrinkage bias & requires unit-$\rho$ \\ +\bottomrule +\end{tabular} +\end{center} + +% =============================================================== +\section{Headline Deliverable} +\label{sec:deliverable} +% =============================================================== + +The duration study is summarized by curves of the retained primary metrics +\begin{equation} + \{\,J,\; \rho_S^{\,\mathrm{E}},\; \rho_S^{\,\mathrm{I}},\; + \mathrm{RFE}_{\mathrm{supp}},\; \beta^{\,\mathrm{E}},\; \beta^{\,\mathrm{I}}\,\} + \quad\text{vs.}\quad D\in\{10,\dots,60,120\}\text{ min}, +\end{equation} +each normalized to its $120$-minute value and averaged over seeds with error +bars. The magnitude-error curve is overlaid with the +$\sqrt{6}\approx 2.45$ sampling-noise reference (Section~\ref{sec:neff}); +deviations from this reference, together with the step-like behavior of recall +in the stratified heatmap (Section~\ref{sec:strat}), constitute the principal +findings on information loss between the $10$- and $60$-minute records. + +\end{document} diff --git a/causal_net/nonStation_ver3a_states_simu/docs/fitDale_states_ver3.pdf b/causal_net/nonStation_ver3a_states_simu/docs/fitDale_states_ver3.pdf new file mode 100644 index 00000000..b1ba5e69 Binary files /dev/null and b/causal_net/nonStation_ver3a_states_simu/docs/fitDale_states_ver3.pdf differ diff --git a/causal_net/nonStation_ver3a_states_simu/docs/fitDale_states_ver3.tex b/causal_net/nonStation_ver3a_states_simu/docs/fitDale_states_ver3.tex new file mode 100644 index 00000000..b60d15eb --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/docs/fitDale_states_ver3.tex @@ -0,0 +1,392 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} + +\geometry{margin=0.8in} + +\title{PRISM-EM for Non-Stationary Poisson GLMs with states from Spike Trains:\\ +Model, Objective, and Implementation. Ver3} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +This document describes the current \texttt{prism\_EM\_train3.py} implementation +(script header: \texttt{prism\_EM\_train.py}) for fitting a non-stationary +Poisson GLM with latent state mixing (PRISM-EM)~\footnote{PRISM: PRobabilistic Inference of Sparse Mixed-state dynamics; ~~EM: Expectations Maximization}. Input spike trains are read from \texttt{spikesData/.spikes.npz} (non-stationary output of \texttt{gen\_nonStationarySpikes3.py}). The model uses a shared connectivity matrix across states, state-specific bias vectors, and time-varying simplex coefficients. Training is performed by block coordinate descent: a sequential projected-gradient E-step for coefficients followed by Viterbi decoding to one-hot state assignments, and an Adam-based M-step for model parameters with off-diagonal L1 shrinkage and spectral-radius projection. Spectral-radius enforcement, L1 pruning, and learning-rate decay are each activated after a configurable number of EM iterations. Distributed training (\texttt{torchrun}) shards the E-step over time bins and synchronizes \(A\), \(B\), and \(c\) across GPUs. We summarize the exact objective used in code, the decoding method, all current default hyperparameters, initialization modes (data-driven vs random) for \(A\), \(B\), and \(C\), and the saved fit file layout. +\end{abstract} + +\tableofcontents +\newpage + +% --------------------------------------------------------------- +\section{EM Training for Non-Stationary Poisson GLM with states} +\label{sec:em} +% --------------------------------------------------------------- + +\subsection{Introduction and Model Assumptions} + +We consider a population of \(N\) neurons observed over \(T\) discrete time +bins of width \(\Delta t\) (seconds). The population activity is governed +by \(M\) latent dynamical states that switch over time. The key +structural assumption is that all states share a single +connectivity matrix \(A \in \mathbb{R}^{N \times N}\) but differ in +baseline excitability through per-state bias vectors +\(B_m \in \mathbb{R}^N,\; m = 0, \dots, M-1\) (0-based state indices in code). + +At each time bin \(t\), a coefficient vector +\(c_t \in \mathbb{R}^M\), constrained to the probability simplex +(\(c_{mt}\ge 0,\; \sum_m c_{mt}=1\)), selects the effective bias as a +convex combination: +\begin{equation} + B_t = \sum_{m=1}^{M} c_{mt}\, B_m = c_t^\top \mathbf{B}, +\end{equation} +where \(\mathbf{B} \in \mathbb{R}^{M \times N}\) stacks the \(M\) bias +vectors row-wise. + +Spike counts follow +\begin{equation} + \eta_t = A\,Y_{t-1} + B_t,\qquad + \tilde{\eta}_t = \min(\eta_t,\eta_{\mathrm{clip}}),\qquad + \mu_t = \exp(\tilde{\eta}_t)\,\Delta t,\qquad + Y_t \sim \mathrm{Poisson}(\mu_t). + \label{eq:glm} +\end{equation} + +\paragraph{Notation summary.} +\(N\): number of neurons; +\(M\): number of latent states; +\(T\): number of time bins; +\(\Delta t\): time-bin width (from metadata); +\(\eta_{\mathrm{clip}}\): clipping threshold (from metadata); +\(Y_t \in \mathbb{N}_0^N\): observed spikes; +\(\mu_t \in \mathbb{R}_+^N\): expected spike count per bin. + +% --------------------------------------------------------------- +\subsection{Objective Function} + +The training objective uses uniformly weighted (unweighted) Poisson NLL in both the E-step and M-step: +\begin{equation} +\mathcal{J} += +\underbrace{\sum_{t=2}^{T}\sum_{i=1}^{N} +\left[\mu_{t,i}-Y_{t,i}\left(\tilde{\eta}_{t,i}+\log\Delta t\right)\right]}_{\text{Poisson NLL}} ++ +\underbrace{\lambda_2\sum_{t=2}^{T}\|\vec c_t- \vec c_{t-1}\|^2}_{\text{temporal smoothness (E-step only)}} ++ +\underbrace{\lambda_3\,\overline{|A_{\mathrm{off}}|}}_{\text{L1 sparsity on off-diagonal (M-step only)}}, +\label{eq:objective} +\end{equation} +with simplex constraints \(c_t\in\Delta^{M-1}\) and spectral constraint +\(\rho(A)\le \rho_{\max}\), where the off-diagonal L1 term is +\begin{equation} + \overline{|A_{\mathrm{off}}|} + = + \frac{1}{N(N-1)} + \sum_{\substack{i,j=1\\i\neq j}}^{N}|A_{ij}|. + \label{eq:l1_offdiag} +\end{equation} +In the $\lambda_2$ term, the sum over states $m$ +is already implicit in the $L_2$ norm. + + +\subsection{Why Block Coordinate Descent (EM) over Joint Optimization} + +In principle one could jointly optimize \(A\), \(\mathbf{B}\), and +\(c_{1:T}\) by simultaneous gradient descent on the full objective +\(\mathcal{J}\). In practice, the EM-style alternation is preferable +for three structural reasons. + +\paragraph{Breaking the bilinear coupling.} +The coefficients \(c_t\) and the bias matrix \(\mathbf{B}\) enter the +log-rate linearly through the product \(c_t^\top\mathbf{B}\), which +then passes through an exponential. When both are free to move, the +landscape contains saddle points and flat valleys: a large \(c_{tm}\) +paired with a small \(B_m\) produces the same likelihood as the +reverse. Fixing one side at each step eliminates this ambiguity. The +Viterbi hard-assignment variant sharpens the separation further: +with one-hot \(c_t = e_{\hat{S}_t}\), each time bin selects exactly one +bias row, so the M-step reduces to \(M\) independent Poisson GLM +regressions and the bilinear degeneracy vanishes entirely. + +\paragraph{Exploiting temporal structure in the E-step.} +The coefficient vectors \(c_{1:T}\) number \(O(TM)\) and vastly +outnumber the \(O(N^2+MN)\) model parameters. Joint optimization +would need to balance learning rates across these disparate groups +while enforcing a simplex constraint at every time bin. By contrast, +the E-step decomposes into a single forward sweep of cheap +\(M\)-dimensional simplex-projected gradient updates, naturally +incorporating the temporal smoothness penalty +\(\lambda_2\|c_t-c_{t-1}\|^2\) through the sequential dependence on the +previous accepted \(c_{t-1}\). + +\paragraph{Avoiding the log-sum-exp mismatch.} +Under soft (probabilistic) state assignments the M-step must fit +parameters to the mixture rate +\(\sum_m c_{tm}\exp(\eta_{tm})\), whereas the model parameterizes the +log-rate as \(\sum_m c_{tm}\,\eta_{tm}\) inside a single exponential. +These two quantities differ whenever the assignment is not one-hot, +creating a systematic bias in the gradient signal for \(\mathbf{B}\). +The hard-EM strategy—Viterbi decode followed by one-hot +encoding—eliminates this mismatch by construction, at the cost of +losing the monotonic likelihood improvement guarantee of soft EM, a +trade-off that is empirically favorable. + +% --------------------------------------------------------------- +\subsection{Algorithm: Block Coordinate Descent (EM)} + +The non-convex problem is solved by alternating: +\begin{itemize} + \item E-step: update \(c_{1:T}\) with \(A,\mathbf{B}\) fixed, then Viterbi-decode to one-hot state assignments. + \item M-step: update \(A,\mathbf{B}\) with one-hot \(c_{1:T}\) fixed. +\end{itemize} +for \(K_{\mathrm{EM}}\) outer iterations. + +\subsubsection{E-step: Coefficient Inference via Simplex PGD, then One-Hot Encoding} + +Holding \((A,\mathbf{B})\) fixed, coefficients are inferred by one forward +sweep from \(t=2\) to \(T\). At each time bin: +\begin{enumerate} + \item Build \(Z_t[m,:] = A\,Y_{t-1}+B_m\in\mathbb{R}^{M\times N}\). + \item Run \texttt{pgd\_iter} projected-gradient iterations: + \[ + \eta_t = c_t^\top Z_t,\quad + \tilde{\eta}_t = \min(\eta_t,\eta_{\mathrm{clip}}),\quad + \mu_t = \exp(\tilde{\eta}_t)\Delta t, + \] + \[ + g_t = Z_t\left(\mu_t-Y_t\right) + 2\lambda_2(c_t-c_{t-1}), + \qquad + c_t \leftarrow \Pi_{\Delta^{M-1}}(c_t-\alpha_E g_t). + \] + \item Accept \(c_t\), continue to next time bin. +\end{enumerate} +Projection uses the sort-based simplex projection (Duchi et al., \(O(M\log M)\)). +Bin \(t=0\) is not updated; the sweep starts at \(t=1\) (first lag pair +\((Y_0,Y_1)\)), matching the objective sum over \(t=2,\dots,T\) in +1-based notation. +In the distributed setting, each rank updates a disjoint time shard of +\(c_{1:T}\); shard results are averaged where multiple ranks overlap at +shard boundaries, then broadcast so all ranks share the same \(c\). + +After the forward sweep completes, the soft coefficients \(c_{1:T}\) are decoded into a hard state sequence via Viterbi (Section~\ref{sec:decode}). The M-step then receives one-hot encoded state assignments rather than the soft simplex probabilities: +\begin{equation} + \hat{S}_{1:T} = \mathrm{Viterbi}(c_{1:T}), \qquad + c_t^{\mathrm{(M\text{-}step)}} = e_{\hat{S}_t}, +\end{equation} +where \(e_m\) is the \(m\)-th standard basis vector. This ensures the M-step fits each state's parameters from cleanly separated data rather than from soft mixtures. + +\subsubsection{M-step: Dictionary Update via Adam} + +Holding \(c_t\) fixed (as one-hot vectors), update \(A,\mathbf{B}\) for \(K_M\) epochs using +Adam over shuffled mini-batches of \((Y_{t-1},Y_t,c_t)\). + +Per mini-batch: +\begin{equation} + \ell_{\mathrm{M}} = + \frac{1}{|\mathcal{B}|} + \sum_{t\in\mathcal{B}}\sum_{i=1}^{N} + \left[-Y_{t,i}\log(\mu_{t,i}+\varepsilon)+\mu_{t,i}\right] + + \lambda_3\,\overline{|A_{\mathrm{off}}|}, +\end{equation} +with \(\varepsilon=10^{-8}\). + +Then, subject to delay scheduling (operations are activated only after a configurable number of EM iterations have passed): +\begin{enumerate} + \item Off-diagonal soft-thresholding (activated after \(K_{\mathrm{prune}}\) EM iterations): + \[ + A_{ij}\leftarrow \operatorname{sign}(A_{ij}) + \left(|A_{ij}|-\alpha_M\lambda_3\right)_+,\quad i\neq j. + \] + \item Periodic spectral-radius projection (activated after \(K_{\rho}\) EM iterations, applied every \(R\) batches). On each application, in the multi-GPU setting, \(A\) and \(B\) are first averaged across ranks, then the spectral constraint is enforced, and the result is broadcast: + \[ + A \leftarrow A\cdot \min\!\left(1,\frac{\rho_{\max}}{\rho(A)}\right). + \] +\end{enumerate} + +Learning rate follows linear decay from \(\alpha_M^{(0)}\) to +\(\alpha_M^{(0)}f_{\mathrm{end}}\). The decay is activated only after \(K_{\mathrm{lr}}\) EM iterations have passed; the total number of decay steps spans the remaining M-step epochs. + +% --------------------------------------------------------------- +\subsection{State Decoding} +\label{sec:decode} + +Viterbi decoding is applied both within the EM loop (after each E-step, to produce one-hot assignments for the M-step) and after the final EM iteration (to produce the output state sequence). The decoder uses a homogeneous transition prior with +\(\tau_{\mathrm{dwell}}=\texttt{decode\_dwell\_sec}\): +\begin{equation} + p_{\mathrm{stay}}=\exp\!\left(-\frac{\Delta t}{\tau_{\mathrm{dwell}}}\right), + \qquad + p_{\mathrm{switch}}=\frac{1-p_{\mathrm{stay}}}{M-1}. +\end{equation} +The decoded sequence is +\begin{equation} + \hat{S}_{1:T}=\arg\max_{\{s_t\}} + \sum_t\left[\log c_{t,s_t}+\log p(s_t\mid s_{t-1})\right]. +\end{equation} +Per-bin confidence is +\[ +\mathrm{CL}_t = 1-\max_m c_{t,m}. +\] + +% --------------------------------------------------------------- +\section{Default Hyperparameters and Runtime Settings} +\label{sec:defaults} +% --------------------------------------------------------------- + +\begin{table}[h] +\centering +\begin{tabular}{lll} +\toprule +Group & Parameter (\texttt{CLI}) & Default \\ +\midrule +EM structure & \texttt{--num\_em\_iters} & 20 \\ + & \texttt{--m\_epochs} & 100 \\ +E-step & \texttt{--pgd\_iter} & 5 \\ + & \texttt{--lr\_estep} & 0.03 \\ + & \texttt{--lambda2} & 2.0 \\ +M-step & \texttt{--lr\_mstep} & \(3 \times 10^{-3}\) \\ + & \texttt{--lr\_end\_factor} & 0.1 \\ + & \texttt{--lambda3} & 0.02 \\ + & \texttt{--rho\_max} & 0.95 \\ + & \texttt{--prescale\_m\_step\_4\_ArhoMax} & 50 (batches) \\ + & \texttt{--delay\_em\_iter\_4\_ArhoMax} & 3 \\ + & \texttt{--delay\_em\_iter\_4\_lrDecay} & 1 \\ + & \texttt{--delay\_em\_iter\_4\_Aprune} & 1 \\ + & \texttt{--batch\_size} & 2048 \\ + & \texttt{--minW} & 0.01 (reporting only)\\ +Data window & \texttt{--time\_range\_sec} & [0.0, 60.0] s \\ +Decoding & \texttt{--decode\_dwell\_sec} & 0.3 s \\ +Init & \texttt{--init\_A} & data (or rand)\\ + & \texttt{--init\_B} & data (or rand) \\ + & \texttt{--init\_states} & data (or rand) \\ +\bottomrule +\end{tabular} +\caption{Current defaults in \texttt{prism\_EM\_train3.py}.} +\end{table} + +\paragraph{Required input.} +\texttt{--num\_states} (\texttt{-M}) is required (no default). +\texttt{--dataName} and \texttt{--basePath} locate input spikes under +\texttt{/spikesData/}; fits are written to +\texttt{/prismFit/.prismEM.npz}. +Wrapper script \texttt{fitPrismEM.sh} launches distributed training via +\texttt{torchrun} (default 4 GPUs) and sets +\texttt{--time\_range\_sec 0 80} unless overridden. + +\paragraph{Dataset-provided values.} +\(\Delta t\) (\texttt{time\_step\_sec}) and \(\eta_{\mathrm{clip}}\) +(\texttt{poisson\_eta\_clip}) are loaded from spike-file metadata. + +% --------------------------------------------------------------- +\section{Initialization of Model} +\label{sec:init} +% --------------------------------------------------------------- + +Initialization is controlled by: +\[ +\texttt{--init\_states},\quad \texttt{--init\_B},\quad \texttt{--init\_A} +\in \{\texttt{data},\texttt{rand}\}, +\] +with default \texttt{data} for all three. + +\subsection{State-Coefficient Initialization (\(C\equiv\{c_t\}\))} + +\paragraph{random mode:} +\[ +c_t = \frac{1}{M}\mathbf{1},\quad S_t=-1 \;\; \forall t. +\] +So coefficients start uniform over states. + +\paragraph{data mode:} +\begin{enumerate} + \item Aggregate spikes into coarse blocks of length + \(\lceil\tau_{\mathrm{dwell}}/\Delta t\rceil\) bins, where + \(\tau_{\mathrm{dwell}}=\texttt{decode\_dwell\_sec}\) (default 0.3\,s). + \item For each block, compute a scalar mean population firing rate. + \item Find \(M-1\) thresholds minimizing within-cluster variance + (dynamic-programming segmentation on unique rate values). + \item Assign each block to a state by thresholds. + \item Convert block states to soft probabilities: + dominant state gets \(0.8\), remainder \(0.2\) split across others, + with transition smoothing (\(2/3\) vs \(1/3\) at boundaries). + \item Map block-level \(c\) and \(S\) back to original bin resolution. +\end{enumerate} + +\subsection{Bias Initialization (\(B\))} + +\paragraph{random mode:} +No external seed is applied; model keeps random Gaussian initialization from +\texttt{torch.randn(...)*0.1}. + +\paragraph{data mode:} +For each neuron \(n\): +\[ +r_n = \frac{1}{T}\sum_t Y_{t,n}\,/\,\Delta t,\qquad +b_n = \log(\max(r_n,10^{-6})). +\] +The same \(b\)-vector is assigned to every state row of \(B\). + +\subsection{Connectivity Initialization (\(A\))} + +\paragraph{random mode:} +No external seed is applied; model keeps random Gaussian initialization from +\texttt{torch.randn(...)*0.1}. + +\paragraph{data mode (covariance/OLS):} +Using first \(T_{\max}=50{,}000\) bins (or fewer if shorter): +\[ +A_{\mathrm{ols}} += +\left(\sum_t Y_tY_{t-1}^\top\right) +\left(\sum_t Y_{t-1}Y_{t-1}^\top\right)^{\dagger}. +\] +Then all elements of \(A_{\mathrm{ols}}\) are scaled by a factor of 3: +\[ +A \leftarrow 3\,A_{\mathrm{ols}}. +\] +No spectral-radius rescaling is applied at initialization; the spectral constraint is enforced later during training via the delayed projection mechanism. + +Diagnostics (condition number, spectral radius, one-step \(R^2\), Frobenius norm, +bins used) are computed on the pre-scaled OLS estimate and saved in metadata +(\texttt{init\_A} sub-dictionary of the output metadata). + +% --------------------------------------------------------------- +\section{Output File and Evaluation} +\label{sec:output} +% --------------------------------------------------------------- + +After training, rank~0 writes \texttt{/prismFit/.prismEM.npz}. +If \texttt{--fitName} is omitted, the name is +\texttt{-EM-<6hex>}. +Key arrays (shapes for \(T\) time bins, \(N\) neurons, \(M\) states): +\begin{itemize} + \item \texttt{A\_init}, \texttt{A\_hat}: \((N,N)\); + \item \texttt{B\_init}: \((N,)\) --- first row of the initialized \(B\) matrix + (common across states in \texttt{data} mode); + \item \texttt{B\_hat}: \((M,N)\) --- full per-state bias matrix; + \item \texttt{c\_init}, \texttt{c\_hat}: \((T,M)\); + \item \texttt{S\_init}, \texttt{S\_hat}: \((T,)\) integer state labels + (\texttt{S\_init} is \(-1\) in \texttt{rand} init mode); + \item \texttt{S\_hat\_CL}: \((T,)\) per-bin confidence + \(\mathrm{CL}_t = 1-\max_m c_{t,m}\); + \item \texttt{e\_nll\_em}, \texttt{m\_loss\_epoch}, \texttt{m\_nll\_epoch}, + \texttt{m\_l1\_epoch}, \texttt{rho\_epoch}, \texttt{nz\_edges\_epoch}, + \texttt{learning\_rates}: training history. +\end{itemize} +Metadata records all CLI hyperparameters, time-window bins, initialization +dictionaries, and \texttt{mstep\_state\_mode=viterbi\_onehot}. +Suggested evaluation: +\begin{verbatim} +./prism_EM_eval3.py --basePath $basePath --dataName -p a e f g +\end{verbatim} + +\end{document} diff --git a/causal_net/nonStation_ver3a_states_simu/docs/genDale_topo_states_ver3.pdf b/causal_net/nonStation_ver3a_states_simu/docs/genDale_topo_states_ver3.pdf new file mode 100644 index 00000000..14c1a50a Binary files /dev/null and b/causal_net/nonStation_ver3a_states_simu/docs/genDale_topo_states_ver3.pdf differ diff --git a/causal_net/nonStation_ver3a_states_simu/docs/genDale_topo_states_ver3.tex b/causal_net/nonStation_ver3a_states_simu/docs/genDale_topo_states_ver3.tex new file mode 100644 index 00000000..971e6fe8 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/docs/genDale_topo_states_ver3.tex @@ -0,0 +1,681 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} +\usepackage{multirow} +\usepackage{algorithm} +\usepackage{algorithmic} + +\geometry{margin=0.8in} + +\title{Synthetic Non-Stationary Spiking Data Generation:\\ +Dale-Law-Constrained Recurrent Networks with\\ +a Switching Poisson GLM} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +We describe the two-script pipeline that produces synthetic non-stationary +neural population spike trains for benchmarking latent-state inference +algorithms (e.g.\ PRISM-EM). +\textbf{Stage~1} (\texttt{gen\_daleMatrices3.py}) constructs a single +Dale-law-compliant weight matrix $A_{\text{true}}$ with a prescribed +spectral radius $R<1$ and $M$ bias vectors $\{B_m\}_{m=0}^{M-1}$, +one per dynamical state. +Stationary spike trains under each bias are generated as a +firing-rate sanity check. +\textbf{Stage~2} (\texttt{gen\_nonStationarySpikes3.py}) loads the +ground-truth dictionary, schedules transitions among the $M$ states, +tracks a smoothly evolving simplex coefficient $c_t$, and generates +the non-stationary spike train from a Poisson GLM with a +time-varying effective bias $B_{\mathrm{eff},t} = \sum_m c_{t,m} B_m$. +An oracle state decoder (exact Poisson log-likelihood) is also +computed and saved as an upper bound for any inference algorithm. +\end{abstract} + +\tableofcontents +\newpage + +% =============================================================== +\section{Overview} +\label{sec:overview} +% =============================================================== + +The pipeline produces ground-truth data for a network of $N$ neurons +divided into $N_E$ excitatory and $N_I = N - N_E$ inhibitory cells. +The network obeys Dale's principle: each neuron is either purely +excitatory or purely inhibitory. +All neurons share one connectivity matrix $A_{\text{true}}$; what +differs across $M$ states is only the bias (log-baseline firing-rate) +vector $B_m$. + +\paragraph{Stage 1 (Section~\ref{sec:dale_matrices}).} +\texttt{gen\_daleMatrices3.py} builds and saves: +\begin{itemize} + \item a sparse binary connectivity mask $E_{\text{true}} \in + \{0,1\}^{N \times N}$; + \item a single weight matrix $A_{\text{true}}$ with + $\rho(A_{\text{true}}) = R$; + \item $M$ bias vectors $B_{\text{true}} \in \mathbb{R}^{M \times N}$, + one per state. +\end{itemize} +Stationary spikes under each $B_m$ are simulated and their +firing-rate statistics are reported. + +\paragraph{Stage 2 (Section~\ref{sec:nonstat_spikes}).} +\texttt{gen\_nonStationarySpikes3.py} loads the Stage-1 output and: +\begin{itemize} + \item draws a target-state sequence $S_{\text{true}}[t]$ + via a Markov chain or round-robin schedule; + \item maintains a simplex coefficient vector $c_t \in \Delta^{M-1}$ + that blends the $M$ bias vectors smoothly; + \item samples $Y_t \sim \mathrm{Poisson}$ from the resulting + time-varying Poisson GLM; + \item computes the oracle state $S_{\text{oracle}}[t]$ by + maximum Poisson log-likelihood. +\end{itemize} + +\paragraph{Notation.} +$N$: total neurons; $N_E$: excitatory count (default $\lfloor 0.8N \rfloor$); +$N_I = N - N_E$: inhibitory count. +$M$: number of states (= \texttt{len(Boffsets)}). +$\Delta t$: time-bin width in seconds (default $0.01$\,s). +$T$: number of time bins. +$R$: target spectral radius (default $0.3$). +$\eta_{\mathrm{clip}} = 5$: overflow-prevention threshold +($\approx [10^{-3}, 10^3]$\,Hz effective rate range). +Column convention: $(A_{\text{true}})_{ij}$ is the synaptic weight +\emph{from} presynaptic neuron $j$ \emph{to} postsynaptic neuron $i$. + +% =============================================================== +\section{Stage 1: Dale Matrix Generation} +\label{sec:dale_matrices} +% =============================================================== + +\subsection{Command-Line Parameters} + +\begin{table}[h] +\centering +\begin{tabular}{lll} +\toprule +CLI flag & Default & Description \\ +\midrule +\texttt{--num\_neurons} $N$ & 50 & Total neuron count \\ +\texttt{--num\_excite} $N_E$ & $\lfloor 0.8N \rfloor$ & + Excitatory neuron count \\ +\texttt{--edge\_prob} $[p_\ell, p_h]$ & $[0.05,\, 0.20]$ & + Per-neuron connectivity probability range \\ +\texttt{--spectral\_radius} $R$ & 0.3 & + Target spectral radius \\ +\texttt{--idleRate} $[f_{\min}, f_{\max}]$ & $[15,\, 30]$\,Hz & + Idle firing-rate band \\ +\texttt{--Boffsets} $\{\delta_m\}$ & $[0.0]$ & + Per-state log-rate offsets (one bias vector per entry) \\ +\texttt{--num\_steps} $T$ & 10\,001 & + Steps for the stationary sanity-check simulation \\ +\texttt{--step\_size} $\Delta t$ & 0.01\,s & Time-bin width \\ +\texttt{--basePath} & \emph{(scratch dir)} & + Root output directory \\ +\bottomrule +\end{tabular} +\caption{Key command-line arguments of \texttt{gen\_daleMatrices3.py}.} +\label{tab:dale_args} +\end{table} + +\subsection{Algorithm} + +\begin{algorithm}[H] +\caption{Dale Matrix Generation (\texttt{gen\_daleMatrices3.py})} +\begin{algorithmic}[1] +\STATE \textbf{Step 1:} Generate sparse binary connectivity mask + $E_{\text{true}} \in \{0,1\}^{N\times N}$ + (\texttt{generate\_sparse\_mask}) +\STATE \textbf{Step 2:} Generate weight matrix $A_{\text{true}}$ with + Dale E/I balance, negative self-loops, and spectral radius $R$ + (\texttt{init\_W}) +\STATE \textbf{Step 3:} For each offset $\delta_m$ in \texttt{Boffsets}: +\begin{ALC@g} + \STATE Generate bias vector $B_m$ (\texttt{set\_flat\_selfSpiking}) + \STATE Simulate stationary spikes $Y^{(m)}$ + (\texttt{gen\_stationary\_lag1\_poisson}) + \STATE Compute firing-rate statistics (\texttt{estimate\_rates}) +\end{ALC@g} +\STATE \textbf{Step 4:} Save ground-truth and spike files +\end{algorithmic} +\end{algorithm} + +\subsubsection{Step 1: Sparse Connectivity Mask} +\label{sec:mask} + +Each neuron $i$ independently draws its connection probability: +\begin{equation} + q_i \;\sim\; \mathrm{Uniform}(p_\ell,\; p_h). +\end{equation} +The binary mask entry is +\begin{equation} + (E_{\text{true}})_{ij} = + \begin{cases} + 1 & \text{if } u_{ij} < q_i, \quad + u_{ij} \sim \mathrm{Uniform}(0,1) \text{ i.i.d.,} \\ + 0 & \text{otherwise.} + \end{cases} +\end{equation} +Self-loops are \emph{always} included: $(E_{\text{true}})_{ii} = 1$ +for all $i$, encoding the mandatory negative self-connection added below. + +\subsubsection{Step 2: Weight Matrix with Dale's Principle} +\label{sec:weights} + +\paragraph{E/I balance ratio.} +Let $r = N_E / N_I$. +This ratio is used to set the mean inhibitory weight magnitude so +that the expected signed row-sum of the full matrix is +approximately zero. + +\paragraph{Raw weight assignment (row-based, weight spread $v = 0.2$).} +\begin{equation} + W_{ij} = + \begin{cases} + \mathrm{Uniform}(1-v,\; 1+v) & + i \in \{0,\dots,N_E-1\} \quad (\text{excitatory rows}), \\[4pt] + \mathrm{Uniform}(-r(1+v),\; -r(1-v)) & + i \in \{N_E,\dots,N-1\} \quad (\text{inhibitory rows}). + \end{cases} + \label{eq:raw_weights} +\end{equation} +Because the \emph{row} index $i$ determines the neuron's type, every +outgoing weight of an excitatory neuron is positive and every +outgoing weight of an inhibitory neuron is negative. This is +Dale's principle. + +\paragraph{Applying the connectivity mask.} +\begin{equation} + W \;\leftarrow\; W \odot E_{\text{true}}. +\end{equation} + +\paragraph{Enforcing negative self-loops.} +Any positive diagonal entry (excitatory self-loop) is sign-flipped: +\begin{equation} + W_{ii} \;\leftarrow\; -|W_{ii}| \quad \forall i, +\end{equation} +so that all self-connections are inhibitory (stabilising). + +\paragraph{Spectral normalization.} +Let $\rho_0 = \max_k |\lambda_k(W)|$ be the current spectral radius. +The connectivity matrix is set to +\begin{equation} + A_{\text{true}} \;\leftarrow\; \frac{R}{\rho_0}\, W, + \label{eq:spectral} +\end{equation} +so $\rho(A_{\text{true}}) = R$ exactly. +Rescaling preserves all signs, the zero pattern, Dale's law, and +the negative self-loops. +For $R < 1$ the operator $A_{\text{true}}$ is contracting, which is +a sufficient condition for the stationary Poisson VAR(1) to have a +well-defined stationary distribution. + +\subsubsection{Step 3: Bias Vectors and Stationary Spike Check} +\label{sec:bias} + +\paragraph{Bias generation (\texttt{set\_flat\_selfSpiking}).} +A size-scaling factor $s = \sqrt{N/100}$ adjusts the excitatory +correction for network size. +For each state offset $\delta_m$ (index $m = 0,\dots,M-1$), +independent draws +\begin{equation} + B_m^{(i)} \;\sim\; + \mathrm{Uniform}\!\bigl(\ln(R\,f_{\min}),\;\ln(R\,f_{\max})\bigr), + \quad i = 0,\dots,N-1, +\end{equation} +followed by state- and type-dependent shifts: +\begin{equation} + B_m^{(i)} \;\leftarrow\; + \begin{cases} + B_m^{(i)} + 1 + \delta_m - 1.5R - s + & i < N_E \quad (\text{excitatory}), \\[4pt] + B_m^{(i)} + \delta_m + & i \geq N_E \quad (\text{inhibitory}). + \end{cases} + \label{eq:bias} +\end{equation} +The excitatory correction $-(1.5R + s - 1)$ lowers the baseline +excitatory rate relative to inhibitory neurons, preventing runaway +excitation. +Larger $\delta_m$ uniformly raises activity across the network, +creating distinct mean-rate states. + +\paragraph{Stationary spike simulation (evaluation only).} +For each bias vector $B_m$ a lag-1 Poisson VAR(1) process is +simulated over $T$ bins: +\begin{equation} + \eta_t = A_{\text{true}}\,Y_{t-1} + B_m, + \qquad + \tilde{\eta}_t = \min(\eta_t,\;\eta_{\mathrm{clip}}), + \qquad + \lambda_t = \exp(\tilde{\eta}_t)\,\Delta t, + \qquad + Y_t \;\sim\; \mathrm{Poisson}(\lambda_t), + \label{eq:stationary} +\end{equation} +with initial condition $Y_0 \sim \mathrm{Poisson}(\exp(B_m)\,\Delta t)$. +These stationary spike arrays (shape $(T, N)$, one per bias vector) +are used \emph{only} to verify that the firing rates are in a +plausible range; they are \emph{not} the final non-stationary output. + +Firing-rate statistics computed via \texttt{estimate\_rates} +(\texttt{UtilDalePoisson.py}) include: +\begin{itemize} + \item mean per-neuron firing rates (Hz), + \item rate variance over non-overlapping $T_{\mathrm{var}} = 5$\,s windows (Hz$^2$), + \item Fano factors ($\mathrm{Var}[\text{count}] / \mathrm{Mean}[\text{count}]$), + \item consecutive cross-neuron coincidence rate (Hz per neuron). +\end{itemize} + +\subsection{Output Files (Stage 1)} + +All files are written to \texttt{/truthDale/}. + +\begin{table}[h] +\centering +\begin{tabular}{llll} +\toprule +File & Array key & Shape & Description \\ +\midrule +\multirow{3}{*}{\texttt{.simTruth.npz}} + & \texttt{A\_true} & $(N,N)$ float & + Shared weight matrix, $\rho(A_{\text{true}})=R$ \\ + & \texttt{B\_true} & $(M,N)$ float & + Per-state bias vectors \\ + & \texttt{E\_true} & $(N,N)$ int & + Binary connectivity mask \\ +\midrule +\multirow{4}{*}{\texttt{.spikes.npz}} + & \texttt{spikes} & $(M,T,N)$ uint8 & + Stationary spikes per state (sanity check) \\ + & \texttt{single\_rates} & $(M,N)$ float & + Mean firing rates (Hz) per state \\ + & \texttt{sigle\_rates\_var} & $(M,N)$ float & + Rate variance (Hz$^2$) per state \\ + & \texttt{single\_fano\_fact} & $(M,N)$ float & + Fano factors per state \\ +\bottomrule +\end{tabular} +\caption{Output arrays of Stage~1. + $M = \texttt{len(Boffsets)}$; $T = \texttt{num\_steps}$.} +\label{tab:stage1_output} +\end{table} + +The metadata dictionary \texttt{dale\_conf} records $N$, $N_E$, $R$, +edge-probability range, idle-rate range, and offset list. +The metadata dictionary \texttt{evol\_conf} records $T$, $\Delta t$, +total simulated time, and $\eta_{\mathrm{clip}}$. + +% =============================================================== +\section{Stage 2: Non-Stationary Spike Generation} +\label{sec:nonstat_spikes} +% =============================================================== + +\subsection{Command-Line Parameters} + +\begin{table}[h] +\centering +\begin{tabular}{lll} +\toprule +CLI flag & Default & Description \\ +\midrule +\texttt{--inputStates} & --- & + Base name of the \texttt{.simTruth.npz} file (required) \\ +\texttt{--num\_steps} $T$ & from \texttt{evol\_conf} & + Number of time bins \\ +\texttt{--true\_dwell\_sec} $t_d$ & 1.0\,s & + Mean dwell time per state \\ +\texttt{--state\_change\_speed} $\nu$ & 0.33 & + Max simplex-coefficient change per bin \\ +\texttt{--schedule} & \texttt{mc} & + State schedule: Markov chain (\texttt{mc}) or round-robin (\texttt{rr}) \\ +\texttt{--seed} & 42 & Random seed \\ +\bottomrule +\end{tabular} +\caption{Key command-line arguments of \texttt{gen\_nonStationarySpikes3.py}.} +\label{tab:nonstat_args} +\end{table} + +\subsection{Algorithm} + +\begin{algorithm}[H] +\caption{Non-Stationary Spike Generation + (\texttt{gen\_nonStationarySpikes3.py})} +\begin{algorithmic}[1] +\STATE Load $(A_{\text{true}}, B_{\text{true}})$ from + \texttt{.simTruth.npz} +\STATE Build target-state sequence $S_{\text{true}}[t]$ + via selected schedule (Section~\ref{sec:scheduling}) +\STATE Simulate switching spikes and record simplex coefficients $C_{\text{true}}$ + (\texttt{simulate\_switching\_poisson}, Section~\ref{sec:switching}) +\STATE Compute oracle state sequence $S_{\text{oracle}}$ + (\texttt{compute\_oracle\_states\_comA}, Section~\ref{sec:oracle}) +\STATE Compute firing-rate statistics (\texttt{estimate\_rates}) +\STATE Save spike and ground-truth files +\end{algorithmic} +\end{algorithm} + +\subsection{State Scheduling} +\label{sec:scheduling} + +The mean dwell time in bins is +$\tau_d = \lceil t_d / \Delta t \rceil$. + +\paragraph{Markov chain (\texttt{mc}, default).} +A discrete-time Markov chain with geometric dwell times. +The exit probability per bin is $p_{\mathrm{exit}} = 1/\tau_d$; +on exit the chain transitions uniformly to one of the $M-1$ other states. +The symmetric $M \times M$ transition matrix stored in the output is: +\begin{equation} + T_{sm} = + \begin{cases} + 1 - p_{\mathrm{exit}} & s = m \quad (\text{self-transition}), \\ + p_{\mathrm{exit}}/(M-1) & s \neq m \quad (\text{cross-transition}). + \end{cases} + \label{eq:trans} +\end{equation} +A minimum dwell of +$\tau_{\min} = \lceil 0.3\,\tau_d \rceil$ bins is enforced +to prevent spuriously short visits. +State coverage across the full record is uncontrolled and subject +to random fluctuations. + +\paragraph{Round-robin (\texttt{rr}).} +States cycle deterministically as $0, 1, \dots, M-1, 0, 1, \dots$ +Each visit lasts exactly $\tau_d$ bins (the last visit in the +record is trimmed to fit $T$). +This guarantees that all states receive equal coverage +($T/M \pm \tau_d$ bins each), eliminating sampling variability. + +\subsection{Smooth Switching Poisson Simulation} +\label{sec:switching} + +\paragraph{Simplex coefficient vector.} +A vector $c_t \in \Delta^{M-1}$ +(non-negative entries summing to 1) tracks the current mixture of +states. +The one-hot target for the scheduled state at bin $t$ is +$e_m = \mathbf{1}[m = S_{\text{true}}[t]]$. +At each bin, $c_t$ is updated by a clipped gradient step and +renormalized: +\begin{align} + c_t &\;\leftarrow\; c_{t-1} + + \mathrm{clip}\!\bigl(e - c_{t-1},\; -\nu,\; \nu\bigr), \\ + c_t &\;\leftarrow\; c_t \;/\; \|c_t\|_1. + \label{eq:smooth_update} +\end{align} +The speed parameter $\nu \in (0,1]$ controls how quickly the mixture +tracks the target: $\nu = 1$ gives instantaneous switching; +small $\nu$ produces gradual multi-bin transitions. +The transition duration is approximately $1/\nu$ bins. + +\paragraph{Effective bias.} +The convex mixture of per-state bias vectors is +\begin{equation} + B_{\mathrm{eff},t} + = \sum_{m=0}^{M-1} c_{t,m}\, B_m + = \sum_{m=0}^{M-1} c_{t,m}\, B_{\text{true}}[m,:]. + \label{eq:Beff} +\end{equation} +This interpolates smoothly between the per-state firing-rate +baselines as the state label switches. + +\paragraph{Non-stationary spike generation.} +At every bin $t = 0, \dots, T-1$: +\begin{equation} + \eta_t = A_{\text{true}}\,Y_{t-1} + B_{\mathrm{eff},t}, + \qquad + \tilde{\eta}_t = \min(\eta_t,\;\eta_{\mathrm{clip}}), + \qquad + \lambda_t = \exp(\tilde{\eta}_t)\,\Delta t, + \qquad + Y_t \;\sim\; \mathrm{Poisson}(\lambda_t). + \label{eq:nonstat} +\end{equation} +This is the same Poisson VAR(1) model as in Eq.~\eqref{eq:stationary}, +except that the bias $B_{\mathrm{eff},t}$ is now time-varying. +The connectivity matrix $A_{\text{true}}$ is unchanged across all states. + +\subsection{Oracle State Decoder} +\label{sec:oracle} + +An oracle decoder assigns the most likely state to each bin using +the \emph{exact} ground-truth parameters $(A_{\text{true}}, \{B_m\})$: +\begin{equation} + S_{\text{oracle}}[t] + = \arg\max_{m \in \{0,\dots,M-1\}}\; + \sum_{j=0}^{N-1} + \Bigl[ + Y_{t,j}\,\tilde{\eta}_{t,j}^{(m)} + - \exp\!\bigl(\tilde{\eta}_{t,j}^{(m)}\bigr)\,\Delta t + \Bigr], + \label{eq:oracle} +\end{equation} +where +$\eta_{t,j}^{(m)} = \bigl(A_{\text{true}}\,Y_{t-1}\bigr)_j + B_{m,j}$ +and $\tilde{\eta}^{(m)} = \min(\eta^{(m)}, \eta_{\mathrm{clip}})$. +This is the exact Poisson log-likelihood score for state $m$ given +the observed spike vector $Y_t$, computed independently at each bin. +The oracle accuracy is an upper bound for any latent-state inference +algorithm. + +Per-state oracle scores and overall accuracy are printed at the end of +the run and stored in the \texttt{oracle\_eval} sub-dictionary of the +metadata. + +\subsection{Output Files (Stage 2)} + +All files are written to \texttt{/spikesData/}. + +\begin{table}[h] +\centering +\begin{tabular}{llll} +\toprule +File & Array key & Shape & Description \\ +\midrule +\multirow{2}{*}{\texttt{.spikes.npz}} + & \texttt{spikes} & $(T,N)$ int32 & + Non-stationary spike counts \\ + & \texttt{single\_rates} & $(N,)$ float & + Mean firing rates over the full record (Hz) \\ +\midrule +\multirow{7}{*}{\texttt{.prismTruth.npz}} + & \texttt{A\_true} & $(N,N)$ float & + Shared weight matrix (copy from Stage 1) \\ + & \texttt{B\_true} & $(M,N)$ float & + Per-state bias vectors (copy from Stage 1) \\ + & \texttt{S\_true} & $(T,)$ int32 & + Target state index at each bin \\ + & \texttt{C\_true} & $(T,M)$ float32 & + Simplex coefficient vector at each bin \\ + & \texttt{S\_oracle} & $(T,)$ int32 & + Oracle (max-likelihood) state at each bin \\ + & \texttt{state\_transition} & $(M,M)$ float32 & + Transition matrix used/implied by schedule \\ + & \texttt{sigle\_rates\_var} & $(N,)$ float & + Per-neuron rate variance (Hz$^2$) \\ + & \texttt{single\_fano\_fact} & $(N,)$ float & + Per-neuron Fano factors \\ +\bottomrule +\end{tabular} +\caption{Output arrays of Stage~2. + $T = \texttt{num\_steps}$; $M = \texttt{num\_states}$.} +\label{tab:stage2_output} +\end{table} + +The metadata dictionary \texttt{evol\_conf} records $T$, $\Delta t$, +total time, $\nu$, requested and effective dwell times (seconds and bins), +$M$, random seed, and schedule type. +The \texttt{oracle\_eval} sub-dictionary records the overall oracle +accuracy and per-state breakdown. + +% =============================================================== +\section{Notation Summary} +\label{sec:notation} +% =============================================================== + +\begin{table}[h] +\centering +\renewcommand{\arraystretch}{1.2} +\begin{tabular}{lll} +\toprule +Symbol & Domain & Description \\ +\midrule +$N$ & $\mathbb{Z}_{>0}$ & Total number of neurons \\ +$N_E$ & $\mathbb{Z}_{>0}$ & Excitatory neuron count (rows $0\dots N_E-1$) \\ +$N_I = N-N_E$ & $\mathbb{Z}_{>0}$ & Inhibitory neuron count \\ +$M$ & $\mathbb{Z}_{>0}$ & + Number of states ($= \texttt{len(Boffsets)}$) \\ +$T$ & $\mathbb{Z}_{>0}$ & Number of time bins \\ +$\Delta t$ & $\mathbb{R}_{>0}$ & Time-bin width (s); default $0.01$ \\ +$R$ & $(0,1)$ & Target spectral radius \\ +$r = N_E/N_I$ & $\mathbb{R}_{>0}$ & E/I balance ratio \\ +$v = 0.2$ & --- & Weight spread (hard-coded) \\ +$E_{\text{true}}$ & $\{0,1\}^{N\times N}$ & Binary connectivity mask \\ +$A_{\text{true}}$ & $\mathbb{R}^{N\times N}$ & + Shared weight matrix, $\rho(A_{\text{true}}) = R$ \\ +$B_m$ & $\mathbb{R}^N$ & Bias (log-baseline rate) vector for state $m$ \\ +$B_{\text{true}}$ & $\mathbb{R}^{M\times N}$ & Stacked per-state biases \\ +$Y_t$ & $\mathbb{Z}_{\geq 0}^N$ & Spike-count vector at bin $t$ \\ +$\eta_{\mathrm{clip}}$ & $\mathbb{R}_{>0}$ & + Rate clipping threshold; fixed at $5$ \\ +\midrule +$S_{\text{true}}[t]$ & $\{0,\dots,M-1\}$ & + Scheduled target state at bin $t$ \\ +$c_t$ & $\Delta^{M-1}$ & + Simplex coefficient vector at bin $t$ \\ +$\nu$ & $(0,1]$ & + Max coefficient change per bin (\texttt{state\_change\_speed}) \\ +$B_{\mathrm{eff},t}$ & $\mathbb{R}^N$ & + Effective bias: convex mixture $\sum_m c_{t,m} B_m$ \\ +$\tau_d$ & $\mathbb{Z}_{>0}$ & + Mean dwell time in bins \\ +$S_{\text{oracle}}[t]$ & $\{0,\dots,M-1\}$ & + Oracle max-likelihood state at bin $t$ \\ +\bottomrule +\end{tabular} +\caption{Notation used throughout this document.} +\label{tab:notation} +\end{table} + +% =============================================================== +\section{Full Pipeline and Suggested Commands} +\label{sec:pipeline} +% =============================================================== + +\begin{enumerate} + +\item \textbf{Generate ground-truth matrices} (Stage 1): +\begin{verbatim} +./gen_daleMatrices3.py --num_neurons 100 --num_excite 70 \ + --spectral_radius 0.5 --Boffsets 0 10 20 --dataName myData \ + --basePath $basePath +\end{verbatim} +Produces \texttt{myData.simTruth.npz} ($A_{\text{true}}$, +$B_{\text{true}}$, $E_{\text{true}}$; $M=3$ states here) and +\texttt{myData.spikes.npz} (stationary sanity-check arrays, +shape $(3, T, N)$). + +\item \textbf{Inspect matrices and stationary spikes}: +\begin{verbatim} +./view_daleMatrix3.py --basePath $basePath --dataName myData \ + -p b -m 0 -X -p a c d +./view_spikesTrain3.py --basePath $basePath --dataName myData \ + --time_range_sec 0 20 -p b -m 0 -X +\end{verbatim} + +\item \textbf{Generate non-stationary spike train} (Stage 2): +\begin{verbatim} +./gen_nonStationarySpikes3.py --basePath $basePath \ + --inputStates myData --true_dwell_sec 1.0 +\end{verbatim} +Produces \texttt{.spikes.npz} (shape $(T, N)$) and +\texttt{.prismTruth.npz} ($S_{\text{true}}$, +$C_{\text{true}}$, $S_{\text{oracle}}$, transition matrix, and +copies of $A_{\text{true}}$, $B_{\text{true}}$). + +\item \textbf{Inspect non-stationary spikes}: +\begin{verbatim} +./view_spikesTrain3.py --basePath $basePath \ + --dataName --idxState -1 \ + --time_range_sec 0 15 -p b +\end{verbatim} + +\item \textbf{Fit with PRISM-EM}: +\begin{verbatim} +./fitPrismEM.sh --basePath $basePath --dataName \ + --num_states 3 --num_em_iters 4 --m_epochs 16 \ + --time_range_sec 0 80 +\end{verbatim} + +\item \textbf{Fit with Lasso-Poisson}: +\begin{verbatim} +./fit_lassoPoisson3.py --basePath $basePath \ + --dataName --num_epochs 300 +\end{verbatim} + +\end{enumerate} + +% =============================================================== +\section{Key Design Choices} +\label{sec:design} +% =============================================================== + +\paragraph{Shared $A_{\text{true}}$ across all states.} +The synaptic weight matrix is identical in every state; only the +bias vector $B_m$ differs between states. +This models a circuit whose connectivity is fixed but whose +input drive (e.g.\ neuromodulatory tone, sensory stimulus level) +shifts the population-wide activity baseline between states. +The inference task is therefore to recover both $A_{\text{true}}$ +and $\{B_m\}$ from a single non-stationary spike record. + +\paragraph{Row convention for Dale's principle.} +Row $i$ of $A_{\text{true}}$ contains the outgoing weights of +neuron $i$. +Excitatory neurons (rows $0 \dots N_E - 1$) have all non-zero +weights $> 0$; inhibitory neurons (rows $N_E \dots N-1$) have +all non-zero weights $< 0$. +This guarantees that neuron $i$ can only excite \emph{or} only +inhibit all of its postsynaptic partners, never both. + +\paragraph{Spectral radius and stability.} +$\rho(A_{\text{true}}) = R < 1$ is a sufficient condition for +the stationary Poisson VAR(1) to admit a well-defined stationary +distribution. +In the non-stationary regime, stability is maintained instantaneously +because $A_{\text{true}}$ is unchanged; only the bias shifts. + +\paragraph{Smooth simplex interpolation.} +The clipped gradient update (Eq.~\eqref{eq:smooth_update}) prevents +abrupt discontinuities in $B_{\mathrm{eff},t}$ at state-boundary +bins. +This mimics a biological circuit that cannot switch network state +instantaneously. +The mixing speed $\nu$ is independent of the mean dwell time and +can be set to match a desired biophysical transition timescale. + +\paragraph{Oracle score as a decoder upper bound.} +The oracle uses exact ground-truth parameters and per-bin +maximum-likelihood, achieving the best possible state assignment +from the spike observations. +Any EM or variational algorithm must approach, but cannot exceed, +this score. +The oracle score therefore quantifies how much information about +the latent state is present in the spikes at all, distinguishing +fundamental ambiguity from algorithm suboptimality. + +\end{document} diff --git a/causal_net/nonStation_ver3a_states_simu/eval_fitLasso.py b/causal_net/nonStation_ver3a_states_simu/eval_fitLasso.py new file mode 100755 index 00000000..2881b8cb --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/eval_fitLasso.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +""" +Evaluation and visualization tool for LASSO Poisson model fitting results. + +This script loads fitted LASSO connectivity matrices from Poisson GLM training +and provides comprehensive evaluation through statistical analysis and plotting. +Main functionality includes: +- Loading fitted model parameters (A, B matrices) from .lassoFit.npz files +- Computing network connectivity statistics and sparsity metrics +- Generating various visualizations (structure plots, distributions, reconstructions) +- Optionally comparing against ground truth for simulated data + +Usage: + ./eval_fitLasso.py --dataName mydata --basePath /path/to/data/ -p ab +""" + +import numpy as np +import os +import argparse +import sys +from toolbox.Util_NumpyIO import read_data_npz +from PlotterLassoFitEval import Plotter + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + +from pprint import pprint + + +def get_offdiag_triplets(A, is_pos=True, min_w=None): + """Return off-diagonal [i, j, value] triplets filtered by sign and abs(value)>=min_w.""" + A = np.asarray(A) + assert A.ndim == 2 and A.shape[0] == A.shape[1], f"Expected square 2D matrix, got {A.shape}" + assert min_w is not None, "min_w must be provided" + offdiag = ~np.eye(A.shape[0], dtype=bool) + sign_mask = (A > 0) if is_pos else (A < 0) + mag_mask = np.abs(A) >= float(min_w) + mask = offdiag & sign_mask & mag_mask + i_idx, j_idx = np.where(mask) + vals = A[i_idx, j_idx] + return np.column_stack([i_idx, j_idx, vals]) + + +def compare_triplets(rec, truth): + """Return [TP, FP, FN] for triplets; TP carries both reconstructed and true values.""" + rec_dict = {(int(r[0]), int(r[1])): r[2] for r in rec} + tru_dict = {(int(r[0]), int(r[1])): r[2] for r in truth} + + rec_idx = set(rec_dict.keys()) + tru_idx = set(tru_dict.keys()) + + tp_idx = rec_idx & tru_idx + fp_idx = rec_idx - tru_idx + fn_idx = tru_idx - rec_idx + + tp = np.array([[i, j, rec_dict[(i, j)], tru_dict[(i, j)]] for i, j in sorted(tp_idx)]) + fp = np.array([[i, j, rec_dict[(i, j)]] for i, j in sorted(fp_idx)]) + fn = np.array([[i, j, tru_dict[(i, j)]] for i, j in sorted(fn_idx)]) + + if len(tp) == 0: + tp = np.empty((0, 4)) + if len(fp) == 0: + fp = np.empty((0, 3)) + if len(fn) == 0: + fn = np.empty((0, 3)) + + return [tp, fp, fn] + + +def build_residual_eval_data(fitD, trueD, minW, verb=1): + """Build evalD expected by PlotterLassoFitEval.residuals().""" + assert minW is not None, "minW must be provided" + A_fit = np.asarray(fitD['A_lasso']) + A_true = np.asarray(trueD['A_true']) + if A_true.ndim == 3: + A_true = A_true[0] + + assert A_fit.ndim == 2 and A_true.ndim == 2, "A_lasso/A_true must be 2D matrices" + if A_fit.shape != A_true.shape: + n = min(A_fit.shape[0], A_true.shape[0]) + if verb: + print(f"WARNING: A shape mismatch fit={A_fit.shape}, true={A_true.shape}; using top-left {n}x{n}.") + A_fit = A_fit[:n, :n] + A_true = A_true[:n, :n] + + EposT = get_offdiag_triplets(A_true, True, min_w=minW) + EnegT = get_offdiag_triplets(A_true, False, min_w=minW) + EposR = get_offdiag_triplets(A_fit, True, min_w=minW) + EnegR = get_offdiag_triplets(A_fit, False, min_w=minW) + + if verb: + print(f"Residual eval edges (|w|>={minW:g}): true pos/neg={EposT.shape[0]}/{EnegT.shape[0]}, fit pos/neg={EposR.shape[0]}/{EnegR.shape[0]}") + + evalD = { + 'pos': compare_triplets(EposR, EposT), + 'neg': compare_triplets(EnegR, EnegT), + 'diag': np.column_stack([np.diag(A_fit), np.diag(A_true)]), + 'minW': float(minW), + } + + B_fit = fitD.get('B_lasso') + B_true = trueD.get('B_true') + if B_fit is not None and B_true is not None: + B_fit = np.asarray(B_fit).reshape(-1) + B_true = np.asarray(B_true) + if B_true.ndim == 2: + if verb and B_true.shape[0] > 1: + print(f"INFO: B_true has {B_true.shape[0]} rows; using row 0 for residual plot.") + B_true = B_true[0] + B_true = np.asarray(B_true).reshape(-1) + nB = min(B_fit.size, B_true.size) + if B_fit.size != B_true.size and verb: + print(f"WARNING: B length mismatch fit={B_fit.size}, true={B_true.size}; using first {nB}.") + evalD['bterm'] = np.column_stack([B_fit[:nB], B_true[:nB]]) + + return evalD + +######################### +# MAIN +######################### + +def main(): + parser = argparse.ArgumentParser(description="Evaluate and plot results from fit_poisson.py") + parser.add_argument("--dataName", type=str, default='dale_M120_3M', help="Base name for the dataset") + parser.add_argument("--basePath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input/output data") + parser.add_argument('-p',"--showPlots", type=str,nargs='+', default="f", help="Plot types to show: a=summary, b=edge quality (minW=0), c=A-matrix histograms, d=edge quality (minW=arg), e=residuals") + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument('-v',"--verb", type=int, default=1, help="Verbosity level") + parser.add_argument("--minW", type=float, default=0.01, help="Threshold for A-matrix eval, not for fitting") + + args = parser.parse_args() + args.inpPath = os.path.join(args.basePath, 'lassoFdrFit') + args.outPath = os.path.join(args.basePath, 'plots') + os.makedirs(args.outPath, exist_ok=True) + np.set_printoptions(precision=3) + args.showPlots=''.join(args.showPlots) + print(vars(args)) + + # Load fit results + fitFF = os.path.join(args.inpPath, f"{args.dataName}.lassoFit.npz") + fitD, fitMD = read_data_npz(fitFF) + + if args.verb>1: + pprint(fitMD); exit(1) + + # Load spike data for frequency sorting + spikeF = fitMD['provenance']['state_transition_file'] + inpPath2=os.path.join(args.basePath, 'spikesData') + spikesFF = os.path.join(inpPath2, f"{spikeF}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF) + + #pprint(spikeMD) + MD = {**fitMD, 'short_name': args.dataName} + + if 'simDale' in fitMD['data_type']: + truthPath=inpPath2 + truthF=spikeF + if 'simPrism' in fitMD['data_type']: + truthPath= os.path.join(args.basePath, 'truthDale/') + truthF=fitMD['provenance']['state_model_file'] + + truthFF = os.path.join(truthPath, f"{truthF}.simTruth.npz") + trueD,trueMD = read_data_npz(truthFF) + #MD.update( trueMD ) + fitD['E_true']=trueD['E_true'] + fitD['A_true']=trueD['A_true'] + + #- - - - - - - - - - - - - - - - - Setup plotter - - - - - - - - - - - - - - - - + args.prjName = args.dataName + plot = Plotter(args) + + if 'a' in args.showPlots: + plot.summary_fitLasso(fitD,MD,figId=1) + + if 'b' in args.showPlots: + plot.edges_fitLasso(fitD,MD,minW=0.,figId=2) + + if 'c' in args.showPlots: + plot.A_histos(fitD, MD, spikeD, figId=3) + + if 'd' in args.showPlots: + plot.edges_fitLasso(fitD,MD,minW=args.minW,figId=4) + + if 'e' in args.showPlots: + evalD = build_residual_eval_data(fitD, trueD, minW=args.minW, verb=args.verb) + plot.residuals(evalD, MD, figId=5) + + plot.display_all() + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3a_states_simu/fitPrismEM.sh b/causal_net/nonStation_ver3a_states_simu/fitPrismEM.sh new file mode 100755 index 00000000..8ee56d27 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/fitPrismEM.sh @@ -0,0 +1,86 @@ +#!/bin/bash +# salloc -q interactive -C gpu -t 4:00:00 -N 1 -A m2043 +# module load pytorch +# basePath=/pscratch/sd/b/balewski/2026_causalNet_tmp3/ +# Simplified wrapper script for distributed Lasso Poisson training +# +# ./fitPrismEM.sh --basePath $basePath --dataName daleN100_b6ce1e_2ce4d2 --num_states 2 --num_em_iters 2 --m_epochs 16 --fitName abc16 + + +# Capture variable arguments as text +varArgs="$*" + +# Require key variable arguments to be explicitly provided +for required_arg in --dataName --num_states --basePath ; do + if [[ " $varArgs " != *" $required_arg "* ]]; then + echo "ERROR: Missing required argument in varArgs: $required_arg" >&2 + echo "Provided varArgs: $varArgs" >&2 + exit 1 + fi +done + +# Validate that --basePath points to an existing directory +BASE_PATH="" +for ((i=1; i<=$#; i++)); do + arg="${!i}" + if [[ "$arg" == "--basePath" ]]; then + j=$((i + 1)) + if [ "$j" -le "$#" ]; then + BASE_PATH="${!j}" + fi + elif [[ "$arg" == --basePath=* ]]; then + BASE_PATH="${arg#--basePath=}" + fi +done + +if [ -z "$BASE_PATH" ]; then + echo "ERROR: --basePath is required and must include a directory value." >&2 + exit 1 +fi + +if [ ! -d "$BASE_PATH" ]; then + echo "ERROR: --basePath directory does not exist: $BASE_PATH" >&2 + exit 1 +fi + +# Default/fixed arguments +fixArgs=" --time_range_sec 0 80 " +# --dataPath /pscratch/sd/b/balewski/2025_causalNet_tmp/ + +# GPU configuration +NUM_GPUS=4 +CUDA_DEVICES="0,1,2,3" + +# Ensure enough GPUs are available before launching distributed training +HOST_NAME=$(hostname 2>/dev/null || echo "unknown-host") + +if ! command -v nvidia-smi >/dev/null 2>&1; then + echo "ERROR [$HOST_NAME]: nvidia-smi not found; cannot verify GPU availability." >&2 + exit 1 +fi + +AVAILABLE_GPUS=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if [ "$AVAILABLE_GPUS" -lt "$NUM_GPUS" ]; then + echo "ERROR [$HOST_NAME]: Need at least $NUM_GPUS GPUs, but only $AVAILABLE_GPUS detected." >&2 + exit 1 +fi + +# Print configuration +echo "=== Distributed Lasso Poisson Training ===" +echo "Variable args: $varArgs" +echo "Fixed args: $fixArgs" +echo "GPUs: $NUM_GPUS ($CUDA_DEVICES)" +echo "===========================================" +echo "" + +# Set environment variables and run the training +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export CUDA_VISIBLE_DEVICES=$CUDA_DEVICES + +# Execute the training command +echo "Running: ./prism_EM_train.py $fixArgs $varArgs" +echo "" + +time torchrun --standalone --nnodes=1 --nproc_per_node=4 ./prism_EM_train.py $fixArgs $varArgs diff --git a/causal_net/nonStation_ver3a_states_simu/fit_lassoPoisson3.py b/causal_net/nonStation_ver3a_states_simu/fit_lassoPoisson3.py new file mode 100755 index 00000000..ffa31d54 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/fit_lassoPoisson3.py @@ -0,0 +1,134 @@ +#!/usr/bin/env python3 +""" +Single-GPU training of Poisson GLM with LASSO regularization. + +This script fits a Poisson Generalized Linear Model for neural connectivity +inference with L1 (LASSO) regularization on a single GPU. + +Usage: + ./fit_lassoPoisson3.py --dataName mydata --num_epochs 100 +""" + +import os +import time +import secrets +import argparse + + +import numpy as np +from pprint import pprint +import torch + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from PoissonGLModel import PoissonGLModel + +from UtilTorch import check_gpu_availability, preprocess_data, train_Poisson_model, make_loader + +######################### +# MAIN +######################### + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dataName", type=str, default="dale_2aee70") + parser.add_argument("--basePath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input/output data") + parser.add_argument("--inpPath", type=str, default=None, help="alternative location of input, takes precedence") + parser.add_argument("--num_samples", type=int, default=None) + parser.add_argument("--num_epochs", type=int, default=200) + parser.add_argument("--batch_size", type=int, default=2048) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--L1_alpha", type=float, default=0.02, help="L1 regularization strength, higher=more sparse (0: disable soft-thresholding)") + parser.add_argument("--rho_max", type=float, default=0.97, help="Maximum allowed spectral radius of A; projection applied after each batch") + parser.add_argument("--prescale_m_step_4_ArhoMax", type=int, default=20, help="Apply spectral-radius projection every N batches") + parser.add_argument("--L1_prune_epoch", type=int, default=50, help="Delay L1 edge-pruning only; spectral-radius correction starts immediately") + parser.add_argument("--minW", type=float, default=0.01, help="Threshold for A-matrix eval, not for fitting") + parser.add_argument("--fitName", type=str, default=None) + parser.add_argument("--desyncTime", action='store_true', help="If true completely shuffle time axis for input data, independently for all channels") + parser.add_argument("--dropDataFrac", type=float, default=0.0, help="Fraction of training samples to randomly drop (0.0=use all data, 0.3=drop 30%%)") + parser.add_argument("--verb", "-v", type=int, default=1, help="Verbosity level") + + args = parser.parse_args() + if args.inpPath ==None: + args.inpPath = os.path.join(args.basePath, 'spikesData') + outPath = os.path.join(args.basePath, 'lassoFdrFit') + assert os.path.exists(outPath) + + device = check_gpu_availability() + print("\nFitLasso3 Config:", vars(args,), "\n") + # enable fast matmul paths + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision('high') + gpu_name = torch.cuda.get_device_name(device) if isinstance(device, torch.device) and device.type=='cuda' else str(device) + print("Using device %s : %s" % (str(device), gpu_name)) + + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=True) + if args.verb > 1: + pprint(spikeMD) + dataYield = spikeD['spikes'] + dataRates = spikeD['single_rates'] + step_size = spikeMD['time_step_sec'] + eta_clip = spikeMD["poisson_eta_clip"] + if 'simDale' in spikeMD['data_type']: + _, _, Nn = dataYield.shape + dataYield = dataYield[0] + dataRates = dataRates[0] + else: + _, Nn = dataYield.shape + + XY_np = preprocess_data(dataYield, args) + n_pairs = XY_np.shape[0] + print(f"Preprocessed data: XY shape={XY_np.shape}, n_pairs={n_pairs}") + + # Split XY into X and Y + X_np = XY_np[:, 0, :] + Yt_np = XY_np[:, 1, :] + n_pairs = X_np.shape[0] + + assert n_pairs >= args.batch_size, f"ERROR: Not enough samples ({n_pairs}) for batch size ({args.batch_size}) after data dropping." + + train_loader = make_loader(X_np, Yt_np, args, is_dist=False) + print(f"Loaded pairs={n_pairs/1000}k, Nn={Nn}, using {n_pairs/1000}k pairs (all for training), batch_size={args.batch_size}") + + model = PoissonGLModel(Nn, eta_clip=eta_clip).to(device) + start_time = time.time() + losses_total, losses_wo_L1, learning_rates, train_epochs, sparsity_epoch, nz_offdiag_epoch, spectral_radius_epoch = train_Poisson_model( + model, device, train_loader, args.num_epochs, lr=args.lr, L1_alpha=args.L1_alpha, firing_rates=dataRates, use_scheduler=True, + train_sampler=None, apply_prox=(args.L1_alpha > 0), minW=args.minW, rho_max=args.rho_max, prescale_m_step_4_ArhoMax=args.prescale_m_step_4_ArhoMax, L1_prune_epoch=args.L1_prune_epoch + ) + + total_time = time.time() - start_time + print(f"Training completed in {total_time:.1f} seconds") + + if args.fitName is None: + hash6 = secrets.token_hex(3) # 6 hex digits + fit_core = f"{args.dataName}-{hash6}" + else: + fit_core = args.fitName + + # saving from fit_Lasso --- + mdl = model + A_hat=mdl.A.detach().cpu().numpy() + E_hat = (np.abs(A_hat) > 1e-5) + lassoD = { 'A_lasso': A_hat, 'B_lasso': mdl.B.detach().cpu().numpy(), 'E_lasso':E_hat, 'losses_total': np.array(losses_total), 'losses_wo_L1': np.array(losses_wo_L1), 'losses_epochs': np.array(train_epochs, dtype=np.int32), 'learning_rates': np.array(learning_rates), 'sparsity_epoch': np.array(sparsity_epoch, dtype=np.float32), 'nz_offdiag_epoch': np.array(nz_offdiag_epoch, dtype=np.int32), 'spectral_radius_epoch': np.array(spectral_radius_epoch, dtype=np.float32), 'single_rates': dataRates } + lassoMD = {'batch_size': args.batch_size, 'num_samples_used': n_pairs, 'num_epochs': args.num_epochs, 'num_train_samples': n_pairs, 'learning_rate': args.lr, 'L1_alpha': args.L1_alpha, 'rho_max': args.rho_max, 'prescale_m_step_4_ArhoMax': args.prescale_m_step_4_ArhoMax, 'L1_prune_epoch': args.L1_prune_epoch, 'minW': args.minW, 'step_size': step_size, 'training_time_sec': total_time, 'num_neurons': Nn, 'dropDataFrac': args.dropDataFrac } + # 'lassoFit_output_name': fit_core, 'lassoFit_input_name': args.dataName, 'lassoFit_input_path': args.inpPath , + + outMD=spikeMD + outMD['fit_type']='lasso' + outMD['fit_lasso']=lassoMD + outMD['edge_selector']={'selector_type':'None'} + outMD['provenance']['output_lasso_file']=fit_core + + if args.verb>1: pprint(outMD) + fitFF = os.path.join(outPath, f"{fit_core}.lassoFit.npz") + write_data_npz(lassoD, fitFF, metaD=outMD) + + if spikeMD['data_type']=='simDale': flags=' -p a b c ' + else: flags=' -p a d e ' + print(' basePath='+args.basePath) + print(' ./eval_fitLasso.py --basePath $basePath --dataName %s %s \n ' % (fit_core,flags)) + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3a_states_simu/gen_daleMatrices3.py b/causal_net/nonStation_ver3a_states_simu/gen_daleMatrices3.py new file mode 100755 index 00000000..5f1e0197 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/gen_daleMatrices3.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +""" + ./gen_daleMatrices3.py --num_neurons 100 --num_excite 70 --num_steps 10000 --dataName test1 + ./gen_daleMatrices3.py --num_neurons 100 --num_excite 70 --spectral_radius 0.5 --Boffsets 0 10 20 --dataName test2 + +Primary purpose: generate the ground-truth dictionary (A_true, B_true) for use +by gen_nonStationarySpikes3.py. The stationary spike generation performed here +is for evaluation only (firing-rate sanity check per B-vector). + +Dale's principle: each neuron is either excitatory (positive outgoing +weights) or inhibitory (negative outgoing weights), never both. + +Pipeline: +1. Sparse connectivity mask E_true (N, N). + Each neuron draws its own connection probability uniformly from + --edge_prob [lo, hi]. Self-loops (diagonal) are always included. + +2. Generate ONE weight matrix A_true for --spectral_radius R: + - Excitatory rows (0..N_E-1): weights ~ Uniform(1-v, 1+v), v=0.2. + - Inhibitory rows (N_E..N-1): weights ~ Uniform(-r-rv, -r+rv), + r = N_E/N_I (balance ratio). + - Mask with E_true, then rescale so rho(A_true) = R exactly. + +3. For each offset in --Boffsets generate one bias vector B_m (N,): + - idle rate range shifted by offset, converted to log scale, + with separate corrections for excitatory and inhibitory neurons. + - Stationary spikes Y are simulated (for evaluation) via + Y_t ~ Poisson(exp(clip(A @ Y_{t-1} + B_m, max=eta_clip)) * dt). + - Firing-rate statistics (rate, variance, Fano factor) computed. + +Output files saved to /truthDale/: + .simTruth.npz — A_true, B_true, E_true + metadata + .spikes.npz — stationary spikes + rate statistics + +Output shapes (M = len(Boffsets), N = num_neurons, T = num_steps): + E_true (N, N) int — shared binary connectivity mask + A_true (N, N) float — single shared weight matrix + B_true (M, N) float — one bias vector per state/offset + spikes (M, T, N) uint8 — stationary spikes per bias vector + single_rates (M, N) float — mean firing rates (Hz) per bias vector +""" + +import numpy as np +import time,hashlib +import scipy +import os +import sys + +import argparse +from pprint import pprint + +from toolbox.Util_NumpyIO import write_data_npz +from UtilDalePoisson import estimate_rates + +###### Matrix generation ################## + +def generate_sparse_mask(n_units, edge_prob): + """ + Creates a binary mask representing the network topology. + - n_units: Number of neurons + - edge_prob: [lo, hi] range; each neuron gets a random connectivity drawn uniformly from this range + - allows for self-loops (aka non-zero diagonal elements) + """ + prob_lo, prob_hi = edge_prob + conn_per_neuron = np.random.uniform(prob_lo, prob_hi, size=n_units) + E_true = (np.random.rand(n_units, n_units) < conn_per_neuron[:, None]) + np.fill_diagonal(E_true, True) # allow for self-loops + return E_true + +# Generate an initial network connectivity matrix +def init_W(n_units, n_excite, E_true, R, varyW, verb=1): + """ + Generalized weight initialization with E-I row-based balancing, + zero-diagonal, spectral radius conserved. + """ + # 1. Assertions to ensure valid population counts + n_inhib = n_units - n_excite + assert n_excite > 0, "n_excite must be greater than 0" + assert n_inhib > 0, "n_inhib must be greater than 0 (n_units > n_excite)" + + # 2. Balance ratio: ensures the expected sum of the matrix is zero + ie_ratio = n_excite / n_inhib + + # 3. Initialize the weight container + W = np.zeros((n_units, n_units)) + + # 4. Assign Excitatory Weights (Rows 0 to n_excite-1) + W[:n_excite, :] = np.random.uniform(1.0 - varyW, + 1.0 + varyW, + (n_excite, n_units)) + + # 5. Assign Inhibitory Weights (Rows n_excite to n_units-1) + i_center = -ie_ratio + i_vary = varyW * ie_ratio + W[n_excite:, :] = np.random.uniform(i_center - i_vary, + i_center + i_vary, + (n_inhib, n_units)) + + # 6. Apply the connectivity mask (Topology) + W = W * E_true + + # 7. Enforce negative diagonal: flip sign of any positive self-loop. + d_idx = np.diag_indices(n_units) + W[d_idx] = np.where(W[d_idx] > 0.0, -W[d_idx], W[d_idx]) + + # 8. Spectral Normalization to ensure Stability + # This scales the entire cloud of eigenvalues to fit within radius R + current_rho = np.max(np.abs(np.linalg.eigvals(W))) + if verb > 0: + print('initW current_rho :',current_rho ) + if current_rho > 0: + W = W * (R / current_rho) + + return W + +# Generate the full set of matrices for use in subsequent synthetic experiments +def gen_dale_matrics(conf, E_true, verb=1): + """Generates one stable Dale matrix based on configuration.""" + num_neurons, num_excite, Rmax= conf['num_neurons'], conf['num_excite'],conf[ 'spect_radius'] + if verb > 0: + print(f"\n=== Generating Dale Matrix ===") + print(f"Parameters: N={num_neurons}, excit={num_excite}, Rmax={Rmax:.1f}") + + # additional configuration + varyW=0.2 # controls variation of excitatory weights + + A=init_W(num_neurons, num_excite, E_true, Rmax, varyW, verb=verb) + return A + +def spectral_radius_scaling(A, factors): + n = A.shape[0] + off_diag_mask = ~np.eye(n, dtype=bool) + + print(f"{'factor':>8s} {'ρ(all)':>10s} {'ρ(off-diag)':>12s}") + print("-" * 34) + + for c in factors: + # a) scale all elements + rho_all = np.max(np.abs(np.linalg.eigvals(c * A))) + + # b) scale only off-diagonal + A_off = A.copy() + A_off[off_diag_mask] *= c + rho_off = np.max(np.abs(np.linalg.eigvals(A_off))) + + print(f"{c:8.3f} {rho_all:10.4f} {rho_off:12.4f}") + + +#################### Simulation ################## + +def set_flat_selfSpiking(Nn, idleRate, spect_radius, num_excite, boffsets): + """Generate B_idle per B-offset; excitatory idle-rate range is scaled by sqrt(50/Nn).""" + assert 0 < num_excite <= Nn + sizeScale = np.sqrt(float(Nn)/100) + B_all = np.zeros((len(boffsets), Nn)) + R = float(spect_radius) + for ib, offset in enumerate(boffsets): + idle_eff = np.array(idleRate, dtype=float) + Ri_scaled = idle_eff * R + Bi = np.log(Ri_scaled) + B_all[ib] = np.random.uniform(Bi[0], Bi[1], size=(Nn,)) + #B_all+= + float(offset) + if 1: # for ver Mar 11 + B_all[ib, :num_excite] +=1+offset - R*1.5 - sizeScale # reduce excitatory rate + B_all[ib, num_excite:] += offset # reduce inhibitory rate + return B_all + +def gen_stationary_lag1_poisson(num_steps, dt, A, B_intercept, num_excite, eta_clip,verb=0): + """ + Generates a multivariate Poisson VAR(1) process: + Y_t ~ Poisson(exp(A @ Y_{t-1} + B_intercept)) + + Args: + num_steps (int): Number of time steps for the simulation. + dt (float): Integration time step in seconds. + A (np.ndarray): N x N autoregressive coefficient matrix (the Dale matrix). + B_intercept (np.ndarray): N-dimensional intercept vector (bias). + num_excite (int): Number of excitatory neurons. + verb (int): Verbosity level for printing progress. + + Returns: + (Y, A, B_intercept): Tuple containing: + Y (np.ndarray): num_steps x N time series of spike counts. + A (np.ndarray): The input connectivity matrix. + B_intercept (np.ndarray): The input intercept vector. + """ + if verb > 0: + print(f"\n=== Generating Poisson Process ===") + print(f"Simulation parameters: steps={num_steps}, dt={dt:.3f}, neurons={A.shape[0]}, excit={num_excite}") + print(f"Matrix A stats: min={np.min(A):.3f}, max={np.max(A):.3f}, mean={np.mean(A):.3f}") + print(f"Bias B stats: min={np.min(B_intercept):.3f}, max={np.max(B_intercept):.3f}, mean={np.mean(B_intercept):.3f}") + + if A is None: + raise ValueError("Connectivity matrix A cannot be None.") + d=A.shape[0] + + Y = np.zeros((num_steps, d), dtype=int) + Y[0] = np.random.poisson(np.exp(B_intercept)*dt) # initial state + if verb > 0: + print(f"Initial state: total spikes={np.sum(Y[0])}, excit spikes={np.sum(Y[0][:num_excite])}, inhib spikes={np.sum(Y[0][num_excite:])}") + + if verb>0: + print('t=0 Y[t] sum=%d, Excit(first 3):%s, Inhib(first 3):%s'%(np.sum(Y[0]), Y[0][:3], Y[0][num_excite:num_excite+3])) + + kk=5 + # Main simulation loop + if verb > 0: + print("Starting main simulation loop...") + for t in range(1, num_steps): + eta = A @ Y[t-1] + B_intercept + #eta = B_intercept # use it to see idle rate only + lambda_t = np.exp(np.clip(eta, max=eta_clip)) # avoid overflow + Y[t] = np.random.poisson(lambda_t*dt) + + if verb>0 and t<5: + print('t=%d Y[t] sum=%d, Excit:%s, Inhib:%s'%(t, np.sum(Y[t]), Y[t][:kk], Y[t][num_excite:num_excite+kk])) + + if verb > 0 and t % (num_steps // 4) == 0: + print(f" Progress: {t}/{num_steps} steps ({t/num_steps*100:.1f}%) - total spikes in this step: {np.sum(Y[t])}") + + if verb > 0: + print(f"Simulation complete. Final state: total spikes={np.sum(Y[-1])}") + print(f"Spike data stats: min={np.min(Y)}, max={np.max(Y)}, mean={np.mean(Y):.2f}, total spikes={np.sum(Y)}") + + return Y + +######################### +# MAIN +######################### + +def main(): + print("=" * 60) + print("DALE POISSON SIMULATION Lag=1") + print("=" * 60) + + parser = argparse.ArgumentParser(description="Simulate a recurrent neural network with Dale's principle.") + parser.add_argument("--num_neurons", type=int, default=50, help="Total number of neurons in the network.") + parser.add_argument("--num_excite", type=int, default=None, help="Number of excitatory neurons.") + parser.add_argument("--edge_prob", type=float, nargs=2, default=[0.05, 0.2], help="Range of edge probability [min, max]; mean is used as mask connectivity.") + parser.add_argument("--num_steps", type=int, default=10_001, help="Number of time steps for simulation.") + parser.add_argument("--step_size", type=float, default=0.01, help="Integration time step size (dt) in seconds.") + parser.add_argument("--spectral_radius", type=float, default=0.3, help="Target spectral radius value for the connectivity matrix.") + parser.add_argument("--idleRate", type=float, nargs=2, default=[15, 30.], help="Range of idle firing rates [min, max] in Hz.") + parser.add_argument("--Boffsets", type=float, nargs='+', default=[0.0], help="Per-state offsets added to idleRate range.") + parser.add_argument('-v',"--verb", type=int, default=1, help="Verbosity level (0=quiet, 1=normal).") + parser.add_argument("--dataName", type=str, default=None, help="Base name for output files (default: daleN_).") + parser.add_argument("--basePath", type=str, default='/pscratch/sd/b/balewski/2025_causalNet_tmp/', help="Output directory for all files.") + + np.set_printoptions(precision=3, suppress=True) + + args = parser.parse_args() + args.varTwindow=5 #(sec) + args.poisson_eta_clip=5 #~ [1e-3Hz , 1e3Hz] + if args.dataName is None: + args.dataName='daleN%d_'%args.num_neurons+hashlib.md5(os.urandom(32)).hexdigest()[:6] + + outPath=os.path.join(args.basePath, 'truthDale') + if args.num_excite==None : # per Roy & Kris wisdom + args.num_excite=int(0.8*args.num_neurons) + print('gen dale matrices args:', vars(args), '\n') + + # Validation checks + Nn = args.num_neurons + assert Nn >= 10 + assert args.num_excite >= 5 + assert args.num_excite < Nn + assert os.path.exists(args.basePath) + assert os.path.exists(outPath) + assert args.num_steps>=1000 + assert args.step_size>0.001 + assert args.idleRate[0]>=0.5 + assert args.idleRate[1]>args.idleRate[0] + assert len(args.Boffsets) >= 1 + # Generate sparse connectivity mask (per-neuron random connectivity in edge_prob range) + E_true = generate_sparse_mask(Nn, args.edge_prob) + print(f"\n=== Generated sparse mask: edge_prob={args.edge_prob}, actual={np.mean(E_true):.3f}, non-zero={np.sum(E_true)} ===") + + dale_conf = { + 'num_neurons': args.num_neurons, + 'num_excite': args.num_excite, + 'spectral_radius': args.spectral_radius, + 'edge_prob': args.edge_prob, + 'idleRate': args.idleRate, + 'Boffsets': args.Boffsets, + } + + print("Dale configuration:"); pprint(dale_conf) + + evol_conf={ + 'num_steps': args.num_steps, + 'step_size': args.step_size, + 'evol_time': args.num_steps*args.step_size, + 'poisson_eta_clip': args.poisson_eta_clip + } + + B_all = set_flat_selfSpiking(Nn, args.idleRate, args.spectral_radius, args.num_excite, boffsets=args.Boffsets) + + max_samples = 100_000 + + Y_list = [] + rates_list, rates_var_list, fano_list = [], [], [] + stats_list = [] + + verb_r = args.verb + print(f"\n{'='*60}") + print(f" Spectral radius: R={args.spectral_radius:.3f}") + print(f"{'='*60}") + + dale_conf_r = dale_conf.copy() + dale_conf_r['spect_radius'] = args.spectral_radius + + A_dale = gen_dale_matrics(dale_conf_r, E_true, verb=verb_r) + + if 0: # test scaling behavior of spectral radius when scaling A by different factors + factors = [0.2, 0.4, 0.6, 0.8, 0.9, ] + spectral_radius_scaling(A_dale, factors) + exit(1) + + if verb_r > 0: + total_connections = A_dale.size + zero_connections = np.sum(np.abs(A_dale) < 1e-10) + non_zero_connections = total_connections - zero_connections + sparsity = zero_connections / total_connections + print(f'Matrix sparsity: {sparsity*100:.1f}% ({zero_connections}/{total_connections} connections are zero)') + print(f'Non-zero connections: {non_zero_connections} ({(1-sparsity)*100:.1f}%)') + + for ib, offset in enumerate(args.Boffsets): + print(f"\n{'='*60}") + print(f" B offset [{ib+1}/{len(args.Boffsets)}]: offset={offset:.3f}") + print(f"{'='*60}") + + B_idle = B_all[ib] + start_time = time.time() + Y = gen_stationary_lag1_poisson(num_steps=args.num_steps, dt=args.step_size, A=A_dale, B_intercept=B_idle, num_excite=args.num_excite,eta_clip=args.poisson_eta_clip, verb=verb_r if ib == 0 else 0) + sim_time = time.time() - start_time + print("Spike generation completed in %.1f seconds" % sim_time) + + stats_dict, rates_dict, _ = estimate_rates(Y, dt=args.step_size, num_excite=args.num_excite, max_samples=max_samples, varTwindow=args.varTwindow, mxNn=5, verb=0, spect_radius=args.spectral_radius) + stats_dict['var_time_window_sec'] = float(args.varTwindow) + stats_dict['max_samples'] = int(max_samples) + + Y_list.append(np.clip(Y, 0, 255).astype(np.uint8)) + rates_list.append(rates_dict['single_rates']) + rates_var_list.append(rates_dict['sigle_rates_var']) + fano_list.append(rates_dict['single_fano_fact']) + stats_list.append(stats_dict) + + # Stack results along axis=0 (B-offset dimension) + trueD = { + 'A_true': A_dale, + 'B_true': B_all, + 'E_true': E_true + } + + trueMD = {'dale_conf': dale_conf, 'evol_conf': evol_conf, 'short_name': args.dataName, + 'provenance':{'state_model_file':args.dataName}} + # delete , 'dale_simu_stats': stats_list} + + spikeD = { + 'spikes': np.stack(Y_list, axis=0), + 'single_rates': np.stack(rates_list, axis=0), + 'sigle_rates_var': np.stack(rates_var_list, axis=0), + 'single_fano_fact': np.stack(fano_list, axis=0) + } + spikeMD={ 'short_name':args.dataName,'time_step_sec':args.step_size,'data_type':'simDaleStates', 'poisson_eta_clip': args.poisson_eta_clip } + + outFt = os.path.join(outPath,args.dataName + '.simTruth.npz') + write_data_npz(trueD, outFt, metaD=trueMD) + if args.verb>1: pprint(trueMD) + outFs = os.path.join(outPath, args.dataName + '.spikes.npz') + write_data_npz(spikeD, outFs, metaD=spikeMD) + if args.verb>1: pprint(spikeMD) + + print("\nSimulation completed successfully!") + print(f"\nRate Summary {args.dataName} N={args.num_neurons} exc={args.num_excite}, R={args.spectral_radius:.3f} ") + header = f"{'state':>5} {'B offset':>9} {'all rate (Hz)':>14} {'exc rate (Hz)':>14} {'inh rate (Hz)':>14}" + print(header) + print("-" * len(header)) + for ib, offset in enumerate(args.Boffsets): + s = stats_list[ib] + print(f"{ib:5d} {float(offset):9.1f} {s['avg_spike_rate_all']:14.1f} {s['avg_spike_rate_excit']:14.1f} {s['avg_spike_rate_inhib']:14.1f}") + print("\nNext step commands:") + print(" basePath="+args.basePath) + print(" ./view_daleMatrix3.py --basePath $basePath --dataName %s -p b -m 0 -X -p a c d " % args.dataName) + print(" ./view_spikesTrain3.py --basePath $basePath --dataName %s --time_range_sec 0 20 -p b -m 0 -X " % args.dataName) + print(" ./gen_nonStationarySpikes3.py --basePath $basePath --inputStates %s --true_dwell_sec 1.0 " % args.dataName) + print(" ./fit_lassoPoisson.py --basePath $basePath --inpPath ${basePath}/truthDale --dataName %s --num_epochs 50 " % args.dataName) + + +if __name__ == '__main__': + main() + diff --git a/causal_net/nonStation_ver3a_states_simu/gen_nonStationarySpikes3.py b/causal_net/nonStation_ver3a_states_simu/gen_nonStationarySpikes3.py new file mode 100755 index 00000000..03422dfb --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/gen_nonStationarySpikes3.py @@ -0,0 +1,436 @@ +#!/usr/bin/env python3 +""" +Generate one non-stationary spike train using the ground-truth dictionary +(A_true, B_true) produced by gen_daleMatrices3.py. + +The connectivity matrix A is shared across all states. Each state m has its +own bias vector B_m. At every time bin the effective bias is the convex +combination B_eff = sum_m c_mt * B_m, where c_t tracks a slowly-moving +target that aims at the one-hot vector of the current target state S_true[t]. + +State-sequence schedules (--schedule): + mc — Markov chain with geometric dwell (mean = true_dwell_sec / dt bins), + minimum dwell = ceil(0.3 * true_dwell_sec / dt). + rr — Round-robin: states cycle 0,1,...,M-1 with exactly + true_dwell_sec / dt bins per visit. + +Smooth coefficient update at each bin (state_change_speed = nu): + c_t <- Simplex_project( c_{t-1} + clip(e_{S_t} - c_{t-1}, -nu, nu) ) + +Spike generation (Poisson GLM): + eta_t = A @ Y_{t-1} + B_eff + lambda_t = exp( clip(eta_t, max=eta_clip) ) * dt + Y_t ~ Poisson(lambda_t) + +An oracle state sequence S_oracle is also computed: at each t the state with +the highest Poisson log-likelihood given (A, {B_m}) and the observed spikes. + +Output files saved to /spikesData/: + .spikes.npz — spikes (T, N) int32, single_rates (N,) + .prismTruth.npz — S_true, C_true, S_oracle, state_transition, + A_true, B_true, rate variance, Fano factor + +Output shapes (T = num_steps, N = num_neurons, M = num_states): + spikes (T, N) int32 — non-stationary spike counts + single_rates (N,) float — mean firing rates (Hz) + S_true (T,) int32 — target state index per bin + C_true (T, M) float32 — smooth simplex coefficients + S_oracle (T,) int32 — oracle (max-likelihood) state per bin + state_transition (M, M) float32 — Markov transition matrix used/implied + A_true (N, N) float — shared connectivity matrix (copy) + B_true (M, N) float — per-state bias vectors (copy) +""" + +import os +import hashlib +import argparse +from pprint import pprint +import numpy as np + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from gen_daleMatrices3 import estimate_rates + + +def get_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v", "--verbosity", type=int, default=1, dest="verb", help="Verbosity level.") + parser.add_argument("--basePath", default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input data") + parser.add_argument("--inputStates", default=None, help="input simTruth base name") + parser.add_argument("--dataName", type=str, default=None, help="output spikes base name (default: _)") + + parser.add_argument("-t", "--num_steps", type=int, default=None, help="Number of time steps (default: from input evol_conf)") + parser.add_argument("--state_change_speed", type=float, default=0.33, help="Max coefficient change per step.") + parser.add_argument("--true_dwell_sec", type=float, default=1.0, + help="Mean dwell time in seconds to stay in a target state.") + parser.add_argument("--seed", type=int, default=42, help="Optional random seed.") + parser.add_argument("--schedule", choices=["mc", "rr"], default="mc", + help="State schedule: 'mc' = random Markov chain (default), " + "'rr' = deterministic round-robin (equal state coverage).") + + args = parser.parse_args() + args.inpPath = os.path.join(args.basePath, "truthDale") + args.outPath = os.path.join(args.basePath, "spikesData") + if args.dataName is None: + args.dataName = f"{args.inputStates}_{hashlib.md5(os.urandom(32)).hexdigest()[:6]}" + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + assert os.path.exists(args.basePath), f"missing basePath: {args.basePath}" + assert os.path.exists(args.inpPath), f"missing inpPath: {args.inpPath}" + os.makedirs(args.outPath, exist_ok=True) + return args + + +def ensure_state_atoms(A_in, B_in): + """Normalize A/B arrays to A(N,N), B(M,N) with shared A across states.""" + A = A_in + B = B_in + assert A.ndim == 2, f"A_true must be 2D, got shape={A.shape}" + if B.ndim == 1: + B = B[None, ...] + assert B.ndim == 2, f"B_true must be 1D or 2D, got shape={B.shape}" + assert A.shape[0] == A.shape[1], f"A must be square, got {A.shape}" + assert A.shape[0] == B.shape[1], f"neuron mismatch: A N={A.shape[0]}, B N={B.shape[1]}" + return A.astype(float), B.astype(float) + + +def _print_state_stats(S_true, n_states, n_steps): + """Print per-state counts and dwell statistics (shared by both schedules).""" + counts = np.bincount(S_true, minlength=n_states) + fracs = counts / n_steps + parts = ' '.join(f's{m}:{counts[m]}({fracs[m]:.1%})' for m in range(n_states)) + print(f"Target state distribution (T={n_steps}): {parts}") + + dwell_lens = {m: [] for m in range(n_states)} + run_state, run_len = S_true[0], 1 + for t in range(1, n_steps): + if S_true[t] == run_state: + run_len += 1 + else: + dwell_lens[run_state].append(run_len) + run_state, run_len = S_true[t], 1 + dwell_lens[run_state].append(run_len) + for m in range(n_states): + d = dwell_lens[m] + print(f" state {m}: {len(d)} episodes, dwell mean={np.mean(d):.1f} steps, " + f"total={np.sum(d)} steps") + + +def build_target_states_mc(n_steps, n_states, dwell_steps, rng): + """Random Markov chain target-state sequence with geometric dwell time. + + Each time step the chain either stays (prob = 1 - 1/dwell_steps) or + transitions uniformly to one of the other states. State coverage is + uncontrolled — a single unlucky run can give very few bins to one state. + """ + if n_states == 1: + return np.zeros(n_steps, dtype=np.int32), np.ones((1, 1), dtype=np.float32) + + minDwellFrac = 0.3 + min_dwell_steps = max(1, int(np.ceil(minDwellFrac * float(dwell_steps)))) + p_exit = 1.0 / float(dwell_steps) + p_stay = 1.0 - p_exit + trans = np.full((n_states, n_states), p_exit / float(n_states - 1), dtype=float) + np.fill_diagonal(trans, p_stay) + + S_true = np.zeros(n_steps, dtype=np.int32) + state = 0 + run_len = 0 + for t in range(n_steps): + if run_len < min_dwell_steps: + next_state = state + else: + next_state = rng.choice(n_states, p=trans[state]) + S_true[t] = next_state + if next_state == state: + run_len += 1 + else: + state = next_state + run_len = 1 + + _print_state_stats(S_true, n_states, n_steps) + return S_true, trans.astype(np.float32) + + +def build_target_states_rr(n_steps, n_states, dwell_steps, rng): + """Round-robin with fixed dwell_steps per visit. + + States cycle 0,1,2,0,1,2,... Every visit lasts exactly dwell_steps + steps (last visit of each state is trimmed to fit T). Total steps per + state is guaranteed to be T//M ± dwell_steps — no luck involved. + """ + if n_states == 1: + return np.zeros(n_steps, dtype=np.int32), np.ones((1, 1), dtype=np.float32) + + S_true = np.zeros(n_steps, dtype=np.int32) + t = 0 + visit = 0 + while t < n_steps: + state = visit % n_states # strict round-robin + t_end = min(t + dwell_steps, n_steps) # fixed dwell, trim at end + S_true[t:t_end] = state + t = t_end + visit += 1 + + # Transition matrix: same symmetric form as MC for metadata consistency + p_exit = 1.0 / float(dwell_steps) + p_stay = 1.0 - p_exit + trans = np.full((n_states, n_states), p_exit / float(n_states - 1), dtype=float) + np.fill_diagonal(trans, p_stay) + + _print_state_stats(S_true, n_states, n_steps) + return S_true, trans.astype(np.float32) + +def simulate_switching_poisson(n_steps, A, B_atoms, S_true, state_change_speed, dt, eta_clip, rng, verb=1): + """Switching Poisson generator: smooth c_t toward one-hot target state S_true[t].""" + n_states, n_neurons = B_atoms.shape + spikes = np.zeros((n_steps, n_neurons), dtype=np.int32) + C_true = np.zeros((n_steps, n_states), dtype=np.float32) + + c_curr = np.zeros(n_states, dtype=float) + c_curr[S_true[0]] = 1.0 + + for t in range(n_steps): + target = np.zeros(n_states, dtype=float) + target[S_true[t]] = 1.0 + + diff = target - c_curr + c_curr += np.clip(diff, -state_change_speed, state_change_speed) + c_sum = np.sum(c_curr) + if c_sum <= 0: + c_curr = target.copy() + c_sum = 1.0 + c_curr /= c_sum + C_true[t] = c_curr + + if verb > 1 and t < 12: + c_str = ", ".join([f"{x:.2f}" for x in c_curr]) + print(f"{t:<6} | {int(S_true[t]):<6} | [{c_str}]") + + A_eff = A + B_eff = np.einsum("m,mj->j", c_curr, B_atoms) + + prev_y = spikes[t - 1].astype(float) if t > 0 else np.zeros(n_neurons, dtype=float) + eta_t = A_eff @ prev_y + B_eff + lambda_t = np.exp(np.clip(eta_t, max=eta_clip)) + spikes[t] = rng.poisson(lambda_t * dt).astype(np.int32) + + return spikes, C_true + + +def compute_oracle_states_comA(spikes, A, B_atoms, dt, eta_clip): + """Oracle state by max Poisson log-likelihood using shared A and per-state B.""" + n_steps, n_neurons = spikes.shape + n_states = B_atoms.shape[0] + S_oracle = np.zeros((n_steps,), dtype=np.int32) + prev_y = np.zeros(n_neurons, dtype=np.float64) + for t in range(n_steps): + y_curr = spikes[t].astype(np.float64) + base = A @ prev_y + eta = base[None, :] + B_atoms + eta_c = np.clip(eta, max=eta_clip) + lam = np.exp(eta_c) * dt + scores = np.sum(y_curr * eta_c - lam, axis=1) + S_oracle[t] = int(np.argmax(scores)) + prev_y = y_curr + return S_oracle + + +def compute_oracle_score(S_true, S_oracle, n_states): + """Return overall and per-state agreement between oracle and target states.""" + if S_true is None or S_oracle is None: + return None, [] + nmin = min(S_true.shape[0], S_oracle.shape[0]) + if nmin <= 0: + return None, [] + avr_score = float(np.mean(S_true[:nmin] == S_oracle[:nmin])) + score_per_state = [] + for m in range(n_states): + mask = S_true[:nmin] == m + if mask.sum() > 0: + score_per_state.append(float((S_oracle[:nmin][mask] == m).mean())) + else: + score_per_state.append(float("nan")) + return avr_score, score_per_state + + +def main(): + args = get_parser() + print('gen non-stationary spikes args:', vars(args), '\n') + np.set_printoptions(precision=3, suppress=True) + + daleFF = os.path.join(args.inpPath, f"{args.inputStates}.simTruth.npz") + daleD, daleMD = read_data_npz(daleFF, verb=args.verb > 0) + if args.verb > 1: + print("\nInput simTruth metadata:") + pprint(daleMD) + + assert isinstance(daleMD, dict), "Expected dictionary metadata in simTruth file" + dale_conf_in = daleMD["dale_conf"] + evol_conf_in = daleMD["evol_conf"] + proven_in=daleMD['provenance'] + + step_size = evol_conf_in["step_size"] + var_time_window_sec = 5 #(sec) + max_samples = 100_000 # time steps + + if args.num_steps is None: + args.num_steps = int(evol_conf_in["num_steps"]) + + assert args.num_steps >= 100 + assert step_size > 0 + assert args.state_change_speed > 0 + assert args.true_dwell_sec > 0.0 + assert max_samples >= 100 + assert var_time_window_sec > 0 + + true_dwell_steps = max(1, int(np.ceil(float(args.true_dwell_sec) / float(step_size)))) + + A_atoms, B_atoms = ensure_state_atoms(daleD["A_true"], daleD["B_true"]) + n_states, n_neurons = B_atoms.shape + num_excite = int(dale_conf_in["num_excite"]) + num_excite = min(max(1, num_excite), n_neurons - 1) + + rng = np.random.default_rng(args.seed) + + schedule_fn = {"mc": build_target_states_mc, + "rr": build_target_states_rr}[args.schedule] + S_true, transition_matrix = schedule_fn(args.num_steps, n_states, true_dwell_steps, rng) + + spikes, C_true = simulate_switching_poisson( + n_steps=args.num_steps, + A=A_atoms, + B_atoms=B_atoms, + S_true=S_true, + state_change_speed=args.state_change_speed, + dt=step_size, + eta_clip=evol_conf_in['poisson_eta_clip'], + rng=rng, + verb=args.verb, + ) + S_oracle = compute_oracle_states_comA( + spikes=spikes, + A=A_atoms, + B_atoms=B_atoms, + dt=step_size, + eta_clip=evol_conf_in['poisson_eta_clip'], + ) + + stats_dict, rates_dict, _ = estimate_rates( + spikes, + dt=step_size, + num_excite=num_excite, + max_samples=max_samples, + varTwindow=var_time_window_sec, + mxNn=5, + verb=args.verb, + spect_radius=None, + ) + oracle_score, oracle_score_per_state = compute_oracle_score(S_true, S_oracle, n_states) + true_dwell_sec_req = float(args.true_dwell_sec) + true_dwell_sec_eff = float(true_dwell_steps * step_size) + + evol_conf = { + "num_steps": int(args.num_steps), + "step_size": float(step_size), + "evol_time": float(args.num_steps * step_size), + "state_change_speed": float(args.state_change_speed), + "true_dwell_sec": true_dwell_sec_req, + "true_dwell_sec_eff": true_dwell_sec_eff, + "true_dwell_steps": int(true_dwell_steps), + "true_dwell_time_sec": true_dwell_sec_eff, # legacy alias + "num_states": int(n_states), + "seed": args.seed, + "state_schedule": args.schedule, + "num_states": int(n_states), + "max_samples": int(max_samples), + "truth_input_name" : args.inputStates, + } + + dale_conf = dict(dale_conf_in) + dale_conf["num_neurons"] = int(dale_conf["num_neurons"]) + dale_conf["num_excite"] = int(dale_conf["num_excite"]) + + stats_meta = dict(stats_dict) + for key in ("num_excitatory", "num_inhibitory", "num_neurons", "num_steps"): + stats_meta.pop(key, None) + + spikesD = { + "spikes": spikes.astype(np.int32), + "single_rates": rates_dict["single_rates"], + } + spikesMD = { + "data_type": "simPrism", + "time_step_sec": evol_conf_in["step_size"], + 'poisson_eta_clip': evol_conf_in['poisson_eta_clip'], + 'provenance': proven_in + } + proven_in['state_transition_file']= args.dataName + + prismTruthD = { + "A_true": daleD["A_true"], + "B_true": daleD["B_true"], + "S_true": S_true.astype(np.int32), + "C_true": C_true.astype(np.float32), + "S_oracle": S_oracle.astype(np.int32), + "state_transition": transition_matrix.astype(np.float32), + "sigle_rates_var": rates_dict["sigle_rates_var"], + "single_fano_fact": rates_dict["single_fano_fact"], + } + + + prismTruthMD = { + "short_name": args.dataName, + "data_type": "simPrism", + "var_time_window_sec": float(var_time_window_sec), + "dale_conf": dale_conf, + "evol_conf": evol_conf, + } + prismTruthMD['oracle_eval'] = { + "avr_score": oracle_score, + "score_per_state": oracle_score_per_state, + } + + + # "dale_simu_stats": [stats_meta], + outFs = os.path.join(args.outPath, args.dataName + ".spikes.npz") + outFt = os.path.join(args.outPath, args.dataName + ".prismTruth.npz") + write_data_npz(spikesD, outFs, metaD=spikesMD) + write_data_npz(prismTruthD, outFt, metaD=prismTruthMD) + + if args.verb > 1: + print('\nspikes MD:'); pprint(spikesMD) + print('\nprismTruth MD:'); pprint(prismTruthMD) + + if oracle_score is not None: + target_bins_per_state = np.bincount(S_true, minlength=n_states) + switches_to_per_state = np.zeros(n_states, dtype=np.int64) + if S_true.shape[0] > 1: + switch_idx = np.where(S_true[1:] != S_true[:-1])[0] + 1 + if switch_idx.size > 0: + switches_to_per_state = np.bincount( + S_true[switch_idx], minlength=n_states + ).astype(np.int64) + print(f"gen, oracle avr score {oracle_score:.3f}, {args.dataName}") + print(f" {'state':>5s} {'score':>5s} {'bins':>7s} {'switches_to':>11s}") + print(f" {'-----':>5s} {'-----':>5s} {'-------':>7s} {'-----------':>11s}") + for m, sc in enumerate(oracle_score_per_state): + print( + f" {m:5d} {sc:5.3f} {int(target_bins_per_state[m]):7d} " + f"{int(switches_to_per_state[m]):11d}" + ) + + print("\n ./view_spikesTrain3.py --basePath $basePath --dataName %s --idxState -1 --time_range_sec 0 15 -p b " % args.dataName) + print("\n ./prism_Estep_train.py --basePath $basePath --dataName %s " % args.dataName) + print("\n ./prism_Mstep_train3.py --basePath $basePath --dataName %s " % args.dataName) + + print(" ./fit_lassoPoisson3.py --basePath $basePath --dataName %s --num_epochs 300 " % args.dataName) + + print(" ./fitPrismEM.sh --basePath $basePath --dataName %s --num_states %d --num_em_iters 4 --m_epochs 16 --time_range_sec 0 80 " % (args.dataName,n_states)) + print(" ./bigLassoBoots.sh --basePath $basePath --dataName %s --num_epochs 100 --dropDataFrac 0.33 --num_bootstraps 2 --bootsTag b2 --desyncTime " % args.dataName) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3a_states_simu/prism_EM_eval3.py b/causal_net/nonStation_ver3a_states_simu/prism_EM_eval3.py new file mode 100755 index 00000000..a48611d8 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/prism_EM_eval3.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +""" +Evaluation and plotting for prism EM results. +""" + +import os +import argparse +import numpy as np +from pprint import pprint +from toolbox.Util_NumpyIO import read_data_npz +from PlotterPrismEM import Plotter + + +def compute_ll_gap(spikes_sub, A_true, B_true, dt, eta_clip): + """Per-time LL gap between best and 2nd-best true-state likelihood.""" + y_prev = np.asarray(spikes_sub[:-1], dtype=np.float64) + y_curr = np.asarray(spikes_sub[1:], dtype=np.float64) + a_true = np.asarray(A_true, dtype=np.float64) + b_true = np.asarray(B_true, dtype=np.float64) + + t_pairs = y_prev.shape[0] + m_states = b_true.shape[0] + ll_gap = np.zeros((t_pairs,), dtype=np.float64) + if m_states <= 1: + return ll_gap + + for t in range(t_pairs): + yp = y_prev[t] + yc = y_curr[t] + base = a_true @ yp + eta = base[None, :] + b_true + eta_c = np.minimum(eta, float(eta_clip)) + lam = np.exp(eta_c) * float(dt) + scores = np.sum(yc[None, :] * eta_c - lam, axis=1) + top2 = np.partition(scores, -2)[-2:] + ll_gap[t] = float(top2[-1] - top2[-2]) + return ll_gap + + +def eval_em_metrics_time(fitD, md, spikes): + """Compute per-time metrics for -p f canvas.""" + trainMD = md["train"] + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + dt = float(trainMD["time_step_sec"]) + eta_clip = float(trainMD["eta_clip"]) + lambda2 = float(trainMD["lambda2"]) + + spikes_sub = np.asarray(spikes[t0_bin : t1_bin + 1], dtype=np.float64) + yp = spikes_sub[:-1] + yc = spikes_sub[1:] + + a_hat = np.asarray(fitD["A_hat"], dtype=np.float64) + b_hat = np.asarray(fitD["B_hat"], dtype=np.float64) + c_hat = np.asarray(fitD["c_hat"], dtype=np.float64) + n_pairs = spikes_sub.shape[0] - 1 + assert c_hat.shape[0] == spikes_sub.shape[0], "c_hat and spikes_sub must have matching time bins" + c_pairs = c_hat[1:] + c_prev = c_hat[:-1] + + rates = np.asarray(fitD["single_rates"], dtype=np.float64) + w = 1.0 / np.maximum(rates, 0.1) + w /= w.mean() + + eta = yp @ a_hat.T + c_pairs @ b_hat + eta_c = np.minimum(eta, eta_clip) + lam = np.exp(eta_c) * dt + log_dt = np.log(dt) + nll_t = np.sum(w[None, :] * (lam - yc * (eta + log_dt)), axis=1) + + dc = c_pairs - c_prev + l2_t = lambda2 * np.sum(dc * dc, axis=1) + + ll_gap = compute_ll_gap(spikes_sub, md["A_true"], md["B_true"], dt, eta_clip) + + s_true = np.asarray(md["S_true"])[t0_bin : t1_bin + 1] + s_hat = np.asarray(fitD["S_hat"]) + assert s_true.shape[0] == s_hat.shape[0], "S_true and S_hat must have matching length" + acc = float((s_true == s_hat).mean()) + + return { + "acc": acc, + "loss_nll_time": nll_t, + "loss_l2_time": l2_t, + "ll_gap": ll_gap, + "ll_gap_mean": float(np.mean(ll_gap)), + } + + +def eval_state_recovery(fitD, md): + """Compute state recovery table on full training window.""" + trainMD = md["train"] + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + s_true = np.asarray(md["S_true"], dtype=np.int64)[t0_bin : t1_bin + 1] + s_hat = np.asarray(fitD["S_hat"], dtype=np.int64) + s_hat_cl = np.asarray(fitD["S_hat_CL"], dtype=np.float64) + m_states = int(trainMD["num_states"]) + + assert s_true.shape[0] == s_hat.shape[0] == s_hat_cl.shape[0], \ + "S_true/S_hat/S_hat_CL length mismatch on training window" + + avg_acc = float((s_true == s_hat).mean()) + state_acc_cl = [] + bins_hat = np.bincount(s_hat, minlength=m_states).astype(np.int64) + enter_hat = np.zeros(m_states, dtype=np.int64) + if s_hat.shape[0] > 1: + enter_idx = np.where(s_hat[1:] != s_hat[:-1])[0] + 1 + if enter_idx.size > 0: + enter_hat = np.bincount(s_hat[enter_idx], minlength=m_states).astype(np.int64) + + trans_hat = np.zeros((m_states, m_states), dtype=np.int64) + if s_hat.shape[0] > 1: + np.add.at(trans_hat, (s_hat[:-1], s_hat[1:]), 1) + + for m in range(m_states): + mask = (s_hat == m) + cnt = int(mask.sum()) + if cnt > 0: + acc_m = float((s_true[mask] == m).mean()) + cl_m = float(s_hat_cl[mask].mean()) + else: + acc_m = float("nan") + cl_m = float("nan") + state_acc_cl.append([acc_m, cl_m, int(bins_hat[m]), int(enter_hat[m])]) # [acc, CL, bins_hat, enter_hat] + + return { + "avg_acc": avg_acc, + "state_acc_cl": state_acc_cl, + "trans_hat": trans_hat, + } + + +def main(): + parser = argparse.ArgumentParser( + description="Evaluate and plot prism EM results") + parser.add_argument("--dataName", type=str, required=True, + help="Base name for prismEM file") + parser.add_argument("--basePath", type=str, + default="/pscratch/sd/b/balewski/2026_causalNet_tmp3/", + help="Head dir for input/output data") + parser.add_argument("-p", "--showPlots", type=str, nargs='+', + default="a", + help="Plot types: a=EM convergence summary, b=init-vs-truth states, c=A_init-vs-truth, d=A_hat-vs-truth, e=A_hat edge recovery, f=state sequence, g=2D correlations (A/B), h=A_init TP quality (diag/exc/inh)") + parser.add_argument("--minW", type=float, default=0.02, + help="Threshold for A-matrix edge eval") + parser.add_argument("--timeReb", type=int, default=20, + help="Time rebin factor for time-axis plots") + g = parser.add_argument_group("data") + g.add_argument("-T", "--time_range_sec", default=[0.0, 15.0], + nargs=2, type=float, + help="Time window [t0, t1] in seconds") + parser.add_argument("-X", "--noXterm", action="store_true", + help="Disable X terminal for plotting") + parser.add_argument("-v", "--verb", type=int, default=1, + help="Verbosity level") + args = parser.parse_args() + + args.inpPath = os.path.join(args.basePath, "prismFit") + args.outPath = os.path.join(args.basePath, "plots") + os.makedirs(args.outPath, exist_ok=True) + args.showPlots = ''.join(args.showPlots) + + print("EM-eval args:", vars(args), "\n") + + # ── load EM fit ────────────────────────────────────────────────── + fitFF = os.path.join(args.inpPath, f"{args.dataName}.prismEM.npz") + fitD, fitMD = read_data_npz(fitFF) + assert isinstance(fitMD, dict), "Expected metadata dict in prismEM file" + + if args.verb > 1: pprint(fitMD) + + MD = {**fitMD, "short_name": args.dataName} + + # ── load ground truth if available ─────────────────────────────── + prov = fitMD["provenance"] + truth_name = prov["state_model_file"] + truthPath = os.path.join(args.basePath, "truthDale") + truthFF = os.path.join(truthPath, f"{truth_name}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 1) + if args.verb > 1: pprint(trueMD) + MD.update(trueMD) + MD["A_true"] = trueD["A_true"] + MD["B_true"] = trueD["B_true"] + MD["E_true"] = trueD["E_true"] + + st_name = prov["state_transition_file"] + ptFF = os.path.join(args.basePath, "spikesData", f"{st_name}.prismTruth.npz") + trD, _ = read_data_npz(ptFF, verb=args.verb > 1) + MD["S_true"] = trD["S_true"] + MD["C_true"] = trD["C_true"] + + spikesFF = os.path.join(args.basePath, "spikesData", f"{st_name}.spikes.npz") + spikeD, _ = read_data_npz(spikesFF, verb=args.verb > 1) + spikes = spikeD["spikes"] + + MD["eval_f"] = eval_em_metrics_time(fitD, MD, spikes) + reco = eval_state_recovery(fitD, MD) + MD["states_recovery_eval"]["avg_acc"] = reco["avg_acc"] + MD["states_recovery_eval"]["state_acc_cl"] = reco["state_acc_cl"] + + print(f"state reco avr acc {reco['avg_acc']:.3f}, {args.dataName}") + print(f" {'state':>5s} {'acc':>5s} {'CL':>6s} {'bins_hat':>8s} {'enter_hat':>9s}") + print(f" {'-----':>5s} {'-----':>5s} {'------':>6s} {'--------':>8s} {'---------':>9s}") + for m, (acc_m, cl_m, bins_m, enter_m) in enumerate(reco["state_acc_cl"]): + print(f" {m:5d} {acc_m:5.3f} {cl_m:6.3f} {bins_m:8d} {enter_m:9d}") + + trans_hat = np.asarray(reco["trans_hat"], dtype=np.int64) + n_states = trans_hat.shape[0] + print("\nS_hat transition counts (from row -> to col):") + hdr = " from\\to" + "".join(f"{j:8d}" for j in range(n_states)) + print(hdr) + print(" " + "-" * (len(hdr) - 1)) + for i in range(n_states): + row = f"{i:8d}" + "".join(f"{int(trans_hat[i, j]):8d}" for j in range(n_states)) + print(row) + + MD["short_name"] = args.dataName + + # ── plot ────────────────────────────────────────────────────────── + args.prjName = args.dataName + plot = Plotter(args) + + if 'a' in args.showPlots: + plot.summary_prismEM(fitD, MD, figId=1) + + if 'b' in args.showPlots: + plot.state_init_prismEM(fitD, MD, figId=2, time_reb=args.timeReb) + + if 'c' in args.showPlots: + plot.matrix_init_prismEM(fitD, MD, figId=3, est_key="A_init", est_label="A_init") + + if 'd' in args.showPlots: + plot.matrix_init_prismEM(fitD, MD, figId=3, est_key="A_hat", est_label="A_hat") + + if 'e' in args.showPlots: + plot.edge_recovery_prismEM( + fitD, MD, minW=args.minW, figId=4, est_key="A_hat", est_label="A_hat" + ) + + if 'f' in args.showPlots: + plot.state_seq_prismEM( + fitD, MD, figId=5, time_range_sec=args.time_range_sec + ) + + if 'g' in args.showPlots: + plot.eval_ABcorr_prismEM(fitD, MD, figId=6) + + if 'h' in args.showPlots: + plot.initA_quality_prismEM(fitD, MD, minW=args.minW, figId=7) + + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3a_states_simu/prism_EM_train3.py b/causal_net/nonStation_ver3a_states_simu/prism_EM_train3.py new file mode 100755 index 00000000..ab7d744d --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/prism_EM_train3.py @@ -0,0 +1,851 @@ +#!/usr/bin/env python3 +""" +prism_EM_train.py — EM training for non-stationary Poisson GLM. + +Jointly infers from observed spikes only (no ground truth): + - shared connectivity matrix A (N×N) + - per-state bias vectors B (M×N) + - time-varying coefficients c_hat (T×M) on the probability simplex + +Algorithm: Block Coordinate Descent + E-step: single forward sweep PGD over time bins → update c_hat + M-step: m_epochs of Adam with DataLoader → update (A, B) + +Uses uniformly weighted Poisson NLL in both steps, L1 off-diagonal +soft-thresholding, and hard spectral radius projection on A. +Numerical tricks and defaults follow prism_Mstep_train3.py (M-step) +and prism_Estep_train.py (E-step). + +Usage: + ./prism_EM_train.py --dataName --basePath $basePath -M 3 +""" + +import os +import time +import math +import secrets +import argparse + +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader +from torch.utils.data.distributed import DistributedSampler +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from Util_PrismEM import init_states_vs_time, init_B_from_spikes, init_A_from_spikes + + +# ═══════════════════════════════════════════════════════════════════════ +# Model (from prism_Mstep_train3.py) +# ═══════════════════════════════════════════════════════════════════════ + +class SwitchingPoissonGLM(nn.Module): + """Shared A (N×N), per-state B (M×N), mixed by coefficients c.""" + + def __init__(self, N, M, eta_clip): + super().__init__() + self.N = N + self.M = M + self.eta_clip = eta_clip + self.A = nn.Parameter(torch.randn(N, N) * 0.1) + self.B = nn.Parameter(torch.randn(M, N) * 0.1) + + def forward(self, y_prev, c, dt): + """y_prev (batch, N), c (batch, M) → predicted rate (batch, N).""" + eta = y_prev @ self.A.t() + c @ self.B + return torch.exp(torch.clamp(eta, max=self.eta_clip)) * dt + + +# ═══════════════════════════════════════════════════════════════════════ +# Dataset (c array updated in-place after each E-step) +# ═══════════════════════════════════════════════════════════════════════ + +class PairDataset(Dataset): + """(y_prev, y_curr, c) triplets. + + self.c is a mutable numpy array updated in-place after each E-step. + """ + + def __init__(self, y_prev, y_curr, c): + self.y_prev = y_prev + self.y_curr = y_curr + self.c = c + + def __len__(self): + return self.y_prev.shape[0] + + def __getitem__(self, idx): + return ( + torch.from_numpy(self.y_prev[idx]).float(), + torch.from_numpy(self.y_curr[idx]).float(), + torch.from_numpy(self.c[idx]).float(), + ) + + +def init_distributed(): + """Initialize torch.distributed from torchrun environment.""" + is_dist = (int(os.environ.get("WORLD_SIZE", "1")) > 1) or ("RANK" in os.environ) + if is_dist: + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + rank = 0 + world_size = 1 + local_rank = 0 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + return is_dist, rank, world_size, local_rank, device + + +def cleanup_distributed(is_dist): + if is_dist and dist.is_initialized(): + dist.destroy_process_group() + + +def pair_range_for_rank(t_pairs, world_size, rank): + """Return inclusive pair-index range [p0, p1] assigned to given rank.""" + base = t_pairs // world_size + rem = t_pairs % world_size + n_loc = base + (1 if rank < rem else 0) + p0 = rank * base + min(rank, rem) + p1 = p0 + n_loc - 1 + return p0, p1, n_loc + + +# ═══════════════════════════════════════════════════════════════════════ +# E-step: simplex PGD (from prism_Estep_train.py) +# ═══════════════════════════════════════════════════════════════════════ + +def project_to_simplex(v): + """Project 1-D vector onto the probability simplex (Duchi et al.).""" + if v.numel() == 1: + return torch.ones_like(v) + u, _ = torch.sort(v, descending=True) + cssv = torch.cumsum(u, dim=0) - 1.0 + ind = torch.arange(1, v.numel() + 1, device=v.device, dtype=v.dtype) + cond = u - cssv / ind > 0 + if torch.any(cond): + rho = torch.nonzero(cond, as_tuple=False)[-1, 0] + theta = cssv[rho] / (rho + 1.0) + else: + theta = cssv[-1] / v.numel() + return torch.clamp(v - theta, min=0.0) + + +def run_estep(Y_prev, Y_curr, A, B, c_hat, + dt, eta_clip, lambda2, lr, pgd_iter): + """Single forward sweep E-step. Updates c_hat in-place on GPU. + + Args: + Y_prev, Y_curr: (T_pairs, N) GPU tensors + A: (N, N) detached parameter tensor + B: (M, N) detached parameter tensor + c_hat: (T_full, M) GPU tensor, modified in-place + + Returns: + mean NLL over all time pairs + """ + T_eff = c_hat.shape[0] + log_dt = math.log(dt) + c_prev = c_hat[0].clone() + nll_sum = 0.0 + + for t in range(1, T_eff): + y_p = Y_prev[t - 1] + y_c = Y_curr[t - 1] + + Z = (A @ y_p)[None, :] + B # (M, N) + + c_t = c_hat[t].clone() + for _ in range(pgd_iter): + eta = c_t @ Z # (N,) + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + g = Z @ (lam - y_c) + 2.0 * lambda2 * (c_t - c_prev) + c_t = project_to_simplex(c_t - lr * g) + + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + nll_sum += float((lam - y_c * (eta_c + log_dt)).sum().item()) + + c_hat[t] = c_t + c_prev = c_t + + return nll_sum / max(1, T_eff - 1) + + +def run_estep_shard(Y_prev, Y_curr, A, B, c_hat, + dt, eta_clip, lambda2, lr, pgd_iter, + t0, t1): + """E-step update on a shard of time bins t in [t0, t1], inclusive. + + Returns: + nll_sum_local, n_pairs_local + """ + if t1 < t0: + return 0.0, 0 + log_dt = math.log(dt) + c_prev = c_hat[t0 - 1].clone() + nll_sum = 0.0 + n_pairs = 0 + + for t in range(t0, t1 + 1): + y_p = Y_prev[t - 1] + y_c = Y_curr[t - 1] + + Z = (A @ y_p)[None, :] + B + + c_t = c_hat[t].clone() + for _ in range(pgd_iter): + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + g = Z @ (lam - y_c) + 2.0 * lambda2 * (c_t - c_prev) + c_t = project_to_simplex(c_t - lr * g) + + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + nll_sum += float((lam - y_c * (eta_c + log_dt)).sum().item()) + + c_hat[t] = c_t + c_prev = c_t + n_pairs += 1 + + return nll_sum, n_pairs + + +# ═══════════════════════════════════════════════════════════════════════ +# M-step helpers (from prism_Mstep_train3.py) +# ═══════════════════════════════════════════════════════════════════════ + +def offdiag_soft_threshold_(A, lr, lam): + """In-place proximal L1 on off-diagonal A elements.""" + if lam <= 0: + return + with torch.no_grad(): + n = A.shape[0] + mask = ~torch.eye(n, dtype=torch.bool, device=A.device) + thresh = lr * lam + v = A[mask] + A[mask] = v.sign() * (v.abs() - thresh).clamp(min=0.0) + + +def enforce_spectral_radius_(A, rho_max): + """Hard in-place spectral projection. Returns rho before projection.""" + with torch.no_grad(): + rho = torch.linalg.eigvals(A).abs().max().item() + if rho > rho_max: + A.mul_(rho_max / rho) + return rho + + +def sync_AB_via_rank0_avg_then_broadcast_(mdl, rho_max): + """Rank-0 averages A/B, enforces rho on averaged A, then broadcasts A/B.""" + if not (dist.is_available() and dist.is_initialized()): + enforce_spectral_radius_(mdl.A, rho_max) + return + + rank = dist.get_rank() + world = float(dist.get_world_size()) + with torch.no_grad(): + A_buf = mdl.A.data.clone() + B_buf = mdl.B.data.clone() + + dist.reduce(A_buf, dst=0, op=dist.ReduceOp.SUM) + dist.reduce(B_buf, dst=0, op=dist.ReduceOp.SUM) + + if rank == 0: + A_buf.div_(world) + B_buf.div_(world) + enforce_spectral_radius_(A_buf, rho_max) + + dist.broadcast(A_buf, src=0) + dist.broadcast(B_buf, src=0) + mdl.A.data.copy_(A_buf) + mdl.B.data.copy_(B_buf) + + +def run_mstep_epoch(model, loader, optimizer, device, dt, + l1_wt, lambda3, rho_max, + rho_every, apply_prune, apply_rho, off_mask, minW): + """One DataLoader pass updating A and B. Returns metrics dict. + """ + model.train() + mdl = model.module if hasattr(model, "module") else model + s_tot = s_nll = s_l1 = 0.0 + eps = 1e-8 + + for bi, (yp, yc, cc) in enumerate(loader): + yp = yp.to(device, non_blocking=True) + yc = yc.to(device, non_blocking=True) + cc = cc.to(device, non_blocking=True) + + optimizer.zero_grad(set_to_none=True) + pred = model(yp, cc, dt) + nll = (-yc * torch.log(pred + eps) + pred).mean() + if lambda3 > 0: + n = mdl.A.shape[0] + off_abs_mean = (mdl.A.abs() * l1_wt).sum() / float(n * (n - 1)) + l1 = lambda3 * off_abs_mean + else: + l1 = torch.tensor(0.0, device=device) + loss = nll + l1 + loss.backward() + optimizer.step() + + if apply_prune and lambda3 > 0: + offdiag_soft_threshold_(mdl.A, optimizer.param_groups[0]["lr"], lambda3) + if apply_rho and (bi % rho_every == 0): + sync_AB_via_rank0_avg_then_broadcast_(mdl, rho_max) + + s_tot += loss.item() + s_nll += nll.item() + s_l1 += l1.item() + + nb = max(1, len(loader)) + if dist.is_available() and dist.is_initialized(): + v = torch.tensor([s_tot, s_nll, s_l1, float(nb)], + dtype=torch.float64, device=device) + dist.all_reduce(v, op=dist.ReduceOp.SUM) + s_tot = float(v[0].item()) + s_nll = float(v[1].item()) + s_l1 = float(v[2].item()) + nb = int(v[3].item()) + + with torch.no_grad(): + nz = int((mdl.A[off_mask].abs() > minW).sum().item()) + rho = float(torch.linalg.eigvals(mdl.A).abs().max().item()) + + return dict(loss=s_tot / nb, nll=s_nll / nb, l1=s_l1 / nb, rho=rho, nz=nz) + + +# ═══════════════════════════════════════════════════════════════════════ +# Viterbi (from prism_Estep_train.py) +# ═══════════════════════════════════════════════════════════════════════ + +def viterbi_decode(c_hat, p_stay): + """Viterbi decode: find most likely state path from simplex coefficients.""" + eps = 1e-12 + T, M = c_hat.shape + if M == 1: + return np.zeros(T, dtype=np.int64) + p_sw = (1.0 - p_stay) / (M - 1) + trans = np.full((M, M), p_sw, dtype=np.float64) + np.fill_diagonal(trans, p_stay) + lt = np.log(np.clip(trans, eps, 1.0)) + le = np.log(np.clip(c_hat, eps, 1.0)) + + delta = np.zeros((T, M), dtype=np.float64) + psi = np.zeros((T, M), dtype=np.int64) + delta[0] = le[0] + for t in range(1, T): + sc = delta[t - 1][:, None] + lt + psi[t] = np.argmax(sc, axis=0) + delta[t] = sc[psi[t], np.arange(M)] + le[t] + + path = np.zeros(T, dtype=np.int64) + path[-1] = int(np.argmax(delta[-1])) + for t in range(T - 2, -1, -1): + path[t] = psi[t + 1, path[t + 1]] + return path + + +# ═══════════════════════════════════════════════════════════════════════ +# Args +# ═══════════════════════════════════════════════════════════════════════ + +def parse_args(): + p = argparse.ArgumentParser( + description="EM training: non-stationary Poisson GLM") + + p.add_argument("--dataName", required=True) + p.add_argument("--basePath", + default="/pscratch/sd/b/balewski/2026_causalNet_tmp3/") + p.add_argument("--num_states", "-M", type=int, required=True, + help="Number of latent states") + + g = p.add_argument_group("EM structure") + g.add_argument("--num_em_iters", type=int, default=20, + help="Outer EM iterations") + g.add_argument("--m_epochs", type=int, default=100, + help="M-step Adam epochs per EM iteration") + + g = p.add_argument_group("E-step") + g.add_argument("--pgd_iter", type=int, default=5, + help="PGD iterations per time bin") + g.add_argument("--lr_estep", type=float, default=0.03, + help="PGD step size") + g.add_argument("--lambda2", type=float, default=2.0, + help="Temporal smoothness weight") + + g = p.add_argument_group("M-step") + g.add_argument("--lr_mstep", type=float, default=0.003, + help="Adam learning rate (initial)") + g.add_argument("--lr_end_factor", type=float, default=0.1, + help="LR decays to lr_mstep * lr_end_factor") + g.add_argument("--lambda3", type=float, default=0.02, + help="L1 penalty on off-diagonal A") + g.add_argument("--rho_max", type=float, default=0.95, + help="Hard spectral radius ceiling for A") + g.add_argument("--prescale_m_step_4_ArhoMax", type=int, default=50, + help="Spectral projection frequency (batches)") + g.add_argument("--delay_em_iter_4_ArhoMax", type=int, default=3, + help="EM iter to start rho_max enforcement " + "(default: int(num_em_iters * 0.6))") + g.add_argument("--delay_em_iter_4_lrDecay", type=int, default=1, + help="EM iter to start LR decay " + "(default: int(num_em_iters * 0.7))") + g.add_argument("--delay_em_iter_4_Aprune", type=int, default=1, + help="EM iter to start L1 pruning " + "(default: num_em_iters // 3)") + g.add_argument("--batch_size", type=int, default=2048) + g.add_argument("--minW", type=float, default=0.01, + help="Threshold for edge counting (reporting only)") + g = p.add_argument_group("data") + g.add_argument("-T", "--time_range_sec", default=[0.0, 60.0], + nargs=2, type=float, + help="Time window [t0, t1] in seconds") + + g = p.add_argument_group("decode") + g.add_argument("--decode_dwell_sec", type=float, default=0.3, + help="Viterbi dwell prior τ_dwell") + + g = p.add_argument_group("misc") + g.add_argument("--seed", type=int, default=42) + g.add_argument("--fitName", type=str, default=None, + help="Optional output fit name; if omitted a random name is generated") + g.add_argument("--init_states", type=str, default="data", + choices=["data", "rand"], + help="Initial latent-state mode: 'data' or 'rand'") + g.add_argument("--init_B", type=str, default="data", + choices=["data", "rand"], + help="Initial B mode: 'data' or 'rand'") + g.add_argument("--init_A", type=str, default="data", + choices=["data", "rand"], + help="Initial A mode: 'data' or 'rand'") + g.add_argument("-v", "--verb", type=int, default=1) + + return p.parse_args() + + +# ═══════════════════════════════════════════════════════════════════════ +# Main +# ═══════════════════════════════════════════════════════════════════════ + +def main(): + args = parse_args() + is_dist, rank, world_size, local_rank, device = init_distributed() + inpPath = os.path.join(args.basePath, "spikesData") + outPath = os.path.join(args.basePath, "prismFit") + out_ok = True + out_err = "" + if rank == 0: + out_ok = os.path.exists(outPath) + if not out_ok: + out_err = f"Output dir missing: {outPath}" + if is_dist: + msg = [out_ok, out_err] + dist.broadcast_object_list(msg, src=0) + out_ok, out_err = msg + if not out_ok: + raise FileNotFoundError(out_err) + + if args.delay_em_iter_4_Aprune is None: + args.delay_em_iter_4_Aprune = int(args.num_em_iters * 0.3) + if args.delay_em_iter_4_ArhoMax is None: + args.delay_em_iter_4_ArhoMax = int(args.num_em_iters * 0.6) + if args.delay_em_iter_4_lrDecay is None: + args.delay_em_iter_4_lrDecay = int(args.num_em_iters * 0.7) + + if rank == 0: + print("\nEM-train args:", vars(args), "\n") + print(f"world_size={world_size}") + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision("high") + + # ── load spikes (rank 0 only), then broadcast ─────────────── + if rank == 0: + spikesFF = os.path.join(inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 0) + spikes = spikeD["spikes"] + single_rates = spikeD["single_rates"] + else: + spikeMD = None + spikes = None + single_rates = None + + if is_dist: + obj = [spikeMD] + dist.broadcast_object_list(obj, src=0) + spikeMD = obj[0] + + if rank == 0: + shp = torch.tensor(spikes.shape, dtype=torch.int64, device=device) + else: + shp = torch.zeros(2, dtype=torch.int64, device=device) + dist.broadcast(shp, src=0) + T_raw = int(shp[0].item()) + N = int(shp[1].item()) + + if rank == 0: + spikes_t = torch.as_tensor(spikes, dtype=torch.int32, device=device).contiguous() + else: + spikes_t = torch.empty((T_raw, N), dtype=torch.int32, device=device) + dist.broadcast(spikes_t, src=0) + spikes = spikes_t.cpu().numpy() + else: + T_raw, N = spikes.shape + + dt = float(spikeMD["time_step_sec"]) + eta_clip = float(spikeMD["poisson_eta_clip"]) + prov = dict(spikeMD["provenance"]) + M = args.num_states + + # ── time range selection ───────────────────────────────────── + t0_sec, t1_sec = float(args.time_range_sec[0]), float(args.time_range_sec[1]) + if t1_sec < t0_sec: + t0_sec, t1_sec = t1_sec, t0_sec + start_bin = max(0, int(math.floor(t0_sec / dt))) + end_bin = min(T_raw - 1, int(math.floor(t1_sec / dt))) + if end_bin <= start_bin: + raise ValueError("time_range_sec too small; need at least two bins") + spikes = spikes[start_bin : end_bin + 1] + T_full = spikes.shape[0] + T_pairs = T_full - 1 + if rank == 0: + print(f"N={N} EM={args.num_em_iters} M={M} T={T_full} pairs={T_pairs} " + f"dt={dt} eta_clip={eta_clip} " + f"time=[{t0_sec:.1f}, {t1_sec:.1f}]s bins=[{start_bin}, {end_bin}]") + + # ── time pairs: GPU for E-step, CPU numpy for DataLoader ──── + yp_np = spikes[:-1].astype(np.float32) + yc_np = spikes[1:].astype(np.float32) + Yp_gpu = torch.tensor(yp_np, device=device) + Yc_gpu = torch.tensor(yc_np, device=device) + + # ── initial states and parameter seeds (rank 0 computes, all ranks receive) ── + if rank == 0: + c_init_np, S_init, init_state_md, freq_h1d = init_states_vs_time(spikes, dt, args) + A_seed_np, init_A_md = init_A_from_spikes(spikes, args) + B_seed_np, init_B_md = init_B_from_spikes(spikes, dt, args) + else: + c_init_np = np.empty((T_full, M), dtype=np.float32) + S_init = np.empty((T_full,), dtype=np.int64) + freq_h1d = np.empty((0,), dtype=np.float32) + A_seed_np = None + B_seed_np = None + init_state_md = None + init_A_md = None + init_B_md = None + + if is_dist: + c_t = torch.as_tensor(c_init_np, dtype=torch.float32, device=device) + s_t = torch.as_tensor(S_init, dtype=torch.int64, device=device) + dist.broadcast(c_t, src=0) + dist.broadcast(s_t, src=0) + c_init_np = c_t.cpu().numpy() + S_init = s_t.cpu().numpy() + + f_len_t = torch.tensor([int(freq_h1d.shape[0]) if rank == 0 else 0], + dtype=torch.int64, device=device) + dist.broadcast(f_len_t, src=0) + f_len = int(f_len_t.item()) + if rank != 0: + freq_h1d = np.empty((f_len,), dtype=np.float32) + f_t = (torch.as_tensor(freq_h1d, dtype=torch.float32, device=device) + if rank == 0 else + torch.empty((f_len,), dtype=torch.float32, device=device)) + if f_len > 0: + dist.broadcast(f_t, src=0) + freq_h1d = f_t.cpu().numpy() + + has_A_t = torch.tensor([1 if (rank == 0 and A_seed_np is not None) else 0], + dtype=torch.int64, device=device) + dist.broadcast(has_A_t, src=0) + if int(has_A_t.item()) == 1: + A_t = (torch.as_tensor(A_seed_np, dtype=torch.float32, device=device) + if rank == 0 else + torch.empty((N, N), dtype=torch.float32, device=device)) + dist.broadcast(A_t, src=0) + A_seed_np = A_t.cpu().numpy() + else: + A_seed_np = None + + has_B_t = torch.tensor([1 if (rank == 0 and B_seed_np is not None) else 0], + dtype=torch.int64, device=device) + dist.broadcast(has_B_t, src=0) + if int(has_B_t.item()) == 1: + B_t = (torch.as_tensor(B_seed_np, dtype=torch.float32, device=device) + if rank == 0 else + torch.empty((M, N), dtype=torch.float32, device=device)) + dist.broadcast(B_t, src=0) + B_seed_np = B_t.cpu().numpy() + else: + B_seed_np = None + + c_hat_gpu = torch.tensor(c_init_np, dtype=torch.float32, device=device) + c_pairs_np = c_init_np[1:].copy() + + dataset = PairDataset(yp_np, yc_np, c_pairs_np) + if is_dist: + sampler = DistributedSampler( + dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False + ) + loader = DataLoader( + dataset, batch_size=args.batch_size, shuffle=False, sampler=sampler, + drop_last=False, pin_memory=True, num_workers=0 + ) + else: + sampler = None + loader = DataLoader( + dataset, batch_size=args.batch_size, shuffle=True, + drop_last=False, pin_memory=True, num_workers=0 + ) + + # ── model + optimizer + scheduler ──────────────────────────── + model = SwitchingPoissonGLM(N, M, eta_clip).to(device) + if A_seed_np is not None: + with torch.no_grad(): + model.A.copy_(torch.tensor(A_seed_np, dtype=torch.float32, device=device)) + if B_seed_np is not None: + with torch.no_grad(): + model.B.copy_(torch.tensor(B_seed_np, dtype=torch.float32, device=device)) + A_init_np = model.A.detach().cpu().numpy().copy() + B_init_np = model.B.detach().cpu().numpy().copy() + B_init_out_np = B_init_np[0] if B_init_np.ndim == 2 and B_init_np.shape[0] > 1 else B_init_np + + if is_dist: + model = DDP(model, device_ids=[local_rank]) + + total_m_epochs = args.num_em_iters * args.m_epochs + decay_start_m_epoch = args.delay_em_iter_4_lrDecay * args.m_epochs + decay_m_epochs = max(1, total_m_epochs - decay_start_m_epoch) + optimizer = optim.Adam(model.parameters(), lr=args.lr_mstep, fused=True) + scheduler = optim.lr_scheduler.LinearLR( + optimizer, start_factor=1.0, end_factor=args.lr_end_factor, + total_iters=decay_m_epochs, last_epoch=-1 + ) + scheduler._step_count = 1 + + mdl = model.module if hasattr(model, "module") else model + off_mask = ~torch.eye(N, dtype=torch.bool, device=device) + l1_wt = torch.ones(N, N, device=device) + l1_wt[torch.eye(N, dtype=torch.bool, device=device)] = 0.0 + + # ── history ────────────────────────────────────────────────── + h_e_nll = [] + h_m_loss, h_m_nll, h_m_l1 = [], [], [] + h_rho, h_nz, h_lr = [], [], [] + + t_start = time.time() + m_epoch_global = 0 + p_stay = math.exp(-dt / args.decode_dwell_sec) + onehot_states = np.eye(M, dtype=np.float32) + + # ══════════════════════════════════════════════════════════════ + # EM loop + # ══════════════════════════════════════════════════════════════ + for em in range(1, args.num_em_iters + 1): + if rank == 0: + if em == args.delay_em_iter_4_ArhoMax + 1: + print(f"threshold passed: 'delay_em_iter_4_ArhoMax': {args.delay_em_iter_4_ArhoMax} (em={em})") + if em == args.delay_em_iter_4_lrDecay + 1: + print(f"threshold passed: 'delay_em_iter_4_lrDecay': {args.delay_em_iter_4_lrDecay} (em={em})") + if em == args.delay_em_iter_4_Aprune + 1: + print(f"threshold passed: 'delay_em_iter_4_Aprune': {args.delay_em_iter_4_Aprune} (em={em})") + apply_prune = em > args.delay_em_iter_4_Aprune + apply_rho = em > args.delay_em_iter_4_ArhoMax + + # ── E-step ─────────────────────────────────────────────── + te0 = time.time() + with torch.no_grad(): + if is_dist: + p0, p1, n_loc = pair_range_for_rank(T_pairs, world_size, rank) + t0 = p0 + 1 + t1 = p1 + 1 + nll_local, n_pair_local = run_estep_shard( + Yp_gpu, Yc_gpu, mdl.A.data, mdl.B.data, c_hat_gpu, + dt, eta_clip, args.lambda2, args.lr_estep, args.pgd_iter, + t0, t1 + ) + + c_upd = torch.zeros_like(c_hat_gpu) + c_msk = torch.zeros((T_full, 1), dtype=torch.float32, device=device) + c_upd[0] = c_hat_gpu[0] + c_msk[0] = 1.0 + if n_loc > 0: + c_upd[t0:t1 + 1] = c_hat_gpu[t0:t1 + 1] + c_msk[t0:t1 + 1] = 1.0 + + dist.all_reduce(c_upd, op=dist.ReduceOp.SUM) + dist.all_reduce(c_msk, op=dist.ReduceOp.SUM) + c_avg = c_upd / torch.clamp(c_msk, min=1.0) + c_hat_gpu.copy_(torch.where(c_msk > 0.0, c_avg, c_hat_gpu)) + + ev = torch.tensor([nll_local, float(n_pair_local)], + dtype=torch.float64, device=device) + dist.all_reduce(ev, op=dist.ReduceOp.SUM) + e_nll = float(ev[0].item() / max(1.0, ev[1].item())) + else: + e_nll = run_estep( + Yp_gpu, Yc_gpu, mdl.A.data, mdl.B.data, + c_hat_gpu, dt, eta_clip, + args.lambda2, args.lr_estep, args.pgd_iter + ) + + te = time.time() - te0 + h_e_nll.append(e_nll) + + # sync E-step states to CPU dataset (in-place update) for M-step: + # Viterbi decode -> one-hot rows (classification-style M-step). + c_hat_np_iter = c_hat_gpu.cpu().numpy() + if is_dist: + if rank == 0: + s_iter_np = viterbi_decode(c_hat_np_iter, p_stay).astype(np.int64, copy=False) + s_iter_t = torch.as_tensor(s_iter_np, dtype=torch.int64, device=device) + else: + s_iter_t = torch.empty((T_full,), dtype=torch.int64, device=device) + dist.broadcast(s_iter_t, src=0) + s_iter_np = s_iter_t.cpu().numpy() + else: + s_iter_np = viterbi_decode(c_hat_np_iter, p_stay).astype(np.int64, copy=False) + np.copyto(c_pairs_np, onehot_states[s_iter_np[1:]]) + + # ── M-step: m_epochs of Adam ───────────────────────────── + tm0 = time.time() + for _ in range(args.m_epochs): + m_epoch_global += 1 + if sampler is not None: + sampler.set_epoch(m_epoch_global) + met = run_mstep_epoch( + model, loader, optimizer, device, dt, + l1_wt, args.lambda3, args.rho_max, + args.prescale_m_step_4_ArhoMax, apply_prune, apply_rho, + off_mask, args.minW + ) + if em > args.delay_em_iter_4_lrDecay: + scheduler.step() + + h_m_loss.append(met["loss"]) + h_m_nll.append(met["nll"]) + h_m_l1.append(met["l1"]) + h_rho.append(met["rho"]) + h_nz.append(met["nz"]) + h_lr.append(float(optimizer.param_groups[0]["lr"])) + tm = time.time() - tm0 + + if rank == 0 and args.verb > 0: + n_off = int(off_mask.sum().item()) + sp = 1.0 - met["nz"] / max(1, n_off) + print( + f"EM {em:3d}/{args.num_em_iters} " + f"E_nll={e_nll:.4e}({te:.1f}s) " + f"M_nll={met['nll']:.4e} l1={met['l1']:.4e} " + f"rho={met['rho']:.4f} sp={sp:.3f} nz={met['nz']} " + f"lr={h_lr[-1]:.2e} ({tm:.1f}s) " + f"tot={time.time() - t_start:.0f}s" + ) + + if rank == 0: + print(f"\nTotal EM time: {time.time() - t_start:.1f}s") + + # ── Final Viterbi decode ───────────────────────────────── + c_hat_np = c_hat_gpu.cpu().numpy() + S_hat = viterbi_decode(c_hat_np, p_stay) + S_hat_CL = (1.0 - c_hat_np.max(axis=1)).astype(np.float32) + + mdl = model.module if hasattr(model, "module") else model + A_hat = mdl.A.detach().cpu().numpy() + B_hat = mdl.B.detach().cpu().numpy() + + if args.fitName is None: + h6 = secrets.token_hex(3) + outF = f"{args.dataName}-EM-{h6}" + else: + outF = args.fitName + outFF = os.path.join(outPath, f"{outF}.prismEM.npz") + prov["EMtrain_file"] = outF + + outD = { + "A_init": A_init_np.astype(np.float32), + "A_hat": A_hat.astype(np.float32), + "B_init": B_init_out_np.astype(np.float32), + "B_hat": B_hat.astype(np.float32), + "freq_h1d": np.asarray(freq_h1d, dtype=np.float32), + "c_init": c_init_np.astype(np.float32), + "c_hat": c_hat_np.astype(np.float32), + "S_init": S_init.astype(np.int64), + "S_hat": S_hat.astype(np.int64), + "S_hat_CL": S_hat_CL, + "single_rates": np.asarray(single_rates), + "e_nll_em": np.asarray(h_e_nll, dtype=np.float64), + "m_loss_epoch": np.asarray(h_m_loss, dtype=np.float64), + "m_nll_epoch": np.asarray(h_m_nll, dtype=np.float64), + "m_l1_epoch": np.asarray(h_m_l1, dtype=np.float64), + "rho_epoch": np.asarray(h_rho, dtype=np.float64), + "nz_edges_epoch": np.asarray(h_nz, dtype=np.int64), + "learning_rates": np.asarray(h_lr, dtype=np.float64), + } + + outMD = dict(spikeMD) + outMD["fit_type"] = "prismEM" + outMD["train"] = { + "num_em_iters": args.num_em_iters, + "m_epochs": args.m_epochs, + "total_m_epochs": total_m_epochs, + "pgd_iter": args.pgd_iter, + "lr_estep": args.lr_estep, + "lr_mstep": args.lr_mstep, + "lr_end_factor": args.lr_end_factor, + "lambda2": args.lambda2, + "lambda3": args.lambda3, + "rho_max": args.rho_max, + "prescale_m_step_4_ArhoMax": args.prescale_m_step_4_ArhoMax, + "delay_em_iter_4_ArhoMax": args.delay_em_iter_4_ArhoMax, + "rho_sync_mode": "broadcast", + "delay_em_iter_4_lrDecay": args.delay_em_iter_4_lrDecay, + "delay_em_iter_4_Aprune": args.delay_em_iter_4_Aprune, + "mstep_state_mode": "viterbi_onehot", + "batch_size": args.batch_size, + "minW": args.minW, + "num_states": M, + "num_neurons": N, + "num_time_bins": T_full, + "time_step_sec": dt, + "eta_clip": eta_clip, + "time_range_sec": [t0_sec, t1_sec], + "time_range_bins": [start_bin, end_bin], + "seed": args.seed, + } + outMD["states_recovery_eval"] = { + "decode": "viterbi", + "decode_dwell_sec": args.decode_dwell_sec, + } + outMD["init_A"] = init_A_md + outMD["init_state"] = init_state_md + outMD["init_B"] = init_B_md + outMD["provenance"] = prov + + write_data_npz(outD, outFF, metaD=outMD) + print(f"\nSaved: {outFF}") + print(f" basePath={args.basePath}") + print(f" ./prism_EM_eval.py --basePath $basePath " + f"--dataName {outF} -p a e f g \n") + cleanup_distributed(is_dist) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3a_states_simu/prism_Estep_eval.py b/causal_net/nonStation_ver3a_states_simu/prism_Estep_eval.py new file mode 100755 index 00000000..76f70627 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/prism_Estep_eval.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +""" +Evaluation and plotting for prism E-step results. +""" + +import os +import argparse +import numpy as np +from pprint import pprint + +from toolbox.Util_NumpyIO import read_data_npz +from PlotterPrismEstep import Plotter + + +def compute_ll_gap(spikes, A_true, B_true, t0_bin, t1_bin, dt, eta_clip): + spikes_sub = spikes[t0_bin : t1_bin + 1].astype(np.float64) + y_prev = spikes_sub[:-1] + y_curr = spikes_sub[1:] + T_pairs = y_prev.shape[0] + M = B_true.shape[0] + ll_gap = np.zeros((T_pairs,), dtype=np.float64) + if M == 1: + return ll_gap + for t in range(T_pairs): + yp = y_prev[t] + yc = y_curr[t] + base = A_true @ yp + eta = base[None, :] + B_true + eta_c = np.minimum(eta, eta_clip) + lam = np.exp(eta_c) * dt + scores = np.sum(yc * eta_c - lam, axis=1) + top2 = np.partition(scores, -2)[-2:] + ll_gap[t] = top2[-1] - top2[-2] + return ll_gap + + +def eval_estep_metrics(fitD, md, spikes): + trainMD = md["train"] + t0_bin, t1_bin = trainMD["time_range_bins"] + s_true = md["S_true"][t0_bin : t1_bin + 1] + s_hat = fitD["S_hat"] + s_hat_cl = fitD["S_hat_CL"] + n_cmp = min(s_true.shape[0], s_hat.shape[0]) + if n_cmp > 0: + acc = float((s_true[:n_cmp] == s_hat[:n_cmp]).mean()) + else: + acc = float("nan") + + # Per-state average S_hat_CL and accuracy + M = int(trainMD["num_states"]) + avg_cl_per_state = [] + acc_per_state = [] + for m in range(M): + mask = s_hat[:n_cmp] == m + cnt = int(mask.sum()) + if cnt > 0: + avg_cl_per_state.append(float(s_hat_cl[:n_cmp][mask].mean())) + acc_per_state.append(float((s_true[:n_cmp][mask] == m).mean())) + else: + avg_cl_per_state.append(float("nan")) + acc_per_state.append(float("nan")) + + # Store per-state metrics in decode_eval metadata + md["decode_eval"]["avg_cl_per_state"] = avg_cl_per_state + md["decode_eval"]["acc_per_state"] = acc_per_state + + A_true = md["A_true"] + B_true = md["B_true"] + dt = float(trainMD["time_step_sec"]) + eta_clip = float(trainMD["eta_clip"]) + ll_gap = compute_ll_gap(spikes, A_true, B_true, t0_bin, t1_bin, dt, eta_clip) + return { + "acc": acc, + "ll_gap": ll_gap, + "ll_gap_mean": float(np.mean(ll_gap)) if ll_gap.size else float("nan"), + } + + +def main(): + parser = argparse.ArgumentParser(description="Evaluate and plot prism E-step results") + parser.add_argument("--dataName", type=str, required=True, help="Base name for prismEstep file") + parser.add_argument("--basePath", type=str, default="/pscratch/sd/b/balewski/2026_causalNet_tmp2/", help="head dir for input/output data") + parser.add_argument("-p", "--showPlots", type=str, nargs='+', default="a", help="Plot types: a=training summary, b=state sequence (fit vs truth)") + parser.add_argument("--time_bin_merge", type=int, default=2, help="Merge factor for loss(time) x-axis") + parser.add_argument("-X", "--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument("-v", "--verb", type=int, default=1, help="Verbosity level") + args = parser.parse_args() + + + truthPath = os.path.join(args.basePath,"spikesData") + args.inpPath = os.path.join(args.basePath, "prismFit") + args.outPath = os.path.join(args.basePath, "plots") + os.makedirs(args.outPath, exist_ok=True) + args.showPlots = ''.join(args.showPlots) + + print(vars(args)) + + fitFF = os.path.join(args.inpPath, f"{args.dataName}.prismEstep.npz") + fitD, fitMD = read_data_npz(fitFF) + if args.verb > 1: pprint(fitMD) + + truthF = fitMD["provenance"]['state_transition_file'] + truthFF = os.path.join(truthPath, f"{truthF}.prismTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 0) + A_true = trueD["A_true"] + B_true = trueD["B_true"] + C_true = trueD["C_true"] + S_true = trueD["S_true"] + if args.verb > 1: pprint(trueMD) + + spikesFF = os.path.join(truthPath, f"{truthF}.spikes.npz") + spikesD, spikesMD = read_data_npz(spikesFF, verb=args.verb > 0) + spikes = spikesD["spikes"] + + MD = {**fitMD, "short_name": args.dataName} + MD["A_true"] = A_true + MD["B_true"] = B_true + MD["C_true"] = C_true + MD["S_true"] = S_true + MD["eval_Estep"] = eval_estep_metrics(fitD, MD, spikes) + decE = MD["decode_eval"] + print( + f"E-step eval avr acc {MD['eval_Estep']['acc']:.3f}, {args.dataName}" + ) + print(f" {'state':>5s} {'CL':>6s} {'acc':>5s}") + print(f" {'-----':>5s} {'------':>6s} {'-----':>5s}") + for m, (cl, ac) in enumerate(zip(decE["avg_cl_per_state"], decE["acc_per_state"])): + print(f" {m:5d} {cl:6.3f} {ac:5.3f}") + + #pprint(decE) + args.prjName = args.dataName + plot = Plotter(args) + + if 'a' in args.showPlots: + plot.summary_prismEstep(fitD, MD, figId=1, time_bin_merge=args.time_bin_merge) + + if 'b' in args.showPlots: + plot.state_seq_prismEstep(fitD, MD, figId=2, time_reb=args.time_bin_merge) + + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3a_states_simu/prism_Estep_train.py b/causal_net/nonStation_ver3a_states_simu/prism_Estep_train.py new file mode 100755 index 00000000..6746a358 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/prism_Estep_train.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python3 +""" + salloc -q shared_interactive -C gpu -t 4:00:00 -N 1 -A m2043 + +E-step feasibility test for non-stationary Poisson dLDS. + +Given the ground-truth dictionary (A_true, B_true) from gen_daleMatrices3.py, +infer the simplex-valued state coefficients c_t in [Delta^{M-1}] for every +time bin using projected gradient descent (PGD) on the probability simplex. + +This isolates the E-step of the EM framework to assess how well the latent +state sequence can be recovered when the dictionary is perfectly known. + +Model (shared A, per-state B): + Z_t[m, :] = A @ Y_{t-1} + B_m predictor matrix (M x N) + eta_t = c_t @ Z_t linear predictor (N,) + lambda_t = exp( clip(eta_t, max=eta_clip) ) * dt + Y_t ~ Poisson(lambda_t) + +Per-step objective (minimised over c_t in Delta^{M-1}): + L_t(c_t) = sum_i [ exp(clip(eta_i)) * dt - Y_i * (eta_i + log dt) ] + + lambda2 * || c_t - c_{t-1} ||^2 + +Optimisation: + - Single forward sweep PGD: process bins t = 1 ... T in order. + - At each bin: pgd_iter gradient steps followed by simplex projection. + - The smoothness coupling is forward-only (c_t depends on c_{t-1}), + so a single pass suffices. + +Initialisation: c_t = 1/M (uniform) for all t. + +State decoding (post-hoc, no gradient): + S_hat[t] = argmax_m c_hat[t, m] + Optionally apply a minimum-dwell filter (--decode_inertia_sec) to + suppress short state flips. + +Confidence: S_hat_CL[t] = 1 - max_m c_hat[t, m] (0 = fully confident). + +Hyperparameter defaults: + lambda2 = 2.0 temporal smoothness weight + lr = 0.03 PGD step size + pgd_iter = 40 PGD iterations per time bin + chunk_size = 2048 bins per progress-print block + decode_inertia_sec = 0.0 minimum dwell filter (disabled) +""" + +import os +import time +import math +import argparse +import secrets +from pprint import pprint + +import numpy as np +import torch + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from UtilTorch import check_gpu_availability + + +def parse_args(): + parser = argparse.ArgumentParser(description="E-step training for prism model (simplex PGD)") + parser.add_argument("--dataName", type=str, required=True, help="Base name of spikes file in spikesData/") + parser.add_argument("--basePath", type=str, default="/pscratch/sd/b/balewski/2026_causalNet_tmp2/", help="Head dir for input/output data") + parser.add_argument("--lambda2", type=float, default=2.0, help="Temporal smoothness strength") + parser.add_argument("--pgd_iter", type=int, default=40, help="PGD iterations per time step") + parser.add_argument("--lr", type=float, default=0.03, help="PGD step size") + parser.add_argument("--decode_dwell_sec", type=float, default=0.1, help="Expected state dwell time in seconds for Viterbi decoding") + parser.add_argument("--chunk_size", type=int, default=2 * 1024, help="Time chunk size for progress") + parser.add_argument("-T", "--time_range_sec", default=[0.0, 10.0], nargs=2, type=float, help="display data time range in seconds") + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("-v", "--verb", type=int, default=1) + return parser.parse_args() + + +def make_prism_out_base(data_name: str) -> str: + hash6 = secrets.token_hex(3) # 6 hex digits + return f"{data_name}-Estep-{hash6}" + + +def project_to_simplex(v: torch.Tensor) -> torch.Tensor: + """Project 1D tensor v onto the probability simplex.""" + if v.numel() == 1: + return torch.ones_like(v) + u, _ = torch.sort(v, descending=True) + cssv = torch.cumsum(u, dim=0) - 1.0 + ind = torch.arange(1, v.numel() + 1, device=v.device, dtype=v.dtype) + cond = u - cssv / ind > 0 + if torch.any(cond): + rho = torch.nonzero(cond, as_tuple=False)[-1, 0] + theta = cssv[rho] / (rho + 1.0) + else: + theta = cssv[-1] / v.numel() + w = torch.clamp(v - theta, min=0.0) + return w + + +def viterbi_decode(c_hat, p_stay): + """Viterbi decode most likely state sequence from c_hat with stay/switch prior.""" + """ Viterbi is a globally optimal probabilistic decoder. It operates on the soft coefficients c_hat (the simplex weights) and finds the single state sequence hat{S}_{1:T} that maximizes the joint probability of emissions and transitions simultaneously across all time bins. + """ + eps = 1e-12 + T, M = c_hat.shape + if M == 1: + return np.zeros((T,), dtype=np.int64), np.ones((T,), dtype=np.float32) + + p_switch = (1.0 - p_stay) / float(M - 1) + trans = np.full((M, M), p_switch, dtype=np.float64) + np.fill_diagonal(trans, p_stay) + log_trans = np.log(np.clip(trans, eps, 1.0)) + + log_emit = np.log(np.clip(c_hat, eps, 1.0)) + delta = np.zeros((T, M), dtype=np.float64) + psi = np.zeros((T, M), dtype=np.int64) + + delta[0] = log_emit[0] + for t in range(1, T): + scores = delta[t - 1][:, None] + log_trans + psi[t] = np.argmax(scores, axis=0) + delta[t] = scores[psi[t], np.arange(M)] + log_emit[t] + + path = np.zeros((T,), dtype=np.int64) + path[T - 1] = int(np.argmax(delta[T - 1])) + for t in range(T - 2, -1, -1): + path[t] = psi[t + 1, path[t + 1]] + + return path + + +def main(): + args = parse_args() + inpPath = os.path.join(args.basePath, "spikesData") + truthPath = os.path.join(args.basePath, "truthDale") + outPath = os.path.join(args.basePath, "prismFit") + assert os.path.exists(outPath) + + print('\nE-train args:', vars(args), '\n') + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + + device = check_gpu_availability() + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision('high') + + spikesFF = os.path.join(inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 0) + spikes = spikeD["spikes"] + single_rates = spikeD.get("single_rates") + + args.time_step_sec = float(spikeMD["time_step_sec"]) + args.eta_clip = float(spikeMD["poisson_eta_clip"]) + + prov = spikeMD["provenance"] + truthF = prov["state_model_file"] + + truthFF = os.path.join(truthPath, f"{truthF}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 1) + A_true = trueD["A_true"] + B_true = trueD["B_true"] + + assert A_true.ndim == 2, "A_true must have shape (N,N)" + assert B_true.ndim == 2, "B_true must have shape (M,N)" + N, N2 = A_true.shape + assert N == N2, "A_true must be square" + M = B_true.shape[0] + assert B_true.shape[1] == N, "B_true shape mismatch" + + T_full, N_spk = spikes.shape + assert N_spk == N, "Spike data N does not match A_true/B_true" + + # Time range in bins + t0_sec, t1_sec = float(args.time_range_sec[0]), float(args.time_range_sec[1]) + if t1_sec < t0_sec: + t0_sec, t1_sec = t1_sec, t0_sec + start_bin = max(0, int(math.floor(t0_sec / args.time_step_sec))) + end_bin = min(T_full - 1, int(math.floor(t1_sec / args.time_step_sec))) + if end_bin <= start_bin: + raise ValueError("time_range_sec too small; need at least two bins") + + # Slice spikes to selected range (inclusive end_bin) + spikes_sub = spikes[start_bin:end_bin + 1] + T_eff = spikes_sub.shape[0] + T_pairs = T_eff - 1 + if args.verb > 0: + print(f"time bins: start_bin={start_bin} num_bins={T_eff} (pairs={T_pairs})") + + # Build per-time data (pairs) + Y_prev = torch.tensor(spikes_sub[:-1], dtype=torch.float32, device=device) + Y_curr = torch.tensor(spikes_sub[1:], dtype=torch.float32, device=device) + + # Move dictionaries to GPU + A_t = torch.tensor(A_true, dtype=torch.float32, device=device) + B_t = torch.tensor(B_true, dtype=torch.float32, device=device) + + # Initialize c_hat on GPU + c_hat = torch.full((T_eff, M), 1.0 / float(M), dtype=torch.float32, device=device) + + loss_time = np.zeros((T_pairs,), dtype=np.float64) + loss_nll_time = np.zeros((T_pairs,), dtype=np.float64) + loss_l2_time = np.zeros((T_pairs,), dtype=np.float64) + log_dt = math.log(args.time_step_sec) + + t_start = time.time() + torch.set_grad_enabled(False) + + c_prev = c_hat[0].clone() + nll_sum = 0.0 + l2_sum = 0.0 + g_data_sum = 0.0 + g_smooth_sum = 0.0 + g_count = 0 + + for t0_idx in range(1, T_eff, args.chunk_size): + t1_idx = min(T_eff, t0_idx + args.chunk_size) + for t in range(t0_idx, t1_idx): + y_prev = Y_prev[t - 1] + y_curr = Y_curr[t - 1] + + base = torch.matmul(A_t, y_prev) + z = base[None, :] + B_t + + c_t = c_hat[t].clone() + for _ in range(args.pgd_iter): + eta = torch.matmul(c_t, z) + eta_x = torch.clamp(eta, max=args.eta_clip) + lam = torch.exp(eta_x) * args.time_step_sec + resid = lam - y_curr + g_data = torch.matmul(z, resid) + g_smooth = 2.0 * args.lambda2 * (c_t - c_prev) + g = g_data + g_smooth + g_data_sum += float(torch.linalg.norm(g_data).item()) + g_smooth_sum += float(torch.linalg.norm(g_smooth).item()) + g_count += 1 + c_t = c_t - args.lr * g + c_t = project_to_simplex(c_t) + + eta = torch.matmul(c_t, z) + eta_y = torch.clamp(eta, max=args.eta_clip) + lam = torch.exp(eta_y) * args.time_step_sec + dev = lam - y_curr * (eta + log_dt) + loss_nll = dev.sum() + loss_l2 = args.lambda2 * torch.sum((c_t - c_prev) ** 2) + loss_t = loss_nll + loss_l2 + + nll_sum += float(loss_nll.item()) + l2_sum += float(loss_l2.item()) + loss_time[t - 1] = float(loss_t.item()) + loss_nll_time[t - 1] = float(loss_nll.item()) + loss_l2_time[t - 1] = float(loss_l2.item()) + + c_hat[t] = c_t + c_prev = c_t + + nll_mean = nll_sum / float(T_pairs) + l2_mean = l2_sum / float(T_pairs) + loss_mean = nll_mean + l2_mean + + elaT = time.time() - t_start + g_data_mean = g_data_sum / float(max(1, g_count)) + g_smooth_mean = g_smooth_sum / float(max(1, g_count)) + if args.verb > 0: + print( + f"PGD sweep nll={nll_mean:.4e} l2={l2_mean:.4e} " + f"loss={loss_mean:.4e} |g_data|={g_data_mean:.3e} " + f"|g_smooth|={g_smooth_mean:.3e} elaT={elaT:.1f}s" + ) + print(f"Total training time: {elaT:.1f}s") + + # Decode most probable state (Viterbi) and confidence level + p_stay = float(math.exp(-args.time_step_sec / float(args.decode_dwell_sec))) + c_hat_np = c_hat.detach().cpu().numpy() + S_hat_np = viterbi_decode(c_hat_np, p_stay) + if M == 1: + S_hat_CL = torch.zeros((T_eff,), dtype=torch.float32, device=c_hat.device) + else: + S_hat_CL = (1.0 - c_hat.max(dim=1).values).to(dtype=torch.float32) + S_hat = torch.from_numpy(S_hat_np).to(dtype=torch.int64, device=c_hat.device) + + # Save results + out_base = make_prism_out_base(args.dataName) + out_name = f"{out_base}.prismEstep.npz" + outFF = os.path.join(outPath, out_name) + + outD = { + "c_hat": c_hat.detach().cpu().numpy(), + "S_hat": S_hat.detach().cpu().numpy().astype(np.int64), + "S_hat_CL": S_hat_CL.detach().cpu().numpy().astype(np.float32), + "loss_time": np.asarray(loss_time, dtype=np.float64), + "loss_nll_time": np.asarray(loss_nll_time, dtype=np.float64), + "loss_l2_time": np.asarray(loss_l2_time, dtype=np.float64), + } + if single_rates is not None: + outD["single_rates"] = single_rates + + outMD = dict(spikeMD) + outMD["fit_type"] = "prismEstep" + outMD["train"] = { + "lambda2": float(args.lambda2), + "pgd_iter": int(args.pgd_iter), + "lr": float(args.lr), + "chunk_size": int(args.chunk_size), + "seed": int(args.seed), + "time_step_sec": float(args.time_step_sec), + "eta_clip": float(args.eta_clip), + "num_states": int(M), + "num_neurons": int(N), + "num_steps": int(T_eff), + "time_range_sec": [float(t0_sec), float(t1_sec)], + "time_range_bins": [int(start_bin), int(end_bin)], + "loss_mean": float(loss_mean), + "loss_nll_mean": float(nll_mean), + "loss_l2_mean": float(l2_mean), + } + outMD["decode_eval"] = { + "decode": "viterbi", + "decode_dwell_sec": float(args.decode_dwell_sec), + } + outMD["provenance"] = prov + + write_data_npz(outD, outFF, metaD=outMD) + print(f"Saved prism E-step fit to: {outFF}") + print(' basePath=' + args.basePath) + print(' ./prism_Estep_eval.py --basePath $basePath --dataName %s -p b a -X \n ' % (out_base)) + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3a_states_simu/prism_Mstep_eval.py b/causal_net/nonStation_ver3a_states_simu/prism_Mstep_eval.py new file mode 100755 index 00000000..2019e2cd --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/prism_Mstep_eval.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +""" +Evaluation and plotting for prism M-step results. +""" + +import os +import argparse +from pprint import pprint +import numpy as np +from toolbox.Util_NumpyIO import read_data_npz +from PlotterPrismMstep import Plotter + + +def compute_edge_eval(A_hat, E_true, minW): + """Compute edge-detection metrics/masks once; plotter only displays them.""" + A_hat = np.asarray(A_hat) + E_true = np.asarray(E_true).astype(bool) + + if A_hat.ndim == 2: + A_hat = A_hat[None, :, :] + if E_true.ndim == 2: + E_true = np.repeat(E_true[None, :, :], A_hat.shape[0], axis=0) + + n_states, N, N2 = A_hat.shape + assert N == N2, "A_hat must be square per state" + assert E_true.shape == (n_states, N, N), "E_true shape must match A_hat" + + off_diag = ~np.eye(N, dtype=bool) + out = [] + for m in range(n_states): + E_hat = (np.abs(A_hat[m]) > minW) & off_diag + E_t = E_true[m] & off_diag + + TP = E_t & E_hat + FP = (~E_t) & E_hat + FN = E_t & (~E_hat) + TN = (~E_t) & (~E_hat) + + tp = int(TP.sum()) + fp = int(FP.sum()) + fn = int(FN.sum()) + tn = int(TN.sum()) + precision = tp / max(1, tp + fp) + recall = tp / max(1, tp + fn) + f1 = 2 * precision * recall / max(1e-12, precision + recall) + acc = (tp + tn) / max(1, tp + tn + fp + fn) + + conf_map = np.zeros((N, N), dtype=np.int8) + conf_map[FN] = 1 + conf_map[FP] = 2 + conf_map[TP] = 3 + + out.append({ + "state": int(m), + "minW": float(minW), + "n_hat": int(E_hat.sum()), + "n_true": int(E_t.sum()), + "tp": tp, + "fp": fp, + "fn": fn, + "tn": tn, + "precision": float(precision), + "recall": float(recall), + "f1": float(f1), + "acc": float(acc), + "E_hat": E_hat.astype(np.uint8), + "E_true_offdiag": E_t.astype(np.uint8), + "conf_map": conf_map, + }) + return out + + +def main(): + parser = argparse.ArgumentParser(description="Evaluate and plot prism M-step results") + parser.add_argument("--dataName", type=str, required=True, help="Base name for prismMstep file") + parser.add_argument("--basePath", type=str, default="/pscratch/sd/b/balewski/2026_causalNet_tmp2/", help="head dir for input/output data") + parser.add_argument("-p", "--showPlots", type=str, nargs='+', default="a", help=("Plot types: a=training summary, b=A-matrix edge evaluation (1-row: E_true|A_hat|confusion|stats), c=2D correlations (A_true vs A_hat, B_true vs B_hat), d=state correlations (fit), e=state correlations (truth), f=edge detection vs truth")) + parser.add_argument("--minW", type=float, default=0.02, help="Threshold for A-matrix eval, not for fitting") + parser.add_argument("--timeReb", type=int, default=20, help="Time rebin factor for plots with time axis") + parser.add_argument("-X", "--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument("-v", "--verb", type=int, default=1, help="Verbosity level") + args = parser.parse_args() + + args.inpPath = os.path.join(args.basePath, "prismFit") + args.outPath = os.path.join(args.basePath, "plots") + os.makedirs(args.outPath, exist_ok=True) + args.showPlots = ''.join(args.showPlots) + + print(vars(args)) + + fitFF = os.path.join(args.inpPath, f"{args.dataName}.prismMstep.npz") + fitD, fitMD = read_data_npz(fitFF) + assert isinstance(fitMD, dict), "Expected metadata dict in prismMstep file" + + if args.verb > 1: + pprint(fitMD) + + MD = {**fitMD, "short_name": args.dataName} + + # Load ground truth (A_true/B_true) if available + prov = fitMD["provenance"] + truth_name = prov["state_model_file"] + truthPath = os.path.join(args.basePath, "truthDale") + truthFF = os.path.join(truthPath, f"{truth_name}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF) + if args.verb > 1: + print("\nsimTruth metadata:"); pprint(trueMD) + + MD.update(trueMD) + MD["A_true"] = trueD["A_true"] + MD["B_true"] = trueD["B_true"] + MD["E_true"] = trueD["E_true"] + MD["short_name"] = args.dataName + + st_name = prov["state_transition_file"] + prismTruthFF = os.path.join(args.basePath, "spikesData", f"{st_name}.prismTruth.npz") + trD, trMD = read_data_npz(prismTruthFF, verb=args.verb > 1) + MD["S_true"] = trD["S_true"] + MD["C_true"] = trD["C_true"] + edge_eval_states = compute_edge_eval(fitD["A_hat"], MD["E_true"], args.minW) + MD["edge_eval_states"] = edge_eval_states + MD["edge_eval_main"] = edge_eval_states[0] + + + args.prjName = args.dataName + plot = Plotter(args) + + if 'a' in args.showPlots: + plot.summary_prismMstep(fitD, MD, figId=1) + + if 'b' in args.showPlots: + plot.eval_Amatrix_prismMstep(fitD, MD, figId=2) + + if 'c' in args.showPlots: + plot.eval_ABcorr_prismMstep(fitD, MD, figId=3) + + if 'd' in args.showPlots: + plot.state_Bcorr_prismMstep(fitD, MD, figId=4) + + + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3a_states_simu/prism_Mstep_train3.py b/causal_net/nonStation_ver3a_states_simu/prism_Mstep_train3.py new file mode 100755 index 00000000..19c1e5eb --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/prism_Mstep_train3.py @@ -0,0 +1,473 @@ +#!/usr/bin/env python3 +""" +Single-GPU prism M-step training for non-stationary Poisson GLM. + +Fits a shared connectivity matrix A and per-state bias vectors B_hat[m] +using ground-truth state targets from prismTruth: + - C_true(t) simplex coefficients (soft labels only) +""" + +import os +import time +import secrets +import argparse +import random + +import numpy as np +from pprint import pprint +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from UtilTorch import check_gpu_availability + + +class SwitchingPoissonGLModel(nn.Module): + """Shared A, per-state B[m], mixed by per-sample coefficients.""" + def __init__(self, n_neurons, n_states, eta_clip): + super().__init__() + self.n_neurons = int(n_neurons) + self.n_states = int(n_states) + self.eta_clip = float(eta_clip) + self.A = nn.Parameter(torch.randn(self.n_neurons, self.n_neurons) * 0.1) + self.B = nn.Parameter(torch.randn(self.n_states, self.n_neurons) * 0.1) + + def forward(self, y_prev, c_coeff, dt): + base = y_prev @ self.A.t() + b_eff = c_coeff @ self.B + linear = base + b_eff + linear = torch.clamp(linear, max=self.eta_clip) + return torch.exp(linear) * dt + + +class NumpyQuartetDataset(Dataset): + def __init__(self, x_np, y_np, s_np, c_np): + assert x_np.shape[0] == y_np.shape[0] == s_np.shape[0] == c_np.shape[0] + self.x = x_np + self.y = y_np + self.s = s_np + self.c = c_np + + def __len__(self): + return self.x.shape[0] + + def __getitem__(self, idx): + return ( + torch.from_numpy(self.x[idx]).to(dtype=torch.float32), + torch.from_numpy(self.y[idx]).to(dtype=torch.float32), + torch.tensor(self.s[idx], dtype=torch.long), + torch.from_numpy(self.c[idx]).to(dtype=torch.float32), + ) + + +def make_loader_xysc(x_np, y_np, s_np, c_np, batch_size, shuffle=True): + ds = NumpyQuartetDataset(x_np, y_np, s_np, c_np) + return DataLoader( + ds, + batch_size=batch_size, + shuffle=shuffle, + drop_last=shuffle, + pin_memory=True, + pin_memory_device="cuda", + num_workers=8, + persistent_workers=True, + prefetch_factor=8, + ) + + +def poisson_nll_weighted(spikes, targets, firing_rates): + eps = 1e-8 + firing_rates_safe = torch.maximum(firing_rates, torch.tensor(0.1, device=firing_rates.device)) + weights = 1.0 / firing_rates_safe + weights = weights / torch.mean(weights) + weights = weights.unsqueeze(0) + loss = -weights * targets * torch.log(spikes + eps) + weights * spikes + return loss.mean() + + +def poisson_nll_weighted_per_sample(spikes, targets, firing_rates): + eps = 1e-8 + firing_rates_safe = torch.maximum(firing_rates, torch.tensor(0.1, device=firing_rates.device)) + weights = 1.0 / firing_rates_safe + weights = weights / torch.mean(weights) + weights = weights.unsqueeze(0) + loss = -weights * targets * torch.log(spikes + eps) + weights * spikes + return loss.mean(dim=1) + + +def offdiag_soft_threshold_(A, lr, lam): + if lam <= 0.0: + return + with torch.no_grad(): + n = A.shape[0] + off_diag = ~torch.eye(n, dtype=torch.bool, device=A.device) + t = lr * lam + A_off = A[off_diag] + A[off_diag] = A_off.sign() * (A_off.abs() - t).clamp(min=0.0) + + +def enforce_spectral_radius_(A, rho_max, eps=1e-12): + if rho_max is None: + return + if rho_max <= 0.0: + raise ValueError(f"rho_max must be > 0, got {rho_max}") + with torch.no_grad(): + rho = torch.linalg.eigvals(A).abs().max() + if torch.isfinite(rho) and (rho > rho_max): + A.mul_(float(rho_max) / float(rho + eps)) + + +def preprocess_data_with_truth(spikes, s_true, c_true, n_states, args): + y = np.asarray(spikes) + s_true = np.asarray(s_true).astype(np.int64) + c_true = np.asarray(c_true).astype(np.float32) + nt, nn = y.shape + assert s_true.shape[0] == nt, "S_true length mismatch with spikes length" + assert c_true.shape[0] == nt, "C_true length mismatch with spikes length" + assert c_true.shape[1] == n_states, f"C_true second dim mismatch: got {c_true.shape[1]} expected {n_states}" + + if args.desyncTime: + seed = int(time.time() * 1000) % 1000000 + np.random.seed(seed) + shift_amounts = np.random.randint(1, nt // 4, size=nn, dtype=np.int32) + y_shifted = np.zeros_like(y) + for neuron_idx in range(nn): + shift_amount = int(shift_amounts[neuron_idx]) + y_shifted[:, neuron_idx] = np.roll(y[:, neuron_idx], shift_amount) + y = y_shifted + print(f"Applied time decorrelation shifts (seed={seed})") + + max_pairs = nt - 1 + num_samples = args.num_samples + if num_samples is None or num_samples > max_pairs: + num_samples = max_pairs + x_np = y[:num_samples] + yt_np = y[1:num_samples + 1] + s_np = s_true[1:num_samples + 1] + c_np = c_true[1:num_samples + 1] + + if args.dropDataFrac > 0: + n_pairs = x_np.shape[0] + drop_seed = (int(time.time() * 1000) + np.random.randint(0, 1000)) % (2**32) + np.random.seed(drop_seed) + random.seed(drop_seed) + n_keep = int(n_pairs * (1.0 - args.dropDataFrac)) + keep_indices = np.random.choice(n_pairs, size=n_keep, replace=False) + keep_indices = np.sort(keep_indices) + x_np = x_np[keep_indices] + yt_np = yt_np[keep_indices] + s_np = s_np[keep_indices] + c_np = c_np[keep_indices] + print(f"Dropped {args.dropDataFrac:.1%} of data, keeping {x_np.shape[0]} samples (seed={drop_seed})") + + return ( + x_np.astype(np.float32), + yt_np.astype(np.float32), + s_np.astype(np.int64), + c_np.astype(np.float32), + ) + + +def train_switching_mstep_model( + model, + device, + train_loader, + n_epochs, + lr, + dt, + firing_rates, + L1_alpha=0.0, + use_scheduler=True, + lr_end_factor=0.03, + rho_max=0.97, + prescale_m_step_4_ArhoMax=20, + L1_prune_epoch=0, + minW=1e-6, +): + if L1_prune_epoch < 0: + raise ValueError(f"L1_prune_epoch must be >= 0, got {L1_prune_epoch}") + if prescale_m_step_4_ArhoMax < 1: + raise ValueError(f"prescale_m_step_4_ArhoMax must be >= 1, got {prescale_m_step_4_ArhoMax}") + use_fused = isinstance(device, torch.device) and device.type == "cuda" and torch.cuda.is_available() + assert use_fused, "This trainer expects CUDA fused Adam" + + optimizer = optim.Adam(model.parameters(), lr=lr, fused=True) + scheduler = ( + optim.lr_scheduler.LinearLR( + optimizer, start_factor=1.0, end_factor=lr_end_factor, total_iters=n_epochs + ) + if use_scheduler + else None + ) + + n = model.n_neurons + diag_mask = torch.eye(n, device=device, dtype=torch.bool) + off_diag_mask = ~diag_mask + l1_weight_matrix = torch.ones(n, n, device=device) + l1_weight_matrix[diag_mask] = 0.0 + firing_rates_t = torch.tensor(firing_rates, dtype=torch.float32, device=device) + + loss_epoch, loss_nll_epoch, loss_l1_epoch = [], [], [] + learning_rates, rho_epoch, nz_edges_epoch, sparsity_epoch = [], [], [], [] + loss_state_total_epoch = [] + + t_start = time.time() + for epoch in range(n_epochs): + model.train() + apply_prune = epoch >= L1_prune_epoch + sum_tot = 0.0 + sum_nll = 0.0 + sum_l1 = 0.0 + state_nll_sum = np.zeros((model.n_states,), dtype=np.float64) + state_n_count = np.zeros((model.n_states,), dtype=np.float64) + + for batch_idx, (y_prev, y_curr, s_idx, c_coeff) in enumerate(train_loader): + y_prev = y_prev.float().to(device, non_blocking=True) + y_curr = y_curr.float().to(device, non_blocking=True) + s_idx = s_idx.long().to(device, non_blocking=True) + c_coeff = c_coeff.float().to(device, non_blocking=True) + + optimizer.zero_grad(set_to_none=True) + spikes = model(y_prev, c_coeff, dt=dt) + nll_per_sample = poisson_nll_weighted_per_sample(spikes, y_curr, firing_rates_t) + base_loss = nll_per_sample.mean() + if L1_alpha > 0: + l1_term = L1_alpha * torch.mean(torch.abs(model.A) * l1_weight_matrix) + else: + l1_term = torch.tensor(0.0, device=device) + total_loss = base_loss + l1_term + total_loss.backward() + optimizer.step() + + if apply_prune and (L1_alpha > 0): + offdiag_soft_threshold_(model.A, optimizer.param_groups[0]["lr"], L1_alpha) + if batch_idx % prescale_m_step_4_ArhoMax == 0: + enforce_spectral_radius_(model.A, rho_max) + + sum_tot += float(total_loss.item()) + sum_nll += float(base_loss.item()) + sum_l1 += float(l1_term.item()) + for m in range(model.n_states): + w = c_coeff[:, m] + w_sum = float(w.sum().item()) + if w_sum <= 0.0: + continue + state_nll_sum[m] += float((nll_per_sample * w).sum().item()) + state_n_count[m] += w_sum + + n_batches = float(len(train_loader)) + loss_epoch.append(sum_tot / n_batches) + loss_nll_epoch.append(sum_nll / n_batches) + loss_l1_epoch.append(sum_l1 / n_batches) + state_nll_epoch = state_nll_sum / np.maximum(1, state_n_count) + state_total_epoch = state_nll_epoch + loss_l1_epoch[-1] + loss_state_total_epoch.append(state_total_epoch) + learning_rates.append(float(optimizer.param_groups[0]["lr"])) + + with torch.no_grad(): + a_off = model.A[off_diag_mask] + nz_off = int((a_off.abs() > minW).sum().item()) + n_off = int(off_diag_mask.sum().item()) + sparsity = 1.0 - nz_off / max(1, n_off) + rho = float(torch.linalg.eigvals(model.A).abs().max().item()) + nz_edges_epoch.append(nz_off) + sparsity_epoch.append(float(sparsity)) + rho_epoch.append(rho) + + if scheduler is not None: + scheduler.step() + + if (epoch + 1) % 5 == 0 or (epoch + 1) == n_epochs: + print( + f"Epoch {epoch+1}/{n_epochs}: NLL={loss_nll_epoch[-1]:.5g}, " + f"L1={loss_l1_epoch[-1]:.5g}, A_sparsity={sparsity:.3f}, " + f"nz_offdiag={nz_off}, rho(A)={rho:.4f}, Elapsed={(time.time() - t_start):.1f}s" + ) + + return { + "loss_epoch": np.asarray(loss_epoch, dtype=np.float64), + "loss_nll_epoch": np.asarray(loss_nll_epoch, dtype=np.float64), + "loss_l1_epoch": np.asarray(loss_l1_epoch, dtype=np.float64), + "loss_state_total_epoch": np.asarray(loss_state_total_epoch, dtype=np.float64), + "learning_rates": np.asarray(learning_rates, dtype=np.float64), + "rho_epoch": np.asarray(rho_epoch, dtype=np.float64), + "nz_edges_epoch": np.asarray(nz_edges_epoch, dtype=np.int64), + "sparsity_epoch": np.asarray(sparsity_epoch, dtype=np.float64), + } + + +######################### +# MAIN +######################### + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dataName", type=str, default="dale_2aee70") + parser.add_argument("--basePath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input/output data") + parser.add_argument("--num_samples", type=int, default=None) + parser.add_argument("--num_epochs", type=int, default=100) + parser.add_argument("--batch_size", type=int, default=2048) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--L1_alpha", type=float, default=0.02, help="L1 regularization strength, higher=more sparse (0: disable soft-thresholding)") + parser.add_argument("--rho_max", type=float, default=0.92, help="Maximum allowed spectral radius of A; projection applied after each batch") + parser.add_argument("--prescale_m_step_4_ArhoMax", type=int, default=20, help="Apply spectral-radius projection every N batches") + parser.add_argument("--L1_prune_epoch", type=int, default=None, help="Delay L1 edge-pruning only; spectral-radius correction starts immediately") + parser.add_argument("--minW", type=float, default=0.01, help="Threshold for A-matrix eval, not for fitting") + parser.add_argument("--fitName", type=str, default=None) + parser.add_argument("--desyncTime", action='store_true', help="If true completely shuffle time axis for input data, independently for all channels") + parser.add_argument("--dropDataFrac", type=float, default=0.0, help="Fraction of training samples to randomly drop (0.0=use all data, 0.3=drop 30%%)") + parser.add_argument("--verb", "-v", type=int, default=1, help="Verbosity level") + + args = parser.parse_args() + + inpPath = os.path.join(args.basePath, 'spikesData') + outPath = os.path.join(args.basePath, 'prismFit') + if args.L1_prune_epoch is None: + args.L1_prune_epoch = args.num_epochs // 3 + + device = check_gpu_availability() + print("\nPrism M-step Config:", vars(args), "\n") + assert os.path.exists(outPath) + + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision('high') + gpu_name = torch.cuda.get_device_name(device) if isinstance(device, torch.device) and device.type == 'cuda' else str(device) + print("Using device %s : %s" % (str(device), gpu_name)) + + spikesFF = os.path.join(inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=True) + if args.verb > 1: + pprint(spikeMD) + dataYield = spikeD['spikes'] + dataRates = spikeD['single_rates'] + step_size = float(spikeMD['time_step_sec']) + eta_clip = float(spikeMD["poisson_eta_clip"]) + prov = spikeMD["provenance"] + + assert spikeMD['data_type'] == 'simPrism' + _, Nn = dataYield.shape + + truthF = prov["state_transition_file"] + truthFF = os.path.join(inpPath, f"{truthF}.prismTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 0) + if args.verb > 1: + print("\nprismTruth metadata:") + pprint(trueMD) + S_true = trueD["S_true"].astype(np.int64) + C_true = trueD["C_true"].astype(np.float32) + evol_conf = trueMD.get("evol_conf", {}) + if "num_states" in evol_conf: + n_states = int(evol_conf["num_states"]) + elif "B_true" in trueD: + n_states = int(np.asarray(trueD["B_true"]).shape[0]) + else: + n_states = int(np.max(S_true) + 1) + + x_np, yt_np, s_np, c_np = preprocess_data_with_truth( + dataYield, S_true, C_true, n_states=n_states, args=args + ) + n_pairs = x_np.shape[0] + print(f"Preprocessed data: X={x_np.shape}, Y={yt_np.shape}, S={s_np.shape}, C={c_np.shape}, n_pairs={n_pairs}") + + assert n_pairs >= args.batch_size, ( + f"ERROR: Not enough samples ({n_pairs}) for batch size ({args.batch_size}) after data dropping." + ) + + if np.any(s_np < 0) or np.any(s_np >= n_states): + raise ValueError(f"S_true contains state id outside [0,{n_states-1}]") + + state_counts = np.bincount(s_np, minlength=n_states) + state_mass = np.sum(c_np, axis=0) + if args.verb > 0: + frac = state_counts / max(1, state_counts.sum()) + msg = " ".join(f"s{m}:{state_counts[m]}({frac[m]:.1%})" for m in range(n_states)) + frac_mass = state_mass / max(1e-12, state_mass.sum()) + msg_mass = " ".join(f"s{m}:{state_mass[m]:.1f}({frac_mass[m]:.1%})" for m in range(n_states)) + print(f"Training hard coverage: {msg}") + print(f"Training coeff mass (soft/C_true): {msg_mass}") + + train_loader = make_loader_xysc(x_np, yt_np, s_np, c_np, batch_size=args.batch_size, shuffle=True) + print(f"Loaded pairs={n_pairs/1000:.3f}k, Nn={Nn}, M={n_states}, batch_size={args.batch_size}, labels=soft/C_true") + + model = SwitchingPoissonGLModel(Nn, n_states=n_states, eta_clip=eta_clip).to(device) + start_time = time.time() + train_hist = train_switching_mstep_model( + model=model, + device=device, + train_loader=train_loader, + n_epochs=args.num_epochs, + lr=args.lr, + dt=step_size, + firing_rates=dataRates, + L1_alpha=args.L1_alpha, + use_scheduler=True, + rho_max=args.rho_max, + prescale_m_step_4_ArhoMax=args.prescale_m_step_4_ArhoMax, + L1_prune_epoch=args.L1_prune_epoch, + minW=args.minW, + ) + + total_time = time.time() - start_time + print(f"Training completed in {total_time:.1f} seconds") + + if args.fitName is None: + hash6 = secrets.token_hex(3) + fit_core = f"{args.dataName}-Mstep-{hash6}" + else: + fit_core = args.fitName + + A_shared = model.A.detach().cpu().numpy() + B_hat = model.B.detach().cpu().numpy() + A_hat = A_shared + E_hat = (np.abs(A_hat) > 1e-5).astype(np.uint8) + + outD = { + "A_hat": A_hat, + "A_shared": A_shared, + "B_hat": B_hat, + "E_hat": E_hat, + "single_rates": dataRates, + **train_hist, + } + + outMD = dict(spikeMD) + outMD["fit_type"] = "prismMstep" + outMD["train"] = { + "num_epochs": int(args.num_epochs), + "batch_size": int(args.batch_size), + "num_samples_used": int(n_pairs), + "lr": float(args.lr), + "learning_rate": float(args.lr), + "L1_alpha": float(args.L1_alpha), + "lambda3": float(args.L1_alpha), + "rho_max": float(args.rho_max), + "prescale_m_step_4_ArhoMax": int(args.prescale_m_step_4_ArhoMax), + "L1_prune_epoch": int(args.L1_prune_epoch), + "minW": float(args.minW), + "dropDataFrac": float(args.dropDataFrac), + "state_label_mode": "soft", + "time_step_sec": float(step_size), + "eta_clip": float(eta_clip), + "num_neurons": int(Nn), + "num_states": int(n_states), + "training_time_sec": float(total_time), + } + outMD["evol_conf"] = trueMD.get("evol_conf", {}) + outMD["dale_conf"] = trueMD.get("dale_conf", {}) + outMD["provenance"] = dict(spikeMD.get("provenance", {})) + outMD["provenance"]["output_mstep_file"] = fit_core + + if args.verb > 1: + pprint(outMD) + fitFF = os.path.join(outPath, f"{fit_core}.prismMstep.npz") + write_data_npz(outD, fitFF, metaD=outMD) + + print(' basePath=' + args.basePath) + print(' ./prism_Mstep_eval.py --basePath $basePath --dataName %s -p a b c d \n ' % (fit_core,)) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3a_states_simu/toolbox b/causal_net/nonStation_ver3a_states_simu/toolbox new file mode 120000 index 00000000..c9203e9b --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/toolbox @@ -0,0 +1 @@ +../toolbox/ \ No newline at end of file diff --git a/causal_net/nonStation_ver3a_states_simu/view_daleMatrix3.py b/causal_net/nonStation_ver3a_states_simu/view_daleMatrix3.py new file mode 100755 index 00000000..f155847a --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/view_daleMatrix3.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +""" +Visualization tool for simulated Dale Poisson network data. + +Loads simulation output (simTruth.npz and spikes.npz) produced by +gen_daleMatrices.py and plots analysis figures for a selected state +index (--idxState). All arrays use natural neuron indexing +(first num_excite neurons are excitatory, remainder inhibitory). + +Available plots (-p flag): + a Dale connectivity matrix (color-coded) + eigenvalue scatter + with spectral-radius circle + b Weight histogram, firing-rate histogram, outgoing-edge count + per neuron, and per-neuron firing-rate bar chart + d B_idle vs firing rate / SNR scatter, plus excitatory and + inhibitory rate histograms + e Pseudospectral contour plot with eigenvalue overlay + +Usage: + ./view_daleMatrix.py --dataName daleN100_9fbe7f -i 0 -p a b d e +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os +from pprint import pprint +import numpy as np +from PlotterDaleMatrix import Plotter +from toolbox.Util_NumpyIO import read_data_npz +import argparse + +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize simulated Dale Poisson network data") + parser.add_argument("-v","--verbosity",type=int, help="increase output verbosity", default=1, dest='verb') + parser.add_argument("-p", "--showPlots", default='a b', nargs='+',help="abc-string listing shown plots: a=Dale_matrix_and_eigen, b=histo_weights_rates, d=rates_study, e=pseudospectra") + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument("--basePath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for input data") + parser.add_argument("--dataName", default='daleN150_448b86',help='simulated Dale network base name') + #parser.add_argument("--outPath", type=str, default=None, help="Output path for plots (defaults to basePath)") + parser.add_argument('-m', '--idxState', type=int, default=0, help="Index into state list, selects which state to plot") + + args = parser.parse_args() + + # make arguments more flexible + args.inpPath = os.path.join(args.basePath,'truthDale') + args.outPath = os.path.join(args.basePath,'plots') + args.showPlots=''.join(args.showPlots) + + print( 'myArg-program:',parser.prog) + for arg in vars(args): print( 'myArg:',arg, getattr(args, arg)) + + assert os.path.exists(args.basePath) + return args + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__=="__main__": + args=get_parser() + np.set_printoptions(precision=3) + + # Load simulation truth data (Dale matrices, biases, etc.) + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb>0) + if args.verb>1: + print("\nSimulation Truth Metadata:"); pprint(trueMD) + + # Load spike data (generated spike counts and rates) + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb>0) + if args.verb>1: + print("\nSpike Data Metadata:"); pprint(spikeMD) + + # Merge metadata for plotting + trueMD['short_name'] = args.dataName + + # Select B-offset slice + ib = args.idxState + boffsets = trueMD['dale_conf']['Boffsets'] + assert ib < len(boffsets), f"idxState={ib} out of range, only {len(boffsets)} states available" + B_offset = boffsets[ib] + R_sel = trueMD['dale_conf']['spectral_radius'] + print(f"\nSelected B offset [{ib}]: offset={B_offset:.3f} (out of {boffsets})") + print(f"Spectral radius: R={R_sel:.3f}") + + # Slice stacked arrays along axis 0 for the selected B offset + trueD_r = { 'A_true': trueD['A_true'], 'B_true': trueD['B_true'][ib], 'E_true': trueD['E_true'] } + spikeD_r = { k: spikeD[k][ib] for k in spikeD } + trueMD['sel_spect_radius'] = R_sel + trueMD['sel_Boffset'] = B_offset + trueMD['sel_state'] = ib + + #-------------------------------- + # .... plotting ........ + args.prjName=args.dataName+'_view%d'%ib + plot=Plotter(args) + + if 'a' in args.showPlots: + plot.Dale_matrix_and_eigen(trueD_r['A_true'],trueMD,trueD_r,figId=1) + + if 'b' in args.showPlots: + plot.histo_weights_rates(trueD_r,spikeD_r,trueMD,figId=2) + + if 'c' in args.showPlots: + plot.rates_study(trueD_r,spikeD_r,trueMD,figId=3) + + if 'd' in args.showPlots: + plot.Dale_matrix_pseudospectra(trueD_r['A_true'],trueMD,trueD_r,figId=4) + + plot.display_all() + print('M:done - view_dalePoisson completed successfully!') diff --git a/causal_net/nonStation_ver3a_states_simu/view_spikesTrain3.py b/causal_net/nonStation_ver3a_states_simu/view_spikesTrain3.py new file mode 100755 index 00000000..31d6f5c9 --- /dev/null +++ b/causal_net/nonStation_ver3a_states_simu/view_spikesTrain3.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +""" +Visualization tool for biological experiment input features and data quality. + +This script provides comprehensive visualization and analysis of experimental +neural data, focusing on data quality assessment and feature exploration. +Main functionality includes: +- Time series visualization of neural recordings with customizable time ranges +- Data quality metrics and statistical summaries +- Interactive plotting with configurable display options + +Used primarily for exploratory data analysis of biological neural recordings +before further processing and connectivity analysis. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os + +from pprint import pprint +import numpy as np +from PlotterSpikesTrain import Plotter +from toolbox.Util_NumpyIO import read_data_npz +import argparse +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize spike-train data from simulated Dale network outputs") + parser.add_argument("-v","--verbosity",type=int, help="increase output verbosity", default=1, dest='verb') + parser.add_argument("-p", "--showPlots", default='a', nargs='+',help="abcd-string listing shown plots") + + + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + + parser.add_argument("--basePath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for input data") + parser.add_argument("--dataName", default=None,help='simulated Dale network base name') + + parser.add_argument('-m', '--idxState', type=int, default=0, help="Index into state list; if idxState<0 read spikes from spikesData/") + + parser.add_argument('-T','--time_range_sec' , default=[0., 50], nargs=2, type=float, help='display data time range in seconds') + parser.add_argument('-r','--time_rebin2', default=10, type=int, help='rebin current time axis') + + args = parser.parse_args() + # make arguments more flexible + if args.idxState >= 0: + args.inpPath = os.path.join(args.basePath, 'truthDale') + else: + args.inpPath = os.path.join(args.basePath, 'spikesData') + args.outPath = os.path.join(args.basePath,'plots') + args.showPlots=''.join(args.showPlots) + + print( 'myArg-program:',parser.prog) + for arg in vars(args): print( 'myArg:',arg, getattr(args, arg)) + + if args.time_range_sec!=None: assert args.time_range_sec[0] < args.time_range_sec[1] + assert os.path.exists(args.basePath) + assert os.path.exists(args.inpPath), f"missing inpPath: {args.inpPath}" + return args + + +#...!...!.................... +def rebin_spike_rates(spikeYield, md, tReb2): + """Rebin 2D spike train (time x neurons) and return rate summaries over rebinned time.""" + assert spikeYield.ndim == 2, f"Expected 2D spike train, got shape={spikeYield.shape}" + assert tReb2<101 # this would exceed 1 seconds + + #.... rebin time axis + ntime,nchan=spikeYield.shape + if ntime % tReb2 != 0: + ntime_c = ntime - (ntime % tReb2) + spikeYield = spikeYield[:ntime_c] # Clip the data + #print('ccl',ntime_c , ntime ,ntime % tReb2, tReb2) + #... sum yileds over tReb2 time bins + spikeYieldR= np.sum( spikeYield.reshape(-1, tReb2, nchan),axis=1) + time_step=md['time_step_sec'] + time_step2=time_step*tReb2 + #print('rr1',time_step2,spikeYieldR.shape,spikeYield.shape) + + rate2D=spikeYieldR/time_step2 + pop_spike_count = np.sum(spikeYieldR, axis=1) + pop_rate_hz = pop_spike_count / time_step2 + + rebD={'time_step2':time_step2} + rebD['rate2D']=rate2D + rebD['pop_spike_count'] = pop_spike_count + rebD['pop_rate_hz'] = pop_rate_hz + ntime//=tReb2 + rebD['timeV']= np.linspace(0, (ntime - 1) * time_step2, ntime) + + print('rebinned rates: nchan=%d dt=%.3f sec rebin=%d' % (nchan, time_step2, tReb2)) + return rebD + +#...!...!.................... +def XXXselect_radius_slice(spikeD, data_path, data_name, idxState, verb=1): + """Select one spectral-radius slice from stacked spikes arrays if present.""" + spikes = spikeD['spikes'] + if spikes.ndim == 2: + return spikeD, None + assert spikes.ndim == 3, f"Expected spikes ndim 2 or 3, got shape={spikes.shape}" + + nR = spikes.shape[0] + assert 0 <= idxState < nR, f"idxState={idxState} out of range for spikes with nR={nR}" + + truthFF = os.path.join(data_path, f"{data_name}.simTruth.npz") + assert os.path.exists(truthFF), f"missing simTruth file: {truthFF}" + _, trueMD = read_data_npz(truthFF, verb=verb>0) + spect_radii = trueMD['dale_conf']['spect_radius'] + assert idxState < len(spect_radii), f"idxState={idxState} out of range, only {len(spect_radii)} states available" + R_sel = spect_radii[idxState] + print(f"\nSelected spectral radius [{idxState}]: R={R_sel:.3f} (out of {spect_radii})") + + spikeD_r = {} + for key, arr in spikeD.items(): + if isinstance(arr, np.ndarray) and arr.ndim > 0 and arr.shape[0] == nR: + spikeD_r[key] = arr[idxState] + else: + spikeD_r[key] = arr + return spikeD_r, R_sel + + + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__=="__main__": + args=get_parser() + np.set_printoptions(precision=3) + + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb>0) + if args.verb>1: pprint(spikeMD) + + S_oracle = None + if args.idxState >=0: # per state information + # Load simulation truth data (Dale matrices, biases, etc.) + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb>0) + if args.verb>1: + print("\nSimulation Truth Metadata:"); pprint(trueMD) + + dataYield = spikeD['spikes'] + Mstate,Nt, Nn = dataYield.shape + m=args.idxState + assert m0) + S_true = prismD['S_true'] + S_oracle = prismD['S_oracle'] + if args.verb > 0: print('loaded prismTruth S_true:', S_true.shape); + if args.verb > 1: pprint(prismMD) + spikeMD['sel_spect_radius'] =-77 + oraE = prismMD['oracle_eval'] + spikeMD['oracle_score'] = oraE['avr_score'] + print(f"gen, oracle avr score {oraE['avr_score']:.3f}, {args.dataName}") + print(f" {'state':>5s} {'score':>5s}") + print(f" {'-----':>5s} {'-----':>5s}") + for m, sc in enumerate(oraE['score_per_state']): + print(f" {m:5d} {sc:5.3f}") + + #-------------------------------- + # .... plotting ........ + spikeMD['short_name']=f"{args.dataName}_state{args.idxState}" + spikeMD['sel_state'] = int(args.idxState) + args.prjName=spikeMD['short_name'] + spikeMD['plot']={} + spikeMD['plot']['time_rangeLR']=np.array(args.time_range_sec) + + plot=Plotter(args) + + if 'a' in args.showPlots: + plot.freq_histo(spikeD,spikeMD,figId=1) + if 'b' in args.showPlots: + rebD=rebin_spike_rates(spikeD['spikes'], spikeMD, args.time_rebin2) + if S_oracle is not None: + rebD['S_oracle'] = S_oracle + plot.freq_vs_time(rebD,spikeMD,figId=2, S_true=S_true, S_oracle=S_oracle) + + plot.display_all() + print('M:done') + #pprint(expMD) #tmp diff --git a/causal_net/nonStation_ver3b_states_expr/PlotterBioExp.py b/causal_net/nonStation_ver3b_states_expr/PlotterBioExp.py new file mode 100644 index 00000000..a555658f --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/PlotterBioExp.py @@ -0,0 +1,373 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for biological experiment data visualization. + +This module provides specialized plotting capabilities for experimental +neural data analysis. The Plotter class extends PlotterBackbone to create +visualizations tailored for biological neural recordings including: +- Time series plots of neural activity patterns +- Statistical summaries and distribution analysis +- Data quality assessment plots +- Comparative analysis between experimental conditions + +Designed specifically for processing and visualizing data from biological +neural experiments, with automatic adaptation to experimental metadata. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from UtilBioExp import clip_rebD_time +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec + +from matplotlib.colors import LinearSegmentedColormap + +METRICS_HIST_COLUMNS = [ + "num_spikes", + "firing_rate", + "presence_ratio", + "snr", + "isi_violations_ratio", + "isi_violations_count", + "rp_contamination", + "rp_violations", + "sliding_rp_violation", + "amplitude_cutoff", + "amplitude_median", + "amplitude_cv_median", + "amplitude_cv_range", + "sync_spike_2", + "sync_spike_4", + "sync_spike_8", + "firing_range", + "sd_ratio", + "noise_cutoff", + "noise_ratio", +] + +METRICS_HIST_TITLE_UNITS = { + "firing_rate": "Hz", + "amplitude_median": "uV", +} + +METRICS_HIST_XMIN_ZERO = { + "snr", + "firing_rate", + "sync_spike_2", + "sync_spike_4", + "sync_spike_8", + "firing_range", + "sd_ratio", +} + +METRICS_HIST_LOG_Y = { + "presence_ratio", + "isi_violations_ratio", + "isi_violations_count", + "rp_violations", + "sliding_rp_violation", + "amplitude_cutoff", +} + +#...!...!.................... +def summary_column(md): + #print(sorted(md)) + rs=md['rate_summary'] + ds=md['data_selector'] + #pprint(md) + txt='dataset: '+md['short_name'] + txt += '\nnum acc neurons: %d' % md['num_neurons'] + txt += '\ndrop neur: %d <%.1f Hz, %d >%.1f Hz' % (ds['num_drop_neur_lo_hi_freq'][0],ds['freq_range'][0],ds['num_drop_neur_lo_hi_freq'][1],ds['freq_range'][1]) + txt+='\ndata type: %s\ntime_step=%.2f sec'%(md['data_type'],md['time_step_sec']) + txt += '\nMedian rate: %.1f Hz' % rs['median_spike_rate'] + txt += '\nrate range [%.1f, %.1f Hz]' % (rs['min_spike_rate'], rs['max_spike_rate']) + txt += '\nAvg Rate: %.1f±%.1f Hz' % (rs['avg_spike_rate'], rs['std_spike_rate']) + txt += '\nAvg Fano: %.2f±%.2f' % (rs['avg_fano_factor'], rs['std_fano_factor']) + + return txt + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def freq_histo(self,spikeD,md,figId=1): + + tit='dataset '+md['short_name'] + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,4)) + + # Create gridspec with 7:3 width ratio + gs = gridspec.GridSpec(1, 2, width_ratios=[7, 3]) + + dataYield, dataRates = spikeD['spikes'], spikeD['single_rates'] + + # Compute median + median_val = np.median(dataRates) + txtM= f'median : {median_val:.2f} Hz' + + #.... freq per channel + ax = self.plt.subplot(gs[0, 0]) + chanV=np.arange(dataRates.shape[0]) + ax.bar(chanV, dataRates , width=0.8, color='g', align='center', alpha=0.7) + + ax.axhline(median_val, color='red', linestyle='--', linewidth=1) + if md['data_type']=='simDale': + ax.set_yscale('log') + ax.grid() + ax.set(xlabel='input neuron index', ylabel='avr frequency (Hz)',title=tit) + ax.text( ax.get_xlim()[1]*0.1,median_val,txtM, + color='red', ha='center', va='bottom', fontsize=10) + + txt=summary_column(md) + ax.text(0.05, 0.3, txt, transform=ax.transAxes, fontsize=12,color='b') + + #........ freq histo ..... + ax = self.plt.subplot(gs[0, 1]) + ax.hist(dataRates,bins=20) + #ax.set_yscale('log') + ax.set(ylabel='num channels', xlabel='avr frequency (Hz)',title=tit) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + # Add median text annotation + ax.text(median_val+1, ax.get_ylim()[1]*0.5,txtM, + color='red', ha='left', va='bottom', fontsize=10,rotation=90) + ax.text(0.25, 0.3, txt, transform=ax.transAxes, fontsize=11,color='m') + + ax.grid() + + def plot_bioexp_sum_rate(self, ax, rebD, clip): + timeV = clip["timeV"] + rate1D = clip["rate1D"] + Tbin = clip["Tbin"] + tL, tR = clip["tL"], clip["tR"] + nchan = clip["nchan"] + ax.bar(timeV + Tbin * 0.5, rate1D, width=Tbin, color='orange', align='center', alpha=0.7) + ax.set(ylabel='sum rate (Hz)', title=f'sum rate from all {nchan} neurons') + ax.set_xlim(tL, tR) + ax.grid() + + def plot_bioexp_heatmap(self, fig, ax, rebD, dataset_name, clip): + rate2D = clip["rate2D"] + tL, tR = clip["tL"], clip["tR"] + nchan = clip["nchan"] + time_step = clip["time_step"] + tit0 = f'dataset: {dataset_name} nchan={nchan} Tbin={time_step:.2f} sec' + + original_cmap = self.plt.get_cmap('tab20c') + custom_cmap = original_cmap.copy() + custom_cmap.set_under('white') + + cax = ax.imshow( + rate2D.T, aspect='auto', cmap=custom_cmap, origin='lower', + extent=[tL, tR, 0, nchan], interpolation='nearest', + vmin=1, vmax=rate2D.max(), + ) + ax.set_ylabel('Neuron index') + ax.set_xlabel('Time (sec)') + ax.set_title(tit0) + cbar = fig.colorbar( + cax, ax=ax, orientation='horizontal', pad=0.17, + label=f'Instantanous freq (Hz) per neuron over Tbin={time_step:.1f} sec', + shrink=0.5, + ) + cbar.ax.tick_params(labelsize=10) + ax.grid() + + def freq_vs_time(self,rebD,md,figId=2): + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,11)) + + rateThr2=rebD['rate_thres2'] + clip = clip_rebD_time(rebD, md['plot']['time_rangeLR']) + print('iTL,R', clip['itL'], clip['itR']) + + highChanCnt = rebD['highChanCnt'][clip['itL']:clip['itR']] + timeV = clip['timeV'] + Tbin = clip['Tbin'] + tL, tR = clip['tL'], clip['tR'] + nchan = clip['nchan'] + + gs = fig.add_gridspec(4, 1, height_ratios=[0.15, 0.15, 0.54, 0.01]) + + ax = fig.add_subplot(gs[0, 0]) + ax.bar(timeV+Tbin*.5, highChanCnt, width=Tbin, color='forestgreen', align='center', alpha=0.7) + tit = (f'num neurons with instantaneous rate > thres={rateThr2:.0f} (Hz), ' + f'nchan={nchan}') + ax.set(ylabel='num neurons', title=tit) + ax.set_xlim(tL, tR) + ax.grid() + + ax = fig.add_subplot(gs[1, 0]) + self.plot_bioexp_sum_rate(ax, rebD, clip) + + ax = fig.add_subplot(gs[2, 0]) + self.plot_bioexp_heatmap(fig, ax, rebD, md['short_name'], clip) + + def _metrics_column(self, bioD, md, col_name): + cols = md["metrics_curated_columns"] + metrics = np.asarray(bioD["metrics_curated"]) + col_map = {str(c): i for i, c in enumerate(cols)} + assert col_name in col_map, ( + f"metrics_curated missing column {col_name!r}; have {cols}" + ) + return metrics[:, col_map[col_name]].astype(np.float64) + + def neuron_spatial(self, bioD, md, figId=3): + """MEA layout: loc_x/loc_y scatter colored by firing_rate (metrics_curated).""" + loc_x = self._metrics_column(bioD, md, "loc_x") + loc_y = self._metrics_column(bioD, md, "loc_y") + rate = self._metrics_column(bioD, md, "firing_rate") + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(8, 9)) + gs = gridspec.GridSpec(2, 1, height_ratios=[1.0, 0.06], hspace=0.12) + + ax = fig.add_subplot(gs[0, 0]) + ax.set_facecolor("white") + sc = ax.scatter( + loc_x, loc_y, c=rate, s=12, marker="s", + cmap="Blues", vmin=0.0, vmax=float(np.max(rate)), + linewidths=0.15, edgecolors="k", alpha=0.95, + ) + ax.set_aspect("equal", adjustable="box") + ax.set_xlabel("loc_x (um)", fontsize=11) + ax.set_ylabel("loc_y (um)", fontsize=11) + ax.tick_params(labelsize=9) + ax.grid(True, alpha=0.35) + tit = f"dataset: {md['short_name']} n={loc_x.shape[0]} neurons" + ax.set_title(tit, fontsize=11, pad=8) + + cax = fig.add_subplot(gs[1, 0]) + cb = fig.colorbar(sc, cax=cax, orientation="horizontal") + cb.set_label("Firing Rate [Hz]", fontsize=11) + cb.ax.tick_params(labelsize=9) + + fig.subplots_adjust(left=0.08, right=0.98, top=0.94, bottom=0.10) + + def _hist_panel_title(self, col_name): + unit = METRICS_HIST_TITLE_UNITS[col_name] + return f"{col_name} [{unit}]" + + def _fmt_metric_val(self, val): + av = abs(float(val)) + if av >= 100: + return f"{val:.0f}" + if av >= 1: + return f"{val:.3f}" + if av >= 0.01: + return f"{val:.4f}" + return f"{val:.2e}" + + def metrics_histograms(self, bioD, md, figId=4): + """1D histograms for curated unit-quality metrics (metrics_curated columns).""" + nrow, ncol = 3, 7 + n_panels = nrow * ncol + assert len(METRICS_HIST_COLUMNS) == n_panels - 1, ( + f"expected {n_panels - 1} metric columns, got {len(METRICS_HIST_COLUMNS)}" + ) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(18, 8)) + n_neur = np.asarray(bioD["metrics_curated"]).shape[0] + fr_lo, fr_hi = md["data_selector"]["freq_range"] + + for i, col_name in enumerate(METRICS_HIST_COLUMNS): + ax = fig.add_subplot(nrow, ncol, i + 1) + vals = self._metrics_column(bioD, md, col_name) + vals = vals[np.isfinite(vals)] + assert vals.size > 0, f"no finite values for metrics column {col_name!r}" + n_bins = min(40, max(10, int(np.sqrt(vals.size)))) + ax.hist(vals, bins=n_bins, color="steelblue", alpha=0.85, edgecolor="white") + if col_name in METRICS_HIST_LOG_Y: + ax.set_yscale("log") + + if col_name in METRICS_HIST_XMIN_ZERO: + ax.set_xlim(left=0.0) + + p16, p50, p84 = np.percentile(vals, [16, 50, 84]) + ylo, yhi = ax.get_ylim() + y_mark = 0.5 * (ylo + yhi) + ax.errorbar( + p50, y_mark, + xerr=[[p50 - p16], [p84 - p50]], + fmt="o", color="k", markersize=9, + capsize=4, capthick=1.2, elinewidth=1.2, zorder=5, + ) + txt = ( + f"p16={self._fmt_metric_val(p16)}\n" + f"med={self._fmt_metric_val(p50)}\n" + f"p84={self._fmt_metric_val(p84)}" + ) + ax.text( + 0.98, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=11, color="k", + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=2), + ) + + title = ( + self._hist_panel_title(col_name) + if col_name in METRICS_HIST_TITLE_UNITS else col_name + ) + ax.set_title(title, fontsize=14) + ax.tick_params(labelsize=7) + ax.grid(True, alpha=0.35) + + fig.add_subplot(nrow, ncol, n_panels).axis("off") + + fig.suptitle( + f"curated metrics, {md['short_name']}, N={n_neur} neurons, " + f"{fr_lo:g}< freq <{fr_hi:g} Hz", + fontsize=12, + ) + fig.tight_layout(rect=[0, 0, 1, 0.94], h_pad=1.62, w_pad=1.62) + + def metrics_correlations(self, bioD, md, figId=5): + """3x2 panel grid; panel 0: firing_rate vs single_rates scatter.""" + rate_sheet = self._metrics_column(bioD, md, "firing_rate") + rate_spikes = np.asarray(bioD["single_rates"], dtype=np.float64).ravel() + assert rate_sheet.shape[0] == rate_spikes.shape[0], ( + "firing_rate and single_rates length mismatch" + ) + + n = rate_sheet.shape[0] + r = float(np.corrcoef(rate_sheet, rate_spikes)[0, 1]) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(12, 10)) + fr_lo, fr_hi = md["data_selector"]["freq_range"] + + for k in range(6): + ax = fig.add_subplot(3, 2, k + 1) + if k == 0: + ax.scatter( + rate_sheet, rate_spikes, s=14, alpha=0.65, + color="steelblue", edgecolors="k", linewidths=0.2, + ) + lo = float(min(rate_sheet.min(), rate_spikes.min())) + hi = float(max(rate_sheet.max(), rate_spikes.max())) + pad = 0.05 * max(1e-9, hi - lo) + ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], + "k--", lw=0.9, alpha=0.5) + ax.set_xlabel("firing_rate (metrics_curated)", fontsize=10) + ax.set_ylabel("single_rates (bioExp)", fontsize=10) + ax.set_title(f"r={r:.3f}, n={n}", fontsize=11) + ax.grid(True, alpha=0.35) + else: + ax.axis("off") + + fig.suptitle( + f"metrics correlations, {md['short_name']}, N={n} neurons, " + f"{fr_lo:g}< freq <{fr_hi:g} Hz", + fontsize=12, + ) + fig.tight_layout(rect=[0, 0, 1, 0.94]) + diff --git a/causal_net/nonStation_ver3b_states_expr/PlotterPrismEM.py b/causal_net/nonStation_ver3b_states_expr/PlotterPrismEM.py new file mode 100644 index 00000000..bdaf8970 --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/PlotterPrismEM.py @@ -0,0 +1,1767 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for prism EM evaluation. +""" + +from toolbox.PlotterBackbone import PlotterBackbone +from PlotterBioExp import Plotter as PlotterBioExp +from UtilBioExp import clip_rebD_time +import numpy as np +from matplotlib.ticker import MaxNLocator +import matplotlib.colors as colors + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def _rebin_1d(self, x, merge): + merge = int(max(1, merge)) + if merge <= 1: + return np.asarray(x) + x = np.asarray(x) + n = x.shape[0] // merge + if n <= 0: + return x + return x[: n * merge].reshape(n, merge).mean(axis=1) + + def _rebin_2d(self, x, merge): + merge = int(max(1, merge)) + if merge <= 1: + return np.asarray(x) + x = np.asarray(x) + n = x.shape[0] // merge + if n <= 0: + return x + return x[: n * merge].reshape(n, merge, x.shape[1]).mean(axis=1) + + def _rebin_mode_1d(self, x, merge): + merge = int(max(1, merge)) + x = np.asarray(x) + if merge <= 1: + return x.astype(np.int64, copy=False) + n = x.shape[0] // merge + if n <= 0: + return x.astype(np.int64, copy=False) + x = x[: n * merge].reshape(n, merge) + out = np.zeros((n,), dtype=np.int64) + for i in range(n): + vals = x[i].astype(np.int64, copy=False) + vals = vals[vals >= 0] + if vals.size == 0: + out[i] = -1 + else: + out[i] = np.bincount(vals).argmax() + return out + + def summary_prismEM(self, fitD, md, figId=1): + """EM convergence overview: 2 rows × 3 columns. + + Row 1: E-step NLL vs EM iter | M-step NLL+L1 vs M-epoch | spectral radius vs M-epoch + Row 2: nz off-diag edges vs M-epoch | mean occupancy | A off-diag weight histogram + """ + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 6)) + fig.subplots_adjust(hspace=0.45, wspace=0.35) + + trainMD = md["train"] + short_name = md["short_name"] + + e_nll = np.asarray(fitD["e_nll_em"]) + m_nll = np.asarray(fitD["m_nll_epoch"]) + m_l1 = np.asarray(fitD["m_l1_epoch"]) + m_loss = np.asarray(fitD["m_loss_epoch"]) + rho = np.asarray(fitD["rho_epoch"]) + nz = np.asarray(fitD["nz_edges_epoch"]) + lr = np.asarray(fitD["learning_rates"]) + + n_em = trainMD["num_em_iters"] + m_per_em = trainMD["m_epochs"] + n_m_total = len(m_nll) + m_epochs = np.arange(1, n_m_total + 1) + em_iters = np.arange(1, len(e_nll) + 1) + + prune_em = trainMD["delay_em_iter_4_Aprune"] + rho_start_em = trainMD["delay_em_iter_4_ArhoMax"] + lr_drop_em = trainMD["delay_em_iter_4_lrDecay"] + prune_m_epoch = prune_em * m_per_em + rho_start_m_epoch = rho_start_em * m_per_em + lr_drop_m_epoch = lr_drop_em * m_per_em + + def draw_threshold_marker(ax, x_pos, x_max, txt, color,yFac=0.02): + if not (0 < x_pos <= x_max): + return + ax.axvline(x_pos, color=color, ls='--', lw=0.9, alpha=0.95) + y0, y1 = ax.get_ylim() + x0, x1 = ax.get_xlim() + x_off = 0.01 * max(1e-9, x1 - x0) + y_txt = y1 - yFac * (y1 - y0) + ax.text( + x_pos + x_off, y_txt, txt, rotation=90, color=color, fontsize=7, + ha='left', va='top', + bbox=dict(facecolor='white', alpha=0.55, edgecolor='none', pad=0.2) + ) + + jSkipEM = 0 + jSkipM = int(jSkipEM * m_per_em) + # ── Row 1, Col 1: E-step NLL vs EM iteration ──────────────── + ax = self.plt.subplot(2, 3, 1) + ax.plot(em_iters[jSkipEM:], e_nll[jSkipEM:], 'o-', color='tab:blue', markersize=3, + linewidth=1.2) + ax.set(title="E-step NLL", xlabel="EM iteration", + ylabel="NLL / bin") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_em, len(e_nll), "start Aprune", "k") + draw_threshold_marker(ax, rho_start_em, len(e_nll), "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_em, len(e_nll), "start_lrDrop", "tab:gray", yFac=0.6) + txt = (f"pgd_iter={trainMD['pgd_iter']}\n" + f"lr_E={trainMD['lr_estep']}\n" + f"λ₂={trainMD['lambda2']}") + ax.text(0.97, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=8, + bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')) + + # ── Row 1, Col 2: M-step NLL + L1 vs M-epoch ──────────────── + ax = self.plt.subplot(2, 3, 2) + ax.plot(m_epochs[jSkipM:], m_nll[jSkipM:], color='tab:blue', linewidth=1, label='NLL') + ax.set_ylabel('NLL', color='tab:blue') + ax.tick_params(axis='y', labelcolor='tab:blue') + ax.set(title="M-step loss", xlabel="M-epoch (global)") + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + ax2.plot(m_epochs[jSkipM:], m_l1[jSkipM:] , color='tab:red', linewidth=1, + linestyle='--', label='L1') + ax2.set_ylabel('L1', color='tab:red') + ax2.tick_params(axis='y', labelcolor='tab:red') + + lines1, lab1 = ax.get_legend_handles_labels() + lines2, lab2 = ax2.get_legend_handles_labels() + ax.legend(lines1 + lines2, lab1 + lab2, fontsize=7, loc='upper right') + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + + txt = (f"lr_M={trainMD['lr_mstep']}\n" + f"L1 λ3={trainMD['lambda3']}\n" + f"batch={trainMD['batch_size']}") + ax.text(0.03, 0.03, txt, transform=ax.transAxes, + va="bottom", ha="left", fontsize=8, + bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')) + + # ── Row 1, Col 3: spectral radius vs M-epoch ──────────────── + ax = self.plt.subplot(2, 3, 3) + rho_max = trainMD["rho_max"] + ax.plot(m_epochs[jSkipM:], rho[jSkipM:] , color='tab:green', linewidth=1, + label='ρ(A)') + ax.axhline(rho_max, color='red', ls='--', lw=1, + label=f'ρ_max={rho_max}') + ax.set(title="Spectral radius ρ(A)", xlabel="M-epoch (global)", + ylabel="ρ") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + ax.legend(fontsize=8) + + # ── Row 2, Col 1: non-zero off-diag edges vs M-epoch ──────── + ax = self.plt.subplot(2, 3, 4) + ax.plot(m_epochs[jSkipM:], nz[jSkipM:], color='tab:purple', linewidth=1) + ax.set(title=f"Non-zero off-diag edges (|A|>{trainMD['minW']})", + xlabel="M-epoch (global)", ylabel="count") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + + # learning rate on twin axis + ax2 = ax.twinx() + ax2.plot(m_epochs[jSkipM:], lr[jSkipM:], color='tab:orange', linewidth=0.8, + linestyle='--', alpha=0.6) + ax2.set_ylabel('learning rate', color='tab:orange') + ax2.tick_params(axis='y', labelcolor='tab:orange') + + # ── Row 2, Col 2: mean occupancy ───────────────────────────── + ax = self.plt.subplot(2, 3, 5) + c_hat = np.asarray(fitD["c_hat"]) + occ = c_hat.mean(axis=0) + M = len(occ) + ax.bar(np.arange(M), occ, color='tab:blue', alpha=0.7) + ax.set(title="Mean occupancy", xlabel="state", ylabel="mean c") + ax.set_ylim(0.0, 1.0) + ax.set_xticks(np.arange(M)) + ax.grid(True, alpha=0.3, axis='y') + + t0s, t1s = trainMD["time_range_sec"] + txt = (f"N={trainMD['num_neurons']} M={M}\n" + f"T=[{t0s:.0f},{t1s:.0f}]s\n" + f"bins={trainMD['num_time_bins']}") + ax.text(0.97, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=8, + bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')) + + # ── Row 2, Col 3: A off-diagonal weight histogram ──────────── + ax = self.plt.subplot(2, 3, 6) + A_fit = np.asarray(fitD["A_hat"]) + Nn = A_fit.shape[0] + diag_mask = np.eye(Nn, dtype=bool) + A_off = A_fit[~diag_mask] + valid = np.abs(A_off) > 1e-10 + A_off_nz = A_off[valid] + n_edges = int(A_off_nz.size) + ax.hist(A_off_nz, bins=100, color='g', alpha=0.8) + ax.set_yscale('log') + ax.set(title=f"A off-diagonal, {n_edges} edges", + xlabel="edge value", ylabel="edges") + ax.grid(True, alpha=0.3) + + fig.suptitle(f"Prism EM: {short_name}, " + f"K_EM={n_em}×K_M={m_per_em}", + fontsize=12) + fig.tight_layout(rect=[0, 0, 1, 0.95]) + + def state_init_prismEM(self, fitD, md, figId=2, time_reb=20): + """Two-panel plot: initialization vs truth state trajectories.""" + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(13.5, 8.0)) + gs = fig.add_gridspec(3, 3, height_ratios=[1.0, 1.0, 0.85], hspace=0.75, wspace=0.35) + + trainMD = md["train"] + t0_bin, t1_bin = trainMD["time_range_bins"] + dt = float(trainMD["time_step_sec"]) + + S_init = np.asarray(fitD["S_init"]) + C_init = np.asarray(fitD["c_init"]) + S_true = np.asarray(md["S_true"])[t0_bin : t1_bin + 1] + C_true = np.asarray(md["C_true"])[t0_bin : t1_bin + 1] + + assert S_init.shape[0] == S_true.shape[0] == C_init.shape[0] == C_true.shape[0], \ + "S_init/S_true/C_init/C_true must have matching lengths" + n_cmp = S_init.shape[0] + + # Use original EM time bins (dt), not coarse rate-bin or plotting rebinning. + S_init_cl = 1.0 - np.max(C_init, axis=1) + t = (t0_bin + np.arange(n_cmp)) * dt + n_states = C_init.shape[1] + + ax = fig.add_subplot(gs[0, :]) + ax.step(t, S_init, where="post", color="k", linewidth=1.0, label="S_init") + ax.plot(t, S_init, linestyle="none", marker=".", markersize=1.5, color="k", alpha=0.6) + lo = np.clip(S_init - S_init_cl, -1.0, float(n_states - 1)) + hi = np.clip(S_init + S_init_cl, -1.0, float(n_states - 1)) + ax.fill_between(t, lo, hi, color="gray", alpha=0.3, label="S_init_CL") + ax.set(title="Initialization", xlabel="time (s)", ylabel="state") + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + for m in range(C_init.shape[1]): + ax2.step(t, C_init[:, m], where="post", linewidth=0.8, alpha=0.85, label=f"C_init[{m}]") + ax2.plot(t, C_init[:, m], linestyle="none", marker=".", markersize=1.4, alpha=0.55) + ax2.set_ylabel("C_init") + ax2.set_ylim(0.0, 1.0) + + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + ax = fig.add_subplot(gs[1, :]) + ax.step(t, S_true, where="post", color="k", linewidth=1.0, label="S_true") + ax.plot(t, S_true, linestyle="none", marker=".", markersize=1.5, color="k", alpha=0.6) + ax.set(title="Truth", xlabel="time (s)", ylabel="state") + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + for m in range(C_true.shape[1]): + ax2.step(t, C_true[:, m], where="post", linewidth=0.8, alpha=0.85, label=f"C_true[{m}]") + ax2.plot(t, C_true[:, m], linestyle="none", marker=".", markersize=1.4, alpha=0.55) + ax2.set_ylabel("C_true") + + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + # Bottom row, left: histogram of classification frequency with thresholds. + ax = fig.add_subplot(gs[2, 0]) + init_state_md = md["init_state"] + + freq_h1d = np.asarray(fitD["freq_h1d"], dtype=np.float64).ravel() + bins = min(40, max(10, int(np.sqrt(max(1, freq_h1d.size))))) + ax.hist(freq_h1d, bins=bins, color='tab:blue', alpha=0.75) + rate_bin_sec = float(init_state_md["rate_bin_sec"]) + ax.set(title=f"Rate (time bin {rate_bin_sec:g} s)", xlabel="rate", ylabel="count") + ax.grid(True, alpha=0.3) + + thres = [float(x) for x in init_state_md["rate_thres"]] + for thr in thres: + ax.axvline(float(thr), color='k', linestyle='--', linewidth=1.0, alpha=0.9) + + # Label state regions 0,1,... inside histogram intervals. + xlo, xhi = ax.get_xlim() + ylo, yhi = ax.get_ylim() + edges = [xlo] + thres + [xhi] + for i in range(max(0, len(edges) - 1)): + xc = 0.5 * (edges[i] + edges[i + 1]) + ax.text( + xc, ylo + 0.88 * (yhi - ylo), f"{i}", + ha="center", va="center", fontsize=10, color="k", + bbox=dict(facecolor="white", alpha=0.6, edgecolor="none"), + ) + + # Bottom row, middle: B_init correlation vs log(single_rates). + ax_mid = fig.add_subplot(gs[2, 1]) + b_init = np.asarray(fitD["B_init"], dtype=np.float64).ravel() + single_rates = np.asarray(fitD["single_rates"], dtype=np.float64).ravel() + assert b_init.ndim == 1 and b_init.size > 0, "Expected B_init to be 1D and non-empty" + assert b_init.size == single_rates.size, "B_init and single_rates must have matching lengths" + + n = b_init.size + assert n > 1, "Need at least 2 points for correlation plot" + x = np.log(np.clip(single_rates, 1e-12, None)) + y = b_init + ax_mid.scatter(x, y, s=10, alpha=0.75, color='tab:green', edgecolors='none') + corr = float(np.corrcoef(x, y)[0, 1]) + lo = float(min(np.min(x), np.min(y))) + hi = float(max(np.max(x), np.max(y))) + ax_mid.plot([lo, hi], [lo, hi], color='k', linestyle='--', linewidth=0.8, alpha=0.6) + ax_mid.set(title=f"B_init r={corr:.3f}", + xlabel="log(single_rates)", ylabel="B_init") + ax_mid.set_aspect("equal", adjustable="box") + ax_mid.grid(True, alpha=0.3) + + # Bottom row, right: reserved cell for future diagnostics. + ax_right = fig.add_subplot(gs[2, 2]) + ax_right.axis("off") + + t0s, t1s = trainMD["time_range_sec"] + fig.suptitle( + f"Prism EM Init: {md['short_name']} T=[{t0s:.1f}, {t1s:.1f}] s", + fontsize=12, + ) + + def _draw_A_matrix( + self, ax, A, title, num_exc=None, norm_map=None, + exc_inh_from_offdiag=False, exc_inh_min_abs=0.0, + ): + A = np.asarray(A) + if norm_map is None: + vmin = float(np.min(A)) + vmax = float(np.max(A)) + if np.isclose(vmin, vmax): + vmax = vmin + 1e-9 + norm_map = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + im = ax.imshow(A, aspect=1.0, origin='lower', cmap='bwr', + norm=norm_map, interpolation='nearest') + N = A.shape[0] + n_exc = n_inh = n_und = None + if num_exc is not None: + n_exc = int(num_exc) + n_inh = N - n_exc + ax.axhline(n_exc - 0.5, color='k', ls='--', lw=0.8) + ax.axvline(n_exc - 0.5, color='k', ls='--', lw=0.8) + elif exc_inh_from_offdiag: + offdiag = A.astype(np.float64).copy() + np.fill_diagonal(offdiag, 0.0) + row_sum = offdiag.sum(axis=1) + min_abs = float(exc_inh_min_abs) + n_exc = int(np.sum(row_sum > min_abs)) + n_inh = int(np.sum(row_sum < -min_abs)) + n_und = int(np.sum(np.abs(row_sum) < min_abs)) + if n_exc is not None and n_inh is not None and (n_exc + n_inh) > 0: + txt = f"exc={n_exc}\ninh={n_inh}" + if n_und is not None: + txt += f"\nund={n_und}" + ax.text( + 0.03, 0.97, txt, + transform=ax.transAxes, va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=0.3), + ) + ax.plot([0, N], [0, N], '--', lw=0.8, color='magenta') + ax.set_xlim(-0.5, N + 0.5) + ax.set_ylim(-0.5, N + 0.5) + ax.grid(True, alpha=0.25) + ax.set_title(title) + ax.set_xlabel('presyn. neuron index (output)') + ax.set_ylabel('postsyn. neuron index (input)') + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + def _get_A_true_info(self, md): + A_true = np.asarray(md["A_true"]) + if A_true.ndim == 3: + A_true = A_true[0] + assert A_true.ndim == 2 and A_true.shape[0] == A_true.shape[1], "A_true must be square 2D" + N = A_true.shape[0] + num_exc = md["dale_conf"]["num_excite"] + off_mask = ~np.eye(N, dtype=bool) + n_diag = int(np.eye(N, dtype=bool).sum()) + n_off = int(np.count_nonzero(A_true[off_mask])) + vmin = float(np.min(A_true)) + vmax = float(np.max(A_true)) + if np.isclose(vmin, vmax): + vmax = vmin + 1e-9 + shared_norm = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + return A_true, N, num_exc, n_diag, n_off, shared_norm + + def _scatter_true_vs_cmp(self, ax, x_true, y_cmp, title, color, split_by_true_sign=False): + x_true = np.asarray(x_true, dtype=np.float64).ravel() + y_cmp = np.asarray(y_cmp, dtype=np.float64).ravel() + n = min(x_true.size, y_cmp.size) + x = x_true[:n] + y = y_cmp[:n] + ax.scatter(x, y, alpha=0.7, color=color, marker='.', s=9) + + if n > 0: + lo = float(min(np.min(x), np.min(y))) + hi = float(max(np.max(x), np.max(y))) + pad = 0.05 * max(1e-9, hi - lo) + lo -= pad + hi += pad + ax.plot([lo, hi], [lo, hi], color='k', linestyle='--', linewidth=1.0, alpha=0.6) + ax.axhline(0.0, color='k', linestyle='--', linewidth=0.8, alpha=0.6) + ax.axvline(0.0, color='k', linestyle='--', linewidth=0.8, alpha=0.6) + ax.set_xlim(lo, hi) + ax.set_ylim(lo, hi) + + if split_by_true_sign: + m_neg = x < 0.0 + m_pos = x > 0.0 + if np.any(m_neg): + xm = float(np.mean(x[m_neg])) + ym = float(np.mean(y[m_neg])) + ax.plot([xm], [ym], marker='+', markersize=14, markeredgewidth=2.5, color='k') + if np.any(m_pos): + xm = float(np.mean(x[m_pos])) + ym = float(np.mean(y[m_pos])) + ax.plot([xm], [ym], marker='+', markersize=14, markeredgewidth=2.5, color='k') + else: + xm = float(np.mean(x)) + ym = float(np.mean(y)) + ax.plot([xm], [ym], marker='+', markersize=14, markeredgewidth=2.5, color='k') + + ax.set_aspect('equal', adjustable='box') + ax.grid(True, alpha=0.4) + ax.set_title(title) + ax.set_xlabel('true weight') + ax.set_ylabel('fitted') + + def _count_A_edges(self, A, minW=0.0): + A = np.asarray(A) + N = A.shape[0] + off_mask = ~np.eye(N, dtype=bool) + n_diag = int(np.eye(N, dtype=bool).sum()) + if minW > 0: + n_off = int(np.sum(off_mask & (np.abs(A) >= minW))) + else: + n_off = int(np.count_nonzero(A[off_mask])) + return N, n_diag, n_off + + def _A_for_display_minW(self, A, minW): + """Zero off-diagonal entries with |A_ij| < minW; keep diagonal.""" + A = np.asarray(A, dtype=np.float64).copy() + N = A.shape[0] + off_mask = ~np.eye(N, dtype=bool) + A[off_mask & (np.abs(A) < float(minW))] = 0.0 + return A + + def _A_prune_from_fitD(self, fitD): + """Saved A_prune matrix.""" + return np.asarray(fitD["A_prune"], dtype=np.float64) + + def _postsyn_nz_edge_stats(self, A, minW): + """Per postsynaptic row: count and sum of off-diagonal edges with |A| >= minW.""" + A = np.asarray(A, dtype=np.float64) + N = A.shape[0] + minW = float(minW) + strong = ~np.eye(N, dtype=bool) & (np.abs(A) >= minW) + cnt = strong.sum(axis=1).astype(np.int64) + sum_post = np.where(strong, A, 0.0).sum(axis=1) + return cnt, sum_post + + def _overlay_single_rates(self, ax, single_rates, x_post): + rate = np.asarray(single_rates, dtype=np.float64).ravel() + ax2 = ax.twinx() + ax2.plot(x_post, rate, color='tab:orange', linewidth=0.9, alpha=0.75) + ax2.set_ylabel("frequency (Hz)", color='tab:orange') + ax2.tick_params(axis='y', labelcolor='tab:orange') + return ax2 + + def _draw_A_offdiag_hist(self, ax, A, title, minW=0.02, diagnostics_txt=None): + A = np.asarray(A) + N = A.shape[0] + off_mask = ~np.eye(N, dtype=bool) + A_off = A[off_mask] + A_off_nz = A_off[np.abs(A_off) > 1e-12] + ax.hist(A_off_nz, bins=120, color='saddlebrown', alpha=0.85) + ax.set_yscale('log') + ax.grid(True, alpha=0.35) + ax.axvline(float(minW), color='green', ls='--', lw=1.0, alpha=0.9) + ax.axvline(-float(minW), color='green', ls='--', lw=1.0, alpha=0.9) + ax.set_title(title) + ax.set_xlabel("edge value") + ax.set_ylabel("edges") + if diagnostics_txt is not None: + ax.text( + 0.05, 0.70, diagnostics_txt, transform=ax.transAxes, + va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"), + family="monospace", + ) + + def _A_init_diagnostics_txt(self, md): + initA = md["init_A"] + return ( + "A-init diagnostics:\n" + f" cond(YpYp) = {initA['cond_YpYp']:.2e}\n" + f" rho(A_ols) = {initA['rho_A_init']:.3f}\n" + f" R2 = {initA['R2_1step']:.3f}\n" + f" ||A||_F = {initA['fro_A_init']:.3f}\n" + f" bins_used = {initA['num_bins_used']}/{initA['num_bins_total']}" + ) + + def A_fitted_prismEM(self, fitD, md, single_rates, minW=0.02, figId=8): + """Fitted A only: 2 rows x 4 cols. + + Cols: heatmap | off-diagonal histogram | # non-zero edges vs postsyn | + sum non-zero edge vs postsyn. + Non-zero edge: off-diagonal with |A_ij| >= minW. + Neuron axes use frequency-sorted indices as stored in the fit (no truth). + Bottom-row cols 3-4 overlay single_rates (Hz) on right y-axis. + """ + minW = float(minW) + A_init = np.asarray(fitD["A_init"]) + A_hat = np.asarray(fitD["A_hat"]) + assert A_init.shape == A_hat.shape and A_init.ndim == 2, "A_init and A_hat must match" + N = A_init.shape[0] + x_post = np.arange(N, dtype=np.int64) + single_rates = np.asarray(single_rates, dtype=np.float64).ravel() + assert single_rates.shape[0] == N, \ + f"single_rates length {single_rates.shape[0]} != N={N}" + + num_exc = None + dale_conf = md.get("dale_conf") + if dale_conf is not None: + num_exc = dale_conf.get("num_excite") + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(18, 8)) + gs = fig.add_gridspec(2, 4, hspace=0.55, wspace=0.45) + + for row, (A, label, diag_txt) in enumerate([ + (A_init, "A_init", self._A_init_diagnostics_txt(md)), + (A_hat, "A_hat", None), + ]): + _, n_diag, n_off = self._count_A_edges(A, minW=minW) + mat_title = f"{label}, nEdges={n_diag}+{n_off}" + A_disp = self._A_for_display_minW(A, minW) + cnt_post, sum_post = self._postsyn_nz_edge_stats(A, minW) + + ax = fig.add_subplot(gs[row, 0]) + self._draw_A_matrix(ax, A_disp, mat_title, num_exc=num_exc) + + ax = fig.add_subplot(gs[row, 1]) + self._draw_A_offdiag_hist( + ax, A, f"{label} off-diagonal", + minW=minW, diagnostics_txt=diag_txt + ) + + ax = fig.add_subplot(gs[row, 2]) + ax.plot(x_post, cnt_post, color='tab:blue', linewidth=1.0) + med_cnt = int(np.median(cnt_post)) + ax.axhline(med_cnt, color='green', ls='--', lw=1.0, alpha=0.9) + ax.text( + 0.02, 0.95, f"median={med_cnt:d}", transform=ax.transAxes, + va="top", ha="left", fontsize=9, color="green", + ) + ax.set_title(f"{label}: # non-zero edges") + ax.set_xlabel("postsyn. neuron index") + ax.set_ylabel("# non-zero edges") + ax.grid(True, alpha=0.35) + if row == 1: + self._overlay_single_rates(ax, single_rates, x_post) + + ax = fig.add_subplot(gs[row, 3]) + ax.plot(x_post, sum_post, color='tab:green', linewidth=1.0) + ax.axhline(0.0, color='k', ls='--', lw=1.0, alpha=0.9) + ax.set_title(f"{label}: sum non-zero edges") + ax.set_xlabel("postsyn. neuron index") + ax.set_ylabel("sum edge value") + ax.grid(True, alpha=0.35) + if row == 1: + self._overlay_single_rates(ax, single_rates, x_post) + + fig.suptitle( + f"A fitted summary: {md['short_name']}, minW={minW:g}, neur. freq. sorted", + fontsize=13, + ) + fig.tight_layout(rect=[0, 0, 1, 0.95]) + + def compare_A_vs_truth(self, A_cmp, md, cmp_label='A_init', figId=3): + """Reusable canvas to compare an A-matrix candidate against A_true.""" + A_true, N, num_exc, n_diag, n_off, shared_norm = self._get_A_true_info(md) + A_cmp = np.asarray(A_cmp) + assert A_cmp.ndim == 2 and A_cmp.shape == A_true.shape, "A_compare must match A_true shape" + + off_mask = ~np.eye(N, dtype=bool) + neg_mask = off_mask & (A_true < 0) + pos_mask = off_mask & (A_true > 0) + diag_mask = np.eye(N, dtype=bool) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(14, 8)) + gs = fig.add_gridspec(2, 3, hspace=0.55, wspace=0.45) + + ax = fig.add_subplot(gs[0, 0]) + self._draw_A_matrix( + ax, A_true, f"True Dale, N{N}, nEdges={n_diag}+{n_off}", + num_exc=num_exc, norm_map=shared_norm + ) + + ax = fig.add_subplot(gs[0, 1]) + self._draw_A_matrix( + ax, A_cmp, f"{cmp_label}, nEdges={n_diag}+{n_off}", + num_exc=num_exc, norm_map=shared_norm + ) + + ax = fig.add_subplot(gs[0, 2]) + A_off = A_cmp[off_mask] + A_off_nz = A_off[np.abs(A_off) > 1e-12] + ax.hist(A_off_nz, bins=120, color='saddlebrown', alpha=0.85) + ax.set_yscale('log') + ax.grid(True, alpha=0.35) + ax.set_title(f"{cmp_label} off-diagonal") + ax.set_xlabel("edge value") + ax.set_ylabel("edges") + + txt = self._A_init_diagnostics_txt(md) + ax.text( + 0.05, 0.70, txt, transform=ax.transAxes, + va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"), + family="monospace", + ) + + ax = fig.add_subplot(gs[1, 0]) + self._scatter_true_vs_cmp( + ax, + A_true[neg_mask], + A_cmp[neg_mask], + f"{cmp_label} :neg TP", + color='tab:green', + ) + + ax = fig.add_subplot(gs[1, 1]) + self._scatter_true_vs_cmp( + ax, + A_true[pos_mask], + A_cmp[pos_mask], + f"{cmp_label} :pos TP", + color='tab:green', + ) + + ax = fig.add_subplot(gs[1, 2]) + self._scatter_true_vs_cmp( + ax, + A_true[diag_mask], + A_cmp[diag_mask], + f"{cmp_label} :diag", + color='salmon', + split_by_true_sign=True, + ) + + fig.suptitle(f"A_true vs. {cmp_label} comparison: {md['short_name']}", fontsize=13) + fig.tight_layout(rect=[0, 0, 1, 0.95]) + + def matrix_init_prismEM(self, fitD, md, figId=3, est_key="A_init", est_label="A_init"): + """Wrapper for current EM use case; reusable for A_hat later.""" + self.compare_A_vs_truth(fitD[est_key], md, cmp_label=est_label, figId=figId) + + def _node_outgoing_edge_stats(self, A, minW): + """Per postsynaptic row: Nedge count and Sedge sum, |A|>minW off-diag.""" + A = np.asarray(A, dtype=np.float64) + N = A.shape[0] + minW = float(minW) + strong = ~np.eye(N, dtype=bool) & (np.abs(A) > minW) + Nedge = strong.sum(axis=1).astype(np.int64) + Sedge = np.where(strong, A, 0.0).sum(axis=1) + return Nedge, Sedge + + def edge_recovery_prismEM(self, fitD, md, minW=0.02, figId=4, est_key="A_hat", est_label="A_hat"): + """One-row, four-panel edge recovery: A_true | A_est edges | confusion | stats.""" + A_true, N, num_exc, n_diag, n_off, shared_norm = self._get_A_true_info(md) + A_est = np.asarray(fitD[est_key]) + assert A_est.ndim == 2 and A_est.shape == A_true.shape, "A_est must match A_true shape" + + E_true = np.asarray(md["E_true"]).astype(bool) + if E_true.ndim == 3: + E_true = E_true[0] + assert E_true.shape == A_true.shape, "E_true shape must match A_true" + + off_diag = ~np.eye(N, dtype=bool) + E_t = E_true & off_diag + E_hat = (np.abs(A_est) > float(minW)) & off_diag + + TP = E_t & E_hat + FP = (~E_t) & E_hat + FN = E_t & (~E_hat) + TN = (~E_t) & (~E_hat) + + tp = int(TP.sum()); fp = int(FP.sum()); fn = int(FN.sum()); tn = int(TN.sum()) + precision = tp / max(1, tp + fp) + recall = tp / max(1, tp + fn) + f1 = 2 * precision * recall / max(1e-12, precision + recall) + acc = (tp + tn) / max(1, tp + tn + fp + fn) + + conf_map = np.zeros((N, N), dtype=np.int8) + conf_map[FN] = 1 + conf_map[FP] = 2 + conf_map[TP] = 3 + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 3.8)) + kw = dict(origin='lower', interpolation='nearest') + + ax = self.plt.subplot(1, 4, 1) + self._draw_A_matrix( + ax, A_true, f"True Dale, N{N}, nEdges={n_diag}+{n_off}", + num_exc=num_exc, norm_map=shared_norm + ) + + ax = self.plt.subplot(1, 4, 2) + ax.imshow(E_hat.astype(float), cmap='Greys', vmin=0, vmax=1, **kw) + ax.set(title=f"{est_label} edges (n={int(E_hat.sum())}, minW={float(minW):g})", + xlabel='presyn. neuron (output)', ylabel='postsyn. neuron (input)') + ax.plot([0, N-1], [0, N-1], '--', lw=0.8, color='magenta') + + cmap_conf = colors.ListedColormap(['white', 'magenta', 'red', 'green']) + norm_conf = colors.BoundaryNorm([-0.5, 0.5, 1.5, 2.5, 3.5], cmap_conf.N) + ax = self.plt.subplot(1, 4, 3) + im = ax.imshow(conf_map, cmap=cmap_conf, norm=norm_conf, **kw) + cbar = self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + cbar.set_ticks([0, 1, 2, 3]) + cbar.set_ticklabels(['TN', 'FN', 'FP', 'TP']) + ax.set(title='Confusion map', + xlabel='presyn. neuron', ylabel='postsyn. neuron') + ax.plot([0, N-1], [0, N-1], '--', lw=0.8, color='k') + + ax = self.plt.subplot(1, 4, 4) + vals = [tp, fp, fn] + ax.bar(['TP', 'FP', 'FN'], vals, color=['green', 'red', 'magenta']) + ax.set(title='stats', ylabel='count') + ax.grid(axis='y', alpha=0.4) + y_pos = 0.25 * max(1, max(vals)) + for name, val in zip(['TP', 'FP', 'FN'], vals): + ax.text(name, y_pos, str(val), ha='center', va='center', fontsize=10) + txt = (f"precision={precision:.3f}\nrecall={recall:.3f}\n" + f"f1={f1:.3f}\nacc={acc:.3f}") + ax.text(0.55, 0.60, txt, transform=ax.transAxes, fontsize=9) + + fig.suptitle( + f"A-matrix edge recovery, minW={float(minW):g} (off-diag only): {md['short_name']}", + fontsize=12) + fig.tight_layout() + + def _draw_A_matrix_summary_3cols( + self, fig, gs, row, A, label, minW, single_rates, num_exc=None, + ): + """Heatmap, off-diagonal hist, # non-zero edges (as in -p i bottom row, cols 0-2).""" + A = np.asarray(A, dtype=np.float64) + N = A.shape[0] + minW = float(minW) + x_post = np.arange(N, dtype=np.int64) + single_rates = np.asarray(single_rates, dtype=np.float64).ravel() + assert single_rates.shape[0] == N, ( + f"single_rates length {single_rates.shape[0]} != N={N}" + ) + + _, n_diag, n_off = self._count_A_edges(A, minW=minW) + A_disp = self._A_for_display_minW(A, minW) + cnt_post, _ = self._postsyn_nz_edge_stats(A, minW) + sum_Nedge = int(np.sum(cnt_post)) + + ax = fig.add_subplot(gs[row, 0]) + self._draw_A_matrix( + ax, A_disp, f"{label}, nEdges={n_diag}+{n_off}", num_exc=num_exc, + exc_inh_from_offdiag=(num_exc is None), exc_inh_min_abs=minW, + ) + + ax = fig.add_subplot(gs[row, 1]) + self._draw_A_offdiag_hist(ax, A, f"{label} off-diagonal", minW=minW) + ax.text( + 0.98, 0.97, f"sum Nedge={sum_Nedge:d}", transform=ax.transAxes, + va="top", ha="right", fontsize=9, color="k", + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=2), + ) + + ax = fig.add_subplot(gs[row, 2]) + ax.plot(x_post, cnt_post, color='tab:blue', linewidth=1.0) + med_cnt = int(np.median(cnt_post)) + ax.axhline(med_cnt, color='green', ls='--', lw=1.0, alpha=0.9) + ax.text( + 0.02, 0.95, f"median={med_cnt:d}", transform=ax.transAxes, + va="top", ha="left", fontsize=9, color="green", + ) + ax.set_title(f"{label}: # non-zero edges") + ax.set_xlabel("postsyn. neuron index") + ax.set_ylabel("# non-zero edges") + ax.grid(True, alpha=0.35) + self._overlay_single_rates(ax, single_rates, x_post) + + def node_outgoing_edge_stats_prismEM( + self, fitD, md, single_rates, neuron_type, + minW=0.02, figId=12, est_key="A_hat", est_label="A_hat", + ): + """2x3 canvas: row0 Nedge/Sedge stats; row1 A_prune summary (-p i style).""" + A_est = np.asarray(fitD[est_key], dtype=np.float64) + A_prune = self._A_prune_from_fitD(fitD) + assert A_est.ndim == 2 and A_est.shape[0] == A_est.shape[1], "A_est must be square" + assert A_prune.shape == A_est.shape, "A_prune shape must match A_hat" + N = A_est.shape[0] + minW = float(minW) + + Nedge, Sedge = self._node_outgoing_edge_stats(A_est, minW) + r_ns = self._corrcoef_safe(Nedge, Sedge) + n_bins = max(10, min(40, int(np.sqrt(N)) * 2)) + neuron_type = np.asarray(neuron_type, dtype=np.int8) + assert neuron_type.shape[0] == N, "neuron_type length must match A_hat rows" + n_exc = int(np.sum(neuron_type > 0)) + n_inh = int(np.sum(neuron_type < 0)) + n_und = int(np.sum(neuron_type == 0)) + + def add_hist_percentile_marker(ax, vals, fmt=".3f"): + vals = np.asarray(vals, dtype=np.float64) + vals = vals[np.isfinite(vals)] + p16, p50, p84 = np.percentile(vals, [16, 50, 84]) + ylo, yhi = ax.get_ylim() + y_mark = 0.5 * (ylo + yhi) + ax.errorbar( + p50, y_mark, + xerr=[[p50 - p16], [p84 - p50]], + fmt="o", color="k", markersize=8, + capsize=4, capthick=1.2, elinewidth=1.2, zorder=5, + ) + p16_txt = format(p16, fmt) + p50_txt = format(p50, fmt) + p84_txt = format(p84, fmt) + txt = f"p16={p16_txt}\nmed={p50_txt}\np84={p84_txt}" + ax.text( + 0.98, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=9, color="k", + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=2), + ) + + num_exc = None + dale_conf = md.get("dale_conf") + if dale_conf is not None: + num_exc = dale_conf.get("num_excite") + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(14, 8)) + gs = fig.add_gridspec(2, 3, height_ratios=[1.0, 1.15], hspace=0.45, wspace=0.35) + + ax = fig.add_subplot(gs[0, 0]) + for cls, color, label in [ + (1, "magenta", f"exc={n_exc}"), + (-1, "#39FF14", f"inh={n_inh}"), + (0, "salmon", f"und={n_und}"), + ]: + mask = neuron_type == cls + if cls == 0: + ax.scatter( + Nedge[mask], Sedge[mask], s=18, marker='o', + alpha=0.80, color=color, edgecolors="k", + linewidths=0.35, label=label, + ) + else: + ax.scatter( + Nedge[mask], Sedge[mask], s=16, marker='.', + alpha=0.80, color=color, label=label, + ) + ax.set( + title=f"Nedge vs Sedge r={r_ns:.3f}", + xlabel="Nedge (# outgoing edges per row)", + ylabel="Sedge", + ) + ax.legend(loc="best", fontsize=9) + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[0, 1]) + ax.hist(Nedge, bins=n_bins, color='tab:blue', alpha=0.85) + add_hist_percentile_marker(ax, Nedge, fmt=".0f") + ax.set( + title=f"Nedge, N={N}, sum Nedge={int(np.sum(Nedge))}", + xlabel="Nedge (# outgoing edges per row)", + ylabel="nodes", + ) + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[0, 2]) + ax.hist(Sedge, bins=n_bins, color='tab:green', alpha=0.85) + add_hist_percentile_marker(ax, Sedge) + ax.set( + title=f"Sedge distribution, N={N}", + xlabel="Sedge (sum outgoing edge values per row)", + ylabel="nodes", + ) + ax.grid(True, alpha=0.35) + + self._draw_A_matrix_summary_3cols( + fig, gs, 1, A_prune, "A_prune", minW, single_rates, num_exc=num_exc, + ) + + fig.suptitle( + f"{est_label} per-node outgoing edges, minW={minW:g} (row=postsynaptic): " + f"{md['short_name']}", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.94]) + + def _scatter_tp_true_vs_est(self, ax, a_true_tp, a_est_tp, title, color): + x_true = np.asarray(a_true_tp, dtype=np.float64).ravel() + y_init = np.asarray(a_est_tp, dtype=np.float64).ravel() + n = min(x_true.size, y_init.size) + if n <= 0: + ax.set(title=f"{title}\n(no TP points)", xlabel="A_init", ylabel="A_true") + ax.grid(True, alpha=0.35) + return + x_true = x_true[:n] + y_init = y_init[:n] + ax.scatter(y_init, x_true, s=10, marker='.', alpha=0.70, color=color, zorder=4) + self._add_x45_lins(ax, only45=True) + ax.axvline(0.0, linestyle='--', color='0.35', linewidth=0.9, alpha=0.7, zorder=1) + ax.axhline(0.0, linestyle='--', color='0.35', linewidth=0.9, alpha=0.7, zorder=1) + + lo = float(min(np.min(x_true), np.min(y_init))) + hi = float(max(np.max(x_true), np.max(y_init))) + span = max(1e-6, hi - lo) + pad = 0.05 * span + ax.set_xlim(lo - pad, hi + pad) + ax.set_ylim(lo - pad, hi + pad) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.35) + ax.set(title=f"{title} (n={n})", xlabel="A_init", ylabel="A_true") + + def _hist_tp_residuals(self, ax, a_true_tp, a_est_tp, title, color): + x = np.asarray(a_true_tp, dtype=np.float64).ravel() + y = np.asarray(a_est_tp, dtype=np.float64).ravel() + n = min(x.size, y.size) + if n <= 0: + ax.set(title=f"{title}\n(no TP points)", xlabel="A_init - A_true", ylabel="count") + ax.grid(True, alpha=0.35) + return + resid = y[:n] - x[:n] + bins = min(80, max(15, int(np.sqrt(n)))) + ax.hist(resid, bins=bins, color=color, alpha=0.8) + ax.axvline(0.0, color='k', linestyle='--', linewidth=1.0, alpha=0.8) + ax.grid(True, alpha=0.35) + ax.set(title=f"{title} (n={n})", xlabel="A_init - A_true", ylabel="count") + + mae = float(np.mean(np.abs(resid))) + rmse = float(np.sqrt(np.mean(resid ** 2))) + ax.text( + 0.04, 0.96, + f"MAE={mae:.3f}\nRMSE={rmse:.3f}", + transform=ax.transAxes, va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none"), + ) + + def _hist_all_values(self, ax, vals, title, color, tp_vals=None): + v = np.asarray(vals, dtype=np.float64).ravel() + if v.size == 0: + ax.set(title=f"{title}\n(no values)", xlabel="A_init", ylabel="count") + ax.grid(True, alpha=0.35) + return + bins = min(120, max(25, int(np.sqrt(v.size)))) + counts, edges, _ = ax.hist(v, bins=bins, color=color, alpha=0.8, label="all") + if tp_vals is not None: + tp = np.asarray(tp_vals, dtype=np.float64).ravel() + if tp.size > 0: + ax.hist( + tp, bins=edges, histtype='step', color='k', linewidth=1.8, + label=f"TP (n={tp.size})" + ) + ax.axvline(0.0, color='k', linestyle='--', linewidth=1.0, alpha=0.9) + ax.grid(True, alpha=0.35) + ax.set(title=f"{title} (n={v.size})", xlabel="A_init", ylabel="count") + ax.text( + 0.04, 0.96, + f"mean={float(np.mean(v)):.3f}\nstd={float(np.std(v)):.3f}", + transform=ax.transAxes, va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none"), + ) + if tp_vals is not None: + ax.legend(loc="upper right", fontsize=8, framealpha=0.8) + + def _hist2d_values_vs_row(self, ax, A, row_lo, row_hi, title, cmap="bwr"): + A = np.asarray(A, dtype=np.float64) + n_rows = int(A.shape[0]) + n_cols = int(A.shape[1]) + r0 = max(0, int(row_lo)) + r1 = min(n_rows, int(row_hi)) + if r1 <= r0: + ax.set(title=f"{title}\n(empty row range)", xlabel="A_init", ylabel="row neuron index") + ax.grid(True, alpha=0.35) + return + + sub = A[r0:r1, :] + xx = sub.ravel() + yy = np.repeat(np.arange(r0, r1, dtype=np.float64), n_cols) + xbins = min(140, max(40, int(np.sqrt(xx.size)))) + ybins = max(12, min(80, r1 - r0)) + + x_lo = float(np.min(xx)) + x_hi = float(np.max(xx)) + if np.isclose(x_lo, x_hi): + x_hi = x_lo + 1e-9 + x_edges = np.linspace(x_lo, x_hi, xbins + 1) + y_edges = np.linspace(r0 - 0.5, r1 - 0.5, ybins + 1) + + H, _, _ = np.histogram2d(xx, yy, bins=[x_edges, y_edges]) + x_cent = 0.5 * (x_edges[:-1] + x_edges[1:]) + sign_x = np.sign(x_cent) + # Normalize by neuron count (N columns) per user request. + H_signed = (H * sign_x[:, None]) / float(n_cols) + + vmin = float(np.min(H_signed)) + vmax = float(np.max(H_signed)) + if vmin >= 0.0: + vmin = -1e-9 + if vmax <= 0.0: + vmax = 1e-9 + norm = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + + im = ax.pcolormesh(x_edges, y_edges, H_signed.T, shading='auto', cmap=cmap, norm=norm) + ax.axvline(0.0, color='k', linestyle='--', linewidth=1.0, alpha=0.8) + ax.set(title=title, xlabel="A_init", ylabel="row neuron index") + ax.set_ylim(r0 - 0.5, r1 - 0.5) + ax.grid(False) + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + def _hist2d_values_vs_row_mask(self, ax, A, mask, title, cmap="bwr"): + A = np.asarray(A, dtype=np.float64) + M = np.asarray(mask, dtype=bool) + assert A.shape == M.shape, "A/mask shape mismatch" + n_rows = int(A.shape[0]) + n_cols = int(A.shape[1]) + + rr, cc = np.nonzero(M) + if rr.size <= 0: + ax.set(title=f"{title}\n(no values)", xlabel="A_init", ylabel="row neuron index") + ax.grid(True, alpha=0.35) + return + + xx = A[rr, cc] + yy = rr.astype(np.float64, copy=False) + xbins = min(140, max(40, int(np.sqrt(xx.size)))) + ybins = max(12, min(80, n_rows)) + + x_lo = float(np.min(xx)) + x_hi = float(np.max(xx)) + if np.isclose(x_lo, x_hi): + x_hi = x_lo + 1e-9 + x_edges = np.linspace(x_lo, x_hi, xbins + 1) + y_edges = np.linspace(-0.5, n_rows - 0.5, ybins + 1) + + H, _, _ = np.histogram2d(xx, yy, bins=[x_edges, y_edges]) + x_cent = 0.5 * (x_edges[:-1] + x_edges[1:]) + sign_x = np.sign(x_cent) + H_signed = (H * sign_x[:, None]) / float(max(1, n_cols)) + + vmin = float(np.min(H_signed)) + vmax = float(np.max(H_signed)) + if vmin >= 0.0: + vmin = -1e-9 + if vmax <= 0.0: + vmax = 1e-9 + norm = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + + im = ax.pcolormesh(x_edges, y_edges, H_signed.T, shading='auto', cmap=cmap, norm=norm) + ax.axvline(0.0, color='k', linestyle='--', linewidth=1.0, alpha=0.8) + ax.set(title=f"{title} (n={xx.size})", xlabel="A_init", ylabel="row neuron index") + ax.set_ylim(-0.5, n_rows - 0.5) + ax.grid(False) + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + def initA_quality_prismEM(self, fitD, md, minW=0.02, figId=7): + """Three-row diagnostics for A_init vs A_true with diag/exc/inh categories.""" + A_true, N, num_exc, n_diag, n_off, shared_norm = self._get_A_true_info(md) + A_init = np.asarray(fitD["A_init"], dtype=np.float64) + assert A_init.ndim == 2 and A_init.shape == A_true.shape, "A_init must match A_true" + + E_true = np.asarray(md["E_true"]).astype(bool) + if E_true.ndim == 3: + E_true = E_true[0] + assert E_true.shape == A_true.shape, "E_true shape must match A_true" + + diag_mask = np.eye(N, dtype=bool) + off_diag = ~diag_mask + E_t = E_true & off_diag + E_hat = (np.abs(A_init) > float(minW)) & off_diag + TP = E_t & E_hat + FP = (~E_t) & E_hat + FN = E_t & (~E_hat) + + # Exclusive categories: diagonal | excitatory off-diagonal rows | inhibitory off-diagonal rows. + presyn_is_exc = (np.arange(N, dtype=np.int64)[:, None] < int(num_exc)) + cat_exc = off_diag & presyn_is_exc + cat_inh = off_diag & (~presyn_is_exc) + cat_diag = diag_mask + tp_exc = TP & cat_exc + tp_inh = TP & cat_inh + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(14.0, 8.0)) + gs = fig.add_gridspec(3, 4, hspace=0.42, wspace=0.38) + + # Top-left: A_true matrix in the same style as -p c. + ax = fig.add_subplot(gs[0, 0]) + self._draw_A_matrix( + ax, A_true, f"True Dale, N{N}, nEdges={n_diag}+{n_off}", + num_exc=num_exc, norm_map=shared_norm + ) + + ax = fig.add_subplot(gs[0, 1]) + self._scatter_tp_true_vs_est( + ax, A_true[tp_exc], A_init[tp_exc], + "excitatory TP scatter", color="red" + ) + x0, x1 = ax.get_xlim() + ax.set_xlim(min(x0, 0.0), max(x1, 0.0)) + + ax = fig.add_subplot(gs[0, 2]) + self._scatter_tp_true_vs_est( + ax, A_true[tp_inh], A_init[tp_inh], + "inhibitory TP scatter", color="tab:blue" + ) + x0, x1 = ax.get_xlim() + ax.set_xlim(min(x0, 0.0), max(x1, 0.0)) + + ax = fig.add_subplot(gs[0, 3]) + self._scatter_tp_true_vs_est( + ax, A_true[cat_diag], A_init[cat_diag], + "diagonal scatter", color="magenta" + ) + x0, x1 = ax.get_xlim() + ax.set_xlim(min(x0, 0.0), max(x1, 0.0)) + + # Middle-left: A_init matrix in the same style as A_true. + ax = fig.add_subplot(gs[1, 0]) + self._draw_A_matrix( + ax, A_init, "A_init", num_exc=num_exc, norm_map=shared_norm + ) + + ax = fig.add_subplot(gs[1, 1]) + self._hist_all_values( + ax, A_init[cat_exc], + "A_init: excit off-diag", color="red", + tp_vals=A_init[tp_exc] + ) + + ax = fig.add_subplot(gs[1, 2]) + self._hist_all_values( + ax, A_init[cat_inh], + "A_init: inhib off-diag", color="tab:blue", + tp_vals=A_init[tp_inh] + ) + + ax = fig.add_subplot(gs[1, 3]) + self._hist_all_values( + ax, A_init[cat_diag], + "A_init: all diag", color="magenta", + tp_vals=None + ) + + # Third row: A_init value vs presyn neuron index. + ax = fig.add_subplot(gs[2, 0]) + self._hist2d_values_vs_row( + ax, A_init, 0, N, "A_init value vs row idx (all)", cmap="bwr" + ) + + ax = fig.add_subplot(gs[2, 1]) + self._hist2d_values_vs_row_mask( + ax, A_init, cat_exc, "excit off-diag", cmap="bwr" + ) + + ax = fig.add_subplot(gs[2, 2]) + self._hist2d_values_vs_row_mask( + ax, A_init, cat_inh, "inhib off-diag", cmap="bwr" + ) + + ax = fig.add_subplot(gs[2, 3]) + self._hist2d_values_vs_row_mask( + ax, A_init, cat_diag, "diagonal", cmap="bwr" + ) + + tp = int(TP.sum()) + fp = int(FP.sum()) + fn = int(FN.sum()) + tp_e = int(tp_exc.sum()) + tp_i = int(tp_inh.sum()) + n_d = int(cat_diag.sum()) + fig.suptitle( + f"A_init TP quality vs A_true, minW={float(minW):g}: {md['short_name']}\n" + f"off-diag TP={tp} (exc={tp_e}, inh={tp_i}), FP={fp}, FN={fn}; diag n={n_d}", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.93]) + + def _add_x45_lins(self, ax, only45=False): + lims = [ + np.min([ax.get_xlim(), ax.get_ylim()]), + np.max([ax.get_xlim(), ax.get_ylim()]), + ] + ax.plot(lims, lims, '--', color='k', linewidth=0.8) + if only45: + return + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + + def _corrcoef_safe(self, x, y): + x = np.asarray(x) + y = np.asarray(y) + if x.size == 0 or y.size == 0: + return np.nan + if np.std(x) == 0 or np.std(y) == 0: + return np.nan + return float(np.corrcoef(x, y)[0, 1]) + + def _find_bimodal_divider(self, xV): + x = np.asarray(xV, dtype=float).ravel() + assert x.size >= 2, f"Need >=2 samples for bimodal split, got {x.size}" + assert np.isfinite(x).all(), "xV contains non-finite values" + assert np.std(x) > 0.0, "xV must have non-zero variance for bimodal split" + + q1, q3 = np.quantile(x, [0.25, 0.75]) + c1, c2 = float(q1), float(q3) + if c1 == c2: + c1 = float(np.min(x)) + c2 = float(np.max(x)) + assert c1 != c2, "Failed to initialize two distinct mode centers" + + for _ in range(32): + d1 = np.abs(x - c1) + d2 = np.abs(x - c2) + left = d1 <= d2 + n_left = int(left.sum()) + n_right = int((~left).sum()) + assert n_left > 0 and n_right > 0, "Bimodal split produced empty cluster" + c1_new = float(np.mean(x[left])) + c2_new = float(np.mean(x[~left])) + if abs(c1_new - c1) < 1e-10 and abs(c2_new - c2) < 1e-10: + c1, c2 = c1_new, c2_new + break + c1, c2 = c1_new, c2_new + + if c1 > c2: + c1, c2 = c2, c1 + divide = 0.5 * (c1 + c2) + assert np.isfinite(divide), "Computed non-finite bimodal divider" + return float(divide) + + def _plot_corr_A_regions(self, ax, xV, yV, minW, title, xlab, ylab, s=6, alpha=0.5, color=None): + x = np.asarray(xV) + y = np.asarray(yV) + ax.scatter(x, y, s=s, alpha=alpha, color=color) + self._add_x45_lins(ax, only45=True) + ax.axvline(-minW, color="red", linestyle="--", linewidth=1) + ax.axvline(minW, color="red", linestyle="--", linewidth=1) + ax.set(title=title, xlabel=xlab, ylabel=ylab) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.4) + + mask_left = x < -minW + mask_mid = (x >= -minW) & (x <= minW) + mask_right = x > minW + r_left = self._corrcoef_safe(x[mask_left], y[mask_left]) + r_mid = self._corrcoef_safe(x[mask_mid], y[mask_mid]) + r_right = self._corrcoef_safe(x[mask_right], y[mask_right]) + n_left = int(mask_left.sum()) + n_mid = int(mask_mid.sum()) + n_right = int(mask_right.sum()) + + ax.text(0.20, 0.05, f"rL={r_left:.3f}\nnL={n_left}", transform=ax.transAxes) + ax.text(0.50, 0.40, f"rM={r_mid:.3f}\nnM={n_mid}", transform=ax.transAxes) + ax.text(0.65, 0.60, f"rR={r_right:.3f}\nnR={n_right}", transform=ax.transAxes) + + def _plot_corr_B_divisor(self, ax, xV, yV, title, xlab, ylab, s=6, alpha=0.5, color=None): + x = np.asarray(xV) + y = np.asarray(yV) + assert x.shape == y.shape, f"x and y shape mismatch: {x.shape} vs {y.shape}" + divide = self._find_bimodal_divider(x) + + ax.scatter(x, y, s=s, alpha=alpha, color=color) + self._add_x45_lins(ax, only45=True) + ax.axvline(divide, color="red", linestyle="--", linewidth=1) + ax.set(title=title, xlabel=xlab, ylabel=ylab) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.4) + + mask_left = x < divide + mask_right = x >= divide + r_left = self._corrcoef_safe(x[mask_left], y[mask_left]) + r_right = self._corrcoef_safe(x[mask_right], y[mask_right]) + n_left = int(mask_left.sum()) + n_right = int(mask_right.sum()) + ax.text(0.05, 0.45, f"rL={r_left:.3f}\nnL={n_left}", transform=ax.transAxes, va="top") + ax.text(0.65, 0.55, f"rR={r_right:.3f}\nnR={n_right}", transform=ax.transAxes, va="top") + + def eval_ABcorr_prismEM(self, fitD, md, figId=6): + A_hat = np.asarray(fitD["A_hat"]) + B_hat = np.asarray(fitD["B_hat"]) + A_true = np.asarray(md["A_true"]) + B_true = np.asarray(md["B_true"]) + minW = float(self.args.minW) + + assert A_true.ndim == 2, f"A_true must be 2D, got shape={A_true.shape}" + assert A_hat.ndim == 2, f"A_hat must be 2D, got shape={A_hat.shape}" + + if B_hat.ndim == 1: + B_hat = B_hat[None, :] + if B_true.ndim == 1: + B_true = B_true[None, :] + n_states = min(B_hat.shape[0], B_true.shape[0]) + assert n_states >= 1, "Need at least one B-state for correlation plot" + + figId = self.smart_append(figId) + ncol = 1 + n_states + fig_w = max(10.0, 3.0 * ncol) + fig = self.plt.figure(figId, facecolor='white', figsize=(fig_w, 3.8)) + + ax = self.plt.subplot(1, ncol, 1) + self._plot_corr_A_regions( + ax, A_true.ravel(), A_hat.ravel(), minW, + "A fit, non-zero ,", "A_true", "A_hat", s=6, alpha=0.4, color="green" + ) + + for m in range(n_states): + ax = self.plt.subplot(1, ncol, 2 + m) + self._plot_corr_B_divisor( + ax, B_true[m], B_hat[m], + f"B fit, state {m}", "B_true", "B_hat", + s=8, alpha=0.5, color="blue" + ) + + fig.suptitle(f"2D correlations, minW={minW:g}: {md['short_name']}", fontsize=12) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.92]) + + def state_seq_prismEM(self, fitD, md, figId=5, time_range_sec=None): + """4-row state-sequence canvas in the style of prism_Estep_eval -p b.""" + trainMD = md["train"] + dt = float(trainMD["time_step_sec"]) + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + + t0_req, t1_req = [float(x) for x in time_range_sec] + if t1_req < t0_req: + t0_req, t1_req = t1_req, t0_req + b0 = max(t0_bin, int(np.floor(t0_req / dt))) + b1 = min(t1_bin, int(np.floor(t1_req / dt))) + if b1 <= b0: + raise ValueError("Requested --time_range_sec leaves no bins in fitted window") + + i0 = b0 - t0_bin + i1 = b1 - t0_bin + p0 = i0 + p1 = i1 + + S_hat = np.asarray(fitD["S_hat"])[i0:i1 + 1] + C_hat = np.asarray(fitD["c_hat"])[i0:i1 + 1] + S_hat_CL = np.asarray(fitD["S_hat_CL"])[i0:i1 + 1] + + S_true = np.asarray(md["S_true"])[b0:b1 + 1] + C_true = np.asarray(md["C_true"])[b0:b1 + 1] + + nll_t = np.asarray(md["eval_f"]["loss_nll_time"])[p0:p1] + l2_t = np.asarray(md["eval_f"]["loss_l2_time"])[p0:p1] + t_bins = np.arange(b0, b1 + 1, dtype=np.float64) * dt + t_pairs = np.arange(b0, b1, dtype=np.float64) * dt + x0, x1 = float(t_bins[0]), float(t_bins[-1]) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 9)) + # Add an explicit spacer row between diagnostics (top row) and Fit row + # so top-row x labels do not collide with Fit legends. + gs = fig.add_gridspec( + 5, 3, + height_ratios=[0.95, 0.18, 1.0, 1.0, 0.55], + hspace=0.75, + wspace=0.35, + ) + + # Row 2: Fit + ax = fig.add_subplot(gs[2, :]) + ax.plot(t_bins, S_hat, color="k", linewidth=2.5, label="S_hat") + ax.plot(t_bins, S_true, color="lime", linestyle="--", linewidth=1.5, label="S_true") + lo = np.clip(S_hat - S_hat_CL, 0.0, float(C_hat.shape[1] - 1)) + hi = np.clip(S_hat + S_hat_CL, 0.0, float(C_hat.shape[1] - 1)) + ax.fill_between(t_bins, lo, hi, color="gray", alpha=0.3, label="S_hat_CL") + ax.set(title=f"Fit (acc={float(md['eval_f']['acc']):.3f})", xlabel="time (s)", ylabel="state") + ax.set_xlim(x0, x1) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax2 = ax.twinx() + for m in range(C_hat.shape[1]): + ax2.plot(t_bins, C_hat[:, m], linewidth=0.8, alpha=0.85, label=f"C_hat[{m}]") + ax2.set_ylabel("C_hat") + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + # Row 3: Truth + ax = fig.add_subplot(gs[3, :]) + ax.plot(t_bins, S_true, color="lime", linestyle="--",linewidth=1.5, label="S_true") + ax.set(title="Truth", xlabel="time (s)", ylabel="state") + ax.set_xlim(x0, x1) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax2 = ax.twinx() + for m in range(C_true.shape[1]): + ax2.plot(t_bins, C_true[:, m], linewidth=0.8, alpha=0.85, label=f"C_true[{m}]") + ax2.set_ylabel("C_true") + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + # Row 1: diagnostics panels (occupancy, state accuracy, confidence) + srec = md["states_recovery_eval"] + acc_avg = float(srec["avg_acc"]) + state_acc_cl = np.asarray(srec["state_acc_cl"], dtype=np.float64) + acc_ps = state_acc_cl[:, 0] + cl_ps = state_acc_cl[:, 1] + + ax = fig.add_subplot(gs[0, 0]) + occ = C_hat.mean(axis=0) + x = np.arange(C_hat.shape[1], dtype=np.int64) + ax.bar(x, occ, color="tab:blue", alpha=0.65) + ax.set(title="Mean occupancy", xlabel="state", ylabel="mean c") + ax.set_ylim(0.0, 1.0) + ax.set_xticks(x) + ax.grid(True, alpha=0.3) + ax.set_box_aspect(0.9) + + ax = fig.add_subplot(gs[0, 1]) + x = np.arange(acc_ps.size, dtype=np.int64) + ax.bar(x, acc_ps, color="tab:green", alpha=0.7) + ax.errorbar(x, acc_ps, yerr=cl_ps, fmt='o', color='k', capsize=5, markersize=4) + ax.axhline(1.0, color="green", linestyle="--", linewidth=1) + ymin = min(0.5, float(np.nanmin(acc_ps - cl_ps)) - 0.05) + ax.set_ylim(ymin, None) + ax.set(title=f"State accuracy, avr={acc_avg:.3f}", xlabel="state", ylabel="accuracy") + ax.set_xticks(x) + ax.grid(True, alpha=0.3, axis="y") + ylo, yhi = ax.get_ylim() + y_txt = ylo + 0.80 * (yhi - ylo) + for i, v in enumerate(acc_ps): + ax.text(i + 0.3, y_txt, f"{v:.2f}", ha="center", va="center", fontsize=9) + ax.set_box_aspect(0.9) + + ax = fig.add_subplot(gs[0, 2]) + cl = np.asarray(S_hat_CL, dtype=np.float64) + ax.hist(cl, bins=30, color="tab:purple", alpha=0.7) + ax.axvline(np.mean(cl), color="k", linestyle="--", linewidth=1.0) + ax.set(title=f"Confidence (acc={acc_avg:.3f})", xlabel="S_hat_CL", ylabel="count") + ax.grid(True, alpha=0.3) + ax.set_box_aspect(0.9) + + # Row 4: Loss(time) + ax = fig.add_subplot(gs[4, :]) + ax.plot(t_pairs, nll_t, color="tab:blue", linewidth=0.9, label="nll") + ax.plot(t_pairs, l2_t, color="tab:orange", linewidth=0.9, label="l2") + ax.set_yscale("log") + ax.set(title="Loss (time)", xlabel="time (s)", ylabel="value") + ax.set_xlim(x0, x1) + ax.grid(True, alpha=0.35) + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), ncol=2, fontsize=8) + ax2 = ax.twinx() + ax2.plot(t_bins, S_true, color="k", linewidth=0.8, alpha=0.6, label="S_true") + ax2.set_ylabel("state") + ax2.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), fontsize=8) + + fig.suptitle( + f"Dataset: {md['short_name']} | λ₂={trainMD['lambda2']}, lr={trainMD['lr_estep']}, " + f"pgd_iter={trainMD['pgd_iter']}, decode_dwell={srec['decode_dwell_sec']}s", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.965]) + + def _state_seq_time_slice(self, trainMD, time_range_sec): + """Return index slice [i0:i1+1] and time axis for --time_range_sec.""" + dt = float(trainMD["time_step_sec"]) + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + t0_req, t1_req = [float(x) for x in time_range_sec] + if t1_req < t0_req: + t0_req, t1_req = t1_req, t0_req + b0 = max(t0_bin, int(np.floor(t0_req / dt))) + b1 = min(t1_bin, int(np.floor(t1_req / dt))) + if b1 <= b0: + raise ValueError("Requested --time_range_sec leaves no bins in fitted window") + i0 = b0 - t0_bin + i1 = b1 - t0_bin + t_bins = np.arange(b0, b1 + 1, dtype=np.float64) * dt + return i0, i1, t_bins + + def _draw_state_c_panel( + self, ax, t_bins, S, S_cl, C, title, + s_label, cl_label, c_ylabel, show_xlabel=True, + ): + n_states = C.shape[1] + ax.plot(t_bins, S, color="k", linewidth=2.5, label=s_label, zorder=4) + lo = np.clip(S - S_cl, 0.0, float(n_states - 1)) + hi = np.clip(S + S_cl, 0.0, float(n_states - 1)) + ax.fill_between( + t_bins, lo, hi, color="gray", alpha=0.3, label=cl_label, + zorder=3, + ) + ax.set(title=title, ylabel="state") + if show_xlabel: + ax.set_xlabel("time (s)") + ax.set_xlim(float(t_bins[0]), float(t_bins[-1])) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + c_colors = ["tab:blue", "tab:orange", "tab:green", "tab:red", "tab:purple"] + ax2 = ax.twinx() + ax2.set_zorder(1) + for m in range(n_states): + ax2.plot( + t_bins, C[:, m], linewidth=0.8, alpha=0.85, + color=c_colors[m % len(c_colors)], + label=f"{c_ylabel}[{m}]", + zorder=2, + ) + ax.set_zorder(ax2.get_zorder() + 1) + ax.patch.set_visible(False) + + ax2.set_ylabel(c_ylabel) + ax2.set_ylim(0.0, 1.0) + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + def state_seq_fitonly_prismEM( + self, fitD, md, rebD, bio_plot_md, figId=9, time_range_sec=None, + ): + """4-row canvas: init/fit states + bioExp sum-rate + heatmap (no truth).""" + if time_range_sec is None: + time_range_sec = md["train"]["time_range_sec"] + trainMD = md["train"] + i0, i1, t_bins = self._state_seq_time_slice(trainMD, time_range_sec) + + S_init = np.asarray(fitD["S_init"], dtype=np.float64)[i0:i1 + 1] + C_init = np.asarray(fitD["c_init"], dtype=np.float64)[i0:i1 + 1] + S_init_cl = 1.0 - np.max(C_init, axis=1) + + S_hat = np.asarray(fitD["S_hat"], dtype=np.float64)[i0:i1 + 1] + C_hat = np.asarray(fitD["c_hat"], dtype=np.float64)[i0:i1 + 1] + S_hat_cl = np.asarray(fitD["S_hat_CL"], dtype=np.float64)[i0:i1 + 1] + + clip = clip_rebD_time(rebD, time_range_sec) + dataset_name = bio_plot_md["short_name"] + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(12, 14)) + gs = fig.add_gridspec(5, 1, height_ratios=[0.14, 0.14, 0.12, 0.50, 0.01]) + + ax = fig.add_subplot(gs[0, 0]) + self._draw_state_c_panel( + ax, t_bins, S_init, S_init_cl, C_init, "Initialization", + s_label="S_init", cl_label="S_init_CL", c_ylabel="C_init", + show_xlabel=False, + ) + + ax = fig.add_subplot(gs[1, 0]) + self._draw_state_c_panel( + ax, t_bins, S_hat, S_hat_cl, C_hat, "Fit", + s_label="S_hat", cl_label="S_hat_CL", c_ylabel="C_hat", + show_xlabel=False, + ) + + ax = fig.add_subplot(gs[2, 0]) + PlotterBioExp.plot_bioexp_sum_rate(self, ax, rebD, clip) + + ax = fig.add_subplot(gs[3, 0]) + PlotterBioExp.plot_bioexp_heatmap(self, fig, ax, rebD, dataset_name, clip) + + t0s, t1s = float(time_range_sec[0]), float(time_range_sec[1]) + fig.suptitle( + f"State sequences: {md['short_name']} T=[{t0s:g}, {t1s:g}] s", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.94]) + + def _bioexp_Ahat_topology_ax( + self, ax, loc_x, loc_y, A_hat, minW, sign=None, neuron_type=None, + node_size=16, row_select=None, und_color="salmon", + ): + """Neuron layout with A_hat edges; sign is 'pos', 'neg', or None for both.""" + ax.set_facecolor("white") + if neuron_type is None: + ax.scatter( + loc_x, loc_y, s=node_size, marker="o", facecolors="white", + edgecolors="k", linewidths=0.4, alpha=0.95, zorder=3, + ) + else: + neuron_type = np.asarray(neuron_type) + node_color = np.full(loc_x.shape, und_color, dtype=object) + node_color[neuron_type > 0] = "magenta" + node_color[neuron_type < 0] = "#39FF14" + ax.scatter( + loc_x, loc_y, s=node_size, marker="o", c=node_color, + edgecolors="k", linewidths=0.4, alpha=0.95, zorder=3, + ) + edge_mask = np.abs(A_hat) > float(minW) + np.fill_diagonal(edge_mask, False) + if sign == "pos": + edge_mask &= A_hat > 0 + edge_color = "red" + elif sign == "neg": + edge_mask &= A_hat < 0 + edge_color = "blue" + post_idx, pre_idx = np.where(edge_mask) + n_total = int(post_idx.size) + if row_select is not None: + row_select = np.asarray(row_select, dtype=bool) + draw_mask = row_select[post_idx] + post_draw = post_idx[draw_mask] + pre_draw = pre_idx[draw_mask] + else: + post_draw = post_idx + pre_draw = pre_idx + if sign is None: + for i_post, j_pre in zip(post_draw, pre_draw): + color = "red" if A_hat[i_post, j_pre] > 0 else "blue" + ax.plot( + [loc_x[j_pre], loc_x[i_post]], + [loc_y[j_pre], loc_y[i_post]], + color=color, alpha=0.7, linewidth=0.7, zorder=2, + ) + else: + for i_post, j_pre in zip(post_draw, pre_draw): + ax.plot( + [loc_x[j_pre], loc_x[i_post]], + [loc_y[j_pre], loc_y[i_post]], + color=edge_color, alpha=0.7, linewidth=0.7, zorder=2, + ) + ax.set_aspect("equal", adjustable="box") + ax.set_xlabel("loc_x (um)", fontsize=11) + ax.set_ylabel("loc_y (um)", fontsize=11) + ax.tick_params(labelsize=9) + ax.grid(True, alpha=0.35) + return n_total, int(post_draw.size) + + def neuron_spatial_Ahat_edges(self, fitD, bioD, bioMD, minW=0.02, figId=10): + """MEA layout plus fitted A_hat edges for bioExp data.""" + loc_x = PlotterBioExp._metrics_column(self, bioD, bioMD, "loc_x") + loc_y = PlotterBioExp._metrics_column(self, bioD, bioMD, "loc_y") + rate = PlotterBioExp._metrics_column(self, bioD, bioMD, "firing_rate") + A_hat = np.asarray(fitD["A_hat"], dtype=np.float64) + n_neur = loc_x.shape[0] + assert A_hat.shape == (n_neur, n_neur), ( + f"A_hat shape {A_hat.shape} does not match {n_neur} neuron locations" + ) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(15, 10)) + gs = fig.add_gridspec(1, 3, width_ratios=[1.0, 0.06, 1.0], wspace=0.28) + + ax = fig.add_subplot(gs[0, 0]) + ax.set_facecolor("white") + sc = ax.scatter( + loc_x, loc_y, c=rate, s=12, marker="s", + cmap="Blues", vmin=0.0, vmax=float(np.max(rate)), + linewidths=0.15, edgecolors="k", alpha=0.95, + ) + ax.set_aspect("equal", adjustable="box") + ax.set_xlabel("loc_x (um)", fontsize=11) + ax.set_ylabel("loc_y (um)", fontsize=11) + ax.tick_params(labelsize=9) + ax.grid(True, alpha=0.35) + ax.set_title(f"dataset: {bioMD['short_name']} n={n_neur} neurons", fontsize=11, pad=8) + + cax = fig.add_subplot(gs[0, 1]) + cb = fig.colorbar(sc, cax=cax, orientation="vertical") + cb.set_label("Firing Rate [Hz]", fontsize=11) + cb.ax.tick_params(labelsize=9) + + ax = fig.add_subplot(gs[0, 2]) + n_edges, _ = self._bioexp_Ahat_topology_ax(ax, loc_x, loc_y, A_hat, minW, sign=None) + ax.set_title(f"A_hat edges, |A|>{float(minW):g}, n={n_edges}", fontsize=11, pad=8) + ax.plot([], [], color="red", alpha=0.7, label="A_hat > 0") + ax.plot([], [], color="blue", alpha=0.7, label="A_hat < 0") + ax.legend(loc="best", fontsize=9) + + fig.suptitle(f"Neuron spatial A_hat topology: {bioMD['short_name']}", fontsize=12) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.94]) + + def neuron_spatial_Ahat_edges_split( + self, fitD, bioD, bioMD, neuron_type, neuron_Sedge, + minW=0.02, maxNeurons=30, figId=11, + ): + """Two panels: positive-only and negative-only A_prune edges (bioExp).""" + loc_x = PlotterBioExp._metrics_column(self, bioD, bioMD, "loc_x") + loc_y = PlotterBioExp._metrics_column(self, bioD, bioMD, "loc_y") + A_prune = self._A_prune_from_fitD(fitD) + n_neur = loc_x.shape[0] + assert A_prune.shape == (n_neur, n_neur), ( + f"A_prune shape {A_prune.shape} does not match {n_neur} neuron locations" + ) + neuron_type = np.asarray(neuron_type, dtype=np.int8) + assert neuron_type.shape[0] == n_neur, "neuron_type length must match neuron locations" + neuron_Sedge = np.asarray(neuron_Sedge, dtype=np.float64) + assert neuron_Sedge.shape[0] == n_neur, "neuron_Sedge length must match neuron locations" + n_exc = int(np.sum(neuron_type > 0)) + n_inh = int(np.sum(neuron_type < 0)) + n_und = int(np.sum(neuron_type == 0)) + + def ranked_row_select(type_value): + idx = np.where(neuron_type == int(type_value))[0] + idx = idx[np.argsort(neuron_Sedge[idx])] + n_side = max(1, int(maxNeurons) // 2) + if idx.size <= int(maxNeurons): + sel = idx + else: + sel = np.unique(np.concatenate([idx[:n_side], idx[-n_side:]])) + row_select = np.zeros((n_neur,), dtype=bool) + row_select[sel] = True + return row_select, int(sel.size) + + exc_rows, n_exc_show = ranked_row_select(1) + inh_rows, n_inh_show = ranked_row_select(-1) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(14, 4.8)) + gs = fig.add_gridspec( + 1, 2, left=0.06, right=0.99, bottom=0.13, top=0.84, wspace=0.18 + ) + + ax = fig.add_subplot(gs[0, 0]) + n_pos, n_pos_show = self._bioexp_Ahat_topology_ax( + ax, loc_x, loc_y, A_prune, 0.0, sign="pos", + neuron_type=neuron_type, node_size=32, row_select=exc_rows, + und_color="yellow", + ) + ax.set_title( + f"A_prune>0, shown edges {n_pos_show}/{n_pos}, neurons {n_exc_show}/{n_exc}", + fontsize=11, pad=8, + ) + ax.scatter([], [], s=32, marker="o", c="magenta", edgecolors="k", label=f"exc={n_exc}") + ax.scatter([], [], s=32, marker="o", c="#39FF14", edgecolors="k", label=f"inh={n_inh}") + ax.scatter([], [], s=32, marker="o", c="yellow", edgecolors="k", label=f"und={n_und}") + ax.legend(loc="best", fontsize=9) + + ax = fig.add_subplot(gs[0, 1]) + n_neg, n_neg_show = self._bioexp_Ahat_topology_ax( + ax, loc_x, loc_y, A_prune, 0.0, sign="neg", + neuron_type=neuron_type, node_size=32, row_select=inh_rows, + und_color="yellow", + ) + ax.set_title( + f"A_prune<0, shown edges {n_neg_show}/{n_neg}, neurons {n_inh_show}/{n_inh}", + fontsize=11, pad=8, + ) + ax.scatter([], [], s=32, marker="o", c="magenta", edgecolors="k", label=f"exc={n_exc}") + ax.scatter([], [], s=32, marker="o", c="#39FF14", edgecolors="k", label=f"inh={n_inh}") + ax.scatter([], [], s=32, marker="o", c="yellow", edgecolors="k", label=f"und={n_und}") + ax.legend(loc="best", fontsize=9) + + data_name = bioMD.get("dataName", bioMD["short_name"]) + fig.suptitle( + f"Neuron spatial A_prune topology (signed): {data_name} " + f"max neurons per panel={int(maxNeurons)}", + fontsize=12, y=0.97, + ) diff --git a/causal_net/nonStation_ver3b_states_expr/Readme b/causal_net/nonStation_ver3b_states_expr/Readme new file mode 100644 index 00000000..063ca9b2 --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/Readme @@ -0,0 +1,28 @@ +Then this is how I read experimental data from Mandar + +Use this version: +https://github.com/BouchardLab/pyuoi/tree/uoi-var/causal_net/nonStation_ver3_states_expr + +INPUT: +expPath=/global/cfs/cdirs/m2043/causal_inference/Canine_Organoids_PVS/ +sesN=Analysis/CanineOrganoids_03112026_PVS/260324/M08092/Network/000023/well000; +shortN=Canine_260324_r23_w0_1hz + +The shortN is my compact data idetniffier: day, run#, well#, and what is the cut-off frequency (I use 1Hz + +OUTPUT: +this is where I want to uncompressed data to land: +basePath=/pscratch/sd/b/balewski/2026_causalNet_exp_ver3 +dataPath=${basePath}/spikesData/ + +This is unpackimg command: +./prep_bioexp3b.py --sessionName $sesN --dataPath $dataPath --expPath $expPath --freqRange 1 50 --shortName $shortN + +This is how you can inspect some aspects of raw data: + ./view_bioexp.py --dataPath $dataPath --dataName Canine_260324_r23_w0_1hz -p a b -T 0 3550 + +login09:/global/cfs/cdirs/m2043/causal_inference/Canine_Organoids_PVS/Analysis/CanineOrganoids_03112026_PVS> ls */*/*/* +260324/M08092/Network/000021: +well000/ well001/ well002/ well003/ well004/ well005/ + + diff --git a/causal_net/nonStation_ver3b_states_expr/UtilBioExp.py b/causal_net/nonStation_ver3b_states_expr/UtilBioExp.py new file mode 100644 index 00000000..50019e00 --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/UtilBioExp.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +""" +Utility functions for biological experiment data processing. + +This module provides specialized functions for analyzing biological neural data, +particularly focused on identifying and processing clusters of neural activity. +Main functionality includes: +- Cluster detection in time series data using connectivity analysis +- Threshold-based cluster validation and masking +- Signal processing utilities for experimental neural recordings + +Used primarily in conjunction with experimental data preprocessing and +analysis pipelines for biological neural network studies. +""" + +import numpy as np + + +def detect_spike_bursts(spikeD, spikeMD, time_rebin2, burst_freq_thres): + """Rebin spikes in time; return rate panels for bioExp freq-vs-time plots.""" + rateThr2 = float(burst_freq_thres) + spikeYield = spikeD["spikes"] + time_step = float(spikeMD["time_step_sec"]) + tReb2 = int(time_rebin2) + assert tReb2 < 101 + + ntime, nchan = spikeYield.shape + if ntime % tReb2 != 0: + ntime_c = ntime - (ntime % tReb2) + spikeYield = spikeYield[:ntime_c] + + spikeYieldR = np.sum(spikeYield.reshape(-1, tReb2, nchan), axis=1) + time_step2 = time_step * tReb2 + + rate2D = spikeYieldR / time_step2 + mask2D = rate2D > rateThr2 + highChan = np.sum(mask2D, axis=1) + + rebD = { + "time_step2": time_step2, + "rate_thres2": rateThr2, + "rate2D": rate2D, + "mask2D": mask2D, + "highChanCnt": highChan, + "rate1D": np.sum(rate2D, axis=1), + } + ntime = spikeYieldR.shape[0] + rebD["timeV"] = np.linspace(0, (ntime - 1) * time_step2, ntime) + + print("nchan=%d thr=%.1f Hz" % (nchan, rateThr2)) + return rebD + + +def clip_rebD_time(rebD, time_range_lr): + """Clip rebD arrays to a time window in seconds.""" + time_step = rebD["time_step2"] + tL, tR = [float(x) for x in time_range_lr] + itL, itR = (np.asarray(time_range_lr, dtype=np.float64) / time_step).astype(int) + itR = min(itR, rebD["rate2D"].shape[0]) + if itR <= itL: + raise ValueError("time_range_lr leaves no rebinned bins") + + rate2D = rebD["rate2D"][itL:itR] + timeV = rebD["timeV"][itL:itR] + rate1D = rebD["rate1D"][itL:itR] + _, nchan = rate2D.shape + Tbin = float(timeV[1] - timeV[0]) + return { + "tL": tL, "tR": tR, "itL": itL, "itR": itR, + "time_step": time_step, "Tbin": Tbin, "nchan": nchan, + "rate2D": rate2D, "timeV": timeV, "rate1D": rate1D, + } + + +if __name__=="__main__": + print("Utility module for biological experiment spike-rate rebinning.") diff --git a/causal_net/nonStation_ver3b_states_expr/Util_PrismEM.py b/causal_net/nonStation_ver3b_states_expr/Util_PrismEM.py new file mode 100644 index 00000000..16b1fa08 --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/Util_PrismEM.py @@ -0,0 +1,309 @@ +#!/usr/bin/env python3 +""" +Utilities for prism EM initialization. +""" + +import numpy as np +import time + + +def _normalize_init_opt(opt): + s = str(opt).strip().lower() + if s not in ("data", "rand"): + raise ValueError(f"Unsupported init option: {opt}") + return s + + +def _rate_thresholds_min_var(rate_t, n_state): + """Find up to (n_state-1) thresholds minimizing within-group variance.""" + x = np.asarray(rate_t, dtype=np.float64).ravel() + if x.size == 0 or n_state <= 1: + return [] + + uniq, counts = np.unique(x, return_counts=True) + u = uniq.size + if u <= 1: + return [] + + k = min(int(n_state), u) + w = counts.astype(np.float64) + wx = w * uniq + wx2 = wx * uniq + + c_w = np.concatenate(([0.0], np.cumsum(w))) + c_s = np.concatenate(([0.0], np.cumsum(wx))) + c_q = np.concatenate(([0.0], np.cumsum(wx2))) + + # dp[c, j]: minimal weighted SSE for first (j+1) unique values into c groups. + dp = np.full((k + 1, u), np.inf, dtype=np.float64) + ptr = np.full((k + 1, u), -1, dtype=np.int32) + + j_all = np.arange(u, dtype=np.int32) + 1 + w0 = c_w[j_all] - c_w[0] + s0 = c_s[j_all] - c_s[0] + q0 = c_q[j_all] - c_q[0] + dp[1, :] = q0 - (s0 * s0) / np.maximum(w0, 1e-12) + ptr[1, :] = 0 + + for c in range(2, k + 1): + for j in range(c - 1, u): + i = np.arange(c - 1, j + 1, dtype=np.int32) + jj = j + 1 + + ww = c_w[jj] - c_w[i] + ss = c_s[jj] - c_s[i] + qq = c_q[jj] - c_q[i] + seg = qq - (ss * ss) / np.maximum(ww, 1e-12) + cand = dp[c - 1, i - 1] + seg + + ib = int(np.argmin(cand)) + dp[c, j] = float(cand[ib]) + ptr[c, j] = int(i[ib]) + + starts = [] + j = u - 1 + for c in range(k, 1, -1): + i = int(ptr[c, j]) + starts.append(i) + j = i - 1 + starts.reverse() + + thres = [] + for i in starts: + l = uniq[i - 1] + r = uniq[i] + thres.append(float(0.5 * (l + r))) + return thres + + +def init_states_vs_time(spikes, dt, args): + """Prepare initial c_hat probabilities and state labels from spikes. + + Returns: + c_init: float32 array (T, M) + S_init: int64 array (T,) + init_meta: None or {"rate_thres": [..], "rate_bin_sec": float, "rate_bin_bins": int} + freq_h1d: float32 array used for state classification + """ + spikes = np.asarray(spikes) + if spikes.ndim != 2: + raise ValueError("spikes must be 2D array (T, N)") + + T_full = spikes.shape[0] + n_state = int(args.num_states) + opt = _normalize_init_opt(args.init_states) + + if opt == "rand": + c_init = np.full((T_full, n_state), 1.0 / n_state, dtype=np.float32) + s_init = np.full((T_full,), -1, dtype=np.int64) + freq_h1d = np.zeros((0,), dtype=np.float32) + return c_init, s_init, None, freq_h1d + if opt != "data": + raise ValueError(f"Unsupported --init_states option: {opt}") + + dwell_sec = float(args.decode_dwell_sec) + dwell_bins = max(1, int(np.ceil(dwell_sec / float(dt)))) + + # Aggregate spikes over coarse dwell bins, then compute mean firing rate. + T = T_full + n_blk = (T + dwell_bins - 1) // dwell_bins + rate_blk = np.zeros((n_blk,), dtype=np.float64) + blk_sizes = np.zeros((n_blk,), dtype=np.int64) + for b in range(n_blk): + i0 = b * dwell_bins + i1 = min(T, i0 + dwell_bins) + blk = spikes[i0:i1] + nb = i1 - i0 + blk_sizes[b] = nb + rate_per_neuron = blk.sum(axis=0).astype(np.float64) / (float(nb) * float(dt)) + rate_blk[b] = rate_per_neuron.mean() + + thres = _rate_thresholds_min_var(rate_blk, n_state) + th = np.asarray(thres, dtype=np.float64) + s_blk = np.searchsorted(th, rate_blk, side="right").astype(np.int64) + s_blk = np.clip(s_blk, 0, max(0, n_state - 1)) + + c_blk, s_blk_out = init_probs_from_states(s_blk, n_state) + + # Map coarse-bin initialization back to original binning. + c_init = np.zeros((T_full, n_state), dtype=np.float32) + s_init = np.zeros((T_full,), dtype=np.int64) + for b in range(n_blk): + i0 = b * dwell_bins + i1 = i0 + int(blk_sizes[b]) + c_init[i0:i1] = c_blk[b] + s_init[i0:i1] = s_blk_out[b] + + freq_h1d = rate_blk.astype(np.float32, copy=False) + init_meta = { + "rate_thres": [float(x) for x in thres], + "rate_bin_sec": float(dwell_bins * float(dt)), + "rate_bin_bins": int(dwell_bins), + } + return c_init, s_init, init_meta, freq_h1d + + +def init_selfspikingB(spikes, dt, args): + """Initialize common per-state B from observed mean firing rates. + + For each neuron n: + rate_hz[n] = mean_t spikes[t, n] / dt + B[:, n] = log(max(rate_hz[n], eps)) + """ + spikes = np.asarray(spikes) + if spikes.ndim != 2: + raise ValueError("spikes must be 2D array (T, N)") + + n_state = int(args.num_states) + eps_rate_hz = 1e-6 + + rate_hz = spikes.mean(axis=0).astype(np.float64) / float(dt) + rate_hz_clip = np.maximum(rate_hz, eps_rate_hz) + b_vec = np.log(rate_hz_clip).astype(np.float32) + b_init = np.tile(b_vec[None, :], (n_state, 1)) + + init_meta = { + "method": "rate", + "common_across_states": True, + "rate_floor_hz": float(eps_rate_hz), + "formula": "B=log(max(rate_hz,eps))", + } + return b_init, init_meta + + +def init_B_from_spikes(spikes, dt, args): + """Select B initialization mode based on args.init_B.""" + opt = _normalize_init_opt(args.init_B) + if opt == "rand": + return None, None + if opt != "data": + raise ValueError(f"Unsupported --init_B option: {opt}") + return init_selfspikingB(spikes, dt, args) + + +def init_edgesA(spikes, Tmax=50000, verbose=True): + """OLS/covariance initialization of A from spike data Y (T x N).""" + t0 = time.perf_counter() + Y = np.asarray(spikes, dtype=np.float64) + if Y.ndim != 2: + raise ValueError("spikes must be 2D array (T, N)") + T_full, _ = Y.shape + T_use = min(int(Tmax), int(T_full)) + if T_use < 2: + raise ValueError("Need at least 2 time bins to initialize A") + Y = Y[:T_use] + + YpYp = Y[:-1].T @ Y[:-1] # Gram matrix + YYp = Y[1:].T @ Y[:-1] # one-step cross-correlation + + A_ols = YYp @ np.linalg.pinv(YpYp) + + kappa = float(np.linalg.cond(YpYp)) + rho = float(np.max(np.abs(np.linalg.eigvals(A_ols)))) + Y_pred = Y[:-1] @ A_ols.T + var_y = float(np.var(Y[1:])) + if var_y <= 0.0: + R2 = float("nan") + else: + R2 = float(1.0 - np.var(Y[1:] - Y_pred) / var_y) + frob = float(np.linalg.norm(A_ols, ord="fro")) + elapsed_sec = float(time.perf_counter() - t0) + + if verbose: + print("A-init OLS diagnostics:") + print(f" cond(YpYp) = {kappa:.2e}") + print(f" rho(A_ols) = {rho:.3f}") + print(f" R2 = {R2:.3f}") + print(f" ||A||_F = {frob:.3f}") + print(f" bins_used = {T_use}/{T_full}") + print(f" elapsed_s = {elapsed_sec:.3f}") + + meta = { + "method": "cov_ols", + "Tmax": int(Tmax), + "num_bins_used": int(T_use), + "num_bins_total": int(T_full), + "cond_YpYp": kappa, + "rho_A_init": rho, + "R2_1step": R2, + "fro_A_init": frob, + "elapsed_init_edgesA_sec": elapsed_sec, + } + return A_ols.astype(np.float32), meta + + +def init_A_from_spikes(spikes, args): + """Select A initialization mode based on args.init_A.""" + + opt = _normalize_init_opt(args.init_A) + if opt == "rand": + return None, None + if opt != "data": + raise ValueError(f"Unsupported --init_A option: {opt}") + A_init, meta = init_edgesA(spikes, verbose=False) + + if 1: # rescale A + offDiagFact = 3 + meta["offDiag_A_init_fact"] = offDiagFact + A_init = A_init * offDiagFact + + if 0: # thresholded off-diagonal renormalization + off_diag = ~np.eye(A_init.shape[0], A_init.shape[1], dtype=bool) + above_thr = (np.abs(A_init) > float(args.minW)) + below_thr = (np.abs(A_init) < float(args.minW)) + strong_off_diag = off_diag & above_thr + weak_off_diag = off_diag & below_thr + A_init[strong_off_diag] *= offDiagFact + A_init[weak_off_diag] = 0.0 + + + rho_max = float(args.rho_max) + if 0: # global spectral radius rescaling + rho_before = float(np.max(np.abs(np.linalg.eigvals(A_init)))) + if rho_before > rho_max: + A_init = A_init * (rho_max / rho_before) + rho_after = float(np.max(np.abs(np.linalg.eigvals(A_init)))) + meta["rho_A_init_raw"] = rho_before + meta["rho_A_init"] = rho_after + meta["rho_max_target"] = rho_max + + if args.verb > 0: + print("A-init OLS diagnostics:") + print(f" cond(YpYp) = {meta['cond_YpYp']:.1f}") + print(f" rho(A_ols) = {meta['rho_A_init']:.3f}") + print(f" R2 = {meta['R2_1step']:.3f}") + print(f" ||A||_F = {meta['fro_A_init']:.3f}") + print(f" bins_used = {meta['num_bins_used']}/{meta['num_bins_total']}") + print(f" elapsed_s = {meta['elapsed_init_edgesA_sec']:.3f}") + + return A_init, meta + + +def init_probs_from_states(S_rate, M): + """Build c_init from S_rate with transition smoothing.""" + S_rate = np.asarray(S_rate, dtype=np.int64) + T = S_rate.size + if T == 0: + return np.zeros((0, M), dtype=np.float32), np.zeros((0,), dtype=np.int64) + if M <= 1: + return np.ones((T, 1), dtype=np.float32), S_rate.copy() + + # Base: dominant state at 0.8, remaining 0.2 spread over all others. + other = 0.2 / float(M - 1) + c_init = np.full((T, M), other, dtype=np.float32) + c_init[np.arange(T), S_rate] = 0.8 + S_out = S_rate.copy() + + # Transition smoothing: (t-1) gets 2/3 old + 1/3 new, t gets 1/3 old + 2/3 new. + tr_idx = np.where(S_rate[1:] != S_rate[:-1])[0] + 1 + for t in tr_idx: + a = int(S_rate[t - 1]) + b = int(S_rate[t]) + c_init[t - 1, :] = 0.0 + c_init[t, :] = 0.0 + c_init[t - 1, a] = 2.0 / 3.0 + c_init[t - 1, b] = 1.0 / 3.0 + c_init[t, a] = 1.0 / 3.0 + c_init[t, b] = 2.0 / 3.0 + + return c_init, S_out diff --git a/causal_net/nonStation_ver3b_states_expr/prep_bioexp3b.py b/causal_net/nonStation_ver3b_states_expr/prep_bioexp3b.py new file mode 100755 index 00000000..b6d8c0d7 --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/prep_bioexp3b.py @@ -0,0 +1,377 @@ +#!/usr/bin/env python3 +""" +Preprocessing pipeline for experimental neural data from Roy/Mandar laboratory. + +This script processes raw experimental neural recordings into standardized format +for connectivity analysis. The preprocessing pipeline includes: +- Raw data loading and format conversion +- Temporal binning and spike count extraction +- Data quality assessment and filtering +- Metadata extraction and session identification +- Metrics spreadsheet ingestion (metrics_curated.xlsx) +- Output formatting for downstream analysis tools + +Session naming convention: +- B6J: cell line name +- 250619: recording date (YYMMDD) +- M08020: chip identifier +- 000093: run number +- Well000: well number + +The script generates .spikes.npz files with standardized spike count matrices +and associated metadata for further analysis. + +Usage: + ./prep_bioexp.py --sessionName B6J_250619_M08020_000093_Well000 --inputPath /path/to/raw/data/ +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" +import sys, os, hashlib +import numpy as np +import pandas as pd +from pprint import pprint +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + +import argparse + + +#...!...!.................. +def commandline_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v", "--verb", type=int, help="increase debug verbosity", default=1) + parser.add_argument("--expPath", required=True, help="raw experimnetal data on CFS") + + parser.add_argument("--sessionName", default='celllinename/dateofrecording/chipID/Assaytype/runnumber/wellnumber', help='raw data session name') + parser.add_argument("--dataPath", default='/pscratch/sd/b/balewski/2025_causalNet_tmp/', help="head dir for any further data processing") + parser.add_argument("--shortName", default=None, help='(optional) output file name - Is it needed?') + parser.add_argument("--inpFormat", default=1, type=int, choices=[1, 2], + help="input format selector: 1=metrics_curated.xlsx, 2=quality_metrics.xlsx") + + # .... activity speciffic speciffic, + parser.add_argument('--samp_freq', default=100, type=int, help='sets binning of time axis') + + parser.add_argument('--freqRange', default=[1., 50], type=float, nargs=2, help='rebin of raw time axis') + + args = parser.parse_args() + + for arg in vars(args): + print('myArgs:', arg, getattr(args, arg)) + + return args + + +#...!...!.................... +def buildBioMeta(args): + pd = {} # payload + pd['raw_bioexp_path'] = args.expPath + pd['session_name'] = args.sessionName + pd['inp_format'] = args.inpFormat + txtL = args.sessionName.split('/') + #print('tt',txtL); aa + pd['cell_line_name'] = txtL[0] + pd['recording_date'] = txtL[1] + pd['chip_ID'] = txtL[2] + pd['run_num'] = txtL[4] + pd['well_num'] = txtL[5] + + sel = {'freq_range': args.freqRange} + md = {'bioexp': pd, 'data_selector': sel} + myHN = hashlib.md5(os.urandom(32)).hexdigest()[:6] + md['hash'] = myHN + if args.shortName == None: + md['short_name'] = '%s-%s' % (pd['recording_date'], md['hash']) + else: + md['short_name'] = args.shortName + + if args.verb > 1: + print('\nBMD:') + pprint(md) + return md + + +def read_spike_npy(md, args): + pmd = md['bioexp'] + inpF = os.path.join(args.expPath, args.sessionName, 'spike_times.npy') + print('inpF:', inpF) + assert os.path.exists(inpF) + # Load the dictionary from the .npy file + spike_dict = np.load(inpF, allow_pickle=True).item() + + pmd['sampling_freq'] = args.samp_freq + + # neuron ID MEA chip + meaIdL = np.array(sorted(spike_dict)) # here order of feature_id is settled + maxFeat = len(meaIdL) + + if args.verb > 1: + print('RSN: meaID list:', meaIdL) + pmd['num_feature'] = len(meaIdL) + + spikeT = {} # spike times + spikeCntL = np.zeros(pmd['num_feature'], dtype=int) # num spikes per neuron + maxTbin = 0 + + j = 0 + for k in meaIdL: + rec = np.array(spike_dict[k]) * args.samp_freq + #print(rec[:100],len(rec)) + spikeT[k] = rec.astype(int) # time-bin may repeat + spikeCntL[j] = len(rec) + j += 1 + if len(rec) == 0: + continue + mxTb = np.max(rec) + if maxTbin < mxTb: + maxTbin = mxTb + #exit(0) + pmd['num_time_bin'] = int(maxTbin) + 1 + pmd['max_time'] = pmd['num_time_bin'] / pmd['sampling_freq'] + chanFreq = spikeCntL / pmd['max_time'] + rawD = {'spikeT': spikeT, 'chanFreq': chanFreq, + 'spike_key': meaIdL, 'MEA_idx': meaIdL.copy()} + return rawD + + +def _mea_match_key(val): + """Canonical key for matching MEA_idx across npy dict keys and spreadsheet.""" + if isinstance(val, (bytes, np.bytes_)): + val = val.decode() + if isinstance(val, (int, np.integer)): + return str(int(val)) + if isinstance(val, (float, np.floating)): + if np.isnan(val): + raise ValueError("MEA_idx is NaN in metrics spreadsheet") + f = float(val) + return str(int(f)) if f == int(f) else str(f) + s = str(val).strip() + try: + f = float(s) + return str(int(f)) if f == int(f) else s + except ValueError: + return s + + +def metrics_xlsx_name(args): + return "quality_metrics.xlsx" if args.inpFormat == 2 else "metrics_curated.xlsx" + + +def _read_metrics_frame(args): + metrics_name = metrics_xlsx_name(args) + xlsxF = os.path.join(args.expPath, args.sessionName, metrics_name) + print("metrics xlsx:", xlsxF) + assert os.path.exists(xlsxF), f"missing metrics spreadsheet: {xlsxF}" + + df = pd.read_excel(xlsxF, engine="openpyxl") + cols = list(df.columns) + cols[0] = "MEA_idx" + df.columns = cols + + if args.inpFormat == 2: + rename_map = { + "location_X": "loc_x", + "location_Y": "loc_y", + "location_Z": "loc_z", + } + df = df.rename(columns={k: v for k, v in rename_map.items() if k in df.columns}) + else: + assert "MEA_idx" in df.columns, "first spreadsheet column must be MEA_idx" + + return df, metrics_name + + +def _is_integral_key(val): + try: + fval = float(val) + except (TypeError, ValueError): + return False + return np.isfinite(fval) and fval == int(fval) + + +def _integer_if_possible(vals): + arr = np.asarray(vals) + try: + farr = arr.astype(np.float64) + except (TypeError, ValueError): + return arr + if np.all(np.isfinite(farr)) and np.all(farr == farr.astype(np.int64)): + return farr.astype(np.int64) + return arr + + +def load_metrics_curated(args, mea_idx_order, spike_key_order): + """Load metrics spreadsheet; filter and reorder rows to match output neuron order.""" + df, metrics_name = _read_metrics_frame(args) + + mea_keys = [_mea_match_key(v) for v in df["MEA_idx"].values] + if len(mea_keys) != len(set(mea_keys)): + raise ValueError(f"duplicate MEA_idx rows in {metrics_name}") + + row_by_mea = {_k: i for i, _k in enumerate(mea_keys)} + order_idx = [] + for mea_id in np.asarray(mea_idx_order).ravel(): + key = _mea_match_key(mea_id) + assert key in row_by_mea, ( + f"MEA_idx {mea_id!r} (key={key!r}) not found in {metrics_name}" + ) + order_idx.append(row_by_mea[key]) + + df_ord = df.iloc[order_idx].reset_index(drop=True) + col_names = [str(c) for c in df_ord.columns] + metrics_2d = df_ord.to_numpy() + mea_idx_order = _integer_if_possible(df_ord["MEA_idx"].to_numpy()) + + n_match = len(order_idx) + n_sheet = len(df) + print( + "%s: kept %d/%d rows, shape=%s, cols=%d" + % (metrics_name, n_match, n_sheet, metrics_2d.shape, len(col_names)) + ) + if args.verb > 1: + print(" columns:", col_names) + return metrics_2d, col_names, mea_idx_order + + +#...!...!.................... +def unroll_bioexp(rawD, bioMD, args): + pmd = bioMD['bioexp'] + sel = bioMD['data_selector'] + frLo, frHi = sel['freq_range'] + print('frLo, frHi', frLo, frHi) + assert frLo < frHi + chanFreq = np.asarray(rawD['chanFreq'], dtype=float) + # vectorized boolean mask for channels within (frLo, frHi) range + freqMask = (chanFreq >= frLo) & (chanFreq <= frHi) + sel['num_drop_neur_lo_hi_freq'] = [int(np.sum(chanFreq < frLo)), int(np.sum(chanFreq > frHi))] + print('freqMask all=%d , passed=%d' % (freqMask.shape[0], np.sum(freqMask))) + #print(sel);aaa + # --- drop channles out of freq range + chanFreq = rawD['chanFreq'][freqMask] + MEA_idx = rawD['MEA_idx'][freqMask] + spike_key = rawD['spike_key'][freqMask] + + # .... REMAP MATRICES TO FREQUENCY-SORTED ORDER (PRIMARY INDEX) + neur_freqIdx = np.argsort(chanFreq) # indices that sort chanFreq by value + neur_revFreqIdx = np.empty(len(neur_freqIdx), dtype=int) # natural_index -> freq_sorted_position + neur_revFreqIdx[neur_freqIdx] = np.arange(len(neur_freqIdx)) + + #--- reorder channles by frequency + chanFreq = chanFreq[neur_freqIdx] + MEA_idx = MEA_idx[neur_freqIdx] + spike_key = spike_key[neur_freqIdx] + print('chanFreq', chanFreq[:5], '...', chanFreq[-5:], 'Hz') + + # create spike matrix: rows=time bins, cols=accepted channels + ntime = pmd['num_time_bin'] + nchan = spike_key.shape[0] + spikes2D = np.zeros((ntime, nchan), dtype=np.int32) + spikeT = rawD['spikeT'] + for ic, key in enumerate(spike_key): + tV = spikeT[key] + if len(tV) == 0: + continue + # tV holds time-bin indices where this channel fired one or more spikes + # bincount returns a length-ntime vector with spike counts per bin (zeros elsewhere) + # minlength=ntime guarantees the vector spans the full recording duration + cnt = np.bincount(tV, minlength=ntime) + spikes2D[:, ic] = cnt.astype(np.int32) + print('spikes2D shape', spikes2D.shape) + + # keep handy in meta for downstream + sel['num_chan'] = nchan + sel['max_spike_per_bin'] = int(np.max(spikes2D)) + + Y_uchar = np.clip(spikes2D, 0, 255).astype(np.uint8) + spikeD = {'spikes': Y_uchar, + 'single_rates': chanFreq + } + + bioD = {} + bioD['neur_freqIdx'] = neur_freqIdx + bioD['neur_revFreqIdx'] = neur_revFreqIdx + bioD['single_rates'] = np.asarray(chanFreq, dtype=np.float64) + + metrics_2d, metrics_cols, MEA_idx = load_metrics_curated(args, MEA_idx, spike_key) + bioD['MEA_idx'] = MEA_idx + bioD['spike_key'] = spike_key + bioD['metrics_curated'] = metrics_2d + bioMD['metrics_curated_columns'] = metrics_cols + + #.... compute neural statistics for spikeMD + num_neurons = nchan + avg_rate = float(np.mean(chanFreq)) + std_rate = float(np.std(chanFreq)) + median_rate = float(np.median(chanFreq)) + min_rate = float(np.min(chanFreq)) + max_rate = float(np.max(chanFreq)) + + # Compute Fano factor (variance/mean) for each neuron + mean_counts_per_bin = np.mean(spikes2D, axis=0) + spike_variance = np.var(spikes2D, axis=0) + fano_factor = np.divide(spike_variance, mean_counts_per_bin, out=np.zeros_like(spike_variance), where=mean_counts_per_bin != 0) + avg_fano = float(np.mean(fano_factor)) + std_fano = float(np.std(fano_factor)) + + # Print summary statistics + print('Neural Statistics Summary:') + print('num neurons: %d, Avg Rate= %.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_neurons, avg_rate, std_rate, avg_fano, std_fano)) + print('Median rate %.2f Hz' % median_rate) + + #.... extract spikeMD for fitter + spikeMD = {'time_step_sec': 1. / pmd['sampling_freq'], + 'provenance': {'experiment_name': bioMD['short_name']}, + 'poisson_eta_clip': 5, # expected by fitter + 'data_type': 'bioExp', 'num_neurons': num_neurons} + + bioMD['rate_summary'] = { + 'avg_spike_rate': avg_rate, + 'std_spike_rate': std_rate, + 'avg_fano_factor': avg_fano, + 'std_fano_factor': std_fano, + 'median_spike_rate': median_rate, + 'min_spike_rate': min_rate, + 'max_spike_rate': max_rate + } + return bioD, spikeD, spikeMD + + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__ == "__main__": + + args = commandline_parser() + np.set_printoptions(precision=3) + bioMD = buildBioMeta(args) + + # read raw data + #rawD=read_spike_dict(bioMD,args) + rawD = read_spike_npy(bioMD, args) + + #.... filter & unroll data + bioD, spikeD, spikeMD = unroll_bioexp(rawD, bioMD, args) + + #...... WRITE OUTPUT ......... + outFt = os.path.join(args.dataPath, bioMD['short_name'] + '.bioExp.npz') + write_data_npz(bioD, outFt, metaD=bioMD) + if args.verb > 2: + print('\n bioD:', sorted(bioD)) + pprint(bioMD) + + outFs = outFt.replace('.bioExp.', '.spikes.') + write_data_npz(spikeD, outFs, metaD=spikeMD) + if args.verb > 2: + print('\nspikeD:', sorted(spikeD)) + pprint(spikeMD) + + print("\nNext step command:") + print(' ./view_bioexp.py --dataPath $dataPath --dataName %s -p a b c -T 0 3550 ' % (bioMD['short_name'])) + print(" ./fit_lassoPoisson.py --dataPath $dataPath --dataName %s --num_epochs 10 " % bioMD['short_name']) + print(" ./fitLasso4GPU.sh --dataPath $dataPath --dataName %s --num_epochs 200 " % bioMD['short_name']) + print(" ./bootsFit.sh --dataName %s --num_epochs 250 --dropDataFrac 0.5 --num_bootstraps 7 " % bioMD['short_name']) + + print(" ./selectEdges_FDR.py --dataName %s --num_bootstraps 6 10 -p a c d " % bioMD['short_name']) + + print(' --dataPath ' + args.dataPath) diff --git a/causal_net/nonStation_ver3b_states_expr/prism_EM_eval3b.py b/causal_net/nonStation_ver3b_states_expr/prism_EM_eval3b.py new file mode 100755 index 00000000..b5b299c7 --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/prism_EM_eval3b.py @@ -0,0 +1,357 @@ +#!/usr/bin/env python3 +""" +Evaluation and plotting for prism EM results. +""" + +import os +import argparse +import numpy as np +from pprint import pprint +from toolbox.Util_NumpyIO import read_data_npz +from PlotterPrismEM import Plotter +from UtilBioExp import detect_spike_bursts + + +def compute_ll_gap(spikes_sub, A_true, B_true, dt, eta_clip): + """Per-time LL gap between best and 2nd-best true-state likelihood.""" + y_prev = np.asarray(spikes_sub[:-1], dtype=np.float64) + y_curr = np.asarray(spikes_sub[1:], dtype=np.float64) + a_true = np.asarray(A_true, dtype=np.float64) + b_true = np.asarray(B_true, dtype=np.float64) + + t_pairs = y_prev.shape[0] + m_states = b_true.shape[0] + ll_gap = np.zeros((t_pairs,), dtype=np.float64) + if m_states <= 1: + return ll_gap + + for t in range(t_pairs): + yp = y_prev[t] + yc = y_curr[t] + base = a_true @ yp + eta = base[None, :] + b_true + eta_c = np.minimum(eta, float(eta_clip)) + lam = np.exp(eta_c) * float(dt) + scores = np.sum(yc[None, :] * eta_c - lam, axis=1) + top2 = np.partition(scores, -2)[-2:] + ll_gap[t] = float(top2[-1] - top2[-2]) + return ll_gap + + +def eval_em_metrics_time(fitD, md, spikes): + """Compute per-time metrics for -p f canvas.""" + trainMD = md["train"] + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + dt = float(trainMD["time_step_sec"]) + eta_clip = float(trainMD["eta_clip"]) + lambda2 = float(trainMD["lambda2"]) + + spikes_sub = np.asarray(spikes[t0_bin : t1_bin + 1], dtype=np.float64) + yp = spikes_sub[:-1] + yc = spikes_sub[1:] + + a_hat = np.asarray(fitD["A_hat"], dtype=np.float64) + b_hat = np.asarray(fitD["B_hat"], dtype=np.float64) + c_hat = np.asarray(fitD["c_hat"], dtype=np.float64) + n_pairs = spikes_sub.shape[0] - 1 + assert c_hat.shape[0] == spikes_sub.shape[0], "c_hat and spikes_sub must have matching time bins" + c_pairs = c_hat[1:] + c_prev = c_hat[:-1] + + rates = np.asarray(fitD["single_rates"], dtype=np.float64) + w = 1.0 / np.maximum(rates, 0.1) + w /= w.mean() + + eta = yp @ a_hat.T + c_pairs @ b_hat + eta_c = np.minimum(eta, eta_clip) + lam = np.exp(eta_c) * dt + log_dt = np.log(dt) + nll_t = np.sum(w[None, :] * (lam - yc * (eta + log_dt)), axis=1) + + dc = c_pairs - c_prev + l2_t = lambda2 * np.sum(dc * dc, axis=1) + + ll_gap = compute_ll_gap(spikes_sub, md["A_true"], md["B_true"], dt, eta_clip) + + s_true = np.asarray(md["S_true"])[t0_bin : t1_bin + 1] + s_hat = np.asarray(fitD["S_hat"]) + assert s_true.shape[0] == s_hat.shape[0], "S_true and S_hat must have matching length" + acc = float((s_true == s_hat).mean()) + + return { + "acc": acc, + "loss_nll_time": nll_t, + "loss_l2_time": l2_t, + "ll_gap": ll_gap, + "ll_gap_mean": float(np.mean(ll_gap)), + } + + +def eval_true_state_recovery(fitD, md): + """Compute state recovery table on full training window.""" + trainMD = md["train"] + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + s_true = np.asarray(md["S_true"], dtype=np.int64)[t0_bin : t1_bin + 1] + s_hat = np.asarray(fitD["S_hat"], dtype=np.int64) + s_hat_cl = np.asarray(fitD["S_hat_CL"], dtype=np.float64) + m_states = int(trainMD["num_states"]) + + assert s_true.shape[0] == s_hat.shape[0] == s_hat_cl.shape[0], \ + "S_true/S_hat/S_hat_CL length mismatch on training window" + + avg_acc = float((s_true == s_hat).mean()) + state_acc_cl = [] + bins_hat = np.bincount(s_hat, minlength=m_states).astype(np.int64) + enter_hat = np.zeros(m_states, dtype=np.int64) + if s_hat.shape[0] > 1: + enter_idx = np.where(s_hat[1:] != s_hat[:-1])[0] + 1 + if enter_idx.size > 0: + enter_hat = np.bincount(s_hat[enter_idx], minlength=m_states).astype(np.int64) + + trans_hat = np.zeros((m_states, m_states), dtype=np.int64) + if s_hat.shape[0] > 1: + np.add.at(trans_hat, (s_hat[:-1], s_hat[1:]), 1) + + for m in range(m_states): + mask = (s_hat == m) + cnt = int(mask.sum()) + if cnt > 0: + acc_m = float((s_true[mask] == m).mean()) + cl_m = float(s_hat_cl[mask].mean()) + else: + acc_m = float("nan") + cl_m = float("nan") + state_acc_cl.append([acc_m, cl_m, int(bins_hat[m]), int(enter_hat[m])]) # [acc, CL, bins_hat, enter_hat] + + return { + "avg_acc": avg_acc, + "state_acc_cl": state_acc_cl, + "trans_hat": trans_hat, + } + + +def compute_neuron_type(fitD, minW): + """Compute and store neuron_type vector: -1=inh, 0=und, +1=exc.""" + A_hat = np.asarray(fitD["A_hat"], dtype=np.float64) + assert A_hat.ndim == 2 and A_hat.shape[0] == A_hat.shape[1], "A_hat must be square" + N = A_hat.shape[0] + minW = float(minW) + off_mask = ~np.eye(N, dtype=bool) + A_thr = A_hat.copy() + A_thr[off_mask & (np.abs(A_thr) < minW)] = 0.0 + np.fill_diagonal(A_thr, 0.0) + Sedge = A_thr.sum(axis=1) + neuron_type = np.zeros((N,), dtype=np.int8) + neuron_type[Sedge > minW] = 1 + neuron_type[Sedge < -minW] = -1 + fitD["neuron_Sedge"] = Sedge.astype(np.float32) + fitD["neuron_type"] = neuron_type + return neuron_type + + +def compute_A_prune(fitD, neuron_type, minW): + """Compute and store A_prune from A_hat and neuron_type.""" + A_hat = np.asarray(fitD["A_hat"], dtype=np.float64) + neuron_type = np.asarray(neuron_type, dtype=np.int8) + assert neuron_type.shape[0] == A_hat.shape[0], "neuron_type length must match A_hat rows" + minW = float(minW) + N = A_hat.shape[0] + off_mask = ~np.eye(N, dtype=bool) + A_prune = A_hat.copy() + A_prune[off_mask & (np.abs(A_prune) < minW)] = 0.0 + diag_A = np.diag(A_hat).copy() + exc_rows = neuron_type > 0 + inh_rows = neuron_type < 0 + A_prune[exc_rows, :] = np.where(A_prune[exc_rows, :] > 0, A_prune[exc_rows, :], 0.0) + A_prune[inh_rows, :] = np.where(A_prune[inh_rows, :] < 0, A_prune[inh_rows, :], 0.0) + np.fill_diagonal(A_prune, diag_A) + fitD["A_prune"] = A_prune.astype(np.float32) + return fitD["A_prune"] + + +def main(): + parser = argparse.ArgumentParser( + description="Evaluate and plot prism EM results") + parser.add_argument("--dataName", type=str, required=True, + help="Base name for prismEM file") + parser.add_argument("--basePath", type=str, + default="/pscratch/sd/b/balewski/2026_causalNet_tmp3/", + help="Head dir for input/output data") + parser.add_argument("-p", "--showPlots", type=str, nargs='+', + default="a", + help="Plot types: a=EM convergence summary, b=init-vs-truth states, c=A_init-vs-truth, d=A_hat-vs-truth, e=A_hat edge recovery, f=state sequence, g=2D correlations (A/B), h=A_init TP quality (diag/exc/inh), i=A_init/A_hat fitted-only, j=state+bioExp rates (no truth), k=bioExp spatial A_hat edges, m=Nedge/Sedge node stats, n=bioExp A_prune pos/neg edges") + parser.add_argument("--minW", type=float, default=0.02, + help="Threshold for A-matrix edge eval") + parser.add_argument("--timeReb", type=int, default=20, + help="Time rebin factor for time-axis plots") + g = parser.add_argument_group("data") + g.add_argument("-T", "--time_range_sec", default=[0.0, 65.0], + nargs=2, type=float, + help="Time window [t0, t1] in seconds") + g.add_argument("-R", "--time_rebin2", default=50, type=int, + help="bioExp burst panels: rebin factor on spike time axis") + g.add_argument("--burst_freq_thres", default=5.0, type=float, + help="bioExp burst panels: per-neuron rate threshold (Hz)") + parser.add_argument("-X", "--noXterm", action="store_true", + help="Disable X terminal for plotting") + parser.add_argument("-v", "--verb", type=int, default=1, + help="Verbosity level") + args = parser.parse_args() + + args.inpPath = os.path.join(args.basePath, "prismFit") + args.outPath = os.path.join(args.basePath, "plots") + os.makedirs(args.outPath, exist_ok=True) + args.showPlots = ''.join(args.showPlots) + + print("EM-eval args:", vars(args), "\n") + + # ── load EM fit ────────────────────────────────────────────────── + fitFF = os.path.join(args.inpPath, f"{args.dataName}.prismEM.npz") + fitD, fitMD = read_data_npz(fitFF) + assert isinstance(fitMD, dict), "Expected metadata dict in prismEM file" + + if args.verb > 1: pprint(fitMD) + + MD = {**fitMD, "short_name": args.dataName} + + is_bioexp = fitMD.get("data_type") == "bioExp" + prov = fitMD["provenance"] + spikesPath = os.path.join(args.basePath, "spikesData") + + if is_bioexp: + st_name = prov["experiment_name"] + spikesFF = os.path.join(spikesPath, f"{st_name}.spikes.npz") + print(f"bioExp data: skipped simTruth/prismTruth, spikes={spikesFF}") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 1) + spikes = spikeD["spikes"] + bioFF = spikesFF.replace('spikes.npz', 'bioExp.npz') + bioD, bioMD = read_data_npz(bioFF, verb=args.verb > 1) + bio_plot_md = {**spikeMD, **bioMD} + bio_plot_md["dataName"] = args.dataName + if args.verb > 1: pprint(bioMD) + else: + truth_name = prov["state_model_file"] + truthPath = os.path.join(args.basePath, "truthDale") + truthFF = os.path.join(truthPath, f"{truth_name}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 1) + if args.verb > 1: + pprint(trueMD) + MD.update(trueMD) + MD["A_true"] = trueD["A_true"] + MD["B_true"] = trueD["B_true"] + MD["E_true"] = trueD["E_true"] + + st_name = prov["state_transition_file"] + ptFF = os.path.join(spikesPath, f"{st_name}.prismTruth.npz") + trD, _ = read_data_npz(ptFF, verb=args.verb > 1) + MD["S_true"] = trD["S_true"] + MD["C_true"] = trD["C_true"] + + spikesFF = os.path.join(spikesPath, f"{st_name}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 1) + spikes = spikeD["spikes"] + + MD["eval_f"] = eval_em_metrics_time(fitD, MD, spikes) + + if not is_bioexp: + reco = eval_true_state_recovery(fitD, MD) + MD["states_recovery_eval"]["avg_acc"] = reco["avg_acc"] + MD["states_recovery_eval"]["state_acc_cl"] = reco["state_acc_cl"] + + print(f"state reco avr acc {reco['avg_acc']:.3f}, {args.dataName}") + print(f" {'state':>5s} {'acc':>5s} {'CL':>6s} {'bins_hat':>8s} {'enter_hat':>9s}") + print(f" {'-----':>5s} {'-----':>5s} {'------':>6s} {'--------':>8s} {'---------':>9s}") + for m, (acc_m, cl_m, bins_m, enter_m) in enumerate(reco["state_acc_cl"]): + print(f" {m:5d} {acc_m:5.3f} {cl_m:6.3f} {bins_m:8d} {enter_m:9d}") + + trans_hat = np.asarray(reco["trans_hat"], dtype=np.int64) + n_states = trans_hat.shape[0] + print("\nS_hat transition counts (from row -> to col):") + hdr = " from\\to" + "".join(f"{j:8d}" for j in range(n_states)) + print(hdr) + print(" " + "-" * (len(hdr) - 1)) + for i in range(n_states): + row = f"{i:8d}" + "".join(f"{int(trans_hat[i, j]):8d}" for j in range(n_states)) + print(row) + + MD["short_name"] = args.dataName + neuron_type = compute_neuron_type(fitD, args.minW) + compute_A_prune(fitD, neuron_type, args.minW) + + # ── plot ─────────────────────────── + if is_bioexp: + assert not any(c in args.showPlots for c in "bcdefgh"), \ + "bioExp: plots b-h require ground truth; use -p a i j k m n" + + args.prjName = args.dataName + plot = Plotter(args) + + if 'a' in args.showPlots: + plot.summary_prismEM(fitD, MD, figId=1) + + if 'b' in args.showPlots: # need truth + plot.state_init_prismEM(fitD, MD, figId=2, time_reb=args.timeReb) + + if 'c' in args.showPlots: + plot.matrix_init_prismEM(fitD, MD, figId=3, est_key="A_init", est_label="A_init") + + if 'd' in args.showPlots: + plot.matrix_init_prismEM(fitD, MD, figId=3, est_key="A_hat", est_label="A_hat") + + if 'e' in args.showPlots: + plot.edge_recovery_prismEM( + fitD, MD, minW=args.minW, figId=4, est_key="A_hat", est_label="A_hat" + ) + + if 'f' in args.showPlots: + plot.state_seq_prismEM( + fitD, MD, figId=5, time_range_sec=args.time_range_sec + ) + + if 'g' in args.showPlots: + plot.eval_ABcorr_prismEM(fitD, MD, figId=6) + + if 'h' in args.showPlots: + plot.initA_quality_prismEM(fitD, MD, minW=args.minW, figId=7) + + if 'i' in args.showPlots: + plot.A_fitted_prismEM( + fitD, MD, spikeD["single_rates"], minW=args.minW, figId=8 + ) + + if 'j' in args.showPlots: + assert is_bioexp, "plot j requires bioExp data (experiment_name in provenance)" + rebD = detect_spike_bursts( + spikeD, spikeMD, args.time_rebin2, + args.burst_freq_thres, + ) + plot.state_seq_fitonly_prismEM( + fitD, MD, rebD, bio_plot_md, figId=9, + time_range_sec=args.time_range_sec, + ) + + if 'k' in args.showPlots: + assert is_bioexp, "plot k requires bioExp data (experiment_name in provenance)" + plot.neuron_spatial_Ahat_edges( + fitD, bioD, bio_plot_md, minW=args.minW, figId=10 + ) + + if 'm' in args.showPlots: + plot.node_outgoing_edge_stats_prismEM( + fitD, MD, spikeD["single_rates"], neuron_type, + minW=args.minW, figId=11, + est_key="A_hat", est_label="A_hat", + ) + + if 'n' in args.showPlots: + assert is_bioexp, "plot n requires bioExp data (experiment_name in provenance)" + plot.neuron_spatial_Ahat_edges_split( + fitD, bioD, bio_plot_md, neuron_type, fitD["neuron_Sedge"], + minW=args.minW, maxNeurons=24, figId=12 + ) + + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3b_states_expr/prism_EM_train3b.py b/causal_net/nonStation_ver3b_states_expr/prism_EM_train3b.py new file mode 100755 index 00000000..63454579 --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/prism_EM_train3b.py @@ -0,0 +1,888 @@ +#!/usr/bin/env python3 +""" +prism_EM_train.py — EM training for non-stationary Poisson GLM. + +Jointly infers from observed spikes only (no ground truth): + - shared connectivity matrix A (N×N) + - per-state bias vectors B (M×N) + - time-varying coefficients c_hat (T×M) on the probability simplex + +Algorithm: Block Coordinate Descent + E-step: single forward sweep PGD over time bins → update c_hat + M-step: m_epochs of Adam with DataLoader → update (A, B) + +Uses uniformly weighted Poisson NLL in both steps, L1 off-diagonal +soft-thresholding, and delayed soft spectral radius correction on A. +Numerical tricks and defaults follow prism_Mstep_train3.py (M-step) +and prism_Estep_train.py (E-step). + +Usage: + ./prism_EM_train.py --dataName --basePath $basePath -M 3 +""" + +import os +import time +import math +import secrets +import argparse +from pprint import pprint + +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader +from torch.utils.data.distributed import DistributedSampler +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from Util_PrismEM import init_states_vs_time, init_B_from_spikes, init_A_from_spikes + + +# ═══════════════════════════════════════════════════════════════════════ +# Model (from prism_Mstep_train3.py) +# ═══════════════════════════════════════════════════════════════════════ + +class SwitchingPoissonGLM(nn.Module): + """Shared A (N×N), per-state B (M×N), mixed by coefficients c.""" + + def __init__(self, N, M, eta_clip): + super().__init__() + self.N = N + self.M = M + self.eta_clip = eta_clip + self.A = nn.Parameter(torch.randn(N, N) * 0.1) + self.B = nn.Parameter(torch.randn(M, N) * 0.1) + + def forward(self, y_prev, c, dt): + """y_prev (batch, N), c (batch, M) → predicted rate (batch, N).""" + eta = y_prev @ self.A.t() + c @ self.B + return torch.exp(torch.clamp(eta, max=self.eta_clip)) * dt + + +# ═══════════════════════════════════════════════════════════════════════ +# Dataset (c array updated in-place after each E-step) +# ═══════════════════════════════════════════════════════════════════════ + +class PairDataset(Dataset): + """(y_prev, y_curr, c) triplets. + + self.c is a mutable numpy array updated in-place after each E-step. + """ + + def __init__(self, y_prev, y_curr, c): + self.y_prev = y_prev + self.y_curr = y_curr + self.c = c + + def __len__(self): + return self.y_prev.shape[0] + + def __getitem__(self, idx): + return ( + torch.from_numpy(self.y_prev[idx]).float(), + torch.from_numpy(self.y_curr[idx]).float(), + torch.from_numpy(self.c[idx]).float(), + ) + + +def init_distributed(): + """Initialize torch.distributed from torchrun environment.""" + is_dist = (int(os.environ.get("WORLD_SIZE", "1")) > 1) or ("RANK" in os.environ) + if is_dist: + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + rank = 0 + world_size = 1 + local_rank = 0 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + return is_dist, rank, world_size, local_rank, device + + +def cleanup_distributed(is_dist): + if is_dist and dist.is_initialized(): + dist.destroy_process_group() + + +def pair_range_for_rank(t_pairs, world_size, rank): + """Return inclusive pair-index range [p0, p1] assigned to given rank.""" + base = t_pairs // world_size + rem = t_pairs % world_size + n_loc = base + (1 if rank < rem else 0) + p0 = rank * base + min(rank, rem) + p1 = p0 + n_loc - 1 + return p0, p1, n_loc + + +# ═══════════════════════════════════════════════════════════════════════ +# E-step: simplex PGD (from prism_Estep_train.py) +# ═══════════════════════════════════════════════════════════════════════ + +def project_to_simplex(v): + """Project 1-D vector onto the probability simplex (Duchi et al.).""" + if v.numel() == 1: + return torch.ones_like(v) + u, _ = torch.sort(v, descending=True) + cssv = torch.cumsum(u, dim=0) - 1.0 + ind = torch.arange(1, v.numel() + 1, device=v.device, dtype=v.dtype) + cond = u - cssv / ind > 0 + if torch.any(cond): + rho = torch.nonzero(cond, as_tuple=False)[-1, 0] + theta = cssv[rho] / (rho + 1.0) + else: + theta = cssv[-1] / v.numel() + return torch.clamp(v - theta, min=0.0) + + +def run_estep(Y_prev, Y_curr, A, B, c_hat, + dt, eta_clip, lambda2, lr, pgd_iter): + """Single forward sweep E-step. Updates c_hat in-place on GPU. + + Args: + Y_prev, Y_curr: (T_pairs, N) GPU tensors + A: (N, N) detached parameter tensor + B: (M, N) detached parameter tensor + c_hat: (T_full, M) GPU tensor, modified in-place + + Returns: + mean NLL over all time pairs + """ + T_eff = c_hat.shape[0] + log_dt = math.log(dt) + c_prev = c_hat[0].clone() + nll_sum = 0.0 + + for t in range(1, T_eff): + y_p = Y_prev[t - 1] + y_c = Y_curr[t - 1] + + Z = (A @ y_p)[None, :] + B # (M, N) + + c_t = c_hat[t].clone() + for _ in range(pgd_iter): + eta = c_t @ Z # (N,) + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + g = Z @ (lam - y_c) + 2.0 * lambda2 * (c_t - c_prev) + c_t = project_to_simplex(c_t - lr * g) + + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + nll_sum += float((lam - y_c * (eta_c + log_dt)).sum().item()) + + c_hat[t] = c_t + c_prev = c_t + + return nll_sum / max(1, T_eff - 1) + + +def run_estep_shard(Y_prev, Y_curr, A, B, c_hat, + dt, eta_clip, lambda2, lr, pgd_iter, + t0, t1): + """E-step update on a shard of time bins t in [t0, t1], inclusive. + + Returns: + nll_sum_local, n_pairs_local + """ + if t1 < t0: + return 0.0, 0 + log_dt = math.log(dt) + c_prev = c_hat[t0 - 1].clone() + nll_sum = 0.0 + n_pairs = 0 + + for t in range(t0, t1 + 1): + y_p = Y_prev[t - 1] + y_c = Y_curr[t - 1] + + Z = (A @ y_p)[None, :] + B + + c_t = c_hat[t].clone() + for _ in range(pgd_iter): + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + g = Z @ (lam - y_c) + 2.0 * lambda2 * (c_t - c_prev) + c_t = project_to_simplex(c_t - lr * g) + + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + nll_sum += float((lam - y_c * (eta_c + log_dt)).sum().item()) + + c_hat[t] = c_t + c_prev = c_t + n_pairs += 1 + + return nll_sum, n_pairs + + +# ═══════════════════════════════════════════════════════════════════════ +# M-step helpers (from prism_Mstep_train3.py) +# ═══════════════════════════════════════════════════════════════════════ + +def offdiag_soft_threshold_(A, lr, lam): + """In-place proximal L1 on off-diagonal A elements.""" + if lam <= 0: + return + with torch.no_grad(): + n = A.shape[0] + mask = ~torch.eye(n, dtype=torch.bool, device=A.device) + thresh = lr * lam + v = A[mask] + A[mask] = v.sign() * (v.abs() - thresh).clamp(min=0.0) + + +def enforce_spectral_radius_(A, rho_max, correction_strength=1.0): + """In-place spectral-radius correction. Returns rho before correction.""" + with torch.no_grad(): + rho = torch.linalg.eigvals(A).abs().max().item() + alpha = min(1.0, max(0.0, float(correction_strength))) + if rho > rho_max and alpha > 0.0: + hard_scale = float(rho_max) / float(rho) + scale = 1.0 - alpha * (1.0 - hard_scale) + A.mul_(scale) + return rho + + +def sync_AB_via_rank0_avg_then_broadcast_(mdl, rho_max, correction_strength=1.0): + """Rank-0 averages A/B, enforces rho on averaged A, then broadcasts A/B.""" + if not (dist.is_available() and dist.is_initialized()): + enforce_spectral_radius_(mdl.A, rho_max, correction_strength) + return + + rank = dist.get_rank() + world = float(dist.get_world_size()) + with torch.no_grad(): + A_buf = mdl.A.data.clone() + B_buf = mdl.B.data.clone() + + dist.reduce(A_buf, dst=0, op=dist.ReduceOp.SUM) + dist.reduce(B_buf, dst=0, op=dist.ReduceOp.SUM) + + if rank == 0: + A_buf.div_(world) + B_buf.div_(world) + enforce_spectral_radius_(A_buf, rho_max, correction_strength) + + dist.broadcast(A_buf, src=0) + dist.broadcast(B_buf, src=0) + mdl.A.data.copy_(A_buf) + mdl.B.data.copy_(B_buf) + + +def run_mstep_epoch(model, loader, optimizer, device, dt, + l1_wt, lambda3, rho_max, + rho_every, apply_prune, apply_rho, rho_correction_strength, + off_mask, minW): + """One DataLoader pass updating A and B. Returns metrics dict. + """ + model.train() + mdl = model.module if hasattr(model, "module") else model + s_tot = s_nll = s_l1 = 0.0 + eps = 1e-8 + + for bi, (yp, yc, cc) in enumerate(loader): + yp = yp.to(device, non_blocking=True) + yc = yc.to(device, non_blocking=True) + cc = cc.to(device, non_blocking=True) + + optimizer.zero_grad(set_to_none=True) + pred = model(yp, cc, dt) + nll = (-yc * torch.log(pred + eps) + pred).mean() + if lambda3 > 0: + n = mdl.A.shape[0] + off_abs_mean = (mdl.A.abs() * l1_wt).sum() / float(n * (n - 1)) + l1 = lambda3 * off_abs_mean + else: + l1 = torch.tensor(0.0, device=device) + loss = nll + l1 + loss.backward() + optimizer.step() + + if apply_prune and lambda3 > 0: + offdiag_soft_threshold_(mdl.A, optimizer.param_groups[0]["lr"], lambda3) + if apply_rho and (bi % rho_every == 0): + sync_AB_via_rank0_avg_then_broadcast_( + mdl, rho_max, rho_correction_strength + ) + + s_tot += loss.item() + s_nll += nll.item() + s_l1 += l1.item() + + nb = max(1, len(loader)) + if dist.is_available() and dist.is_initialized(): + v = torch.tensor([s_tot, s_nll, s_l1, float(nb)], + dtype=torch.float64, device=device) + dist.all_reduce(v, op=dist.ReduceOp.SUM) + s_tot = float(v[0].item()) + s_nll = float(v[1].item()) + s_l1 = float(v[2].item()) + nb = int(v[3].item()) + + with torch.no_grad(): + nz = int((mdl.A[off_mask].abs() > minW).sum().item()) + rho = float(torch.linalg.eigvals(mdl.A).abs().max().item()) + + return dict(loss=s_tot / nb, nll=s_nll / nb, l1=s_l1 / nb, rho=rho, nz=nz) + + +# ═══════════════════════════════════════════════════════════════════════ +# Viterbi (from prism_Estep_train.py) +# ═══════════════════════════════════════════════════════════════════════ + +def viterbi_decode(c_hat, p_stay): + """Viterbi decode: find most likely state path from simplex coefficients.""" + eps = 1e-12 + T, M = c_hat.shape + if M == 1: + return np.zeros(T, dtype=np.int64) + p_sw = (1.0 - p_stay) / (M - 1) + trans = np.full((M, M), p_sw, dtype=np.float64) + np.fill_diagonal(trans, p_stay) + lt = np.log(np.clip(trans, eps, 1.0)) + le = np.log(np.clip(c_hat, eps, 1.0)) + + delta = np.zeros((T, M), dtype=np.float64) + psi = np.zeros((T, M), dtype=np.int64) + delta[0] = le[0] + for t in range(1, T): + sc = delta[t - 1][:, None] + lt + psi[t] = np.argmax(sc, axis=0) + delta[t] = sc[psi[t], np.arange(M)] + le[t] + + path = np.zeros(T, dtype=np.int64) + path[-1] = int(np.argmax(delta[-1])) + for t in range(T - 2, -1, -1): + path[t] = psi[t + 1, path[t + 1]] + return path + + +# ═══════════════════════════════════════════════════════════════════════ +# Args +# ═══════════════════════════════════════════════════════════════════════ + +def parse_args(): + p = argparse.ArgumentParser( + description="EM training: non-stationary Poisson GLM") + + p.add_argument("--dataName", required=True) + p.add_argument("--basePath", + default="/pscratch/sd/b/balewski/2026_causalNet_tmp3/") + p.add_argument("--num_states", "-M", type=int, required=True, + help="Number of latent states") + + g = p.add_argument_group("EM structure") + g.add_argument("--num_em_iters", type=int, default=20, + help="Outer EM iterations") + g.add_argument("--m_epochs", type=int, default=100, + help="M-step Adam epochs per EM iteration") + + g = p.add_argument_group("E-step") + g.add_argument("--pgd_iter", type=int, default=5, + help="PGD iterations per time bin") + g.add_argument("--lr_estep", type=float, default=0.03, + help="PGD step size") + g.add_argument("--lambda2", type=float, default=2.0, + help="Temporal smoothness weight") + + g = p.add_argument_group("M-step") + g.add_argument("--lr_mstep", type=float, default=0.003, + help="Adam learning rate (initial)") + g.add_argument("--lr_end_factor", type=float, default=0.1, + help="LR decays to lr_mstep * lr_end_factor") + g.add_argument("--lambda3", type=float, default=0.02, + help="L1 penalty on off-diagonal A") + g.add_argument("--rho_max", type=float, default=0.95, + help="Spectral radius target for delayed soft correction of A") + g.add_argument("--prescale_m_step_4_ArhoMax", type=int, default=50, + help="Spectral projection frequency (batches)") + g.add_argument("--delay_em_iter_4_ArhoMax", type=int, default=3, + help="EM iter to start rho_max enforcement " + "(default: int(num_em_iters * 0.6))") + g.add_argument("--delay_em_iter_4_lrDecay", type=int, default=1, + help="EM iter to start LR decay " + "(default: int(num_em_iters * 0.7))") + g.add_argument("--delay_em_iter_4_Aprune", type=int, default=1, + help="EM iter to start L1 pruning " + "(default: num_em_iters // 3)") + g.add_argument("--batch_size", type=int, default=2048) + g.add_argument("--minW", type=float, default=0.01, + help="Threshold for edge counting (reporting only)") + g = p.add_argument_group("data") + g.add_argument("-T", "--time_range_sec", default=[0.0, 60.0], + nargs=2, type=float, + help="Time window [t0, t1] in seconds") + + g = p.add_argument_group("decode") + g.add_argument("--decode_dwell_sec", type=float, default=0.3, + help="Viterbi dwell prior τ_dwell") + + g = p.add_argument_group("misc") + g.add_argument("--seed", type=int, default=42) + g.add_argument("--fitName", type=str, default=None, + help="Optional output fit name; if omitted a random name is generated") + g.add_argument("--init_states", type=str, default="data", + choices=["data", "rand"], + help="Initial latent-state mode: 'data' or 'rand'") + g.add_argument("--init_B", type=str, default="data", + choices=["data", "rand"], + help="Initial B mode: 'data' or 'rand'") + g.add_argument("--init_A", type=str, default="data", + choices=["data", "rand"], + help="Initial A mode: 'data' or 'rand'") + g.add_argument("-v", "--verb", type=int, default=1) + + return p.parse_args() + + +# ═══════════════════════════════════════════════════════════════════════ +# Main +# ═══════════════════════════════════════════════════════════════════════ + +def main(): + args = parse_args() + is_dist, rank, world_size, local_rank, device = init_distributed() + inpPath = os.path.join(args.basePath, "spikesData") + outPath = os.path.join(args.basePath, "prismFit") + out_ok = True + out_err = "" + if rank == 0: + out_ok = os.path.exists(outPath) + if not out_ok: + out_err = f"Output dir missing: {outPath}" + if is_dist: + msg = [out_ok, out_err] + dist.broadcast_object_list(msg, src=0) + out_ok, out_err = msg + if not out_ok: + raise FileNotFoundError(out_err) + + if args.delay_em_iter_4_Aprune is None: + args.delay_em_iter_4_Aprune = int(args.num_em_iters * 0.3) + if args.delay_em_iter_4_ArhoMax is None: + args.delay_em_iter_4_ArhoMax = int(args.num_em_iters * 0.6) + if args.delay_em_iter_4_lrDecay is None: + args.delay_em_iter_4_lrDecay = int(args.num_em_iters * 0.7) + + if rank == 0: + print("\nEM-train args:", vars(args), "\n") + print(f"world_size={world_size}") + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision("high") + + # ── load spikes (rank 0 only), then broadcast ─────────────── + if rank == 0: + spikesFF = os.path.join(inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 0) + spikes = spikeD["spikes"] + single_rates = spikeD["single_rates"] + if args.verb >1: pprint(spikeMD) + + if 0: # patch MD for experiments + if "poisson_eta_clip" not in spikeMD: spikeMD["poisson_eta_clip"]=5 + else: + spikeMD = None + spikes = None + single_rates = None + + if is_dist: + obj = [spikeMD] + dist.broadcast_object_list(obj, src=0) + spikeMD = obj[0] + + if rank == 0: + shp = torch.tensor(spikes.shape, dtype=torch.int64, device=device) + else: + shp = torch.zeros(2, dtype=torch.int64, device=device) + dist.broadcast(shp, src=0) + T_raw = int(shp[0].item()) + N = int(shp[1].item()) + + if rank == 0: + spikes_t = torch.as_tensor(spikes, dtype=torch.int32, device=device).contiguous() + else: + spikes_t = torch.empty((T_raw, N), dtype=torch.int32, device=device) + dist.broadcast(spikes_t, src=0) + spikes = spikes_t.cpu().numpy() + else: + T_raw, N = spikes.shape + + dt = float(spikeMD["time_step_sec"]) + eta_clip = float(spikeMD["poisson_eta_clip"]) + prov = dict(spikeMD["provenance"]) + M = args.num_states + + # ── time range selection ───────────────────────────────────── + t0_sec, t1_sec = float(args.time_range_sec[0]), float(args.time_range_sec[1]) + if t1_sec < t0_sec: + t0_sec, t1_sec = t1_sec, t0_sec + start_bin = max(0, int(math.floor(t0_sec / dt))) + end_bin = min(T_raw - 1, int(math.floor(t1_sec / dt))) + if end_bin <= start_bin: + raise ValueError("time_range_sec too small; need at least two bins") + spikes = spikes[start_bin : end_bin + 1] + T_full = spikes.shape[0] + T_pairs = T_full - 1 + if rank == 0: + print(f"N={N} EM={args.num_em_iters} M={M} T={T_full} pairs={T_pairs} " + f"dt={dt} eta_clip={eta_clip} " + f"time=[{t0_sec:.1f}, {t1_sec:.1f}]s bins=[{start_bin}, {end_bin}]") + + # ── time pairs: GPU for E-step, CPU numpy for DataLoader ──── + yp_np = spikes[:-1].astype(np.float32) + yc_np = spikes[1:].astype(np.float32) + Yp_gpu = torch.tensor(yp_np, device=device) + Yc_gpu = torch.tensor(yc_np, device=device) + + # ── initial states and parameter seeds (rank 0 computes, all ranks receive) ── + if rank == 0: + c_init_np, S_init, init_state_md, freq_h1d = init_states_vs_time(spikes, dt, args) + A_seed_np, init_A_md = init_A_from_spikes(spikes, args) + B_seed_np, init_B_md = init_B_from_spikes(spikes, dt, args) + else: + c_init_np = np.empty((T_full, M), dtype=np.float32) + S_init = np.empty((T_full,), dtype=np.int64) + freq_h1d = np.empty((0,), dtype=np.float32) + A_seed_np = None + B_seed_np = None + init_state_md = None + init_A_md = None + init_B_md = None + + if is_dist: + c_t = torch.as_tensor(c_init_np, dtype=torch.float32, device=device) + s_t = torch.as_tensor(S_init, dtype=torch.int64, device=device) + dist.broadcast(c_t, src=0) + dist.broadcast(s_t, src=0) + c_init_np = c_t.cpu().numpy() + S_init = s_t.cpu().numpy() + + f_len_t = torch.tensor([int(freq_h1d.shape[0]) if rank == 0 else 0], + dtype=torch.int64, device=device) + dist.broadcast(f_len_t, src=0) + f_len = int(f_len_t.item()) + if rank != 0: + freq_h1d = np.empty((f_len,), dtype=np.float32) + f_t = (torch.as_tensor(freq_h1d, dtype=torch.float32, device=device) + if rank == 0 else + torch.empty((f_len,), dtype=torch.float32, device=device)) + if f_len > 0: + dist.broadcast(f_t, src=0) + freq_h1d = f_t.cpu().numpy() + + has_A_t = torch.tensor([1 if (rank == 0 and A_seed_np is not None) else 0], + dtype=torch.int64, device=device) + dist.broadcast(has_A_t, src=0) + if int(has_A_t.item()) == 1: + A_t = (torch.as_tensor(A_seed_np, dtype=torch.float32, device=device) + if rank == 0 else + torch.empty((N, N), dtype=torch.float32, device=device)) + dist.broadcast(A_t, src=0) + A_seed_np = A_t.cpu().numpy() + else: + A_seed_np = None + + has_B_t = torch.tensor([1 if (rank == 0 and B_seed_np is not None) else 0], + dtype=torch.int64, device=device) + dist.broadcast(has_B_t, src=0) + if int(has_B_t.item()) == 1: + B_t = (torch.as_tensor(B_seed_np, dtype=torch.float32, device=device) + if rank == 0 else + torch.empty((M, N), dtype=torch.float32, device=device)) + dist.broadcast(B_t, src=0) + B_seed_np = B_t.cpu().numpy() + else: + B_seed_np = None + + c_hat_gpu = torch.tensor(c_init_np, dtype=torch.float32, device=device) + c_pairs_np = c_init_np[1:].copy() + + dataset = PairDataset(yp_np, yc_np, c_pairs_np) + if is_dist: + sampler = DistributedSampler( + dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False + ) + loader = DataLoader( + dataset, batch_size=args.batch_size, shuffle=False, sampler=sampler, + drop_last=False, pin_memory=True, num_workers=0 + ) + else: + sampler = None + loader = DataLoader( + dataset, batch_size=args.batch_size, shuffle=True, + drop_last=False, pin_memory=True, num_workers=0 + ) + + # ── model + optimizer + scheduler ──────────────────────────── + model = SwitchingPoissonGLM(N, M, eta_clip).to(device) + if A_seed_np is not None: + with torch.no_grad(): + model.A.copy_(torch.tensor(A_seed_np, dtype=torch.float32, device=device)) + if B_seed_np is not None: + with torch.no_grad(): + model.B.copy_(torch.tensor(B_seed_np, dtype=torch.float32, device=device)) + A_init_np = model.A.detach().cpu().numpy().copy() + B_init_np = model.B.detach().cpu().numpy().copy() + B_init_out_np = B_init_np[0] if B_init_np.ndim == 2 and B_init_np.shape[0] > 1 else B_init_np + + if is_dist: + model = DDP(model, device_ids=[local_rank]) + + total_m_epochs = args.num_em_iters * args.m_epochs + decay_start_m_epoch = args.delay_em_iter_4_lrDecay * args.m_epochs + decay_m_epochs = max(1, total_m_epochs - decay_start_m_epoch) + optimizer = optim.Adam(model.parameters(), lr=args.lr_mstep, fused=True) + scheduler = optim.lr_scheduler.LinearLR( + optimizer, start_factor=1.0, end_factor=args.lr_end_factor, + total_iters=decay_m_epochs, last_epoch=-1 + ) + scheduler._step_count = 1 + + mdl = model.module if hasattr(model, "module") else model + off_mask = ~torch.eye(N, dtype=torch.bool, device=device) + l1_wt = torch.ones(N, N, device=device) + l1_wt[torch.eye(N, dtype=torch.bool, device=device)] = 0.0 + + # ── history ────────────────────────────────────────────────── + h_e_nll = [] + h_m_loss, h_m_nll, h_m_l1 = [], [], [] + h_rho, h_nz, h_lr, h_rho_alpha = [], [], [], [] + + t_start = time.time() + m_epoch_global = 0 + p_stay = math.exp(-dt / args.decode_dwell_sec) + onehot_states = np.eye(M, dtype=np.float32) + + # ══════════════════════════════════════════════════════════════ + # EM loop + # ══════════════════════════════════════════════════════════════ + for em in range(1, args.num_em_iters + 1): + if rank == 0: + if em == args.delay_em_iter_4_ArhoMax + 1: + print(f"threshold passed: 'delay_em_iter_4_ArhoMax': {args.delay_em_iter_4_ArhoMax} (em={em})") + if em == args.delay_em_iter_4_lrDecay + 1: + print(f"threshold passed: 'delay_em_iter_4_lrDecay': {args.delay_em_iter_4_lrDecay} (em={em})") + if em == args.delay_em_iter_4_Aprune + 1: + print(f"threshold passed: 'delay_em_iter_4_Aprune': {args.delay_em_iter_4_Aprune} (em={em})") + apply_prune = em > args.delay_em_iter_4_Aprune + apply_rho = em > args.delay_em_iter_4_ArhoMax + em_iters_left = max(1, args.num_em_iters - em + 1) + rho_correction_strength = 1.0 / float(em_iters_left) if apply_rho else 0.0 + + # ── E-step ─────────────────────────────────────────────── + te0 = time.time() + with torch.no_grad(): + if is_dist: + p0, p1, n_loc = pair_range_for_rank(T_pairs, world_size, rank) + t0 = p0 + 1 + t1 = p1 + 1 + nll_local, n_pair_local = run_estep_shard( + Yp_gpu, Yc_gpu, mdl.A.data, mdl.B.data, c_hat_gpu, + dt, eta_clip, args.lambda2, args.lr_estep, args.pgd_iter, + t0, t1 + ) + + c_upd = torch.zeros_like(c_hat_gpu) + c_msk = torch.zeros((T_full, 1), dtype=torch.float32, device=device) + c_upd[0] = c_hat_gpu[0] + c_msk[0] = 1.0 + if n_loc > 0: + c_upd[t0:t1 + 1] = c_hat_gpu[t0:t1 + 1] + c_msk[t0:t1 + 1] = 1.0 + + dist.all_reduce(c_upd, op=dist.ReduceOp.SUM) + dist.all_reduce(c_msk, op=dist.ReduceOp.SUM) + c_avg = c_upd / torch.clamp(c_msk, min=1.0) + c_hat_gpu.copy_(torch.where(c_msk > 0.0, c_avg, c_hat_gpu)) + + ev = torch.tensor([nll_local, float(n_pair_local)], + dtype=torch.float64, device=device) + dist.all_reduce(ev, op=dist.ReduceOp.SUM) + e_nll = float(ev[0].item() / max(1.0, ev[1].item())) + else: + e_nll = run_estep( + Yp_gpu, Yc_gpu, mdl.A.data, mdl.B.data, + c_hat_gpu, dt, eta_clip, + args.lambda2, args.lr_estep, args.pgd_iter + ) + + te = time.time() - te0 + h_e_nll.append(e_nll) + + # sync E-step states to CPU dataset (in-place update) for M-step: + # Viterbi decode -> one-hot rows (classification-style M-step). + c_hat_np_iter = c_hat_gpu.cpu().numpy() + if is_dist: + if rank == 0: + s_iter_np = viterbi_decode(c_hat_np_iter, p_stay).astype(np.int64, copy=False) + s_iter_t = torch.as_tensor(s_iter_np, dtype=torch.int64, device=device) + else: + s_iter_t = torch.empty((T_full,), dtype=torch.int64, device=device) + dist.broadcast(s_iter_t, src=0) + s_iter_np = s_iter_t.cpu().numpy() + else: + s_iter_np = viterbi_decode(c_hat_np_iter, p_stay).astype(np.int64, copy=False) + np.copyto(c_pairs_np, onehot_states[s_iter_np[1:]]) + + # ── M-step: m_epochs of Adam ───────────────────────────── + tm0 = time.time() + for _ in range(args.m_epochs): + m_epoch_global += 1 + if sampler is not None: + sampler.set_epoch(m_epoch_global) + met = run_mstep_epoch( + model, loader, optimizer, device, dt, + l1_wt, args.lambda3, args.rho_max, + args.prescale_m_step_4_ArhoMax, apply_prune, apply_rho, + rho_correction_strength, + off_mask, args.minW + ) + if em > args.delay_em_iter_4_lrDecay: + scheduler.step() + + h_m_loss.append(met["loss"]) + h_m_nll.append(met["nll"]) + h_m_l1.append(met["l1"]) + h_rho.append(met["rho"]) + h_nz.append(met["nz"]) + h_lr.append(float(optimizer.param_groups[0]["lr"])) + h_rho_alpha.append(rho_correction_strength) + tm = time.time() - tm0 + + if rank == 0 and args.verb > 0: + n_off = int(off_mask.sum().item()) + sp = 1.0 - met["nz"] / max(1, n_off) + print( + f"EM {em:3d}/{args.num_em_iters} " + f"E_nll={e_nll:.4e}({te:.1f}s) " + f"M_nll={met['nll']:.4e} l1={met['l1']:.4e} " + f"rho={met['rho']:.4f} sp={sp:.3f} nz={met['nz']} " + f"lr={h_lr[-1]:.2e} ({tm:.1f}s) " + f"tot={time.time() - t_start:.0f}s" + ) + + if rank == 0: + print(f"\nTotal EM time: {time.time() - t_start:.1f}s") + + # ── Final Viterbi decode ───────────────────────────────── + c_hat_np = c_hat_gpu.cpu().numpy() + S_hat = viterbi_decode(c_hat_np, p_stay) + S_hat_CL = (1.0 - c_hat_np.max(axis=1)).astype(np.float32) + + mdl = model.module if hasattr(model, "module") else model + A_hat = mdl.A.detach().cpu().numpy() + B_hat = mdl.B.detach().cpu().numpy() + A_thr = A_hat.copy() + Nn = A_thr.shape[0] + off_mask = ~np.eye(Nn, dtype=bool) + A_thr[off_mask & (np.abs(A_thr) < float(args.minW))] = 0.0 + np.fill_diagonal(A_thr, 0.0) + neuron_Sedge = A_thr.sum(axis=1) + neuron_type = np.zeros((Nn,), dtype=np.int8) + neuron_type[neuron_Sedge > float(args.minW)] = 1 + neuron_type[neuron_Sedge < -float(args.minW)] = -1 + A_prune = A_hat.copy() + A_prune[off_mask & (np.abs(A_prune) < float(args.minW))] = 0.0 + diag_A = np.diag(A_hat).copy() + exc_rows = neuron_type > 0 + inh_rows = neuron_type < 0 + A_prune[exc_rows, :] = np.where(A_prune[exc_rows, :] > 0, A_prune[exc_rows, :], 0.0) + A_prune[inh_rows, :] = np.where(A_prune[inh_rows, :] < 0, A_prune[inh_rows, :], 0.0) + np.fill_diagonal(A_prune, diag_A) + + if args.fitName is None: + h6 = secrets.token_hex(3) + outF = f"{args.dataName}-EM-{h6}" + else: + outF = args.fitName + outFF = os.path.join(outPath, f"{outF}.prismEM.npz") + prov["EMtrain_file"] = outF + + outD = { + "A_init": A_init_np.astype(np.float32), + "A_hat": A_hat.astype(np.float32), + "A_prune": A_prune.astype(np.float32), + "neuron_type": neuron_type, + "neuron_Sedge": neuron_Sedge.astype(np.float32), + "B_init": B_init_out_np.astype(np.float32), + "B_hat": B_hat.astype(np.float32), + "freq_h1d": np.asarray(freq_h1d, dtype=np.float32), + "c_init": c_init_np.astype(np.float32), + "c_hat": c_hat_np.astype(np.float32), + "S_init": S_init.astype(np.int64), + "S_hat": S_hat.astype(np.int64), + "S_hat_CL": S_hat_CL, + "single_rates": np.asarray(single_rates), + "e_nll_em": np.asarray(h_e_nll, dtype=np.float64), + "m_loss_epoch": np.asarray(h_m_loss, dtype=np.float64), + "m_nll_epoch": np.asarray(h_m_nll, dtype=np.float64), + "m_l1_epoch": np.asarray(h_m_l1, dtype=np.float64), + "rho_epoch": np.asarray(h_rho, dtype=np.float64), + "rho_correction_strength_epoch": np.asarray(h_rho_alpha, dtype=np.float64), + "nz_edges_epoch": np.asarray(h_nz, dtype=np.int64), + "learning_rates": np.asarray(h_lr, dtype=np.float64), + } + + outMD = dict(spikeMD) + outMD["fit_type"] = "prismEM" + outMD["train"] = { + "num_em_iters": args.num_em_iters, + "m_epochs": args.m_epochs, + "total_m_epochs": total_m_epochs, + "pgd_iter": args.pgd_iter, + "lr_estep": args.lr_estep, + "lr_mstep": args.lr_mstep, + "lr_end_factor": args.lr_end_factor, + "lambda2": args.lambda2, + "lambda3": args.lambda3, + "rho_max": args.rho_max, + "rho_projection_mode": "soft_inverse_em_iters_left", + "prescale_m_step_4_ArhoMax": args.prescale_m_step_4_ArhoMax, + "delay_em_iter_4_ArhoMax": args.delay_em_iter_4_ArhoMax, + "rho_sync_mode": "broadcast", + "delay_em_iter_4_lrDecay": args.delay_em_iter_4_lrDecay, + "delay_em_iter_4_Aprune": args.delay_em_iter_4_Aprune, + "mstep_state_mode": "viterbi_onehot", + "batch_size": args.batch_size, + "minW": args.minW, + "num_states": M, + "num_neurons": N, + "num_time_bins": T_full, + "time_step_sec": dt, + "eta_clip": eta_clip, + "time_range_sec": [t0_sec, t1_sec], + "time_range_bins": [start_bin, end_bin], + "seed": args.seed, + } + outMD["states_recovery_eval"] = { + "decode": "viterbi", + "decode_dwell_sec": args.decode_dwell_sec, + } + outMD["init_A"] = init_A_md + outMD["init_state"] = init_state_md + outMD["init_B"] = init_B_md + outMD["provenance"] = prov + + write_data_npz(outD, outFF, metaD=outMD) + print(f"\nSaved: {outFF}") + print(f" basePath={args.basePath}") + print(f" ./prism_EM_eval3b.py --basePath $basePath " + f"--dataName {outF} -p a m n j i \n") + cleanup_distributed(is_dist) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3b_states_expr/toolbox b/causal_net/nonStation_ver3b_states_expr/toolbox new file mode 120000 index 00000000..c9203e9b --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/toolbox @@ -0,0 +1 @@ +../toolbox/ \ No newline at end of file diff --git a/causal_net/nonStation_ver3b_states_expr/view_bioexp.py b/causal_net/nonStation_ver3b_states_expr/view_bioexp.py new file mode 100755 index 00000000..0003e913 --- /dev/null +++ b/causal_net/nonStation_ver3b_states_expr/view_bioexp.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +""" +Visualization tool for biological experiment input features and data quality. + +This script provides comprehensive visualization and analysis of experimental +neural data, focusing on data quality assessment and feature exploration. +Main functionality includes: +- Time series visualization of neural recordings with customizable time ranges +- Data quality metrics and statistical summaries +- Cluster detection and activity pattern analysis +- Interactive plotting with configurable display options + +Used primarily for exploratory data analysis of biological neural recordings +before further processing and connectivity analysis. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os + +from pprint import pprint +import numpy as np +from PlotterBioExp import Plotter +from toolbox.Util_NumpyIO import read_data_npz +from UtilBioExp import detect_spike_bursts +import argparse + + +def _require_bioexp(bioD, bioMD): + for key in ("metrics_curated", "MEA_idx", "single_rates"): + assert key in bioD, f"bioExp.npz missing {key!r}; rerun prep_bioexp3b.py" + for key in ("metrics_curated_columns", "short_name", "data_selector", "rate_summary"): + assert key in bioMD, f"bioExp metadata missing {key!r}" + assert "freq_range" in bioMD["data_selector"] + + +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v", "--verbosity", type=int, + help="increase output verbosity", default=1, dest="verb") + parser.add_argument("-p", "--showPlots", default="a", nargs="+", + help="plots: a=freq histo, b=freq vs time, c=MEA layout, " + "d=metrics histograms, e=metrics correlations") + + parser.add_argument("-X", "--noXterm", action="store_true", + help="Disable X terminal for plotting") + + parser.add_argument("--dataPath", + default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", + help="head dir for any further data processing") + + parser.add_argument("-T", "--time_range", default=[0., 60], nargs=2, type=float, + help="display data time range in seconds") + parser.add_argument("--burst_freq_thres", default=5., type=float, + help="tags high freq channels for burst detection") + parser.add_argument("--dataName", default="HET_80k_1-fc62ef", + help="preprocessed session name") + + parser.add_argument("-R", "--time_rebin2", default=50, type=int, + help="rebin current time axis for burst panels") + + args = parser.parse_args() + args.outPath = "out/" + args.showPlots = "".join(args.showPlots) + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + assert args.time_range[0] < args.time_range[1] + assert os.path.exists(args.dataPath) + assert os.path.exists(args.outPath) + return args + + +#================================= +# M A I N +#================================= +if __name__ == "__main__": + args = get_parser() + np.set_printoptions(precision=3) + + bioFF = os.path.join(args.dataPath, f"{args.dataName}.bioExp.npz") + print("bioExp:", bioFF) + bioD, bioMD = read_data_npz(bioFF) + if args.verb > 1: pprint(bioMD) + _require_bioexp(bioD, bioMD) + + plotMD = dict(bioMD) + args.prjName = bioMD["short_name"] + + spikeD = None + rebD = None + if "a" in args.showPlots or "b" in args.showPlots: + spikesFF = os.path.join(args.dataPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF) + ''' + for key in ("spikes", "single_rates", "provenance", "data_type", + "time_step_sec", "num_neurons"): + assert key in spikeMD, f"spikes.npz metadata missing {key!r}" + ''' + assert "experiment_name" in spikeMD["provenance"] + plotMD = {**bioMD, **spikeMD} + args.prjName = spikeMD["provenance"]["experiment_name"] + if args.verb > 1: + pprint(spikeMD) + + if "b" in args.showPlots: + plotMD["plot"] = {"time_rangeLR": np.array(args.time_range)} + rebD = detect_spike_bursts( + spikeD, spikeMD, args.time_rebin2, + args.burst_freq_thres, + ) + + plot = Plotter(args) + + if "a" in args.showPlots: + plot.freq_histo(spikeD, plotMD, figId=1) + if "b" in args.showPlots: + plot.freq_vs_time(rebD, plotMD, figId=2) + if "c" in args.showPlots: + plot.neuron_spatial(bioD, plotMD, figId=3) + if "d" in args.showPlots: + plot.metrics_histograms(bioD, plotMD, figId=4) + if "e" in args.showPlots: + plot.metrics_correlations(bioD, plotMD, figId=5) + + plot.display_all() + print("M:done") + + diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/PlotterBioExp.py b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterBioExp.py new file mode 100644 index 00000000..a555658f --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterBioExp.py @@ -0,0 +1,373 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for biological experiment data visualization. + +This module provides specialized plotting capabilities for experimental +neural data analysis. The Plotter class extends PlotterBackbone to create +visualizations tailored for biological neural recordings including: +- Time series plots of neural activity patterns +- Statistical summaries and distribution analysis +- Data quality assessment plots +- Comparative analysis between experimental conditions + +Designed specifically for processing and visualizing data from biological +neural experiments, with automatic adaptation to experimental metadata. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from UtilBioExp import clip_rebD_time +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec + +from matplotlib.colors import LinearSegmentedColormap + +METRICS_HIST_COLUMNS = [ + "num_spikes", + "firing_rate", + "presence_ratio", + "snr", + "isi_violations_ratio", + "isi_violations_count", + "rp_contamination", + "rp_violations", + "sliding_rp_violation", + "amplitude_cutoff", + "amplitude_median", + "amplitude_cv_median", + "amplitude_cv_range", + "sync_spike_2", + "sync_spike_4", + "sync_spike_8", + "firing_range", + "sd_ratio", + "noise_cutoff", + "noise_ratio", +] + +METRICS_HIST_TITLE_UNITS = { + "firing_rate": "Hz", + "amplitude_median": "uV", +} + +METRICS_HIST_XMIN_ZERO = { + "snr", + "firing_rate", + "sync_spike_2", + "sync_spike_4", + "sync_spike_8", + "firing_range", + "sd_ratio", +} + +METRICS_HIST_LOG_Y = { + "presence_ratio", + "isi_violations_ratio", + "isi_violations_count", + "rp_violations", + "sliding_rp_violation", + "amplitude_cutoff", +} + +#...!...!.................... +def summary_column(md): + #print(sorted(md)) + rs=md['rate_summary'] + ds=md['data_selector'] + #pprint(md) + txt='dataset: '+md['short_name'] + txt += '\nnum acc neurons: %d' % md['num_neurons'] + txt += '\ndrop neur: %d <%.1f Hz, %d >%.1f Hz' % (ds['num_drop_neur_lo_hi_freq'][0],ds['freq_range'][0],ds['num_drop_neur_lo_hi_freq'][1],ds['freq_range'][1]) + txt+='\ndata type: %s\ntime_step=%.2f sec'%(md['data_type'],md['time_step_sec']) + txt += '\nMedian rate: %.1f Hz' % rs['median_spike_rate'] + txt += '\nrate range [%.1f, %.1f Hz]' % (rs['min_spike_rate'], rs['max_spike_rate']) + txt += '\nAvg Rate: %.1f±%.1f Hz' % (rs['avg_spike_rate'], rs['std_spike_rate']) + txt += '\nAvg Fano: %.2f±%.2f' % (rs['avg_fano_factor'], rs['std_fano_factor']) + + return txt + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def freq_histo(self,spikeD,md,figId=1): + + tit='dataset '+md['short_name'] + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,4)) + + # Create gridspec with 7:3 width ratio + gs = gridspec.GridSpec(1, 2, width_ratios=[7, 3]) + + dataYield, dataRates = spikeD['spikes'], spikeD['single_rates'] + + # Compute median + median_val = np.median(dataRates) + txtM= f'median : {median_val:.2f} Hz' + + #.... freq per channel + ax = self.plt.subplot(gs[0, 0]) + chanV=np.arange(dataRates.shape[0]) + ax.bar(chanV, dataRates , width=0.8, color='g', align='center', alpha=0.7) + + ax.axhline(median_val, color='red', linestyle='--', linewidth=1) + if md['data_type']=='simDale': + ax.set_yscale('log') + ax.grid() + ax.set(xlabel='input neuron index', ylabel='avr frequency (Hz)',title=tit) + ax.text( ax.get_xlim()[1]*0.1,median_val,txtM, + color='red', ha='center', va='bottom', fontsize=10) + + txt=summary_column(md) + ax.text(0.05, 0.3, txt, transform=ax.transAxes, fontsize=12,color='b') + + #........ freq histo ..... + ax = self.plt.subplot(gs[0, 1]) + ax.hist(dataRates,bins=20) + #ax.set_yscale('log') + ax.set(ylabel='num channels', xlabel='avr frequency (Hz)',title=tit) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + # Add median text annotation + ax.text(median_val+1, ax.get_ylim()[1]*0.5,txtM, + color='red', ha='left', va='bottom', fontsize=10,rotation=90) + ax.text(0.25, 0.3, txt, transform=ax.transAxes, fontsize=11,color='m') + + ax.grid() + + def plot_bioexp_sum_rate(self, ax, rebD, clip): + timeV = clip["timeV"] + rate1D = clip["rate1D"] + Tbin = clip["Tbin"] + tL, tR = clip["tL"], clip["tR"] + nchan = clip["nchan"] + ax.bar(timeV + Tbin * 0.5, rate1D, width=Tbin, color='orange', align='center', alpha=0.7) + ax.set(ylabel='sum rate (Hz)', title=f'sum rate from all {nchan} neurons') + ax.set_xlim(tL, tR) + ax.grid() + + def plot_bioexp_heatmap(self, fig, ax, rebD, dataset_name, clip): + rate2D = clip["rate2D"] + tL, tR = clip["tL"], clip["tR"] + nchan = clip["nchan"] + time_step = clip["time_step"] + tit0 = f'dataset: {dataset_name} nchan={nchan} Tbin={time_step:.2f} sec' + + original_cmap = self.plt.get_cmap('tab20c') + custom_cmap = original_cmap.copy() + custom_cmap.set_under('white') + + cax = ax.imshow( + rate2D.T, aspect='auto', cmap=custom_cmap, origin='lower', + extent=[tL, tR, 0, nchan], interpolation='nearest', + vmin=1, vmax=rate2D.max(), + ) + ax.set_ylabel('Neuron index') + ax.set_xlabel('Time (sec)') + ax.set_title(tit0) + cbar = fig.colorbar( + cax, ax=ax, orientation='horizontal', pad=0.17, + label=f'Instantanous freq (Hz) per neuron over Tbin={time_step:.1f} sec', + shrink=0.5, + ) + cbar.ax.tick_params(labelsize=10) + ax.grid() + + def freq_vs_time(self,rebD,md,figId=2): + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,11)) + + rateThr2=rebD['rate_thres2'] + clip = clip_rebD_time(rebD, md['plot']['time_rangeLR']) + print('iTL,R', clip['itL'], clip['itR']) + + highChanCnt = rebD['highChanCnt'][clip['itL']:clip['itR']] + timeV = clip['timeV'] + Tbin = clip['Tbin'] + tL, tR = clip['tL'], clip['tR'] + nchan = clip['nchan'] + + gs = fig.add_gridspec(4, 1, height_ratios=[0.15, 0.15, 0.54, 0.01]) + + ax = fig.add_subplot(gs[0, 0]) + ax.bar(timeV+Tbin*.5, highChanCnt, width=Tbin, color='forestgreen', align='center', alpha=0.7) + tit = (f'num neurons with instantaneous rate > thres={rateThr2:.0f} (Hz), ' + f'nchan={nchan}') + ax.set(ylabel='num neurons', title=tit) + ax.set_xlim(tL, tR) + ax.grid() + + ax = fig.add_subplot(gs[1, 0]) + self.plot_bioexp_sum_rate(ax, rebD, clip) + + ax = fig.add_subplot(gs[2, 0]) + self.plot_bioexp_heatmap(fig, ax, rebD, md['short_name'], clip) + + def _metrics_column(self, bioD, md, col_name): + cols = md["metrics_curated_columns"] + metrics = np.asarray(bioD["metrics_curated"]) + col_map = {str(c): i for i, c in enumerate(cols)} + assert col_name in col_map, ( + f"metrics_curated missing column {col_name!r}; have {cols}" + ) + return metrics[:, col_map[col_name]].astype(np.float64) + + def neuron_spatial(self, bioD, md, figId=3): + """MEA layout: loc_x/loc_y scatter colored by firing_rate (metrics_curated).""" + loc_x = self._metrics_column(bioD, md, "loc_x") + loc_y = self._metrics_column(bioD, md, "loc_y") + rate = self._metrics_column(bioD, md, "firing_rate") + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(8, 9)) + gs = gridspec.GridSpec(2, 1, height_ratios=[1.0, 0.06], hspace=0.12) + + ax = fig.add_subplot(gs[0, 0]) + ax.set_facecolor("white") + sc = ax.scatter( + loc_x, loc_y, c=rate, s=12, marker="s", + cmap="Blues", vmin=0.0, vmax=float(np.max(rate)), + linewidths=0.15, edgecolors="k", alpha=0.95, + ) + ax.set_aspect("equal", adjustable="box") + ax.set_xlabel("loc_x (um)", fontsize=11) + ax.set_ylabel("loc_y (um)", fontsize=11) + ax.tick_params(labelsize=9) + ax.grid(True, alpha=0.35) + tit = f"dataset: {md['short_name']} n={loc_x.shape[0]} neurons" + ax.set_title(tit, fontsize=11, pad=8) + + cax = fig.add_subplot(gs[1, 0]) + cb = fig.colorbar(sc, cax=cax, orientation="horizontal") + cb.set_label("Firing Rate [Hz]", fontsize=11) + cb.ax.tick_params(labelsize=9) + + fig.subplots_adjust(left=0.08, right=0.98, top=0.94, bottom=0.10) + + def _hist_panel_title(self, col_name): + unit = METRICS_HIST_TITLE_UNITS[col_name] + return f"{col_name} [{unit}]" + + def _fmt_metric_val(self, val): + av = abs(float(val)) + if av >= 100: + return f"{val:.0f}" + if av >= 1: + return f"{val:.3f}" + if av >= 0.01: + return f"{val:.4f}" + return f"{val:.2e}" + + def metrics_histograms(self, bioD, md, figId=4): + """1D histograms for curated unit-quality metrics (metrics_curated columns).""" + nrow, ncol = 3, 7 + n_panels = nrow * ncol + assert len(METRICS_HIST_COLUMNS) == n_panels - 1, ( + f"expected {n_panels - 1} metric columns, got {len(METRICS_HIST_COLUMNS)}" + ) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(18, 8)) + n_neur = np.asarray(bioD["metrics_curated"]).shape[0] + fr_lo, fr_hi = md["data_selector"]["freq_range"] + + for i, col_name in enumerate(METRICS_HIST_COLUMNS): + ax = fig.add_subplot(nrow, ncol, i + 1) + vals = self._metrics_column(bioD, md, col_name) + vals = vals[np.isfinite(vals)] + assert vals.size > 0, f"no finite values for metrics column {col_name!r}" + n_bins = min(40, max(10, int(np.sqrt(vals.size)))) + ax.hist(vals, bins=n_bins, color="steelblue", alpha=0.85, edgecolor="white") + if col_name in METRICS_HIST_LOG_Y: + ax.set_yscale("log") + + if col_name in METRICS_HIST_XMIN_ZERO: + ax.set_xlim(left=0.0) + + p16, p50, p84 = np.percentile(vals, [16, 50, 84]) + ylo, yhi = ax.get_ylim() + y_mark = 0.5 * (ylo + yhi) + ax.errorbar( + p50, y_mark, + xerr=[[p50 - p16], [p84 - p50]], + fmt="o", color="k", markersize=9, + capsize=4, capthick=1.2, elinewidth=1.2, zorder=5, + ) + txt = ( + f"p16={self._fmt_metric_val(p16)}\n" + f"med={self._fmt_metric_val(p50)}\n" + f"p84={self._fmt_metric_val(p84)}" + ) + ax.text( + 0.98, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=11, color="k", + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=2), + ) + + title = ( + self._hist_panel_title(col_name) + if col_name in METRICS_HIST_TITLE_UNITS else col_name + ) + ax.set_title(title, fontsize=14) + ax.tick_params(labelsize=7) + ax.grid(True, alpha=0.35) + + fig.add_subplot(nrow, ncol, n_panels).axis("off") + + fig.suptitle( + f"curated metrics, {md['short_name']}, N={n_neur} neurons, " + f"{fr_lo:g}< freq <{fr_hi:g} Hz", + fontsize=12, + ) + fig.tight_layout(rect=[0, 0, 1, 0.94], h_pad=1.62, w_pad=1.62) + + def metrics_correlations(self, bioD, md, figId=5): + """3x2 panel grid; panel 0: firing_rate vs single_rates scatter.""" + rate_sheet = self._metrics_column(bioD, md, "firing_rate") + rate_spikes = np.asarray(bioD["single_rates"], dtype=np.float64).ravel() + assert rate_sheet.shape[0] == rate_spikes.shape[0], ( + "firing_rate and single_rates length mismatch" + ) + + n = rate_sheet.shape[0] + r = float(np.corrcoef(rate_sheet, rate_spikes)[0, 1]) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(12, 10)) + fr_lo, fr_hi = md["data_selector"]["freq_range"] + + for k in range(6): + ax = fig.add_subplot(3, 2, k + 1) + if k == 0: + ax.scatter( + rate_sheet, rate_spikes, s=14, alpha=0.65, + color="steelblue", edgecolors="k", linewidths=0.2, + ) + lo = float(min(rate_sheet.min(), rate_spikes.min())) + hi = float(max(rate_sheet.max(), rate_spikes.max())) + pad = 0.05 * max(1e-9, hi - lo) + ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], + "k--", lw=0.9, alpha=0.5) + ax.set_xlabel("firing_rate (metrics_curated)", fontsize=10) + ax.set_ylabel("single_rates (bioExp)", fontsize=10) + ax.set_title(f"r={r:.3f}, n={n}", fontsize=11) + ax.grid(True, alpha=0.35) + else: + ax.axis("off") + + fig.suptitle( + f"metrics correlations, {md['short_name']}, N={n} neurons, " + f"{fr_lo:g}< freq <{fr_hi:g} Hz", + fontsize=12, + ) + fig.tight_layout(rect=[0, 0, 1, 0.94]) + diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/PlotterDaleMatrix.py b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterDaleMatrix.py new file mode 100644 index 00000000..563b5c12 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterDaleMatrix.py @@ -0,0 +1,426 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for simulated Poisson process visualization. + +This module provides specialized plotting capabilities for analyzing +simulated Dale's principle neural networks. The Plotter class extends +PlotterBackbone to create visualizations including: +- Dale connectivity matrix plots with excitatory/inhibitory separation +- Eigenvalue analysis and network stability visualization +- Poisson process statistics and firing rate distributions +- Network dynamics and temporal evolution plots + +Designed specifically for validating and analyzing simulated neural +networks that follow Dale's principle with Poisson spiking dynamics. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec +import matplotlib.colors as colors +import numpy as np +import matplotlib.pyplot as plt +from scipy.optimize import curve_fit +from Util_pseudospectra import compute_pseudospectrum +from UtilDalePoisson import get_offdiag_triplets + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def Dale_matrix_and_eigen(self, A, md, trueD, spikeD, figId=3): + + figId=self.smart_append(figId) + nrow,ncol=1,3 + fig=self.plt.figure(figId,facecolor='white', figsize=(17,5.)) + + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel if R_sel is not None else '' + + vmin = A.min() + vmax = A.max() + normMap = colors.TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax) + + #.... Position 1: Dale matrix (natural order) ...... + ax = self.plt.subplot(nrow,ncol,1) + im=ax.imshow(A, aspect=1., origin='lower', cmap='bwr', norm=normMap, interpolation='nearest') + ax.set( + xlabel='source / presynaptic neuron index (column)', + ylabel='target / postsynaptic neuron index (row)', + ) + ax.set_aspect(1.0) + ax.grid() + cbar = fig.colorbar(im, ax=ax, extend="both", shrink=0.7) + + tit='True Dale, N%d%s, %s'%(A.shape[0], R_tag, md['short_name']) + ax.set(title=tit) + ax.axvline(numExc-0.5,color='k',ls='--', label='source E/I boundary') + ax.axhline(numExc-0.5,color='0.5',ls=':', label='same index boundary') + ax.plot([0,numNeur],[0,numNeur],'--',lw=0.8,color='magenta') + + #..... Position 2: Eigenvalues...... + ax = self.plt.subplot(nrow,ncol,2) + Eigen=np.linalg.eigvals(A) + real_parts = np.real(Eigen) + imag_parts = np.imag(Eigen) + ax.scatter(real_parts, imag_parts, color='blue', marker='o') + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + tit3='Eigenvalues, N%d%s, %s'%(A.shape[0], R_tag, md['short_name']) + ax.set_title(tit3) + ax.axhline(0, color='black', lw=0.5) + ax.axvline(0, color='black', lw=0.5) + ax.axvline(0,color='red', linestyle='--') + ax.set_aspect('equal') + if R_sel is not None: + theta = np.linspace(0, 2*np.pi, 200) + ax.plot(R_sel*np.cos(theta), R_sel*np.sin(theta), color='magenta', linestyle='--', lw=1.5, label='R=%.3f'%R_sel) + ax.legend() + ax.grid(True) + + #..... Position 3: spatial node locations, signed by Dale type ...... + assert 'node_positions' in trueD, ( + 'Dale_matrix_and_eigen requires node_positions in simTruth; regenerate with gen_daleMatrices3c.py' + ) + node_pos = np.asarray(trueD['node_positions'], dtype=np.float64) + assert node_pos.shape == (numNeur, 2), ( + f"node_positions shape {node_pos.shape} does not match expected {(numNeur, 2)}" + ) + single_rates = np.asarray(spikeD['single_rates'], dtype=np.float64).reshape(-1) + assert single_rates.shape[0] == numNeur, ( + f"single_rates length {single_rates.shape[0]} does not match N={numNeur}" + ) + signed_rates = single_rates.copy() + signed_rates[numExc:] *= -1.0 + vmax_rate = float(np.max(np.abs(signed_rates))) + assert vmax_rate > 0.0, "signed firing rates are all zero" + rate_norm = colors.TwoSlopeNorm(vmin=-vmax_rate, vcenter=0.0, vmax=vmax_rate) + + ax = self.plt.subplot(nrow, ncol, 3) + sc = ax.scatter( + node_pos[:, 0], node_pos[:, 1], + c=signed_rates, cmap='bwr', norm=rate_norm, + s=14, marker='s', edgecolors='k', linewidths=0.15, alpha=0.95, + ) + ax.set_xlabel('node x') + ax.set_ylabel('node y') + ax.set_title('Node locations, signed rate') + ax.grid(True, alpha=0.35) + ax.set_aspect('equal', adjustable='box') + cbar = fig.colorbar(sc, ax=ax, extend='both', shrink=0.8) + cbar.set_label('signed firing rate (Hz)') + + def Dale_matrix_pseudospectra(self, A, md, trueD, figId=5): # p=e + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(8, 6)) + + npts = 80 + minY = -0.2 + epsMin = 0.03 + + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + ax = self.plt.subplot(1, 1, 1) + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel if R_sel is not None else '' + title = 'Pseudospectra, true N%d%s, %s' % (A.shape[0], R_tag, md['short_name']) + + levels = np.logspace(-2.5, -0.5, 10) + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + ax.text(0.05, 0.85, 'green: ε<%.3f'%epsMin, transform=ax.transAxes) + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) + + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(minY,1.1) + ax.set_xlim(-1.1,1.1) + ax.set_aspect(1.) + ax.axhline(1,color='m',linestyle='--') + ax.axvline(1,color='m',linestyle='--') + ax.axvline(-1,color='m',linestyle='--') + +#...!...!.................. + def histo_weights_rates(self,trueD,spikeD,md,figId=3): + figId=self.smart_append(figId) + nrow,ncol=2,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,7.2)) + fig.subplots_adjust(hspace=0.42, wspace=0.34) + data_name=md['short_name'] + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.2f' % R_sel + + A=trueD['A_true'] + single_rates = np.asarray(spikeD['single_rates'], dtype=np.float64).reshape(-1) + + # output: np.column_stack([i_indices, j_indices, values]) + EposT=get_offdiag_triplets(A,isPos=True) + EnegT=get_offdiag_triplets(A,isPos=False) + print('True0 num edges pos=%d neg=%d'%(EposT.shape[0],EnegT.shape[0])) + + A_diag = np.diag(A) + wMin = min(np.min(EnegT[:,2]), np.min(A_diag)) + wMax = max(np.max(EposT[:,2]), np.max(A_diag)) + + def count_elements(E, Nn=numNeur): + source_indices = E[:, 1].astype(int) + counts = np.bincount(source_indices, minlength=Nn) + return counts + + edgeCount=count_elements(EnegT) + count_elements(EposT) + + # Natural neuron indexing: first numExc are excitatory, rest inhibitory + inh_mask = np.zeros(numNeur, dtype=bool) + inh_mask[numExc:] = True + exc_mask = ~inh_mask + neurXlab = 'source neuron index' + x_vals = np.arange(numNeur) + + A_off = np.asarray(A, dtype=np.float64).copy() + np.fill_diagonal(A_off, 0.0) + edge_mask = np.abs(A_off) > 1e-12 + nedge = np.count_nonzero(edge_mask, axis=0) + sedge = np.sum(A_off, axis=0) + med_weight = np.zeros((numNeur,), dtype=np.float64) + for j in range(numNeur): + vals = A_off[:, j] + vals = vals[np.abs(vals) > 1e-12] + if vals.size: + med_weight[j] = float(np.median(vals)) + n_exc = int(np.count_nonzero(exc_mask)) + n_inh = int(np.count_nonzero(inh_mask)) + exc_color = 'red' + inh_color = 'blue' + + #.... weights histogram + ax = self.plt.subplot(nrow,ncol,1) + binX= np.linspace(wMin, wMax, 100) + ax.hist(EnegT[:,2], bins=binX, color='blue', alpha=0.7, edgecolor=None,label='neg:%d'%EnegT.shape[0]) + ax.hist(EposT[:,2], bins=binX, color='red', alpha=0.7, edgecolor=None,label='pos:%d'%EposT.shape[0]) + ax.hist(A_diag, bins=binX, color='cyan', alpha=0.65, edgecolor=None, + label='diag:%d'%A_diag.shape[0]) + + ax.legend(loc='upper left') + tit='True Dale, N%d%s, %s'%(A.shape[0], R_tag, md['short_name']) + ax.set(title=tit, xlabel='True weight value',ylabel='num edges') + ax.axvline(0,color='k',ls='--') + ax.grid(True, alpha=0.3) + + #.... : histogram of rates + ax = self.plt.subplot(nrow,ncol,3) + ax.hist(single_rates, bins=20) + ax.set_xlabel('Firing rate (Hz)') + ax.set_ylabel('num neurons') + ax.grid(True, alpha=0.3) + ax.set_title('Single rates spectrum') + ax.set_xlim(0,) + median_val = np.median(single_rates) + ax.axvline(median_val, color='r', linestyle='--', linewidth=1.5) + y_max = ax.get_ylim()[1] + median_text = f"median: {median_val:.2f} (Hz)\n N={single_rates.shape[0]}" + ax.text( x=median_val * 1.1, y=y_max * 0.7, s=median_text, color='red') + + #.... edge count + ax = self.plt.subplot(nrow,ncol,2) + ax.fill_between(x_vals, edgeCount, step='mid', color='salmon', alpha=0.7) + ax.set_xlabel(neurXlab) + ax.set_ylabel('num true edges') + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + probLo, probHi = dmd['edge_prob'] + rho_title = f'outgoing edges, true, prob=[{probLo:.2f}, {probHi:.2f}]' + ax.set_title(rho_title) + + + #.... firing rates ..... + ax = self.plt.subplot(nrow,ncol,4) + chanW=0.9 + + ax.bar(x_vals[exc_mask], single_rates[exc_mask], width=chanW, color='red', align='center', alpha=0.7, label='Excitatory') + ax.bar(x_vals[inh_mask], single_rates[inh_mask], width=chanW, color='blue', align='center', alpha=0.7, label='Inhibitory') + + ax.set_xlabel(neurXlab) + ax.set_ylabel('Firing rate (Hz)') + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + state_tag = md.get('sel_state', None) + if state_tag is None: + ax.set_title('Single Neurons') + else: + ax.set_title(f'Single Neurons, state={state_tag}') + ax.legend() + + ax = self.plt.subplot(nrow,ncol,5) + ax.scatter(nedge[exc_mask], sedge[exc_mask], s=16, marker='.', alpha=0.80, + color=exc_color, label=f'exc={n_exc}') + ax.scatter(nedge[inh_mask], sedge[inh_mask], s=16, marker='.', alpha=0.80, + color=inh_color, label=f'inh={n_inh}') + ax.axhline(0.0, color='tab:blue', ls='--', lw=0.8, alpha=0.9) + ax.set(title='True Nedge vs Sedge', + xlabel='Nedge (# outgoing edges per source column)', ylabel='Sedge') + ax.legend(loc='best', fontsize=9) + ax.grid(True, alpha=0.35) + + ax = self.plt.subplot(nrow,ncol,6) + ax.scatter(nedge[exc_mask], single_rates[exc_mask], s=16, marker='.', alpha=0.80, + color=exc_color, label=f'exc={n_exc}') + ax.scatter(nedge[inh_mask], single_rates[inh_mask], s=16, marker='.', alpha=0.80, + color=inh_color, label=f'inh={n_inh}') + ax.set(title='True Nedge vs frequency', + xlabel='Nedge (# outgoing edges per source column)', ylabel='frequency (Hz)') + ax.legend(loc='best', fontsize=9) + ax.grid(True, alpha=0.35) + + ax = self.plt.subplot(nrow,ncol,7) + ax.scatter(med_weight[exc_mask], single_rates[exc_mask], s=16, marker='.', alpha=0.80, + color=exc_color, label=f'exc={n_exc}') + ax.scatter(med_weight[inh_mask], single_rates[inh_mask], s=16, marker='.', alpha=0.80, + color=inh_color, label=f'inh={n_inh}') + ax.axvline(0.0, color='tab:blue', ls='--', lw=0.8, alpha=0.9) + ax.set(title='True median weight vs frequency', + xlabel='median outgoing edge weight', ylabel='frequency (Hz)') + ax.legend(loc='best', fontsize=9) + ax.grid(True, alpha=0.35) + + ax = self.plt.subplot(nrow,ncol,8) + bins = np.linspace(A_diag.min(), A_diag.max(), 80) + nonneg_counts = [] + for mask, color, label in [ + (exc_mask, exc_color, f'exc={n_exc}'), + (inh_mask, inh_color, f'inh={n_inh}'), + ]: + name = label.split('=')[0] + vals = A_diag[mask] + nonneg_counts.append((name, int(np.sum(vals >= 0.0)))) + ax.hist(vals, bins=bins, color=color, alpha=0.75, label=label, edgecolor=None) + ax.axvline(0.0, color='k', ls='--', lw=0.8) + ax.legend(loc='best', fontsize=9) + tbl = ax.table( + cellText=[[name, f'{cnt:d}'] for name, cnt in nonneg_counts], + colLabels=['type', '>=0'], + cellLoc='center', + colLoc='center', + bbox=[0.70, 0.46, 0.27, 0.28], + ) + tbl.auto_set_font_size(False) + tbl.set_fontsize(8) + for cell in tbl.get_celld().values(): + cell.set_edgecolor('0.65') + cell.set_linewidth(0.5) + cell.set_facecolor((1.0, 1.0, 1.0, 0.78)) + ax.set(title='True A diagonal', xlabel='A_true diagonal value', ylabel='neurons') + ax.grid(True, alpha=0.35) + + +#............................ +#............................ +#............................ + def rates_study(self,trueD,spikeD,md,figId=4): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,4)) + data_name=md['short_name'] + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel if R_sel is not None else '' + + B_idle = trueD['B_true'] + single_rates = spikeD['single_rates'] + + # Natural indexing: first numExc are excitatory, rest inhibitory + exc_mask = np.zeros(numNeur, dtype=bool) + exc_mask[:numExc] = True + inh_mask = ~exc_mask + + # 1) Scatter: x=B_idle, y=log(single_rates) + ax = self.plt.subplot(nrow,ncol,1) + ax.scatter(B_idle[exc_mask], single_rates[exc_mask], s=14, alpha=0.6, marker='^', facecolors='none', edgecolors='red', label='Excitatory') + ax.scatter(B_idle[inh_mask], single_rates[inh_mask], s=14, alpha=0.6, marker='^', facecolors='none', edgecolors='blue', label='Inhibitory') + ax.set_xlabel('true B_idle ') + ax.set_ylabel('single_rates (Hz)') + ax.set_yscale('log') + ax.grid(True, alpha=0.3) + tit='True Dale, N%d%s, %s'%(numNeur, R_tag, data_name) + ax.set_title(tit) + + # reference line: y = exp(B), clipped to central 80% of B range + bmin = float(np.min(B_idle)) + bmax = float(np.max(B_idle)) + bLo = bmin + 0.1 * (bmax - bmin) + bHi = bmax - 0.1 * (bmax - bmin) + bx = np.linspace(bLo, bHi, 100) + by = np.exp(bx) + ax.plot(bx, by, linestyle='--', color='black', linewidth=0.8, label='y=exp(x)') + ax.legend() + + # 2) Histogram: true B_idle (all neurons) + ax2 = self.plt.subplot(nrow,ncol,2) + b_bins = min(60, max(20, int(np.sqrt(numNeur) * 3))) + ax2.hist(B_idle, bins=b_bins, color='dimgray', alpha=0.8, edgecolor=None) + ax2.set_xlabel('true B_idle') + ax2.set_ylabel('num neurons') + ax2.grid(True, alpha=0.3) + state_tag = md.get('sel_state', None) + if state_tag is None: + ax2.set_title('true B_idle') + else: + ax2.set_title(f'true B_idle, state={state_tag}') + + # 3) Histogram: single_rates for excitatory + ax3 = self.plt.subplot(nrow,ncol,3) + exc_vals = single_rates[exc_mask] + inh_vals = single_rates[inh_mask] + # compute common bins and range (start at 0) + exc_max = np.max(exc_vals) if exc_vals.size > 0 else 0.0 + inh_max = np.max(inh_vals) if inh_vals.size > 0 else 0.0 + x_max = max(exc_max, inh_max) + if x_max <= 0: x_max = 1.0 + #common_bins = np.linspace(0, x_max, 21) + common_bins = np.linspace(0, x_max, int(2*x_max)) + ax3.hist(exc_vals, bins=common_bins, color='red', alpha=0.7, edgecolor=None) + ax3.set_xlabel('single_rates (Hz)') + ax3.set_ylabel('num neurons') + ax3.grid(True, alpha=0.3) + ax3.set_title('Rates: Excitatory (N=%d)' % (numExc)) + + # 4) Histogram: single_rates for inhibitory + ax4 = self.plt.subplot(nrow,ncol,4) + ax4.hist(inh_vals, bins=common_bins, color='blue', alpha=0.7, edgecolor=None) + ax4.set_xlabel('single_rates (Hz)') + ax4.set_ylabel('num neurons') + ax4.grid(True, alpha=0.3) + ax4.set_title('Rates: Inhibitory (N=%d)' % (numNeur-numExc)) + # unify x-range starting at 0 for both histograms + ax3.set_xlim(0, x_max) + ax4.set_xlim(0, x_max) diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/PlotterEdgeMeterAccuracy.py b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterEdgeMeterAccuracy.py new file mode 100644 index 00000000..28aae48e --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterEdgeMeterAccuracy.py @@ -0,0 +1,244 @@ +#!/usr/bin/env python3 +"""Plotting utilities for absolute graph-recovery metrics.""" + +import numpy as np + +from toolbox.PlotterBackbone import PlotterBackbone + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def figId2name(self, fid): + if isinstance(fid, str): + return f"{self.jobName}_{fid}" + return f"{self.jobName}_f{fid}" + + def smart_append(self, fig_id): + if fig_id in self.figL: + raise ValueError(f"Figure id {fig_id!r} is already in use") + self.figL.append(fig_id) + return fig_id + + def _x(self, metricD): + x = np.asarray(metricD["duration_min"], dtype=np.float64) + return x, "duration (min)" + + def _run_labels(self, metricD): + tags = np.asarray(metricD["fit_tag"], dtype=object) + return [str(x) for x in tags] + + def _plot_open_first(self, ax, x, y, marker, linestyle, color, label, lw=1.4, ms=4): + y = np.asarray(y) + ax.plot(x, y, marker + linestyle, lw=lw, ms=ms, label=label, color=color) + if len(x) > 0: + ax.plot( + [x[0]], [y[0]], marker=marker, linestyle="None", + ms=ms, mfc="white", mec=color, mew=1.5, color=color, + ) + + def recovery_summary(self, metricD, md, figId="a"): + """Main recovery curves and source-type diagnostics.""" + figId = self.smart_append(figId) + fig, axs = self.plt.subplots(2, 3, num=figId, facecolor="white", figsize=(15, 8)) + fig.subplots_adjust(hspace=0.38, wspace=0.32) + x, xlabel = self._x(metricD) + + ax = axs[0, 0] + for key, label, color in ( + ("precision", "precision", "tab:blue"), + ("recall", "recall", "tab:orange"), + ("f1", "F1", "tab:green"), + ("mcc", "MCC", "tab:red"), + ): + lw = 2.5 if key == "mcc" else 1.4 + ms = 6 if key == "mcc" else 4 + self._plot_open_first(ax, x, metricD[key], "o", "-", color, label, lw=lw, ms=ms) + ax.set(title="edge is present (metrics)", xlabel=xlabel, ylabel="metric") + ax.set_ylim(-0.05, 1.05) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + + ax = axs[0, 1] + for key, label, color in ( + ("signed_precision", "signed precision", "tab:blue"), + ("signed_recall", "signed recall", "tab:orange"), + ("signed_f1", "signed F1", "tab:green"), + ("signed_mcc", "signed MCC", "tab:red"), + ("sign_flip_rate", "sign flip rate", "tab:purple"), + ): + lw = 2.5 if key == "signed_mcc" else 1.4 + ms = 6 if key == "signed_mcc" else 4 + self._plot_open_first(ax, x, metricD[key], "o", "-", color, label, lw=lw, ms=ms) + ax.set(title="edge sign is correct (metrics)", xlabel=xlabel, ylabel="metric") + ax.set_ylim(-0.05, 1.05) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + + ax = axs[0, 2] + labels = self._run_labels(metricD) + xpos = np.arange(len(labels)) + bot = np.zeros(len(labels), dtype=np.float64) + for key, label, color in ( + ("num_reco_exc_neurons", "exc", "magenta"), + ("num_reco_inh_neurons", "inh", "limegreen"), + ("num_reco_und_neurons", "und", "salmon"), + ): + vals = np.asarray(metricD[key], dtype=np.float64) + ax.bar(xpos, vals, bottom=bot, color=color, alpha=0.75, label=label) + bot += vals + ax.set(title="Recovered source type", xlabel="fit tag", ylabel="neurons") + ax.set_xticks(xpos) + ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=8) + ax.grid(True, alpha=0.3, axis="y") + ax.legend(fontsize=8) + + ax = axs[1, 0] + self._plot_open_first( + ax, x, metricD["fdp"], "o", "-", "tab:red", + "realized FDP", lw=2.5, ms=6, + ) + self._plot_open_first( + ax, x, metricD["stability_bound_fdp"], "s", "--", "tab:purple", + "stability bound / selected", lw=1.1, ms=4, + ) + ax.set(title="False discoveries", xlabel=xlabel, ylabel="fraction") + ax.set_ylim(bottom=0.0) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + + ax = axs[1, 1] + self._plot_open_first(ax, x, metricD["density_true"], "o", "-", "black", "true", lw=1.4, ms=4) + self._plot_open_first(ax, x, metricD["density_reco"], "o", "-", "tab:cyan", "reco", lw=1.4, ms=4) + ax.set(title="Graph density", xlabel=xlabel, ylabel="edges / candidates") + ax.set_ylim(bottom=-0.002) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + + ax = axs[1, 2] + ax.plot(x, metricD["frac_reco_und_neurons"], "o-", lw=1.4, ms=4, + label="und neurons / N", color="tab:gray") + ax.plot(x, metricD["frac_reco_edges_from_reco_und_src"], "o-", lw=1.4, ms=4, + label="reco edges from und src", color="tab:purple") + ax.plot(x, metricD["frac_true_edges_from_reco_und_src"], "o-", lw=1.4, ms=4, + label="true edges from und src", color="tab:brown") + ax.set(title="Undecided-source diagnostics", xlabel=xlabel, ylabel="fraction") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + + fig.suptitle( + f"{md['fitNameTrunk']}, absolute edge-meter diagnostics for A_prune", + fontsize=13, + ) + + def _annotate_matrix(self, ax, mat): + title_fs = self.plt.rcParams.get("axes.titlesize", 12) + for i in range(mat.shape[0]): + for j in range(mat.shape[1]): + ax.text( + j, i, f"{int(mat[i, j])}", + ha="center", va="center", fontsize=title_fs, + color="magenta", + ) + + def count_summary(self, metricD, md, figId="b"): + """Raw presence and signed count bars.""" + figId = self.smart_append(figId) + fig, axs = self.plt.subplots(1, 2, num=figId, facecolor="white", figsize=(13, 5)) + fig.subplots_adjust(wspace=0.28) + labels = self._run_labels(metricD) + xpos = np.arange(len(labels)) + + ax = axs[0] + bot = np.zeros(len(labels), dtype=np.float64) + for key, label, color in ( + ("TP", "TP", "tab:green"), + ("FN", "FN", "tab:orange"), + ("FP", "FP", "tab:red"), + ): + vals = np.asarray(metricD[key], dtype=np.float64) + ax.bar(xpos, vals, bottom=bot, label=label, color=color, alpha=0.75) + bot += vals + ax.set(title="edge is present (counts)", xlabel="fit tag", ylabel="off-diagonal pairs") + ax.set_xticks(xpos) + ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=8) + ax.grid(True, alpha=0.3, axis="y") + ax.legend(fontsize=8) + + ax = axs[1] + bot = np.zeros(len(labels), dtype=np.float64) + for key, label, color in ( + ("signed_TP", "signed TP", "tab:green"), + ("signed_FN", "signed FN", "tab:orange"), + ("signed_FP", "signed FP", "tab:red"), + ): + vals = np.asarray(metricD[key], dtype=np.float64) + ax.bar(xpos, vals, bottom=bot, label=label, color=color, alpha=0.75) + bot += vals + ax.set(title="edge sign is correct (counts)", xlabel="fit tag", ylabel="edges") + ax.set_xticks(xpos) + ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=8) + ax.grid(True, alpha=0.3, axis="y") + ax.legend(fontsize=8) + + fig.suptitle( + f"{md['fitNameTrunk']}, absolute edge-meter raw counts", + fontsize=13, + ) + + def _resolve_tag_indices(self, metricD, tagIdxL): + n_tag = len(metricD["fit_tag"]) + idx_l = [] + for idx in tagIdxL: + idx = int(idx) + if idx < 0: + idx = n_tag + idx + if idx < 0 or idx >= n_tag: + raise IndexError(f"tagIdx {idx} outside 0..{n_tag - 1}") + idx_l.append(idx) + return idx_l + + def confusion_summary(self, metricD, md, tagIdxL=None, figId="c"): + """Confusion matrices for selected fit tags.""" + if tagIdxL is None: + tagIdxL = [-1] + idx_l = self._resolve_tag_indices(metricD, tagIdxL) + labels = self._run_labels(metricD) + + figId = self.smart_append(figId) + fig, axs = self.plt.subplots( + len(idx_l), 2, num=figId, facecolor="white", + figsize=(12, 4.5 * len(idx_l)), squeeze=False, + ) + fig.subplots_adjust(hspace=0.42, wspace=0.38) + + for row, tag_idx in enumerate(idx_l): + tag = labels[tag_idx] + + ax = axs[row, 0] + cm = np.asarray(metricD["confusion3"][tag_idx], dtype=np.int64) + im = ax.imshow(cm, cmap="Blues") + self._annotate_matrix(ax, cm) + ax.set(title=f"Edge-label confusion, tag={tag}", xlabel="reco label", ylabel="true label") + ax.set_xticks(np.arange(3)) + ax.set_yticks(np.arange(3)) + ax.set_xticklabels(["-", "0", "+"]) + ax.set_yticklabels(["-", "0", "+"]) + fig.colorbar(im, ax=ax, shrink=0.82) + + ax = axs[row, 1] + scm = np.asarray(metricD["source_type_confusion"][tag_idx], dtype=np.int64) + im = ax.imshow(scm, cmap="Greens") + self._annotate_matrix(ax, scm) + ax.set(title=f"Neuron type confusion, tag={tag}", xlabel="reco source", ylabel="true source") + ax.set_xticks(np.arange(3)) + ax.set_yticks(np.arange(2)) + ax.set_xticklabels(["exc", "inh", "und"]) + ax.set_yticklabels(["exc", "inh"]) + fig.colorbar(im, ax=ax, shrink=0.82) + + fig.suptitle( + f"{md['fitNameTrunk']}, selected confusion matrices", + fontsize=13, + ) diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/PlotterEdgeMeterFidelity.py b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterEdgeMeterFidelity.py new file mode 100644 index 00000000..e8655cfc --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterEdgeMeterFidelity.py @@ -0,0 +1,219 @@ +#!/usr/bin/env python3 +"""Plotting utilities for PRISM-FDR graph stability metrics (edgeMaterFidelity3c).""" + +import numpy as np + +from toolbox.PlotterBackbone import PlotterBackbone + +SCOPES = ["all", "exc", "inh"] +SCOPE_TITLES = {"all": "all edges", "exc": "excitatory-source edges", "inh": "inhibitory-source edges"} +SCOPE_COLORS = {"all": "tab:blue", "exc": "tab:red", "inh": "tab:green"} + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def figId2name(self, fid): + if isinstance(fid, str): + return f"{self.jobName}_{fid}" + return f"{self.jobName}_f{fid}" + + def smart_append(self, fig_id): + if fig_id in self.figL: + raise ValueError(f"Figure id {fig_id!r} already in use") + self.figL.append(fig_id) + return fig_id + + # ------------------------------------------------------------------ + # shared helpers + # ------------------------------------------------------------------ + + def _x_and_label(self, out_d): + x = np.asarray(out_d["x"], dtype=np.float64) + mode = self._mode_str(out_d) + if np.all(np.isfinite(x)): + xlabel = "duration (min)" if mode == "last" else "duration midpoint (min)" + return x, xlabel + return np.arange(len(x), dtype=np.float64), "comparison index" + + def _subset_x_and_label(self, out_d): + x = np.asarray(out_d["subset_x"], dtype=np.float64) + if np.all(np.isfinite(x)): + return x, "duration (min)" + return np.arange(len(x), dtype=np.float64), "data subset index" + + def _xtick_labels(self, out_d): + return [str(v) for v in np.asarray(out_d["labels"], dtype=object)] + + def _mode_str(self, out_d): + return str(np.asarray(out_d["compare_mode"]).ravel()[0]) + + def _suptitle(self, out_md, out_d, metric_name): + mode = out_md.get("compare_mode", "?") + trunk = out_md.get("fitNameTrunk", "") + if mode == "last": + ref_tag = str(np.asarray(out_d["subset_tags"]).ravel()[-1]) + mode_str = f"last={ref_tag}" + else: + mode_str = mode + return f"{trunk} — {metric_name} [{mode_str}]" + + def _apply_int_xticks(self, ax, x): + """Format x-axis ticks as integers (no decimal point).""" + ax.xaxis.set_major_formatter(self.plt.FuncFormatter(lambda v, _: f"{int(round(v))}")) + ax.set_xticks(x) + + def _bin_edges_from_centers(self, x): + x = np.asarray(x, dtype=np.float64) + if x.size == 1: + return np.array([x[0] - 0.5, x[0] + 0.5], dtype=np.float64) + edges = np.empty(x.size + 1, dtype=np.float64) + edges[1:-1] = 0.5 * (x[:-1] + x[1:]) + edges[0] = x[0] - 0.5 * (x[1] - x[0]) + edges[-1] = x[-1] + 0.5 * (x[-1] - x[-2]) + return edges + + def _plot_scope_rows(self, fig, axs, x, xlabel, key, ylabel, ylim, out_d, out_md): + """Fill a 3-row axes array, one row per scope, for a single scalar metric key.""" + for row, scope in enumerate(SCOPES): + ax = axs[row] + y = np.asarray(out_d[f"{scope}_{key}"], dtype=np.float64) + color = SCOPE_COLORS[scope] + ax.plot(x, y, "o-", lw=1.6, ms=5, color=color) + ax.set(title=SCOPE_TITLES[scope], xlabel=xlabel, ylabel=ylabel) + if ylim is not None: + ax.set_ylim(ylim) + ax.grid(True, alpha=0.3) + + # ------------------------------------------------------------------ + # plot group a — compact topology/sign/Spearman summary (3 rows x 2 cols) + # ------------------------------------------------------------------ + + def plot_jaccard_smr(self, out_d, out_md, figId="a"): + figId = self.smart_append(figId) + fig, axs = self.plt.subplots(3, 2, num=figId, facecolor="white", figsize=(7.2, 6.6)) + fig.subplots_adjust(hspace=0.45, wspace=0.28) + x, xlabel = self._x_and_label(out_d) + + ax = axs[0, 0] + for key, label, color in ( + ("all_J_pos", "J_exc", "tab:red"), + ("all_J_neg", "J_inh", "tab:blue"), + ): + y = np.asarray(out_d[key], dtype=np.float64) + ax.plot(x, y, "o-", lw=1.6, ms=5, color=color, label=label) + ax.set(title="Jaccard - all edges", xlabel=xlabel, ylabel="Jaccard index") + ax.set_ylim(-0.02, 1.05) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=9) + self._apply_int_xticks(ax, x) + + ax = axs[0, 1] + y = np.asarray(out_d["all_SMR"], dtype=np.float64) + ax.plot(x, y, "o-", lw=1.6, ms=5, color=SCOPE_COLORS["all"]) + ax.set(title="Signal Match Rate - all edges", xlabel=xlabel, ylabel="SMR") + ax.grid(True, alpha=0.3) + self._apply_int_xticks(ax, x) + + ax = axs[1, 1] + y = np.asarray(out_d["diag_mae"], dtype=np.float64) + ax.plot(x, y, "o-", lw=1.6, ms=5, color="tab:purple") + ax.set(title="Mean |Delta A_ii|", xlabel=xlabel, ylabel="mean abs diagonal diff") + ax.grid(True, alpha=0.3) + self._apply_int_xticks(ax, x) + + axs[2, 1].set_axis_off() + + for row, key, title, ylabel in ( + (1, "r_spearman", "Spearman r - source types", "Spearman r"), + (2, "r_spearman_w", "Weighted Spearman r - source types", "weighted Spearman r"), + ): + ax = axs[row, 0] + for scope, label, color in (("exc", "exc", "tab:red"), ("inh", "inh", "tab:blue")): + y = np.asarray(out_d[f"{scope}_{key}"], dtype=np.float64) + ax.plot(x, y, "o-", lw=1.6, ms=5, color=color, label=label) + ax.set(title=title, xlabel=xlabel, ylabel=ylabel) + ax.set_ylim(-0.02, 1.05) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + self._apply_int_xticks(ax, x) + + fig.suptitle(self._suptitle(out_md, out_d, "Edge Fidelity Summary"), fontsize=13) + + # ------------------------------------------------------------------ + # plot group b — weight thresholds min_posW / max_negW vs duration + # ------------------------------------------------------------------ + + def plot_weight_thresholds(self, out_d, out_md, figId="b"): + figId = self.smart_append(figId) + fig, axs = self.plt.subplots(1, 4, num=figId, facecolor="white", figsize=(16, 4)) + fig.subplots_adjust(wspace=0.38, left=0.06, right=0.97, top=0.88) + + sx, sxlabel = self._subset_x_and_label(out_d) + min_posW = np.asarray(out_d["min_posW"], dtype=np.float64) + max_negW = np.asarray(out_d["max_negW"], dtype=np.float64) + + ax = axs[0] + ax.plot(sx, min_posW, "o-", lw=1.6, ms=5, color="tab:red", label="min_posW") + ax.plot(sx, max_negW, "s-", lw=1.6, ms=5, color="tab:green", label="max_negW") + ax.axhline(0.0, color="k", ls="--", lw=0.8) + ax.set(title="Weight thresholds vs duration", xlabel=sxlabel, ylabel="weight") + ax.legend(fontsize=9) + ax.grid(True, alpha=0.3) + self._apply_int_xticks(ax, sx) + + ax = axs[1] + for key, color, label in [ + ("exc", "magenta", "exc"), + ("inh", "#39FF14", "inh"), + ]: + lo = np.asarray(out_d[f"w_min_{key}"], dtype=np.float64) + hi = np.asarray(out_d[f"w_max_{key}"], dtype=np.float64) + ax.fill_between(sx, lo, hi, alpha=0.35, color=color, label=label) + ax.plot(sx, lo, "-", lw=0.8, color=color) + ax.plot(sx, hi, "-", lw=0.8, color=color) + if key in ("exc", "inh"): + med = np.asarray(out_d[f"w_med_{key}"], dtype=np.float64) + ax.plot(sx, med, "--", lw=1.2, color="tab:blue") + q4_idx = len(sx) * 3 // 4 + tx = sx[q4_idx] + ty = med[q4_idx] + va = "top" if key == "inh" else "bottom" + ax.text(tx, ty, f"med {label}", fontsize=8, color="tab:blue", + va=va, ha="center") + ax.axhline(0.0, color="k", ls="--", lw=0.8) + ax.set(title="Off-diag weight range by type", xlabel=sxlabel, ylabel="weight") + ax.legend(fontsize=9) + ax.grid(True, alpha=0.3) + self._apply_int_xticks(ax, sx) + + ax = axs[2] + for key, color, label in [ + ("exc", "magenta", "exc"), + ("inh", "#39FF14", "inh"), + ("und", "salmon", "und"), + ]: + n = np.asarray(out_d[f"n_edges_{key}"], dtype=np.float64) + ax.plot(sx, n, "o-", lw=1.6, ms=5, color=color, label=label) + ax.set(title="# off-diag edges by type", xlabel=sxlabel, ylabel="# edges") + ax.legend(fontsize=9) + ax.grid(True, alpha=0.3) + self._apply_int_xticks(ax, sx) + + ax = axs[3] + rate_edges = np.asarray(out_d["rate_bin_edges"], dtype=np.float64) + edge_rate_hist = np.asarray(out_d["edge_count_rate_hist"], dtype=np.float64) + x_edges = self._bin_edges_from_centers(sx) + from matplotlib.colors import LinearSegmentedColormap + cmap = LinearSegmentedColormap.from_list( + "YlOrBl", ["#fff7bc", "#fec44f", "#d95f0e", "#2b8cbe", "#08519c"] + ) + cmap.set_bad("white") + plot_hist = np.ma.masked_where(edge_rate_hist <= 0.0, edge_rate_hist) + pcm = ax.pcolormesh(x_edges, rate_edges, plot_hist, shading="auto", cmap=cmap) + fig.colorbar(pcm, ax=ax, label="num edges") + ax.set(title="recovered exc+inh off-diag edges", xlabel=sxlabel, ylabel="single rate (Hz)") + self._apply_int_xticks(ax, sx) + + fig.suptitle(self._suptitle(out_md, out_d, "Weight Stats"), fontsize=13) diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/PlotterPrismEM3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterPrismEM3c.py new file mode 100755 index 00000000..82499508 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterPrismEM3c.py @@ -0,0 +1,1837 @@ +#!/usr/bin/env python3 +"""Plotting utilities for prism EM 3c evaluation.""" + +import numpy as np + +from toolbox.PlotterBackbone import PlotterBackbone +from UtilBioExp import clip_rebD_time + + +def real_fit_metadata(md): + """Return the real-fit metadata block for ordinary or FDR-aggregated fits.""" + if is_stage_c_metadata(md): + return md + if "bagsFDR_stageA" in md and "real_fit" in md["bagsFDR_stageA"]: + return md["bagsFDR_stageA"]["real_fit"] + return md + + +def is_stage_c_metadata(md): + return md.get("fit_type") == "prismEM_deBias_stageC" or "deBias_stageC" in md + + +def fit_stage_label(md): + if is_stage_c_metadata(md): + mode = real_fit_metadata(md).get("train", {}).get("state_mode", "?") + return f"Stage (c) deBias, state_mode={mode}" + if "bagsFDR_stageB" in md: + return "Stage (b) FDR aggregate" + if md.get("fit_type", "").startswith("prismEM_FDRbag") or "bagsFDR_stageA" in md: + return "Stage (a) FDR bag" + return "PRISM-EM" + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def figId2name(self, fid): + if isinstance(fid, str): + return f"{self.jobName}_{fid}" + return f"{self.jobName}_f{fid}" + + def smart_append(self, fig_id): + if fig_id in self.figL: + raise ValueError(f"Figure id {fig_id!r} is already in use") + self.figL.append(fig_id) + return fig_id + + def canvas_title_prefix(self, md): + data_label = "exper" if md.get("data_type") == "bioExp" else "simu" + return f"{data_label} {md['short_name']}, {fit_stage_label(md)}" + + def _display_A_label(self, md, prune=False): + if is_stage_c_metadata(md): + return "A_debias" + return "A_prune" if prune else "A_hat" + + def _display_B_label(self, md): + return "B_debias" if is_stage_c_metadata(md) else "B_hat" + + def summary(self, fitD, md, figId="a"): + """EM convergence overview.""" + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(12, 6)) + fig.subplots_adjust(hspace=0.45, wspace=0.35) + + trainMD = real_fit_metadata(md)["train"] + title_prefix = self.canvas_title_prefix(md) + is_stage_c = is_stage_c_metadata(md) + + e_nll = np.asarray(fitD["e_nll_em"]) + m_nll = np.asarray(fitD["m_nll_epoch"]) + m_l1 = np.asarray(fitD["m_l1_epoch"]) + rho = np.asarray(fitD["rho_epoch"]) + nz = np.asarray(fitD["nz_edges_epoch"]) + lr = np.asarray(fitD["learning_rates"]) + + n_em = int(trainMD.get("num_em_iters", 0)) + m_per_em = int(trainMD["m_epochs"]) + n_m_total = len(m_nll) + m_epochs = np.arange(1, n_m_total + 1) + em_iters = np.arange(1, len(e_nll) + 1) + is_locked_fdr = n_em < 1 or md.get("fit_type", "").startswith("prismEM_FDRbags") + + prune_em = int(trainMD["delay_em_iter_4_Aprune"]) + rho_start_em = int(trainMD["delay_em_iter_4_ArhoMax"]) + lr_drop_em = int(trainMD["delay_em_iter_4_lrDecay"]) + prune_m_epoch = prune_em * m_per_em + rho_start_m_epoch = rho_start_em * m_per_em + lr_drop_m_epoch = lr_drop_em * m_per_em + + def draw_threshold_marker(ax, x_pos, x_max, txt, color, yFac=0.02): + if not (0 < x_pos <= x_max): + return + ax.axvline(x_pos, color=color, ls="--", lw=0.9, alpha=0.95) + y0, y1 = ax.get_ylim() + x0, x1 = ax.get_xlim() + x_off = 0.01 * max(1e-9, x1 - x0) + y_txt = y1 - yFac * (y1 - y0) + ax.text( + x_pos + x_off, y_txt, txt, rotation=90, color=color, fontsize=7, + ha="left", va="top", + bbox=dict(facecolor="white", alpha=0.55, edgecolor="none", pad=0.2), + ) + + ax = self.plt.subplot(2, 3, 1) + if len(e_nll) > 0: + ax.plot(em_iters, e_nll, "o-", color="tab:blue", markersize=3, linewidth=1.2) + ax.set(title="E-step NLL", xlabel="EM iteration", ylabel="NLL / bin") + else: + if is_stage_c: + txt = ( + "Stage (c) de-biased fit\n" + f"state_mode={trainMD.get('state_mode', '?')}\n" + "display: A_debias, B_debias" + ) + else: + txt = "FDR locked M-step fit\nstate labels from reference EM\nno bag E-step" + if "bagsFDR_stageB" in md: + stg = md["bagsFDR_stageB"] + txt += f"\nbags={stg.get('num_bags', '?')} stab_sel_thresh={stg.get('stab_sel_thresh', '?')}" + ax.text(0.5, 0.55, txt, transform=ax.transAxes, ha="center", va="center", + fontsize=10, bbox=dict(facecolor="white", alpha=0.75, edgecolor="0.8")) + ax.set(title="Fit mode", xlabel="", ylabel="") + ax.set_xticks([]) + ax.set_yticks([]) + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_em, len(e_nll), "start Aprune", "k") + draw_threshold_marker(ax, rho_start_em, len(e_nll), "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_em, len(e_nll), "start_lrDrop", "tab:gray", yFac=0.6) + txt = (f"pgd_iter={trainMD['pgd_iter']}\n" + f"lr_E={trainMD['lr_estep']}\n" + f"lambda2={trainMD['lambda2']}") + ax.text(0.97, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=8, + bbox=dict(facecolor="white", alpha=0.7, edgecolor="none")) + + ax = self.plt.subplot(2, 3, 2) + ax.plot(m_epochs, m_nll, color="tab:blue", linewidth=1, label="NLL") + ax.set_ylabel("NLL", color="tab:blue") + ax.tick_params(axis="y", labelcolor="tab:blue") + ax.set(title="M-step loss", xlabel="M-epoch (global)") + ax.grid(True, alpha=0.3) + ax2 = ax.twinx() + ax2.plot(m_epochs, m_l1, color="tab:red", linewidth=1, linestyle="--", label="L1") + ax2.set_ylabel("L1", color="tab:red") + ax2.tick_params(axis="y", labelcolor="tab:red") + lines1, lab1 = ax.get_legend_handles_labels() + lines2, lab2 = ax2.get_legend_handles_labels() + ax.legend(lines1 + lines2, lab1 + lab2, fontsize=7, loc="upper right") + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + txt = (f"lr_M={trainMD['lr_mstep']}\n" + f"L1 lambda3={trainMD['lambda3']}\n" + f"batch={trainMD['batch_size']}") + ax.text(0.03, 0.03, txt, transform=ax.transAxes, + va="bottom", ha="left", fontsize=8, + bbox=dict(facecolor="white", alpha=0.7, edgecolor="none")) + + ax = self.plt.subplot(2, 3, 3) + rho_max = float(trainMD["rho_max"]) + ax.plot(m_epochs, rho, color="tab:green", linewidth=1, label="rho(A)") + ax.axhline(rho_max, color="red", ls="--", lw=1, label=f"rho_max={rho_max:g}") + ax.set(title="Spectral radius rho(A)", xlabel="M-epoch (global)", ylabel="rho") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + ax.legend(fontsize=8) + + ax = self.plt.subplot(2, 3, 4) + ax.plot(m_epochs, nz, color="tab:purple", linewidth=1) + ax.set(title="Non-zero off-diag edges", xlabel="M-epoch (global)", ylabel="count") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + ax2 = ax.twinx() + ax2.plot(m_epochs, lr, color="tab:orange", linewidth=0.8, linestyle="--", alpha=0.6) + ax2.set_ylabel("learning rate", color="tab:orange") + ax2.tick_params(axis="y", labelcolor="tab:orange") + + ax = self.plt.subplot(2, 3, 5) + c_hat = np.asarray(fitD["c_hat"]) + occ = c_hat.mean(axis=0) + n_state = len(occ) + ax.bar(np.arange(n_state), occ, color="tab:blue", alpha=0.7) + ax.set(title="Mean occupancy", xlabel="state", ylabel="mean c") + ax.set_ylim(0.0, 1.0) + ax.set_xticks(np.arange(n_state)) + ax.grid(True, alpha=0.3, axis="y") + t0s, t1s = trainMD["time_range_sec"] + txt = (f"N={trainMD['num_neurons']} M={n_state}\n" + f"T=[{t0s:.0f},{t1s:.0f}]s\n" + f"bins={trainMD['num_time_bins']}") + ax.text(0.97, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=8, + bbox=dict(facecolor="white", alpha=0.7, edgecolor="none")) + + ax = self.plt.subplot(2, 3, 6) + A_fit = np.asarray(fitD["A_hat"]) + diag_mask = np.eye(A_fit.shape[0], dtype=bool) + A_off = A_fit[~diag_mask] + A_off_nz = A_off[np.abs(A_off) > 1e-10] + ax.hist(A_off_nz, bins=100, color="g", alpha=0.8) + ax.set_yscale("log") + ax.set(title=f"A off-diagonal, {A_off_nz.size} edges", + xlabel="edge value", ylabel="edges") + ax.grid(True, alpha=0.3) + + if is_stage_c: + title_tail = ( + f"Stage (c) deBias: state_mode={trainMD.get('state_mode', '?')}, " + f"K_epoch={n_m_total}" + ) + elif is_locked_fdr: + title_tail = f"FDR locked M-step: K_epoch={n_m_total}" + else: + title_tail = f"Prism EM: K_EM={n_em} x K_M={m_per_em}" + fig.suptitle(f"{title_prefix}, {title_tail}", fontsize=12) + fig.tight_layout(rect=[0, 0, 1, 0.95]) + + def _draw_A_matrix(self, ax, A, title, num_exc=None, norm_map=None): + import matplotlib.colors as colors + + A = np.asarray(A) + if norm_map is None: + vmin = float(np.min(A)) + vmax = float(np.max(A)) + if np.isclose(vmin, vmax): + vmax = vmin + 1e-9 + norm_map = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + im = ax.imshow(A, aspect=1.0, origin="lower", cmap="bwr", + norm=norm_map, interpolation="nearest") + n_neuron = A.shape[0] + if num_exc is not None: + n_exc = int(num_exc) + n_inh = n_neuron - n_exc + ax.axhline(n_exc - 0.5, color="k", ls="--", lw=0.8) + ax.axvline(n_exc - 0.5, color="k", ls="--", lw=0.8) + ax.text( + 0.03, 0.97, f"exc={n_exc}\ninh={n_inh}", + transform=ax.transAxes, va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=0.3), + ) + ax.plot([0, n_neuron], [0, n_neuron], "--", lw=0.8, color="magenta") + ax.set_xlim(-0.5, n_neuron + 0.5) + ax.set_ylim(-0.5, n_neuron + 0.5) + ax.grid(True, alpha=0.25) + ax.set_title(title) + ax.set_xlabel("source / presynaptic neuron index (column)") + ax.set_ylabel("target / postsynaptic neuron index (row)") + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + def _count_A_edges(self, A): + A = np.asarray(A) + n_neuron = A.shape[0] + off_mask = ~np.eye(n_neuron, dtype=bool) + n_diag = int(np.eye(n_neuron, dtype=bool).sum()) + n_off = int(np.count_nonzero(A[off_mask])) + return n_neuron, n_diag, n_off + + def _A_prune_from_fitD(self, fitD): + return np.asarray(fitD["A_prune"], dtype=np.float64) + + def _source_nz_edge_stats(self, A): + A = np.asarray(A, dtype=np.float64) + n_neuron = A.shape[0] + strong = ~np.eye(n_neuron, dtype=bool) & (A != 0) + cnt = strong.sum(axis=0).astype(np.int64) + sum_src = np.where(strong, A, 0.0).sum(axis=0) + return cnt, sum_src + + def _node_outgoing_edge_stats(self, A): + return self._source_nz_edge_stats(A) + + def _corrcoef_safe(self, x, y): + x = np.asarray(x) + y = np.asarray(y) + if x.size == 0 or y.size == 0: + return np.nan + if np.std(x) == 0 or np.std(y) == 0: + return np.nan + return float(np.corrcoef(x, y)[0, 1]) + + def _add_x45_line(self, ax): + lo = float(min(ax.get_xlim()[0], ax.get_ylim()[0])) + hi = float(max(ax.get_xlim()[1], ax.get_ylim()[1])) + ax.plot([lo, hi], [lo, hi], "--", color="k", linewidth=0.8, alpha=0.75) + ax.set_xlim(lo, hi) + ax.set_ylim(lo, hi) + + def _find_bimodal_divider(self, x_vals): + x = np.asarray(x_vals, dtype=np.float64).ravel() + x = x[np.isfinite(x)] + assert x.size >= 2, "Need at least two finite values for B split" + if np.std(x) == 0: + return float(np.median(x)) + c1, c2 = [float(v) for v in np.quantile(x, [0.25, 0.75])] + if c1 == c2: + c1, c2 = float(np.min(x)), float(np.max(x)) + for _ in range(32): + left = np.abs(x - c1) <= np.abs(x - c2) + if int(left.sum()) == 0 or int((~left).sum()) == 0: + return float(np.median(x)) + c1_new = float(np.mean(x[left])) + c2_new = float(np.mean(x[~left])) + if abs(c1_new - c1) < 1e-10 and abs(c2_new - c2) < 1e-10: + c1, c2 = c1_new, c2_new + break + c1, c2 = c1_new, c2_new + if c1 > c2: + c1, c2 = c2, c1 + return float(0.5 * (c1 + c2)) + + def _plot_corr_A_regions(self, ax, A_true, A_hat, y_label="A_hat"): + x = np.asarray(A_true, dtype=np.float64).ravel() + y = np.asarray(A_hat, dtype=np.float64).ravel() + assert x.shape == y.shape, f"A_true/A_hat shape mismatch after ravel: {x.shape} vs {y.shape}" + ax.scatter(x, y, s=6, alpha=0.35, color="green", edgecolors="none") + self._add_x45_line(ax) + ax.axhline(0.0, color="k", linestyle="--", linewidth=0.7, alpha=0.55) + ax.axvline(0.0, color="k", linestyle="--", linewidth=0.7, alpha=0.55) + ax.set(title="A fit", xlabel="A_true", ylabel=y_label) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.35) + + for mask, tag, tx, ty in [ + (x < 0.0, "L", 0.08, 0.10), + (x == 0.0, "M", 0.43, 0.42), + (x > 0.0, "R", 0.67, 0.66), + ]: + r = self._corrcoef_safe(x[mask], y[mask]) + n = int(mask.sum()) + ax.text(tx, ty, f"r{tag}={r:.3f}\nn{tag}={n}", transform=ax.transAxes, fontsize=9) + + def _scatter_true_vs_fit(self, ax, x_true, y_fit, title, color, split_by_true_sign=False): + x = np.asarray(x_true, dtype=np.float64).ravel() + y = np.asarray(y_fit, dtype=np.float64).ravel() + n = min(x.size, y.size) + x = x[:n] + y = y[:n] + ax.scatter(x, y, alpha=0.7, color=color, marker=".", s=9) + if n > 0: + lo = float(min(np.min(x), np.min(y))) + hi = float(max(np.max(x), np.max(y))) + pad = 0.05 * max(1e-9, hi - lo) + lo -= pad + hi += pad + ax.plot([lo, hi], [lo, hi], color="k", linestyle="--", linewidth=1.0, alpha=0.6) + ax.axhline(0.0, color="k", linestyle="--", linewidth=0.8, alpha=0.6) + ax.axvline(0.0, color="k", linestyle="--", linewidth=0.8, alpha=0.6) + ax.set_xlim(lo, hi) + ax.set_ylim(lo, hi) + if split_by_true_sign: + for mask in (x < 0.0, x > 0.0): + if np.any(mask): + ax.plot( + [float(np.mean(x[mask]))], [float(np.mean(y[mask]))], + marker="+", markersize=14, markeredgewidth=2.5, color="k", + ) + else: + ax.plot( + [float(np.mean(x))], [float(np.mean(y))], + marker="+", markersize=14, markeredgewidth=2.5, color="k", + ) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.4) + ax.set_title(title) + ax.set_xlabel("true weight") + ax.set_ylabel("fitted") + + def _plot_corr_B_divisor(self, ax, B_true, B_hat, state_idx, y_label="B_hat"): + x = np.asarray(B_true, dtype=np.float64).ravel() + y = np.asarray(B_hat, dtype=np.float64).ravel() + assert x.shape == y.shape, f"B_true/B_hat state {state_idx} shape mismatch: {x.shape} vs {y.shape}" + divider = self._find_bimodal_divider(x) + ax.scatter(x, y, s=8, alpha=0.45, color="tab:blue", edgecolors="none") + self._add_x45_line(ax) + ax.axvline(divider, color="red", linestyle="--", linewidth=1.0) + ax.set(title=f"B fit, state {state_idx}", xlabel="B_true", ylabel=y_label) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.35) + + left = x < divider + right = ~left + r_left = self._corrcoef_safe(x[left], y[left]) + r_right = self._corrcoef_safe(x[right], y[right]) + ax.text(0.05, 0.45, f"rL={r_left:.3f}\nnL={int(left.sum())}", transform=ax.transAxes, va="top", fontsize=9) + ax.text(0.64, 0.58, f"rR={r_right:.3f}\nnR={int(right.sum())}", transform=ax.transAxes, va="top", fontsize=9) + + def _edge_recovery_stats_ax(self, ax, fitD, md): + A_true = np.asarray(md["A_true"]) + if A_true.ndim == 3: + A_true = A_true[0] + A_hat = np.asarray(fitD["A_hat"]) + assert A_hat.shape == A_true.shape, f"A_hat shape {A_hat.shape} != A_true {A_true.shape}" + + E_true = np.asarray(md["E_true"]) != 0 + if E_true.ndim == 3: + E_true = E_true[0] != 0 + assert E_true.shape == A_true.shape, f"E_true shape {E_true.shape} != A_true {A_true.shape}" + + n_neur = A_true.shape[0] + off_diag = ~np.eye(n_neur, dtype=bool) + E_t = E_true & off_diag + E_hat = (A_hat != 0) & off_diag + tp = int(np.sum(E_t & E_hat)) + fp = int(np.sum((~E_t) & E_hat)) + fn = int(np.sum(E_t & (~E_hat))) + tn = int(np.sum((~E_t) & (~E_hat))) + precision = tp / max(1, tp + fp) + recall = tp / max(1, tp + fn) + f1 = 2.0 * precision * recall / max(1e-12, precision + recall) + acc = (tp + tn) / max(1, tp + tn + fp + fn) + + vals = [tp, fp, fn] + ax.bar(["TP", "FP", "FN"], vals, color=["green", "red", "magenta"]) + ax.set(title="edge stats", ylabel="count") + ax.grid(axis="y", alpha=0.4) + y_pos = 0.25 * max(1, max(vals)) + for name, val in zip(["TP", "FP", "FN"], vals): + ax.text(name, y_pos, str(val), ha="center", va="center", fontsize=10) + txt = f"precision={precision:.3f}\nrecall={recall:.3f}\nf1={f1:.3f}\nacc={acc:.3f}" + ax.text(0.50, 0.62, txt, transform=ax.transAxes, fontsize=9) + + def edge_stats_and_ABcorr(self, fitD, md, figId="n"): + """Edge-recovery stats plus A/B truth-vs-fit correlations for simulations.""" + A_true = np.asarray(md["A_true"]) + if A_true.ndim == 3: + A_true = A_true[0] + A_hat = np.asarray(fitD["A_hat"]) + B_true = np.asarray(md["B_true"]) + B_hat = np.asarray(fitD["B_hat"]) + A_label = self._display_A_label(md) + B_label = self._display_B_label(md) + if B_true.ndim == 1: + B_true = B_true[None, :] + if B_hat.ndim == 1: + B_hat = B_hat[None, :] + n_state = min(B_true.shape[0], B_hat.shape[0]) + assert n_state >= 1, "Need at least one B state for correlation plot" + + ncol = 2 + n_state + fig_w = max(13.0, 3.1 * ncol) + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(fig_w, 3.9)) + gs = fig.add_gridspec(1, ncol, left=0.045, right=0.99, bottom=0.18, top=0.80, wspace=0.38) + + ax = fig.add_subplot(gs[0, 0]) + self._edge_recovery_stats_ax(ax, fitD, md) + + ax = fig.add_subplot(gs[0, 1]) + self._plot_corr_A_regions(ax, A_true, A_hat, y_label=A_label) + + for m in range(n_state): + ax = fig.add_subplot(gs[0, 2 + m]) + self._plot_corr_B_divisor(ax, B_true[m], B_hat[m], m, y_label=B_label) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle(f"{title_prefix}, edge recovery and A/B correlations", fontsize=12) + + def matrix_init(self, fitD, md, figId="o", est_key="A_init", est_label="A_init"): + """Compare a fitted A-like matrix against stored simulation A_true.""" + import matplotlib.colors as colors + + A_true = np.asarray(md["A_true"], dtype=np.float64) + if A_true.ndim == 3: + A_true = A_true[0] + A_cmp = np.asarray(fitD[est_key], dtype=np.float64) + assert A_cmp.shape == A_true.shape, f"{est_key} shape {A_cmp.shape} != A_true {A_true.shape}" + n_neur = A_true.shape[0] + num_exc = md["dale_conf"]["num_excite"] + off_mask = ~np.eye(n_neur, dtype=bool) + neg_mask = off_mask & (A_true < 0.0) + pos_mask = off_mask & (A_true > 0.0) + diag_mask = np.eye(n_neur, dtype=bool) + n_diag = int(np.sum(diag_mask)) + n_off = int(np.count_nonzero(A_true[off_mask])) + + vmin = float(np.min(A_true)) + vmax = float(np.max(A_true)) + if np.isclose(vmin, vmax): + vmax = vmin + 1e-9 + shared_norm = colors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(14, 8)) + gs = fig.add_gridspec(2, 3, hspace=0.55, wspace=0.45) + + ax = fig.add_subplot(gs[0, 0]) + self._draw_A_matrix( + ax, A_true, f"True Dale, N{n_neur}, nEdges={n_diag}+{n_off}", + num_exc=num_exc, norm_map=shared_norm, + ) + + ax = fig.add_subplot(gs[0, 1]) + _, cmp_diag, cmp_off = self._count_A_edges(A_cmp) + self._draw_A_matrix( + ax, A_cmp, f"{est_label}, nEdges={cmp_diag}+{cmp_off}", + num_exc=num_exc, norm_map=shared_norm, + ) + + ax = fig.add_subplot(gs[0, 2]) + self._draw_A_offdiag_hist(ax, A_cmp, f"{est_label} off-diagonal") + if est_key == "A_init": + ax.text( + 0.05, 0.70, self._A_init_diagnostics_txt(md), transform=ax.transAxes, + va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"), + family="monospace", + ) + + ax = fig.add_subplot(gs[1, 0]) + self._scatter_true_vs_fit( + ax, A_true[neg_mask], A_cmp[neg_mask], + f"{est_label}: neg TP", color="tab:green", + ) + + ax = fig.add_subplot(gs[1, 1]) + self._scatter_true_vs_fit( + ax, A_true[pos_mask], A_cmp[pos_mask], + f"{est_label}: pos TP", color="tab:green", + ) + + ax = fig.add_subplot(gs[1, 2]) + self._scatter_true_vs_fit( + ax, A_true[diag_mask], A_cmp[diag_mask], + f"{est_label}: diag", color="salmon", split_by_true_sign=True, + ) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle(f"{title_prefix}, A_true vs {est_label} comparison", fontsize=12) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.94]) + + def _overlay_single_rates(self, ax, single_rates, x_src): + rate = np.asarray(single_rates, dtype=np.float64).ravel() + ax2 = ax.twinx() + ax2.plot(x_src, rate, color="tab:orange", linewidth=0.9, alpha=0.75) + ax2.set_ylabel("frequency (Hz)", color="tab:orange") + ax2.tick_params(axis="y", labelcolor="tab:orange") + return ax2 + + def _draw_A_offdiag_hist(self, ax, A, title, diagnostics_txt=None, weight_lines=None): + A = np.asarray(A) + n_neuron = A.shape[0] + off_mask = ~np.eye(n_neuron, dtype=bool) + A_off = A[off_mask] + A_off_nz = A_off[np.abs(A_off) > 1e-12] + ax.hist(A_off_nz, bins=120, color="saddlebrown", alpha=0.85) + ax.set_yscale("log") + ax.grid(True, alpha=0.35) + ax.set_title(title) + ax.set_xlabel("edge value") + ax.set_ylabel("edges") + if weight_lines is not None: + for y_pos, w in zip([0.88, 0.78], weight_lines): + if np.isfinite(w): + ax.axvline(w, color="tab:blue", ls="--", lw=0.8, alpha=0.9) + ax.text( + w, y_pos, f"{w:.3f}", transform=ax.get_xaxis_transform(), + va="top", ha="center", fontsize=8, color="tab:blue", + rotation=90, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=1), + ) + if diagnostics_txt is not None: + ax.text( + 0.05, 0.70, diagnostics_txt, transform=ax.transAxes, + va="top", ha="left", fontsize=9, + bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"), + family="monospace", + ) + + def _draw_A_matrix_summary_3cols(self, fig, gs, row, A, label, single_rates, num_exc=None, weight_lines=None): + A = np.asarray(A, dtype=np.float64) + n_neuron = A.shape[0] + x_src = np.arange(n_neuron, dtype=np.int64) + single_rates = np.asarray(single_rates, dtype=np.float64).ravel() + assert single_rates.shape[0] == n_neuron, ( + f"single_rates length {single_rates.shape[0]} != N={n_neuron}" + ) + + _, n_diag, n_off = self._count_A_edges(A) + cnt_src, _ = self._source_nz_edge_stats(A) + sum_nedge = int(np.sum(cnt_src)) + + ax = fig.add_subplot(gs[row, 0]) + self._draw_A_offdiag_hist(ax, A, f"{label} off-diagonal", weight_lines=weight_lines) + ax.text( + 0.98, 0.97, f"sum Nedge={sum_nedge:d}", transform=ax.transAxes, + va="top", ha="right", fontsize=9, color="k", + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=2), + ) + + ax = fig.add_subplot(gs[row, 1]) + self._draw_A_matrix(ax, A, f"{label}, nEdges={n_diag}+{n_off}", num_exc=num_exc) + + ax = fig.add_subplot(gs[row, 2]) + ax.plot(x_src, cnt_src, color="tab:blue", linewidth=1.0) + med_cnt = int(np.median(cnt_src)) + ax.axhline(med_cnt, color="green", ls="--", lw=1.0, alpha=0.9) + ax.text( + 0.02, 0.95, f"median={med_cnt:d}", transform=ax.transAxes, + va="top", ha="left", fontsize=9, color="green", + ) + ax.set_title(f"{label}: # non-zero edges") + ax.set_xlabel("source neuron index (column)") + ax.set_ylabel("# outgoing non-zero edges") + ax.grid(True, alpha=0.35) + self._overlay_single_rates(ax, single_rates, x_src) + + def node_outgoing_edge_stats( + self, fitD, md, single_rates, neuron_type, + figId="e", est_key="A_hat", est_label="A_hat", + ): + A_est = np.asarray(fitD[est_key], dtype=np.float64) + A_prune = self._A_prune_from_fitD(fitD) + prune_label = self._display_A_label(md, prune=True) + assert A_est.ndim == 2 and A_est.shape[0] == A_est.shape[1], "A_est must be square" + assert A_prune.shape == A_est.shape, "A_prune shape must match A_hat" + n_neuron = A_est.shape[0] + + nedge, sedge = self._node_outgoing_edge_stats(A_est) + n_bins = max(10, min(40, int(np.sqrt(n_neuron)) * 2)) + stage_b = md["bagsFDR_stageB"] + min_posW = float(stage_b["min_posW"]) + max_negW = float(stage_b["max_negW"]) + neuron_type = np.asarray(neuron_type, dtype=np.int8) + assert neuron_type.shape[0] == n_neuron, "neuron_type length must match A_hat columns" + n_exc = int(np.sum(neuron_type > 0)) + n_inh = int(np.sum(neuron_type < 0)) + n_und = int(np.sum(neuron_type == 0)) + + def add_hist_percentile_marker(ax, vals, fmt=".3f"): + vals = np.asarray(vals, dtype=np.float64) + vals = vals[np.isfinite(vals)] + p16, p50, p84 = np.percentile(vals, [16, 50, 84]) + ylo, yhi = ax.get_ylim() + y_mark = 0.5 * (ylo + yhi) + ax.errorbar( + p50, y_mark, + xerr=[[p50 - p16], [p84 - p50]], + fmt="o", color="k", markersize=8, + capsize=4, capthick=1.2, elinewidth=1.2, zorder=5, + ) + txt = f"p16={format(p16, fmt)}\nmed={format(p50, fmt)}\np84={format(p84, fmt)}" + ax.text( + 0.98, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=9, color="k", + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=2), + ) + + num_exc = None + dale_conf = md.get("dale_conf") + if dale_conf is not None: + num_exc = dale_conf.get("num_excite") + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(14, 8)) + gs = fig.add_gridspec(2, 3, height_ratios=[1.0, 1.15], hspace=0.45, wspace=0.35) + + _, n_diag_prune, n_off_prune = self._count_A_edges(A_prune) + cnt_src_prune, _ = self._source_nz_edge_stats(A_prune) + sum_nedge_prune = int(np.sum(cnt_src_prune)) + x_src = np.arange(n_neuron, dtype=np.int64) + + ax = fig.add_subplot(gs[0, 0]) + for cls, color, label in [ + (1, "magenta", f"exc={n_exc}"), + (-1, "#39FF14", f"inh={n_inh}"), + (0, "salmon", f"und={n_und}"), + ]: + mask = neuron_type == cls + if cls == 0: + ax.scatter( + nedge[mask], sedge[mask], s=18, marker="o", + alpha=0.80, color=color, edgecolors="k", + linewidths=0.35, label=label, + ) + else: + ax.scatter(nedge[mask], sedge[mask], s=16, marker=".", alpha=0.80, color=color, label=label) + ax.axhline(min_posW, color="tab:blue", ls="--", lw=0.8, alpha=0.9) + ax.axhline(max_negW, color="tab:blue", ls="--", lw=0.8, alpha=0.9) + ax.set( + title="Reco Nedge vs Sedge", + xlabel="Nedge (# outgoing edges per source column)", + ylabel="Sedge", + ) + ax.legend(loc="best", fontsize=9) + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[0, 1]) + self._draw_A_offdiag_hist(ax, A_prune, f"Reco {prune_label} off-diagonal", weight_lines=(max_negW, min_posW)) + ax.text( + 0.98, 0.97, f"sum Nedge={sum_nedge_prune:d}", transform=ax.transAxes, + va="top", ha="right", fontsize=9, color="k", + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=2), + ) + + ax = fig.add_subplot(gs[0, 2]) + blk = 10 + n_blk = n_neuron // blk + blk_edges = np.arange(n_blk + 1) * blk + blk_cnt = cnt_src_prune[:n_blk * blk].reshape(n_blk, blk).mean(axis=1) + ax.stairs(blk_cnt, blk_edges, color="tab:blue", linewidth=1.2) + med_cnt = float(np.median(blk_cnt)) + ax.axhline(med_cnt, color="green", ls="--", lw=1.0, alpha=0.9) + ax.text( + 0.02, 0.95, f"median={med_cnt:.1f}", transform=ax.transAxes, + va="top", ha="left", fontsize=9, color="green", + ) + ax.set_title(f"Reco {prune_label}: median non-zero edges={med_cnt:.1f}") + ax.set_xlabel(f"source neuron index (column), step={blk}") + ax.set_ylabel("# outgoing non-zero edges") + ax.grid(True, alpha=0.35) + rate = np.asarray(single_rates, dtype=np.float64).ravel() + blk_rate = rate[:n_blk * blk].reshape(n_blk, blk).mean(axis=1) + ax2 = ax.twinx() + ax2.stairs(blk_rate, blk_edges, color="tab:orange", linewidth=1.2, alpha=0.85) + ax2.set_ylabel("frequency (Hz)", color="tab:orange") + ax2.tick_params(axis="y", labelcolor="tab:orange") + + ax = fig.add_subplot(gs[1, 0]) + single_rates = np.asarray(single_rates, dtype=np.float64).ravel() + for cls, color, label in [ + (1, "magenta", f"exc={n_exc}"), + (-1, "#39FF14", f"inh={n_inh}"), + (0, "salmon", f"und={n_und}"), + ]: + mask = neuron_type == cls + if cls == 0: + ax.scatter( + nedge[mask], single_rates[mask], s=18, marker="o", + alpha=0.80, color=color, edgecolors="k", + linewidths=0.35, label=label, + ) + else: + ax.scatter(nedge[mask], single_rates[mask], s=16, marker=".", alpha=0.80, color=color, label=label) + ax.set(title="Reco Nedge vs frequency", xlabel="Nedge (# outgoing edges per source column)", ylabel="frequency (Hz)") + ax.legend(loc="best", fontsize=9) + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[1, 1]) + off_mask = ~np.eye(n_neuron, dtype=bool) + med_weight = np.zeros((n_neuron,), dtype=np.float64) + for j in range(n_neuron): + vals = A_est[off_mask[:, j], j] + vals = vals[vals != 0.0] + if vals.size: + med_weight[j] = float(np.median(vals)) + for cls, color, label in [ + (1, "magenta", f"exc={n_exc}"), + (-1, "#39FF14", f"inh={n_inh}"), + (0, "salmon", f"und={n_und}"), + ]: + mask = neuron_type == cls + if cls == 0: + ax.scatter( + med_weight[mask], single_rates[mask], s=18, marker="o", + alpha=0.80, color=color, edgecolors="k", + linewidths=0.35, label=label, + ) + else: + ax.scatter(med_weight[mask], single_rates[mask], s=16, marker=".", alpha=0.80, color=color, label=label) + ax.axvline(0.0, color="tab:blue", ls="--", lw=0.8, alpha=0.9) + ax.set(title="Reco median weight vs frequency", xlabel="median outgoing edge weight", ylabel="frequency (Hz)") + ax.legend(loc="best", fontsize=9) + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[1, 2]) + self._draw_A_matrix( + ax, A_prune, f"Reco {prune_label}, nEdges={n_diag_prune}+{n_off_prune}", + num_exc=num_exc, + ) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle( + f"{title_prefix}, {est_label} per-source outgoing edges (column=source)", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.94]) + + def final_weight_distributions(self, fitD, md, figId="h"): + """Final non-truth weight summaries for FDR aggregate outputs.""" + required = ("A_prune", "B_hat", "neuron_type", "A_diag_mean", "A_diag_stderr") + missing = [key for key in required if key not in fitD] + assert not missing, f"Plot h requires regenerated aggregate arrays; missing {missing}" + + A_prune = np.asarray(fitD["A_prune"], dtype=np.float64) + B_hat = np.asarray(fitD["B_hat"], dtype=np.float64) + neuron_type = np.asarray(fitD["neuron_type"], dtype=np.int8) + diag_mean = np.asarray(fitD["A_diag_mean"], dtype=np.float64).ravel() + diag_stderr = np.asarray(fitD["A_diag_stderr"], dtype=np.float64).ravel() + A_label = self._display_A_label(md, prune=True) + B_label = self._display_B_label(md) + assert A_prune.ndim == 2 and A_prune.shape[0] == A_prune.shape[1], "A_prune must be square" + n_neuron = A_prune.shape[0] + assert neuron_type.shape[0] == n_neuron, "neuron_type length must match A_prune" + assert diag_mean.shape[0] == n_neuron and diag_stderr.shape[0] == n_neuron, ( + "A_diag_mean/A_diag_stderr length must match A_prune" + ) + if B_hat.ndim == 1: + B_hat = B_hat[None, :] + assert B_hat.ndim == 2 and B_hat.shape[1] == n_neuron, "B_hat must have shape (M,N)" + + off_mask = ~np.eye(n_neuron, dtype=bool) + + def outgoing_values_for_type(type_value): + src_mask = neuron_type == int(type_value) + vals = A_prune[:, src_mask][off_mask[:, src_mask]] + return vals[np.abs(vals) > 1e-12] + + exc_w = outgoing_values_for_type(1) + inh_w = outgoing_values_for_type(-1) + diag_vals = np.diag(A_prune) + n_exc = int(np.sum(neuron_type > 0)) + n_inh = int(np.sum(neuron_type < 0)) + n_und = int(np.sum(neuron_type == 0)) + + def mark_edge(ax, val): + if np.isfinite(val): + ax.axvline(val, color="tab:blue", ls="--", lw=0.8, alpha=0.9) + ax.text( + val, 0.92, f"{val:.3f}", transform=ax.get_xaxis_transform(), + va="top", ha="center", fontsize=8, color="tab:blue", + rotation=90, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=1), + ) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(14, 8)) + gs = fig.add_gridspec(2, 3, hspace=0.45, wspace=0.35) + + ax = fig.add_subplot(gs[0, 0]) + ax.hist(exc_w, bins=80, color="magenta", alpha=0.75) + ax.axvline(0.0, color="k", ls="--", lw=0.8) + mark_edge(ax, float(np.min(exc_w[exc_w > 0.0])) if np.any(exc_w > 0.0) else float("nan")) + ax.set(title=f"Reco exc outgoing weights, neurons={n_exc}", xlabel=f"{A_label} weight", ylabel="edges") + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[0, 1]) + ax.hist(inh_w, bins=80, color="#39FF14", alpha=0.75) + ax.axvline(0.0, color="k", ls="--", lw=0.8) + mark_edge(ax, float(np.max(inh_w[inh_w < 0.0])) if np.any(inh_w < 0.0) else float("nan")) + ax.set(title=f"Reco inh outgoing weights, neurons={n_inh}", xlabel=f"{A_label} weight", ylabel="edges") + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[0, 2]) + bins = np.linspace(diag_vals.min(), diag_vals.max(), 81) + nonneg_counts = [] + for cls, color, label in [ + (1, "magenta", f"exc={n_exc}"), + (-1, "#39FF14", f"inh={n_inh}"), + (0, "salmon", f"und={n_und}"), + ]: + mask = neuron_type == cls + name = label.split("=")[0] + nonneg_counts.append((name, int(np.sum(diag_vals[mask] >= 0.0)))) + if np.any(mask): + ec = "k" if cls == 0 else "none" + ax.hist(diag_vals[mask], bins=bins, color=color, alpha=0.75, label=label, edgecolor=ec, linewidth=0.5) + ax.axvline(0.0, color="k", ls="--", lw=0.8) + ax.legend(loc="best", fontsize=9) + tbl = ax.table( + cellText=[[name, f"{cnt:d}"] for name, cnt in nonneg_counts], + colLabels=["type", ">=0"], + cellLoc="center", + colLoc="center", + bbox=[0.70, 0.46, 0.27, 0.28], + ) + tbl.auto_set_font_size(False) + tbl.set_fontsize(8) + for cell in tbl.get_celld().values(): + cell.set_edgecolor("0.65") + cell.set_linewidth(0.5) + cell.set_facecolor((1.0, 1.0, 1.0, 0.78)) + ax.set(title="Reco A diagonal", xlabel=f"{A_label} diagonal value", ylabel="neurons") + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[1, 0]) + ax.hist(B_hat[0], bins=80, color="tab:orange", alpha=0.8) + ax.axvline(0.0, color="k", ls="--", lw=0.8) + ax.set(title=f"{B_label} mode 0", xlabel="B value", ylabel="neurons") + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[1, 1]) + if B_hat.shape[0] >= 2: + ax.hist(B_hat[1], bins=80, color="tab:purple", alpha=0.8) + ax.axvline(0.0, color="k", ls="--", lw=0.8) + ax.set(title=f"{B_label} mode 1", xlabel="B value", ylabel="neurons") + else: + ax.set_axis_off() + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[1, 2]) + for cls, color, label in [ + (1, "magenta", f"exc={n_exc}"), + (-1, "#39FF14", f"inh={n_inh}"), + (0, "salmon", f"und={n_und}"), + ]: + mask = neuron_type == cls + if not np.any(mask): + continue + kw = dict(s=18, marker="o", edgecolors="k", linewidths=0.5) if cls == 0 else dict(s=16, marker=".") + ax.scatter(diag_mean[mask], diag_stderr[mask], alpha=0.80, color=color, label=label, **kw) + ax.axvline(0.0, color="k", ls="--", lw=0.8) + ax.legend(loc="best", fontsize=9) + ax.set(title="A-diag summary", xlabel="mean diagonal A", ylabel="stderr diagonal A") + ax.grid(True, alpha=0.35) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle( + f"{title_prefix}, final weight distributions ({A_label}, {B_label}): " + f"exc={n_exc} inh={n_inh} und={n_und}", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.94]) + + def offdiag_weight_investigation(self, fitD, md, figId="i"): + """Off-diagonal A_prune weight investigation: histogram and weight-vs-frequency heatmap.""" + A_prune = np.asarray(fitD["A_prune"], dtype=np.float64) + neuron_type = np.asarray(fitD["neuron_type"], dtype=np.int8) + single_rates = np.asarray(fitD["single_rates"], dtype=np.float64).ravel() + A_label = self._display_A_label(md, prune=True) + n_neuron = A_prune.shape[0] + off_mask = ~np.eye(n_neuron, dtype=bool) + n_exc = int(np.sum(neuron_type > 0)) + n_inh = int(np.sum(neuron_type < 0)) + n_und = int(np.sum(neuron_type == 0)) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(18, 4)) + gs = fig.add_gridspec(1, 4, wspace=0.38, left=0.05, right=0.97) + + # --- plot 1: off-diagonal weight histogram by type --- + ax = fig.add_subplot(gs[0, 0]) + all_off = A_prune[off_mask] + bins = np.linspace(all_off.min(), all_off.max(), 81) + edge_counts = {} + for cls, color, label in [ + (1, "magenta", f"exc={n_exc}"), + (-1, "#39FF14", f"inh={n_inh}"), + (0, "salmon", f"und={n_und}"), + ]: + src_mask = neuron_type == cls + if not np.any(src_mask): + continue + vals = A_prune[:, src_mask][off_mask[:, src_mask]] + vals = vals[np.abs(vals) > 1e-12] + edge_counts[label.split("=")[0]] = len(vals) + ec = "k" if cls == 0 else "none" + ax.hist(vals, bins=bins, color=color, alpha=0.75, label=label, edgecolor=ec, linewidth=0.5) + ax.axvline(0.0, color="k", ls="--", lw=0.8) + leg = ax.legend(loc="upper left", fontsize=9, title="neurons") + leg.get_title().set_fontsize(9) + edge_txt = "\n".join(f"{k}: {v}" for k, v in edge_counts.items()) + ax.text(0.98, 0.97, f"edges:\n{edge_txt}", transform=ax.transAxes, + va="top", ha="right", fontsize=9, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=2)) + ax.set(title="Off-diag weights by type", xlabel=f"{A_label} weight", ylabel="edges") + ax.grid(True, alpha=0.35) + + # --- plot 2: weight vs spike frequency heatmap (column = source neuron) --- + import matplotlib.colors as mcolors + ax = fig.add_subplot(gs[0, 1]) + col_idx = np.where(off_mask)[1] # source column for each off-diag entry + w_vals = A_prune[off_mask] + nz = np.abs(w_vals) > 1e-12 + freq_rep = single_rates[col_idx[nz]] + w_nz = w_vals[nz] + h, xedges, yedges = np.histogram2d(w_nz, freq_rep, bins=[80, 40]) + pcm = ax.pcolormesh(xedges, yedges, h.T, cmap="Blues", shading="flat", + norm=mcolors.LogNorm(vmin=1, vmax=h.max())) + fig.colorbar(pcm, ax=ax, label="edges (log)") + ax.axvline(0.0, color="k", ls="--", lw=0.8) + ax.set(title="Weight vs source frequency", xlabel=f"{A_label} weight", ylabel="spike frequency (Hz)") + ax.grid(True, alpha=0.2) + + # --- plot 3: diagonal weight histogram by type --- + diag_vals = np.diag(A_prune) + ax = fig.add_subplot(gs[0, 2]) + bins_diag = np.linspace(diag_vals.min(), diag_vals.max(), 81) + for cls, color, label in [ + (1, "magenta", f"exc={n_exc}"), + (-1, "#39FF14", f"inh={n_inh}"), + (0, "salmon", f"und={n_und}"), + ]: + mask = neuron_type == cls + if not np.any(mask): + continue + vals = diag_vals[mask] + vals = vals[np.abs(vals) > 1e-12] + ec = "k" if cls == 0 else "none" + ax.hist(vals, bins=bins_diag, color=color, alpha=0.75, label=label, edgecolor=ec, linewidth=0.5) + ax.axvline(0.0, color="k", ls="--", lw=0.8) + leg = ax.legend(loc="upper center", fontsize=9, title="neurons") + leg.get_title().set_fontsize(9) + ax.set(title="Diagonal weights by type", xlabel=f"{A_label} diagonal", ylabel="neurons") + ax.grid(True, alpha=0.35) + + # --- plot 4: diagonal weight vs spike frequency heatmap --- + ax = fig.add_subplot(gs[0, 3]) + diag_nz_mask = np.abs(diag_vals) > 1e-12 + h2, xedges2, yedges2 = np.histogram2d( + diag_vals[diag_nz_mask], single_rates[diag_nz_mask], bins=[80, 40]) + pcm2 = ax.pcolormesh(xedges2, yedges2, h2.T, cmap="Blues", shading="flat", + norm=mcolors.LogNorm(vmin=1, vmax=h2.max())) + fig.colorbar(pcm2, ax=ax, label="neurons (log)") + ax.axvline(0.0, color="k", ls="--", lw=0.8) + ax.set(title="Diagonal vs source frequency", xlabel=f"{A_label} diagonal", ylabel="spike frequency (Hz)") + ax.grid(True, alpha=0.2) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle(f"{title_prefix}, {A_label} off-diagonal weight investigation", fontsize=12) + fig.subplots_adjust(top=0.88) + + def fdr_selection_summary(self, fitD, md, figId="g"): + """Summarize Stage (b) FDR bag filtering and cross-bag stability.""" + import matplotlib.colors as colors + + stage_b = md.get("bagsFDR_stageB") + assert stage_b is not None, "Plot g requires a Stage (b) aggregate with bagsFDR_stageB metadata" + required = ( + "A_bag", "selected_in_bag", "selection_frequency", "selected_mask", + "src_null_tau_bag", "selected_edges_per_bag", "selected_edges_per_bag_src", + "A_hat", "A_mean_selected", "single_rates", + ) + missing = [key for key in required if key not in fitD] + assert not missing, f"Plot g requires Stage (b) diagnostic arrays; missing {missing}" + + A_bag = np.asarray(fitD["A_bag"], dtype=np.float64) + selected_in_bag = np.asarray(fitD["selected_in_bag"], dtype=bool) + sel_freq = np.asarray(fitD["selection_frequency"], dtype=np.float64) + final_sel = np.asarray(fitD["selected_mask"], dtype=bool) + tau = np.asarray(fitD["src_null_tau_bag"], dtype=np.float64) + selected_edges_per_bag = np.asarray(fitD["selected_edges_per_bag"], dtype=np.int64) + A_hat = np.asarray(fitD["A_hat"], dtype=np.float64) + + assert A_bag.ndim == 3 and A_bag.shape[1] == A_bag.shape[2], "A_bag must have shape (bag,N,N)" + assert selected_in_bag.shape == A_bag.shape, "selected_in_bag shape must match A_bag" + n_bag, n_neur, _ = A_bag.shape + assert tau.shape == (n_bag, n_neur), f"src_null_tau_bag shape {tau.shape} != {(n_bag, n_neur)}" + + off_mask = ~np.eye(n_neur, dtype=bool) + off3 = np.broadcast_to(off_mask, A_bag.shape) + p_cand = int(np.sum(off_mask)) + per_bag_quantile = float(stage_b["per_bag_quantile"]) + stab_sel_thresh = float(stage_b["stab_sel_thresh"]) + final_count = int(np.sum(final_sel & off_mask)) + ever_count = int(np.sum((sel_freq > 0.0) & off_mask)) + mean_per_bag = float(np.mean(selected_edges_per_bag)) + median_per_bag = float(np.median(selected_edges_per_bag)) + + with np.errstate(invalid="ignore", divide="ignore"): + ratio = np.abs(A_bag) / tau[:, None, :] + ratio_all = ratio[off3] + ratio_sel = ratio[selected_in_bag & off3] + ratio_all = ratio_all[np.isfinite(ratio_all)] + ratio_sel = ratio_sel[np.isfinite(ratio_sel)] + log_ratio_all = np.log10(np.clip(ratio_all, 1e-8, 1e8)) + log_ratio_sel = np.log10(np.clip(ratio_sel, 1e-8, 1e8)) + + freq_off = sel_freq[off_mask] + freq_nonzero = freq_off[freq_off > 0.0] + thresholds = np.arange(1, n_bag + 1, dtype=np.float64) / float(n_bag) + survivor_counts = np.array([np.sum(freq_off >= th) for th in thresholds], dtype=np.int64) + + A_mean_selected = np.asarray(fitD["A_mean_selected"], dtype=np.float64) + amp = np.abs(A_mean_selected[off_mask]) + final_for_scatter = (final_sel & off_mask)[off_mask] + finite_amp = np.isfinite(amp) + + src_tau_mean = np.mean(tau, axis=0) + src_tau_sd = np.std(tau, axis=0, ddof=1) if n_bag > 1 else np.zeros_like(src_tau_mean) + src_selected_mean = np.mean(np.asarray(fitD["selected_edges_per_bag_src"]), axis=0) + single_rates = np.asarray(fitD["single_rates"], dtype=np.float64).ravel() + assert single_rates.shape[0] == n_neur, "single_rates length must match A_bag columns" + src_order = np.argsort(single_rates) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(17, 12)) + gs = fig.add_gridspec(3, 3, hspace=0.48, wspace=0.36) + + ax = fig.add_subplot(gs[0, 0]) + funnel_labels = ["all\npairs", "mean\nbag pass", "ever\npass", "stable\nfinal"] + funnel_counts = [p_cand, mean_per_bag, ever_count, final_count] + bar_colors = ["0.72", "tab:blue", "tab:orange", "tab:green"] + ax.bar(np.arange(4), funnel_counts, color=bar_colors, alpha=0.85) + ax.set_yscale("log") + ax.set_xticks(np.arange(4)) + ax.set_xticklabels(funnel_labels) + ax.set_ylabel("off-diagonal edges") + ax.set_title("FDR reduction funnel") + ax.grid(True, axis="y", alpha=0.35) + for i, val in enumerate(funnel_counts): + ax.text(i, max(val, 1.0), f"{val:.0f}", ha="center", va="bottom", fontsize=9) + + ax = fig.add_subplot(gs[0, 1]) + xbag = np.arange(n_bag) + ax.bar(xbag, selected_edges_per_bag, color="tab:blue", alpha=0.8) + ax.axhline(mean_per_bag, color="k", ls="--", lw=1, label=f"mean={mean_per_bag:.1f}") + ax.axhline(median_per_bag, color="tab:orange", ls=":", lw=1.4, label=f"median={median_per_bag:.1f}") + ax.set(title="Per-bag edges passing source null threshold", xlabel="bag index", ylabel="edges") + ax.grid(True, axis="y", alpha=0.35) + ax.legend(fontsize=8) + + ax = fig.add_subplot(gs[0, 2]) + ax.step(thresholds, survivor_counts, where="post", color="tab:green", linewidth=2) + ax.axvline(stab_sel_thresh, color="red", ls="--", lw=1.2, label=f"stab_sel_thresh={stab_sel_thresh:g}") + ax.axhline(final_count, color="k", ls=":", lw=1.0, label=f"final={final_count}") + ax.set(title="Stable edge count", xlabel="stability sel. thres.", ylabel="stable edges") + ax.set_xlim(0.0, 1.02) + ax.set_ylim(0, max(1, int(np.max(survivor_counts)) + 1)) + ax.grid(True, alpha=0.35) + ax.legend(fontsize=8) + + ax = fig.add_subplot(gs[1, 0]) + bins = np.linspace(0.05, 1.05, 11) + if freq_nonzero.size: + ax.hist(freq_nonzero, bins=bins, color="tab:purple", alpha=0.75) + ax.axvline(stab_sel_thresh, color="red", ls="--", lw=1.2) + ax.set(title=f"Ever-pass edges, n={freq_nonzero.size}", + xlabel="stability sel. thres.", ylabel="edges") + ax.set_xlim(0.05, 1.05) + ax.set_xticks(np.arange(0.1, 1.01, 0.1)) + ax.grid(True, axis="y", alpha=0.35) + + ax = fig.add_subplot(gs[1, 1]) + im = ax.imshow(sel_freq, origin="lower", aspect="equal", interpolation="nearest", + cmap="viridis", vmin=0.0, vmax=1.0) + if final_count > 0: + ax.contour(final_sel.astype(float), levels=[0.5], colors="red", linewidths=0.6) + ax.plot([0, n_neur], [0, n_neur], "--", lw=0.7, color="white", alpha=0.9) + ax.set(title="Selection frequency matrix", xlabel="source neuron", ylabel="target neuron") + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="selected fraction") + + ax = fig.add_subplot(gs[1, 2]) + vmax = float(np.nanmax(np.abs(A_hat))) if np.any(np.isfinite(A_hat)) else 1.0 + if vmax <= 0: + vmax = 1.0 + norm = colors.TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax) + self._draw_A_matrix(ax, A_hat, f"Final A_hat, stable offdiag={final_count}", norm_map=norm) + + ax = fig.add_subplot(gs[2, 0]) + bins = np.linspace(-3.0, 3.0, 100) + if log_ratio_all.size: + ax.hist(log_ratio_all, bins=bins, color="0.65", alpha=0.75, label="all bag-edge tests") + if log_ratio_sel.size: + ax.hist(log_ratio_sel, bins=bins, color="tab:blue", alpha=0.65, label="per-bag pass") + ax.axvline(0.0, color="red", ls="--", lw=1.2, label="|A| = tau") + ax.set_yscale("log") + ax.set(title=f"Per-bag threshold pressure, q={per_bag_quantile:g}", + xlabel="log10(|A_bag| / source tau)", ylabel="bag-edge tests") + ax.grid(True, axis="y", alpha=0.35) + ax.legend(fontsize=8) + + ax = fig.add_subplot(gs[2, 1]) + x = np.arange(n_neur) + ax.plot(x, src_tau_mean[src_order], color="tab:red", linewidth=1.0, label="mean source tau") + ax.fill_between( + x, + np.maximum(0.0, src_tau_mean[src_order] - src_tau_sd[src_order]), + src_tau_mean[src_order] + src_tau_sd[src_order], + color="tab:red", alpha=0.18, linewidth=0, + ) + ax.set(title="Source null threshold and outgoing pass count", + xlabel="source neuron sorted by firing rate", ylabel="source tau") + ax.grid(True, alpha=0.35) + ax2 = ax.twinx() + ax2.plot(x, src_selected_mean[src_order], color="tab:blue", linewidth=0.9, alpha=0.9, label="mean outgoing pass count") + ax2.set_ylabel("mean per-bag outgoing pass count") + lines1, lab1 = ax.get_legend_handles_labels() + lines2, lab2 = ax2.get_legend_handles_labels() + ax.legend(lines1 + lines2, lab1 + lab2, fontsize=8, loc="upper right") + rate_min = float(single_rates[src_order[0]]) + rate_max = float(single_rates[src_order[-1]]) + ax.text( + 0.02, 0.96, f"rate: {rate_min:.2g} to {rate_max:.2g} Hz", + transform=ax.transAxes, va="top", ha="left", fontsize=8, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none", pad=2), + ) + + ax = fig.add_subplot(gs[2, 2]) + positive_amp = finite_amp & (amp > 0.0) + nonfinal = positive_amp & ~final_for_scatter & (freq_off > 0.0) + final = positive_amp & final_for_scatter + ax.scatter(freq_off[nonfinal], amp[nonfinal], s=12, color="0.55", alpha=0.45, label="rejected after bag pass") + ax.scatter(freq_off[final], amp[final], s=18, color="tab:green", alpha=0.8, label="final stable") + ax.axvline(stab_sel_thresh, color="red", ls="--", lw=1.2) + ax.set_yscale("log") + ax.set(title="Amplitude vs stability", xlabel="stability sel. thres.", ylabel="|A_mean_selected|") + ax.set_xlim(-0.02, 1.02) + ax.grid(True, alpha=0.35) + ax.legend(fontsize=8) + + false_bound = float(stage_b.get("stability_false_edge_bound", np.nan)) + title_prefix = self.canvas_title_prefix(md) + fig.suptitle( + f"{title_prefix}, FDR bag selection: q={per_bag_quantile:g}, " + f"stab_sel_thresh={stab_sel_thresh:g}, bags={n_bag}, final={final_count}, " + f"bound={false_bound:.3g}", + fontsize=13, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.955]) + + def fdr_acceptance_truth(self, accD, md, figId="o"): + """Display simulation-truth FDR/bag acceptance diagnostics precomputed by eval.""" + weight = accD["weight"] + rate = accD["rate"] + summary = accD["summary"] + + wc = np.asarray(weight["center"], dtype=np.float64) + wprob = np.asarray(weight["prob"], dtype=np.float64) + wtotal = np.asarray(weight["total"], dtype=np.int64) + wpass = np.asarray(weight["passed"], dtype=np.int64) + wbin = float(weight["bin_width"]) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(18, 4.2)) + gs = fig.add_gridspec(1, 4, left=0.055, right=0.99, bottom=0.18, top=0.78, wspace=0.34) + + ax = fig.add_subplot(gs[0, 0]) + m = np.isfinite(wprob) + ax.bar(wc[m], wprob[m], width=0.92 * wbin, color="tab:purple", alpha=0.85, align="center") + ax.set( + title="Acceptance vs true edge value", + xlabel="A_true edge value", + ylabel="P(A_prune edge accepted)", + ylim=(-0.03, 1.03), + ) + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[0, 1]) + for label, color in (("exc", "red"), ("inh", "blue")): + d = rate[label] + x = np.asarray(d["center"], dtype=np.float64) + y = np.asarray(d["prob"], dtype=np.float64) + mm = np.isfinite(y) + ax.plot(x[mm], y[mm], "o-", color=color, lw=1.2, ms=4, label=label) + if rate["scale"] == "log": + ax.set_xscale("log") + ax.set( + title="Acceptance vs source frequency", + xlabel="source neuron frequency (Hz)", + ylabel="P(A_prune edge accepted)", + ylim=(-0.03, 1.03), + ) + ax.legend(fontsize=9) + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[0, 2]) + mt = wtotal > 0 + mp = wpass > 0 + ax.bar(wc[mt], wtotal[mt], width=0.92 * wbin, color="0.75", alpha=0.8, label="true edges") + ax.bar(wc[mp], wpass[mp], width=0.55 * wbin, color="tab:green", alpha=0.85, label="accepted") + ax.set_yscale("log") + ax.set(title="True-edge counts vs value", xlabel="A_true edge value", ylabel="edges") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[0, 3]) + for label, color, ls in (("exc", "red", "-"), ("inh", "blue", "--")): + d = rate[label] + x = np.asarray(d["center"], dtype=np.float64) + total = np.asarray(d["total"], dtype=np.int64) + passed = np.asarray(d["passed"], dtype=np.int64) + mt = total > 0 + mp = passed > 0 + ax.plot(x[mt], total[mt], ls=ls, color=color, lw=1.0, alpha=0.45, label=f"{label} true edges") + ax.plot(x[mp], passed[mp], "o-", color=color, lw=1.2, ms=4, label=f"{label} accepted") + if rate["scale"] == "log": + ax.set_xscale("log") + ax.set_yscale("log") + ax.set(title="True-edge counts vs source frequency", xlabel="source neuron frequency (Hz)", ylabel="edges") + ax.legend(fontsize=8, ncol=1) + ax.grid(True, alpha=0.35) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle( + f"{title_prefix}, FDR/bag A_prune acceptance vs truth: " + f"{summary['num_accepted']}/{summary['num_candidates']} true edges accepted, " + f"truth exc={summary['num_exc']} inh={summary['num_inh']}", + fontsize=12, + ) + + def _A_init_diagnostics_txt(self, md): + initA = real_fit_metadata(md).get("init_A", {}) + if initA is None: + initA = {} + method = initA.get("method", "unknown") + lines = ["A-init diagnostics:", f" method = {method}"] + if "source_key" in initA: + lines.append(f" source = {initA['source_key']}") + if "cond_YpYp" in initA: + lines.extend([ + f" cond(YpYp) = {initA['cond_YpYp']:.2e}", + f" rho(A_ols) = {initA['rho_A_init']:.3f}", + f" R2 = {initA['R2_1step']:.3f}", + f" ||A||_F = {initA['fro_A_init']:.3f}", + f" bins_used = {initA['num_bins_used']}/{initA['num_bins_total']}", + ]) + if "reference_file" in initA: + lines.append(f" ref = {initA['reference_file']}") + return "\n".join(lines) + + def _debias_init_diagnostics_txt(self, md): + stg = md.get("deBias_stageC", {}) + lines = [ + "Stage (c) init:", + " A = Stage (b) A_hat", + " B = Stage (b) B_hat", + " support = selected_mask + diag", + ] + if "num_dale_active" in stg: + lines.append(f" Dale edges = {stg['num_dale_active']}") + if "num_free_active" in stg: + lines.append(f" free A params = {stg['num_free_active']}") + return "\n".join(lines) + + def A_fitted(self, fitD, md, single_rates, figId="b"): + """Fitted A summary using only arrays stored in the prismEM file.""" + A_init = np.asarray(fitD["A_init"]) + A_hat = np.asarray(fitD["A_hat"]) + assert A_init.shape == A_hat.shape and A_init.ndim == 2, "A_init and A_hat must match" + n_neuron = A_init.shape[0] + x_src = np.arange(n_neuron, dtype=np.int64) + single_rates = np.asarray(single_rates, dtype=np.float64).ravel() + assert single_rates.shape[0] == n_neuron, ( + f"single_rates length {single_rates.shape[0]} != N={n_neuron}" + ) + + num_exc = None + dale_conf = md.get("dale_conf") + if dale_conf is not None: + num_exc = dale_conf.get("num_excite") + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(18, 8)) + gs = fig.add_gridspec(2, 4, hspace=0.55, wspace=0.45) + + if is_stage_c_metadata(md): + init_label = "A_debias_init" + fit_label = "A_debias" + init_txt = self._debias_init_diagnostics_txt(md) + else: + init_label = "A_init" + fit_label = "A_hat" + init_txt = self._A_init_diagnostics_txt(md) + rows = [ + (A_init, init_label, init_txt), + (A_hat, fit_label, None), + ] + for row, (A, label, diag_txt) in enumerate(rows): + _, n_diag, n_off = self._count_A_edges(A) + cnt_src, sum_src = self._source_nz_edge_stats(A) + + ax = fig.add_subplot(gs[row, 0]) + self._draw_A_matrix(ax, A, f"{label}, nEdges={n_diag}+{n_off}", num_exc=num_exc) + + ax = fig.add_subplot(gs[row, 1]) + self._draw_A_offdiag_hist(ax, A, f"{label} off-diagonal", diagnostics_txt=diag_txt) + + ax = fig.add_subplot(gs[row, 2]) + ax.plot(x_src, cnt_src, color="tab:blue", linewidth=1.0) + med_cnt = int(np.median(cnt_src)) + ax.axhline(med_cnt, color="green", ls="--", lw=1.0, alpha=0.9) + ax.text( + 0.02, 0.95, f"median={med_cnt:d}", transform=ax.transAxes, + va="top", ha="left", fontsize=9, color="green", + ) + ax.set_title(f"{label}: # non-zero edges") + ax.set_xlabel("source neuron index (column)") + ax.set_ylabel("# outgoing non-zero edges") + ax.grid(True, alpha=0.35) + if row == 1: + self._overlay_single_rates(ax, single_rates, x_src) + + ax = fig.add_subplot(gs[row, 3]) + ax.plot(x_src, sum_src, color="tab:green", linewidth=1.0) + ax.axhline(0.0, color="k", ls="--", lw=1.0, alpha=0.9) + ax.set_title(f"{label}: sum non-zero edges") + ax.set_xlabel("source neuron index (column)") + ax.set_ylabel("sum outgoing edge value") + ax.grid(True, alpha=0.35) + if row == 1: + self._overlay_single_rates(ax, single_rates, x_src) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle(f"{title_prefix}, {fit_label} fitted summary: neur. freq. sorted", fontsize=13) + fig.tight_layout(rect=[0, 0, 1, 0.95]) + + def _state_seq_window(self, trainMD, time_range_sec): + dt = float(trainMD["time_step_sec"]) + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + + assert time_range_sec is not None, "state sequence plots require time_range_sec" + t0_req, t1_req = [float(x) for x in time_range_sec] + if t1_req < t0_req: + t0_req, t1_req = t1_req, t0_req + b0 = max(t0_bin, int(np.floor(t0_req / dt))) + b1 = min(t1_bin, int(np.floor(t1_req / dt))) + if b1 <= b0: + raise ValueError("Requested --time_range_sec leaves no bins in fitted window") + + i0 = b0 - t0_bin + i1 = b1 - t0_bin + t_bins = np.arange(b0, b1 + 1, dtype=np.float64) * dt + t_pairs = np.arange(b0, b1, dtype=np.float64) * dt + return i0, i1, t_bins, t_pairs + + def _fit_loss_time(self, fitD, md, p0, p1): + if "loss_nll_time" in fitD and "loss_l2_time" in fitD: + return ( + np.asarray(fitD["loss_nll_time"])[p0:p1], + np.asarray(fitD["loss_l2_time"])[p0:p1], + ) + eval_f = md.get("eval_f") + if eval_f is not None and "loss_nll_time" in eval_f and "loss_l2_time" in eval_f: + return ( + np.asarray(eval_f["loss_nll_time"])[p0:p1], + np.asarray(eval_f["loss_l2_time"])[p0:p1], + ) + raise AssertionError( + "state_seq_fit requires precomputed loss_nll_time/loss_l2_time in the input npz; " + "move time-loss computation upstream" + ) + + def state_seq_fit(self, fitD, md, figId="c", time_range_sec=None): + from matplotlib.ticker import MaxNLocator + + trainMD = real_fit_metadata(md)["train"] + i0, i1, t_bins, t_pairs = self._state_seq_window(trainMD, time_range_sec) + p0, p1 = i0, i1 + + S_hat = np.asarray(fitD["S_hat"])[i0:i1 + 1] + C_hat = np.asarray(fitD["c_hat"])[i0:i1 + 1] + S_hat_CL = np.asarray(fitD["S_hat_CL"])[i0:i1 + 1] + nll_t, l2_t = self._fit_loss_time(fitD, md, p0, p1) + + eval_f = md.get("eval_f", {}) + acc_txt = f"acc={float(eval_f['acc']):.3f}" if "acc" in eval_f else "fit only" + x0, x1 = float(t_bins[0]), float(t_bins[-1]) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(12, 7.2)) + gs = fig.add_gridspec( + 4, 3, + height_ratios=[0.95, 0.22, 1.15, 0.70], + hspace=0.78, + wspace=0.38, + ) + + ax = fig.add_subplot(gs[0, 0]) + occ = C_hat.mean(axis=0) + x = np.arange(C_hat.shape[1], dtype=np.int64) + ax.bar(x, occ, color="tab:blue", alpha=0.65) + ax.set(title="Mean occupancy", xlabel="state", ylabel="mean c") + ax.set_ylim(0.0, 1.0) + ax.set_xticks(x) + ax.grid(True, alpha=0.3) + ax.set_box_aspect(0.9) + + ax = fig.add_subplot(gs[0, 1]) + ax.axis("off") + + ax = fig.add_subplot(gs[0, 2]) + cl = np.asarray(S_hat_CL, dtype=np.float64) + ax.hist(cl, bins=30, color="tab:purple", alpha=0.7) + ax.axvline(np.mean(cl), color="k", linestyle="--", linewidth=1.0) + ax.set(title=f"Confidence ({acc_txt})", xlabel="S_hat_CL", ylabel="count") + ax.grid(True, alpha=0.3) + ax.set_box_aspect(0.9) + + ax = fig.add_subplot(gs[2, :]) + line_s, = ax.plot(t_bins, S_hat, color="k", linewidth=2.0, label="S_hat", zorder=4) + lo = np.clip(S_hat - S_hat_CL, 0.0, float(C_hat.shape[1] - 1)) + hi = np.clip(S_hat + S_hat_CL, 0.0, float(C_hat.shape[1] - 1)) + fill = ax.fill_between(t_bins, lo, hi, color="gray", alpha=0.25, label="S_hat_CL", zorder=3) + ax.set(title=f"Fit ({acc_txt})", xlabel="time (s)", ylabel="state") + ax.set_xlim(x0, x1) + ax.set_ylim(-0.03, float(C_hat.shape[1] - 1) + 0.03) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + c_lines = [] + for m in range(C_hat.shape[1]): + line, = ax2.plot(t_bins, C_hat[:, m], linewidth=0.8, alpha=0.85, label=f"C_hat[{m}]") + c_lines.append(line) + ax2.set_ylabel("C_hat") + ax2.set_ylim(ax.get_ylim()) + handles = [line_s, fill] + c_lines + labels = [h.get_label() for h in handles] + ax.legend( + handles, labels, + loc="lower center", bbox_to_anchor=(0.5, 1.18), + ncol=len(handles), fontsize=8, frameon=True, + ) + + ax = fig.add_subplot(gs[3, :]) + ax.plot(t_pairs, nll_t, color="tab:blue", linewidth=0.9, label="nll") + ax.plot(t_pairs, l2_t, color="tab:orange", linewidth=0.9, label="l2") + ax.set_yscale("log") + ax.set(title="Loss (time)", xlabel="time (s)", ylabel="value") + ax.set_xlim(x0, x1) + ax.grid(True, alpha=0.35) + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), ncol=2, fontsize=8) + + ax2 = ax.twinx() + ax2.plot(t_bins, S_hat, color="k", linewidth=0.8, alpha=0.65, label="S_hat") + ax2.set_ylabel("state") + ax2.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), fontsize=8) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle( + f"{title_prefix}, fitted state sequence: lambda2={trainMD['lambda2']}, " + f"lr={trainMD['lr_estep']}, pgd_iter={trainMD['pgd_iter']}", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.94]) + + def spike_bursts(self, rebD, md, figId="d", time_range_sec=None): + from PlotterBioExp import Plotter as PlotterBioExp + + if time_range_sec is None: + time_range_sec = md["plot"]["time_rangeLR"] + clip = clip_rebD_time(rebD, time_range_sec) + rate_thr = float(rebD["rate_thres2"]) + high_chan = np.asarray(rebD["highChanCnt"])[clip["itL"]:clip["itR"]] + timeV = clip["timeV"] + Tbin = clip["Tbin"] + tL, tR = clip["tL"], clip["tR"] + nchan = clip["nchan"] + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(16, 11)) + gs = fig.add_gridspec(4, 1, height_ratios=[0.15, 0.15, 0.54, 0.01]) + + ax = fig.add_subplot(gs[0, 0]) + ax.bar(timeV + Tbin * 0.5, high_chan, width=Tbin, + color="forestgreen", align="center", alpha=0.7) + ax.set( + ylabel="num neurons", + title=f"num neurons with instantaneous rate > thres={rate_thr:.0f} Hz, nchan={nchan}", + ) + ax.set_xlim(tL, tR) + ax.grid(True, alpha=0.35) + + ax = fig.add_subplot(gs[1, 0]) + PlotterBioExp.plot_bioexp_sum_rate(self, ax, rebD, clip) + + ax = fig.add_subplot(gs[2, 0]) + PlotterBioExp.plot_bioexp_heatmap(self, fig, ax, rebD, md["short_name"], clip) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle( + f"{title_prefix}, spike bursts: T=[{float(tL):g}, {float(tR):g}] s", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.94]) + + def _node_Ahat_topology_ax( + self, ax, loc_x, loc_y, A_hat, sign=None, neuron_type=None, + node_size=16, source_select=None, und_color="salmon", + ): + ax.set_facecolor("white") + if neuron_type is None: + ax.scatter( + loc_x, loc_y, s=node_size, marker="o", facecolors="white", + edgecolors="k", linewidths=0.4, alpha=0.95, zorder=3, + ) + else: + neuron_type = np.asarray(neuron_type) + node_color = np.full(loc_x.shape, und_color, dtype=object) + node_color[neuron_type > 0] = "magenta" + node_color[neuron_type < 0] = "#39FF14" + ax.scatter( + loc_x, loc_y, s=node_size, marker="o", c=node_color, + edgecolors="k", linewidths=0.4, alpha=0.95, zorder=3, + ) + + edge_mask = A_hat != 0 + np.fill_diagonal(edge_mask, False) + edge_color = None + if sign == "pos": + edge_mask &= A_hat > 0 + edge_color = "red" + elif sign == "neg": + edge_mask &= A_hat < 0 + edge_color = "blue" + + post_idx, pre_idx = np.where(edge_mask) + n_total = int(post_idx.size) + if source_select is not None: + source_select = np.asarray(source_select, dtype=bool) + draw_mask = source_select[pre_idx] + post_draw = post_idx[draw_mask] + pre_draw = pre_idx[draw_mask] + else: + post_draw = post_idx + pre_draw = pre_idx + + for i_post, j_pre in zip(post_draw, pre_draw): + color = edge_color + if color is None: + color = "red" if A_hat[i_post, j_pre] > 0 else "blue" + ax.plot( + [loc_x[j_pre], loc_x[i_post]], + [loc_y[j_pre], loc_y[i_post]], + color=color, alpha=0.7, linewidth=0.7, zorder=2, + ) + + ax.set_aspect("equal", adjustable="box") + ax.set_xlabel("node x", fontsize=11) + ax.set_ylabel("node y", fontsize=11) + ax.tick_params(labelsize=9) + ax.grid(True, alpha=0.35) + return n_total, int(post_draw.size) + + def neuron_spatial_edges_split( + self, fitD, nodeD, nodeMD, neuron_type, neuron_Sedge, + maxNeurons=24, figId="f", + ): + node_pos = np.asarray(nodeD["node_positions"], dtype=np.float64) + assert node_pos.ndim == 2 and node_pos.shape[1] == 2, ( + f"node_positions must have shape (N, 2), got {node_pos.shape}" + ) + loc_x = node_pos[:, 0] + loc_y = node_pos[:, 1] + A_prune = self._A_prune_from_fitD(fitD) + n_neur = loc_x.shape[0] + assert A_prune.shape == (n_neur, n_neur), ( + f"A_prune shape {A_prune.shape} does not match {n_neur} neuron locations" + ) + neuron_type = np.asarray(neuron_type, dtype=np.int8) + assert neuron_type.shape[0] == n_neur, "neuron_type length must match neuron locations" + neuron_Sedge = np.asarray(neuron_Sedge, dtype=np.float64) + assert neuron_Sedge.shape[0] == n_neur, "neuron_Sedge length must match neuron locations" + n_exc = int(np.sum(neuron_type > 0)) + n_inh = int(np.sum(neuron_type < 0)) + n_und = int(np.sum(neuron_type == 0)) + + def ranked_source_select(type_value): + idx = np.where(neuron_type == int(type_value))[0] + idx = idx[np.argsort(neuron_Sedge[idx])] + n_side = max(1, int(maxNeurons) // 2) + if idx.size <= int(maxNeurons): + sel = idx + else: + sel = np.unique(np.concatenate([idx[:n_side], idx[-n_side:]])) + source_select = np.zeros((n_neur,), dtype=bool) + source_select[sel] = True + return source_select, int(sel.size) + + exc_sources, n_exc_show = ranked_source_select(1) + inh_sources, n_inh_show = ranked_source_select(-1) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(14, 4.8)) + gs = fig.add_gridspec( + 1, 2, left=0.06, right=0.99, bottom=0.13, top=0.84, wspace=0.18 + ) + + ax = fig.add_subplot(gs[0, 0]) + n_pos, n_pos_show = self._node_Ahat_topology_ax( + ax, loc_x, loc_y, A_prune, sign="pos", + neuron_type=neuron_type, node_size=32, source_select=exc_sources, + und_color="yellow", + ) + ax.set_title( + f"A_prune>0, shown edges {n_pos_show}/{n_pos}, neurons {n_exc_show}/{n_exc}", + fontsize=11, pad=8, + ) + ax.scatter([], [], s=32, marker="o", c="magenta", edgecolors="k", label=f"exc={n_exc}") + ax.scatter([], [], s=32, marker="o", c="#39FF14", edgecolors="k", label=f"inh={n_inh}") + ax.scatter([], [], s=32, marker="o", c="yellow", edgecolors="k", label=f"und={n_und}") + ax.legend(loc="best", fontsize=9) + + ax = fig.add_subplot(gs[0, 1]) + n_neg, n_neg_show = self._node_Ahat_topology_ax( + ax, loc_x, loc_y, A_prune, sign="neg", + neuron_type=neuron_type, node_size=32, source_select=inh_sources, + und_color="yellow", + ) + ax.set_title( + f"A_prune<0, shown edges {n_neg_show}/{n_neg}, neurons {n_inh_show}/{n_inh}", + fontsize=11, pad=8, + ) + ax.scatter([], [], s=32, marker="o", c="magenta", edgecolors="k", label=f"exc={n_exc}") + ax.scatter([], [], s=32, marker="o", c="#39FF14", edgecolors="k", label=f"inh={n_inh}") + ax.scatter([], [], s=32, marker="o", c="yellow", edgecolors="k", label=f"und={n_und}") + ax.legend(loc="best", fontsize=9) + + title_prefix = self.canvas_title_prefix(nodeMD) + fig.suptitle( + f"{title_prefix}, neuron spatial A_prune topology: max neurons per panel={int(maxNeurons)}", + fontsize=12, y=0.97, + ) + + def state_seq_simu(self, fitD, md, figId="m", time_range_sec=None): + from matplotlib.ticker import MaxNLocator + + trainMD = real_fit_metadata(md)["train"] + srec = real_fit_metadata(md)["states_recovery_eval"] + assert "eval_f" in md, ( + "state_seq_simu requires precomputed eval_f in the input metadata; " + "move the old evaluator's eval_em_metrics_time computation upstream" + ) + assert "avg_acc" in srec and "state_acc_cl" in srec, ( + "state_seq_simu requires precomputed states_recovery_eval avg_acc/state_acc_cl; " + "move the old evaluator's eval_true_state_recovery computation upstream" + ) + dt = float(trainMD["time_step_sec"]) + t0_bin = int(trainMD["time_range_bins"][0]) + i0, i1, t_bins, t_pairs = self._state_seq_window(trainMD, time_range_sec) + b0 = t0_bin + i0 + b1 = t0_bin + i1 + p0 = i0 + p1 = i1 + + S_hat = np.asarray(fitD["S_hat"])[i0:i1 + 1] + C_hat = np.asarray(fitD["c_hat"])[i0:i1 + 1] + S_hat_CL = np.asarray(fitD["S_hat_CL"])[i0:i1 + 1] + + S_true = np.asarray(md["S_true"])[b0:b1 + 1] + C_true = np.asarray(md["C_true"])[b0:b1 + 1] + + eval_f = md["eval_f"] + nll_t, l2_t = self._fit_loss_time(fitD, md, p0, p1) + acc = float(eval_f["acc"]) + + x0, x1 = float(t_bins[0]), float(t_bins[-1]) + + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(12, 9)) + gs = fig.add_gridspec( + 5, 3, + height_ratios=[0.95, 0.18, 1.0, 1.0, 0.55], + hspace=0.75, + wspace=0.35, + ) + + ax = fig.add_subplot(gs[2, :]) + ax.plot(t_bins, S_hat, color="k", linewidth=2.5, label="S_hat") + ax.plot(t_bins, S_true, color="lime", linestyle="--", linewidth=1.5, label="S_true") + lo = np.clip(S_hat - S_hat_CL, 0.0, float(C_hat.shape[1] - 1)) + hi = np.clip(S_hat + S_hat_CL, 0.0, float(C_hat.shape[1] - 1)) + ax.fill_between(t_bins, lo, hi, color="gray", alpha=0.3, label="S_hat_CL") + ax.set(title=f"Fit (acc={acc:.3f})", xlabel="time (s)", ylabel="state") + ax.set_xlim(x0, x1) + ax.set_ylim(-0.03, float(C_hat.shape[1] - 1) + 0.03) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax2 = ax.twinx() + for m in range(C_hat.shape[1]): + ax2.plot(t_bins, C_hat[:, m], linewidth=0.8, alpha=0.85, label=f"C_hat[{m}]") + ax2.set_ylabel("C_hat") + ax2.set_ylim(ax.get_ylim()) + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + ax = fig.add_subplot(gs[3, :]) + ax.plot(t_bins, S_true, color="lime", linestyle="--", linewidth=1.5, label="S_true") + ax.set(title="Truth", xlabel="time (s)", ylabel="state") + ax.set_xlim(x0, x1) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax2 = ax.twinx() + for m in range(C_true.shape[1]): + ax2.plot(t_bins, C_true[:, m], linewidth=0.8, alpha=0.85, label=f"C_true[{m}]") + ax2.set_ylabel("C_true") + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + acc_avg = float(srec["avg_acc"]) + state_acc_cl = np.asarray(srec["state_acc_cl"], dtype=np.float64) + acc_ps = state_acc_cl[:, 0] + cl_ps = state_acc_cl[:, 1] + + ax = fig.add_subplot(gs[0, 0]) + occ = C_hat.mean(axis=0) + x = np.arange(C_hat.shape[1], dtype=np.int64) + ax.bar(x, occ, color="tab:blue", alpha=0.65) + ax.set(title="Mean occupancy", xlabel="state", ylabel="mean c") + ax.set_ylim(0.0, 1.0) + ax.set_xticks(x) + ax.grid(True, alpha=0.3) + ax.set_box_aspect(0.9) + + ax = fig.add_subplot(gs[0, 1]) + x = np.arange(acc_ps.size, dtype=np.int64) + ax.bar(x, acc_ps, color="tab:green", alpha=0.7) + ax.errorbar(x, acc_ps, yerr=cl_ps, fmt="o", color="k", capsize=5, markersize=4) + ax.axhline(1.0, color="green", linestyle="--", linewidth=1) + ymin = min(0.5, float(np.nanmin(acc_ps - cl_ps)) - 0.05) + ax.set_ylim(ymin, None) + ax.set(title=f"State accuracy, avr={acc_avg:.3f}", xlabel="state", ylabel="accuracy") + ax.set_xticks(x) + ax.grid(True, alpha=0.3, axis="y") + ylo, yhi = ax.get_ylim() + y_txt = ylo + 0.80 * (yhi - ylo) + for i, v in enumerate(acc_ps): + ax.text(i + 0.3, y_txt, f"{v:.2f}", ha="center", va="center", fontsize=9) + ax.set_box_aspect(0.9) + + ax = fig.add_subplot(gs[0, 2]) + cl = np.asarray(S_hat_CL, dtype=np.float64) + ax.hist(cl, bins=30, color="tab:purple", alpha=0.7) + ax.axvline(np.mean(cl), color="k", linestyle="--", linewidth=1.0) + ax.set(title=f"Confidence (acc={acc_avg:.3f})", xlabel="S_hat_CL", ylabel="count") + ax.grid(True, alpha=0.3) + ax.set_box_aspect(0.9) + + ax = fig.add_subplot(gs[4, :]) + ax.plot(t_pairs, nll_t, color="tab:blue", linewidth=0.9, label="nll") + ax.plot(t_pairs, l2_t, color="tab:orange", linewidth=0.9, label="l2") + ax.set_yscale("log") + ax.set(title="Loss (time)", xlabel="time (s)", ylabel="value") + ax.set_xlim(x0, x1) + ax.grid(True, alpha=0.35) + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), ncol=2, fontsize=8) + ax2 = ax.twinx() + ax2.plot(t_bins, S_true, color="k", linewidth=0.8, alpha=0.6, label="S_true") + ax2.set_ylabel("state") + ax2.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), fontsize=8) + + title_prefix = self.canvas_title_prefix(md) + fig.suptitle( + f"{title_prefix}, state sequence: lambda2={trainMD['lambda2']}, " + f"lr={trainMD['lr_estep']}, pgd_iter={trainMD['pgd_iter']}, " + f"decode_dwell={srec['decode_dwell_sec']}s", + fontsize=12, + ) + fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.965]) diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/PlotterSpikesTrain.py b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterSpikesTrain.py new file mode 100644 index 00000000..a2ba787b --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/PlotterSpikesTrain.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for biological experiment data visualization. + +This module provides specialized plotting capabilities for experimental +neural data analysis. The Plotter class extends PlotterBackbone to create +visualizations tailored for biological neural recordings including: +- Time series plots of neural activity patterns +- Statistical summaries and distribution analysis +- Data quality assessment plots +- Comparative analysis between experimental conditions + +Designed specifically for processing and visualizing data from biological +neural experiments, with automatic adaptation to experimental metadata. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +import numpy as np +import matplotlib.gridspec as gridspec +from matplotlib.ticker import MaxNLocator + +#...!...!.................... +def summary_column(md): + return 'fix-me 32678' + #print(sorted(md)) + rs=md['rate_summary'] + ds=md['data_selector'] + #pprint(md) + txt='dataset: '+md['short_name'] + txt += '\nnum acc neurons: %d' % md['num_neurons'] + txt += '\ndrop neur: %d <%.1f Hz, %d >%.1f Hz' % (ds['drop_neur_by_freq_range'][0],ds['freq_range'][0],ds['drop_neur_by_freq_range'][1],ds['freq_range'][1]) + txt+='\ndata type: %s\ntime_step=%.2f sec'%(md['data_type'],md['time_step_sec']) + txt += '\nMedian rate: %.1f Hz' % rs['median_spike_rate'] + txt += '\nrate range [%.1f, %.1f Hz]' % (rs['min_spike_rate'], rs['max_spike_rate']) + txt += '\nAvg Rate: %.1f±%.1f Hz' % (rs['avg_spike_rate'], rs['std_spike_rate']) + txt += '\nAvg Fano: %.2f±%.2f' % (rs['avg_fano_factor'], rs['std_fano_factor']) + + return txt + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def freq_histo(self,spikeD,md,figId=1): + + tit='dataset '+md['short_name'] + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,4)) + + # Create gridspec with 7:3 width ratio + gs = gridspec.GridSpec(1, 2, width_ratios=[7, 3]) + + dataYield, dataRates = spikeD['spikes'], spikeD['single_rates'] + + # Compute median + median_val = np.median(dataRates) + txtM= f'median : {median_val:.2f} Hz' + + #.... freq per channel + ax = self.plt.subplot(gs[0, 0]) + chanV=np.arange(dataRates.shape[0]) + ax.bar(chanV, dataRates , width=0.8, color='g', align='center', alpha=0.7) + + ax.axhline(median_val, color='red', linestyle='--', linewidth=1) + if md['data_type']=='simDale': + ax.set_yscale('log') + ax.grid() + ax.set(xlabel='input neuron index', ylabel='avr frequency (Hz)',title=tit) + ax.text( ax.get_xlim()[1]*0.1,median_val,txtM, + color='red', ha='center', va='bottom', fontsize=10) + + txt=summary_column(md) + ax.text(0.05, 0.3, txt, transform=ax.transAxes, fontsize=12,color='b') + + #........ freq histo ..... + ax = self.plt.subplot(gs[0, 1]) + ax.hist(dataRates,bins=20) + #ax.set_yscale('log') + ax.set(ylabel='num channels', xlabel='avr frequency (Hz)',title=tit) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + # Add median text annotation + ax.text(median_val+1, ax.get_ylim()[1]*0.5,txtM, + color='red', ha='left', va='bottom', fontsize=10,rotation=90) + ax.text(0.25, 0.3, txt, transform=ax.transAxes, fontsize=11,color='m') + + ax.grid() + +#...!...!.................. + def freq_vs_time(self,rebD,md,figId=2,S_true=None,S_oracle=None): + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,11)) + if S_oracle is None and isinstance(rebD, dict) and 'S_oracle' in rebD: + S_oracle = rebD['S_oracle'] + + tit='dataset '+md['short_name'] + time_step=rebD['time_step2'] + R_sel = md['sel_spect_radius'] + state_tag = md['sel_state'] + + # clip time data for display + tL,tR=md['plot']['time_rangeLR'] + itL,itR=(md['plot']['time_rangeLR']/time_step).astype(int) + ntime_tot = rebD['rate2D'].shape[0] + itL = max(0, min(itL, ntime_tot - 1)) + itR = max(itL + 1, min(itR, ntime_tot)) + print('iTL,R', itL,itR) + + # unpack data and clip the time range + rate2D=rebD['rate2D'][itL:itR] + timeV=rebD['timeV'][itL:itR] + + _,nchan=rate2D.shape + Tbin = time_step + xL = float(timeV[0]) + xR = float(timeV[-1] + Tbin) + if isinstance(rebD, dict) and 'pop_rate_hz' in rebD: + pop_rate_hz = np.asarray(rebD['pop_rate_hz'][itL:itR], dtype=np.float64) + else: + pop_rate_hz = np.sum(rate2D, axis=1) + medRateDisp = float(np.median(rate2D)) + cntAboveMed = np.sum(rate2D > medRateDisp, axis=1) + + tit0='dataset: %s state=%d R=%.3f nchan=%d Tbin=%.2f sec'%(md['short_name'], state_tag, R_sel, nchan, time_step) + + # Layout: top-count trace, heatmap, synchronicity trace, target-state, oracle-state. + gs = fig.add_gridspec(5, 1, height_ratios=[0.14, 0.52, 0.14, 0.10, 0.10], hspace=0.36) + + # ..... top plot + ax = fig.add_subplot(gs[0, 0]) + ax.bar(timeV+Tbin*.5, cntAboveMed, width=Tbin, color='teal', align='center', alpha=0.7) + ax.set(ylabel='neurons > median', title=f'neurons above displayed-time median rate ({medRateDisp:.2f} Hz)') + ax.set_xlim(xL,xR) + ax.tick_params(axis='x', labelbottom=False) + ax.grid() + + # ....... main heatmap + ax = fig.add_subplot(gs[1, 0]) + #cmap='tab20c', 'Oranges' + # - - - Get the colormap and modify it --- + original_cmap = self.plt.get_cmap('tab20c') + custom_cmap = original_cmap.copy() + custom_cmap.set_under('white') + + cax = ax.imshow(rate2D.T, + aspect='auto', + cmap=custom_cmap, # Use our modified colormap + origin='lower', + extent=[xL, xR, 0, nchan], + interpolation='nearest', + vmin=1, # Set the lower bound of the colormap + vmax=rate2D.max()) # Optional: ensure the upper bound is set + + ax.set_ylabel('Neuron index') + ax.set_title(tit0 ) + ax.tick_params(axis='x', labelbottom=False) + cbar = fig.colorbar(cax, ax=ax, orientation='horizontal', pad=0.02, label=f'Instantanous freq (Hz) per neuron over Tbin={time_step:.1f} sec', shrink=0.55) + cbar.ax.tick_params(labelsize=10) + ax.grid() + + # ..... synchronicity (yellow) below heatmap + ax = fig.add_subplot(gs[2, 0]) + ax.bar(timeV+Tbin*.5, pop_rate_hz, width=Tbin, color='gold', align='center', alpha=0.8) + ax.set(ylabel='synchronicity (Hz)', title=f'Network synchronicity from {nchan} neurons, Tbin={Tbin:.2f} sec') + ax.set_xlim(xL,xR) + ax.tick_params(axis='x', labelbottom=True) + ax.grid() + + # ..... target state + ax = fig.add_subplot(gs[3, 0]) + ax.set_xlim(xL, xR) + if S_true is not None: + assert S_true.ndim == 1, f"S_true must be 1D, got shape={S_true.shape}" + dt0 = float(md['time_step_sec']) + t_state = np.arange(S_true.shape[0], dtype=float) * dt0 + sel = (t_state >= xL) & (t_state < xR + dt0) + if np.any(sel): + t_sel = t_state[sel] + s_sel = S_true[sel] + ax.step(t_sel, s_sel, where='post', color='k', linewidth=1.0) + smin = int(np.min(s_sel)) + smax = int(np.max(s_sel)) + ax.set_ylim(smin - 0.5, smax + 0.5) + ax.set_yticks(np.arange(smin, smax + 1, 1)) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.set_ylabel('state S') + ax.set_title(f"target state — {md['short_name']}") + ax.grid() + else: + ax.set_axis_off() + ax.tick_params(axis='x', labelbottom=False) + + # ..... oracle state + ax = fig.add_subplot(gs[4, 0]) + ax.set_xlim(xL, xR) + if S_oracle is not None: + assert S_oracle.ndim == 1, f"S_oracle must be 1D, got shape={S_oracle.shape}" + score_txt = f", avr_score={md['oracle_score']:.3f}" + dt0 = float(md['time_step_sec']) + t_state = np.arange(S_oracle.shape[0], dtype=float) * dt0 + sel = (t_state >= xL) & (t_state < xR + dt0) + if np.any(sel): + t_sel = t_state[sel] + s_sel = S_oracle[sel] + ax.step(t_sel, s_sel, where='post', color='g', linewidth=1.0) + smin = int(np.min(s_sel)) + smax = int(np.max(s_sel)) + ax.set_ylim(smin - 0.5, smax + 0.5) + ax.set_yticks(np.arange(smin, smax + 1, 1)) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.set_ylabel('oracle S') + ax.set_title(f"oracle state{score_txt}") + ax.grid() + else: + ax.set_axis_off() + ax.tick_params(axis='x', labelbottom=True) + ax.set_xlabel('Time (s)') diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/PrismDeBias_Workhorse3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/PrismDeBias_Workhorse3c.py new file mode 100755 index 00000000..317dd4a4 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/PrismDeBias_Workhorse3c.py @@ -0,0 +1,624 @@ +#!/usr/bin/env python3 +"""Stage (c) de-biased active-set fitting for PRISM-EM 3c.""" + +import math +import os +import time + +os.environ.setdefault("TORCH_CPP_LOG_LEVEL", "ERROR") + +import numpy as np +import torch +import torch.nn as nn +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + +from PrismEM_Workhorse3c import ( + _history_dict, + _history_arrays, + is_rank0, + make_optimizer_and_scheduler, + make_pair_loader, + onehot_from_curr_states, + pair_range_for_rank, + run_estep, + run_estep_shard, + viterbi_decode, +) + + +class DebiasActiveSetGLM(nn.Module): + """Poisson GLM with active-vector parameterization for constrained A.""" + + def __init__( + self, + n_neuron, + n_state, + eta_clip, + dale_i, + dale_j, + dale_sign, + free_i, + free_j, + u_init, + v_init, + b_init, + ): + super().__init__() + self.N = int(n_neuron) + self.M = int(n_state) + self.eta_clip = float(eta_clip) + + self.register_buffer("dale_i", torch.as_tensor(dale_i, dtype=torch.long)) + self.register_buffer("dale_j", torch.as_tensor(dale_j, dtype=torch.long)) + self.register_buffer("dale_sign", torch.as_tensor(dale_sign, dtype=torch.float32)) + self.register_buffer("free_i", torch.as_tensor(free_i, dtype=torch.long)) + self.register_buffer("free_j", torch.as_tensor(free_j, dtype=torch.long)) + + self.u = nn.Parameter(torch.as_tensor(u_init, dtype=torch.float32).clone()) + self.v = nn.Parameter(torch.as_tensor(v_init, dtype=torch.float32).clone()) + self.B = nn.Parameter(torch.as_tensor(b_init, dtype=torch.float32).clone()) + + def effective_A(self): + A = torch.zeros((self.N, self.N), dtype=self.B.dtype, device=self.B.device) + if self.u.numel() > 0: + A[self.dale_i, self.dale_j] = self.dale_sign * self.u.pow(2) + if self.v.numel() > 0: + A[self.free_i, self.free_j] = self.v + return A + + def forward(self, y_prev, c, dt): + A = self.effective_A() + eta = y_prev @ A.t() + c @ self.B + return torch.exp(torch.clamp(eta, max=self.eta_clip)) * float(dt) + + +def build_debias_active_set(A_stage, selected_mask, neuron_type, eps_u): + """Build active index arrays and initial u/v from Stage (b) arrays.""" + A_stage = np.asarray(A_stage, dtype=np.float32) + selected_mask = np.asarray(selected_mask, dtype=bool) + neuron_type = np.asarray(neuron_type, dtype=np.int8) + if A_stage.ndim != 2 or A_stage.shape[0] != A_stage.shape[1]: + raise ValueError("A_stage must be square") + if selected_mask.shape != A_stage.shape: + raise ValueError("selected_mask must have same shape as A_stage") + n_neuron = A_stage.shape[0] + if neuron_type.shape[0] != n_neuron: + raise ValueError("neuron_type length must match A_stage columns") + + dale_i = [] + dale_j = [] + dale_sign = [] + free_i = [] + free_j = [] + u_init = [] + v_init = [] + eps_u = float(eps_u) + + for i in range(n_neuron): + for j in range(n_neuron): + if i == j: + free_i.append(i) + free_j.append(j) + v_init.append(float(A_stage[i, j])) + continue + if not bool(selected_mask[i, j]): + continue + src_type = int(neuron_type[j]) + if src_type > 0: + dale_i.append(i) + dale_j.append(j) + dale_sign.append(1.0) + u_init.append(max(math.sqrt(abs(float(A_stage[i, j]))), eps_u)) + elif src_type < 0: + dale_i.append(i) + dale_j.append(j) + dale_sign.append(-1.0) + u_init.append(max(math.sqrt(abs(float(A_stage[i, j]))), eps_u)) + else: + free_i.append(i) + free_j.append(j) + v_init.append(float(A_stage[i, j])) + + return { + "dale_i": np.asarray(dale_i, dtype=np.int64), + "dale_j": np.asarray(dale_j, dtype=np.int64), + "dale_sign": np.asarray(dale_sign, dtype=np.float32), + "free_i": np.asarray(free_i, dtype=np.int64), + "free_j": np.asarray(free_j, dtype=np.int64), + "u_init": np.asarray(u_init, dtype=np.float32), + "v_init": np.asarray(v_init, dtype=np.float32), + "num_dale": int(len(dale_i)), + "num_free": int(len(free_i)), + "num_diag": int(n_neuron), + "num_selected_offdiag": int(np.sum(selected_mask & ~np.eye(n_neuron, dtype=bool))), + } + + +def init_debias_model(active, B_init, eta_clip, ctx): + B_init = np.asarray(B_init, dtype=np.float32) + if B_init.ndim == 1: + B_init = B_init[None, :] + n_state, n_neuron = B_init.shape + model = DebiasActiveSetGLM( + n_neuron, + n_state, + eta_clip, + active["dale_i"], + active["dale_j"], + active["dale_sign"], + active["free_i"], + active["free_j"], + active["u_init"], + active["v_init"], + B_init, + ).to(ctx.device) + if ctx.is_dist: + model = DDP(model, device_ids=[ctx.local_rank]) + return model + + +def project_active_spectral_(mdl, rho_max, correction_strength=1.0): + """Scale u/v so the reconstructed A obeys the spectral-radius limit.""" + with torch.no_grad(): + A = mdl.effective_A() + rho = float(torch.linalg.eigvals(A).abs().max().item()) + alpha = min(1.0, max(0.0, float(correction_strength))) + if rho > float(rho_max) and alpha > 0.0: + hard_scale = float(rho_max) / float(rho) + scale = 1.0 - alpha * (1.0 - hard_scale) + if mdl.u.numel() > 0: + mdl.u.mul_(math.sqrt(scale)) + if mdl.v.numel() > 0: + mdl.v.mul_(scale) + return rho + + +def run_debias_mstep_epoch( + model, + loader, + optimizer, + device, + dt, + rho_max, + rho_every, + apply_rho, + rho_correction_strength, + off_mask, + progress_every_batches=0, + progress_prefix="M-step", +): + model.train() + mdl = model.module if hasattr(model, "module") else model + s_tot = 0.0 + s_nll = 0.0 + eps = 1e-8 + rho_every = max(1, int(rho_every)) + progress_every_batches = int(progress_every_batches) + n_batch = len(loader) + if progress_every_batches > 0: + print("%s start batches=%d" % (progress_prefix, n_batch), flush=True) + + for bi, (yp, yc, cc) in enumerate(loader): + yp = yp.to(device, non_blocking=True) + yc = yc.to(device, non_blocking=True) + cc = cc.to(device, non_blocking=True) + + optimizer.zero_grad(set_to_none=True) + pred = model(yp, cc, dt) + nll = (-yc * torch.log(pred + eps) + pred).mean() + nll.backward() + optimizer.step() + + if apply_rho and (bi % rho_every == 0): + project_active_spectral_(mdl, rho_max, rho_correction_strength) + + s_tot += float(nll.item()) + s_nll += float(nll.item()) + batch_done = bi + 1 + if progress_every_batches > 0 and ( + batch_done % progress_every_batches == 0 or batch_done == n_batch + ): + print( + "%s batch %d/%d nll=%.4e lr=%.2e" + % (progress_prefix, batch_done, n_batch, float(nll.item()), optimizer.param_groups[0]["lr"]), + flush=True, + ) + + nb = max(1, len(loader)) + if dist.is_available() and dist.is_initialized(): + v = torch.tensor([s_tot, s_nll, float(nb)], dtype=torch.float64, device=device) + dist.all_reduce(v, op=dist.ReduceOp.SUM) + s_tot = float(v[0].item()) + s_nll = float(v[1].item()) + nb = int(v[2].item()) + + with torch.no_grad(): + A = mdl.effective_A() + nz = int((A[off_mask] != 0).sum().item()) + rho = float(torch.linalg.eigvals(A).abs().max().item()) + return dict(loss=s_tot / nb, nll=s_nll / nb, l1=0.0, rho=rho, nz=nz) + + +def _decode_and_broadcast(c_hat_np, p_stay, ctx): + if ctx.is_dist: + if is_rank0(ctx): + s_np = viterbi_decode(c_hat_np, p_stay).astype(np.int64, copy=False) + s_t = torch.as_tensor(s_np, dtype=torch.int64, device=ctx.device) + else: + s_t = torch.empty((c_hat_np.shape[0],), dtype=torch.int64, device=ctx.device) + dist.broadcast(s_t, src=0) + return s_t.cpu().numpy() + return viterbi_decode(c_hat_np, p_stay).astype(np.int64, copy=False) + + +def _run_one_estep(Yp_gpu, Yc_gpu, model, c_hat_gpu, dt, eta_clip, args, ctx): + mdl = model.module if hasattr(model, "module") else model + progress_every = getattr(args, "progress_every_pairs", 0) if is_rank0(ctx) else 0 + with torch.no_grad(): + A_cur = mdl.effective_A() + B_cur = mdl.B.data + if ctx.is_dist: + t_pairs = Yp_gpu.shape[0] + p0, p1, n_loc = pair_range_for_rank(t_pairs, ctx.world_size, ctx.rank) + t0 = p0 + 1 + t1 = p1 + 1 + nll_local, n_pair_local = run_estep_shard( + Yp_gpu, + Yc_gpu, + A_cur, + B_cur, + c_hat_gpu, + dt, + eta_clip, + float(args.lambda2), + float(args.lr_estep), + int(args.pgd_iter), + t0, + t1, + progress_every=progress_every, + progress_prefix="E-step", + ) + c_upd = torch.zeros_like(c_hat_gpu) + c_msk = torch.zeros((c_hat_gpu.shape[0], 1), dtype=torch.float32, device=ctx.device) + c_upd[0] = c_hat_gpu[0] + c_msk[0] = 1.0 + if n_loc > 0: + c_upd[t0:t1 + 1] = c_hat_gpu[t0:t1 + 1] + c_msk[t0:t1 + 1] = 1.0 + dist.all_reduce(c_upd, op=dist.ReduceOp.SUM) + dist.all_reduce(c_msk, op=dist.ReduceOp.SUM) + c_avg = c_upd / torch.clamp(c_msk, min=1.0) + c_hat_gpu.copy_(torch.where(c_msk > 0.0, c_avg, c_hat_gpu)) + ev = torch.tensor([nll_local, float(n_pair_local)], dtype=torch.float64, device=ctx.device) + dist.all_reduce(ev, op=dist.ReduceOp.SUM) + return float(ev[0].item() / max(1.0, ev[1].item())) + return run_estep( + Yp_gpu, + Yc_gpu, + A_cur, + B_cur, + c_hat_gpu, + dt, + eta_clip, + float(args.lambda2), + float(args.lr_estep), + int(args.pgd_iter), + progress_every=progress_every, + progress_prefix="E-step", + ) + + +def _rho_correction_for_iter(iter_idx, args): + delay = int(args.delay_iter_4_ArhoMax) + target = int(args.target_iter_4_ArhoMax) + apply_rho = int(iter_idx) > delay + iters_left = max(1, target - int(iter_idx) + 1) + strength = 1.0 / float(iters_left) if apply_rho else 0.0 + return apply_rho, strength + + +def _should_log_m_epoch(epoch_idx, total_epochs, every): + every = int(every) + if every <= 0: + return False + epoch_idx = int(epoch_idx) + total_epochs = int(total_epochs) + return epoch_idx == 1 or epoch_idx == total_epochs or (epoch_idx % every) == 0 + + +def _append_m_history(h, met, optimizer, rho_strength): + h["m_loss_epoch"].append(met["loss"]) + h["m_nll_epoch"].append(met["nll"]) + h["m_l1_epoch"].append(0.0) + h["rho_epoch"].append(met["rho"]) + h["nz_edges_epoch"].append(met["nz"]) + h["learning_rates"].append(float(optimizer.param_groups[0]["lr"])) + h["rho_correction_strength_epoch"].append(float(rho_strength)) + + +def run_debias_fit(spikes, spikeMD, A_stage, B_stage, selected_mask, + neuron_type, c_stageB, S_stageB, args, ctx): + """Run Stage (c) active-set de-biased fitting on all distributed ranks.""" + spikes = np.asarray(spikes) + if spikes.ndim != 2 or spikes.shape[0] < 2: + raise ValueError("spikes must have shape (T,N) with at least two bins") + T_full, N = spikes.shape + M = int(args.num_states) + dt = float(spikeMD["time_step_sec"]) + eta_clip = float(spikeMD["poisson_eta_clip"]) + p_stay = math.exp(-dt / float(args.decode_dwell_sec)) + do_log = is_rank0(ctx) and int(args.verb) > 0 + progress_batches = getattr(args, "progress_every_batches", 0) if is_rank0(ctx) else 0 + progress_epochs = getattr(args, "progress_every_epochs", 10) + if do_log: + print( + "deBias start T=%d N=%d M=%d dt=%.6g device=%s " + "state_mode=%s m_epochs=%d num_debias_iters=%d batch=%d" + % ( + T_full, + N, + M, + dt, + ctx.device, + args.state_mode, + int(args.m_epochs), + int(args.num_debias_iters), + int(args.batch_size), + ), + flush=True, + ) + + c_stageB = np.asarray(c_stageB, dtype=np.float32) + S_stageB = np.asarray(S_stageB, dtype=np.int64) + if c_stageB.shape != (T_full, M): + raise ValueError("c_stageB shape must match spike window and num_states") + if S_stageB.shape[0] != T_full: + raise ValueError("S_stageB length must match spike window") + + active = build_debias_active_set(A_stage, selected_mask, neuron_type, args.eps_u) + if do_log: + print( + "deBias active set: dale=%d free=%d diag=%d selected_offdiag=%d" + % ( + active["num_dale"], + active["num_free"], + active["num_diag"], + active["num_selected_offdiag"], + ), + flush=True, + ) + model = init_debias_model(active, B_stage, eta_clip, ctx) + mdl = model.module if hasattr(model, "module") else model + A_init = mdl.effective_A().detach().cpu().numpy().astype(np.float32) + if do_log: + print( + "deBias model initialized A_init_shape=%s B_shape=%s" + % (A_init.shape, tuple(mdl.B.shape)), + flush=True, + ) + + yp_np = spikes[:-1].astype(np.float32) + yc_np = spikes[1:].astype(np.float32) + Yp_gpu = torch.tensor(yp_np, dtype=torch.float32, device=ctx.device) + Yc_gpu = torch.tensor(yc_np, dtype=torch.float32, device=ctx.device) + c_hat_gpu = torch.tensor(c_stageB, dtype=torch.float32, device=ctx.device) + + c_pairs_np = onehot_from_curr_states(S_stageB, M) + loader, sampler = make_pair_loader(yp_np, yc_np, c_pairs_np, args, ctx, shuffle=True) + if do_log: + print( + "deBias tensors ready Yp=%s Yc=%s c_hat=%s loader_batches=%d" + % ( + tuple(Yp_gpu.shape), + tuple(Yc_gpu.shape), + tuple(c_hat_gpu.shape), + len(loader), + ), + flush=True, + ) + + if str(args.state_mode) == "locked": + total_m_epochs = int(args.warmup_locked_epochs) + int(args.m_epochs) + else: + total_m_epochs = int(args.warmup_locked_epochs) + int(args.num_debias_iters) * int(args.m_epochs) + total_m_epochs = max(1, total_m_epochs) + optimizer, scheduler = make_optimizer_and_scheduler( + model, args, total_m_epochs, decay_from_epoch=0 + ) + + off_mask = ~torch.eye(N, dtype=torch.bool, device=ctx.device) + h = _history_dict() + m_epoch_global = 0 + t_start = time.time() + if do_log: + print( + "deBias optimizer ready total_m_epochs=%d lr=%.3e progress_every_epochs=%d" + % (total_m_epochs, float(args.lr_mstep), int(progress_epochs)), + flush=True, + ) + + for _ in range(int(args.warmup_locked_epochs)): + m_epoch_global += 1 + epoch_t0 = time.time() + if sampler is not None: + sampler.set_epoch(m_epoch_global) + apply_rho, rho_strength = _rho_correction_for_iter(1, args) + log_m_epoch = do_log and _should_log_m_epoch(m_epoch_global, total_m_epochs, progress_epochs) + met = run_debias_mstep_epoch( + model, loader, optimizer, ctx.device, dt, args.rho_max, + args.prescale_m_step_4_ArhoMax, apply_rho, rho_strength, off_mask, + progress_every_batches=progress_batches if log_m_epoch else 0, + progress_prefix=( + "warmup M-epoch %d/%d" + % (m_epoch_global, total_m_epochs) + ), + ) + scheduler.step() + _append_m_history(h, met, optimizer, rho_strength) + if log_m_epoch: + print( + "warmup M-epoch %d/%d done nll=%.4e rho=%.4f nz=%d lr=%.2e ela=%.1fs" + % ( + m_epoch_global, + total_m_epochs, + met["nll"], + met["rho"], + met["nz"], + h["learning_rates"][-1], + time.time() - epoch_t0, + ), + flush=True, + ) + + if str(args.state_mode) == "locked": + for _ in range(int(args.m_epochs)): + m_epoch_global += 1 + epoch_t0 = time.time() + if sampler is not None: + sampler.set_epoch(m_epoch_global) + apply_rho, rho_strength = _rho_correction_for_iter(1, args) + log_m_epoch = do_log and _should_log_m_epoch(m_epoch_global, total_m_epochs, progress_epochs) + met = run_debias_mstep_epoch( + model, loader, optimizer, ctx.device, dt, args.rho_max, + args.prescale_m_step_4_ArhoMax, apply_rho, rho_strength, off_mask, + progress_every_batches=progress_batches if log_m_epoch else 0, + progress_prefix=( + "locked M-epoch %d/%d" + % (m_epoch_global, total_m_epochs) + ), + ) + scheduler.step() + _append_m_history(h, met, optimizer, rho_strength) + if log_m_epoch: + print( + "locked M-epoch %d/%d done nll=%.4e rho=%.4f nz=%d lr=%.2e ela=%.1fs" + % ( + m_epoch_global, + total_m_epochs, + met["nll"], + met["rho"], + met["nz"], + h["learning_rates"][-1], + time.time() - epoch_t0, + ), + flush=True, + ) + S_hat = S_stageB.astype(np.int64, copy=True) + else: + S_hat = S_stageB.astype(np.int64, copy=True) + for deb_iter in range(1, int(args.num_debias_iters) + 1): + te0 = time.time() + if do_log: + print( + "deBias iter %d/%d E-step start" + % (deb_iter, int(args.num_debias_iters)), + flush=True, + ) + e_nll = _run_one_estep( + Yp_gpu, Yc_gpu, model, c_hat_gpu, dt, eta_clip, args, ctx + ) + if do_log: + print( + "deBias iter %d/%d E-step done e_nll=%.4e ela=%.1fs" + % ( + deb_iter, + int(args.num_debias_iters), + e_nll, + time.time() - te0, + ), + flush=True, + ) + h["e_nll_em"].append(e_nll) + c_np_iter = c_hat_gpu.detach().cpu().numpy() + S_hat = _decode_and_broadcast(c_np_iter, p_stay, ctx) + np.copyto(c_pairs_np, onehot_from_curr_states(S_hat, M)) + if do_log: + occ = np.bincount(S_hat, minlength=M).astype(np.float64) / float(max(1, S_hat.size)) + print( + "deBias iter %d decoded state occupancy=%s" + % (deb_iter, np.array2string(occ, precision=4)), + flush=True, + ) + + for _ in range(int(args.m_epochs)): + m_epoch_global += 1 + epoch_t0 = time.time() + if sampler is not None: + sampler.set_epoch(m_epoch_global) + apply_rho, rho_strength = _rho_correction_for_iter(deb_iter, args) + log_m_epoch = do_log and _should_log_m_epoch(m_epoch_global, total_m_epochs, progress_epochs) + met = run_debias_mstep_epoch( + model, loader, optimizer, ctx.device, dt, args.rho_max, + args.prescale_m_step_4_ArhoMax, apply_rho, rho_strength, off_mask, + progress_every_batches=progress_batches if log_m_epoch else 0, + progress_prefix=( + "iter %d/%d M-epoch %d/%d" + % ( + deb_iter, + int(args.num_debias_iters), + m_epoch_global, + total_m_epochs, + ) + ), + ) + scheduler.step() + _append_m_history(h, met, optimizer, rho_strength) + if log_m_epoch: + print( + "iter %d M-epoch %d/%d done nll=%.4e rho=%.4f nz=%d lr=%.2e ela=%.1fs" + % ( + deb_iter, + m_epoch_global, + total_m_epochs, + met["nll"], + met["rho"], + met["nz"], + h["learning_rates"][-1], + time.time() - epoch_t0, + ), + flush=True, + ) + + if do_log: + print( + "deBias iter %3d/%d E_nll=%.4e M_nll=%.4e rho=%.4f nz=%d lr=%.2e elaT=%.1fs" + % ( + deb_iter, + int(args.num_debias_iters), + e_nll, + met["nll"], + met["rho"], + met["nz"], + h["learning_rates"][-1], + time.time() - te0, + ), + flush=True, + ) + + c_hat_np = c_hat_gpu.detach().cpu().numpy().astype(np.float32) + if str(args.state_mode) == "locked": + S_hat = S_stageB.astype(np.int64, copy=True) + S_hat_CL = (1.0 - c_hat_np.max(axis=1)).astype(np.float32) + A_debias = mdl.effective_A().detach().cpu().numpy().astype(np.float32) + B_debias = mdl.B.detach().cpu().numpy().astype(np.float32) + + hist = _history_arrays(h) + if do_log: + print( + "deBias final tensors ready A=%s B=%s train_elapsed=%.1fs" + % (A_debias.shape, B_debias.shape, time.time() - t_start), + flush=True, + ) + return { + "A_debias": A_debias, + "B_debias": B_debias, + "A_debias_init": A_init, + "B_debias_init": np.asarray(B_stage, dtype=np.float32), + "c_hat": c_hat_np, + "S_hat": np.asarray(S_hat, dtype=np.int64), + "S_hat_CL": S_hat_CL, + "active": active, + "history": hist, + "elapsed_train_sec": float(time.time() - t_start), + } diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/PrismEM_Workhorse3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/PrismEM_Workhorse3c.py new file mode 100755 index 00000000..e3772728 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/PrismEM_Workhorse3c.py @@ -0,0 +1,1015 @@ +#!/usr/bin/env python3 +"""Reusable PRISM-EM training workhorse for full EM and FDR bagging.""" + +import math +import os +import time +from types import SimpleNamespace + +os.environ.setdefault("TORCH_CPP_LOG_LEVEL", "ERROR") + +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import Dataset, DataLoader +from torch.utils.data.distributed import DistributedSampler + +from Util_PrismEM import init_states_vs_time, init_B_from_spikes, init_A_from_spikes + + +class DistContext: + def __init__(self, is_dist, rank, world_size, local_rank, device): + self.is_dist = is_dist + self.rank = rank + self.world_size = world_size + self.local_rank = local_rank + self.device = device + + +class SwitchingPoissonGLM(nn.Module): + """Shared A (N x N), per-state B (M x N), mixed by state coefficients.""" + + def __init__(self, N, M, eta_clip): + super().__init__() + self.N = N + self.M = M + self.eta_clip = eta_clip + self.A = nn.Parameter(torch.randn(N, N) * 0.1) + self.B = nn.Parameter(torch.randn(M, N) * 0.1) + + def forward(self, y_prev, c, dt): + eta = y_prev @ self.A.t() + c @ self.B + return torch.exp(torch.clamp(eta, max=self.eta_clip)) * dt + + +class PairDataset(Dataset): + """(y_prev, y_curr, c) triplets; c may be updated in-place between epochs.""" + + def __init__(self, y_prev, y_curr, c): + self.y_prev = y_prev + self.y_curr = y_curr + self.c = c + + def __len__(self): + return self.y_prev.shape[0] + + def __getitem__(self, idx): + return ( + torch.from_numpy(self.y_prev[idx]).float(), + torch.from_numpy(self.y_curr[idx]).float(), + torch.from_numpy(self.c[idx]).float(), + ) + + +def add_prism_em_args(parser, include_time_range=True): + """Add ordinary PRISM-EM training arguments to an argparse parser.""" + parser.add_argument("--dataName", required=True) + parser.add_argument("--basePath", + default="/pscratch/sd/b/balewski/2026_causalNet_tmp3/") + parser.add_argument("--num_states", "-M", type=int, required=True, + help="Number of latent states") + + g = parser.add_argument_group("EM structure") + g.add_argument("--num_em_iters", type=int, default=20, + help="Outer EM iterations") + g.add_argument("--m_epochs", type=int, default=2, + help="M-step Adam epochs per EM iteration") + + g = parser.add_argument_group("E-step") + g.add_argument("--pgd_iter", type=int, default=5, + help="PGD iterations per time bin") + g.add_argument("--lr_estep", type=float, default=0.03, + help="PGD step size") + g.add_argument("--lambda2", type=float, default=2.0, + help="Temporal smoothness weight") + + g = parser.add_argument_group("M-step") + g.add_argument("--lr_mstep", type=float, default=0.003, + help="Adam learning rate (initial)") + g.add_argument("--lr_end_factor", type=float, default=0.1, + help="LR decays to lr_mstep * lr_end_factor") + g.add_argument("--lambda3", type=float, default=0.02, + help="L1 penalty on off-diagonal A") + g.add_argument("--rho_max", type=float, default=0.95, + help="Spectral radius target for soft correction of A") + g.add_argument("--prescale_m_step_4_ArhoMax", type=int, default=120, + help="Spectral projection frequency in M-step batches") + g.add_argument("--delay_em_iter_4_ArhoMax", type=int, nargs="+", default=[3], + help="One or two EM iters: start rho_max enforcement; optional target iter for full correction") + g.add_argument("--delay_em_iter_4_lrDecay", type=int, nargs="+", default=[1], + help="One or two EM iters: start LR decay; optional target iter for final LR") + g.add_argument("--delay_em_iter_4_Aprune", type=int, default=1, + help="EM iter after which L1 proximal pruning starts") + g.add_argument("--batch_size", type=int, default=2048) + + g = parser.add_argument_group("initialization") + g.add_argument("--init_states", type=str, default="data", + choices=["data", "rand"], + help="Initial latent-state mode for the real fit") + g.add_argument("--init_B", type=str, default="data", + choices=["data", "rand"], + help="Initial B mode for the real fit") + g.add_argument("--init_A", type=str, default="data", + choices=["data", "rand"], + help="Initial A mode for the real fit") + g.add_argument("--init_A_Tmax", type=int, default=50000, + help="Maximum bins used by OLS/covariance A initialization") + + if include_time_range: + g = parser.add_argument_group("data") + g.add_argument("-T", "--time_range_sec", default=[0.0, 60.0], + nargs=2, type=float, + help="Time window [t0, t1] in seconds") + + g = parser.add_argument_group("decode") + g.add_argument("--decode_dwell_sec", type=float, default=0.3, + help="Viterbi dwell prior tau_dwell") + + g = parser.add_argument_group("misc") + g.add_argument("--seed", type=int, default=42) + g.add_argument("--fitName", type=str, default=None, + help="Optional output fit name") + g.add_argument("-v", "--verb", type=int, default=1) + return parser + + +def normalize_delay_args(args): + if args.delay_em_iter_4_Aprune is None: + args.delay_em_iter_4_Aprune = int(args.num_em_iters * 0.3) + if not hasattr(args, "target_em_iter_4_ArhoMax"): + args.delay_em_iter_4_ArhoMax, args.target_em_iter_4_ArhoMax = normalize_start_target_arg( + args.delay_em_iter_4_ArhoMax, + int(args.num_em_iters * 0.6), + int(args.num_em_iters), + "--delay_em_iter_4_ArhoMax", + ) + if not hasattr(args, "target_em_iter_4_lrDecay"): + args.delay_em_iter_4_lrDecay, args.target_em_iter_4_lrDecay = normalize_start_target_arg( + args.delay_em_iter_4_lrDecay, + int(args.num_em_iters * 0.7), + int(args.num_em_iters), + "--delay_em_iter_4_lrDecay", + ) + return args + + +def normalize_start_target_arg(value, default_start, default_target, name): + if value is None: + start = int(default_start) + target = int(default_target) + elif isinstance(value, (list, tuple)): + if len(value) < 1 or len(value) > 2: + raise ValueError(f"{name} expects 1 or 2 integers") + start = int(value[0]) + target = int(value[1]) if len(value) == 2 else int(default_target) + else: + start = int(value) + target = int(default_target) + + if start < 0: + raise ValueError(f"{name} start must be non-negative") + if target <= start: + raise ValueError(f"{name} target must be greater than start") + if target > int(default_target): + raise ValueError(f"{name} target must be <= num_em_iters ({default_target})") + return start, target + + +def init_distributed(verb=0, program_name=None, rank0_only=False): + is_dist = (int(os.environ.get("WORLD_SIZE", "1")) > 1) or ("RANK" in os.environ) + if verb > 0: + env_rank = os.environ.get("RANK", "?") + env_local = os.environ.get("LOCAL_RANK", "?") + env_world = os.environ.get("WORLD_SIZE", "1") + tag = program_name or "PRISM" + if (not rank0_only) or env_rank in ("0", "?"): + print( + f"[{tag} pre-init rank_env={env_rank} local_env={env_local}/{env_world}] " + f"is_dist={is_dist} MASTER_ADDR={os.environ.get('MASTER_ADDR', '?')} " + f"MASTER_PORT={os.environ.get('MASTER_PORT', '?')} " + f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}", + flush=True, + ) + if is_dist: + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + rank = 0 + world_size = 1 + local_rank = 0 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + if verb > 0 and ((not rank0_only) or rank == 0): + tag = program_name or "PRISM" + dev_name = torch.cuda.get_device_name(local_rank) if device.type == "cuda" else "cpu" + print( + f"[{tag} rank={rank}/{world_size} local_rank={local_rank}] " + f"distributed initialized device={device} cuda_available={torch.cuda.is_available()} " + f"device_name={dev_name}", + flush=True, + ) + return DistContext(is_dist, rank, world_size, local_rank, device) + + +def cleanup_distributed(ctx): + if ctx.is_dist and dist.is_initialized(): + dist.destroy_process_group() + + +def seed_everything(seed): + torch.manual_seed(int(seed)) + np.random.seed(int(seed)) + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision("high") + + +def is_rank0(ctx): + return ctx.rank == 0 + + +def barrier(ctx): + if ctx.is_dist: + if ctx.device.type == "cuda": + dist.barrier(device_ids=[ctx.local_rank]) + else: + dist.barrier() + + +def runtime_summary(ctx): + if ctx.device.type == "cuda": + dev_idx = ctx.local_rank if ctx.is_dist else torch.cuda.current_device() + dev_name = torch.cuda.get_device_name(dev_idx) + else: + dev_name = "cpu" + return ( + f"distributed={ctx.is_dist} rank={ctx.rank}/{ctx.world_size} " + f"local_rank={ctx.local_rank} device={ctx.device} " + f"cuda_available={torch.cuda.is_available()} device_name={dev_name}" + ) + + +def broadcast_object(obj, ctx, src=0): + if not ctx.is_dist: + return obj + box = [obj] + dist.broadcast_object_list(box, src=src) + return box[0] + + +def _torch_dtype(np_dtype): + dt = np.dtype(np_dtype) + if dt == np.dtype("float32"): + return torch.float32 + if dt == np.dtype("float64"): + return torch.float64 + if dt == np.dtype("int64"): + return torch.int64 + if dt == np.dtype("int32"): + return torch.int32 + if dt == np.dtype("int16"): + return torch.int16 + if dt == np.dtype("int8"): + return torch.int8 + if dt == np.dtype("uint8"): + return torch.uint8 + if dt == np.dtype("bool"): + return torch.bool + raise TypeError(f"Unsupported broadcast dtype: {dt}") + + +def broadcast_array(arr, ctx, src=0): + if not ctx.is_dist: + return arr + if is_rank0(ctx): + arr = np.asarray(arr) + info = (tuple(arr.shape), str(arr.dtype)) + else: + info = None + shape, dtype_s = broadcast_object(info, ctx, src=src) + dtype = np.dtype(dtype_s) + torch_dtype = _torch_dtype(dtype) + if is_rank0(ctx): + tens = torch.as_tensor(arr, dtype=torch_dtype, device=ctx.device).contiguous() + else: + tens = torch.empty(shape, dtype=torch_dtype, device=ctx.device) + dist.broadcast(tens, src=src) + return tens.cpu().numpy() + + +def broadcast_optional_array(arr, ctx, src=0): + if not ctx.is_dist: + return arr + has_arr = bool(arr is not None) if is_rank0(ctx) else False + has_arr = broadcast_object(has_arr, ctx, src=src) + if not has_arr: + return None + return broadcast_array(arr, ctx, src=src) + + +def pair_range_for_rank(t_pairs, world_size, rank): + base = t_pairs // world_size + rem = t_pairs % world_size + n_loc = base + (1 if rank < rem else 0) + p0 = rank * base + min(rank, rem) + p1 = p0 + n_loc - 1 + return p0, p1, n_loc + + +def project_to_simplex(v): + if v.numel() == 1: + return torch.ones_like(v) + if v.numel() == 2: + x0 = torch.clamp(0.5 * (v[0] - v[1] + 1.0), min=0.0, max=1.0) + return torch.stack((x0, 1.0 - x0)) + u, _ = torch.sort(v, descending=True) + cssv = torch.cumsum(u, dim=0) - 1.0 + ind = torch.arange(1, v.numel() + 1, device=v.device, dtype=v.dtype) + cond = u - cssv / ind > 0 + if torch.any(cond): + rho = torch.nonzero(cond, as_tuple=False)[-1, 0] + theta = cssv[rho] / (rho + 1.0) + else: + theta = cssv[-1] / v.numel() + return torch.clamp(v - theta, min=0.0) + + +def run_estep(Y_prev, Y_curr, A, B, c_hat, dt, eta_clip, lambda2, lr, pgd_iter, + progress_every=0, progress_prefix="E-step"): + T_eff = c_hat.shape[0] + log_dt = math.log(dt) + c_prev = c_hat[0].clone() + nll_sum = torch.zeros((), dtype=torch.float32, device=c_hat.device) + progress_every = int(progress_every) + + for t in range(1, T_eff): + y_p = Y_prev[t - 1] + y_c = Y_curr[t - 1] + Z = (A @ y_p)[None, :] + B + c_t = c_hat[t].clone() + for _ in range(pgd_iter): + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + g = Z @ (lam - y_c) + 2.0 * lambda2 * (c_t - c_prev) + c_t = project_to_simplex(c_t - lr * g) + + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + nll_sum = nll_sum + (lam - y_c * (eta_c + log_dt)).sum() + c_hat[t] = c_t + c_prev = c_t + if progress_every > 0 and (t % progress_every == 0 or t == T_eff - 1): + print( + "%s pair %d/%d nll_avg=%.4e" + % (progress_prefix, t, T_eff - 1, float((nll_sum / max(1, t)).item())), + flush=True, + ) + + return float((nll_sum / max(1, T_eff - 1)).item()) + + +def run_estep_shard(Y_prev, Y_curr, A, B, c_hat, + dt, eta_clip, lambda2, lr, pgd_iter, t0, t1, + progress_every=0, progress_prefix="E-step shard"): + if t1 < t0: + return 0.0, 0 + log_dt = math.log(dt) + c_prev = c_hat[t0 - 1].clone() + nll_sum = torch.zeros((), dtype=torch.float32, device=c_hat.device) + n_pairs = 0 + progress_every = int(progress_every) + n_total = int(t1 - t0 + 1) + + for t in range(t0, t1 + 1): + y_p = Y_prev[t - 1] + y_c = Y_curr[t - 1] + Z = (A @ y_p)[None, :] + B + c_t = c_hat[t].clone() + for _ in range(pgd_iter): + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + g = Z @ (lam - y_c) + 2.0 * lambda2 * (c_t - c_prev) + c_t = project_to_simplex(c_t - lr * g) + + eta = c_t @ Z + eta_c = torch.clamp(eta, max=eta_clip) + lam = torch.exp(eta_c) * dt + nll_sum = nll_sum + (lam - y_c * (eta_c + log_dt)).sum() + c_hat[t] = c_t + c_prev = c_t + n_pairs += 1 + if progress_every > 0 and (n_pairs % progress_every == 0 or n_pairs == n_total): + print( + "%s local_pair %d/%d global_t=%d nll_avg=%.4e" + % (progress_prefix, n_pairs, n_total, t, float((nll_sum / max(1, n_pairs)).item())), + flush=True, + ) + + return float(nll_sum.item()), n_pairs + + +def offdiag_soft_threshold_(A, lr, lam): + if lam <= 0: + return + with torch.no_grad(): + n = A.shape[0] + mask = ~torch.eye(n, dtype=torch.bool, device=A.device) + thresh = lr * lam + v = A[mask] + A[mask] = v.sign() * (v.abs() - thresh).clamp(min=0.0) + + +def enforce_spectral_radius_(A, rho_max, correction_strength=1.0): + with torch.no_grad(): + rho = torch.linalg.eigvals(A).abs().max().item() + alpha = min(1.0, max(0.0, float(correction_strength))) + if rho > rho_max and alpha > 0.0: + hard_scale = float(rho_max) / float(rho) + scale = 1.0 - alpha * (1.0 - hard_scale) + A.mul_(scale) + return rho + + +def sync_AB_via_rank0_avg_then_broadcast_(mdl, rho_max, correction_strength=1.0): + if not (dist.is_available() and dist.is_initialized()): + enforce_spectral_radius_(mdl.A, rho_max, correction_strength) + return + + rank = dist.get_rank() + world = float(dist.get_world_size()) + with torch.no_grad(): + A_buf = mdl.A.data.clone() + B_buf = mdl.B.data.clone() + dist.reduce(A_buf, dst=0, op=dist.ReduceOp.SUM) + dist.reduce(B_buf, dst=0, op=dist.ReduceOp.SUM) + if rank == 0: + A_buf.div_(world) + B_buf.div_(world) + enforce_spectral_radius_(A_buf, rho_max, correction_strength) + dist.broadcast(A_buf, src=0) + dist.broadcast(B_buf, src=0) + mdl.A.data.copy_(A_buf) + mdl.B.data.copy_(B_buf) + + +def run_mstep_epoch(model, loader, optimizer, device, dt, + l1_wt, lambda3, rho_max, rho_every, + apply_prune, apply_rho, rho_correction_strength, + off_mask): + model.train() + mdl = model.module if hasattr(model, "module") else model + s_tot = s_nll = s_l1 = 0.0 + eps = 1e-8 + + for bi, (yp, yc, cc) in enumerate(loader): + yp = yp.to(device, non_blocking=True) + yc = yc.to(device, non_blocking=True) + cc = cc.to(device, non_blocking=True) + + optimizer.zero_grad(set_to_none=True) + pred = model(yp, cc, dt) + nll = (-yc * torch.log(pred + eps) + pred).mean() + if lambda3 > 0: + n = mdl.A.shape[0] + off_abs_mean = (mdl.A.abs() * l1_wt).sum() / float(n * (n - 1)) + l1 = lambda3 * off_abs_mean + else: + l1 = torch.tensor(0.0, device=device) + loss = nll + l1 + loss.backward() + optimizer.step() + + if apply_prune and lambda3 > 0: + offdiag_soft_threshold_(mdl.A, optimizer.param_groups[0]["lr"], lambda3) + if apply_rho and (bi % rho_every == 0): + sync_AB_via_rank0_avg_then_broadcast_( + mdl, rho_max, rho_correction_strength + ) + + s_tot += loss.item() + s_nll += nll.item() + s_l1 += l1.item() + + nb = max(1, len(loader)) + if dist.is_available() and dist.is_initialized(): + v = torch.tensor([s_tot, s_nll, s_l1, float(nb)], + dtype=torch.float64, device=device) + dist.all_reduce(v, op=dist.ReduceOp.SUM) + s_tot = float(v[0].item()) + s_nll = float(v[1].item()) + s_l1 = float(v[2].item()) + nb = int(v[3].item()) + + with torch.no_grad(): + nz = int((mdl.A[off_mask] != 0).sum().item()) + rho = float(torch.linalg.eigvals(mdl.A).abs().max().item()) + return dict(loss=s_tot / nb, nll=s_nll / nb, l1=s_l1 / nb, rho=rho, nz=nz) + + +def viterbi_decode(c_hat, p_stay): + eps = 1e-12 + T, M = c_hat.shape + if M == 1: + return np.zeros(T, dtype=np.int64) + p_sw = (1.0 - p_stay) / (M - 1) + trans = np.full((M, M), p_sw, dtype=np.float64) + np.fill_diagonal(trans, p_stay) + lt = np.log(np.clip(trans, eps, 1.0)) + le = np.log(np.clip(c_hat, eps, 1.0)) + + delta = np.zeros((T, M), dtype=np.float64) + psi = np.zeros((T, M), dtype=np.int64) + delta[0] = le[0] + for t in range(1, T): + sc = delta[t - 1][:, None] + lt + psi[t] = np.argmax(sc, axis=0) + delta[t] = sc[psi[t], np.arange(M)] + le[t] + + path = np.zeros(T, dtype=np.int64) + path[-1] = int(np.argmax(delta[-1])) + for t in range(T - 2, -1, -1): + path[t] = psi[t + 1, path[t + 1]] + return path + + +def onehot_from_prev_states(S, M): + """Return one-hot rows aligned as S[t-1] for lag pair (Y[t-1], Y[t]).""" + S = np.asarray(S, dtype=np.int64) + if S.ndim != 1 or S.shape[0] < 2: + raise ValueError("S must be 1D with at least two time bins") + eye = np.eye(M, dtype=np.float32) + return eye[S[:-1]] + + +def onehot_from_curr_states(S, M): + """Return one-hot rows aligned as S[t] for lag pair (Y[t-1], Y[t]).""" + S = np.asarray(S, dtype=np.int64) + if S.ndim != 1 or S.shape[0] < 2: + raise ValueError("S must be 1D with at least two time bins") + eye = np.eye(M, dtype=np.float32) + return eye[S[1:]] + + +def onehot_from_pair_states(S_pair, M): + """Return one-hot rows from one state label per supplied lag pair.""" + S_pair = np.asarray(S_pair, dtype=np.int64) + if S_pair.ndim != 1: + raise ValueError("S_pair must be 1D") + eye = np.eye(M, dtype=np.float32) + return eye[S_pair] + + +def make_pair_loader(yp_np, yc_np, c_pairs_np, args, ctx, shuffle=True): + dataset = PairDataset(yp_np, yc_np, c_pairs_np) + if ctx.is_dist: + sampler = DistributedSampler( + dataset, num_replicas=ctx.world_size, rank=ctx.rank, + shuffle=shuffle, drop_last=False + ) + loader = DataLoader( + dataset, batch_size=args.batch_size, shuffle=False, sampler=sampler, + drop_last=False, pin_memory=True, num_workers=0 + ) + else: + sampler = None + loader = DataLoader( + dataset, batch_size=args.batch_size, shuffle=shuffle, + drop_last=False, pin_memory=True, num_workers=0 + ) + return loader, sampler + + +def make_optimizer_and_scheduler(model, args, total_epochs, decay_from_epoch=0, decay_to_epoch=None): + first_param = next(model.parameters()) + fused_ok = bool(first_param.is_cuda) + optimizer = optim.Adam(model.parameters(), lr=args.lr_mstep, fused=fused_ok) + if decay_to_epoch is None: + decay_to_epoch = int(total_epochs) + decay_epochs = max(1, int(decay_to_epoch) - int(decay_from_epoch)) + scheduler = optim.lr_scheduler.LinearLR( + optimizer, start_factor=1.0, end_factor=args.lr_end_factor, + total_iters=decay_epochs, last_epoch=-1 + ) + scheduler._step_count = 1 + return optimizer, scheduler + + +def init_model(N, M, eta_clip, A_init, B_init, ctx): + model = SwitchingPoissonGLM(N, M, eta_clip).to(ctx.device) + if A_init is not None: + with torch.no_grad(): + model.A.copy_(torch.tensor(A_init, dtype=torch.float32, device=ctx.device)) + if B_init is not None: + with torch.no_grad(): + model.B.copy_(torch.tensor(B_init, dtype=torch.float32, device=ctx.device)) + if ctx.is_dist: + model = DDP(model, device_ids=[ctx.local_rank]) + return model + + +def source_type_prune(A_hat): + """Compute source-neuron signs and Dale-style pruned A using columns.""" + A_hat = np.asarray(A_hat, dtype=np.float32) + N = A_hat.shape[0] + A_thr = A_hat.copy() + np.fill_diagonal(A_thr, 0.0) + + neuron_Sedge = A_thr.sum(axis=0) + neuron_type = np.zeros((N,), dtype=np.int8) + neuron_type[neuron_Sedge > 0.0] = 1 + neuron_type[neuron_Sedge < 0.0] = -1 + + A_prune = A_hat.copy() + diag_A = np.diag(A_hat).copy() + exc_cols = neuron_type > 0 + inh_cols = neuron_type < 0 + A_prune[:, exc_cols] = np.where(A_prune[:, exc_cols] > 0, A_prune[:, exc_cols], 0.0) + A_prune[:, inh_cols] = np.where(A_prune[:, inh_cols] < 0, A_prune[:, inh_cols], 0.0) + np.fill_diagonal(A_prune, diag_A) + return A_prune.astype(np.float32), neuron_type, neuron_Sedge.astype(np.float32) + + +def _history_dict(): + return dict( + e_nll_em=[], + m_loss_epoch=[], + m_nll_epoch=[], + m_l1_epoch=[], + rho_epoch=[], + rho_correction_strength_epoch=[], + nz_edges_epoch=[], + learning_rates=[], + ) + + +def _history_arrays(h): + return dict( + e_nll_em=np.asarray(h["e_nll_em"], dtype=np.float64), + m_loss_epoch=np.asarray(h["m_loss_epoch"], dtype=np.float64), + m_nll_epoch=np.asarray(h["m_nll_epoch"], dtype=np.float64), + m_l1_epoch=np.asarray(h["m_l1_epoch"], dtype=np.float64), + rho_epoch=np.asarray(h["rho_epoch"], dtype=np.float64), + rho_correction_strength_epoch=np.asarray(h["rho_correction_strength_epoch"], dtype=np.float64), + nz_edges_epoch=np.asarray(h["nz_edges_epoch"], dtype=np.int64), + learning_rates=np.asarray(h["learning_rates"], dtype=np.float64), + ) + + +def run_full_fit(spikes, spikeMD, single_rates, args, ctx, + time_range_sec=None, time_range_bins=None, + fit_name=None, provenance_update=None): + """Run full E/M training on an in-memory spike matrix on all ranks.""" + args = normalize_delay_args(args) + spikes = np.asarray(spikes) + if spikes.ndim != 2: + raise ValueError("spikes must be shaped (T, N)") + T_full, N = spikes.shape + if T_full < 2: + raise ValueError("Need at least two spike bins") + M = int(args.num_states) + T_pairs = T_full - 1 + dt = float(spikeMD["time_step_sec"]) + eta_clip = float(spikeMD["poisson_eta_clip"]) + p_stay = math.exp(-dt / float(args.decode_dwell_sec)) + + yp_np = spikes[:-1].astype(np.float32) + yc_np = spikes[1:].astype(np.float32) + Yp_gpu = torch.tensor(yp_np, device=ctx.device) + Yc_gpu = torch.tensor(yc_np, device=ctx.device) + + if is_rank0(ctx): + c_init_np, S_init, init_state_md, freq_h1d = init_states_vs_time(spikes, dt, args) + A_seed_np, init_A_md = init_A_from_spikes(spikes, args) + B_seed_np, init_B_md = init_B_from_spikes(spikes, dt, args) + else: + c_init_np = S_init = freq_h1d = None + A_seed_np = B_seed_np = None + init_state_md = init_A_md = init_B_md = None + + c_init_np = broadcast_array(c_init_np, ctx) + S_init = broadcast_array(S_init, ctx) + freq_h1d = broadcast_array(freq_h1d, ctx) + A_seed_np = broadcast_optional_array(A_seed_np, ctx) + B_seed_np = broadcast_optional_array(B_seed_np, ctx) + init_state_md = broadcast_object(init_state_md, ctx) + init_A_md = broadcast_object(init_A_md, ctx) + init_B_md = broadcast_object(init_B_md, ctx) + + c_hat_gpu = torch.tensor(c_init_np, dtype=torch.float32, device=ctx.device) + c_pairs_np = onehot_from_prev_states(S_init, M) + loader, sampler = make_pair_loader(yp_np, yc_np, c_pairs_np, args, ctx, shuffle=True) + + model = init_model(N, M, eta_clip, A_seed_np, B_seed_np, ctx) + mdl = model.module if hasattr(model, "module") else model + A_init_np = mdl.A.detach().cpu().numpy().copy() + B_init_np = mdl.B.detach().cpu().numpy().copy() + B_init_out_np = B_init_np[0] if B_init_np.ndim == 2 and B_init_np.shape[0] > 1 else B_init_np + + total_m_epochs = int(args.num_em_iters) * int(args.m_epochs) + decay_start_m_epoch = int(args.delay_em_iter_4_lrDecay) * int(args.m_epochs) + decay_target_m_epoch = int(args.target_em_iter_4_lrDecay) * int(args.m_epochs) + optimizer, scheduler = make_optimizer_and_scheduler( + model, args, total_m_epochs, + decay_from_epoch=decay_start_m_epoch, + decay_to_epoch=decay_target_m_epoch, + ) + + off_mask = ~torch.eye(N, dtype=torch.bool, device=ctx.device) + l1_wt = torch.ones(N, N, device=ctx.device) + l1_wt[torch.eye(N, dtype=torch.bool, device=ctx.device)] = 0.0 + h = _history_dict() + m_epoch_global = 0 + t_start = time.time() + reported_rho_activation = False + reported_lr_activation = False + + for em in range(1, int(args.num_em_iters) + 1): + apply_prune = em > int(args.delay_em_iter_4_Aprune) + apply_rho = em > int(args.delay_em_iter_4_ArhoMax) + em_iters_left = max(1, int(args.target_em_iter_4_ArhoMax) - em + 1) + rho_correction_strength = 1.0 / float(em_iters_left) if apply_rho else 0.0 + apply_lr_decay = int(args.delay_em_iter_4_lrDecay) < em <= int(args.target_em_iter_4_lrDecay) + + if is_rank0(ctx) and args.verb > 0 and apply_rho and not reported_rho_activation: + print( + f"ArhoMax activated at EM iter {em} " + f"(start_after={int(args.delay_em_iter_4_ArhoMax)}, " + f"target_iter={int(args.target_em_iter_4_ArhoMax)}, " + f"target_rho={float(args.rho_max):.6g})" + ) + reported_rho_activation = True + if is_rank0(ctx) and args.verb > 0 and apply_lr_decay and not reported_lr_activation: + target_lr = float(args.lr_mstep) * float(args.lr_end_factor) + print( + f"lrDecay activated at EM iter {em} " + f"(start_after={int(args.delay_em_iter_4_lrDecay)}, " + f"target_iter={int(args.target_em_iter_4_lrDecay)}, " + f"target_lr={target_lr:.6g})" + ) + reported_lr_activation = True + + te0 = time.time() + with torch.no_grad(): + if ctx.is_dist: + p0, p1, n_loc = pair_range_for_rank(T_pairs, ctx.world_size, ctx.rank) + t0 = p0 + 1 + t1 = p1 + 1 + nll_local, n_pair_local = run_estep_shard( + Yp_gpu, Yc_gpu, mdl.A.data, mdl.B.data, c_hat_gpu, + dt, eta_clip, args.lambda2, args.lr_estep, args.pgd_iter, + t0, t1 + ) + c_upd = torch.zeros_like(c_hat_gpu) + c_msk = torch.zeros((T_full, 1), dtype=torch.float32, device=ctx.device) + c_upd[0] = c_hat_gpu[0] + c_msk[0] = 1.0 + if n_loc > 0: + c_upd[t0:t1 + 1] = c_hat_gpu[t0:t1 + 1] + c_msk[t0:t1 + 1] = 1.0 + dist.all_reduce(c_upd, op=dist.ReduceOp.SUM) + dist.all_reduce(c_msk, op=dist.ReduceOp.SUM) + c_avg = c_upd / torch.clamp(c_msk, min=1.0) + c_hat_gpu.copy_(torch.where(c_msk > 0.0, c_avg, c_hat_gpu)) + ev = torch.tensor([nll_local, float(n_pair_local)], + dtype=torch.float64, device=ctx.device) + dist.all_reduce(ev, op=dist.ReduceOp.SUM) + e_nll = float(ev[0].item() / max(1.0, ev[1].item())) + else: + e_nll = run_estep( + Yp_gpu, Yc_gpu, mdl.A.data, mdl.B.data, c_hat_gpu, + dt, eta_clip, args.lambda2, args.lr_estep, args.pgd_iter + ) + te = time.time() - te0 + h["e_nll_em"].append(e_nll) + + c_hat_np_iter = c_hat_gpu.cpu().numpy() + if ctx.is_dist: + if is_rank0(ctx): + s_iter_np = viterbi_decode(c_hat_np_iter, p_stay).astype(np.int64, copy=False) + s_iter_t = torch.as_tensor(s_iter_np, dtype=torch.int64, device=ctx.device) + else: + s_iter_t = torch.empty((T_full,), dtype=torch.int64, device=ctx.device) + dist.broadcast(s_iter_t, src=0) + s_iter_np = s_iter_t.cpu().numpy() + else: + s_iter_np = viterbi_decode(c_hat_np_iter, p_stay).astype(np.int64, copy=False) + np.copyto(c_pairs_np, onehot_from_prev_states(s_iter_np, M)) + + tm0 = time.time() + for _ in range(int(args.m_epochs)): + m_epoch_global += 1 + if sampler is not None: + sampler.set_epoch(m_epoch_global) + met = run_mstep_epoch( + model, loader, optimizer, ctx.device, dt, + l1_wt, args.lambda3, args.rho_max, + args.prescale_m_step_4_ArhoMax, apply_prune, apply_rho, + rho_correction_strength, off_mask + ) + if apply_lr_decay: + scheduler.step() + h["m_loss_epoch"].append(met["loss"]) + h["m_nll_epoch"].append(met["nll"]) + h["m_l1_epoch"].append(met["l1"]) + h["rho_epoch"].append(met["rho"]) + h["nz_edges_epoch"].append(met["nz"]) + h["learning_rates"].append(float(optimizer.param_groups[0]["lr"])) + h["rho_correction_strength_epoch"].append(rho_correction_strength) + tm = time.time() - tm0 + + if is_rank0(ctx) and args.verb > 0: + n_off = int(off_mask.sum().item()) + sp = 1.0 - met["nz"] / max(1, n_off) + print( + f"EM {em:3d}/{args.num_em_iters} " + f"E_nll={e_nll:.4e}({te:.1f}s) " + f"M_nll={met['nll']:.4e} l1={met['l1']:.4e} " + f"rho={met['rho']:.4f} sp={sp:.3f} nz={met['nz']} " + f"lr={h['learning_rates'][-1]:.2e} ({tm:.1f}s) " + f"tot={time.time() - t_start:.0f}s" + ) + + c_hat_np = c_hat_gpu.cpu().numpy() + S_hat = viterbi_decode(c_hat_np, p_stay) + S_hat_CL = (1.0 - c_hat_np.max(axis=1)).astype(np.float32) + A_hat = mdl.A.detach().cpu().numpy().astype(np.float32) + B_hat = mdl.B.detach().cpu().numpy().astype(np.float32) + A_prune, neuron_type, neuron_Sedge = source_type_prune(A_hat) + + hist = _history_arrays(h) + outD = { + "A_init": A_init_np.astype(np.float32), + "A_hat": A_hat, + "A_prune": A_prune, + "neuron_type": neuron_type, + "neuron_Sedge": neuron_Sedge, + "B_init": B_init_out_np.astype(np.float32), + "B_hat": B_hat, + "freq_h1d": np.asarray(freq_h1d, dtype=np.float32), + "c_init": c_init_np.astype(np.float32), + "c_hat": c_hat_np.astype(np.float32), + "S_init": S_init.astype(np.int64), + "S_hat": S_hat.astype(np.int64), + "S_hat_CL": S_hat_CL, + "single_rates": np.asarray(single_rates), + } + outD.update(hist) + + prov = dict(spikeMD.get("provenance", {})) + if provenance_update: + prov.update(provenance_update) + if fit_name: + prov["EMtrain_file"] = fit_name + + if time_range_sec is None: + time_range_sec = [0.0, float((T_full - 1) * dt)] + if time_range_bins is None: + time_range_bins = [0, T_full - 1] + + outMD = dict(spikeMD) + outMD["fit_type"] = "prismEM" + outMD["train"] = { + "num_em_iters": int(args.num_em_iters), + "m_epochs": int(args.m_epochs), + "total_m_epochs": int(total_m_epochs), + "pgd_iter": int(args.pgd_iter), + "lr_estep": float(args.lr_estep), + "lr_mstep": float(args.lr_mstep), + "lr_end_factor": float(args.lr_end_factor), + "lambda2": float(args.lambda2), + "lambda3": float(args.lambda3), + "rho_max": float(args.rho_max), + "rho_projection_mode": "soft_inverse_em_iters_left", + "prescale_m_step_4_ArhoMax": int(args.prescale_m_step_4_ArhoMax), + "delay_em_iter_4_ArhoMax": int(args.delay_em_iter_4_ArhoMax), + "target_em_iter_4_ArhoMax": int(args.target_em_iter_4_ArhoMax), + "rho_sync_mode": "broadcast", + "delay_em_iter_4_lrDecay": int(args.delay_em_iter_4_lrDecay), + "target_em_iter_4_lrDecay": int(args.target_em_iter_4_lrDecay), + "delay_em_iter_4_Aprune": int(args.delay_em_iter_4_Aprune), + "mstep_state_mode": "viterbi_onehot_prevbin", + "batch_size": int(args.batch_size), + "num_states": int(M), + "num_neurons": int(N), + "num_time_bins": int(T_full), + "time_step_sec": float(dt), + "eta_clip": float(eta_clip), + "time_range_sec": [float(time_range_sec[0]), float(time_range_sec[1])], + "time_range_bins": [int(time_range_bins[0]), int(time_range_bins[1])], + "seed": int(args.seed), + "init_A_Tmax": int(getattr(args, "init_A_Tmax", 50000)), + } + outMD["states_recovery_eval"] = { + "decode": "viterbi", + "decode_dwell_sec": float(args.decode_dwell_sec), + } + outMD["init_A"] = init_A_md + outMD["init_state"] = init_state_md + outMD["init_B"] = init_B_md + outMD["provenance"] = prov + return outD, outMD + + +def run_locked_mstep(Y_prev, Y_curr, S_lock, A_init, B_init, spikeMD, args, ctx): + """Run locked M-step with fixed destination-bin state labels, no E-step/Viterbi. + + S_lock may contain one state label per pair, or one full timeline whose + destination-bin labels S[1:] match the supplied lag pairs. + """ + Y_prev = np.asarray(Y_prev, dtype=np.float32) + Y_curr = np.asarray(Y_curr, dtype=np.float32) + S_lock = np.asarray(S_lock, dtype=np.int64) + if Y_prev.shape != Y_curr.shape: + raise ValueError("Y_prev and Y_curr must have matching shape") + T_pairs, N = Y_prev.shape + if S_lock.shape[0] not in (T_pairs, T_pairs + 1): + raise ValueError("S_lock must have length T_pairs or T_pairs + 1") + + M = int(args.num_states) + dt = float(spikeMD["time_step_sec"]) + eta_clip = float(spikeMD["poisson_eta_clip"]) + if S_lock.shape[0] == T_pairs: + c_pairs_np = onehot_from_pair_states(S_lock, M) + else: + c_pairs_np = onehot_from_curr_states(S_lock, M) + loader, sampler = make_pair_loader(Y_prev, Y_curr, c_pairs_np, args, ctx, shuffle=True) + model = init_model(N, M, eta_clip, A_init, B_init, ctx) + mdl = model.module if hasattr(model, "module") else model + + total_epochs = int(getattr(args, "epochs", int(args.num_em_iters) * int(args.m_epochs))) + optimizer, scheduler = make_optimizer_and_scheduler( + model, args, total_epochs, decay_from_epoch=0 + ) + off_mask = ~torch.eye(N, dtype=torch.bool, device=ctx.device) + l1_wt = torch.ones(N, N, device=ctx.device) + l1_wt[torch.eye(N, dtype=torch.bool, device=ctx.device)] = 0.0 + + h = _history_dict() + for epoch in range(1, total_epochs + 1): + if sampler is not None: + sampler.set_epoch(epoch) + met = run_mstep_epoch( + model, loader, optimizer, ctx.device, dt, + l1_wt, args.lambda3, args.rho_max, + args.prescale_m_step_4_ArhoMax, + apply_prune=True, apply_rho=True, rho_correction_strength=1.0, + off_mask=off_mask + ) + scheduler.step() + h["m_loss_epoch"].append(met["loss"]) + h["m_nll_epoch"].append(met["nll"]) + h["m_l1_epoch"].append(met["l1"]) + h["rho_epoch"].append(met["rho"]) + h["nz_edges_epoch"].append(met["nz"]) + h["learning_rates"].append(float(optimizer.param_groups[0]["lr"])) + h["rho_correction_strength_epoch"].append(1.0) + + A_hat = mdl.A.detach().cpu().numpy().astype(np.float32) + B_hat = mdl.B.detach().cpu().numpy().astype(np.float32) + return A_hat, B_hat, _history_arrays(h) + + +def copy_args_with_init(args, init_A=None, init_B=None): + d = vars(args).copy() + if init_A is not None: + d["init_A"] = init_A + if init_B is not None: + d["init_B"] = init_B + return SimpleNamespace(**d) + + +def init_null_params(scrambled_spikes, spikeMD, args, ctx, null_A_init, null_B_init): + """Fresh null A/B initialization on rank 0, broadcast to all ranks.""" + if is_rank0(ctx): + if null_A_init == "scrambled_data": + a_args = copy_args_with_init(args, init_A="data") + A_init, A_md = init_A_from_spikes(scrambled_spikes, a_args) + elif null_A_init == "rand": + A_init, A_md = None, {"method": "rand"} + else: + raise ValueError(f"Unsupported null_A_init={null_A_init}") + + if null_B_init == "scrambled_data": + b_args = copy_args_with_init(args, init_B="data") + B_init, B_md = init_B_from_spikes(scrambled_spikes, float(spikeMD["time_step_sec"]), b_args) + elif null_B_init == "rand": + B_init, B_md = None, {"method": "rand"} + else: + raise ValueError(f"Unsupported null_B_init={null_B_init}") + else: + A_init = B_init = None + A_md = B_md = None + + A_init = broadcast_optional_array(A_init, ctx) + B_init = broadcast_optional_array(B_init, ctx) + A_md = broadcast_object(A_md, ctx) + B_md = broadcast_object(B_md, ctx) + return A_init, B_init, A_md, B_md diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/README.md b/causal_net/nonStation_ver3c_states_EM_FDR/README.md new file mode 100644 index 00000000..f57a9af7 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/README.md @@ -0,0 +1,314 @@ +# PRISM-EM 3c Training + +This directory contains the lag-1 PRISM-EM trainer and the FDR bagging +fitters. + +The two entry points share the same core implementation in +`PrismEM_Workhorse3c.py`: + +- `prism_EM_train3c.py`: ordinary time-ordered EM fit used for state + discovery. +- `prism_FDR_Bags_train3c.py`: one FDR-bagging Stage (a) job, i.e. one + reference-locked A/B refit plus its per-neuron time-shuffle null refits. +- `prism_EM_FDR_Bags_aggregate3c.py`: EM-FDR-bagging Stage (b), i.e. + aggregate all bag files into one eval-compatible fit. + +The code expects input spike files under: + +```bash +$basePath/spikesData/.spikes.npz +``` + +On Perlmutter, load PyTorch before running: + +```bash +module load pytorch +``` + +For reproducible multi-GPU runs, it is also useful to cap CPU thread pools: + +```bash +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +``` + +## Shell Macros + +The `.sh` files in this directory are convenience launch macros. They are +intended to be edited near the top before use, not treated as stable command +line interfaces. + +- `fitPrismEM.sh`: thin wrapper around `prism_EM_train3c.py` for a single + distributed EM fit. It checks that `--basePath`, `--dataName`, and + `--num_states` were supplied, fixes a short default `--time_range_sec 0 80`, + checks that four GPUs are visible, and launches `torchrun`. This is mostly + useful for quick trainer smoke tests. + +- `big_fit_bags.sh`: end-to-end example for one complete FDR-bagging run on + a chosen dataset and time range. It runs one reference EM fit, loops over + `numBags` calls to `prism_FDR_Bags_train3c.py`, and optionally runs the + Stage (b) aggregator. Edit `basePath`, `shortN`, `numStates`, `timeRange`, + and the training/FDR hyperparameters at the top. + +- `big_scanTime_FDR.sh`: end-to-end run for a single data-duration point in + a scan. It takes three positional arguments: + `./big_scanTime_FDR.sh `. The macro sets + `timeRange=(start_min*60 stop_min*60)` seconds, runs reference EM, all FDR + bags, and aggregation. The final aggregate name is + `__tomin`, which is convenient for + later metric scans over duration or time window. + +- `scan_FDR_BAG_hpar.sh`: cheap Stage (b)-only hyperparameter scan for an + existing set of bag files. It does not rerun EM or bag fitting. It reuses + `$basePath/prismFDR/${fdrBagsName}.bagNNN.prismFDRbag.npz` and sweeps + `stabSelScan` at fixed `base_per_bag_quantile`, then sweeps + `perBagQuantileScan` at fixed `base_stab_sel_thresh`. It writes aggregate + files named from `outAgrName` and prints ready-to-run `edgeMaterAbs3c.py` + metric commands for the two scans. + +- `docs/build3gen.sh`, `docs/build3EM.sh`, `docs/build3FDR.sh`, + `docs/build3exper.sh`, and `docs/build3edgeMeter.sh`: LaTeX build helpers + for the matching documentation files. They load `texlive` if needed, run + `pdflatex` twice, and try to open the produced PDF on supported systems. + +## Plain EM Fit + +Run the ordinary time-ordered EM fit with `torchrun`. This writes one fit file +to: + +```bash +$basePath/prismFit/.prismEM.npz +``` + +Example using 4 GPUs: + +```bash +basePath=/path/to/run +shortN=myDataset + +time torchrun --standalone --nnodes=1 --nproc_per_node=4 \ + ./prism_EM_train3c.py \ + --basePath $basePath \ + --dataName $shortN \ + --num_states 2 \ + --num_em_iters 50 \ + --batch_size 4096 \ + --delay_em_iter_4_lrDecay 5 \ + --delay_em_iter_4_ArhoMax 20 \ + --delay_em_iter_4_Aprune 30 \ + --time_range_sec 300 900 +``` + +Useful optional controls: + +- `--fitName`: output stem. If omitted, the trainer writes + `_` with a random four-character hex suffix. +- `--init_states {data,rand}`: initialize state probabilities from the data + or randomly. +- `--init_A {data,rand}`: initialize lag-1 connectivity from spike data or + randomly. +- `--init_B {data,rand}`: initialize state biases from spike rates or + randomly. +- `--init_A_Tmax`: maximum number of time bins used for data-driven `A` + initialization. + +The reference EM fit should be long enough to assign reliable states. Final +conditional `A,B` precision is controlled by the locked-M-step FDR bag jobs. + +## FDR-Bags Stage (a) + +First create the reference fit with `prism_EM_train3c.py`. In production runs +use an explicit tag, for example `--fitName _jXXXX`, so this file +exists: + +```bash +$basePath/prismFit/_jXXXX.prismEM.npz +``` + +Then run `prism_FDR_Bags_train3c.py` once per bag. Each invocation: + +1. Loads the EM-train output selected by `--emFitName`. +2. Recovers the original spike-file stem from that EM fit's provenance + metadata. +3. Forms lag pairs `(S_hat[t], spikes[t-1], spikes[t])`. +4. Samples `--bag_frac` of those pairs without replacement. +5. Runs a locked M-step to refit only `A,B`. +6. Runs `--num_scrambles` per-neuron time-shuffle null refits using the same + selected pair indices and locked labels. +7. Writes one self-contained bag file. + +The output file is: + +```bash +$basePath/prismFDR/.bag.prismFDRbag.npz +``` + +The bag trainer does not accept `--dataName`. This is intentional: the source +spike file is derived from the reference EM metadata, which prevents mixing +states from one dataset with spikes from another. + +All bag indices, including `bag000`, are equivalent random pair subsets. + +Example using 4 GPUs for bag 3: + +```bash +basePath=/path/to/run +shortN=myDataset + +time torchrun --standalone --nnodes=1 --nproc_per_node=4 \ + ./prism_FDR_Bags_train3c.py \ + --basePath $basePath \ + --emFitName ${shortN}_emj1234 \ + --outFitName ${shortN}_emj1234_rmfA \ + --bag_idx 3 \ + --bag_frac 0.8 \ + --epochs 240 \ + --num_scrambles 4 \ + --batch_size 4096 +``` + +The bag driver loads locked-M-step defaults from the reference EM metadata, +but these can be overridden for FDR fitting: + +- `--bag_idx`: integer bag index used in the output filename and RNG stream. +- `--emFitName`: EM-train input fit stem in `prismFit/`. +- `--outFitName`: Stage (a) output bag stem in `prismFDR/`. + If omitted, it defaults to `_` with a random + four-character hex suffix. +- `--fdr_out_dir`: optional output directory; default is `$basePath/prismFDR`. +- `--bag_frac`: fraction of reference lag pairs selected without replacement. +- `--epochs`: locked M-step epochs for real and null refits. +- `--lr_mstep`, `--lr_end_factor`, `--lambda3`, `--rho_max`, + `--prescale_m_step_4_ArhoMax`, `--batch_size`: locked-M-step overrides. +- `--init_A {data,ref,rand}` and `--init_B {data,ref,rand}`: real locked-fit + initialization; default `data` to let each bag initialize from its sampled + lag pairs. +- `--num_scrambles`: number of per-neuron time-shuffle null refits. +- `--min_roll_shift_sec`: exclusion zone around zero shift for each neuron; + default `2`. +- `--null_A_init {scrambled_data,rand}`: fresh `A` initialization mode for + each null refit; default `scrambled_data`. +- `--null_B_init {scrambled_data,rand}`: fresh `B` initialization mode for + each null refit; default `scrambled_data`. + +The bag trainer does not accept EM state-discovery controls such as +`--time_range_sec`, `--num_states`, `--num_em_iters`, `--pgd_iter`, +`--lr_estep`, `--lambda2`, `--init_states`, or `--decode_dwell_sec`. + +## Running Multiple Bags + +Stage (a) bags are independent jobs. For a small manual launch, vary +`--bag_idx`: + +```bash +for bag in 0 1 2 3 4; do + torchrun --standalone --nnodes=1 --nproc_per_node=4 \ + ./prism_FDR_Bags_train3c.py \ + --basePath $basePath \ + --emFitName ${shortN}_emj1234 \ + --outFitName ${shortN}_emj1234_rmfA \ + --bag_idx $bag \ + --bag_frac 0.8 \ + --epochs 240 \ + --num_scrambles 4 +done +``` + +For production, submit one SLURM job per `bag_idx`, each using one node and +4 GPUs. The bag trainer itself does not perform final pooled FDR selection. + +## EM-FDR-Bags Stage (b) + +After Stage (a) finishes for `bag000` through `bag`, aggregate +the bag files with: + +```bash +basePath=/path/to/run +shortN=myDataset +fdrBagsName=${shortN}_emj1234_rmfA +agrFitName=${fdrBagsName}_agrA +num_bags=5 + +./prism_EM_FDR_Bags_aggregate3c.py \ + --basePath $basePath \ + --dataName $fdrBagsName \ + --outAgrName $agrFitName \ + --num_bags $num_bags \ + --per_bag_quantile 0.99 \ + --stab_sel_thresh 0.7 +``` + +The aggregator reads: + +```bash +$basePath/prismFDR/.bag000.prismFDRbag.npz +$basePath/prismFDR/.bag001.prismFDRbag.npz +... +``` + +and writes: + +```bash +$basePath/prismFit/.prismEM.npz +``` + +If `--outAgrName` is omitted, the output stem defaults to `_` +with a random four-character hex suffix. + +For example, with `fdrBagsName=daleN100_2ba29b_c47b43_emj1234_rmfA` and +`outAgrName=${fdrBagsName}_agrA`, the output stem is: + +```bash +daleN100_2ba29b_c47b43_emj1234_rmfA_agrA +``` + +The output keeps the normal PRISM-EM fields expected by `prism_EM_eval3c.py`. +The connectivity fields are aggregated Stage (b) results, while the time +series fields (`S_hat`, `c_hat`, and `S_hat_CL`) come from the reference EM +fit used to lock the bags. + +You can then run: + +```bash + ./prism_EM_eval3c.py \ + --basePath $basePath \ + --dataName $agrFitName \ + -p a e f g +``` + +To inspect one Stage (a) locked bag fit directly, pass the bag stem as +`--dataName`; names containing `bagNNN` are loaded from `prismFDR`: + +```bash +./prism_EM_eval3c.py \ + --basePath $basePath \ + --dataName ${shortN}.bag000 \ + -p a e f g +``` + +Stage (b) also saves FDR diagnostics such as `selection_frequency`, +`selected_mask`, `A_mean_selected`, `A_sd_selected`, `A_mean_all`, `A_sd_all`, +`src_null_mean`, `src_null_sd`, `src_null_tau_mean`, `z_null`, and a compact +selected-edge table (`edge_i`, `edge_j`, `edge_sel_freq`, `edge_A_mean`, +`edge_A_sd_boot`, `edge_z_null`, `edge_src_null_mean`, `edge_src_null_sd`, +`edge_src_tau_mean`). + +## Saved Bag Contents + +Each `.prismFDRbag.npz` contains the real fit fields in the same style as +`prism_EM_train3c.py`, plus FDR-bagging fields including: + +- `A_null`: null connectivity stack with shape `(num_scrambles, N, N)`. +- `B_null`: null bias stack with shape `(num_scrambles, num_states, N)`. +- `roll_shifts_bin`: realized circular shifts per scramble and neuron. +- `bag_pair_indices`: selected lag-pair row indices in the reference window. +- `bag_pair_dest_bins`: selected destination-bin indices in original time. +- `bag_pair_state`: locked state label for each selected pair. +- `null_*_epoch`: locked M-step histories for the null refits. + +Metadata from the bagging workflow are grouped by stage. `bagsFDR_stageA` +contains pair-sampling settings, locked-M-step hyperparameters, null-refit +settings, per-scramble null initialization diagnostics, and reference-fit +provenance. The aggregate file also adds `bagsFDR_stageB` for the cross-bag +selection and output settings. diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/UtilBioExp.py b/causal_net/nonStation_ver3c_states_EM_FDR/UtilBioExp.py new file mode 100644 index 00000000..50019e00 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/UtilBioExp.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +""" +Utility functions for biological experiment data processing. + +This module provides specialized functions for analyzing biological neural data, +particularly focused on identifying and processing clusters of neural activity. +Main functionality includes: +- Cluster detection in time series data using connectivity analysis +- Threshold-based cluster validation and masking +- Signal processing utilities for experimental neural recordings + +Used primarily in conjunction with experimental data preprocessing and +analysis pipelines for biological neural network studies. +""" + +import numpy as np + + +def detect_spike_bursts(spikeD, spikeMD, time_rebin2, burst_freq_thres): + """Rebin spikes in time; return rate panels for bioExp freq-vs-time plots.""" + rateThr2 = float(burst_freq_thres) + spikeYield = spikeD["spikes"] + time_step = float(spikeMD["time_step_sec"]) + tReb2 = int(time_rebin2) + assert tReb2 < 101 + + ntime, nchan = spikeYield.shape + if ntime % tReb2 != 0: + ntime_c = ntime - (ntime % tReb2) + spikeYield = spikeYield[:ntime_c] + + spikeYieldR = np.sum(spikeYield.reshape(-1, tReb2, nchan), axis=1) + time_step2 = time_step * tReb2 + + rate2D = spikeYieldR / time_step2 + mask2D = rate2D > rateThr2 + highChan = np.sum(mask2D, axis=1) + + rebD = { + "time_step2": time_step2, + "rate_thres2": rateThr2, + "rate2D": rate2D, + "mask2D": mask2D, + "highChanCnt": highChan, + "rate1D": np.sum(rate2D, axis=1), + } + ntime = spikeYieldR.shape[0] + rebD["timeV"] = np.linspace(0, (ntime - 1) * time_step2, ntime) + + print("nchan=%d thr=%.1f Hz" % (nchan, rateThr2)) + return rebD + + +def clip_rebD_time(rebD, time_range_lr): + """Clip rebD arrays to a time window in seconds.""" + time_step = rebD["time_step2"] + tL, tR = [float(x) for x in time_range_lr] + itL, itR = (np.asarray(time_range_lr, dtype=np.float64) / time_step).astype(int) + itR = min(itR, rebD["rate2D"].shape[0]) + if itR <= itL: + raise ValueError("time_range_lr leaves no rebinned bins") + + rate2D = rebD["rate2D"][itL:itR] + timeV = rebD["timeV"][itL:itR] + rate1D = rebD["rate1D"][itL:itR] + _, nchan = rate2D.shape + Tbin = float(timeV[1] - timeV[0]) + return { + "tL": tL, "tR": tR, "itL": itL, "itR": itR, + "time_step": time_step, "Tbin": Tbin, "nchan": nchan, + "rate2D": rate2D, "timeV": timeV, "rate1D": rate1D, + } + + +if __name__=="__main__": + print("Utility module for biological experiment spike-rate rebinning.") diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/UtilDalePoisson.py b/causal_net/nonStation_ver3c_states_EM_FDR/UtilDalePoisson.py new file mode 100644 index 00000000..de4ab1a0 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/UtilDalePoisson.py @@ -0,0 +1,240 @@ +""" +Utility functions for Dale Poisson simulation data processing + +""" + +import numpy as np +import os +import time +from pprint import pprint + +def compute_consecutive_coincidence_rate(Y,time_evol): + """ + Compute the frequency of coincidences for 2 consecutive time bins for 2 different channels. + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + + Returns: + float: conc_rate_all - the overall coincidence rate + """ + num_steps, num_neurons = Y.shape + + # We need at least 2 time steps for consecutive bins + if num_steps < 2: + return 0.0 + + # Get consecutive time slices using vectorized operations + Y_t = Y[:-1, :] # Y[0:T-1, :] - current time bins + Y_t_plus_1 = Y[1:, :] # Y[1:T, :] - next time bins + + # Create boolean masks for non-zero values + mask_t = (Y_t != 0) # Shape: (T-1, N) + mask_t_plus_1 = (Y_t_plus_1 != 0) # Shape: (T-1, N) + + # Use broadcasting to compute all channel pairs at once + # mask_t[:, :, None] has shape (T-1, N, 1) + # mask_t_plus_1[:, None, :] has shape (T-1, 1, N) + # Broadcasting gives shape (T-1, N, N) for all pairs + coincidences = mask_t[:, :, None] & mask_t_plus_1[:, None, :] + + # Remove diagonal (same channel pairs) using boolean indexing + diagonal_mask = np.eye(num_neurons, dtype=bool) + coincidences[:, diagonal_mask] = False + + # Count total coincidences across all time steps + coincidence_count = np.sum(coincidences) + + conc_rate = coincidence_count / time_evol/num_neurons + return conc_rate # Hz, per neuron + +def estimate_rates(Y, dt, num_excite, max_samples, varTwindow=5, mxNn=5, verb=1, spect_radius=None): + """ + Evaluates spike statistics and estimates firing rates and Fano factors. + + Computes statistics over non-overlapping time windows of length varTwindow: + - Fano Factor: Var[spike count] / Mean[spike count] per neuron (dimensionless) + - Rate variance: Var[rate] per neuron (Hz^2) + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + dt (float): Time bin size in seconds. + num_excite (int): Number of excitatory neurons. + max_samples (int): The maximum number of time samples to use for calculation. + varTwindow (float, optional): Time window length in seconds for variance computation. Defaults to 5. + mxNn (int, optional): Max number of neurons to show in detailed stats. Defaults to 5. + + Returns: + tuple: A tuple containing (stats_dict, rates_dict, neur_freq_index): + - stats_dict (dict): Contains population-level spike statistics. + - rates_dict (dict): Contains per-neuron firing rates, Fano factors, and rate variance. + - neur_freq_index (np.ndarray): Neuron indices sorted by firing rate (low to high). + """ + # 1. Clip data to max_samples_for_rates + if verb > 0: + print("\n=== Estimating Rates & Stats ===") + if Y.shape[0] > max_samples: + if verb > 0: + print("Using %d samples (out of %d) for rate computation" % (max_samples, Y.shape[0])) + Y = Y[:max_samples] + + # Part 1: Compute all basic statistics once + num_steps_sim, Nn_sim = Y.shape + num_inhib = Nn_sim - num_excite + time_evol = num_steps_sim * dt + if verb > 0: + print('steps num_steps=%d, time_evol=%.1f sec, Nn=%d (%d Excit, %d Inhib)' % (num_steps_sim, time_evol, Nn_sim, num_excite, num_inhib)) + + # Compute windowed spike counts and statistics over non-overlapping time windows of length varTwindow (sec), per neuron + window_size = max(1, int(round(varTwindow / dt))) + num_windows = num_steps_sim // window_size + if verb > 0: + print('variance computation: num_windows=%d, window_size=%d' % (num_windows, window_size)) + if num_windows <= 1: + raise ValueError("varTwindow=%s is too large for data (num_windows=%d). Need at least 2 windows for variance." % (varTwindow, num_windows)) + Y_trim = Y[:num_windows * window_size] + Y_win = Y_trim.reshape(num_windows, window_size, Nn_sim) + + # Sum spike counts over each window: shape (num_windows, n_neurons) + window_spike_counts = np.sum(Y_win, axis=1) + + # Fano Factor: Var[spike count] / Mean[spike count] over windows + mean_window_spike_counts = np.mean(window_spike_counts, axis=0) + var_window_spike_counts = np.var(window_spike_counts, axis=0) + fano_factor = np.divide(var_window_spike_counts, mean_window_spike_counts, + out=np.zeros_like(var_window_spike_counts), + where=mean_window_spike_counts != 0) + + # Convert window spike counts to rates (Hz) for rate variance + window_rates = window_spike_counts / (window_size * dt) # shape: (num_windows, n_neurons), Hz + single_rates_var = np.var(window_rates, axis=0) # variance of rate across windows, per neuron (Hz^2) + + # Compute raw arrays (full time series) + spike_counts = np.sum(Y, axis=0) + spike_rates = spike_counts / time_evol + mean_counts_per_bin = np.mean(Y, axis=0) + + # Compute consecutive coincidence rate + start_time = time.time() + conc_rate_per_neuron = compute_consecutive_coincidence_rate(Y,time_evol) + elapsed_time = time.time() - start_time + if verb > 0: + print('Coincidence rate %.2g Hz, Y.shape=%s elaT %.3f sec' % (conc_rate_per_neuron, Y.shape, elapsed_time)) + + # Compute all population statistics once + med_rate_all = np.median(spike_rates) + avg_rate_all = float(np.mean(spike_rates)) + std_rate_all = float(np.std(spike_rates)) + avg_fano_all = float(np.mean(fano_factor)) + std_fano_all = float(np.std(fano_factor)) + avg_rate_excit = float(np.mean(spike_rates[:num_excite])) + std_rate_excit = float(np.std(spike_rates[:num_excite])) + avg_fano_excit = float(np.mean(fano_factor[:num_excite])) + std_fano_excit = float(np.std(fano_factor[:num_excite])) + avg_rate_inhib = float(np.mean(spike_rates[num_excite:])) + std_rate_inhib = float(np.std(spike_rates[num_excite:])) + avg_fano_inhib = float(np.mean(fano_factor[num_excite:])) + std_fano_inhib = float(np.std(fano_factor[num_excite:])) + + # Build stats dictionary with computed values + stats_dict = { + 'num_steps': num_steps_sim, + 'time_evol_sec': time_evol, + 'time_step_sec': dt, + 'num_neurons': Nn_sim, + 'num_excitatory': num_excite, + 'num_inhibitory': num_inhib, + 'avg_spike_rate_all': avg_rate_all, + 'std_spike_rate_all': std_rate_all, + 'avg_fano_factor_all': avg_fano_all, + 'std_fano_factor_all': std_fano_all, + 'avg_spike_rate_excit': avg_rate_excit, + 'std_spike_rate_excit': std_rate_excit, + 'avg_fano_factor_excit': avg_fano_excit, + 'std_fano_factor_excit': std_fano_excit, + 'avg_spike_rate_inhib': avg_rate_inhib, + 'std_spike_rate_inhib': std_rate_inhib, + 'avg_fano_factor_inhib': avg_fano_inhib, + 'std_fano_factor_inhib': std_fano_inhib, + 'conc_rate_per_neuron': float(conc_rate_per_neuron), + 'median_spike_rate_all': med_rate_all + } + + # Printing detailed stats for individual neurons + mxE = min(mxNn, num_excite) + mxI = min(mxNn, num_inhib) + + if verb > 0: + print('\n--- Stats for first %d Excitatory Neurons ---' % mxE) + np.set_printoptions(precision=2) + print('Total Spike Counts: %s' % spike_counts[:mxE]) + print('Mean Firing Rate (Hz): %s' % spike_rates[:mxE]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[:mxE])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[:mxE])) + print('Fano Factor (Var/Mean): %s' % fano_factor[:mxE]) + + print('\n--- Stats for first %d Inhibitory Neurons ---' % mxI) + np.set_printoptions(precision=2) + inhib_slice = slice(num_excite, num_excite + mxI) + print('Total Spike Counts: %s' % spike_counts[inhib_slice]) + print('Mean Firing Rate (Hz): %s' % spike_rates[inhib_slice]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[inhib_slice])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[inhib_slice])) + print('Fano Factor (Var/Mean): %s' % fano_factor[inhib_slice]) + + # Print population summary using dictionary values (always printed) + R_tag = ', R=%.3f' % spect_radius if spect_radius is not None else '' + print('\n--- Population Summary Statistics%s ---' % R_tag) + print('All (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (Nn_sim, stats_dict['avg_spike_rate_all'], stats_dict['std_spike_rate_all'], stats_dict['avg_fano_factor_all'], stats_dict['std_fano_factor_all'])) + print('Excit (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_excite, stats_dict['avg_spike_rate_excit'], stats_dict['std_spike_rate_excit'], stats_dict['avg_fano_factor_excit'], stats_dict['std_fano_factor_excit'])) + if num_inhib > 0: + print('Inhib (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_inhib, stats_dict['avg_spike_rate_inhib'], stats_dict['std_spike_rate_inhib'], stats_dict['avg_fano_factor_inhib'], stats_dict['std_fano_factor_inhib'])) + + # Print key summary values using dictionary + summary_keys = ['conc_rate_per_neuron', 'median_spike_rate_all'] + for key in summary_keys: + if key == 'conc_rate_per_neuron': + print('Coincidence rate per neuron %.2g Hz' % stats_dict[key]) + elif key == 'median_spike_rate_all': + print('Median rate %.2f Hz\n' % stats_dict[key]) + + # Part 2: Compute frequency sorting + neur_freq_index = np.argsort(spike_rates) + + rates_dict = { + 'single_rates': spike_rates, + 'neur_freqIdx':neur_freq_index, + 'sigle_rates_var': single_rates_var, + 'single_fano_fact': fano_factor + } + + return stats_dict, rates_dict, neur_freq_index + +def get_offdiag_triplets(A, isPos=True): + """ + Returns positive/negative, off-diagonal elements of 2D array A as array of [i, j, value] triplets. + + Parameters: + A : 2D numpy array + isPos : bool, if True only return positive values (>0), if False return all negative values + + Returns: + Array of shape (n, 3) where n is the number of valid off-diagonal elements + Each row is [row_index, col_index, value] + """ + # Create mask for off-diagonal elements + offdiag_mask = ~np.eye(A.shape[0], A.shape[1], dtype=bool) + + # Add value condition based on isPos + if isPos: + value_mask = A > 0 + else: + value_mask = A < 0 + + # Combine masks + mask = value_mask & offdiag_mask + + i_indices, j_indices = np.where(mask) + values = A[i_indices, j_indices] + + return np.column_stack([i_indices, j_indices, values]) diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/Util_PrismEM.py b/causal_net/nonStation_ver3c_states_EM_FDR/Util_PrismEM.py new file mode 100644 index 00000000..47baaee5 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/Util_PrismEM.py @@ -0,0 +1,300 @@ +#!/usr/bin/env python3 +""" +Utilities for prism EM initialization. +""" + +import numpy as np +import time + + +def _normalize_init_opt(opt): + s = str(opt).strip().lower() + if s not in ("data", "rand"): + raise ValueError(f"Unsupported init option: {opt}") + return s + + +def _rate_thresholds_min_var(rate_t, n_state): + """Find up to (n_state-1) thresholds minimizing within-group variance.""" + x = np.asarray(rate_t, dtype=np.float64).ravel() + if x.size == 0 or n_state <= 1: + return [] + + uniq, counts = np.unique(x, return_counts=True) + u = uniq.size + if u <= 1: + return [] + + k = min(int(n_state), u) + w = counts.astype(np.float64) + wx = w * uniq + wx2 = wx * uniq + + c_w = np.concatenate(([0.0], np.cumsum(w))) + c_s = np.concatenate(([0.0], np.cumsum(wx))) + c_q = np.concatenate(([0.0], np.cumsum(wx2))) + + # dp[c, j]: minimal weighted SSE for first (j+1) unique values into c groups. + dp = np.full((k + 1, u), np.inf, dtype=np.float64) + ptr = np.full((k + 1, u), -1, dtype=np.int32) + + j_all = np.arange(u, dtype=np.int32) + 1 + w0 = c_w[j_all] - c_w[0] + s0 = c_s[j_all] - c_s[0] + q0 = c_q[j_all] - c_q[0] + dp[1, :] = q0 - (s0 * s0) / np.maximum(w0, 1e-12) + ptr[1, :] = 0 + + for c in range(2, k + 1): + for j in range(c - 1, u): + i = np.arange(c - 1, j + 1, dtype=np.int32) + jj = j + 1 + + ww = c_w[jj] - c_w[i] + ss = c_s[jj] - c_s[i] + qq = c_q[jj] - c_q[i] + seg = qq - (ss * ss) / np.maximum(ww, 1e-12) + cand = dp[c - 1, i - 1] + seg + + ib = int(np.argmin(cand)) + dp[c, j] = float(cand[ib]) + ptr[c, j] = int(i[ib]) + + starts = [] + j = u - 1 + for c in range(k, 1, -1): + i = int(ptr[c, j]) + starts.append(i) + j = i - 1 + starts.reverse() + + thres = [] + for i in starts: + l = uniq[i - 1] + r = uniq[i] + thres.append(float(0.5 * (l + r))) + return thres + + +def init_states_vs_time(spikes, dt, args): + """Prepare initial c_hat probabilities and state labels from spikes. + + Returns: + c_init: float32 array (T, M) + S_init: int64 array (T,) + init_meta: None or {"rate_thres": [..], "rate_bin_sec": float, "rate_bin_bins": int} + freq_h1d: float32 array used for state classification + """ + spikes = np.asarray(spikes) + if spikes.ndim != 2: + raise ValueError("spikes must be 2D array (T, N)") + + T_full = spikes.shape[0] + n_state = int(args.num_states) + opt = _normalize_init_opt(args.init_states) + + if opt == "rand": + c_init = np.full((T_full, n_state), 1.0 / n_state, dtype=np.float32) + s_init = np.full((T_full,), -1, dtype=np.int64) + freq_h1d = np.zeros((0,), dtype=np.float32) + return c_init, s_init, None, freq_h1d + if opt != "data": + raise ValueError(f"Unsupported --init_states option: {opt}") + + dwell_sec = float(args.decode_dwell_sec) + dwell_bins = max(1, int(np.ceil(dwell_sec / float(dt)))) + + # Aggregate spikes over coarse dwell bins, then compute mean firing rate. + T = T_full + n_blk = (T + dwell_bins - 1) // dwell_bins + rate_blk = np.zeros((n_blk,), dtype=np.float64) + blk_sizes = np.zeros((n_blk,), dtype=np.int64) + for b in range(n_blk): + i0 = b * dwell_bins + i1 = min(T, i0 + dwell_bins) + blk = spikes[i0:i1] + nb = i1 - i0 + blk_sizes[b] = nb + rate_per_neuron = blk.sum(axis=0).astype(np.float64) / (float(nb) * float(dt)) + rate_blk[b] = rate_per_neuron.mean() + + thres = _rate_thresholds_min_var(rate_blk, n_state) + th = np.asarray(thres, dtype=np.float64) + s_blk = np.searchsorted(th, rate_blk, side="right").astype(np.int64) + s_blk = np.clip(s_blk, 0, max(0, n_state - 1)) + + c_blk, s_blk_out = init_probs_from_states(s_blk, n_state) + + # Map coarse-bin initialization back to original binning. + c_init = np.zeros((T_full, n_state), dtype=np.float32) + s_init = np.zeros((T_full,), dtype=np.int64) + for b in range(n_blk): + i0 = b * dwell_bins + i1 = i0 + int(blk_sizes[b]) + c_init[i0:i1] = c_blk[b] + s_init[i0:i1] = s_blk_out[b] + + freq_h1d = rate_blk.astype(np.float32, copy=False) + init_meta = { + "rate_thres": [float(x) for x in thres], + "rate_bin_sec": float(dwell_bins * float(dt)), + "rate_bin_bins": int(dwell_bins), + } + return c_init, s_init, init_meta, freq_h1d + + +def init_selfspikingB(spikes, dt, args): + """Initialize common per-state B from observed mean firing rates. + + For each neuron n: + rate_hz[n] = mean_t spikes[t, n] / dt + B[:, n] = log(max(rate_hz[n], eps)) + """ + spikes = np.asarray(spikes) + if spikes.ndim != 2: + raise ValueError("spikes must be 2D array (T, N)") + + n_state = int(args.num_states) + eps_rate_hz = 1e-6 + + rate_hz = spikes.mean(axis=0).astype(np.float64) / float(dt) + rate_hz_clip = np.maximum(rate_hz, eps_rate_hz) + b_vec = np.log(rate_hz_clip).astype(np.float32) + b_init = np.tile(b_vec[None, :], (n_state, 1)) + + init_meta = { + "method": "rate", + "common_across_states": True, + "rate_floor_hz": float(eps_rate_hz), + "formula": "B=log(max(rate_hz,eps))", + } + return b_init, init_meta + + +def init_B_from_spikes(spikes, dt, args): + """Select B initialization mode based on args.init_B.""" + opt = _normalize_init_opt(args.init_B) + if opt == "rand": + return None, None + if opt != "data": + raise ValueError(f"Unsupported --init_B option: {opt}") + return init_selfspikingB(spikes, dt, args) + + +def init_edgesA(spikes, Tmax=50000, verbose=True): + """OLS/covariance initialization of A from spike data Y (T x N).""" + t0 = time.perf_counter() + Y = np.asarray(spikes, dtype=np.float64) + if Y.ndim != 2: + raise ValueError("spikes must be 2D array (T, N)") + T_full, _ = Y.shape + T_use = min(int(Tmax), int(T_full)) + if T_use < 2: + raise ValueError("Need at least 2 time bins to initialize A") + Y = Y[:T_use] + + YpYp = Y[:-1].T @ Y[:-1] # Gram matrix + YYp = Y[1:].T @ Y[:-1] # one-step cross-correlation + + A_ols = YYp @ np.linalg.pinv(YpYp) + + kappa = float(np.linalg.cond(YpYp)) + rho = float(np.max(np.abs(np.linalg.eigvals(A_ols)))) + Y_pred = Y[:-1] @ A_ols.T + var_y = float(np.var(Y[1:])) + if var_y <= 0.0: + R2 = float("nan") + else: + R2 = float(1.0 - np.var(Y[1:] - Y_pred) / var_y) + frob = float(np.linalg.norm(A_ols, ord="fro")) + elapsed_sec = float(time.perf_counter() - t0) + + if verbose: + print("A-init OLS diagnostics:") + print(f" cond(YpYp) = {kappa:.2e}") + print(f" rho(A_ols) = {rho:.3f}") + print(f" R2 = {R2:.3f}") + print(f" ||A||_F = {frob:.3f}") + print(f" bins_used = {T_use}/{T_full}") + print(f" elapsed_s = {elapsed_sec:.3f}") + + meta = { + "method": "cov_ols", + "Tmax": int(Tmax), + "num_bins_used": int(T_use), + "num_bins_total": int(T_full), + "cond_YpYp": kappa, + "rho_A_init": rho, + "R2_1step": R2, + "fro_A_init": frob, + "elapsed_init_edgesA_sec": elapsed_sec, + } + return A_ols.astype(np.float32), meta + + +def init_A_from_spikes(spikes, args): + """Select A initialization mode based on args.init_A.""" + + opt = _normalize_init_opt(args.init_A) + if opt == "rand": + return None, None + if opt != "data": + raise ValueError(f"Unsupported --init_A option: {opt}") + init_A_Tmax = int(getattr(args, "init_A_Tmax", 50000)) + A_init, meta = init_edgesA(spikes, Tmax=init_A_Tmax, verbose=False) + + if 1: # rescale A + offDiagFact = 3 + meta["offDiag_A_init_fact"] = offDiagFact + A_init = A_init * offDiagFact + + rho_max = float(args.rho_max) + if 0: # global spectral radius rescaling + rho_before = float(np.max(np.abs(np.linalg.eigvals(A_init)))) + if rho_before > rho_max: + A_init = A_init * (rho_max / rho_before) + rho_after = float(np.max(np.abs(np.linalg.eigvals(A_init)))) + meta["rho_A_init_raw"] = rho_before + meta["rho_A_init"] = rho_after + meta["rho_max_target"] = rho_max + + if args.verb > 0: + print("A-init OLS diagnostics:") + print(f" cond(YpYp) = {meta['cond_YpYp']:.1f}") + print(f" rho(A_ols) = {meta['rho_A_init']:.3f}") + print(f" R2 = {meta['R2_1step']:.3f}") + print(f" ||A||_F = {meta['fro_A_init']:.3f}") + print(f" bins_used = {meta['num_bins_used']}/{meta['num_bins_total']}") + print(f" elapsed_s = {meta['elapsed_init_edgesA_sec']:.3f}") + + return A_init, meta + + +def init_probs_from_states(S_rate, M): + """Build c_init from S_rate with transition smoothing.""" + S_rate = np.asarray(S_rate, dtype=np.int64) + T = S_rate.size + if T == 0: + return np.zeros((0, M), dtype=np.float32), np.zeros((0,), dtype=np.int64) + if M <= 1: + return np.ones((T, 1), dtype=np.float32), S_rate.copy() + + # Base: dominant state at 0.8, remaining 0.2 spread over all others. + other = 0.2 / float(M - 1) + c_init = np.full((T, M), other, dtype=np.float32) + c_init[np.arange(T), S_rate] = 0.8 + S_out = S_rate.copy() + + # Transition smoothing: (t-1) gets 2/3 old + 1/3 new, t gets 1/3 old + 2/3 new. + tr_idx = np.where(S_rate[1:] != S_rate[:-1])[0] + 1 + for t in tr_idx: + a = int(S_rate[t - 1]) + b = int(S_rate[t]) + c_init[t - 1, :] = 0.0 + c_init[t, :] = 0.0 + c_init[t - 1, a] = 2.0 / 3.0 + c_init[t - 1, b] = 1.0 / 3.0 + c_init[t, a] = 1.0 / 3.0 + c_init[t, b] = 2.0 / 3.0 + + return c_init, S_out diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/Util_pseudospectra.py b/causal_net/nonStation_ver3c_states_EM_FDR/Util_pseudospectra.py new file mode 100755 index 00000000..23fc3e1f --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/Util_pseudospectra.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +import numpy as np +import time + +def create_example_matrices(size=100): + A_base = np.diag(np.linspace(-5, -0.5, size)) + np.diag(np.ones(size - 2) * 4, k=2) + sparsity = 0.15 + random_connections = np.random.randn(size, size) + mask = np.random.rand(size, size) > sparsity + random_connections[mask] = 0 + A_true = A_base + random_connections * 0.1 + + noise = np.random.randn(size, size) * 0.10 + noise_mask = np.random.rand(size, size) > sparsity + noise[noise_mask] = 0 + A_fit = A_true + noise + + return A_true, A_fit + +def compute_pseudospectrum(A, npts, minY): + print(f'computing pseudospectrum A:{A.shape} .... ') + eigs = np.linalg.eigvals(A) + + # Calculate bounds + real_min, real_max = np.real(eigs).min(), np.real(eigs).max() + imag_min, imag_max = np.imag(eigs).min(), np.imag(eigs).max() + + real_pad = (real_max - real_min) * 0.2 + imag_pad = (imag_max - imag_min) * 0.2 + + bbox = [ + real_min - real_pad, + min(real_max + real_pad, 1.0), + imag_min, + imag_max + imag_pad + ] + + x_coords = np.linspace(bbox[0], bbox[1], npts) + y_coords = np.linspace(minY, bbox[3], npts) + #y_coords = np.linspace(bbox[2], bbox[3], npts) + X, Y = np.meshgrid(x_coords, y_coords) + Z = X + 1j * Y + + sigma_grid = np.zeros_like(Z, dtype=float) + I = np.eye(A.shape[0]) + + for i in range(npts): + for j in range(npts): + z = Z[i, j] + s = np.linalg.svd(z * I - A, compute_uv=False) + sigma_grid[i, j] = s[-1] + + return X, Y, sigma_grid, eigs + +def plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin): + levels = np.logspace(-2.5, -0.5, 10) # Use 10 contour levels + + # Plot the normal contour lines + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + + # Create a mask for the area greater than epsMin + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) # Masking areas greater than epsMin + + # Fill the area below epsMin + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + # Scatter plot for eigenvalues + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + # Vertical line at Re=0 + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + # Set titles and labels + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + # Additional formatting + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(bottom=minY) # Clip at minY + + +def main(size=100, minY=-0.9, epsMin=0.032): + A_true, A_fit = create_example_matrices(size) + + fig, axes = plt.subplots(1, 2, figsize=(16, 7)) + + matrices_to_plot = [ + ('True Matrix ($A_{true}$)', A_true), + ('Distorted Matrix ($A_{fit}$) ', A_fit) + ] + + print("Starting pseudospectrum calculations...", minY) + start_time = time.time() + + for i, (title, A) in enumerate(matrices_to_plot): + ax = axes[i] + print("Processing '{}'...".format(title)) + + npts = 80 # Grid resolution + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin) # Pass epsMin + + total_time = time.time() - start_time + print(f"Calculation finished in {total_time:.2f} seconds.") + + fig.suptitle('Pseudospectrum Comparison (Clipped Imaginary Part)', fontsize=16) + plt.tight_layout(rect=[0, 0.03, 1, 0.95]) + + outFile = 'out/pseudospectra_comparison.png' + plt.savefig(outFile) + print(f"Plot saved as '{outFile}'") + plt.close(fig) + +# --- Main execution --- +if __name__ == '__main__': + + import matplotlib.pyplot as plt + + main(size=50, minY=-0.5, epsMin=0.024) diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/batchFitBags.slr b/causal_net/nonStation_ver3c_states_EM_FDR/batchFitBags.slr new file mode 100755 index 00000000..dc3b5f99 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/batchFitBags.slr @@ -0,0 +1,141 @@ +#!/bin/bash +# Submit as a SLURM array after the reference EM fit exists: +# sbatch --array=0-5 batchFitBags.slr + +#SBATCH -N 1 +#SBATCH -C gpu +#SBATCH -A m2043 +#SBATCH --time=30:00 +#SBATCH -q debug +# SBATCH --time=2:28:00 +# SBATCH -q regular +#SBATCH --array=0-1 +#SBATCH --output=outj/%A_%a.out +#SBATCH --licenses=scratch + +set -euo pipefail + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +mkdir -p outj +module load pytorch + +bag_idx=${SLURM_ARRAY_TASK_ID:-0} + +# Dataset; EM fit must already be: +# $basePath/prismFit/${emFitName}.prismEM.npz +basePath=/pscratch/sd/b/balewski/2026_causalNet_exp_ver3c +shortN=daleN200_102308_a5d24b + +bagTag="" +emFitName="" +outFitName="" +bagFrac=0.8 +bagEpochs=240 +numScrambles=10 +batchSize=4096 + +while [[ $# -gt 0 ]]; do + case "$1" in + --bagTag) + bagTag="$2" + shift 2 + ;; + --bagTag=*) + bagTag="${1#*=}" + shift + ;; + --emFitName) + emFitName="$2" + shift 2 + ;; + --emFitName=*) + emFitName="${1#*=}" + shift + ;; + --outFitName) + outFitName="$2" + shift 2 + ;; + --outFitName=*) + outFitName="${1#*=}" + shift + ;; + --bag_frac) + bagFrac="$2" + shift 2 + ;; + --bag_frac=*) + bagFrac="${1#*=}" + shift + ;; + --epochs) + bagEpochs="$2" + shift 2 + ;; + --epochs=*) + bagEpochs="${1#*=}" + shift + ;; + --num_scrambles) + numScrambles="$2" + shift 2 + ;; + --num_scrambles=*) + numScrambles="${1#*=}" + shift + ;; + -h|--help) + echo "Usage: sbatch [--array=0-N] batchFitBags.slr [--bagTag TAG] [--emFitName NAME] [--outFitName NAME] [--bag_frac F] [--epochs E] [--num_scrambles P]" + exit 0 + ;; + *) + echo "ERROR: Unknown argument: $1" >&2 + exit 1 + ;; + esac +done + +if [[ -z "${bagTag}" || "${bagTag,,}" == "none" ]]; then + if [[ -n "${SLURM_ARRAY_JOB_ID:-}" ]]; then + bagTag="j${SLURM_ARRAY_JOB_ID: -4}" + else + bagTag="local" + fi +fi +if [[ -z "${emFitName}" ]]; then + emFitName="${shortN}_em${bagTag}" +fi +if [[ -z "${outFitName}" ]]; then + outFitName="${emFitName}_rmf${bagTag}" +fi + +echo "S: job=${SLURM_JOB_ID:-local} host=$(hostname)" +echo "S: array_job=${SLURM_ARRAY_JOB_ID:-none} array_task=${SLURM_ARRAY_TASK_ID:-none}" +echo "S: basePath=${basePath}" +echo "S: dataName=${shortN}" +echo "S: bag_idx=${bag_idx}" +echo "S: bagTag=${bagTag}" +echo "S: emFitName=${emFitName}" +echo "S: outFitName=${outFitName}" +echo "S: bagFrac=${bagFrac}" +echo "S: bagEpochs=${bagEpochs}" +echo "S: numScrambles=${numScrambles}" + +time torchrun --standalone --nnodes=1 --nproc_per_node=4 \ + ./prism_FDR_Bags_train3c.py \ + --basePath "${basePath}" \ + --dataName "${shortN}" \ + --emFitName "${emFitName}" \ + --outFitName "${outFitName}" \ + --bag_idx "${bag_idx}" \ + --bag_frac "${bagFrac}" \ + --epochs "${bagEpochs}" \ + --batch_size "${batchSize}" \ + --num_scrambles "${numScrambles}" \ + --init_A ref \ + --init_B ref \ + --null_A_init scrambled_data \ + --null_B_init scrambled_data + +echo +echo "S: FDR locked-bag fit done for bag_idx=${bag_idx}" diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/batch_scanTime.slr b/causal_net/nonStation_ver3c_states_EM_FDR/batch_scanTime.slr new file mode 100644 index 00000000..bf1fbe11 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/batch_scanTime.slr @@ -0,0 +1,28 @@ +#!/bin/bash +#SBATCH -N 6 +#SBATCH -C gpu +#SBATCH --gpus-per-node=4 +#SBATCH -A m2043 +# SBATCH --time=30:00 +# SBATCH -q debug +#SBATCH --time=16:28:00 +#SBATCH -q regular +#SBATCH --output=outj/%A.out +#SBATCH --licenses=scratch + +set -uo pipefail +SCRIPT_DIR=/global/homes/b/balewski/prjs/2026_UoI-VAR/causal_net/nonStation_ver3c_states_EM_FDR +cd "$SCRIPT_DIR" + +SUFFIX=r5c +LOGDIR="$SCRIPT_DIR" + +srun --exclusive -N1 --gpus-per-node=4 time ./big_scanTime_FDR.sh $SUFFIX 25 35 120 >"${LOGDIR}/L10${SUFFIX}" 2>&1 & +srun --exclusive -N1 --gpus-per-node=4 time ./big_scanTime_FDR.sh $SUFFIX 20 40 100 >"${LOGDIR}/L20${SUFFIX}" 2>&1 & +srun --exclusive -N1 --gpus-per-node=4 time ./big_scanTime_FDR.sh $SUFFIX 15 45 80 >"${LOGDIR}/L30${SUFFIX}" 2>&1 & +srun --exclusive -N1 --gpus-per-node=4 time ./big_scanTime_FDR.sh $SUFFIX 10 50 70 >"${LOGDIR}/L40${SUFFIX}" 2>&1 & +srun --exclusive -N1 --gpus-per-node=4 time ./big_scanTime_FDR.sh $SUFFIX 5 55 60 >"${LOGDIR}/L50${SUFFIX}" 2>&1 & +srun --exclusive -N1 --gpus-per-node=4 time ./big_scanTime_FDR.sh $SUFFIX 0 60 50 >"${LOGDIR}/L60${SUFFIX}" 2>&1 & + +wait +echo "All 6 scanTime tasks finished." diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/big_fit_bags.sh b/causal_net/nonStation_ver3c_states_EM_FDR/big_fit_bags.sh new file mode 100755 index 00000000..2bd6ea4c --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/big_fit_bags.sh @@ -0,0 +1,106 @@ +#!/bin/bash +# Example end-to-end PRISM-EM FDR bagging run. +# salloc -q interactive -C gpu -t 4:00:00 -N 1 -A m2043 + +set -euo pipefail +export OMP_NUM_THREADS=1 +SECONDS=0 + +module load pytorch +cd "$(dirname "$0")" + +# ---------- dataset selection ---------- +# Synthetic example +#basePath=/pscratch/sd/b/balewski/2026_causalNet_exp_ver3c +#shortN=daleN200_55e5a6_ff089c # N200 3/21 Hz +#shortN=daleN200_f33b3b_7d8ff1 # N200 6/11 Hz +#shortN=daleN200_74e6d6_e2b7d7 # N200 12/21 Hz +#shortN=daleN100_d4f303_4abc4c # N100 13/27 Hz +#numStates=2 + +# Experimental example +basePath=/pscratch/sd/b/balewski/2026_causalNet_exp_ver3c +#shortN=Canine_260324_r23_w0_1hz; numStates=2 +shortN=Canine_260324_r21_w0_1hz; numStates=1 +#timeRange=(0 3600) +#timeRange=(0 1800) +timeRange=(1800 3600) +#timeRange=(0 300) + +# ---------- reference EM: state discovery ---------- +numEmIters=12 +numEMepochs=2 +emBatchSize=4096 + +# ---------- FDR bags: locked A/B fitting ---------- +numBags=11 +bagFrac=0.8 +bagEpochs=180 +numScrambles=6 +bagBatchSize=4096 +runAggregate=1 + +runTag="$(python3 -c 'import secrets; print(secrets.token_hex(2))')" + +echo "basePath=$basePath" +echo "dataName=$shortN" +echo "runTag=$runTag" +echo "timeRange=${timeRange[*]} numStates=$numStates" +emFitName="${shortN}_em${runTag}" +fdrBagsName="${emFitName}_fdr${runTag}" +fdrAgrName="${fdrBagsName}_agr${runTag}" + +echo "emFitName=$emFitName" +echo "fdrBagsName=$fdrBagsName" +echo "fdrAgrName=$fdrAgrName" + +echo +echo "=== Reference EM fit: state discovery ===" +time torchrun --standalone --nnodes=1 --nproc_per_node=4 \ + ./prism_EM_train3c.py \ + --basePath "$basePath" \ + --dataName "$shortN" \ + --fitName "$emFitName" \ + --time_range_sec "${timeRange[0]}" "${timeRange[1]}" \ + --num_states "$numStates" \ + --num_em_iters "$numEmIters" \ + --m_epochs "$numEMepochs" \ + --batch_size "$emBatchSize" \ + --delay_em_iter_4_lrDecay 4 "$numEmIters" \ + --delay_em_iter_4_ArhoMax 6 "$numEmIters" \ + --delay_em_iter_4_Aprune 6 + +echo +echo "=== FDR bags: locked M-step A/B fits ===" +for ((bag=0; bag/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $pdf_file" + +# Open the PDF based on the specific environment +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + gv "$pdf_file" +else + echo "Environment not matched. PDF is ready at $pdf_file" +fi diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3FDR.sh b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3FDR.sh new file mode 100755 index 00000000..26f1a794 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3FDR.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +set -e + +doc_name=fitFDR_states_ver3c + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $pdf_file" + +# Open the PDF based on the specific environment +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + gv "$pdf_file" +else + echo "Environment not matched. PDF is ready at $pdf_file" +fi diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3acc.sh b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3acc.sh new file mode 100755 index 00000000..2ed1638e --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3acc.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +set -e + +doc_name=edgeMeterAccuracy_ver3c + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $pdf_file" + +# Open the PDF based on the specific environment +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + gv "$pdf_file" +else + echo "Environment not matched. PDF is ready at $pdf_file" +fi diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3deBias.sh b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3deBias.sh new file mode 100755 index 00000000..88b6128c --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3deBias.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +set -e + +doc_name=deBiasFit_states_ver3c + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $pdf_file" + +# Open the PDF based on the specific environment +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + gv "$pdf_file" +else + echo "Environment not matched. PDF is ready at $pdf_file" +fi diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3exper.sh b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3exper.sh new file mode 100755 index 00000000..87def2a1 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3exper.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +set -e + +doc_name=preproc_exper_ver3c + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $pdf_file" + +# Open the PDF based on the specific environment +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + gv "$pdf_file" +else + echo "Environment not matched. PDF is ready at $pdf_file" +fi diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3gen.sh b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3gen.sh new file mode 100755 index 00000000..ea56d579 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3gen.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +set -e + +doc_name=genDale_topo_states_ver3c + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $pdf_file" + +# Open the PDF based on the specific environment +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + gv "$pdf_file" +else + echo "Environment not matched. PDF is ready at $pdf_file" +fi diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3stability.sh b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3stability.sh new file mode 100755 index 00000000..49615767 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/build3stability.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +set -e + +doc_name=edgeMeterStability_ver3c + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $pdf_file" + +# Open the PDF based on the specific environment +if [ "$(uname -s)" = "Darwin" ]; then + open "$pdf_file" +elif [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + xdg-open "$pdf_file" +elif hostname | grep -qiE 'login|nid0'; then + gv "$pdf_file" +else + echo "Environment not matched. PDF is ready at $pdf_file" +fi diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/deBiasFit_states_ver3c.pdf b/causal_net/nonStation_ver3c_states_EM_FDR/docs/deBiasFit_states_ver3c.pdf new file mode 100644 index 00000000..2a4a6f36 Binary files /dev/null and b/causal_net/nonStation_ver3c_states_EM_FDR/docs/deBiasFit_states_ver3c.pdf differ diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/deBiasFit_states_ver3c.tex b/causal_net/nonStation_ver3c_states_EM_FDR/docs/deBiasFit_states_ver3c.tex new file mode 100644 index 00000000..c4f2f226 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/deBiasFit_states_ver3c.tex @@ -0,0 +1,507 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{tikz} +\usepackage{hyperref} + +\geometry{margin=1in} + +\title{Unpenalized Magnitude Recovery for Signed Biological Networks\\ +\large Stage (c): De-biased Refit via Bound-Constrained Optimization} +\author{Algorithm Specification} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +This document specifies a follow-on ``Stage (c)'' algorithm designed to recover unbiased edge magnitudes in a multi-state Poisson GLM for spiking neural networks. While the preceding EM-FDR-Bagging stage robustly identifies the topological support of the network (detecting true edges and assigning Dale's principle signs), its use of $L_1$ regularization systematically shrinks edge magnitudes. In addition, repeated enforcement of a spectral-radius bound can further scale down the fitted connectivity matrix when the unconstrained update would exceed the stability limit. To isolate magnitude estimation from structural discovery, we define a bound-constrained optimization task that freezes the discovered topological mask, removes the sparsity penalty, and enforces biological signs. The recommended implementation fits compact active-parameter vectors rather than dense mostly-zero matrices: signed Dale-constrained entries are represented by squared parameters, free entries are copied directly, and the full matrix $A$ is rebuilt in the forward pass. +\end{abstract} + +% =============================================================== +\section{Introduction and Background} +% =============================================================== + +The full analysis pipeline is organized into three stages. \textbf{Stage (a)} +is the ordinary PRISM-EM fit on the chosen spike-train window: it estimates a +shared connectivity matrix, state-dependent biases, and a reference decoded +state sequence; this core fitter is specified in +\texttt{fitEM\_states\_ver3c.tex}. \textbf{Stage (b)} is the +EM-FDR-Bagging procedure specified in +\texttt{fitFDR\_states\_ver3c.tex}: it reuses the Stage (a) reference state +labels, performs bagged locked M-step refits, calibrates edges against +time-scrambled null fits, and returns a statistically vetted support mask +with conservative edge magnitudes. \textbf{Stage (c)} is the de-biased +refit described in this document: it freezes the Stage (b) support and +Dale-sign decisions, removes the \(L_1\) sparsity penalty, and refits final +edge magnitudes and state biases under the remaining biological and +stability constraints. + +\subsection{Input Data and Model Summary} +The input consists of a discrete spike-train recording of $N$ neurons over $T$ time bins of width $\Delta t$, yielding a count matrix $Y \in \mathbb{R}^{T \times N}$. The system is modeled as a non-stationary, multi-state Poisson Generalized Linear Model (GLM). The conditional firing rate for neuron $i$ at time $t$ depends on a shared connectivity matrix $A \in \mathbb{R}^{N \times N}$ and a state-dependent bias vector $B_{\hat{S}_t}$: +\begin{equation} +\label{eq:model-rate} +\lambda_t = \exp\big(\min(A Y_{t-1} + B_{\hat{S}_t},\ \eta_{\text{clip}})\big)\,\Delta t, \qquad Y_t \sim \text{Poisson}(\lambda_t) +\end{equation} +where $\hat{S}_{1:T}$ is the hard-decoded state sequence used by the M-step, and $\eta_{\text{clip}}$ prevents numerical overflow. In locked mode $\hat{S}_{1:T}$ is copied from Stage (b); in refit mode it is updated by the Stage (c) E-step and Viterbi decode before the following de-biased M-step. + +\subsection{The Preceding FDR-Bagging Stage} +Prior to this stage, an EM-FDR-Bagging routine was run. It drew subsamples of lag pairs and fit the model using an $L_1$ penalty (Lasso shrinkage) to suppress noise. Edges that survived a source-pooled null hypothesis across multiple bags were retained. Furthermore, source neurons were classified based on their total outgoing edge sum: strictly positive columns were labeled Excitatory, strictly negative columns were labeled Inhibitory, and intermediate columns were labeled Undecided. + +The output of the bagging stage provides: +\begin{enumerate} + \item A vetted topological mask indicating which edges are statistically significant. + \item A categorical label for each source node: \texttt{Excitatory}, \texttt{Inhibitory}, or \texttt{Undecided}. + \item Approximate, heavily $L_1$-shrunk point estimates for the connectivity ($\bar{A}$) and state biases ($\bar{B}$). +\end{enumerate} + +% =============================================================== +\section{Motivation: The Need for a De-biased Refit} +% =============================================================== + +Because the FDR-Bagging stage relies on an $L_1$ penalty to identify structural zeros, all surviving non-zero edges in $\bar{A}$ are systematically biased toward zero. This is a well-known property of Lasso regression. The shrinkage is useful during edge discovery because it suppresses noisy weak coefficients, but it is undesirable once the support has already been selected. + +There is a second, model-specific source of underestimation. PRISM-EM also enforces a spectral-radius limit, $\rho(A)\le\rho_{\text{max}}$, to keep the recurrent GLM numerically stable. If an M-step update produces a matrix with $\rho(A)>\rho_{\text{max}}$, the projection rescales the effective matrix. This operation preserves the discovered edge pattern and signs, but it can reduce many edge magnitudes at once. Therefore, the Stage (b) aggregate $\bar{A}$ should be interpreted as a conservative, regularized support estimate, not as a final magnitude estimate. + +To recover the biological magnitudes of these connections, we perform a ``de-biased Lasso'' step: an unpenalized Maximum Likelihood Estimation (MLE) refit restricted to the selected topological support. The simplest mode keeps the Stage (b) state sequence locked, matching the FDR-bagging refits; the multi-state mode may instead alternate de-biased M-steps with state refitting. + +The state biases should be re-fitted in the same optimization. In Eq.~\eqref{eq:model-rate}, $A Y_{t-1}$ and $B_{\hat{S}_t}$ enter the log-rate additively, so residual coupling between the connectivity and the state-dependent baseline can remain after the support-discovery stage. Once the magnitudes in $A$ are allowed to expand without the $L_1$ penalty, the optimum values of $B$ may also shift. + +This is a bound-constrained optimization problem. We must maximize the Poisson log-likelihood while enforcing four strict rules: +\begin{enumerate} + \item \textbf{Structural Zeros:} Edges pruned in the bagging stage must remain exactly zero. + \item \textbf{Dale's Principle (Strict):} Excitatory nodes must have strictly non-negative outgoing edges. Inhibitory nodes must have strictly non-positive outgoing edges. + \item \textbf{Dale's Principle (Relaxed):} Edges originating from ``undecided'' nodes that survived pruning must be free to adjust their sign as dictated by the unpenalized data. + \item \textbf{Free Diagonal:} The diagonal elements $A_{ii}$ are not treated as tested biological edges and must be left free to vary. They therefore belong to the free parameter set, even when neuron $i$ has an excitatory or inhibitory Dale label. +\end{enumerate} + +Figure~\ref{fig:stage-c-debias-cartoon} summarizes the Stage (c) data flow: +Stage (b) supplies support, signs, initial values, and states; Stage (c) +turns only the surviving entries into active parameters and optimizes the +unpenalized likelihood under fixed structural constraints. + +\begin{figure}[ht] +\centering +\begin{tikzpicture}[x=1cm,y=1cm] + % Stage B aggregate input + \draw[rounded corners, thick, gray!70, fill=gray!8] (0.55,1.10) rectangle (4.00,5.05); + \node[gray!80!black] at (2.28,4.78) {\small Lasso+FDR Mask}; + + \begin{scope}[shift={(0.87,4.35)}, x=0.47cm, y=0.47cm] + \foreach \c in {0,1,2} { + \fill[red!12] (\c,0.10) rectangle ++(1,0.28); + \node[red!75!black] at (\c+0.5,0.24) {\tiny\bfseries +}; + } + \foreach \c in {3,4} { + \fill[blue!12] (\c,0.10) rectangle ++(1,0.28); + \node[blue!70!black] at (\c+0.5,0.24) {\tiny\bfseries -}; + } + \fill[green!14] (5,0.10) rectangle ++(1,0.28); + \node[green!45!black] at (5.5,0.24) {\tiny\bfseries ?}; + + \foreach \r in {0,1,2,3,4,5} { + \foreach \c in {0,1,2,3,4,5} { + \draw[gray!45, fill=white] (\c,-\r) rectangle ++(1,-1); + } + } + \foreach \k in {0,1,2,3,4,5} { + \fill[green!14] (\k,-\k) rectangle ++(1,-1); + \node[green!45!black] at (\k+0.5,-\k-0.5) {\tiny \(\delta_{\k}\)}; + } + \foreach \r/\c/\lab in { + 1/0/\alpha_0, 3/0/\alpha_1, + 0/1/\alpha_2, 4/1/\alpha_3, 5/1/\alpha_4, + 3/2/\alpha_5, 5/2/\alpha_6} { + \fill[red!12] (\c,-\r) rectangle ++(1,-1); + \node[red!75!black] at (\c+0.5,-\r-0.5) {\tiny \(\lab\)}; + } + \foreach \r/\c/\lab in {0/3/\beta_0, 2/3/\beta_1, 1/4/\beta_2} { + \fill[blue!12] (\c,-\r) rectangle ++(1,-1); + \node[blue!70!black] at (\c+0.5,-\r-0.5) {\tiny \(\lab\)}; + } + \foreach \r/\c/\lab in {2/5/\gamma_0, 4/5/\gamma_1} { + \fill[green!14] (\c,-\r) rectangle ++(1,-1); + \node[green!45!black] at (\c+0.5,-\r-0.5) {\tiny \(\lab\)}; + } + \foreach \r in {0,1,2,3,4,5} { + \foreach \c in {0,1,2,3,4,5} { + \draw[gray!45] (\c,-\r) rectangle ++(1,-1); + } + } + \end{scope} + \node[gray!80!black] at (2.28,1.29) {\scriptsize out edges (column)}; + + % Active parameterization + \draw[->, thick] (4.10,3.35) -- (5.65,3.35); + \node at (4.88,3.72) {\scriptsize freeze}; + \draw[rounded corners, thick, blue!65, fill=blue!5] (5.75,1.10) rectangle (9.95,5.05); + \node[blue!65!black] at (7.85,4.78) {\small Trainable Values}; + + \draw[rounded corners, thick, red!65!black, fill=red!7] (5.98,3.80) rectangle (9.72,4.32); + \node[red!70!black, align=center] at (7.85,4.06) + {\scriptsize \(+\)Dale: \(\alpha_0,\ldots,\alpha_6\)}; + \draw[rounded corners, thick, blue!65!black, fill=blue!7] (5.98,3.13) rectangle (9.72,3.65); + \node[blue!70!black, align=center] at (7.85,3.39) + {\scriptsize \(-\)Dale: \(\beta_0,\ldots,\beta_2\)}; + \draw[rounded corners, thick, green!45!black, fill=green!10] (5.98,2.46) rectangle (9.72,2.98); + \node[green!45!black, align=center] at (7.85,2.72) + {\scriptsize free: \(\delta_0,\ldots,\delta_5,\gamma_0,\gamma_1\)}; + \draw[rounded corners, thick, green!45!black, fill=green!10] (5.98,1.83) rectangle (9.72,2.31); + \node[green!45!black, align=center] at (7.85,2.07) + {\scriptsize free state biases: \(B\)}; + + % Forward pass and objective + \draw[->, thick] (10.05,3.35) -- (11.45,3.35); + \node at (10.75,3.72) {\scriptsize scatter}; + \draw[rounded corners, thick, violet!70, fill=violet!6] (11.55,1.10) rectangle (15.15,5.05); + \node[violet!70!black] at (13.35,4.78) {\small Forward pass: \(A\)}; + + \begin{scope}[shift={(11.94,4.35)}, x=0.47cm, y=0.47cm] + \foreach \c in {0,1,2} { + \fill[red!12] (\c,0.10) rectangle ++(1,0.28); + \node[red!75!black] at (\c+0.5,0.24) {\tiny\bfseries +}; + } + \foreach \c in {3,4} { + \fill[blue!12] (\c,0.10) rectangle ++(1,0.28); + \node[blue!70!black] at (\c+0.5,0.24) {\tiny\bfseries -}; + } + \fill[green!14] (5,0.10) rectangle ++(1,0.28); + \node[green!45!black] at (5.5,0.24) {\tiny\bfseries ?}; + + \foreach \r in {0,1,2,3,4,5} { + \foreach \c in {0,1,2,3,4,5} { + \draw[gray!45, fill=white] (\c,-\r) rectangle ++(1,-1); + } + } + \foreach \k in {0,1,2,3,4,5} { + \fill[green!14] (\k,-\k) rectangle ++(1,-1); + \node[green!45!black] at (\k+0.5,-\k-0.5) {\tiny \(\delta_{\k}\)}; + } + \foreach \r/\c/\lab in { + 1/0/\alpha_0, 3/0/\alpha_1, + 0/1/\alpha_2, 4/1/\alpha_3, 5/1/\alpha_4, + 3/2/\alpha_5, 5/2/\alpha_6} { + \fill[red!12] (\c,-\r) rectangle ++(1,-1); + \node[red!75!black] at (\c+0.5,-\r-0.5) {\tiny \(+\lab^2\)}; + } + \foreach \r/\c/\lab in {0/3/\beta_0, 2/3/\beta_1, 1/4/\beta_2} { + \fill[blue!12] (\c,-\r) rectangle ++(1,-1); + \node[blue!70!black] at (\c+0.5,-\r-0.5) {\tiny \(-\lab^2\)}; + } + \foreach \r/\c/\lab in {2/5/\gamma_0, 4/5/\gamma_1} { + \fill[green!14] (\c,-\r) rectangle ++(1,-1); + \node[green!45!black] at (\c+0.5,-\r-0.5) {\tiny \(\lab\)}; + } + \foreach \r in {0,1,2,3,4,5} { + \foreach \c in {0,1,2,3,4,5} { + \draw[gray!45] (\c,-\r) rectangle ++(1,-1); + } + } + \end{scope} + + % Spectral projection feedback + \draw[rounded corners, dashed, very thick, black, fill=white] + (5.55,-0.15) rectangle (15.75,0.80); + \node[black, align=center] at (10.65,0.33) + {\scriptsize spectral guard: if \(\rho(A)>\rho_{\max}\), use + \(s=\rho_{\max}/\rho(A)\)\\[-0.5mm] + \scriptsize \((\alpha_i,\beta_i)\leftarrow\sqrt{s}\,(\alpha_i,\beta_i),\qquad + (\gamma_i,\delta_i)\leftarrow s\,(\gamma_i,\delta_i)\)}; + \draw[->, black, thick] (13.35,1.08) -- (13.35,0.83); + \draw[->, black, thick] (7.85,0.83) -- (7.85,1.08); +\end{tikzpicture} +\caption{Stage (c) de-biased refit cartoon. Stage (b) determines which +off-diagonal edges are allowed to exist and which source columns have fixed +Dale signs. In this toy \(6\times6\) mask, \(\delta_0,\ldots,\delta_5\) +are free diagonal entries, \(\alpha_0,\ldots,\alpha_6\) are edges from +positive Dale columns, \(\beta_0,\ldots,\beta_2\) are edges from negative +Dale columns, and \(\gamma_0,\gamma_1\) are surviving undecided-source +edges. The de-biased trainer carries only active values: \(\alpha_i\) and +\(\beta_i\) for Dale-constrained surviving edges, \(\gamma_i\) for +undecided surviving edges, \(\delta_i\) for all diagonal entries, and \(B\) +for free state biases. Structural zero cells carry no optimizer state. The +forward pass scatters these values into the +effective connectivity matrix \(A\), rebuilds \(A\) for each minibatch, +evaluates the Poisson negative log-likelihood with \(\lambda_3=0\), handles +the states as locked or jointly refit according to the run mode, saves +\texttt{A\_debias} and \texttt{B\_debias}, and applies the spectral-radius +guard in the transformed parameter space.} +\label{fig:stage-c-debias-cartoon} +\end{figure} + +% =============================================================== +\section{Methodology: Active-Set Parameter Transformation} +% =============================================================== + +To enforce the above constraints without relying on manual clamping (which can disrupt the momentum-based Adam optimizer), we use a parameter transformation for $A$: the \textit{squared-parameter trick}. The complete trainable variable set is $(u,v,B)$: the active connectivity vectors defined below plus the state-bias matrix. The biases are re-fitted jointly with $A$ and are initialized from $\bar{B}$, but they do not need the Dale sign transform. The masks below define the mathematical structure of $A$, but the optimizer should not have to carry full $N\times N$ tensors full of zeros. + +We construct two static, mutually exclusive structural masks based on the FDR-bagging output: +\begin{itemize} + \item $\bm{M_{\text{Dale}}} \in \{-1, 0, 1\}^{N \times N}$: This mask handles off-diagonal surviving edges from signed source nodes. It contains $+1$ for surviving excitatory edges, $-1$ for surviving inhibitory edges, and $0$ otherwise. + \item $\bm{M_{\text{Free}}} \in \{0, 1\}^{N \times N}$: This mask handles entries whose sign is not Dale-constrained. It contains $1$ for surviving off-diagonal edges from undecided source nodes, $1$ for every diagonal element $A_{ii}$, and $0$ otherwise. +\end{itemize} + +The dense mathematical form is +\begin{equation} +\label{eq:dense-transform} +A = \big( M_{\text{Dale}} \odot U^{\circ 2} \big) + \big( M_{\text{Free}} \odot V \big) +\end{equation} +where $\odot$ denotes element-wise multiplication (Hadamard product), and $U^{\circ 2}$ denotes the element-wise square of $U$. This form is useful for reasoning about the constraints. For implementation, however, it is better to fit active vectors: +\begin{equation} +\label{eq:active-index-sets} +\mathcal{D}=\{(i,j):(M_{\text{Dale}})_{ij}\ne0\}, \qquad +\mathcal{F}=\{(i,j):(M_{\text{Free}})_{ij}=1\}, +\end{equation} +\begin{equation} +\label{eq:active-vectors} +u\in\mathbb{R}^{|\mathcal{D}|}, \qquad +v\in\mathbb{R}^{|\mathcal{F}|}. +\end{equation} +In each forward pass, initialize $A$ as an all-zero matrix and scatter the active parameters into it: +\begin{equation} +\label{eq:active-scatter} +A_{ij} = +\begin{cases} + (M_{\text{Dale}})_{ij}\,u_k^2, & (i,j)=\mathcal{D}_k,\\ + v_\ell, & (i,j)=\mathcal{F}_\ell,\\ + 0, & \text{otherwise.} +\end{cases} +\end{equation} +Thus each free entry, including every diagonal $A_{ii}$, is copied one-to-one from $v$, while each Dale-constrained edge is generated by squaring the corresponding element of $u$ and applying the fixed sign. + +\textbf{Why this works:} +Because $u$ is squared, $u_k^2$ is non-negative. Multiplying it by the sign stored in $M_{\text{Dale}}$ enforces excitatory or inhibitory signs by construction. Meanwhile, the unconstrained vector $v$ allows undecided off-diagonal edges and diagonal self-history terms to explore positive or negative values freely. The vector formulation is also cleaner numerically: pruned edges have no optimizer state, receive no stale momentum, and cannot accidentally re-enter through implementation mistakes. + +% =============================================================== +\section{Fitting Details and Safeguards} +% =============================================================== + +\subsection{Vanishing Gradients and Initialization} + +The primary pitfall of the squared-parameter trick lies in the chain rule of backpropagation. The gradient of the loss $\mathcal{L}$ with respect to a parameter $U_{ij}$ is: +\begin{equation} +\label{eq:dense-u-gradient} +\frac{\partial \mathcal{L}}{\partial U_{ij}} = \frac{\partial \mathcal{L}}{\partial A_{ij}} \cdot 2 U_{ij} \cdot (M_{\text{Dale}})_{ij} +\end{equation} +In the active-vector notation this becomes +\begin{equation} +\label{eq:active-u-gradient} +\frac{\partial \mathcal{L}}{\partial u_k} += \frac{\partial \mathcal{L}}{\partial A_{ij}} \cdot 2u_k\cdot (M_{\text{Dale}})_{ij}, +\qquad (i,j)=\mathcal{D}_k . +\end{equation} +If the optimizer pushes $u_k$ to exactly zero, or initializes it very close to zero, the gradient vanishes entirely. The parameter becomes ``stuck'' and cannot recover, effectively acting as an unintended structural pruning. + +\textbf{Mitigation via Initialization:} +To prevent vanishing gradients, initialize the learnable signed-edge parameters securely away from zero, using the approximate shrunk values ($\bar{A}$) from the FDR-bagging stage. +\begin{equation} +\label{eq:active-init} +\begin{aligned} +u_k^{\text{init}} &= \max\!\left(\sqrt{\big| \bar{A}_{ij} \big|},\ \epsilon_u\right), +&\qquad (i,j)&=\mathcal{D}_k,\\ +v_\ell^{\text{init}} &= \bar{A}_{ij}, +&\qquad (i,j)&=\mathcal{F}_\ell . +\end{aligned} +\end{equation} +Here $\epsilon_u$ is a small positive floor used only for Dale-constrained parameters whose bagged estimate is extremely close to zero. The bias vectors are initialized directly as $B^{\text{init}} = \bar{B}$. + +\subsection{Spectral-Radius Projection} + +To maintain numerical stability in the Poisson GLM, the effective matrix $A$ must not exceed a maximum spectral radius $\rho_{\text{max}}$ (typically 0.95). + +In a standard network, one applies a scalar scaling factor to $A$. In our transformed parameter space, we must apply the scaling to the underlying learnable variables $u$ and $v$ such that the resulting effective $A$ is correctly scaled. + +At the configured spectral-projection interval during the M-step, we construct the effective $A$ and calculate its spectral radius $\rho(A)$. If $\rho(A) > \rho_{\text{max}}$, we define a scaling factor: +\begin{equation} +\label{eq:spectral-scale} +\begin{aligned} +\gamma &= \frac{\rho_{\text{max}}}{\rho(A)},\\ +u &\leftarrow u \cdot \sqrt{\gamma},\\ +v &\leftarrow v \cdot \gamma . +\end{aligned} +\end{equation} +We then update the underlying active parameters \textit{in-place} according to Eq.~\eqref{eq:spectral-scale}. +Because the forward pass squares $u$, scaling $u$ by $\sqrt{\gamma}$ correctly results in the corresponding effective Dale-constrained elements of $A$ being scaled by exactly $\gamma$. Scaling $v$ by $\gamma$ does the same for free entries. Structural zeros are unaffected, diagonal terms remain free, and biological signs cannot flip. The spectral constraint is therefore compatible with the sign constraints, although it can still limit final magnitudes when the unconstrained MLE wants a matrix outside the stable region. + +% =============================================================== +\section{Implementation Plan for \texttt{prism\_deBiasFit3c.py}} +% =============================================================== + +The production implementation should be a PyTorch distributed trainer, launched +the same way as the existing Stage (a) and Stage (b) jobs: +\begin{verbatim} +module load pytorch +torchrun --standalone --nnodes=1 --nproc_per_node=4 \ + ./prism_deBiasFit3c.py \ + --basePath \ + --fdrFitName \ + --outFitName +\end{verbatim} +On Perlmutter this uses the four A100 GPUs assigned to one node. The script +should reuse distributed utilities from \texttt{PrismEM\_Workhorse3c.py} +(\texttt{init\_distributed}, broadcasts, barriers, seeding, pair loaders, +E-step helpers, simplex projection, and Viterbi decoding). The dense +unconstrained \texttt{SwitchingPoissonGLM}, \texttt{run\_full\_fit}, and +\texttt{run\_locked\_mstep} should not be reused directly because Stage (c) +needs the active-vector parameterization $(u,v,B)$ and no $L_1$ pruning. If +the implementation becomes long, put the Stage (c) model and training loop in +\texttt{PrismDeBias\_Workhorse3c.py}, leaving \texttt{prism\_deBiasFit3c.py} +as the command-line, I/O, and provenance wrapper. + +\subsection{Inputs and Initialization} + +The primary input is the Stage (b) aggregate fit +\texttt{/prismFit/.prismEM.npz}. The script should +recover the original spike-file stem and the reference EM fit name from the +existing metadata provenance, rather than asking the user to provide multiple +file names. It should then load the source spikes from +\texttt{spikesData/}, slice the same \texttt{time\_range\_bins}, and verify +that the spike window length matches the Stage (b) \texttt{S\_hat} and +\texttt{c\_hat} arrays. + +The Stage (c) support and initialization are: +\begin{itemize} + \item Use Stage (b) \texttt{selected\_mask} for the off-diagonal support. + \item Use Stage (b) \texttt{neuron\_type} for Dale-constrained source columns. + \item Put all diagonal entries into the free set $\mathcal{F}$. + \item Initialize Dale-constrained magnitudes from Stage (b) + \texttt{A\_hat}: $u_k=\max(\sqrt{|A_{ij}^{(b)}|},\epsilon_u)$, with the + sign supplied only by $M_{\text{Dale}}$. + \item Initialize free off-diagonal entries and all diagonal entries from + Stage (b) \texttt{A\_hat}: $v_\ell=A_{ij}^{(b)}$. + \item Initialize the bias matrix from Stage (b) \texttt{B\_hat}. +\end{itemize} +Stage (b) \texttt{A\_prune} must be copied to the output unchanged for +comparison, but it is not the fitted Stage (c) matrix. + +\subsection{State Handling} + +Expose a three-way command-line option: +\begin{itemize} + \item \texttt{--state\_mode locked}: keep Stage (b) states fixed and fit only + $(u,v,B)$. + \item \texttt{--state\_mode refit}: alternate E-step updates of $c_t$ and + Viterbi decoding with de-biased M-step updates of $(u,v,B)$. + \item \texttt{--state\_mode auto}: use \texttt{refit} when + \texttt{num\_states>1}, otherwise use \texttt{locked}. +\end{itemize} +All M-step minibatches should use destination-bin state labels: $S_t$ for +the lag pair $(Y_{t-1},Y_t)$. This matches the Stage (b) FDR-bagging +convention and the rate model in Eq.~\eqref{eq:model-rate}. + +\subsection{Training Hyperparameters} + +The script should read defaults from the Stage (b) and reference EM metadata, +but expose the important controls explicitly: +\begin{itemize} + \item \texttt{--num\_debias\_iters}: outer E/M iterations for + \texttt{state\_mode refit}. + \item \texttt{--warmup\_locked\_epochs}: optional fixed-state M-step epochs + before state refitting starts. + \item \texttt{--m\_epochs}, \texttt{--batch\_size}, + \texttt{--lr\_mstep}, and \texttt{--lr\_end\_factor}: Adam M-step controls. + \item \texttt{--pgd\_iter}, \texttt{--lr\_estep}, and \texttt{--lambda2}: + E-step controls used only when states are refit. The parameter + \texttt{--lambda2} is the temporal smoothness weight on the soft state + probabilities, + \[ + \lambda_2 \sum_t \|c_t-c_{t-1}\|_2^2 . + \] + In the projected-gradient E-step this contributes the gradient term + $2\lambda_2(c_t-c_{t-1})$. Larger values make the inferred state + probabilities change more slowly over time; smaller values allow more + rapid state switching. It does not regularize $A$ or $B$ directly. In + \texttt{state\_mode locked}, the states are copied from Stage (b), the + E-step is skipped, and \texttt{--lambda2} is effectively irrelevant. + \item \texttt{--rho\_max}, \texttt{--prescale\_m\_step\_4\_ArhoMax}, + \texttt{--delay\_iter\_4\_ArhoMax}, and + \texttt{--target\_iter\_4\_ArhoMax}: spectral-radius controls. + \item \texttt{--eps\_u} and \texttt{--seed}: numerical floor and RNG seed. +\end{itemize} +There is no Stage (c) $L_1$ penalty; record \texttt{lambda3=0} in metadata. + +\subsection{Output Schema and Provenance} + +The output should be written to +\texttt{/prismFit/.prismEM.npz}. Preserve the original +Stage (b) arrays for comparison and add unambiguous Stage (c) arrays: +\begin{itemize} + \item Keep Stage (b) \texttt{A\_hat}, \texttt{A\_prune}, and + \texttt{B\_hat} unchanged. + \item Save the direct Stage (c) fit as \texttt{A\_debias} and + \texttt{B\_debias}. + \item Save the original Stage (b) states as \texttt{S\_stageB}, + \texttt{c\_stageB}, and \texttt{S\_stageB\_CL}. + \item Save final display states as \texttt{S\_hat}, \texttt{c\_hat}, and + \texttt{S\_hat\_CL}. In locked mode these equal the Stage (b) states; in + refit mode they are the final Stage (c) states. + \item Save histories for E-step NLL, M-step loss/NLL, spectral radius, + spectral correction strength, learning rate, and effective nonzero count. + \item Save support diagnostics: Dale/free index counts, diagonal count, + selected-edge count, and optional arrays mapping active-vector entries + back to $(i,j)$ coordinates. +\end{itemize} +Do not perform a second post-fit pruning pass on \texttt{A\_debias}; the +purpose is to inspect the direct constrained MLE result. + +Metadata should continue the existing provenance scheme. Set +\texttt{fit\_type="prismEM\_deBias\_stageC"} and add a +\texttt{deBias\_stageC} block containing the program name, input Stage (b) +file/name, recovered reference EM name, source spike file/name, +\texttt{state\_mode}, support construction rules, initialization rules, +hyperparameters, output name, elapsed time, and the trainable parameter +counts. The top-level \texttt{provenance} dictionary should retain all +Stage (a)/(b) entries and add \texttt{deBias\_stageC\_input}, +\texttt{deBias\_stageC\_file}, and \texttt{deBias\_stageC\_outFitName}. + +\subsection{Plotting Compatibility} + +\texttt{prism\_EM\_eval3c.py} should detect Stage (c) files from +\texttt{fit\_type} or the presence of \texttt{deBias\_stageC}. For Stage (c) +outputs it should use \texttt{A\_debias}/\texttt{B\_debias} as the primary +display matrices and include ``Stage (c) de-biased refit'' in plot titles. +For Stage (b) files it should keep the current behavior and title plots as +``Stage (b) FDR aggregate''. The Stage (c) output should nevertheless retain +all relevant arrays so that a dedicated plotter can be built later without +rerunning the fit. + +% =============================================================== +\section{Algorithm Summary for Implementation} +% =============================================================== + +An agent tasked with implementing Stage (c) should follow these explicit steps: + +\begin{enumerate} + \item \textbf{Load Data:} Load the Stage (b) aggregate, recover the source + spike file from provenance, slice the Stage (b) time window, and load + Stage (b) states and aggregates. + \item \textbf{Construct Static Masks and Index Sets:} Instantiate constant tensors $M_{\text{Dale}}$ and $M_{\text{Free}}$ based on Stage (b) \texttt{selected\_mask}, \texttt{neuron\_type}, and the diagonal. Put signed off-diagonal surviving edges in $\mathcal{D}$ and put undecided surviving off-diagonal edges plus all diagonal entries in $\mathcal{F}$. + \item \textbf{Initialize Parameters:} + \begin{itemize} + \item $u_k \leftarrow \max(\sqrt{|A^{(b)}_{ij}|},\epsilon_u)$ for $(i,j)=\mathcal{D}_k$ \ (requires gradient) + \item $v_\ell \leftarrow A^{(b)}_{ij}$ for $(i,j)=\mathcal{F}_\ell$ \ (requires gradient) + \item $B \leftarrow B^{(b)}$ \ (requires gradient) + \end{itemize} + \item \textbf{Training Loop (Minibatch Adam):} + \begin{itemize} + \item \textit{Forward Pass:} Construct a zero matrix $A$, scatter $(M_{\text{Dale}})_{ij}u_k^2$ into $(i,j)=\mathcal{D}_k$, and scatter $v_\ell$ into $(i,j)=\mathcal{F}_\ell$. + \item \textit{State Handling:} In locked mode use Stage (b) destination-bin labels. In refit mode update $c_t$ by the E-step, Viterbi-decode $S_t$, and then use destination-bin labels in the M-step. + \item \textit{Loss:} Compute negative Poisson log-likelihood using $A$ and $B_{\hat{S}_t}$, with no $L_1$ penalty ($\lambda_3 = 0$). + \item \textit{Backward Pass:} Call \texttt{loss.backward()} and \texttt{optimizer.step()}. + \end{itemize} + \item \textbf{Spectral Projection (End of interval):} + \begin{itemize} + \item Construct current effective $A$ under \texttt{torch.no\_grad()}. + \item Compute $\rho(A)$. If $\rho(A) > \rho_{\text{max}}$, compute $\gamma = \rho_{\text{max}} / \rho(A)$. + \item In-place multiply $u$ by $\sqrt{\gamma}$ and $v$ by $\gamma$. + \end{itemize} + \item \textbf{Save Output:} Once converged, construct the final effective matrix and save it as \texttt{A\_debias}; save the jointly fitted biases as \texttt{B\_debias}. Preserve Stage (b) \texttt{A\_hat}, \texttt{A\_prune}, and \texttt{B\_hat} unchanged for comparison. Do not run a second pruning pass on \texttt{A\_debias}. +\end{enumerate} + +\end{document} diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterAccuracy_ver3c.pdf b/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterAccuracy_ver3c.pdf new file mode 100644 index 00000000..4e736535 Binary files /dev/null and b/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterAccuracy_ver3c.pdf differ diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterAccuracy_ver3c.tex b/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterAccuracy_ver3c.tex new file mode 100644 index 00000000..aa704c88 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterAccuracy_ver3c.tex @@ -0,0 +1,247 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{geometry} +\usepackage{booktabs} + +\geometry{margin=0.85in} + +\title{Metrics for Accuracy of Reconstructed Connectivity Graph by PRISM-FDR Fits\\for Synthetic data} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +\texttt{edgeMeterAccuracy3c.py} evaluates final aggregate PRISM-FDR connectivity +fits on synthetic data where the true directed graph is known. The evaluated +matrix is \texttt{A\_prune}: the final matrix after bagging, FDR thresholding, +stability selection, source-neuron type classification, and Dale-sign pruning. +The truth edge mask is \texttt{E\_true}; \texttt{A\_true} supplies the true +sign of each true edge. All metrics are computed on directed off-diagonal +pairs only. The default plots \texttt{-p abc} summarize recovery quality, +raw error counts, and selected confusion matrices. +\end{abstract} + +\tableofcontents + + +% =============================================================== +\section{Inputs and Edge Labels} +\label{sec:inputs} +% =============================================================== + +The script reads one or more aggregate fit files from \texttt{prismFit/}, +using a common \texttt{--fitNameTrunk} and a list of \texttt{--fitTags}. Each +fit must point through its provenance metadata to the same synthetic truth +file in \texttt{truthDale/}. The script aborts if different fit tags refer to +different truth graphs. + +The matrix convention is source-column: +\[ + A_{ij} = \hbox{effect of source neuron } j + \hbox{ on target neuron } i . +\] +The candidate graph consists of all off-diagonal directed pairs: +\[ + (i,j), \quad i\ne j . +\] +Diagonal dynamics are useful for the fitted model, but they are not counted +as graph-recovery edges. + +For each off-diagonal candidate, the true label is +\[ +\mathrm{label}^{\mathrm{true}}_{ij} = +\begin{cases} ++ & E^{\mathrm{true}}_{ij}\ne0 \hbox{ and } A^{\mathrm{true}}_{ij}>0,\\ +- & E^{\mathrm{true}}_{ij}\ne0 \hbox{ and } A^{\mathrm{true}}_{ij}<0,\\ +0 & E^{\mathrm{true}}_{ij}=0 . +\end{cases} +\] +The recovered label comes from the final \texttt{A\_prune}: +\[ +\mathrm{label}^{\mathrm{reco}}_{ij} = +\begin{cases} ++ & A^{\mathrm{prune}}_{ij} > \varepsilon,\\ +- & A^{\mathrm{prune}}_{ij} < -\varepsilon,\\ +0 & |A^{\mathrm{prune}}_{ij}|\le\varepsilon . +\end{cases} +\] +The default tolerance is \(\varepsilon\) is the minimal edge weights from bagging algorithm. This keeps the test +effectively binary while avoiding accidental counting of tiny numerical +residue as recovered edges. + +% =============================================================== +\section{Plot Group \texttt{-p a}: Metric Curves} +\label{sec:plot-a} +% =============================================================== + +Plot group \texttt{a} is the main overview canvas. It contains six panels. +The x-axis is the data duration in minutes when the input fits correspond to +a scan over recording length. + +\subsection{Presence recovery} + +The presence panel collapses signed labels to a binary edge-present decision. +It shows: +\[ +\mathrm{precision}=\frac{\mathrm{TP}}{\mathrm{TP}+\mathrm{FP}},\qquad +\mathrm{recall}=\frac{\mathrm{TP}}{\mathrm{TP}+\mathrm{FN}}, +\] +\[ +\mathrm{F1}= +\frac{2\,\mathrm{precision}\,\mathrm{recall}} +{\mathrm{precision}+\mathrm{recall}}, +\] +and the Matthews correlation coefficient +\[ +\mathrm{MCC}= +\frac{\mathrm{TP}\,\mathrm{TN}-\mathrm{FP}\,\mathrm{FN}} +{\sqrt{(\mathrm{TP}+\mathrm{FP})(\mathrm{TP}+\mathrm{FN}) +(\mathrm{TN}+\mathrm{FP})(\mathrm{TN}+\mathrm{FN})}} . +\] +For sparse graphs, MCC is often the best single-number fidelity diagnostic +because it uses all four confusion counts without being dominated by the +large number of true negatives. Precision and recall remain essential +because they identify whether an MCC change is caused mainly by false edges +or missed true edges. + +\subsection{Sign-aware recovery} + +The sign-aware panel asks a stricter question: did the reconstruction assign +the correct signed edge label in \(\{-,0,+\}\)? It shows signed precision, +signed recall, signed F1, sign-aware multiclass MCC, and sign-flip rate. +Wrong-sign recovered true edges count as both signed false positives and +signed false negatives. This panel distinguishes topology recovery from Dale +sign recovery. + +\subsection{Recovered source type} + +The source-type panel shows how many neurons were classified by the aggregate +fit as excitatory, inhibitory, or undecided. This is a neuron-level Dale +diagnostic, not an edge-level metric. A high undecided count means the fit +may have recovered some edges but did not collect enough consistent outgoing +evidence to assign clean source type. + +\subsection{False discoveries} + +The false-discovery panel shows realized false-discovery proportion: +\[ + \mathrm{FDP}=\frac{\mathrm{FP}}{\max(1,\mathrm{TP}+\mathrm{FP})}, +\] +and compares it to the aggregate file's approximate stability-selection +bound normalized by selected-edge count. The stability bound is a useful +conservativeness diagnostic, but it is not a formal proof of FDR control for +these overlapping bags. + +\subsection{Graph density} + +The density panel compares true and recovered graph density: +\[ + \rho_{\mathrm{true}} = + \frac{\#\hbox{ true off-diagonal edges}}{N(N-1)},\qquad + \rho_{\mathrm{reco}} = + \frac{\#\hbox{ recovered off-diagonal edges}}{N(N-1)} . +\] +This panel quickly reveals whether a good recall score is being achieved by +making the recovered graph too dense. + +\subsection{Undecided-source diagnostics} + +The undecided-source panel shows the fraction of neurons classified as +undecided, the fraction of recovered edges emitted by undecided sources, and +the fraction of true edges whose source was classified as undecided. This is +important because undecided neurons are treated as an abstention category: +their nonzero \texttt{A\_prune} edges still count in edge recovery, but they +are flagged as coming from sources with ambiguous Dale evidence. + +% =============================================================== +\section{Plot Group \texttt{-p b}: Raw Count Bars} +\label{sec:plot-b} +% =============================================================== + +Plot group \texttt{b} shows two stacked bar charts. The first chart is titled +\textit{edge is present (counts)}. For each fit tag, the stack is ordered: +\[ + \mathrm{TP}\quad \hbox{bottom}, \qquad + \mathrm{FN}\quad \hbox{middle}, \qquad + \mathrm{FP}\quad \hbox{top}. +\] +True negatives are intentionally omitted. In sparse graphs, the number of +true negatives can be very large and would hide the counts that matter for +reconstruction failures. + +The second chart is titled \textit{edge sign is correct (counts)} and uses +the same stacked layout for signed TP, signed FN, and signed FP. This shows +how much of the remaining error comes from missed signed edges, spurious +signed edges, or wrong signs on recovered true edges. + +Raw counts matter because ratio metrics can hide scale. For example, two +fits can have similar precision but very different numbers of false edges. +The count canvas is therefore the best place to inspect the absolute error +budget. + +% =============================================================== +\section{Plot Group \texttt{-p c}: Selected Confusion Matrices} +\label{sec:plot-c} +% =============================================================== + +Plot group \texttt{c} displays detailed confusion matrices for selected fit +tags. The current plotting call selects \texttt{tagIdxL=[2,-1]}, meaning the +third provided tag and the last provided tag. The value \(-1\) always means +the last tag. + +For each selected tag, the left matrix is the edge-label confusion matrix: +rows are true labels and columns are recovered labels in the order +\(\{-,0,+\}\). It separates: +\begin{itemize} + \item missed inhibitory and excitatory edges: true \(-\) or \(+\) recovered + as \(0\); + \item spurious inhibitory and excitatory edges: true \(0\) recovered as + \(-\) or \(+\); + \item sign flips: true \(+\) recovered as \(-\), or true \(-\) recovered as + \(+\). +\end{itemize} + +For each selected tag, the right matrix is the neuron type confusion matrix. +Rows are true source type, excitatory or inhibitory. Columns are recovered +source type: excitatory, inhibitory, or undecided. This matrix diagnoses +Dale source classification separately from edge topology. + +% =============================================================== +\section{Saved Metric File} +\label{sec:output} +% =============================================================== + +The script always saves the computed metric arrays to +\texttt{edgeReco/\{outName\}.npz}. The default output name is +\texttt{\{fitNameTrunk\}\_ermHASH4}. The saved arrays include: +\begin{itemize} + \item fit names, fit tags, input files, and truth file; + \item duration, number of neurons, number of bags, bag fraction, FDR + quantile, stability threshold, and \(\varepsilon\); + \item TP, FP, FN, TN, precision, recall, F1, MCC, FDP, graph densities; + \item signed precision, signed recall, signed F1, signed MCC, sign-flip + count, and sign-flip rate; + \item edge-label confusion matrices and neuron-type confusion matrices for + every tag. +\end{itemize} + +% =============================================================== +\section{Possible Future Upgrade} +\label{sec:future} +% =============================================================== + +A threshold-free precision--recall study could be added later. The most +natural candidate score is not final \(|\texttt{A\_prune}|\), because +\texttt{A\_prune} has already passed per-bag FDR, stability selection, and +Dale pruning. Better future scores may come from \texttt{selection\_frequency}, +\texttt{A\_mean\_selected}, null thresholds, or per-edge diagnostics saved +before the final stability cut. Such an upgrade would complement the current +absolute operating-point metrics rather than replace them. + +\end{document} diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterStability_ver3c.pdf b/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterStability_ver3c.pdf new file mode 100644 index 00000000..5c39881b Binary files /dev/null and b/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterStability_ver3c.pdf differ diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterStability_ver3c.tex b/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterStability_ver3c.tex new file mode 100644 index 00000000..aa45b286 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/edgeMeterStability_ver3c.tex @@ -0,0 +1,254 @@ +\documentclass[11pt]{article} + +\usepackage[utf8]{inputenc} +\usepackage[T1]{fontenc} +\usepackage{amsmath} +\usepackage{amssymb} +\usepackage{amsfonts} +\usepackage{bm} +\usepackage[margin=1in]{geometry} +\usepackage{enumitem} +\usepackage{hyperref} + +\hypersetup{ + colorlinks=true, + linkcolor=blue, + urlcolor=blue, + citecolor=blue +} + +\title{Stability Metrics for Reconstructed Connectivity Graphs\\by PRISM-FDR Fits} +\author{Jan Balewski, NERSC} +\date{\today} + + +\begin{document} + +\maketitle + +\begin{abstract} +This document defines the mathematical framework and programmatic specifications for evaluating the stability of sparse directed graphs reconstructed by PRISM-FDR from real or simulated data. The framework is ground-truth-free: it measures how consistently the pipeline recovers the same structure and weights across $K \ge 2$ independent data subsets. It therefore cannot quantify what true graph information may be absent from all subsets simultaneously---only how much the subsets agree with one another. Applications include stationarity testing and determining the minimum data footprint $t_{\text{min}}$ for reliable inference. +\end{abstract} + +\tableofcontents + +% =============================================================== +\section{Input Specifications, Conventions, and Baselines} +% =============================================================== + +The diagnostic code processes $K \ge 2$ reconstructed graph files, one per data subset (e.g., consecutive time blocks, disjoint subsets, or expanding recording durations). + +\subsection*{Matrix Convention} + +Each data subset $k$ yields an $N \times N$ estimated weight matrix $\mathbf{A}^{(k)}$, where $N$ is the number of neurons ($50 \le N \le 500$). The matrix follows the source-column convention: +\[ +A_{ij} = \text{directed effect of source neuron } j \text{ on target neuron } i +\] + +\subsection*{Input Data per Data Subset $k$} + +For each data subset $k$, the input must supply: +\begin{itemize} + \item $\mathbf{A}^{(k)}_{\text{prune}}$: The final pruned matrix of mean edge weights after bagging, FDR thresholding, and Dale-sign pruning. The matrix is sparse (typical row density between 5\% and 20\%). Non-zero off-diagonal entries are the \emph{selection-conditional} bagged mean: the arithmetic mean of $\hat{A}^{(b)}_{ij}$ over only the bags $b$ in which edge $(i,j)$ passed the per-bag null threshold. Diagonal entries $A^{(k)}_{ii}$ are retained as fitted values; they are not subject to FDR or bagging pruning and may take any real value. + \item $\boldsymbol{\Sigma}^{(k)}$: An $N \times N$ matrix containing the selection-conditional bootstrap standard deviation $\sigma_{ij}$ of off-diagonal edge weights computed across bags. Specifically, + \[ + \sigma^{(k)}_{ij} = \widehat{\mathrm{SD}}^{(k)}_{ij} + = \sqrt{\frac{1}{B_{ij}-1}\sum_{b \in \mathcal{B}_{ij}} \bigl(\hat{A}^{(b)}_{ij} - \bar{A}^{(k)}_{ij}\bigr)^2}, + \] + where $\mathcal{B}_{ij}$ is the set of bags in which edge $(i,j)$ was selected and $B_{ij} = |\mathcal{B}_{ij}|$. Entries where $B_{ij} \le 1$ are stored as zero. The number of bootstrap bags is fixed at \texttt{numBags = 20}. + \item $\mathbf{V}^{(k)}_{\text{node}}$: An $N \times 1$ vector containing the final categorical assignment for each source neuron: Excitatory ($+$), Inhibitory ($-$), or Undecided ($0$). +\end{itemize} + +\subsection*{Numerical Zero-Floor and Global Active Edge Union} + +To avoid arbitrary hyperparameters, the numerical zero-floor ($\varepsilon$) is defined dynamically as the minimum non-zero absolute off-diagonal weight across all data subsets: +\[ +\varepsilon = \min_{k, i \neq j} \left\{ |A^{(k)}_{\text{prune}, ij}| \;\middle|\; A^{(k)}_{\text{prune}, ij} \neq 0 \right\} +\] +To protect correlation metrics from zero-inflation (where large blocks of matching zeros artificially inflate consistency scores), the framework constructs a \textbf{Global Active Edge Union} ($E_{\text{global}}$) across all $K$ data subsets: +\[ +E_{\text{global}} = \bigcup_{k=1}^{K} \left\{ (i,j) \;\middle|\; i \neq j \text{ and } |A^{(k)}_{\text{prune}, ij}| \ge \varepsilon \right\} +\] +All off-diagonal metrics are evaluated exclusively on pairs in $E_{\text{global}}$. $E_{\text{global}}$ also serves as the reference topology for the multi-subset comparison strategies in Section~\ref{sec:strategies}. + +% =============================================================== +\section{Off-Diagonal: Topology and Sign Alignment Metrics} +% =============================================================== + +When comparing any pair of data subsets, edge existence and sign consistency are assessed using metrics independent of absolute weight scales. + +\subsection*{A. Signed Jaccard Indices} + +To measure topological stability without influence from $0 \to 0$ matches, positive and negative edge sets are isolated for each data subset restricted to $E_{\text{global}}$: +\[ +E_k^+ = \left\{ (i,j) \in E_{\text{global}} \;\middle|\; A^{(k)}_{\text{prune}, ij} > \varepsilon \right\}, \qquad +E_k^- = \left\{ (i,j) \in E_{\text{global}} \;\middle|\; A^{(k)}_{\text{prune}, ij} < -\varepsilon \right\} +\] +Two Jaccard values are reported for every pairwise comparison: +\[ +J_+ = \frac{|E_1^+ \cap E_2^+|}{|E_1^+ \cup E_2^+|}, \qquad +J_- = \frac{|E_1^- \cap E_2^-|}{|E_1^- \cup E_2^-|} +\] +If either denominator is zero (no positive or no negative edges across both data subsets), the corresponding $J$ is set to \texttt{NaN}. + +\subsection*{B. Sign Match Rate (SMR)} + +The SMR measures sign consistency on the subset of edges where both data subsets agree a connection exists, ignoring edges absent in either one: +\[ +\text{SMR} = \frac{|(E_1^+ \cap E_2^+) \cup (E_1^- \cap E_2^-)|}{|(E_1^+ \cup E_1^-) \cap (E_2^+ \cup E_2^-)|} +\] +The denominator counts edges that both data subsets call present (regardless of sign); the numerator counts those where both also agree on sign. + +% =============================================================== +\section{Off-Diagonal: Parametric Magnitude Consistency Metrics} +% =============================================================== + +Edge weight magnitudes are extracted as 1D vectors over $E_{\text{global}}$. For a pair of data subsets, let +\[ + X_m = |A^{(k_1)}_{i_m j_m}|, \qquad + Y_m = |A^{(k_2)}_{i_m j_m}|, +\] +where $(i_m,j_m)$ runs over the selected off-diagonal edge universe. Entries are zero for edges absent in a given subset. For source-type Spearman curves, this universe is additionally restricted by source column: the excitatory scope keeps edges whose source neuron is classified excitatory in either compared subset, and the inhibitory scope keeps edges whose source neuron is classified inhibitory in either compared subset. + +\subsection*{A. Unweighted Spearman Rank Correlation ($r_s$)} + +Because regularization scales coefficients dynamically with data length, rank-order preservation is prioritized over linear agreement. The ranks are ranks of edge magnitude, not signed weight, so excitatory and inhibitory source scopes are treated symmetrically. +\begin{enumerate} + \item Convert the magnitude vectors $X$ and $Y$ into rank vectors $R_X$ and $R_Y$ using ascending ranks, with average ranks for ties as in \texttt{scipy.stats.rankdata}. Thus the weakest magnitudes have the smallest ranks and the strongest magnitudes have the largest ranks. + \item Compute the Pearson correlation of the rank vectors: +\end{enumerate} +\[ +r_s = \frac{\sum_m (R_{X,m} - \bar{R}_X)(R_{Y,m} - \bar{R}_Y)}{\sqrt{\sum_m (R_{X,m} - \bar{R}_X)^2 \sum_m (R_{Y,m} - \bar{R}_Y)^2}} +\] + +\subsection*{B. Error-Aware Weighted Spearman Correlation ($r_s^w$)} + +This metric down-weights edges with high estimation instability. The weight for edge $(i,j)$ when comparing data subsets $k_1$ and $k_2$ is the inverse of the mean selection-conditional bootstrap SD across both subsets, regularized by $\varepsilon$: +\[ +w_{ij} = \frac{1}{\dfrac{\sigma^{(k_1)}_{ij} + \sigma^{(k_2)}_{ij}}{2} + \varepsilon} +\] +where $\sigma^{(k)}_{ij} = 0$ for edges absent in data subset $k$. The $+\varepsilon$ floor prevents division by zero using the same dynamically computed value from Section~1, introducing no new hyperparameter. + +The weighted Spearman replaces the Pearson step with its weighted analogue on $R_X$ and $R_Y$: +\[ +r_s^w = \frac{\sum_m w_m (R_{X,m} - \bar{R}_X^w)(R_{Y,m} - \bar{R}_Y^w)} + {\sqrt{\sum_m w_m (R_{X,m} - \bar{R}_X^w)^2 \;\cdot\; \sum_m w_m (R_{Y,m} - \bar{R}_Y^w)^2}} +\] +where $\bar{R}_X^w = \sum_m w_m R_{X,m} / \sum_m w_m$ and similarly for $\bar{R}_Y^w$. + +% =============================================================== +\section{Diagonal: Self-Coupling Stability Metrics} +\label{sec:diagonal} +% =============================================================== + +The diagonal entries $A^{(k)}_{ii}$ represent self-coupling of each neuron and are fitted without FDR thresholding or bagging pruning; they may take any real value. One magnitude-change metric is defined for the diagonal. + +\subsection*{A. Diagonal Magnitude Change} + +Let $d^{(k)} \in \mathbb{R}^N$ be the vector of diagonal entries from data subset $k$, i.e.\ $d^{(k)}_i = A^{(k)}_{ii}$. For each pairwise comparison of data subsets $k_1$ and $k_2$, compute the mean absolute diagonal difference: +\[ +\Delta_{\text{diag}} = \frac{1}{N}\sum_i \left|d^{(k_1)}_i - d^{(k_2)}_i\right| +\] +This metric is sensitive to absolute changes in fitted self-coupling values and is not invariant to global shifts or scaling. + +% =============================================================== +\section{Per-Neuron-Type Breakdown of Off-Diagonal Metrics} +\label{sec:neurontype} +% =============================================================== + +Off-diagonal Jaccard metrics are computed over three edge scopes: +\begin{enumerate} + \item \textbf{All edges}: the full $E_{\text{global}}$ index. + \item \textbf{Excitatory-source edges}: restrict to $(i,j) \in E_{\text{global}}$ where source neuron $j$ is classified as Excitatory ($V^{(k)}_{\text{node},j} = +$) in \emph{either} compared data subset. + \item \textbf{Inhibitory-source edges}: restrict to $(i,j) \in E_{\text{global}}$ where source neuron $j$ is classified as Inhibitory ($V^{(k)}_{\text{node},j} = -$) in \emph{either} compared data subset. +\end{enumerate} +Undecided-source edges contribute to the all-edges scope but are excluded from the excitatory and inhibitory scopes. + +SMR is computed only for the all-edge scope. Spearman metrics are computed only for the excitatory-source and inhibitory-source scopes; all-edge Spearman is intentionally omitted. The compact plot group \texttt{a} displays all-edge topology/sign metrics and source-type Spearman metrics in one canvas. + +% =============================================================== +\section{Multi-Subset Comparison Strategies} +\label{sec:strategies} +% =============================================================== + +Two comparison strategies are supported, selected by \texttt{--compare\_mode}. + +\subsection*{Strategy A: Each Data Subset vs.\ Last Subset (\texttt{--compare\_mode last})} + +The last data subset in the ordered list ($k = K$) serves as the reference. All other subsets ($k = 1, \dots, K-1$) are compared against it. The reference is not compared against itself. This mode is designed for studies where the last subset represents the most complete recording (e.g., the full time window $[0, 60]$\,min), and earlier subsets represent restricted windows (e.g., $[10, 50]$\,min, $[20, 40]$\,min). The x-axis is the duration of each non-reference subset, so the plot directly shows how shrinking the time window degrades stability relative to the full recording. + +This produces $K-1$ metric values per metric. + +\subsection*{Strategy B: Consecutive Pairwise Comparison (\texttt{--compare\_mode consecutive})} + +Compare ordered adjacent pairs: $(k=1$ vs.\ $k=2)$, $(k=2$ vs.\ $k=3)$, $\ldots$, $(k=K-1$ vs.\ $k=K)$. This produces $K-1$ metric values per metric and directly measures drift between neighboring data subsets. + +% =============================================================== +% =============================================================== +\section{Saved Metric File} +\label{sec:output} +% =============================================================== + +The script always saves the computed metric arrays to \texttt{edgeFidelity/\{outName\}.npz}. The default output name is \texttt{\{fitNameTrunk\}\_efmHASH4}. The saved arrays include: +\begin{itemize} + \item fit name trunk, fit tags, comparison mode, and $\varepsilon$; + \item comparison labels and x-axis values (one per comparison); + \item data subset tags and x-axis values (one per subset); + \item per-comparison pairwise Jaccard metrics for each scope (all, exc, inh): $J_+$, $J_-$; + \item all-edge SMR, stored only for the all-edge scope; + \item source-type Spearman metrics $r_s$ and $r_s^w$, stored only for excitatory-source and inhibitory-source scopes; + \item per-comparison mean absolute diagonal difference $\Delta_{\text{diag}}$. + \item per-subset source-rate heatmap arrays: shared single-rate bin edges and summed excitatory-plus-inhibitory outgoing edge counts per rate bin; +\end{itemize} + +% =============================================================== +\section{Plots} +\label{sec:plots} +% =============================================================== + +Plot groups are selected with \texttt{-p}. The default \texttt{-p a} produces the edge-fidelity summary. + +\paragraph{X-axis convention.} +\begin{itemize} + \item \texttt{consecutive}: x-axis is the midpoint duration between the two compared subsets (e.g., midpoint of 10\,min and 20\,min is 15\,min). Produces $K-1$ points connected by lines. + \item \texttt{last}: x-axis is the duration of the non-reference subset. Produces $K-1$ points connected by lines. The last subset is the reference and does not appear as a data point. +\end{itemize} + +\bigskip +\noindent\textbf{Plot group \texttt{a} --- Edge Fidelity Summary.} +A $3 \times 2$ canvas with the same layout for both \texttt{consecutive} and \texttt{last} comparison modes. Top-left: all-edge Jaccard curves $J_+$ and $J_-$. Top-right: all-edge Signal Match Rate. Middle-left: overlaid excitatory-source and inhibitory-source Spearman $r_s$. Bottom-left: overlaid excitatory-source and inhibitory-source weighted Spearman $r_s^w$. Middle-right: mean absolute diagonal difference $\Delta_{\text{diag}}$. The top-right and middle-right panels use automatic y-axis scaling. The bottom-right panel is intentionally blank. + +\bigskip +\noindent\textbf{Plot group \texttt{b} --- Weight and Edge-Count Stats.} +A $1 \times 4$ canvas using per-subset values rather than pairwise comparison values. Panel 1 shows the positive and negative Dale-pruning weight thresholds, \texttt{min\_posW} and \texttt{max\_negW}. Panel 2 shows off-diagonal weight ranges for excitatory-source and inhibitory-source edges, including median trends. Panel 3 shows the number of off-diagonal edges by source type (excitatory, inhibitory, undecided). Panel 4 is a 2D heat map with duration on the x-axis and neuron single rate on the y-axis; color gives the summed number of outgoing excitatory-plus-inhibitory source edges in each rate bin. The y-axis uses at most 30 shared rate bins, and a color scale is drawn on the right. + +% =============================================================== +\section{Script Interface} +\label{sec:cli} +% =============================================================== + +The script is \texttt{edgeMeterFidelity3c.py}. It is invoked as: +\begin{verbatim} +python edgeMeterFidelity3c.py \ + --basePath \ + --fitNameTrunk \ + --fitTags ... \ + --compare_mode \ + [-p ] +\end{verbatim} + +\paragraph{Required arguments.} +\begin{description} + \item[\texttt{--basePath}] Root directory containing \texttt{prismFit/} and receiving \texttt{edgeFidelity/} and \texttt{plots/}. + \item[\texttt{--fitNameTrunk}] Common stem shared by all fit files. + \item[\texttt{--fitTags}] Ordered list of tags identifying each data subset. The order determines the sequence for consecutive mode and the reference for last mode. + \item[\texttt{--compare\_mode}] Either \texttt{consecutive} (default) or \texttt{last}. +\end{description} + +\paragraph{Optional arguments.} +\begin{description} + \item[\texttt{--outName}] Output metric stem in \texttt{edgeFidelity/}; default is \texttt{\{fitNameTrunk\}\_efmHASH4}. + \item[\texttt{-p}] Plot selector string (default \texttt{a}): \texttt{a} = edge fidelity summary, \texttt{b} = weight and edge-count stats. +\end{description} + +\end{document} diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/fitEM_states_ver3c.tex b/causal_net/nonStation_ver3c_states_EM_FDR/docs/fitEM_states_ver3c.tex new file mode 100644 index 00000000..5adb59ca --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/fitEM_states_ver3c.tex @@ -0,0 +1,481 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} + +\geometry{margin=0.8in} + +\title{PRISM-EM for Non-Stationary Poisson GLMs with States from Spike Trains:\\ +Model, Objective, and Implementation. Ver3c} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +This document describes the current \texttt{prism\_EM\_train3c.py} implementation +for fitting a non-stationary +Poisson GLM with latent state mixing (PRISM-EM)~\footnote{PRISM: PRobabilistic Inference of Sparse Mixed-state dynamics; ~~EM: Expectations Maximization}. Input spike trains are read from \texttt{spikesData/.spikes.npz} (non-stationary output of \texttt{gen\_nonStationarySpikes3c.py}). The model uses a shared connectivity matrix across states, state-specific bias vectors, and time-varying simplex coefficients. Training is performed by block coordinate descent: a sequential projected-gradient E-step for coefficients followed by Viterbi decoding to one-hot state assignments, and an Adam-based M-step for model parameters with off-diagonal L1 shrinkage and spectral-radius projection. Spectral-radius enforcement, L1 pruning, and learning-rate decay are each activated after a configurable number of EM iterations. Distributed training (\texttt{torchrun}) shards the E-step over time bins and synchronizes \(A\), \(B\), and \(c\) across GPUs. We summarize the exact objective used in code, the decoding method, all current default hyperparameters, initialization modes (data-driven vs random) for \(A\), \(B\), and \(C\), and the saved fit file layout. +\end{abstract} + +\tableofcontents +\newpage + +% --------------------------------------------------------------- +\section{EM Training for Non-Stationary Poisson GLM with states} +\label{sec:em} +% --------------------------------------------------------------- + +\subsection{Introduction and Model Assumptions} + +We consider a population of \(N\) neurons observed over \(T\) discrete time +bins of width \(\Delta t\) (seconds). The population activity is governed +by \(M\) latent dynamical states that switch over time. The key +structural assumption is that all states share a single +connectivity matrix \(A \in \mathbb{R}^{N \times N}\) but differ in +baseline excitability through per-state bias vectors +\(B_m \in \mathbb{R}^N,\; m = 0, \dots, M-1\) (0-based state indices in code). + +At each time bin \(t\), a coefficient vector +\(c_t \in \mathbb{R}^M\), constrained to the probability simplex +(\(c_{mt}\ge 0,\; \sum_m c_{mt}=1\)), selects the effective bias as a +convex combination: +\begin{equation} + B_t = \sum_{m=1}^{M} c_{mt}\, B_m = c_t^\top \mathbf{B}, +\end{equation} +where \(\mathbf{B} \in \mathbb{R}^{M \times N}\) stacks the \(M\) bias +vectors row-wise. + +Spike counts follow +\begin{equation} + \eta_t = A\,Y_{t-1} + B_t,\qquad + \tilde{\eta}_t = \min(\eta_t,\eta_{\mathrm{clip}}),\qquad + \mu_t = \exp(\tilde{\eta}_t)\,\Delta t,\qquad + Y_t \sim \mathrm{Poisson}(\mu_t). + \label{eq:glm} +\end{equation} + +\paragraph{Notation summary.} +\(N\): number of neurons; +\(M\): number of latent states; +\(T\): number of time bins; +\(\Delta t\): time-bin width (from metadata); +\(\eta_{\mathrm{clip}}\): clipping threshold (from metadata); +\(Y_t \in \mathbb{N}_0^N\): observed spikes; +\(\mu_t \in \mathbb{R}_+^N\): expected spike count per bin. + +% --------------------------------------------------------------- +\subsection{Objective Function} + +The training objective uses uniformly weighted (unweighted) Poisson NLL in both the E-step and M-step: +\begin{equation} +\mathcal{J} += +\underbrace{\sum_{t=2}^{T}\sum_{i=1}^{N} +\left[\mu_{t,i}-Y_{t,i}\left(\tilde{\eta}_{t,i}+\log\Delta t\right)\right]}_{\text{Poisson NLL}} ++ +\underbrace{\lambda_2\sum_{t=2}^{T}\|\vec c_t- \vec c_{t-1}\|^2}_{\text{temporal smoothness (E-step only)}} ++ +\underbrace{\lambda_3\,\overline{|A_{\mathrm{off}}|}}_{\text{L1 sparsity on off-diagonal (M-step only)}}, +\label{eq:objective} +\end{equation} +with simplex constraints \(c_t\in\Delta^{M-1}\) and spectral constraint +\(\rho(A)\le \rho_{\max}\), where the off-diagonal L1 term is +\begin{equation} + \overline{|A_{\mathrm{off}}|} + = + \frac{1}{N(N-1)} + \sum_{\substack{i,j=1\\i\neq j}}^{N}|A_{ij}|. + \label{eq:l1_offdiag} +\end{equation} +In the $\lambda_2$ term, the sum over states $m$ +is already implicit in the $L_2$ norm. + + +\subsection{Why Block Coordinate Descent (EM) over Joint Optimization} + +In principle one could jointly optimize \(A\), \(\mathbf{B}\), and +\(c_{1:T}\) by simultaneous gradient descent on the full objective +\(\mathcal{J}\). In practice, the EM-style alternation is preferable +for three structural reasons. + +\paragraph{Breaking the bilinear coupling.} +The coefficients \(c_t\) and the bias matrix \(\mathbf{B}\) enter the +log-rate linearly through the product \(c_t^\top\mathbf{B}\), which +then passes through an exponential. When both are free to move, the +landscape contains saddle points and flat valleys: a large \(c_{tm}\) +paired with a small \(B_m\) produces the same likelihood as the +reverse. Fixing one side at each step eliminates this ambiguity. The +Viterbi hard-assignment variant sharpens the separation further: +with one-hot \(c_t = e_{\hat{S}_t}\), each time bin selects exactly one +bias row, so the M-step reduces to \(M\) independent Poisson GLM +regressions and the bilinear degeneracy vanishes entirely. + +\paragraph{Exploiting temporal structure in the E-step.} +The coefficient vectors \(c_{1:T}\) number \(O(TM)\) and vastly +outnumber the \(O(N^2+MN)\) model parameters. Joint optimization +would need to balance learning rates across these disparate groups +while enforcing a simplex constraint at every time bin. By contrast, +the E-step decomposes into a single forward sweep of cheap +\(M\)-dimensional simplex-projected gradient updates, naturally +incorporating the temporal smoothness penalty +\(\lambda_2\|c_t-c_{t-1}\|^2\) through the sequential dependence on the +previous accepted \(c_{t-1}\). + +\paragraph{Avoiding the log-sum-exp mismatch.} +Under soft (probabilistic) state assignments the M-step must fit +parameters to the mixture rate +\(\sum_m c_{tm}\exp(\eta_{tm})\), whereas the model parameterizes the +log-rate as \(\sum_m c_{tm}\,\eta_{tm}\) inside a single exponential. +These two quantities differ whenever the assignment is not one-hot, +creating a systematic bias in the gradient signal for \(\mathbf{B}\). +The hard-EM strategy—Viterbi decode followed by one-hot +encoding—eliminates this mismatch by construction, at the cost of +losing the monotonic likelihood improvement guarantee of soft EM, a +trade-off that is empirically favorable. + +% --------------------------------------------------------------- +\subsection{Algorithm: Block Coordinate Descent (EM)} + +The non-convex problem is solved by alternating: +\begin{itemize} + \item E-step: update \(c_{1:T}\) with \(A,\mathbf{B}\) fixed, then Viterbi-decode to one-hot state assignments. + \item M-step: update \(A,\mathbf{B}\) with one-hot \(c_{1:T}\) fixed. +\end{itemize} +for \(K_{\mathrm{EM}}\) outer iterations. + +\subsubsection{E-step: Coefficient Inference via Simplex PGD, then One-Hot Encoding} + +Holding \((A,\mathbf{B})\) fixed, coefficients are inferred by one forward +sweep from \(t=2\) to \(T\). At each time bin: +\begin{enumerate} + \item Build \(Z_t[m,:] = A\,Y_{t-1}+B_m\in\mathbb{R}^{M\times N}\). + \item Run \texttt{pgd\_iter} projected-gradient iterations: + \[ + \eta_t = c_t^\top Z_t,\quad + \tilde{\eta}_t = \min(\eta_t,\eta_{\mathrm{clip}}),\quad + \mu_t = \exp(\tilde{\eta}_t)\Delta t, + \] + \[ + g_t = Z_t\left(\mu_t-Y_t\right) + 2\lambda_2(c_t-c_{t-1}), + \qquad + c_t \leftarrow \Pi_{\Delta^{M-1}}(c_t-\alpha_E g_t). + \] + \item Accept \(c_t\), continue to next time bin. +\end{enumerate} +Projection uses the sort-based simplex projection (Duchi et al., \(O(M\log M)\)). +Bin \(t=0\) is not updated; the sweep starts at \(t=1\) (first lag pair +\((Y_0,Y_1)\)), matching the objective sum over \(t=2,\dots,T\) in +1-based notation. +In the distributed setting, each rank updates a disjoint time shard of +\(c_{1:T}\); shard results are averaged where multiple ranks overlap at +shard boundaries, then broadcast so all ranks share the same \(c\). + +After the forward sweep completes, the soft coefficients \(c_{1:T}\) are decoded into a hard state sequence via Viterbi (Section~\ref{sec:decode}). The M-step then receives one-hot encoded state assignments rather than the soft simplex probabilities: +\begin{equation} + \hat{S}_{1:T} = \mathrm{Viterbi}(c_{1:T}), \qquad + c_t^{\mathrm{(M\text{-}step)}} = e_{\hat{S}_{t-1}} + \quad\hbox{for lag pair }(Y_{t-1},Y_t), +\end{equation} +where \(e_m\) is the \(m\)-th standard basis vector. In the current ordinary +EM implementation the one-hot row supplied to the M-step for the pair +\((Y_{t-1},Y_t)\) is the previous-bin decoded state \(\hat S_{t-1}\), matching +the saved metadata \texttt{mstep\_state\_mode=viterbi\_onehot\_prevbin}. +This ensures the M-step fits each state's parameters from cleanly separated +data rather than from soft mixtures. + +\subsubsection{M-step: Dictionary Update via Adam} + +Holding \(c_t\) fixed (as one-hot vectors), update \(A,\mathbf{B}\) for \(K_M\) epochs using +Adam over shuffled mini-batches of \((Y_{t-1},Y_t,c_t)\). + +Per mini-batch: +\begin{equation} + \ell_{\mathrm{M}} = + \frac{1}{|\mathcal{B}|} + \sum_{t\in\mathcal{B}}\sum_{i=1}^{N} + \left[-Y_{t,i}\log(\mu_{t,i}+\varepsilon)+\mu_{t,i}\right] + + \lambda_3\,\overline{|A_{\mathrm{off}}|}, +\end{equation} +with \(\varepsilon=10^{-8}\). + +Then, subject to delay scheduling (operations are activated only after a configurable number of EM iterations have passed): +\begin{enumerate} + \item Off-diagonal soft-thresholding (activated after \(K_{\mathrm{prune}}\) EM iterations): + \[ + A_{ij}\leftarrow \operatorname{sign}(A_{ij}) + \left(|A_{ij}|-\alpha_M\lambda_3\right)_+,\quad i\neq j. + \] + \item Periodic spectral-radius projection (activated after \(K_{\rho}\) EM iterations, applied every \(R\) batches). On each application, in the multi-GPU setting, \(A\) and \(B\) are first averaged across ranks, then the spectral constraint is enforced, and the result is broadcast: + \[ + A \leftarrow A\cdot \min\!\left(1,\frac{\rho_{\max}}{\rho(A)}\right). + \] +\end{enumerate} + +Learning rate follows linear decay from \(\alpha_M^{(0)}\) to +\(\alpha_M^{(0)}f_{\mathrm{end}}\). The decay is activated only after \(K_{\mathrm{lr}}\) EM iterations have passed; the total number of decay steps spans the remaining M-step epochs. + +% --------------------------------------------------------------- +\subsection{State Decoding} +\label{sec:decode} + +Viterbi decoding is applied both within the EM loop (after each E-step, to produce one-hot assignments for the M-step) and after the final EM iteration (to produce the output state sequence). The decoder uses a homogeneous transition prior with +\(\tau_{\mathrm{dwell}}=\texttt{decode\_dwell\_sec}\): +\begin{equation} + p_{\mathrm{stay}}=\exp\!\left(-\frac{\Delta t}{\tau_{\mathrm{dwell}}}\right), + \qquad + p_{\mathrm{switch}}=\frac{1-p_{\mathrm{stay}}}{M-1}. +\end{equation} +The decoded sequence is +\begin{equation} + \hat{S}_{1:T}=\arg\max_{\{s_t\}} + \sum_t\left[\log c_{t,s_t}+\log p(s_t\mid s_{t-1})\right]. +\end{equation} +Per-bin confidence is +\[ +\mathrm{CL}_t = 1-\max_m c_{t,m}. +\] + +% --------------------------------------------------------------- +\section{Input and Output File Pattern} +\label{sec:io} +% --------------------------------------------------------------- + +The ordinary EM trainer is \texttt{prism\_EM\_train3c.py}. It is self-contained +through \texttt{PrismEM\_Workhorse3c.py}; no code from older +\texttt{ver3} trainers is required. + +\paragraph{Input.} +The required input spike file is +\[ +\texttt{/spikesData/.spikes.npz}. +\] +The file must contain \texttt{spikes} with shape \((T,N)\) and +\texttt{single\_rates} with shape \((N,)\). Metadata must provide +\texttt{time\_step\_sec} and \texttt{poisson\_eta\_clip}. Synthetic files +from \texttt{gen\_nonStationarySpikes3c.py} also carry provenance metadata, +which is copied into the fit output. + +For real experimental recordings, the preprocessing program is +\texttt{prep\_bioexp3c.py}. Its output arrays are compatible with the EM +trainer, but the output directory must be chosen to match the EM file +layout. Run preprocessing with +\[ +\texttt{--dataPath /spikesData} +\] +and choose \texttt{--shortName } if an explicit EM data name is +desired. The resulting files are +\[ +\texttt{/spikesData/.spikes.npz}, +\qquad +\texttt{/spikesData/.bioExp.npz}. +\] +The \texttt{.spikes.npz} file is the direct EM input. The companion +\texttt{.bioExp.npz} file is not read during training, but it is expected by +the biological-data evaluation and plotting path to sit next to the spike +file with the same stem. In the spike metadata, +\texttt{provenance.experiment\_name} must therefore match +\texttt{}. + +\paragraph{Output.} +Rank~0 creates \texttt{/prismFit/} if needed and writes +\[ +\texttt{/prismFit/.prismEM.npz}. +\] +If \texttt{--fitName} is omitted, the program generates +\[ +\texttt{\_}. +\] +The metadata provenance field \texttt{EMtrain\_file} is set to the same +output stem. + +\paragraph{Canonical commands.} +\begin{verbatim} +torchrun --standalone --nnodes=1 --nproc_per_node=4 \ + ./prism_EM_train3c.py \ + --basePath $basePath \ + --dataName \ + --num_states \ + --time_range_sec 0 80 + +./prism_EM_train3c.py \ + --basePath $basePath \ + --dataName \ + --num_states \ + --fitName +\end{verbatim} + +The first form uses distributed GPU training when launched under +\texttt{torchrun}; the second form runs on a single process, using CUDA if a +GPU is visible and CPU otherwise. + +% --------------------------------------------------------------- +\section{Default Hyperparameters and Runtime Settings} +\label{sec:defaults} +% --------------------------------------------------------------- + +\begin{table}[h] +\centering +\begin{tabular}{lll} +\toprule +Group & Parameter (\texttt{CLI}) & Default \\ +\midrule +EM structure & \texttt{--num\_em\_iters} & 20 \\ + & \texttt{--m\_epochs} & 2 \\ +E-step & \texttt{--pgd\_iter} & 5 \\ + & \texttt{--lr\_estep} & 0.03 \\ + & \texttt{--lambda2} & 2.0 \\ +M-step & \texttt{--lr\_mstep} & \(3 \times 10^{-3}\) \\ + & \texttt{--lr\_end\_factor} & 0.1 \\ + & \texttt{--lambda3} & 0.02 \\ + & \texttt{--rho\_max} & 0.95 \\ + & \texttt{--prescale\_m\_step\_4\_ArhoMax} & 120 (batches) \\ + & \texttt{--delay\_em\_iter\_4\_ArhoMax} & [3] \\ + & \texttt{--delay\_em\_iter\_4\_lrDecay} & [1] \\ + & \texttt{--delay\_em\_iter\_4\_Aprune} & 1 \\ + & \texttt{--batch\_size} & 2048 \\ +Data window & \texttt{--time\_range\_sec} & [0.0, 60.0] s \\ +Decoding & \texttt{--decode\_dwell\_sec} & 0.3 s \\ +Init & \texttt{--init\_A} & data (or rand)\\ + & \texttt{--init\_B} & data (or rand) \\ + & \texttt{--init\_states} & data (or rand) \\ +\bottomrule +\end{tabular} +\caption{Current defaults in \texttt{PrismEM\_Workhorse3c.py}, as used by +\texttt{prism\_EM\_train3c.py}. The two delay arguments may be given as one +or two integers: start iteration, and optionally target iteration. If the +target is omitted, the code uses \texttt{num\_em\_iters}.} +\end{table} + +\paragraph{Required input.} +\texttt{--num\_states} (\texttt{-M}) is required (no default). +\texttt{--dataName} and \texttt{--basePath} locate input spikes under +\texttt{/spikesData/}; fits are written to +\texttt{/prismFit/.prismEM.npz}. + +\paragraph{Dataset-provided values.} +\(\Delta t\) (\texttt{time\_step\_sec}) and \(\eta_{\mathrm{clip}}\) +(\texttt{poisson\_eta\_clip}) are loaded from spike-file metadata. + +% --------------------------------------------------------------- +\section{Initialization of Model} +\label{sec:init} +% --------------------------------------------------------------- + +Initialization is controlled by: +\[ +\texttt{--init\_states},\quad \texttt{--init\_B},\quad \texttt{--init\_A} +\in \{\texttt{data},\texttt{rand}\}, +\] +with default \texttt{data} for all three. + +\subsection{State-Coefficient Initialization (\(C\equiv\{c_t\}\))} + +\paragraph{random mode:} +\[ +c_t = \frac{1}{M}\mathbf{1},\quad S_t=-1 \;\; \forall t. +\] +So coefficients start uniform over states. + +\paragraph{data mode:} +\begin{enumerate} + \item Aggregate spikes into coarse blocks of length + \(\lceil\tau_{\mathrm{dwell}}/\Delta t\rceil\) bins, where + \(\tau_{\mathrm{dwell}}=\texttt{decode\_dwell\_sec}\) (default 0.3\,s). + \item For each block, compute a scalar mean population firing rate. + \item Find \(M-1\) thresholds minimizing within-cluster variance + (dynamic-programming segmentation on unique rate values). + \item Assign each block to a state by thresholds. + \item Convert block states to soft probabilities: + dominant state gets \(0.8\), remainder \(0.2\) split across others, + with transition smoothing (\(2/3\) vs \(1/3\) at boundaries). + \item Map block-level \(c\) and \(S\) back to original bin resolution. +\end{enumerate} + +\subsection{Bias Initialization (\(B\))} + +\paragraph{random mode:} +No external seed is applied; model keeps random Gaussian initialization from +\texttt{torch.randn(...)*0.1}. + +\paragraph{data mode:} +For each neuron \(n\): +\[ +r_n = \frac{1}{T}\sum_t Y_{t,n}\,/\,\Delta t,\qquad +b_n = \log(\max(r_n,10^{-6})). +\] +The same \(b\)-vector is assigned to every state row of \(B\). + +\subsection{Connectivity Initialization (\(A\))} + +\paragraph{random mode:} +No external seed is applied; model keeps random Gaussian initialization from +\texttt{torch.randn(...)*0.1}. + +\paragraph{data mode (covariance/OLS):} +Using first \(T_{\max}=50{,}000\) bins (or fewer if shorter): +\[ +A_{\mathrm{ols}} += +\left(\sum_t Y_tY_{t-1}^\top\right) +\left(\sum_t Y_{t-1}Y_{t-1}^\top\right)^{\dagger}. +\] +Then all elements of \(A_{\mathrm{ols}}\) are scaled by a factor of 3: +\[ +A \leftarrow 3\,A_{\mathrm{ols}}. +\] +No spectral-radius rescaling is applied at initialization; the spectral constraint is enforced later during training via the delayed projection mechanism. + +Diagnostics (condition number, spectral radius, one-step \(R^2\), Frobenius norm, +bins used) are computed on the pre-scaled OLS estimate and saved in metadata +(\texttt{init\_A} sub-dictionary of the output metadata). + +% --------------------------------------------------------------- +\section{Output File and Evaluation} +\label{sec:output} +% --------------------------------------------------------------- + +After training, rank~0 writes \texttt{/prismFit/.prismEM.npz}. +If \texttt{--fitName} is omitted, the name is +\texttt{\_}. +Key arrays (shapes for \(T\) time bins, \(N\) neurons, \(M\) states): +\begin{itemize} + \item \texttt{A\_init}, \texttt{A\_hat}: \((N,N)\); + \item \texttt{B\_init}: \((N,)\) --- first row of the initialized \(B\) matrix + (common across states in \texttt{data} mode); + \item \texttt{B\_hat}: \((M,N)\) --- full per-state bias matrix; + \item \texttt{c\_init}, \texttt{c\_hat}: \((T,M)\); + \item \texttt{S\_init}, \texttt{S\_hat}: \((T,)\) integer state labels + (\texttt{S\_init} is \(-1\) in \texttt{rand} init mode); + \item \texttt{S\_hat\_CL}: \((T,)\) per-bin confidence + \(\mathrm{CL}_t = 1-\max_m c_{t,m}\); + \item \texttt{e\_nll\_em}, \texttt{m\_loss\_epoch}, \texttt{m\_nll\_epoch}, + \texttt{m\_l1\_epoch}, \texttt{rho\_epoch}, + \texttt{rho\_correction\_strength\_epoch}, + \texttt{nz\_edges\_epoch}, \texttt{learning\_rates}: training history; + \item \texttt{freq\_h1d}: coarse firing-rate trace used by + \texttt{init\_states=data}; + \item \texttt{single\_rates}: \((N,)\), copied from the input spike file; + \item \texttt{neuron\_type}, \texttt{neuron\_Sedge}, \texttt{A\_prune}: + source-column Dale-style postprocessing derived from fitted + \texttt{A\_hat}. Here \texttt{neuron\_Sedge} is the off-diagonal + column sum for each source neuron, \texttt{neuron\_type} is inferred + from that source-column sum, and \texttt{A\_prune} preserves + positive outgoing columns for excitatory sources and negative + outgoing columns for inhibitory sources. +\end{itemize} +Metadata records all CLI hyperparameters, time-window bins, initialization +dictionaries, and \texttt{mstep\_state\_mode=viterbi\_onehot\_prevbin}. +Suggested evaluation: +\begin{verbatim} +./prism_EM_eval3c.py --basePath $basePath --dataName -p a e f g +\end{verbatim} + +\end{document} diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/fitFDR_states_ver3c.pdf b/causal_net/nonStation_ver3c_states_EM_FDR/docs/fitFDR_states_ver3c.pdf new file mode 100644 index 00000000..190c89fc Binary files /dev/null and b/causal_net/nonStation_ver3c_states_EM_FDR/docs/fitFDR_states_ver3c.pdf differ diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/fitFDR_states_ver3c.tex b/causal_net/nonStation_ver3c_states_EM_FDR/docs/fitFDR_states_ver3c.tex new file mode 100644 index 00000000..af4534a0 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/fitFDR_states_ver3c.tex @@ -0,0 +1,1537 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} +\usepackage{tikz} + +\geometry{margin=0.8in} + +\title{EM-FDR-Bagging: Robust Estimation of the Connectivity Matrix\\ +from Multi-State Spiking Data\\ +\large A bagged, scramble-calibrated wrapper around PRISM-EM} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +This document describes \textbf{EM-FDR-Bagging}, a statistical wrapper around +the PRISM-EM fitter that aims to recover the shared connectivity matrix +\(A\in\mathbb{R}^{N\times N}\) of a non-stationary, multi-state Poisson GLM +\emph{robustly}, i.e.\ with explicit control of false edges and with +per-edge uncertainty estimates, on real data where no ground truth exists. +The method combines three ideas. First, one ordinary PRISM-EM run is made on +the requested time range to discover a reference hard Viterbi state sequence. +Second, \emph{reference-labeled pair bagging}: every FDR bag draws the same +fraction of lag pairs \((\hat S_t,Y_{t-1},Y_t)\) from that reference window +and re-estimates only \(A\) and \(\mathbf B\) with the state labels locked. +No bag performs a new E-step. Third, a \emph{per-neuron time-shuffle null} +and \emph{per-source pooled significance with cross-bag stability selection}: +inside each bag, independently time-scrambled neurons are refit by the same +locked M-step, edges are tested against a source-column pooled null, and retained only +if they recur reliably across bags. The procedure yields, for every +surviving edge, a selection frequency, a bagged point estimate, a bootstrap +SD, and a null-referenced effect size. The resulting edge masks are expected +to be reliable, but the corresponding fitted magnitudes should be interpreted +as conservative: the locked M-step still uses \(L_1\) shrinkage and may also +be scaled by spectral-radius enforcement. Therefore these outputs are best +treated as Stage (b) support estimates, and a follow-up de-biased refit is +needed for final magnitudes; see the companion writeup +\texttt{deBiasFit\_states\_ver3c.tex} for details. We treat PRISM-EM as a +black box exposed through two entry points and specify everything needed for +an implementer who has the PRISM-EM code in hand. +\end{abstract} + +\tableofcontents +\newpage + +% =============================================================== +\section{Scope and Inputs} +% =============================================================== + +The full analysis pipeline is organized into three stages. \textbf{Stage (a)} +is the ordinary PRISM-EM fit on the chosen spike-train window: it estimates a +shared connectivity matrix, state-dependent biases, and a reference decoded +state sequence; this core fitter is specified in +\texttt{fitEM\_states\_ver3c.tex}. \textbf{Stage (b)} is the +EM-FDR-Bagging procedure described in this document: it reuses the Stage (a) +reference state labels, performs bagged locked M-step refits, calibrates +edges against time-scrambled null fits, and returns a statistically vetted +support mask with conservative edge magnitudes. \textbf{Stage (c)} is the +follow-up de-biased refit specified in +\texttt{deBiasFit\_states\_ver3c.tex}: it freezes the Stage (b) support and +Dale-sign decisions, removes the \(L_1\) sparsity penalty, and refits final +edge magnitudes and state biases under the remaining biological and +stability constraints. + +The input is a single spike-train recording of \(N\) neurons over \(T\) +discrete time bins of width \(\Delta t\), in the format consumed by +PRISM-EM (\texttt{spikesData/.spikes.npz}). The target regime, +which sizes every numerical choice below, is: +\(N=50\)–\(500\) neurons; recordings of \(10\)–\(60\) minutes at +\(\Delta t=10\,\text{ms}\), so \(T\) ranges from \(6\times10^{4}\) to +\(3.6\times10^{5}\) bins; per-neuron firing rates of \(1\)–\(5\,\text{Hz}\); +state-switching rates of \(5\)–\(15\) transitions per minute; and a small +number of latent states, typically \(M=2\). Crucially, the method is +designed for application to \emph{real} data: it never references ground +truth, and every quantity it reports is computed from the data alone. + +\begin{figure}[h] +\centering +\begin{tikzpicture}[x=0.11cm,y=0.55cm] + % axes and neuron labels + \draw[->, thick] (0,0) -- (72,0) node[right] {time}; + \foreach \n/\lab in {1/{neuron 1},2/{neuron 2},3/{neuron 3},4/{neuron 4},5/{neuron 5}} { + \draw[gray!35] (0,\n) -- (70,\n); + \node[left] at (0,\n) {\scriptsize \lab}; + } + \foreach \x/\txt in {0/{0},20/{20 s},40/{40 s},60/{60 s}} { + \draw[gray!60] (\x,-0.08) -- (\x,6.85); + \node[below] at (\x,0) {\scriptsize \txt}; + } + + % hidden two-state timeline + \node[left] at (0,6.55) {\scriptsize hidden state}; + \node[left, green!45!black] at (0,6.25) {\scriptsize 0}; + \node[left, orange!85!black] at (0,6.85) {\scriptsize 1}; + \draw[gray!35] (0,6.25) -- (70,6.25); + \draw[gray!35] (0,6.85) -- (70,6.85); + \draw[green!60!black, very thick] (0,6.25) -- (11,6.25); + \draw[orange!85!black, very thick] (11,6.85) -- (23,6.85); + \draw[green!60!black, very thick] (23,6.25) -- (36,6.25); + \draw[orange!85!black, very thick] (36,6.85) -- (49,6.85); + \draw[green!60!black, very thick] (49,6.25) -- (70,6.25); + \foreach \x/\ya/\yb in {11/6.25/6.85,23/6.85/6.25,36/6.25/6.85,49/6.85/6.25} { + \draw[gray!65, dashed] (\x,0.3) -- (\x,6.95); + \draw[black, thick] (\x,\ya) -- (\x,\yb); + } + \node[align=center, anchor=west] at (47,7.45) + {\scriptsize \(M=2\) hidden state}; + + % spike ticks, deliberately irregular and sparse + \foreach \x in {3,9,18,22,38,44,61,67} \draw[black, line width=0.7pt] (\x,0.82) -- (\x,1.18); + \foreach \x in {5,6,25,31,32,50,57,58,69} \draw[black, line width=0.7pt] (\x,1.82) -- (\x,2.18); + \foreach \x in {11,17,28,41,42,46,63} \draw[black, line width=0.7pt] (\x,2.82) -- (\x,3.18); + \foreach \x in {2,15,16,34,35,36,54,65} \draw[black, line width=0.7pt] (\x,3.82) -- (\x,4.18); + \foreach \x in {8,19,23,39,48,49,53,60,66} \draw[black, line width=0.7pt] (\x,4.82) -- (\x,5.18); + + % analysis window and lag-pair hint + \draw[blue!70, very thick, rounded corners] (14,0.45) rectangle (56,5.45); + \node[blue!70, fill=white, inner sep=1pt] at (35,5.72) {\scriptsize selected time range}; + \draw[->, orange!90!black, thick] (27,-0.55) -- (30,-0.55); + \node[orange!90!black, below] at (28.5,-0.62) + {\scriptsize lag pair \((Y_{t-1},Y_t)\)}; +\end{tikzpicture} +\caption{Cartoon of the spike-train input. Rows are neurons, vertical ticks +are spike events in discrete time bins, and the colored track above the +raster sketches a hidden \(M=2\) state process switching every few seconds. +The blue box marks the reference time window supplied to the initial +PRISM-EM fit. FDR bags later sample lagged pairs +\((\hat S_t,Y_{t-1},Y_t)\) from this same window.} +\label{fig:input-spike-cartoon} +\end{figure} + +The deliverable is not a final unregularized point estimate of \(A\) but a +vetted edge list and support mask. For each off-diagonal entry \((i,j)\) +that survives the procedure we report a point estimate, a statistical +uncertainty, a stability (selection) frequency, and an effect size relative +to the null. These values are useful for deciding which edges exist, but the +edge magnitudes are likely underestimated because of the same \(L_1\) +shrinkage and spectral-radius projection used during fitting. Final +magnitude recovery is handled by the Stage (c) de-biased refit described in +\texttt{deBiasFit\_states\_ver3c.tex}. Self-couplings \(A_{ii}\) are +excluded from testing throughout. + +% =============================================================== +\section{The PRISM-EM Workhorse and Its Two Entry Points} +\label{sec:workhorse} +% =============================================================== + +EM-FDR-Bagging treats PRISM-EM as a callable. We summarize only what the +wrapper relies on; the full model, objective, and defaults are documented +separately in the PRISM-EM specification. + +\subsection{What PRISM-EM fits} + +PRISM-EM fits a non-stationary Poisson GLM in which all \(M\) latent states +share one connectivity matrix \(A\) but carry state-specific bias vectors +\(B_m\). With \(c_t\) the simplex coefficients at bin \(t\) and +\(B_t=\sum_m c_{mt}B_m\), the rate is +\(\mu_t=\exp\!\big(\min(A Y_{t-1}+B_t,\,\eta_{\text{clip}})\big)\,\Delta t\) +and \(Y_t\sim\text{Poisson}(\mu_t)\). Training is block coordinate descent: +an E-step infers \(c_{1:T}\) by simplex-projected gradient descent and +Viterbi-decodes them to a hard one-hot state sequence +\(\hat S_{1:T}\); an M-step then updates \(A,\mathbf{B}\) by Adam under the +frozen one-hot assignment, with off-diagonal \(L_1\) shrinkage (strength +\(\lambda_3\)), a spectral-radius projection +\(\rho(A)\le\rho_{\max}\), and a learning-rate decay. The +shrinkage, projection, and decay each switch on only after a configurable +number of EM iterations. The fitted \(A\) is therefore the output of a +specific, deterministic-in-configuration pipeline of regression plus +shrinkage plus spectral projection --- a fact the null construction below +depends on. + +\begin{figure}[h] +\centering +\begin{tikzpicture}[x=1cm,y=1cm] + % shared A block + \draw[rounded corners, very thick, blue!70] (0,0.4) rectangle (2.6,2.6); + \node[blue!70] at (1.3,2.85) {\small shared connectivity}; + \node at (1.3,1.5) {\Large \(A\)}; + \node[align=center] at (1.3,0.75) {\scriptsize same for\\[-1mm]\scriptsize both states}; + + % two state bias blocks + \draw[rounded corners, thick, green!60!black, fill=green!8] (4.0,1.9) rectangle (6.2,3.1); + \node[green!45!black] at (5.1,3.35) {\small state 0}; + \node at (5.1,2.5) {\Large \(B_0\)}; + + \draw[rounded corners, thick, orange!85!black, fill=orange!10] (4.0,-0.1) rectangle (6.2,1.1); + \node[orange!80!black] at (5.1,1.35) {\small state 1}; + \node at (5.1,0.5) {\Large \(B_1\)}; + + % state timeline + \draw[->, thick] (0,-1.05) -- (7.2,-1.05) node[right] {\small time}; + \foreach \x in {0,1,2,3,4,5,6} { + \draw[gray!45] (\x,-1.18) -- (\x,-0.92); + } + \draw[green!60!black, very thick] (0,-0.75) -- (1.5,-0.75); + \draw[orange!85!black, very thick] (1.5,-1.35) -- (3.7,-1.35); + \draw[green!60!black, very thick] (3.7,-0.75) -- (5.2,-0.75); + \draw[orange!85!black, very thick] (5.2,-1.35) -- (6.7,-1.35); + \node[green!45!black] at (0.75,-0.48) {\scriptsize \(\hat S_t=0\)}; + \node[orange!80!black] at (2.6,-1.62) {\scriptsize \(\hat S_t=1\)}; + + % model arrows + \draw[->, thick] (2.75,1.5) -- (3.75,2.5); + \draw[->, thick] (2.75,1.5) -- (3.75,0.5); + \draw[->, thick, gray!80] (6.35,2.5) -- (7.45,1.78); + \draw[->, thick, gray!80] (6.35,0.5) -- (7.45,1.22); + \draw[rounded corners, thick, gray!80] (7.45,0.85) rectangle (9.75,2.15); + \node[align=center] at (8.6,1.5) + {\small rate for\\[-1mm]\small \(Y_t\mid Y_{t-1}\)}; + + \node[align=center] at (3.6,-1.9) + {\scriptsize hard Viterbi state selects which bias vector is active}; +\end{tikzpicture} +\[ + \eta_t = A Y_{t-1}+B_{\hat S_t}, \qquad + \tilde\eta_t=\min(\eta_t,\eta_{\text{clip}}), \qquad + \lambda_t=\exp(\tilde\eta_t)\Delta t, \qquad + Y_t\sim\mathrm{Poisson}(\lambda_t). +\] +\caption{Two-state PRISM-EM cartoon. This is the EM fitting model, not the +synthetic data generator. The lagged spike vector \(Y_{t-1}\) always enters +through the same shared connectivity matrix \(A\). The decoded state +\(\hat S_t\in\{0,1\}\) selects the active bias vector, \(B_0\) or \(B_1\), +for the conditional rate of \(Y_t\). FDR bagging later locks this state +sequence and refits only \(A\) and the two bias vectors.} +\label{fig:two-state-prism-cartoon} +\end{figure} + +\subsection{Entry point 1: full fit} +\label{sec:ep1} + +\begin{quote} +\(\textsf{full\_fit}(\text{spikes},\,M,\,\theta_{\text{EM}}) +\;\longrightarrow\;(\hat A,\ \hat S_{1:T},\ \hat{\mathbf B},\ \text{history})\) +\end{quote} + +This is the ordinary PRISM-EM run: full E/M alternation on the supplied +spikes, returning the fitted connectivity \(\hat A\), the decoded one-hot +state sequence \(\hat S_{1:T}\), the per-state biases \(\hat{\mathbf B}\), +and the training history. Here \(\theta_{\text{EM}}\) collects all PRISM-EM +hyperparameters (number of EM iterations, M-step epochs, learning rates, +\(\lambda_2,\lambda_3,\rho_{\max}\), the delay schedules, etc.). + +\subsection{Entry point 2: M-step-only refit with locked state sequence} +\label{sec:ep2} + +\begin{quote} +\(\textsf{mstep\_locked}(\,Y_{\text{prev}},\,Y_{\text{curr}},\,\hat S^{\text{lock}},\, +A^{\text{init}},\,\hat{\mathbf B}^{\text{init}},\,\theta_{\text{M}}\,) +\;\longrightarrow\;(A,\mathbf B)\) +\end{quote} + +This entry point runs \emph{only} the M-step. It accepts spike pairs +\((Y_{\text{prev}},Y_{\text{curr}})\), an externally supplied and +\emph{frozen} state sequence \(\hat S^{\text{lock}}\) (used directly as the +one-hot assignment, with no E-step and no Viterbi re-decoding), initial +parameters \((A^{\text{init}},\hat{\mathbf B}^{\text{init}})\), and returns +the re-estimated connectivity \(A\) and the refit state biases +\(\mathbf B\). In the null workflow the initial \(A\) and \(\mathbf B\) are +not copied from the real locked fit. Instead, for every scramble they +are initialized freshly from the scrambled data, or randomly, according to +\texttt{null\_A\_init} and \texttt{null\_B\_init}. Re-fitting +\(\mathbf B\) in the scrambled context lets the null absorb state-specific +baseline effects, so the test of \(A\) is not made artificially harsh by a +bias mismatch; fresh \(A\) initialization avoids seeding the null from the +real edge pattern. + +\paragraph{Locked-state time indexing.} For each lag-1 training pair +\((Y_{t-1},Y_t)\), the locked state used by the M-step is +\(\hat S^{\text{lock}}_{t}\), i.e.\ the state attached to the destination +or current bin whose conditional rate is being modeled. Equivalently, if the +spike pairs have length \(T-1\), the state covariate rows are formed from +\(\hat S^{\text{lock}}_{1:T-1}\), converted to one-hot vectors. This +convention is part of the FDR bagging definition and must be used +identically for every real and scrambled locked refit. + +The locked M-step is configured to reproduce the real +pipeline exactly: the same \(L_1\) strength \(\lambda_3\), the same spectral +projection \(\rho_{\max}\), and the same learning-rate decay as the +corresponding \textsf{full\_fit}. Because there is no E-step here, the EM +iteration counters that gate PRISM-EM's delayed schedules are not +meaningful; instead all schedule-gated operations (shrinkage, spectral +projection, decay) are \textbf{active from the first M-step epoch}. The +rationale is that the null estimate must traverse the same shrinkage and +constraint geometry as the converged real estimate, which lives in the +late-EM regime where all schedules have already fired. + +\paragraph{Implementation note.} In the current \texttt{ver3c} code these +two entry points live in \texttt{PrismEM\_Workhorse3c.py}: +\texttt{run\_full\_fit} and \texttt{run\_locked\_mstep}. The locked M-step +consumes externally supplied spike pairs, an externally supplied +\(\hat S^{\text{lock}}\), and explicit initial \(A\) and \(\mathbf B\). Its +lag-1 state alignment for the new FDR bagging path must be +\(\hat S^{\text{lock}}_{t}\) for pair \((Y_{t-1},Y_t)\). If existing helper +names still refer to ``previous states'', their use must be audited while +implementing this plan. + +% =============================================================== +\section{Method Overview} +% =============================================================== + +EM-FDR-Bagging has two stages, and the document is organized to match the +intended software structure. + +\textbf{Reference EM fit} (performed before Stage~(a)). Run +\texttt{prism\_EM\_train3c.py} once on the desired \texttt{--time\_range\_sec}. +This discovers the reference state sequence +\(\hat S^{\text{ref}}_{0:T_{\rm ref}-1}\) +and writes a tagged reference fit, for example +\texttt{prismFit/\_jXXXX.prismEM.npz}. Its EM hyperparameters +control state discovery; once the reference state labels are good enough, the +final precision of \(A\) and \(\mathbf B\) is delegated to the locked bag +fits. + +\textbf{Stage~(a), per-bag fitting} (Sections~\ref{sec:bagging}--\ref{sec:bagout}). +For a single bag: load the reference EM fit and original spikes, draw a +fixed-size random subset of reference-labeled lag pairs +\((\hat S^{\text{ref}}_t,Y_{t-1},Y_t)\), re-estimate only +\((\hat A^{(b)},\hat{\mathbf B}^{(b)})\) by \textsf{mstep\_locked}, then run +\(P\) independently time-scrambled locked-M-step null refits on the same +selected pair indices. Write everything to disk. This stage is one +self-contained program; it is run \(N_{\text{bag}}\) times, once per bag, as +independent jobs. The bag trainer has its own locked-M-step training +hyperparameters. Their defaults are read from the reference EM metadata for +continuity, but they may be overridden on the Stage~(a) command line and the +effective values are saved in every bag file. + +\textbf{Stage~(b), aggregation and pruning} (Section~\ref{sec:aggregate}). +Read all per-bag outputs from a directory, build per-source pooled null +distributions, perform per-bag selection and cross-bag stability selection, +and emit the final vetted edge list with uncertainties. + +The symbols introduced for the wrapper, kept distinct from PRISM-EM's +\(N,M,T\), are collected in Table~\ref{tab:symbols}. + +\begin{table}[h] +\centering +\begin{tabular}{lll} +\toprule +Symbol & Meaning & Nominal value \\ +\midrule +\(T_{\rm ref}\) & bins in the reference EM time window & from reference fit \\ +\(N_{\rm pair}\) & lag pairs in the reference window, \(T_{\rm ref}-1\) & from reference fit \\ +\(f_{\rm bag}\) & pair fraction per bag, \texttt{--bag\_frac} & \(0.8\) \\ +\(K_{\rm bag}\) & pairs sampled per bag, \(\lfloor f_{\rm bag}N_{\rm pair}\rfloor\) & from \texttt{--bag\_frac} \\ +\(N_{\text{bag}}\) & number of bags & \(5\)--\(10\) \\ +\(P\) & scrambles per bag & \(4\)--\(8\) \\ +\(\delta_{\text{roll}}^{\min}\) & roll exclusion-zone half-width (s) & \(2\) \\ +\(q\) & per-bag per-source selection quantile level & e.g.\ \(0.99\) \\ +\(\pi_{\text{thr}}\) & cross-bag selection-frequency threshold & \(0.6\)--\(0.8\) \\ +\(M,N,T\) & states, neurons, bins (PRISM-EM) & \(\approx3\); \(50\)--\(500\); \(6\!\times\!10^4\)--\(3.6\!\times\!10^5\) \\ +\bottomrule +\end{tabular} +\caption{Symbols introduced by EM-FDR-Bagging and their nominal values. +Every numeric choice in the text is written as ``symbol = value'' so it can +be adjusted.} +\label{tab:symbols} +\end{table} + +% =============================================================== +\section{Stage (a.1): Reference-Labeled Pair Bagging} +\label{sec:bagging} +% =============================================================== + +A reference EM fit is produced before any bag job is launched: +\[ +\texttt{prism\_EM\_train3c.py}\quad +\text{on}\quad \texttt{--time\_range\_sec}=[t_0,t_1]. +\] +The reference output fixes the training window, the fitted soft state +coefficients \(\hat c^{\text{ref}}\), and the hard Viterbi labels +\(\hat S^{\text{ref}}_{0:T_{\rm ref}-1}\). The bag driver does not accept a +new \texttt{--time\_range\_sec}; it inherits the reference fit's saved +window and therefore uses the same source data for simulations and +experiments. + +From the reference window we form a table of lag pairs +\[ +X_t=Y_{t-1},\qquad Z_t=Y_t,\qquad H_t=\hat S^{\text{ref}}_t, +\qquad t=1,\dots,T_{\rm ref}-1. +\] +Here \(X_t\) is the source/history spike vector, \(Z_t\) is the +destination/current spike vector, and \(H_t\) is the destination-bin state +label that controls the conditional rate for \(Z_t\). This table contains +\(N_{\rm pair}=T_{\rm ref}-1\) rows. + +\begin{figure}[h] +\centering +\begin{tikzpicture}[x=0.13cm,y=0.7cm] + % full reference table + \draw[rounded corners, thick, gray!70, fill=gray!8] (0,0.2) rectangle (74,2.0); + \node[above] at (37,2.1) {\small reference-labeled lag-pair table}; + \foreach \x in {2,6,...,72} { + \draw[gray!45] (\x,0.2) -- (\x,2.0); + } + \node[rotate=90] at (-2.6,1.1) {\scriptsize neurons}; + \node[below] at (37,0.18) {\scriptsize time}; + \node[align=center] at (37,1.1) + {\scriptsize \((H_t,X_t,Z_t)=(\hat S^{\text{ref}}_t,Y_{t-1},Y_t)\)}; + + % bag subsets as scattered row selections, sampled without replacement + \node[left] at (0,-1.0) {\scriptsize bag 0}; + \draw[gray!25, line width=1.2pt] (2,-1.0) -- (72,-1.0); + \node[blue!65, right] at (75,-1.0) {\scriptsize \(f_{\rm bag}=0.8\)}; + + \node[left] at (0,-2.0) {\scriptsize bag 1}; + \draw[gray!25, line width=1.2pt] (2,-2.0) -- (72,-2.0); + \node[green!55!black, right] at (75,-2.0) {\scriptsize \(f_{\rm bag}=0.8\)}; + + \node[left] at (0,-3.0) {\scriptsize bag 2}; + \draw[gray!25, line width=1.2pt] (2,-3.0) -- (72,-3.0); + \node[orange!85!black, right] at (75,-3.0) {\scriptsize \(f_{\rm bag}=0.8\)}; + \foreach \x in {2,6,10,14,18,22,30,34,38,42,46,50,54,58,62,66,70} { + \draw[blue!65, line width=1.0pt] (\x,-1.23) -- (\x,-0.77); + } + \draw[->, blue!45, line width=0.8pt, opacity=0.55] + (10,0.17) .. controls (10,-0.18) and (10,-0.55) .. (10,-0.86); + \draw[->, blue!45, line width=0.8pt, opacity=0.55] + (30,0.17) .. controls (30,-0.18) and (30,-0.55) .. (30,-0.86); + \draw[->, blue!45, line width=0.8pt, opacity=0.55] + (50,0.17) .. controls (50,-0.18) and (50,-0.55) .. (50,-0.86); + \foreach \x in {2,6,14,18,22,26,30,34,42,46,50,54,58,62,66,70,72} { + \draw[green!55!black, line width=1.0pt] (\x,-2.23) -- (\x,-1.77); + } + \foreach \x in {2,10,14,18,22,26,30,34,38,42,50,54,58,62,66,70,72} { + \draw[orange!85!black, line width=1.0pt] (\x,-3.23) -- (\x,-2.77); + } + \node[align=center] at (37,-3.85) + {\scriptsize each bag samples scattered rows without replacement\\[-1mm] + \scriptsize selected row indices are sorted, so time order is preserved}; +\end{tikzpicture} +\caption{Reference-labeled pair bagging. The initial EM fit defines one +time window and one table of lagged rows \((\hat S^{\text{ref}}_t,Y_{t-1},Y_t)\). +Each FDR bag samples a different random 80\% subset of rows without +replacement; the selected rows are sparse but kept in chronological order. +The same selected row indices are used for all neurons inside a bag, so +real cross-neuron synchrony is preserved.} +\label{fig:bagging-cartoon} +\end{figure} + +For bag \(b\), draw an index set \(I_b\subset\{1,\dots,N_{\rm pair}\}\) +without replacement, with +\[ +|I_b|=K_{\rm bag}=\lfloor f_{\rm bag}N_{\rm pair}\rfloor, +\qquad f_{\rm bag}=\texttt{--bag\_frac}. +\] +The same row index set is used for all neurons. This preserves the +within-time synchrony between neurons in the real bag fit while giving every +bag the same statistical precision. All bag indices are equivalent: +\texttt{bag\_idx=0} is no longer chronological or special; it is simply the +first random subset in the reproducible bag-indexed RNG stream. + +This design deliberately drops the previous block-concatenation machinery. +There are no block starts, no block boundaries, no minimum block-start +separation, and no bag-specific E-step. The only perturbation defining a +real bag is the random subset of already-labeled lag-pair rows. + +% =============================================================== +\section{Stage (a.2): Real Fit and the Circular-Shift Null} +\label{sec:null} +% =============================================================== + +\subsection{Real locked M-step on the selected pairs} + +For bag \(b\), run entry point~2 on the selected rows: +\[ +(\hat A^{(b)},\ \hat{\mathbf B}^{(b)}) +=\textsf{mstep\_locked}\!\big( +X_{I_b},\ Z_{I_b},\ H_{I_b},\ +A^{\text{init},(b)},\ \mathbf B^{\text{init},(b)},\ \theta_{\text{M}}\big). +\] +This \(\hat A^{(b)}\) is the bag's real connectivity estimate. The state +labels are the locked hard labels from the reference EM fit; there is no +E-step, no Viterbi decoding, and no re-estimation of \(\hat c_t\) inside a +bag. + +\subsection{The null: per-neuron independent circular shifts} + +The null hypothesis for an edge is ``no directed lag-1 coupling from neuron +\(j\) to neuron \(i\) beyond what the bias/state structure already explains.'' +We realize it by destroying cross-neuron temporal alignment while preserving +each neuron's own marginal statistics and the state timeline. Concretely, +for scramble \(p=1,\dots,P\) we draw an independent integer circular shift +\(\delta_n^{(b,p)}\) for each neuron \(n\) and roll that neuron's spike +train over the full reference window: +\[ +\tilde Y^{(b,p)}_{t,n} +=Y_{(t-\delta_n^{(b,p)})\bmod T_{\rm ref},\ n}. +\] + +\begin{figure}[h] +\centering +\begin{tikzpicture}[x=0.12cm,y=0.7cm] + % Grid lines (0s, 20s, 40s, 60s) crossing the timelines + \foreach \x in {0,20,40,60} { + \draw[gray!30, line width=0.5pt] (\x, 0.4) -- (\x, 4.0); + \node[below, black!70] at (\x, 0.4) {\scriptsize \x s}; + } + + % Horizontal lines for the neurons (custom lengths based on overflow spikes) + \draw[gray!35, line width=0.8pt] (0, 3.5) -- (82, 3.5); + \node[left, black!85] at (0, 3.5) {\scriptsize neuron 3}; + + \draw[gray!35, line width=0.8pt] (0, 2.2) -- (74, 2.2); + \node[left, black!85] at (0, 2.2) {\scriptsize neuron 2}; + + \draw[gray!35, line width=0.8pt] (0, 0.9) -- (68, 0.9); + \node[left, black!85] at (0, 0.9) {\scriptsize neuron 1}; + + % Vertical ticks at 0s for each neuron timeline + \foreach \y in {0.9, 2.2, 3.5} { + \draw[gray!45, line width=0.8pt] (0, \y-0.2) -- (0, \y+0.2); + } + + % Original spike positions (light grey ticks) + % neuron 3 (y=3.5) + \foreach \x in {8,16,29,35,48,57,63} { + \draw[gray!40, line width=0.8pt] (\x,3.3) -- (\x,3.7); + } + % neuron 2 (y=2.2) + \foreach \x in {3,18,24,31,43,52,59} { + \draw[gray!40, line width=0.8pt] (\x,2.0) -- (\x,2.4); + } + + % Shift arrows (blue) + % neuron 3 (y=3.5) shift of 18 units + \draw[->, blue, line width=0.6pt] (0, 3.75) -- (18, 3.75); + \node[above, blue] at (9, 3.75) {\scriptsize \(T_{\text{offset},3}\)}; + + % neuron 2 (y=2.2) shift of 11 units + \draw[->, blue, line width=0.6pt] (0, 2.45) -- (11, 2.45); + \node[above, blue] at (5.5, 2.45) {\scriptsize \(T_{\text{offset},2}\)}; + + % Ticks at the offset positions + \draw[gray!45, line width=0.8pt] (18, 3.3) -- (18, 3.7); + \draw[gray!45, line width=0.8pt] (11, 2.0) -- (11, 2.4); + + % Shifted spike positions (black ticks) + % neuron 3 (y=3.5) shifted by 18 units: + \foreach \x in {26,34,47,53,66,75,81} { + \draw[black, line width=1.3pt] (\x,3.3) -- (\x,3.7); + } + + % neuron 2 (y=2.2) shifted by 11 units: + \foreach \x in {14,29,35,42,54,63,70} { + \draw[black, line width=1.3pt] (\x,2.0) -- (\x,2.4); + } + + % neuron 1 (y=0.9) unshifted: + \foreach \x in {5,14,21,37,46,55,64} { + \draw[black, line width=1.3pt] (\x,0.7) -- (\x,1.1); + } + + % Dashed vertical boundary line at 68s + \draw[dashed, black, line width=0.8pt] (68, 0.4) -- (68, 4.0); + + % Wrap time axis around label (moved up and to the right) + \node[blue, align=left, anchor=west] at (71, 1.8) { + \scriptsize wrap time\\ + \scriptsize axis around + }; + + % A hand-drawn/wavy look wrap arrow + \draw[->, blue, thick] (80, 0.9) .. controls (86, 0.9) and (86, 1.9) .. (80, 1.9); + +\end{tikzpicture} +\caption{One circular-shift null scramble for one bag. For each neuron, original spike times (grey ticks) are shifted to the right by an independent offset \(T_{\text{offset},n}\) to produce the scrambled spike train (black ticks). Spikes that are shifted beyond the end of the window (dashed line) wrap around to the beginning.} +\label{fig:circular-shift-null-cartoon} +\end{figure} + +Because every neuron is shifted by a different random amount, the precise +lag-1 relationship that \(A\) captures is broken, while each neuron's firing +rate, burstiness, and count distribution are untouched (a circular roll is a +pure permutation of a neuron's own time axis). The locked reference state +labels \(H_t=\hat S^{\text{ref}}_t\) are \emph{not} rolled: they are held +fixed so that the null asks specifically about coupling, not about state +structure. + +\paragraph{Two-sided exclusion zone on each shift.} A shift near \(0\) leaves +the data almost unshifted and a shift near \(T_b\) wraps back to almost +unshifted; both leave the real lag structure largely intact and would +contaminate the null toward the alternative. We therefore restrict each +shift to +\[ +\delta_n^{(b,p)}\in +\big[\,\delta_{\text{roll}}^{\min,\text{bin}},\ \ +T_{\rm ref}-\delta_{\text{roll}}^{\min,\text{bin}}\,\big], +\qquad +\delta_{\text{roll}}^{\min,\text{bin}}=\delta_{\text{roll}}^{\min}/\Delta t, +\] +with \(\delta_{\text{roll}}^{\min}=2\,\text{s}\) (\(200\) bins). The +exclusion is two-sided by construction: both the small-shift and the +near-full-wrap small-effective-shift regions are forbidden. + +\paragraph{No pairwise separation among neurons.} The shifts of different +neurons need not be separated from one another. The null only requires that +each neuron's shift is individually non-trivial (the two-sided zone above) +and that shifts are drawn independently across neurons; a coincidence in +which two neurons happen to share the same shift is harmless to null +validity. This matters in practice: with a \(10\)-minute bag and a +\(2\,\text{s}\) separation there would be only a finite number of effective +``slots,'' and a bag may contain up to \(500\) neurons, so any attempt to enforce pairwise +distinct shifts would be infeasible by pigeonhole. We deliberately avoid +that requirement. The implemented draw routine simply samples independently +from the allowed two-sided interval for each neuron. + +\subsection{Null re-estimation through the locked M-step} + +For each scramble \(p\) we recompute connectivity with the state sequence +frozen, via entry point~2: +\[ +(A^{\text{null},(b,p)},\mathbf B^{\text{null},(b,p)}) +=\textsf{mstep\_locked}\!\big(\tilde X^{(b,p)}_{I_b},\ \tilde Z^{(b,p)}_{I_b},\ +H_{I_b},\ A^{\text{init},(b,p)},\ +\mathbf B^{\text{init},(b,p)},\ \theta_{\text{M}}\big). +\] +The scrambled pair arrays \(\tilde X^{(b,p)}_t=\tilde Y^{(b,p)}_{t-1}\) and +\(\tilde Z^{(b,p)}_t=\tilde Y^{(b,p)}_t\) are formed from the rolled +reference-window spikes, and the same selected row indices \(I_b\) are used +as in the real fit. The locked state attached to pair +\((\tilde Y^{(b,p)}_{t-1},\tilde Y^{(b,p)}_t)\) is +\(H_t=\hat S^{\text{ref}}_t\). Only the spike counts entering the null +M-step are circularly shifted. The null initialization is fresh for each +scramble: \(A^{\text{init},(b,p)}\) is either computed from the +scrambled spike train (\texttt{null\_A\_init=scrambled\_data}, the default) or drawn +randomly (\texttt{null\_A\_init=rand}), and +\(\mathbf B^{\text{init},(b,p)}\) is likewise initialized from scrambled +data (\texttt{null\_B\_init=scrambled\_data}, the default) or randomly +(\texttt{null\_B\_init=rand}). The real locked +\(\hat A^{(b)}\) is deliberately not used to seed the null refit. + +The real locked fit defaults to data-driven initialization +(\texttt{init\_A=data}, \texttt{init\_B=data}) rather than reference-fit +initialization. This sacrifices some computational convenience, but it makes +each bag start from its own sampled lag-pair evidence and therefore preserves +more of the sampling variance that the FDR/stability procedure is designed +to expose. + +The configuration \(\theta_{\text{M}}\) matches the real locked-M-step +pipeline exactly (Section~\ref{sec:ep2}): identical \(\lambda_3\), +identical \(\rho_{\max}\), identical learning-rate decay, all active from +the first epoch. This is the crux of the calibration: the real +\(\hat A^{(b)}\) and the null +\(A^{\text{null},(b,p)}\) are both outputs of the \emph{same} regression, +shrinkage, and spectral projection, applied to the same state timeline, +differing only in whether cross-neuron alignment was destroyed and in the +resulting refit biases. A +significance statement of the form ``this edge is larger than the same +machinery produces on scrambled data'' is therefore meaningful; using a raw, +unshrunk M-step for the null would compare against a differently biased +quantity and invalidate the test. + +\subsection{How many scrambles are needed} +\label{sec:howmanyscrambles} + +A useful efficiency follows from testing significance per source column +rather than per edge (Section~\ref{sec:aggregate}). For source neuron \(j\) +we pool the \((N-1)\) off-diagonal null weights of that column across all +scrambles. A \emph{single} scramble therefore already contributes +\(\sim(N-1)\) null samples to source \(j\)'s distribution --- of order +\(199\) for \(N=200\). With a modest \(P=4\)--\(8\), the pooled per-source +null holds \((N-1)\times P\approx800\)--\(1600\) samples, ample to resolve a +\(1\%\) tail. By contrast a \emph{per-edge} null would need on the order of +hundreds of scrambles just to populate a single edge's distribution, which +is infeasible at this budget. The two regimes are not equivalent: the +\((N-1)\) samples within one scramble are not independent (they share a +single scrambled realization of the data), so increasing \(P\) still helps +by adding genuinely independent null realizations and by averaging out +scramble-to-scramble variability. The recommendation is to keep \(P\) small +(\(4\)--\(8\)) but greater than one, precisely because per-source pooling +supplies the raw sample count while extra scrambles supply independence. + +% =============================================================== +\section{Stage (a.3): Per-Bag Outputs} +\label{sec:bagout} +% =============================================================== + +Each per-bag job writes a self-contained file to a shared results directory, +with a name carrying the bag index so that \(N_{\text{bag}}\) parallel jobs +do not collide. The file contains, at minimum: the real connectivity +\(\hat A^{(b)}\in\mathbb{R}^{N\times N}\); the stack of null connectivities +\(\{A^{\text{null},(b,p)}\}_{p=1}^{P}\) of shape \((P,N,N)\); the frozen +reference state labels used by the selected pairs; the selected pair indices +\(I_b\); \texttt{bag\_frac}; the reference fit filename and reference +time-range metadata; the realized per-neuron shifts (for reproducibility +and for the data-reuse study); the stack of refit null biases +\(\{\mathbf B^{\text{null},(b,p)}\}_{p=1}^{P}\) of shape \((P,M,N)\); the +fitted off-diagonal sparsity of \(\hat A^{(b)}\) (number and fraction of +surviving edges) as a sanity cross-check; and the locked-M-step history +summary. Stage~(b) consumes only these files and never re-runs PRISM-EM. + +\paragraph{Per-bag program output naming.} One invocation of +\texttt{prism\_FDR\_Bags\_train3c.py} produces one bag file containing +the real locked-M-step fit and all \(P\) scramble refits for that bag. The +filename is deterministic and includes \texttt{outFitName} and +\texttt{bag\_idx}, for example +\texttt{.bag003.prismFDRbag.npz}. The real fit fields are written +in the same outer schema as the standalone \texttt{prism\_EM\_train3c.py} +output where practical, so the usual plotting code can inspect the real +part of a bag. Time-series fields used for plotting are copied from the +reference EM fit or explicitly marked as reference-derived; they are not +newly inferred in the bag job. + +\paragraph{Per-bag driver runtime controls.} The bag driver uses the locked +M-step controls needed to estimate \(A\) and \(\mathbf B\) +(\texttt{m\_epochs}, learning rates, \(\lambda_3,\rho_{\max}\), schedule +controls as applicable, batch size, and initialization modes for the locked +real fit). Its additional controls are: \texttt{bag\_idx}, +\texttt{outFitName}, \texttt{bag\_frac}, \texttt{num\_scrambles}, +\texttt{min\_roll\_shift\_sec}, \texttt{null\_A\_init}, and +\texttt{null\_B\_init}. It does not accept +\texttt{--time\_range\_sec}, \texttt{--num\_blocks}, +\texttt{--block\_len\_sec}, \texttt{--min\_block\_start\_sep\_sec}, or +\texttt{--max\_block\_draw\_trials}. There is no separate +\texttt{null\_m\_epochs} parameter: the locked null M-step inherits the +M-step epoch count and optimizer configuration from the locked real fit, +with pruning, spectral projection, and learning-rate decay active from the +first locked epoch. + +\paragraph{Compute layout.} Each per-bag job is a single SLURM job on a +single node using \(4\times\)A100, matching the existing PRISM-EM launch +(\texttt{torchrun --standalone --nnodes=1 --nproc\_per\_node=4}). The +\(P\) scramble refits reuse the same node and the frozen reference labels, +so they are cheap relative to a full EM fit. The \(N_{\text{bag}}\) jobs are +mutually independent and submitted in parallel. + +% =============================================================== +\section{Stage (b): Aggregation, Significance, and Pruning} +\label{sec:aggregate} +% =============================================================== + +Stage~(b) is a single postprocessing program pointed at the results +directory. It performs per-source null pooling, per-bag selection, and +cross-bag stability selection, then reports the vetted edge list with +uncertainties. + +\subsection{Per-source pooled null and per-bag selection} +\label{sec:persource} + +We test significance \emph{per source column}, exploiting the assumption that +the outgoing edges from a given source neuron \(j\) are exchangeable under +the null: they share that source neuron's spike statistics and shuffled +history stream. This matches the source-column convention +\(A_{ij}\) = effect of source neuron \(j\) on target neuron \(i\), and aligns +the null threshold with source-neuron quantities such as firing rate, Dale +type, outgoing edge count, and \texttt{Sedge}. The per-source exchangeability +assumption is what makes the small scramble budget viable +(Section~\ref{sec:howmanyscrambles}). + +For bag \(b\) and source column \(j\), pool the off-diagonal null weights across all +scrambles of that bag, +\[ +\mathcal{N}^{(b)}_j +=\big\{\,A^{\text{null},(b,p)}_{ij}\ :\ i\neq j,\ p=1,\dots,P\,\big\}, +\qquad |\mathcal{N}^{(b)}_j| = (N-1)\,P, +\] +and form a per-source threshold as the upper-\(q\) quantile of the magnitudes, +\[ +\tau^{(b)}_j = \mathrm{Quantile}_{q}\!\big(\,\{|x| : x\in\mathcal{N}^{(b)}_j\}\,\big), +\qquad q = 0.99 \ \text{(nominal)}. +\] +An edge is \emph{selected in bag \(b\)} if its real magnitude exceeds the +source threshold, +\[ +\text{sel}^{(b)}_{ij} = \mathbb{1}\!\big[\,|\hat A^{(b)}_{ij}| > \tau^{(b)}_j\,\big], +\qquad i\neq j. +\] +See Figure~\ref{fig:edge-selection-cartoon} for a cartoon illustrating this quantile-based selection process. + +\begin{figure}[ht] +\centering +\begin{tikzpicture}[x=1.2cm,y=1.2cm] + % Draw filled areas for null distribution + \path[fill=gray!15] (0,0) -- (0, 2.5) .. controls (1.5, 2.5) and (2.2, 0.5) .. (3.0, 0.2) -- (3.0, 0) -- cycle; + \path[fill=red!20] (3.0,0) -- (3.0, 0.2) .. controls (3.4, 0.1) and (3.8, 0.02) .. (4.2, 0) -- cycle; + + % Draw curve of the null distribution + \draw[thick, gray!80] (0, 2.5) .. controls (1.5, 2.5) and (2.2, 0.5) .. (3.0, 0.2) + .. controls (3.4, 0.1) and (3.8, 0.02) .. (4.2, 0); + + % Axis + \draw[->, thick] (0,0) -- (7.5,0) node[right] {\small weight magnitude \(|A_{ij}|\)}; + \draw[thick] (0,0) -- (0,2.8); + + % Threshold line + \draw[dashed, red, very thick] (3.0, 0) -- (3.0, 2.8) node[above, red, align=center] + {\small threshold \(\tau_j^{(b)}\)\\[-0.5mm]\scriptsize (quantile \(q = 0.99\))}; + + % Labels for regions + \node[gray!80!black] at (1.2, 1.1) {\small Pooled Null}; + \node[gray!80!black] at (1.2, 0.8) {\small Distribution \(\mathcal{N}_j^{(b)}\)}; + + % Moved higher to avoid collision + \node[red!80!black, anchor=west] at (3.2, 1.2) {\scriptsize Significance tail}; + \node[red!80!black, anchor=west] at (3.2, 1.0) {\scriptsize (fraction \(1-q\))}; + \draw[->, red!80!black, line width=0.5pt] (3.3, 0.9) -- (3.6, 0.1); + + % Real weight points and labels + % Discarded points + \fill[red] (0.8, 0) circle (2.5pt) node[above, red] {\scriptsize \(\bm{\times}\)}; + \fill[red] (2.2, 0) circle (2.5pt) node[above, red] {\scriptsize \(\bm{\times}\)}; + \node[below, red, align=center] at (1.5, -0.15) + {\scriptsize Real weights \(|\hat A_{ij}^{(b)}| < \tau_j^{(b)}\)\\[-0.5mm]\scriptsize \textbf{Discarded} (noise floor)}; + + % Selected points + \fill[green!60!black] (4.8, 0) circle (2.5pt) node[above, green!60!black] {\scriptsize \(\bm{\checkmark}\)}; + \fill[green!60!black] (6.5, 0) circle (2.5pt) node[above, green!60!black] {\scriptsize \(\bm{\checkmark}\)}; + \node[below, green!60!black, align=center] at (5.65, -0.15) + {\scriptsize Real weights \(|\hat A_{ij}^{(b)}| > \tau_j^{(b)}\)\\[-0.5mm]\scriptsize \textbf{Selected} (significant)}; + +\end{tikzpicture} +\caption{Cartoon of quantile-based edge selection per bag. The magnitude of each real fitted edge weight (dots) is compared against the threshold \(\tau_j^{(b)}\) derived from the pooled null distribution of source neuron \(j\). Real weights below the threshold (red \(\bm{\times}\)) are discarded as noise; only weights exceeding the threshold (green \(\bm{\checkmark}\)) are selected in this bag.} +\label{fig:edge-selection-cartoon} +\end{figure} + +Equivalently, each real edge can be converted to a within-source \(p\)-value +against the pooled null and the Benjamini--Hochberg procedure applied within +the source column; the quantile-threshold form above is the operational default. The +expected number of edges selected per bag and per source is recorded as a +diagnostic (it is the quantity \(q_\Lambda\) of the stability bound below); +it is \emph{measured}, never prescribed, so the algorithm discovers how many +edges it keeps rather than being told a target sparsity. + +\subsection{Cross-bag stability selection} + +A single bag's selection is noisy. We retain only edges that are selected +reliably across bags. Define the selection frequency +\[ +\pi_{ij} = \frac{1}{N_{\text{bag}}}\sum_{b=1}^{N_{\text{bag}}}\text{sel}^{(b)}_{ij}, +\] +and keep edge \((i,j)\) in the final network if +\(\pi_{ij}\ge\pi_{\text{thr}}\), with +\(\pi_{\text{thr}}=0.6\)--\(0.8\) (nominal). This is Meinshausen--B\"uhlmann +stability selection: the data are perturbed by resampling, the base selector +is run on each perturbation, and only consistently chosen variables survive. +For example, the implementation default \texttt{--stab\_sel\_thresh 0.7} means that +an off-diagonal edge must be selected in at least \(70\%\) of the bags to +survive. The value assigned to a surviving edge is not weighted by the +selection frequency or by any bag quality score; it is the plain arithmetic +mean of \(\hat A^{(b)}_{ij}\) over the bags in which that edge passed the +per-bag null threshold. Bags where the edge did not pass the per-bag selector +do not enter this final \texttt{A\_hat} value, although the all-bag mean is +also saved as a diagnostic. +See Figure~\ref{fig:stability-selection-cartoon} for a cartoon of this +cross-bag pruning step. + +\begin{figure}[ht] +\centering +\begin{tikzpicture}[x=0.95cm,y=0.95cm] + \def\bagcols{10} + \def\colw{1.05} + \def\rowA{1.35} + \def\rowB{-0.55} + \def\xleft{2.8} + \def\xfreq{13.8} + \def\xdec{15.7} + + % column headers + \foreach \b/\lab in {0/0,1/1,2/2,3/3,4/4,5/5,6/6,7/7,8/8,9/9} { + \pgfmathsetmacro{\xc}{\xleft + \b*\colw} + \node[gray!70!black] at (\xc, 2.35) {\scriptsize bag \lab}; + } + \node[gray!70!black, align=center] at (\xfreq, 2.35) + {\scriptsize selection\\[-1mm]\scriptsize frequency \(\pi_{ij}\)}; + \node[gray!70!black, align=center] at (\xdec, 2.35) + {\scriptsize cross-bag\\[-1mm]\scriptsize decision}; + + % threshold annotation + \draw[dashed, orange!85!black, very thick] (\xleft-0.35, 2.95) -- (\xdec+0.55, 2.95); + \node[orange!85!black, anchor=west] at (\xdec+0.62, 2.95) + {\scriptsize \(\pi_{\text{thr}}=0.7\)}; + + % row labels + \node[left, align=right] at (\xleft-0.15, \rowA) + {\scriptsize edge \((i,j)\)}; + \node[left, align=right] at (\xleft-0.15, \rowB) + {\scriptsize edge \((i',j')\)}; + + % per-bag selections for edge (i,j): 8/10 selected + \foreach \b in {0,1,2,3,4,5,6,7} { + \pgfmathsetmacro{\xc}{\xleft + \b*\colw} + \node[green!60!black] at (\xc, \rowA) {\scriptsize \(\bm{\checkmark}\)}; + } + \foreach \b in {8,9} { + \pgfmathsetmacro{\xc}{\xleft + \b*\colw} + \node[red!80!black] at (\xc, \rowA) {\scriptsize \(\bm{\times}\)}; + } + \node at (\xfreq, \rowA) {\scriptsize \(8/10=0.8\)}; + \draw[rounded corners, thick, green!60!black, fill=green!8] + (\xdec-0.55, \rowA-0.28) rectangle (\xdec+0.55, \rowA+0.28); + \node[green!45!black] at (\xdec, \rowA) {\scriptsize \textbf{keep}}; + + % per-bag selections for edge (i',j'): 3/10 selected + \foreach \b in {0,2,5} { + \pgfmathsetmacro{\xc}{\xleft + \b*\colw} + \node[green!60!black] at (\xc, \rowB) {\scriptsize \(\bm{\checkmark}\)}; + } + \foreach \b in {1,3,4,6,7,8,9} { + \pgfmathsetmacro{\xc}{\xleft + \b*\colw} + \node[red!80!black] at (\xc, \rowB) {\scriptsize \(\bm{\times}\)}; + } + \node at (\xfreq, \rowB) {\scriptsize \(3/10=0.3\)}; + \draw[rounded corners, thick, red!75!black, fill=red!8] + (\xdec-0.55, \rowB-0.28) rectangle (\xdec+0.55, \rowB+0.28); + \node[red!75!black] at (\xdec, \rowB) {\scriptsize \textbf{drop}}; + + % legend for per-bag marks + \node[green!60!black] at (\xleft+0.2, -1.45) {\scriptsize \(\bm{\checkmark}\)}; + \node[anchor=west, align=left] at (\xleft+0.55, -1.45) + {\scriptsize selected in bag \(b\): \(|\hat A^{(b)}_{ij}|>\tau^{(b)}_j\)}; + \node[red!80!black] at (\xleft+0.2, -1.95) {\scriptsize \(\bm{\times}\)}; + \node[anchor=west, align=left] at (\xleft+0.55, -1.95) + {\scriptsize not selected in bag \(b\) (below per-bag null threshold)}; +\end{tikzpicture} +\caption{Cartoon of cross-bag stability selection. Each column is one FDR bag; +each row is one candidate off-diagonal edge. A green \(\bm{\checkmark}\) means +the edge passed the per-bag null threshold in that bag (Section~\ref{sec:persource}); +a red \(\bm{\times}\) means it did not. The selection frequency +\(\pi_{ij}\) is the fraction of bags with a checkmark. Edge \((i,j)\) appears +in \(8\) of \(10\) bags (\(\pi_{ij}=0.8\ge\pi_{\text{thr}}\)) and is kept in +the final network; edge \((i',j')\) appears in only \(3\) bags +(\(\pi_{ij}=0.3<\pi_{\text{thr}}\)) and is removed.} +\label{fig:stability-selection-cartoon} +\end{figure} + +\paragraph{Reported stability bound.} The implementation exposes the pair +\((q,\ \pi_{\text{thr}})\): a per-bag per-source quantile level \(q\) and a +cross-bag frequency threshold \(\pi_{\text{thr}}\). It also reports the +Meinshausen--B\"uhlmann-style stability-selection diagnostic +\[ +\mathbb{E}[V]\;\le\;\frac{1}{2\pi_{\text{thr}}-1}\cdot\frac{q_\Lambda^{2}}{p_{\text{cand}}}, +\] +where \(\mathbb{E}[V]\) is the expected number of false edges, +\(q_\Lambda\) is the \emph{measured} average number of edges selected per bag, +and \(p_{\text{cand}}=N(N-1)\) is the number of candidate off-diagonal edges. +The current program does not derive \(\pi_{\text{thr}}\) from a target false-edge +count; the user sets \texttt{--per\_bag\_quantile} and \texttt{--stab\_sel\_thresh} +directly, and the bound is saved as a diagnostic. The honest caveat is that +the bound assumes independent exchangeable bags, whereas all bags are random +subsets from the same finite reference-labeled pair table and can overlap in +selected rows. They are therefore positively correlated, so the reported +value should be read as an approximate stability diagnostic rather than as a +guaranteed FDR. + +% =============================================================== +\section{Per-Edge Uncertainty: Two Distinct Errors} +\label{sec:errors} +% =============================================================== + +The design furnishes two genuinely different per-edge ``errors,'' which +answer different questions and must not be conflated. + +\subsection{Null spread and effect size} + +For source column \(j\), the pooled null +\(\mathcal{N}_j=\bigcup_b\mathcal{N}^{(b)}_j\) +has a mean and standard deviation, +\[ +m^{\text{null}}_j=\operatorname{mean}\big(\mathcal{N}_j\big), +\qquad +s^{\text{null}}_j=\operatorname{SD}\big(\mathcal{N}_j\big), +\] +which characterize the \emph{noise floor}: how large an edge weight looks +purely by chance once cross-neuron alignment is destroyed and the same +machinery is applied. This is not the uncertainty of the estimate; it is the +spread of the null. It does, however, define a natural null-referenced +effect size, +\[ +z_{ij}=\frac{\bar A_{ij}-m^{\text{null}}_j}{s^{\text{null}}_j}, +\] +a per-edge signal-to-noise ratio against the scramble null, where +\(\bar A_{ij}\) is the bagged point estimate defined next. +See Figure~\ref{fig:null-effect-cartoon} for a cartoon of this null-referenced +effect size. + +\begin{figure}[ht] +\centering +\begin{tikzpicture}[x=1.1cm,y=1.1cm] + % symmetric null distribution around m_null (plot occupies x in [0,6]) + \path[fill=gray!15] + (1.0,0) -- (1.0, 2.15) .. controls (2.0, 2.15) and (2.55, 0.55) .. (3.0, 0.12) + .. controls (3.45, 0.55) and (4.0, 2.15) .. (5.0, 2.15) -- (5.0, 0) -- cycle; + \draw[thick, gray!80] + (1.0, 2.15) .. controls (2.0, 2.15) and (2.55, 0.55) .. (3.0, 0.12) + .. controls (3.45, 0.55) and (4.0, 2.15) .. (5.0, 2.15); + + \draw[->, thick] (0,0) -- (6.0,0); + \draw[thick] (0,0) -- (0,2.45); + + \draw[dashed, gray!70, thick] (3.0, 0) -- (3.0, 2.45); + \node[gray!70!black, above] at (3.0, 2.45) {\scriptsize \(m^{\text{null}}_j\)}; + + \draw[<->, gray!70, thick] (2.35, -0.48) -- (3.65, -0.48); + \node[gray!70!black, below] at (3.0, -0.48) {\scriptsize \(s^{\text{null}}_j\)}; + + \node[gray!80!black, align=left, anchor=west] at (0.15, 1.85) + {\small pooled null \(\mathcal{N}_j\)\\[-1mm]\scriptsize (all bags, all scrambles)}; + + \fill[blue!70!black] (5.6, 0) circle (2.5pt); + \node[blue!70!black, above] at (5.6, 0.08) {\scriptsize \(\bar A_{ij}\)}; + + \draw[->, blue!70!black, very thick] (3.12, 0.22) -- (5.45, 0.22); + + \node[blue!70!black, align=left, anchor=west, text width=3.6cm] at (6.75, 1.05) + {\scriptsize \(z_{ij}=\dfrac{\bar A_{ij}-m^{\text{null}}_j}{s^{\text{null}}_j}\)}; + + \node[below, gray!70!black] at (1.4, -0.48) {\small edge weight \(A_{ij}\)}; +\end{tikzpicture} +\caption{Cartoon of null spread and null-referenced effect size. +\textbf{Question answered:} how large does this edge look relative to the +scramble-noise floor? The gray curve is the pooled per-source null +distribution \(\mathcal{N}_j\) from all scrambled locked refits. Its mean +\(m^{\text{null}}_j\) and standard deviation \(s^{\text{null}}_j\) define that +noise floor: how large an edge weight can appear by chance once cross-neuron +alignment is destroyed. The real bagged estimate \(\bar A_{ij}\) (blue dot) is +compared to the floor through \(z_{ij}\), a signal-to-noise ratio against the +null. This is \emph{not} the sampling error of \(\bar A_{ij}\); that distinct +quantity is the across-bag bootstrap SD (Section~\ref{sec:bootstrap-sd}).} +\label{fig:null-effect-cartoon} +\end{figure} + +\subsection{Estimation variability: the across-bag bootstrap SD} +\label{sec:bootstrap-sd} + +The quantity one usually means by ``the statistical error of an edge'' is the +sampling variability of its fitted weight, and bagging estimates it directly. +Each bag yields a real fitted weight \(\hat A^{(b)}_{ij}\) on a different +random subset of the reference-labeled lag pairs; their spread across bags +is a pair-subsampling estimate of the estimator's sampling distribution. We +report the bagged mean and bootstrap SD, +\[ +\bar A_{ij}=\frac{1}{B_{ij}}\sum_{b\in\mathcal{B}_{ij}}\hat A^{(b)}_{ij}, +\qquad +\widehat{\mathrm{SD}}_{ij} +=\sqrt{\frac{1}{B_{ij}-1}\sum_{b\in\mathcal{B}_{ij}}\big(\hat A^{(b)}_{ij}-\bar A_{ij}\big)^{2}}, +\] +where \(\mathcal{B}_{ij}\) is the set of bags entering the average and +\(B_{ij}=|\mathcal{B}_{ij}|\). +See Figure~\ref{fig:bootstrap-sd-cartoon} for a cartoon of this across-bag +estimation variability. + +\begin{figure}[ht] +\centering +\begin{tikzpicture} + \begin{scope}[x=0.82cm, y=0.11cm] + \def\ymean{15} + \def\ysdlo{12} + \def\ysdhi{18} + + \draw[->, thick] (0,0) -- (0,22) node[above, align=center] + {\scriptsize fitted weight\\[-1mm]\scriptsize \(\hat A^{(b)}_{ij}\)}; + \draw[->, thick] (0,0) -- (8.6,0); + + \foreach \b in {0,1,2,3,4,5,6,7} { + \node[below, gray!70!black] at (\b+0.55, 0) {\scriptsize \b}; + } + \node[below, gray!70!black] at (8.35, 0) {\scriptsize bag index \(b\)}; + + \foreach \b/\val in {0/14, 1/17, 2/13, 3/16, 4/15, 5/18, 6/14, 7/16} { + \fill[green!60!black] (\b+0.55, \val) circle (2.2pt); + } + + \draw[dashed, blue!70!black, very thick] (-0.25, \ymean) -- (8.45, \ymean); + \node[left, blue!70!black, align=right] at (-0.3, \ymean) + {\scriptsize \(\bar A_{ij}\)}; + + \draw[<->, orange!85!black, very thick] (8.65, \ysdlo) -- (8.65, \ysdhi); + \node[orange!85!black, anchor=west, align=left] at (8.75, \ymean) + {\scriptsize \(\widehat{\mathrm{SD}}_{ij}\)}; + \end{scope} +\end{tikzpicture} +\caption{Cartoon of across-bag bootstrap SD. +\textbf{Question answered:} how much does the fitted edge weight move when +resampling reference-labeled lag pairs across bags? For a surviving edge +\((i,j)\), each bag produces a real fitted weight \(\hat A^{(b)}_{ij}\) on a +different random subset of reference-labeled lag pairs. The green dots show +eight such bag-level estimates; the dashed horizontal line is their mean +\(\bar A_{ij}\), and the orange bracket is the bootstrap standard deviation +\(\widehat{\mathrm{SD}}_{ij}\). This spread measures \emph{estimation +variability}: how reliably the edge weight is estimated under pair resampling. +It does not describe how large the edge looks under scrambled data; that role +belongs to the null spread \(s^{\text{null}}_j\) and effect size \(z_{ij}\) +(Figure~\ref{fig:null-effect-cartoon}).} +\label{fig:bootstrap-sd-cartoon} +\end{figure} + +\paragraph{Which bags to average over.} A bag in which edge \((i,j)\) was +pruned to exactly zero by \(L_1\) contributes a structural zero, not a noisy +nonzero. Two summaries are therefore reported. The +\emph{selection-conditional} estimate averages only over bags where the edge +survived (\(\mathcal{B}_{ij}=\{b:\text{sel}^{(b)}_{ij}=1\}\)); it reflects the +typical magnitude when the edge is detected but is biased upward in +magnitude. The \emph{all-bag} estimate averages over every bag +(\(\mathcal{B}_{ij}=\{1,\dots,N_{\text{bag}}\}\), structural zeros included); +it is diluted toward zero in proportion to how often the edge is absent. We +recommend reporting the selection-conditional mean and bootstrap SD +\emph{together with} the selection frequency \(\pi_{ij}\), so that the dilution is visible +rather than hidden: \(\pi_{ij}\) and \(\bar A_{ij}\) together convey both how +often and how strongly an edge appears. + +\paragraph{Honest caveat on calibration.} The bags are overlapping random +subsets from one finite reference window, not independent datasets. Shared +reference state labels and repeated use of many of the same lag pairs make +bags more alike than independent draws, so the naive across-bag SD can +\emph{underestimate} the true sampling error; with only +\(N_{\text{bag}}=5\)--\(10\) bags it is additionally coarse (few degrees of +freedom). Accordingly \(\widehat{\mathrm{SD}}_{ij}\) should be used as a +\emph{relative} reliability measure --- excellent for ranking edges and for +flagging unstable ones --- rather than as a calibrated frequentist standard +error supporting a strict confidence interval. The degree of underestimation +can be probed empirically by changing \texttt{--bag\_frac} and the number of +bags. + +\subsection{Reported per-edge record} + +For every edge that passes stability selection +(\(\pi_{ij}\ge\pi_{\text{thr}}\)) the pipeline emits +\[ +\big(\,i,\ j,\ \pi_{ij},\ \bar A_{ij},\ \widehat{\mathrm{SD}}_{ij},\ z_{ij}\,\big), +\] +together with the all-bag mean for comparison and the source null summary +\((m^{\text{null}}_j,s^{\text{null}}_j)\). The selection frequency +\(\pi_{ij}\), the bagged point estimate, the bootstrap SD, and the +null-referenced effect size are complementary: an edge that is frequently +selected, large relative to its bootstrap SD, and far into the null tail is +trustworthy; an edge strong in one bag but rarely selected, or with bootstrap SD +comparable to its magnitude, is flagged as unreliable. + +% =============================================================== +\section{Implementation and Runtime Arguments} +\label{sec:implementation} +% =============================================================== + +This section records the concrete command-line interface used by the current +implementation. Stage~(a) produces one bag file per invocation. Stage~(b) +reads the ordered set of bag files and writes a single aggregate +\texttt{.prismEM.npz} file that can be opened by the existing +\texttt{prism\_EM\_eval3c.py} machinery. + +\subsection{\texttt{prism\_FDR\_Bags\_train3c.py}: Stage (a)} + +This program is run once per bag, normally as one multi-GPU job: +\[ +\texttt{torchrun --standalone --nnodes=1 --nproc\_per\_node=4 +./prism\_FDR\_Bags\_train3c.py ...} +\] +It loads the reference EM output from \texttt{prismFit}, loads the original +spikes from \texttt{spikesData}, samples a fixed-size subset of +reference-labeled lag pairs, performs one real locked-M-step fit on that +subset, runs the requested circular-shift locked-M-step null refits, and +writes one self-contained bag file. + +\paragraph{Input and output naming.} +The reference EM fit is +\[ +\texttt{/prismFit/.prismEM.npz}. +\] +The input spike file is derived from this reference fit's provenance +metadata, so Stage~(a) cannot accidentally combine an EM fit from one +dataset with spikes from another. For simulations this provenance points to +\texttt{state\_transition\_file}; for experiments it points to +\texttt{experiment\_name}. The corresponding spike file is read from +\[ +\texttt{/spikesData/.spikes.npz}. +\] +The output directory is \texttt{fdr\_out\_dir} if supplied, otherwise +\[ +\texttt{/prismFDR}. +\] +The output file is deterministic in \texttt{outFitName} and \texttt{bag\_idx}: +\[ +\texttt{/prismFDR/.bag003.prismFDRbag.npz} +\] +for \(\texttt{bag\_idx}=3\) and default output directory. The bag index is +not only a label: it also enters the random seed stream, so rerunning the +same command with the same \texttt{bag\_idx} reproduces the same bag and +scramble shifts, while changing only \texttt{bag\_idx} produces a different +bag file. There is no special chronological bag; \texttt{bag\_idx=0} is the +first random pair subset. + +A single Stage~(a) bag can be evaluated directly by passing the bag stem to +\texttt{prism\_EM\_eval3c.py}; if \texttt{dataName} contains a +\texttt{bagNNN} token, the evaluator reads +\texttt{/prismFDR/.prismFDRbag.npz} instead of looking +under \texttt{prismFit}. + +\paragraph{Required arguments.} +\begin{description} + \item[\texttt{--basePath}] Root directory containing \texttt{spikesData/} + and receiving \texttt{prismFDR/}. This is the same run directory used by + the plain PRISM-EM trainer. + \item[\texttt{--emFitName}] EM-train input fit stem in \texttt{prismFit/}. + This file supplies the hard Viterbi state labels used by the locked + M-step and the provenance record used to locate the original spike file. + \item[\texttt{--outFitName}] Stage~(a) output stem in \texttt{prismFDR/}. + The program writes + \texttt{.bag.prismFDRbag.npz}. If omitted, the + output stem defaults to \texttt{\_} with a random + four-character hexadecimal suffix. + \item[\texttt{--bag\_idx}] Integer bag index. It is used in the output + filename and in the per-bag RNG seed. All bag indices, including + \texttt{bag\_idx=0}, draw random pair subsets. +\end{description} + +\paragraph{Bagging and scramble arguments.} +\begin{description} + \item[\texttt{--fdr\_out\_dir}] Optional output directory for Stage~(a) + bag files. Default: \texttt{/prismFDR}. + \item[\texttt{--bag\_frac}] Fraction \(f_{\rm bag}\) of reference-window + lag pairs selected without replacement for each bag. Default: \(0.8\). + \item[\texttt{--num\_scrambles}] Number \(P\) of circular-shift null + refits saved in this bag file. Default: \(4\). + \item[\texttt{--min\_roll\_shift\_sec}] Two-sided exclusion zone around + zero circular shift for each neuron. A shift is drawn from the interval + excluding shifts smaller than this value and shifts closer than this value + to a full wrap. Default: \(2\) seconds. + \item[\texttt{--null\_A\_init}] Fresh \(A\) initialization mode for every + null refit. Values: \texttt{scrambled\_data} (default) or \texttt{rand}. + The real locked-fit \(\hat A\) is deliberately not used to seed the null. + \item[\texttt{--null\_B\_init}] Fresh \(\mathbf B\) initialization mode for + every null refit. Values: \texttt{scrambled\_data} (default) or + \texttt{rand}. The null \(\mathbf B\) is refit and saved as + \texttt{B\_null}. +\end{description} + +\paragraph{Inherited locked-M-step arguments.} +The Stage~(a) driver shares the M-step machinery with +\texttt{prism\_EM\_train3c.py}, but it does not run a bag-specific E-step. +The default values are loaded from the reference EM metadata so a bag job is +reproducible from the reference fit alone. Any argument listed here may be +redefined on the \texttt{prism\_FDR\_Bags\_train3c.py} command line; the +resolved values are saved in the bag metadata and used identically for the +real locked fit and all null locked fits in that bag. This lets the +reference EM run be tuned for state assignment while the bag jobs are tuned +for conditional \(A,\mathbf B\) estimation. +\begin{description} + \item[\texttt{--m\_epochs}] M-step Adam epochs for each locked fit. The + null locked M-step uses the same count; there is no separate + \texttt{null\_m\_epochs}. + \item[\texttt{--lr\_mstep}] Adam learning rate for the M-step. Default: + \(0.003\). + \item[\texttt{--lr\_end\_factor}] Final learning-rate multiplier after + linear decay. Default: \(0.1\). + \item[\texttt{--lambda3}] L1 penalty weight on off-diagonal \(A\). Default: + \(0.02\). + \item[\texttt{--rho\_max}] Target spectral radius for soft correction of + \(A\). Default: \(0.95\). + \item[\texttt{--prescale\_m\_step\_4\_ArhoMax}] M-step batch interval for + spectral-radius projection checks. Default: \(120\). + \item[\texttt{--delay\_em\_iter\_4\_ArhoMax}, + \texttt{--delay\_em\_iter\_4\_lrDecay}, + \texttt{--delay\_em\_iter\_4\_Aprune}] Accepted by the shared M-step + configuration. For locked FDR refits, schedule-gated operations are active + from the first locked epoch unless the implementation explicitly maps + these controls to locked-epoch delays. + \item[\texttt{--batch\_size}] M-step minibatch size. Default: \(2048\). + \item[\texttt{--init\_A}] Initial \(A\) mode for the locked real fit: + \texttt{data} (default), \texttt{ref}, or \texttt{rand}. + \item[\texttt{--init\_B}] Initial \(\mathbf B\) mode for the locked real fit: + \texttt{data} (default), \texttt{ref}, or \texttt{rand}. + The default \texttt{data} mode is chosen to capture bag-to-bag sampling + variance in the FDR procedure. Initializing every real bag from the same + reference EM \(\hat A,\hat{\mathbf B}\) is faster and stable, but it can + make bags overly correlated with the reference solution and understate the + variability that the stability filter is meant to measure. + \item[\texttt{--init\_A\_Tmax}] Maximum number of bins used for + data-driven \(A\) initialization. Default: \(50000\). + \item[\texttt{--seed}] Base random seed. Default: \(42\); the effective + seed is offset by \texttt{bag\_idx}. + \item[\texttt{--verb}] Verbosity level. With the default verbosity, rank + \(0\) reports runtime device information, real-fit elapsed time, and + elapsed time for each scramble refit. +\end{description} + +The bag trainer does not expose state-discovery controls such as +\texttt{--num\_em\_iters}, \texttt{--pgd\_iter}, \texttt{--lr\_estep}, +\texttt{--lambda2}, \texttt{--init\_states}, \texttt{--decode\_dwell\_sec}, +or \texttt{--num\_states}. These belong to the reference +\texttt{prism\_EM\_train3c.py} run. In particular, \texttt{--num\_states} +is inferred from the locked reference labels. + +\paragraph{Saved Stage (a) contents.} +The file has the real locked-fit fields (\texttt{A\_hat}, \texttt{B\_hat}) +plus reference-derived time-series fields needed for plotting +(\texttt{S\_hat}, \texttt{c\_hat}, and confidence/loss summaries when saved +upstream). It also contains \texttt{A\_null}, \texttt{B\_null}, +\texttt{roll\_shifts\_bin}, \texttt{bag\_pair\_indices}, +\texttt{bag\_frac}, reference fit provenance, and locked M-step histories. +Stage~(a) metadata are grouped under \texttt{bagsFDR\_stageA}. This includes +the pair-sampling controls, the circular-shift null-refit controls, +per-scramble initialization diagnostics, and a reference-fit sub-dict +recording the \texttt{prismFit} file and time range used to make the locked +state labels. It also includes the effective locked-M-step hyperparameters +after applying command-line overrides to the metadata defaults. + +\paragraph{\texttt{big\_fit\_bags.sh} launch pattern.} +The helper script \texttt{big\_fit\_bags.sh} is a concrete end-to-end launch +example for one dataset. It sets +\[ +\texttt{basePath=/pscratch/sd/b/balewski/2026\_causalNet\_exp\_ver3c} +\] +and a short input spike name such as +\texttt{daleN200\_102308\_a5d24b}. The script derives +\[ +\texttt{runTag=,\qquad emFitName=\_em} +\] +as a random four-character hexadecimal tag. It first runs +\texttt{prism\_EM\_train3c.py} on the configured \texttt{timeRange} and +writes +\[ +\texttt{/prismFit/.prismEM.npz}. +\] +Then all bags launched inside the same script use the same Stage~(a) extended +name +\[ +\texttt{\_em\_fdr}. +\] +For each requested bag index, the script runs +\texttt{torchrun --standalone --nnodes=1 --nproc\_per\_node=4} on +\texttt{prism\_FDR\_Bags\_train3c.py} with +\texttt{--emFitName } and +\texttt{--outFitName \_em\_fdr}. With +\texttt{numBags=6}, the expected Stage~(a) output files are +\begin{verbatim} +/prismFDR/_em_fdr.bag000.prismFDRbag.npz +/prismFDR/_em_fdr.bag001.prismFDRbag.npz +... +/prismFDR/_em_fdr.bag005.prismFDRbag.npz +\end{verbatim} +These are all equivalent random locked-M-step bags; \texttt{bag000} is not a +chronological reference. If aggregation is enabled, Stage~(b) is called with +\texttt{--dataName \_em\_fdr}, \texttt{--outAgrName} +\texttt{\_em\_fdr\_agr}, and +\texttt{--num\_bags} equal to the number of contiguous bags starting from +\texttt{bag000} that are present on disk. + +\subsection{\texttt{prism\_EM\_FDR\_Bags\_aggregate3c.py}: Stage (b)} + +This program is a single CPU postprocessing step. It reads the Stage~(a) +files, performs the per-source null thresholding and cross-bag stability +selection, and writes one aggregate file in the same outer format expected +by \texttt{prism\_EM\_eval3c.py}. + +\paragraph{Input and output naming.} +The default input directory is +\[ +\texttt{/prismFDR}. +\] +For \texttt{num\_bags} \(=B\), the program expects the contiguous file set +\begin{verbatim} +.bag000.prismFDRbag.npz +.bag001.prismFDRbag.npz +... +.bag{num_bags-1, zero padded}.prismFDRbag.npz +\end{verbatim} +Here \texttt{dataName} is the Stage~(a) output stem, i.e. the same value +passed as \texttt{--outFitName} to \texttt{prism\_FDR\_Bags\_train3c.py}. +The output stem is set explicitly by +\texttt{--outAgrName}; this makes it easy to rerun only the aggregation with +different thresholds while preserving the same input bag stem. +If \texttt{--outAgrName} is omitted, the output stem defaults to +\texttt{\_} with a random four-character hexadecimal suffix. +The output file is written to +\[ +\texttt{/prismFit/.prismEM.npz}. +\] +For example, with +\texttt{dataName=daleN100\_2ba29b\_c47b43\_em1234\_fdr1234}, one possible +output stem is +\[ +\texttt{daleN100\_2ba29b\_c47b43\_em1234\_fdr1234\_agrA}. +\] +It can then be opened by +\begin{verbatim} +./prism_EM_eval3c.py \ + --basePath \ + --dataName \ + -p a e f g +\end{verbatim} + +\paragraph{Runtime arguments.} +\begin{description} + \item[\texttt{--basePath}] Root run directory. The default input is + \texttt{/prismFDR}; the aggregate output is written to + \texttt{/prismFit}. + \item[\texttt{--dataName}] Extended Stage~(a) output data name, normally + the Stage~(a) output stem. This exact string is the prefix of + all bag files consumed by Stage~(b), e.g.\ + \texttt{.bag000.prismFDRbag.npz}. + \item[\texttt{--outAgrName}] Aggregate output stem written to + \texttt{prismFit/.prismEM.npz}. This is intentionally separate + from \texttt{--dataName} so repeated aggregation choices can create + differently named outputs from the same bag files. If omitted, the program + appends a random \texttt{\_} suffix to \texttt{dataName}. + \item[\texttt{--num\_bags}] Number of Stage~(a) bag files to read. The + implementation expects bags numbered contiguously from \texttt{bag000}. + \item[\texttt{--per\_bag\_quantile}] Quantile \(q\) used for the per-source + null magnitude threshold \(\tau^{(b)}_j\). Default: \(0.99\). + \item[\texttt{--stab\_sel\_thresh}] Cross-bag stability selection-frequency threshold + \(\pi_{\text{thr}}\). Default: \(0.7\). + \item[\texttt{--fdr\_out\_dir}] Optional override for the Stage~(a) bag + directory. Default: \texttt{/prismFDR}. + \item[\texttt{--verb}] Verbosity level. Default: \(1\). +\end{description} + +\paragraph{Aggregate output schema.} +The output \texttt{A\_hat} is the final stable matrix: off-diagonal entries +are the selection-conditional bagged means for edges with +\(\pi_{ij}\ge\texttt{stab\_sel\_thresh}\), and zero otherwise. With the default +\texttt{--stab\_sel\_thresh 0.7}, this means retaining only edges present in at least +\(70\%\) of bags; the retained edge value is the unweighted mean over only +those bag instances where the edge was selected. Diagonal entries are +the all-bag mean diagonals from the real bag fits; the corresponding +diagonal mean, standard deviation, and standard error are also saved as +\texttt{A\_diag\_mean}, \texttt{A\_diag\_std}, and +\texttt{A\_diag\_stderr}. \texttt{A\_prune}, \texttt{neuron\_type}, and +\texttt{neuron\_Sedge} are recomputed from this aggregate \texttt{A\_hat} +using source columns. First \texttt{min\_posW} and \texttt{max\_negW} are +computed from the stable off-diagonal \texttt{A\_hat} weights. Then +\texttt{neuron\_Sedge} is the off-diagonal column sum for each source +neuron, and source columns are classified as excitatory when +\(\texttt{neuron\_Sedge}>\texttt{min\_posW}\), inhibitory when +\(\texttt{neuron\_Sedge}<\texttt{max\_negW}\), and undecided otherwise. +\texttt{A\_prune} preserves the diagonal, keeps only positive outgoing +weights for excitatory source columns, keeps only negative outgoing weights +for inhibitory source columns, and leaves undecided source columns +sign-unpruned. +The output \texttt{B\_hat} is the mean across bags in the reference state +order. Because every bag uses the hard labels from the same reference EM fit, +state rows are already aligned; no bag-to-bag state permutation is needed +other than an optional sanity check. The time-series fields +\texttt{S\_hat}, \texttt{c\_hat}, \texttt{S\_hat\_CL}, and the EM histories +are copied from the reference EM output to preserve compatibility with the +existing plotting code; those plots should therefore be interpreted as +reference-fit diagnostics, while the connectivity fields are Stage~(b) +aggregates. Stage~(b) metadata are grouped under +\texttt{bagsFDR\_stageB}; reference provenance and sampled-pair metadata +from Stage~(a) remain under \texttt{bagsFDR\_stageA}. + +Additional Stage~(b) arrays include \texttt{A\_bag}, +\texttt{B\_bag}, +\texttt{selected\_mask}, \texttt{selected\_in\_bag}, +\texttt{selection\_count}, \texttt{selection\_frequency}, +\texttt{A\_mean\_selected}, \texttt{A\_sd\_selected}, +\texttt{A\_mean\_all}, \texttt{A\_sd\_all}, \texttt{A\_diag\_mean}, +\texttt{A\_diag\_std}, \texttt{A\_diag\_stderr}, \texttt{src\_null\_tau\_bag}, +\texttt{src\_null\_tau\_mean}, \texttt{src\_null\_mean}, +\texttt{src\_null\_sd}, \texttt{z\_null}, and a compact selected-edge table +with \texttt{edge\_i}, \texttt{edge\_j}, \texttt{edge\_sel\_freq}, +\texttt{edge\_A\_mean}, \texttt{edge\_A\_sd\_boot}, and +\texttt{edge\_z\_null}, plus \texttt{edge\_src\_null\_mean}, +\texttt{edge\_src\_null\_sd}, and \texttt{edge\_src\_tau\_mean}. + +% =============================================================== +\section{End-to-End Summary} +% =============================================================== + +\paragraph{Reference EM fit, run once.} +\begin{enumerate} + \item Run \texttt{prism\_EM\_train3c.py} on the desired + \texttt{--time\_range\_sec}; save + \texttt{prismFit/.prismEM.npz}, for example + \texttt{prismFit/\_emjXXXX.prismEM.npz}. + \item Use its saved reference window, \(\hat S^{\text{ref}}\), and + \(\hat c^{\text{ref}}\) as the only state timeline for all FDR bags. +\end{enumerate} + +\paragraph{Stage (a), run \(N_{\text{bag}}\) times in parallel (one node, +\(4\times\)A100 each).} +\begin{enumerate} + \item Load \texttt{prismFit/.prismEM.npz} and + recover the input spike stem from its provenance metadata. Then load + \texttt{spikesData/.spikes.npz} and form the + reference-labeled pair table + \((H_t,X_t,Z_t)=(\hat S^{\text{ref}}_t,Y_{t-1},Y_t)\). + \item Draw \(K_{\rm bag}=\lfloor\texttt{bag\_frac}(T_{\rm ref}-1)\rfloor\) + pair indices without replacement, using the same selected rows for all + neurons. + \item Run + \(\textsf{mstep\_locked}\to(\hat A^{(b)},\hat{\mathbf B}^{(b)})\) on the + selected real rows with \(H_t=\hat S^{\text{ref}}_t\) fixed. + \item For \(p=1,\dots,P\): draw independent per-neuron circular shifts in + the two-sided zone + \([\delta_{\text{roll}}^{\min,\text{bin}},T_{\rm ref}-\delta_{\text{roll}}^{\min,\text{bin}}]\), + roll each neuron's reference-window spikes, form scrambled pair rows using + the same selected indices, and + \(\textsf{mstep\_locked}\to(A^{\text{null},(b,p)},\mathbf B^{\text{null},(b,p)})\) + with the locked labels \(H_t=\hat S^{\text{ref}}_t\), fresh \(A\) + initialization controlled by + \texttt{null\_A\_init}\(\in\{\texttt{scrambled\_data},\texttt{rand}\}\), + and fresh \(\mathbf B\) initialization controlled by + \texttt{null\_B\_init}\(\in\{\texttt{scrambled\_data},\texttt{rand}\}\). + \item Write one file for bag \(b\), named from \texttt{dataName} and + \texttt{bag\_idx}, containing \(\hat A^{(b)}\), + \(\{A^{\text{null},(b,p)}\}\), \(\{\mathbf B^{\text{null},(b,p)}\}\), + sampled pair indices, realized shifts, fitted sparsity, reference-fit + provenance, and history summary. +\end{enumerate} + +\paragraph{Stage (b), single postprocessing job.} +\begin{enumerate} + \item Per bag and source column, pool the \((N-1)P\) off-diagonal null + weights, form the upper-\(q\) magnitude threshold \(\tau^{(b)}_j\), and + select edges with \(|\hat A^{(b)}_{ij}|>\tau^{(b)}_j\); record + \(q_\Lambda\). + \item Compute selection frequencies \(\pi_{ij}\) and keep edges with + \(\pi_{ij}\ge\pi_{\text{thr}}\), where \(\pi_{\text{thr}}\) is supplied + by \texttt{--stab\_sel\_thresh}. + \item For surviving edges, compute the bagged mean \(\bar A_{ij}\), + bootstrap SD \(\widehat{\mathrm{SD}}_{ij}\) (selection-conditional and + all-bag), and effect size \(z_{ij}\) against the pooled source null. + \item Emit the vetted edge list + \((i,j,\pi_{ij},\bar A_{ij},\widehat{\mathrm{SD}}_{ij},z_{ij})\) plus + diagnostics. +\end{enumerate} + +\paragraph{Closing note.} The method's robustness comes from separating +state discovery from edge testing: connectivity is asserted only when it +survives repeated random subsampling of reference-labeled lag pairs, +exceeds a null built by the identical locked-M-step machinery (scramble +calibration), and recurs across bags (stability selection). Its +two reported uncertainties --- the null spread and the across-bag bootstrap +SD --- are deliberately kept distinct, and their known limitations under +overlapping pair subsets are stated rather than papered over, so that the +resulting edge list can be read with appropriate confidence on real data +where no ground truth is available to check it. + +\end{document} diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/genDale_topo_states_ver3c.pdf b/causal_net/nonStation_ver3c_states_EM_FDR/docs/genDale_topo_states_ver3c.pdf new file mode 100644 index 00000000..06908fd3 Binary files /dev/null and b/causal_net/nonStation_ver3c_states_EM_FDR/docs/genDale_topo_states_ver3c.pdf differ diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/genDale_topo_states_ver3c.tex b/causal_net/nonStation_ver3c_states_EM_FDR/docs/genDale_topo_states_ver3c.tex new file mode 100644 index 00000000..afb0d265 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/genDale_topo_states_ver3c.tex @@ -0,0 +1,797 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} +\usepackage{multirow} +\usepackage{array} +\usepackage{algorithm} +\usepackage{algorithmic} + +\geometry{margin=0.8in} + +\title{Synthetic Non-Stationary Spiking Data Generation:\\ +Dale-Law-Constrained Recurrent Networks with\\ +a Switching Poisson GLM} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +We describe the two-script pipeline that produces synthetic non-stationary +neural population spike trains for benchmarking latent-state inference +algorithms (e.g.\ PRISM-EM). +\textbf{Stage~1} constructs a source-column Dale-style weight matrix +$A_{\text{true}}$ with prescribed spectral radius $R<1$, $M$ bias +vectors $\{B_m\}_{m=0}^{M-1}$, one per dynamical state, and static +two-dimensional node locations. +Stationary spike trains under each bias are generated as a +firing-rate sanity check. +\textbf{Stage~2} loads the ground-truth dictionary, schedules transitions +among the $M$ states, tracks a smoothly evolving simplex coefficient +$c_t$, and generates the non-stationary spike train from a Poisson GLM +with time-varying effective bias +$B_{\mathrm{eff},t} = \sum_m c_{t,m} B_m$. +An oracle state decoder (exact Poisson log-likelihood) is also +computed and saved as an upper bound for any inference algorithm. +\end{abstract} + +\tableofcontents +\newpage + +% =============================================================== +\section{Overview} +\label{sec:overview} +% =============================================================== + +The pipeline produces ground-truth data for a network of $N$ neurons +divided into $N_E$ excitatory and $N_I = N - N_E$ inhibitory source cells. +The network obeys Dale's principle in the natural matrix convention used by +the simulator: each source neuron has outgoing off-diagonal weights of one +sign only. +All neurons share one connectivity matrix $A_{\text{true}}$; what +differs across $M$ states is only the bias (log-baseline firing-rate) +vector $B_m$. + +\paragraph{Stage 1 (Section~\ref{sec:dale_matrices}).} +\texttt{gen\_daleMatrices3c.py} builds and saves: +\begin{itemize} + \item a sparse binary connectivity mask $E_{\text{true}} \in + \{0,1\}^{N \times N}$; + \item a single weight matrix $A_{\text{true}}$ with + $\rho(A_{\text{true}}) = R$; + \item $M$ bias vectors $B_{\text{true}} \in \mathbb{R}^{M \times N}$, + one per state. + \item static node-location arrays: coordinates, Dale type labels, and + pairwise Euclidean distances. +\end{itemize} +Stationary spikes under each $B_m$ are simulated and their +firing-rate statistics are reported. + +\paragraph{Stage 2 (Section~\ref{sec:nonstat_spikes}).} +\texttt{gen\_nonStationarySpikes3c.py} loads the Stage-1 output and: +\begin{itemize} + \item draws a target-state sequence $S_{\text{true}}[t]$ + via a Markov chain or round-robin schedule; + \item maintains a simplex coefficient vector $c_t \in \Delta^{M-1}$ + that blends the $M$ bias vectors smoothly; + \item samples $Y_t \sim \mathrm{Poisson}$ from the resulting + time-varying Poisson GLM; + \item computes the oracle state $S_{\text{oracle}}[t]$ by + maximum Poisson log-likelihood. +\end{itemize} + +\paragraph{Notation.} +$N$: total neurons; $N_E$: excitatory count (default $\lfloor 0.8N \rfloor$); +$N_I = N - N_E$: inhibitory count. +$M$: number of states (= \texttt{len(Boffsets)}). +$\Delta t$: time-bin width in seconds (default $0.01$\,s). +$T$: number of time bins. +$R$: target spectral radius (default $0.3$). +$\eta_{\mathrm{clip}} = 5$: overflow-prevention threshold +($\approx [10^{-3}, 10^3]$\,Hz effective rate range). +Matrix convention in the simulator: +$(A_{\text{true}})_{ij}$ is the lag-1 effect of source neuron \(j\) +on target neuron \(i\), because the code evaluates +\(\eta_t=A_{\text{true}}Y_{t-1}+B_t\). Therefore source neuron \(j\)'s +outgoing weights are column \(j\), and Dale's law is a column-wise +off-diagonal sign constraint. The diagonal \(A_{jj}\) is a self-history term; +it is forced non-positive for stability and is excluded from the source-type +interpretation. + +% =============================================================== +\section{Input and Output File Pattern} +\label{sec:io} +% =============================================================== + +\paragraph{Stage 1 input and output.} +\texttt{gen\_daleMatrices3c.py} takes no input files. The root directory +\[ +\texttt{} +\] +and the output subdirectory +\[ +\texttt{/truthDale} +\] +must already exist; the program asserts both paths and does not create +\texttt{truthDale}. If \texttt{--dataName} is not +provided, the output stem is generated as +\[ +\texttt{daleN\_<6hex>}. +\] +The two Stage~1 files are +\[ +\texttt{/truthDale/.simTruth.npz} +\] +and +\[ +\texttt{/truthDale/.spikes.npz}. +\] +The \texttt{.simTruth.npz} file is the required input for Stage~2. The +\texttt{.spikes.npz} file in \texttt{truthDale} contains stationary +sanity-check spikes, not the final switching dataset. +The static spatial information for the simulated nodes is also stored in +\texttt{.simTruth.npz}; it is not duplicated into the stationary +\texttt{truthDale/*.spikes.npz} file. + +\paragraph{Stage 2 input and output.} +\texttt{gen\_nonStationarySpikes3c.py} reads +\[ +\texttt{/truthDale/.simTruth.npz}. +\] +It creates \texttt{/spikesData} if needed. If +\texttt{--dataName} is not provided, the output stem is +\[ +\texttt{\_<6hex>}. +\] +The two Stage~2 files are +\[ +\texttt{/spikesData/.spikes.npz} +\] +and +\[ +\texttt{/spikesData/.prismTruth.npz}. +\] +The \texttt{spikesData/.spikes.npz} file is the input consumed by +\texttt{prism\_EM\_train3c.py} and by the EM-FDR-bagging driver. + +% =============================================================== +\section{Stage 1: Dale Matrix Generation} +\label{sec:dale_matrices} +% =============================================================== + +\subsection{Command-Line Parameters} + +\begin{table}[htbp] +\centering +\begin{tabular}{lll} +\toprule +CLI flag & Default & Description \\ +\midrule +\texttt{--num\_neurons} $N$ & 50 & Total neuron count \\ +\texttt{--num\_excite} $N_E$ & $\lfloor 0.8N \rfloor$ & + Excitatory neuron count \\ +\texttt{--edge\_prob} $[p_\ell, p_h]$ & $[0.05,\, 0.20]$ & + Per-neuron connectivity probability range \\ +\texttt{--spectral\_radius} $R$ & 0.3 & + Target spectral radius \\ +\texttt{--idleRate} $[f_{\min}, f_{\max}]$ & $[15,\, 30]$\,Hz & + Idle firing-rate band \\ +\texttt{--Boffsets} $\{\delta_m\}$ & $[0.0]$ & + Per-state log-rate offsets (one bias vector per entry) \\ +\texttt{--placement\_H\_L\_delta} & $[1.0,\,2.0,\,2.0]$ & + Node-placement height, width, and recorded distance-kernel exponent \\ +\texttt{--placement\_min\_dist} & $0.01$ & + Grid spacing and minimum inter-node separation \\ +\texttt{--num\_steps} $T$ & 10\,001 & + Steps for the stationary sanity-check simulation \\ +\texttt{--step\_size} $\Delta t$ & 0.01\,s & Time-bin width \\ +\texttt{--basePath} & \emph{(scratch dir)} & + Root output directory \\ +\bottomrule +\end{tabular} +\caption{Key command-line arguments of \texttt{gen\_daleMatrices3c.py}.} +\label{tab:dale_args} +\end{table} + +\subsection{Algorithm} + +\begin{algorithm}[H] +\caption{Dale Matrix Generation (\texttt{gen\_daleMatrices3c.py})} +\begin{algorithmic}[1] +\STATE \textbf{Step 1:} Generate static 2D node coordinates on a rectangular + placement grid without replacement + (\texttt{generate\_node\_locations}) +\STATE \textbf{Step 2:} Generate sparse binary connectivity mask + $E_{\text{true}} \in \{0,1\}^{N\times N}$ + (\texttt{generate\_sparse\_mask}) +\STATE \textbf{Step 3:} Generate weight matrix $A_{\text{true}}$ with + Dale E/I balance, negative self-loops, and spectral radius $R$ + (\texttt{init\_W}) +\STATE \textbf{Step 4:} For each offset $\delta_m$ in \texttt{Boffsets}: +\begin{ALC@g} + \STATE Generate bias vector $B_m$ (\texttt{set\_flat\_selfSpiking}) + \STATE Simulate stationary spikes $Y^{(m)}$ + (\texttt{gen\_stationary\_lag1\_poisson}) + \STATE Compute firing-rate statistics (\texttt{estimate\_rates}) +\end{ALC@g} +\STATE \textbf{Step 5:} Save ground-truth and spike files +\end{algorithmic} +\end{algorithm} + +\subsubsection{Step 1: Static Node Placement} +\label{sec:node_placement} + +The generator samples one fixed two-dimensional location for every neuron +before constructing the graph. The placement domain is a rectangle +\([0,L]\times[0,H]\), where \(L\) and \(H\) are taken from +\texttt{--placement\_H\_L\_delta}. The domain is discretized on a grid with +spacing +\[ +d_{\min}=\texttt{placement\_min\_dist}, +\] +and \(N\) grid points are sampled without replacement. Thus no two neurons +share a coordinate, and the minimum separation is at least \(d_{\min}\) in +the grid metric. The saved coordinate array is +\[ +\texttt{node\_positions}[i] = (x_i,y_i),\qquad +\texttt{node\_positions}\in\mathbb{R}^{N\times 2}. +\] +The script also saves the full Euclidean distance matrix +\[ +D_{ij} = \sqrt{(x_i-x_j)^2 + (y_i-y_j)^2} +\] +as \texttt{node\_distance\_matrix}, with zero diagonal. The third value in +\texttt{--placement\_H\_L\_delta} is recorded as +\texttt{placement\_ker\_delta} for topology provenance; in this generator it +does not alter the Dale matrix, whose connectivity mask is still controlled +by \texttt{--edge\_prob}. Finally, \texttt{node\_is\_inhibitory} records the +Dale type of each source neuron, using \(0\) for excitatory and \(1\) for +inhibitory. + +\subsubsection{Step 2: Sparse Connectivity Mask} +\label{sec:mask} + +Each neuron $i$ independently draws its connection probability: +\begin{equation} + q_i \;\sim\; \mathrm{Uniform}(p_\ell,\; p_h). +\end{equation} +The binary mask entry is +\begin{equation} + (E_{\text{true}})_{ij} = + \begin{cases} + 1 & \text{if } u_{ij} < q_i, \quad + u_{ij} \sim \mathrm{Uniform}(0,1) \text{ i.i.d.,} \\ + 0 & \text{otherwise.} + \end{cases} +\end{equation} +Self-loops are \emph{always} included: $(E_{\text{true}})_{ii} = 1$ +for all $i$, encoding the mandatory negative self-connection added below. + +\subsubsection{Step 3: Weight Matrix with Dale's Principle} +\label{sec:weights} + +\paragraph{E/I balance ratio.} +Let $r = N_E / N_I$. +This ratio is used to set the mean inhibitory weight magnitude so +that the expected signed row-sum of the full matrix is +approximately zero. + +\paragraph{Raw weight assignment (source-column based, weight spread $v = 0.2$).} +\begin{equation} + W_{ij} = + \begin{cases} + \mathrm{Uniform}(1-v,\; 1+v) & + j \in \{0,\dots,N_E-1\} \quad (\text{excitatory source columns}), \\[4pt] + \mathrm{Uniform}(-r(1+v),\; -r(1-v)) & + j \in \{N_E,\dots,N-1\} \quad (\text{inhibitory source columns}). + \end{cases} + \label{eq:raw_weights} +\end{equation} +Because the \emph{column} index $j$ determines the source neuron's type, +every nonzero off-diagonal outgoing weight from an excitatory source is +positive and every nonzero off-diagonal outgoing weight from an inhibitory +source is negative before the special diagonal rule below. This matches the +model equation directly: no transposition is needed anywhere in the +pipeline. + +\paragraph{Applying the connectivity mask.} +\begin{equation} + W \;\leftarrow\; W \odot E_{\text{true}}. +\end{equation} + +\paragraph{Enforcing negative self-loops.} +Any positive diagonal entry (excitatory self-loop) is sign-flipped: +\begin{equation} + W_{ii} \;\leftarrow\; -|W_{ii}| \quad \forall i, +\end{equation} +so that all self-connections are inhibitory (stabilising). + +\paragraph{Spectral normalization.} +Let $\rho_0 = \max_k |\lambda_k(W)|$ be the current spectral radius. +The connectivity matrix is set to +\begin{equation} + A_{\text{true}} \;\leftarrow\; \frac{R}{\rho_0}\, W, + \label{eq:spectral} +\end{equation} +so $\rho(A_{\text{true}}) = R$ exactly. +Rescaling preserves all signs, the zero pattern, Dale's law, and +the negative self-loops. +For $R < 1$ the operator $A_{\text{true}}$ is contracting, which is +a sufficient condition for the stationary Poisson VAR(1) to have a +well-defined stationary distribution. + +\subsubsection{Step 4: Bias Vectors and Stationary Spike Check} +\label{sec:bias} + +\paragraph{Bias generation (\texttt{set\_flat\_selfSpiking}).} +A size-scaling factor $s = \sqrt{N/100}$ adjusts the excitatory +correction for network size. +For each state offset $\delta_m$ (index $m = 0,\dots,M-1$), +independent draws +\begin{equation} + B_m^{(i)} \;\sim\; + \mathrm{Uniform}\!\bigl(\ln(R\,f_{\min}),\;\ln(R\,f_{\max})\bigr), + \quad i = 0,\dots,N-1, +\end{equation} +followed by state- and type-dependent shifts: +\begin{equation} + B_m^{(i)} \;\leftarrow\; + \begin{cases} + B_m^{(i)} + 1 + \delta_m - 1.5R - s + & i < N_E \quad (\text{excitatory}), \\[4pt] + B_m^{(i)} + \delta_m + & i \geq N_E \quad (\text{inhibitory}). + \end{cases} + \label{eq:bias} +\end{equation} +The excitatory correction $-(1.5R + s - 1)$ lowers the baseline +excitatory rate relative to inhibitory neurons, preventing runaway +excitation. +Larger $\delta_m$ uniformly raises activity across the network, +creating distinct mean-rate states. + +\paragraph{Stationary spike simulation (evaluation only).} +For each bias vector $B_m$ a lag-1 Poisson VAR(1) process is +simulated over $T$ bins: +\begin{equation} + \eta_t = A_{\text{true}}\,Y_{t-1} + B_m, + \qquad + \tilde{\eta}_t = \min(\eta_t,\;\eta_{\mathrm{clip}}), + \qquad + \lambda_t = \exp(\tilde{\eta}_t)\,\Delta t, + \qquad + Y_t \;\sim\; \mathrm{Poisson}(\lambda_t), + \label{eq:stationary} +\end{equation} +with initial condition $Y_0 \sim \mathrm{Poisson}(\exp(B_m)\,\Delta t)$. +These stationary spike arrays (shape $(T, N)$, one per bias vector) +are used \emph{only} to verify that the firing rates are in a +plausible range; they are \emph{not} the final non-stationary output. + +Firing-rate statistics computed via \texttt{estimate\_rates} +(\texttt{UtilDalePoisson.py}) include: +\begin{itemize} + \item mean per-neuron firing rates (Hz), + \item rate variance over non-overlapping $T_{\mathrm{var}} = 5$\,s windows (Hz$^2$), + \item Fano factors ($\mathrm{Var}[\text{count}] / \mathrm{Mean}[\text{count}]$), + \item consecutive cross-neuron coincidence rate (Hz per neuron). +\end{itemize} + +\subsection{Output Files (Stage 1)} + +All files are written to \texttt{/truthDale/}. + +\begin{table}[htbp] +\centering +\small +\begin{tabular}{@{}>{\raggedright\arraybackslash}p{0.30\linewidth} + >{\raggedright\arraybackslash}p{0.20\linewidth} + >{\raggedright\arraybackslash}p{0.14\linewidth} + >{\raggedright\arraybackslash}p{0.28\linewidth}@{}} +\toprule +File & Array key & Shape & Description \\ +\midrule +\multirow{6}{*}{\texttt{.simTruth.npz}} + & \texttt{A\_true} & $(N,N)$ float & + Shared weight matrix, $\rho(A_{\text{true}})=R$ \\ + & \texttt{B\_true} & $(M,N)$ float & + Per-state bias vectors \\ + & \texttt{E\_true} & $(N,N)$ int & + Binary connectivity mask \\ + & \texttt{node\_positions} & $(N,2)$ float & + Static node coordinates $(x,y)$ \\ + & \texttt{node\_is\_inhibitory} & $(N,)$ int & + Dale type label: 0 excitatory, 1 inhibitory \\ + & \texttt{node\_distance\_matrix} & $(N,N)$ float & + Pairwise Euclidean distances between node positions \\ +\midrule +\multirow{4}{*}{\texttt{.spikes.npz}} + & \texttt{spikes} & $(M,T,N)$ uint8 & + Stationary spikes per state (sanity check) \\ + & \texttt{single\_rates} & $(M,N)$ float & + Mean firing rates (Hz) per state \\ + & \texttt{sigle\_rates\_var} & $(M,N)$ float & + Rate variance (Hz$^2$) per state \\ + & \texttt{single\_fano\_fact} & $(M,N)$ float & + Fano factors per state \\ +\bottomrule +\end{tabular} +\caption{Output arrays of Stage~1. + $M = \texttt{len(Boffsets)}$; $T = \texttt{num\_steps}$.} +\label{tab:stage1_output} +\end{table} + +The metadata dictionary \texttt{dale\_conf} records $N$, $N_E$, $R$, +edge-probability range, idle-rate range, and offset list. +The metadata dictionary \texttt{evol\_conf} records $T$, $\Delta t$, +total simulated time, and $\eta_{\mathrm{clip}}$. + +% =============================================================== +\section{Stage 2: Non-Stationary Spike Generation} +\label{sec:nonstat_spikes} +% =============================================================== + +\subsection{Command-Line Parameters} + +\begin{table}[htbp] +\centering +\begin{tabular}{lll} +\toprule +CLI flag & Default & Description \\ +\midrule +\texttt{--inputStates} & --- & + Base name of the \texttt{.simTruth.npz} file (required) \\ +\texttt{--num\_steps} $T$ & from \texttt{evol\_conf} & + Number of time bins \\ +\texttt{--true\_dwell\_sec} $t_d$ & 1.0\,s & + Mean dwell time per state \\ +\texttt{--state\_change\_speed} $\nu$ & 0.33 & + Max simplex-coefficient change per bin \\ +\texttt{--schedule} & \texttt{mc} & + State schedule: Markov chain (\texttt{mc}) or round-robin (\texttt{rr}) \\ +\texttt{--seed} & 42 & Random seed \\ +\bottomrule +\end{tabular} +\caption{Key command-line arguments of \texttt{gen\_nonStationarySpikes3.py}.} +\label{tab:nonstat_args} +\end{table} + +\subsection{Algorithm} + +\begin{algorithm}[H] +\caption{Non-Stationary Spike Generation + (\texttt{gen\_nonStationarySpikes3.py})} +\begin{algorithmic}[1] +\STATE Load $(A_{\text{true}}, B_{\text{true}})$ from + \texttt{.simTruth.npz} +\STATE Build target-state sequence $S_{\text{true}}[t]$ + via selected schedule (Section~\ref{sec:scheduling}) +\STATE Simulate switching spikes and record simplex coefficients $C_{\text{true}}$ + (\texttt{simulate\_switching\_poisson}, Section~\ref{sec:switching}) +\STATE Compute oracle state sequence $S_{\text{oracle}}$ + (\texttt{compute\_oracle\_states\_comA}, Section~\ref{sec:oracle}) +\STATE Compute firing-rate statistics (\texttt{estimate\_rates}) +\STATE Save spike and ground-truth files +\end{algorithmic} +\end{algorithm} + +\subsection{State Scheduling} +\label{sec:scheduling} + +The mean dwell time in bins is +$\tau_d = \lceil t_d / \Delta t \rceil$. + +\paragraph{Markov chain (\texttt{mc}, default).} +A discrete-time Markov chain with geometric dwell times. +The exit probability per bin is $p_{\mathrm{exit}} = 1/\tau_d$; +on exit the chain transitions uniformly to one of the $M-1$ other states. +The symmetric $M \times M$ transition matrix stored in the output is: +\begin{equation} + T_{sm} = + \begin{cases} + 1 - p_{\mathrm{exit}} & s = m \quad (\text{self-transition}), \\ + p_{\mathrm{exit}}/(M-1) & s \neq m \quad (\text{cross-transition}). + \end{cases} + \label{eq:trans} +\end{equation} +A minimum dwell of +$\tau_{\min} = \lceil 0.3\,\tau_d \rceil$ bins is enforced +to prevent spuriously short visits. +State coverage across the full record is uncontrolled and subject +to random fluctuations. + +\paragraph{Round-robin (\texttt{rr}).} +States cycle deterministically as $0, 1, \dots, M-1, 0, 1, \dots$ +Each visit lasts exactly $\tau_d$ bins (the last visit in the +record is trimmed to fit $T$). +This guarantees that all states receive equal coverage +($T/M \pm \tau_d$ bins each), eliminating sampling variability. + +\subsection{Smooth Switching Poisson Simulation} +\label{sec:switching} + +\paragraph{Simplex coefficient vector.} +A vector $c_t \in \Delta^{M-1}$ +(non-negative entries summing to 1) tracks the current mixture of +states. +The one-hot target for the scheduled state at bin $t$ is +$e_m = \mathbf{1}[m = S_{\text{true}}[t]]$. +At each bin, $c_t$ is updated by a clipped gradient step and +renormalized: +\begin{align} + c_t &\;\leftarrow\; c_{t-1} + + \mathrm{clip}\!\bigl(e - c_{t-1},\; -\nu,\; \nu\bigr), \\ + c_t &\;\leftarrow\; c_t \;/\; \|c_t\|_1. + \label{eq:smooth_update} +\end{align} +The speed parameter $\nu \in (0,1]$ controls how quickly the mixture +tracks the target: $\nu = 1$ gives instantaneous switching; +small $\nu$ produces gradual multi-bin transitions. +The transition duration is approximately $1/\nu$ bins. + +\paragraph{Effective bias.} +The convex mixture of per-state bias vectors is +\begin{equation} + B_{\mathrm{eff},t} + = \sum_{m=0}^{M-1} c_{t,m}\, B_m + = \sum_{m=0}^{M-1} c_{t,m}\, B_{\text{true}}[m,:]. + \label{eq:Beff} +\end{equation} +This interpolates smoothly between the per-state firing-rate +baselines as the state label switches. + +\paragraph{Non-stationary spike generation.} +At every bin $t = 0, \dots, T-1$: +\begin{equation} + \eta_t = A_{\text{true}}\,Y_{t-1} + B_{\mathrm{eff},t}, + \qquad + \tilde{\eta}_t = \min(\eta_t,\;\eta_{\mathrm{clip}}), + \qquad + \lambda_t = \exp(\tilde{\eta}_t)\,\Delta t, + \qquad + Y_t \;\sim\; \mathrm{Poisson}(\lambda_t). + \label{eq:nonstat} +\end{equation} +This is the same Poisson VAR(1) model as in Eq.~\eqref{eq:stationary}, +except that the bias $B_{\mathrm{eff},t}$ is now time-varying. +The connectivity matrix $A_{\text{true}}$ is unchanged across all states. + +\subsection{Oracle State Decoder} +\label{sec:oracle} + +An oracle decoder assigns the most likely state to each bin using +the \emph{exact} ground-truth parameters $(A_{\text{true}}, \{B_m\})$: +\begin{equation} + S_{\text{oracle}}[t] + = \arg\max_{m \in \{0,\dots,M-1\}}\; + \sum_{j=0}^{N-1} + \Bigl[ + Y_{t,j}\,\tilde{\eta}_{t,j}^{(m)} + - \exp\!\bigl(\tilde{\eta}_{t,j}^{(m)}\bigr)\,\Delta t + \Bigr], + \label{eq:oracle} +\end{equation} +where +$\eta_{t,j}^{(m)} = \bigl(A_{\text{true}}\,Y_{t-1}\bigr)_j + B_{m,j}$ +and $\tilde{\eta}^{(m)} = \min(\eta^{(m)}, \eta_{\mathrm{clip}})$. +This is the exact Poisson log-likelihood score for state $m$ given +the observed spike vector $Y_t$, computed independently at each bin. +The oracle accuracy is an upper bound for any latent-state inference +algorithm. + +Per-state oracle scores and overall accuracy are printed at the end of +the run and stored in the \texttt{oracle\_eval} sub-dictionary of the +metadata. + +\subsection{Output Files (Stage 2)} + +All files are written to \texttt{/spikesData/}. + +\begin{table}[htbp] +\centering +\small +\begin{tabular}{@{}>{\raggedright\arraybackslash}p{0.30\linewidth} + >{\raggedright\arraybackslash}p{0.20\linewidth} + >{\raggedright\arraybackslash}p{0.14\linewidth} + >{\raggedright\arraybackslash}p{0.28\linewidth}@{}} +\toprule +File & Array key & Shape & Description \\ +\midrule +\multirow{2}{*}{\texttt{.spikes.npz}} + & \texttt{spikes} & $(T,N)$ int32 & + Non-stationary spike counts \\ + & \texttt{single\_rates} & $(N,)$ float & + Mean firing rates over the full record (Hz) \\ +\midrule +\multirow{8}{*}{\texttt{.prismTruth.npz}} + & \texttt{A\_true} & $(N,N)$ float & + Shared weight matrix (copy from Stage 1) \\ + & \texttt{B\_true} & $(M,N)$ float & + Per-state bias vectors (copy from Stage 1) \\ + & \texttt{S\_true} & $(T,)$ int32 & + Target state index at each bin \\ + & \texttt{C\_true} & $(T,M)$ float32 & + Simplex coefficient vector at each bin \\ + & \texttt{S\_oracle} & $(T,)$ int32 & + Oracle (max-likelihood) state at each bin \\ + & \texttt{state\_transition} & $(M,M)$ float32 & + Transition matrix used/implied by schedule \\ + & \texttt{sigle\_rates\_var} & $(N,)$ float & + Per-neuron rate variance (Hz$^2$) \\ + & \texttt{single\_fano\_fact} & $(N,)$ float & + Per-neuron Fano factors \\ +\bottomrule +\end{tabular} +\caption{Output arrays of Stage~2. + $T = \texttt{num\_steps}$; $M = \texttt{num\_states}$.} +\label{tab:stage2_output} +\end{table} + +Stage~2 copies the dynamical ground-truth arrays needed for state evaluation +(\(A_{\text{true}}\), \(B_{\text{true}}\), state labels, and oracle +quantities) into \texttt{spikesData}. Static node positions are intentionally +not copied; plotting code that needs spatial coordinates should read +\texttt{truthDale/.simTruth.npz}. + +The metadata dictionary \texttt{evol\_conf} records $T$, $\Delta t$, +total time, $\nu$, requested and effective dwell times (seconds and bins), +$M$, random seed, and schedule type. +The \texttt{oracle\_eval} sub-dictionary records the overall oracle +accuracy and per-state breakdown. + +% =============================================================== +\section{Notation Summary} +\label{sec:notation} +% =============================================================== + +\begin{table}[htbp] +\centering +\renewcommand{\arraystretch}{1.2} +\begin{tabular}{lll} +\toprule +Symbol & Domain & Description \\ +\midrule +$N$ & $\mathbb{Z}_{>0}$ & Total number of neurons \\ +$N_E$ & $\mathbb{Z}_{>0}$ & Excitatory neuron count (rows $0\dots N_E-1$) \\ +$N_I = N-N_E$ & $\mathbb{Z}_{>0}$ & Inhibitory neuron count \\ +$M$ & $\mathbb{Z}_{>0}$ & + Number of states ($= \texttt{len(Boffsets)}$) \\ +$T$ & $\mathbb{Z}_{>0}$ & Number of time bins \\ +$\Delta t$ & $\mathbb{R}_{>0}$ & Time-bin width (s); default $0.01$ \\ +$R$ & $(0,1)$ & Target spectral radius \\ +$r = N_E/N_I$ & $\mathbb{R}_{>0}$ & E/I balance ratio \\ +$v = 0.2$ & --- & Weight spread (hard-coded) \\ +$E_{\text{true}}$ & $\{0,1\}^{N\times N}$ & Binary connectivity mask \\ +$A_{\text{true}}$ & $\mathbb{R}^{N\times N}$ & + Shared weight matrix, $\rho(A_{\text{true}}) = R$ \\ +$B_m$ & $\mathbb{R}^N$ & Bias (log-baseline rate) vector for state $m$ \\ +$B_{\text{true}}$ & $\mathbb{R}^{M\times N}$ & Stacked per-state biases \\ +$Y_t$ & $\mathbb{Z}_{\geq 0}^N$ & Spike-count vector at bin $t$ \\ +$\eta_{\mathrm{clip}}$ & $\mathbb{R}_{>0}$ & + Rate clipping threshold; fixed at $5$ \\ +\midrule +$S_{\text{true}}[t]$ & $\{0,\dots,M-1\}$ & + Scheduled target state at bin $t$ \\ +$c_t$ & $\Delta^{M-1}$ & + Simplex coefficient vector at bin $t$ \\ +$\nu$ & $(0,1]$ & + Max coefficient change per bin (\texttt{state\_change\_speed}) \\ +$B_{\mathrm{eff},t}$ & $\mathbb{R}^N$ & + Effective bias: convex mixture $\sum_m c_{t,m} B_m$ \\ +$\tau_d$ & $\mathbb{Z}_{>0}$ & + Mean dwell time in bins \\ +$S_{\text{oracle}}[t]$ & $\{0,\dots,M-1\}$ & + Oracle max-likelihood state at bin $t$ \\ +\bottomrule +\end{tabular} +\caption{Notation used throughout this document.} +\label{tab:notation} +\end{table} + +% =============================================================== +\section{Full Pipeline and Suggested Commands} +\label{sec:pipeline} +% =============================================================== + +\begin{enumerate} + +\item \textbf{Generate ground-truth matrices} (Stage 1): +\begin{verbatim} +mkdir -p $basePath/truthDale +./gen_daleMatrices3c.py --num_neurons 100 --num_excite 70 \ + --spectral_radius 0.5 --Boffsets 0 10 20 --dataName myData \ + --basePath $basePath +\end{verbatim} +Produces \texttt{myData.simTruth.npz} ($A_{\text{true}}$, +$B_{\text{true}}$, $E_{\text{true}}$, \texttt{node\_positions}, +\texttt{node\_is\_inhibitory}, and \texttt{node\_distance\_matrix}; +$M=3$ states here) and +\texttt{myData.spikes.npz} (stationary sanity-check arrays, +shape $(3, T, N)$), both under \texttt{\$basePath/truthDale/}. + +\item \textbf{Generate non-stationary spike train} (Stage 2): +\begin{verbatim} +./gen_nonStationarySpikes3c.py --basePath $basePath \ + --inputStates myData --dataName myData_sw \ + --true_dwell_sec 1.0 --schedule mc +\end{verbatim} +Produces \texttt{myData\_sw.spikes.npz} (shape $(T, N)$) and +\texttt{myData\_sw.prismTruth.npz} ($S_{\text{true}}$, +$C_{\text{true}}$, $S_{\text{oracle}}$, transition matrix, and +copies of $A_{\text{true}}$, $B_{\text{true}}$), both under +\texttt{\$basePath/spikesData/}. Static node coordinates remain in +\texttt{\$basePath/truthDale/myData.simTruth.npz}. + +\item \textbf{Fit with PRISM-EM}: +\begin{verbatim} +torchrun --standalone --nnodes=1 --nproc_per_node=4 \ + ./prism_EM_train3c.py \ + --basePath $basePath --dataName myData_sw \ + --num_states 3 --num_em_iters 4 --m_epochs 16 \ + --time_range_sec 0 80 +\end{verbatim} +Writes \texttt{\$basePath/prismFit/.prismEM.npz}, where +\texttt{} is either supplied by \texttt{--fitName} or generated as +\texttt{myData\_sw-EM-<6hex>}. + +\end{enumerate} + +% =============================================================== +\section{Key Design Choices} +\label{sec:design} +% =============================================================== + +\paragraph{Shared $A_{\text{true}}$ across all states.} +The synaptic weight matrix is identical in every state; only the +bias vector $B_m$ differs between states. +This models a circuit whose connectivity is fixed but whose +input drive (e.g.\ neuromodulatory tone, sensory stimulus level) +shifts the population-wide activity baseline between states. +The inference task is therefore to recover both $A_{\text{true}}$ +and $\{B_m\}$ from a single non-stationary spike record. + +\paragraph{Source-column Dale convention.} +The simulator uses \(A_{ij}\) as source \(j\) into target \(i\). Therefore +column \(j\) contains source neuron \(j\)'s outgoing effects. The generator +assigns signs by column: columns \(0 \dots N_E - 1\) have positive +off-diagonal entries and columns \(N_E \dots N-1\) have negative +off-diagonal entries. Diagonal entries are forced non-positive as stabilizing +self-history terms and are not part of the excitatory/inhibitory source +classification. This is the same convention used by PRISM-EM and by the +EM-FDR aggregate when they infer source neuron type from fitted columns. + +\paragraph{Spectral radius and stability.} +$\rho(A_{\text{true}}) = R < 1$ is a sufficient condition for +the stationary Poisson VAR(1) to admit a well-defined stationary +distribution. +In the non-stationary regime, stability is maintained instantaneously +because $A_{\text{true}}$ is unchanged; only the bias shifts. + +\paragraph{Smooth simplex interpolation.} +The clipped gradient update (Eq.~\eqref{eq:smooth_update}) prevents +abrupt discontinuities in $B_{\mathrm{eff},t}$ at state-boundary +bins. +This mimics a biological circuit that cannot switch network state +instantaneously. +The mixing speed $\nu$ is independent of the mean dwell time and +can be set to match a desired biophysical transition timescale. + +\paragraph{Oracle score as a decoder upper bound.} +The oracle uses exact ground-truth parameters and per-bin +maximum-likelihood, achieving the best possible state assignment +from the spike observations. +Any EM or variational algorithm must approach, but cannot exceed, +this score. +The oracle score therefore quantifies how much information about +the latent state is present in the spikes at all, distinguishing +fundamental ambiguity from algorithm suboptimality. + +\end{document} diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/preproc_exper_ver3c.pdf b/causal_net/nonStation_ver3c_states_EM_FDR/docs/preproc_exper_ver3c.pdf new file mode 100644 index 00000000..6da107b9 Binary files /dev/null and b/causal_net/nonStation_ver3c_states_EM_FDR/docs/preproc_exper_ver3c.pdf differ diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/docs/preproc_exper_ver3c.tex b/causal_net/nonStation_ver3c_states_EM_FDR/docs/preproc_exper_ver3c.tex new file mode 100644 index 00000000..a7b286c5 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/docs/preproc_exper_ver3c.tex @@ -0,0 +1,295 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} + +\geometry{margin=0.8in} + +\title{Preprocessing Experimental Spike Recordings for PRISM-EM\\ +\large Ver3c biological-data input preparation} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +This document describes the current \texttt{prep\_bioexp3c.py} +preprocessing program for biological spike recordings. The script reads +raw experimental spike times and a per-channel metrics spreadsheet, +converts spike times into a binned spike-count matrix, filters channels by +firing-rate range, orders accepted neurons by firing rate, and writes two +companion files. The first file, \texttt{.spikes.npz}, contains only the +data needed by the EM fitter. The second file, \texttt{.bioExp.npz}, +contains experimental metadata, MEA identifiers, metrics-table rows, and +other biological bookkeeping used by inspection, plotting, and downstream +interpretation. +\end{abstract} + +\tableofcontents +\newpage + +% =============================================================== +\section{Purpose and Scope} +\label{sec:scope} +% =============================================================== + +\texttt{prep\_bioexp3c.py} converts one experimental recording session into +the same spike-file contract used elsewhere in the ver3c pipeline: +rows are time bins and columns are neurons. The preprocessing script does +not fit a model and does not estimate connectivity. Its job is to make the +experimental recording usable by the EM fitter while keeping +experiment-specific information outside the fitter input. + +The key design choice is a two-file output: +\begin{itemize} + \item \texttt{.spikes.npz}: compact fitter input; + \item \texttt{.bioExp.npz}: biological/session metadata and + channel annotations. +\end{itemize} +The files share the same stem and are written to the same directory. + +% =============================================================== +\section{Command-Line Inputs} +\label{sec:cli} +% =============================================================== + +The command-line arguments are: +\begin{itemize} + \item \texttt{--expPath}: root directory containing the raw experimental + session directory; + \item \texttt{--sessionName}: session path below \texttt{--expPath}; + \item \texttt{--dataPath}: output directory for both generated + \texttt{.npz} files; + \item \texttt{--shortName}: optional output stem. If omitted, the script + uses \texttt{-<6hex>}; + \item \texttt{--inpFormat}: metrics-spreadsheet selector, + \texttt{1} for \texttt{metrics\_curated.xlsx} and + \texttt{2} for \texttt{quality\_metrics.xlsx}; + \item \texttt{--samp\_freq}: binning frequency in Hz, default + \texttt{100}; + \item \texttt{--freqRange}: inclusive accepted firing-rate range in Hz, + default \texttt{1 50}; + \item \texttt{--verb}: diagnostic verbosity. +\end{itemize} + +The raw session is expected to contain +\[ +\texttt{//spike\_times.npy} +\] +and one metrics spreadsheet: +\[ +\texttt{metrics\_curated.xlsx} +\quad\hbox{or}\quad +\texttt{quality\_metrics.xlsx}. +\] + +The session name is parsed as a slash-separated path whose fields include +cell line, recording date, chip ID, run number, and well number. These +values are stored in the biological metadata. + +% =============================================================== +\section{Preprocessing Steps} +\label{sec:steps} +% =============================================================== + +\subsection{Read Raw Spike Times} + +The input \texttt{spike\_times.npy} is loaded as a Python dictionary keyed +by MEA feature/channel identifiers. The keys are sorted once, and this +sorted key order defines the initial raw channel order. + +For each channel, spike times are multiplied by \texttt{--samp\_freq} and +cast to integer bin indices. Thus \texttt{--samp\_freq=100} corresponds to +\(\Delta t=0.01\,\mathrm{s}\). Repeated bin indices are retained and later +become spike counts larger than one in a bin. + +The script computes: +\begin{itemize} + \item the total number of time bins; + \item the recording duration in seconds; + \item one firing rate per raw channel. +\end{itemize} + +\subsection{Filter by Firing Rate} + +Channels are retained if their firing rate lies inside the inclusive range +\[ +\texttt{freqRange[0]} \le r_i \le \texttt{freqRange[1]}. +\] +The metadata records how many channels were dropped below and above this +range. The remaining channels define the population used downstream. + +\subsection{Sort Accepted Neurons by Firing Rate} + +After rate filtering, accepted channels are sorted by firing rate in +ascending order. This frequency-sorted order becomes the primary neuron +index used in the saved spike matrix. The \texttt{.bioExp.npz} file stores +both the frequency-sort index and the inverse map so that downstream tools +can relate the saved columns back to the accepted-channel order. + +\subsection{Build the Spike-Count Matrix} + +For each accepted, frequency-sorted channel, the script forms a +length-\(T\) spike-count vector using \texttt{numpy.bincount}. These vectors +are stacked into +\[ +Y \in \mathbb{N}_0^{T\times N}, +\] +where rows are time bins and columns are accepted neurons. + +The maximum pre-clipping count in any bin is recorded in the biological +metadata. The matrix written to the fitter-facing file is clipped to +\([0,255]\) and stored as \texttt{uint8}. This keeps the spike file compact +while preserving ordinary binned spike counts. + +\subsection{Load and Align Metrics} + +The metrics spreadsheet is loaded with \texttt{pandas}. For input format 1, +the file name is \texttt{metrics\_curated.xlsx}. For input format 2, the +file name is \texttt{quality\_metrics.xlsx}; location columns named +\texttt{location\_X}, \texttt{location\_Y}, and \texttt{location\_Z} are +renamed to \texttt{loc\_x}, \texttt{loc\_y}, and \texttt{loc\_z} when +present. + +Rows are matched by \texttt{MEA\_idx} and reordered to match the final +frequency-sorted neuron order. Duplicate \texttt{MEA\_idx} rows are treated +as an error, and every retained channel must have a corresponding metrics +row. + +\subsection{Compute Summary Statistics} + +The script computes summary statistics from the accepted channels: +mean, standard deviation, median, minimum, and maximum firing rate, plus +the mean and standard deviation of the per-neuron Fano factor. These values +are stored in the biological metadata. + +% =============================================================== +\section{Output File Pattern} +\label{sec:files} +% =============================================================== + +Let \texttt{} denote the output stem. It is equal to +\texttt{--shortName} when that argument is supplied; otherwise it is +generated as +\[ +\texttt{-<6hex>}. +\] + +The script writes both output files directly into \texttt{--dataPath}: +\[ +\texttt{/.bioExp.npz}, +\qquad +\texttt{/.spikes.npz}. +\] + +The order is deliberate: the \texttt{.spikes.npz} file is the direct input +to numerical fitters, while the \texttt{.bioExp.npz} file stores +experimental context that should not be required by the EM training code. + +% =============================================================== +\section{Saved Contents} +\label{sec:saved} +% =============================================================== + +\subsection{\texttt{.spikes.npz}} + +The fitter-facing data dictionary contains: +\begin{itemize} + \item \texttt{spikes}: \texttt{uint8} array with shape \((T,N)\); + \item \texttt{single\_rates}: floating-point array with shape \((N,)\), + ordered exactly like the spike-matrix columns. +\end{itemize} + +The metadata contains: +\begin{itemize} + \item \texttt{time\_step\_sec}: bin width, equal to + \(1/\texttt{samp\_freq}\); + \item \texttt{provenance.experiment\_name}: the output stem + \texttt{}; + \item \texttt{poisson\_eta\_clip}: clipping value expected by the EM + fitter, currently \texttt{5}; + \item \texttt{data\_type}: \texttt{bioExp}; + \item \texttt{num\_neurons}: number of accepted channels \(N\). +\end{itemize} + +\subsection{\texttt{.bioExp.npz}} + +The biological data dictionary contains: +\begin{itemize} + \item \texttt{neur\_freqIdx}: index vector that sorts accepted channels + by firing rate; + \item \texttt{neur\_revFreqIdx}: inverse map from accepted-channel order + to frequency-sorted order; + \item \texttt{single\_rates}: firing rates in the final saved neuron + order; + \item \texttt{MEA\_idx}: MEA identifiers reordered to match the final + neuron order; + \item \texttt{spike\_key}: raw \texttt{spike\_times.npy} keys reordered + to match the final neuron order; + \item \texttt{metrics\_curated}: metrics-spreadsheet rows reordered to + match the final neuron order. +\end{itemize} + +The biological metadata contains: +\begin{itemize} + \item \texttt{bioexp}: raw-data provenance such as + \texttt{raw\_bioexp\_path}, \texttt{session\_name}, + \texttt{inp\_format}, \texttt{cell\_line\_name}, + \texttt{recording\_date}, \texttt{chip\_ID}, \texttt{run\_num}, + \texttt{well\_num}, \texttt{sampling\_freq}, + \texttt{num\_feature}, \texttt{num\_time\_bin}, and + \texttt{max\_time}; + \item \texttt{data\_selector}: selection parameters and results, + including \texttt{freq\_range}, + \texttt{num\_drop\_neur\_lo\_hi\_freq}, \texttt{num\_chan}, and + \texttt{max\_spike\_per\_bin}; + \item \texttt{hash}: random six-character hash used when + \texttt{--shortName} is omitted; + \item \texttt{short\_name}: output stem \texttt{}; + \item \texttt{metrics\_curated\_columns}: spreadsheet column names after + any input-format normalization; + \item \texttt{rate\_summary}: firing-rate and Fano-factor summary + statistics for the accepted population. +\end{itemize} + +% =============================================================== +\section{Example Commands} +\label{sec:commands} +% =============================================================== + +For inspection or plotting only, the output directory can be any directory +chosen by \texttt{--dataPath}. For use with the ver3c EM fitter, choose the +same \texttt{spikesData} layout used by \texttt{fitEM\_states\_ver3c.tex}: +\begin{verbatim} +basePath=/path/to/analysis + +./prep_bioexp3c.py \ + --expPath /path/to/raw/experiments \ + --sessionName ///// \ + --dataPath $basePath/spikesData \ + --shortName \ + --inpFormat 1 \ + --samp_freq 100 \ + --freqRange 1 50 +\end{verbatim} + +This produces: +\[ +\texttt{\$basePath/spikesData/.spikes.npz}, +\qquad +\texttt{\$basePath/spikesData/.bioExp.npz}. +\] + +The \texttt{.spikes.npz} file is then the input selected by +\texttt{--basePath \$basePath --dataName } in the EM training +workflow. See \texttt{fitEM\_states\_ver3c.tex} for the EM model, fitting +options, and fit-output file pattern. + +\end{document} diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/edgeMeterAccuracy3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/edgeMeterAccuracy3c.py new file mode 100755 index 00000000..b8f24dba --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/edgeMeterAccuracy3c.py @@ -0,0 +1,482 @@ +#!/usr/bin/env python3 +"""Absolute graph-recovery metrics for final prism EM-FDR aggregate fits.""" + +import argparse +import os +import secrets +import time + +import numpy as np + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + + +LABEL_ORDER = np.asarray([-1, 0, 1], dtype=np.int8) +SOURCE_RECO_ORDER = np.asarray([1, -1, 0], dtype=np.int8) +SOURCE_RECO_NAMES = np.asarray(["exc", "inh", "und"], dtype=object) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Compute absolute graph-recovery metrics for synthetic prism EM-FDR aggregate fits", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--basePath", required=True, + help="Run directory containing prismFit/, truthDale/, edgeReco/, plots/") + parser.add_argument("--fitNameTrunk", required=True, + help="Common prefix of aggregate fit files") + parser.add_argument("--fitTags", nargs="+", required=True, + help="Fit suffixes, e.g. 'em103c_fdr103c_agr103c emfaec_fdrfaec_agrfaec'") + parser.add_argument("--outName", default=argparse.SUPPRESS, + help="Output metric stem in edgeReco/; default is fitNameTrunk_ermHASH4") + parser.add_argument("--epsilon", type=float, default=1e-4, + help="Presence tolerance for A_prune") + parser.add_argument("-p", "--showPlots", nargs="+", default="abc", + help="Plot letters/groups, e.g. -p a b") + parser.add_argument("-X", "--noXterm", action="store_true", + help="Disable X terminal for plotting") + parser.add_argument("-v", "--verb", type=int, default=1, + help="Verbosity level") + args = parser.parse_args() + if not hasattr(args, "outName"): + args.outName = f"{args.fitNameTrunk}_erm{secrets.token_hex(2)}" + return args + + +def flatten_words(items): + words = [] + for item in items: + words.extend(str(item).split()) + return words + + +def normalize_plot_letters(items): + letters = [] + for token in flatten_words(items): + letters.extend(list(token)) + return letters + + +def full_fit_name(trunk, tag): + tag = str(tag) + if tag.startswith(trunk): + return tag + return f"{trunk}_{tag}" + + +def fit_file(base_path, fit_name): + return os.path.join(base_path, "prismFit", f"{fit_name}.prismEM.npz") + + +def truth_file(base_path, truth_name): + return os.path.join(base_path, "truthDale", f"{truth_name}.simTruth.npz") + + +def load_fit(base_path, fit_name, verb=1): + inp_f = fit_file(base_path, fit_name) + fit_d, fit_md = read_data_npz(inp_f, verb=verb > 1) + if not isinstance(fit_md, dict): + raise ValueError(f"Fit file has no metadata: {inp_f}") + if "A_prune" not in fit_d: + raise KeyError(f"{inp_f} must contain A_prune") + if "neuron_type" not in fit_d: + raise KeyError(f"{inp_f} must contain neuron_type") + if "provenance" not in fit_md or "state_model_file" not in fit_md["provenance"]: + raise KeyError(f"{inp_f} metadata must contain provenance.state_model_file") + return fit_d, fit_md, inp_f + + +def load_truth_once(base_path, truth_name, verb=1): + inp_f = truth_file(base_path, truth_name) + truth_d, truth_md = read_data_npz(inp_f, verb=verb > 1) + for key in ("A_true", "E_true", "node_is_inhibitory"): + if key not in truth_d: + raise KeyError(f"{inp_f} must contain {key}; regenerate synthetic truth") + return truth_d, truth_md, inp_f + + +def as_static_2d(arr, name): + arr = np.asarray(arr) + if arr.ndim != 2 or arr.shape[0] != arr.shape[1]: + raise ValueError(f"{name} must be a square 2D matrix, got {arr.shape}") + return arr + + +def safe_div(num, den): + den = float(den) + if den == 0.0: + return float("nan") + return float(num) / den + + +def safe_f1(precision, recall): + if not np.isfinite(precision) or not np.isfinite(recall): + return float("nan") + den = precision + recall + if den == 0.0: + return float("nan") + return float(2.0 * precision * recall / den) + + +def matthews_corrcoef(tp, fp, fn, tn): + den = (tp + fp) * (tp + fn) * (tn + fp) * (tn + fn) + if den <= 0: + return float("nan") + return float((tp * tn - fp * fn) / np.sqrt(den)) + + +def graph_labels(A_true, E_true, A_prune, epsilon): + A_true = as_static_2d(A_true, "A_true").astype(np.float64) + E_true = as_static_2d(E_true, "E_true") != 0 + A_prune = as_static_2d(A_prune, "A_prune").astype(np.float64) + if A_prune.shape != A_true.shape or E_true.shape != A_true.shape: + raise ValueError( + f"Shape mismatch: A_true={A_true.shape} E_true={E_true.shape} A_prune={A_prune.shape}" + ) + + n = A_true.shape[0] + off_mask = ~np.eye(n, dtype=bool) + true_edge = E_true & off_mask + true_edge_sign = A_true[true_edge] + if np.any(true_edge_sign == 0.0): + raise ValueError("Every off-diagonal E_true edge must have nonzero A_true sign") + + true_lab_2d = np.zeros((n, n), dtype=np.int8) + true_lab_2d[true_edge & (A_true > 0.0)] = 1 + true_lab_2d[true_edge & (A_true < 0.0)] = -1 + + est_lab_2d = np.zeros((n, n), dtype=np.int8) + est_lab_2d[off_mask & (A_prune > float(epsilon))] = 1 + est_lab_2d[off_mask & (A_prune < -float(epsilon))] = -1 + + return true_lab_2d[off_mask], est_lab_2d[off_mask], true_lab_2d, est_lab_2d, off_mask + + +def edge_confusion_3x3(true_lab, est_lab): + cm = np.zeros((3, 3), dtype=np.int64) + idx = {-1: 0, 0: 1, 1: 2} + for t, e in zip(true_lab, est_lab): + cm[idx[int(t)], idx[int(e)]] += 1 + return cm + + +def multiclass_mcc(confusion): + """Matthews correlation coefficient for a multiclass confusion matrix.""" + cm = np.asarray(confusion, dtype=np.float64) + n = float(np.sum(cm)) + if n <= 0.0: + return float("nan") + c = float(np.trace(cm)) + row_sum = np.sum(cm, axis=1) + col_sum = np.sum(cm, axis=0) + num = c * n - float(np.dot(row_sum, col_sum)) + den = np.sqrt( + (n * n - float(np.dot(col_sum, col_sum))) * + (n * n - float(np.dot(row_sum, row_sum))) + ) + if den == 0.0: + return float("nan") + return float(num / den) + + +def presence_metrics(true_lab, est_lab): + t = true_lab != 0 + e = est_lab != 0 + tp = int(np.count_nonzero(t & e)) + fp = int(np.count_nonzero(~t & e)) + fn = int(np.count_nonzero(t & ~e)) + tn = int(np.count_nonzero(~t & ~e)) + precision = safe_div(tp, tp + fp) + recall = safe_div(tp, tp + fn) + return { + "TP": tp, + "FP": fp, + "FN": fn, + "TN": tn, + "precision": precision, + "recall": recall, + "f1": safe_f1(precision, recall), + "mcc": matthews_corrcoef(tp, fp, fn, tn), + "fdp": safe_div(fp, tp + fp), + "num_true_edges": int(tp + fn), + "num_reco_edges": int(tp + fp), + "num_candidates": int(true_lab.size), + "density_true": safe_div(tp + fn, true_lab.size), + "density_reco": safe_div(tp + fp, true_lab.size), + } + + +def signed_metrics(true_lab, est_lab): + true_present = true_lab != 0 + est_present = est_lab != 0 + correct = true_present & est_present & (true_lab == est_lab) + sign_flip = true_present & est_present & (true_lab != est_lab) + spurious = ~true_present & est_present + missed = true_present & ~est_present + + signed_tp = int(np.count_nonzero(correct)) + signed_fp = int(np.count_nonzero(spurious) + np.count_nonzero(sign_flip)) + signed_fn = int(np.count_nonzero(missed) + np.count_nonzero(sign_flip)) + recovered_true = int(np.count_nonzero(true_present & est_present)) + n_flip = int(np.count_nonzero(sign_flip)) + precision = safe_div(signed_tp, signed_tp + signed_fp) + recall = safe_div(signed_tp, signed_tp + signed_fn) + return { + "signed_TP": signed_tp, + "signed_FP": signed_fp, + "signed_FN": signed_fn, + "sign_flip_count": n_flip, + "signed_precision": precision, + "signed_recall": recall, + "signed_f1": safe_f1(precision, recall), + "sign_flip_rate": safe_div(n_flip, recovered_true), + } + + +def true_source_type_from_truth(truth_d, A_true, E_true): + inhib = np.asarray(truth_d["node_is_inhibitory"]).astype(bool) + n = A_true.shape[0] + if inhib.shape != (n,): + raise ValueError(f"node_is_inhibitory shape {inhib.shape} incompatible with N={n}") + true_type = np.where(inhib, -1, 1).astype(np.int8) + + off_mask = ~np.eye(n, dtype=bool) + E_true = E_true != 0 + for j in range(n): + vals = A_true[off_mask[:, j] & E_true[:, j], j] + if vals.size == 0: + continue + if true_type[j] == 1 and np.any(vals < 0.0): + raise ValueError(f"Truth source column {j} marked excitatory but has negative true edges") + if true_type[j] == -1 and np.any(vals > 0.0): + raise ValueError(f"Truth source column {j} marked inhibitory but has positive true edges") + return true_type + + +def source_type_metrics(true_source_type, reco_source_type, true_lab_2d, est_lab_2d, off_mask): + reco_source_type = np.asarray(reco_source_type, dtype=np.int8) + n = true_source_type.shape[0] + if reco_source_type.shape != (n,): + raise ValueError(f"neuron_type shape {reco_source_type.shape} incompatible with N={n}") + + counts = { + "num_reco_exc_neurons": int(np.count_nonzero(reco_source_type == 1)), + "num_reco_inh_neurons": int(np.count_nonzero(reco_source_type == -1)), + "num_reco_und_neurons": int(np.count_nonzero(reco_source_type == 0)), + "frac_reco_und_neurons": safe_div(np.count_nonzero(reco_source_type == 0), n), + } + + source_conf = np.zeros((2, 3), dtype=np.int64) + row_idx = {1: 0, -1: 1} + col_idx = {1: 0, -1: 1, 0: 2} + for j in range(n): + source_conf[row_idx[int(true_source_type[j])], col_idx[int(reco_source_type[j])]] += 1 + + src_idx = np.broadcast_to(np.arange(n, dtype=np.int64)[None, :], (n, n)) + true_edge = (true_lab_2d != 0) & off_mask + reco_edge = (est_lab_2d != 0) & off_mask + true_und = true_edge & (reco_source_type[src_idx] == 0) + reco_und = reco_edge & (reco_source_type[src_idx] == 0) + counts.update({ + "frac_true_edges_from_reco_und_src": safe_div(np.count_nonzero(true_und), np.count_nonzero(true_edge)), + "frac_reco_edges_from_reco_und_src": safe_div(np.count_nonzero(reco_und), np.count_nonzero(reco_edge)), + "source_type_confusion": source_conf, + }) + return counts + + +def real_fit_metadata(md): + if "bagsFDR_stageA" not in md: + raise KeyError("fit metadata missing 'bagsFDR_stageA'; was this file produced by prism_EM_FDR_Bags_aggregate3c.py?") + if "real_fit" not in md["bagsFDR_stageA"]: + raise KeyError("fit metadata missing 'bagsFDR_stageA.real_fit'") + return md["bagsFDR_stageA"]["real_fit"] + + +def metric_row(base_path, fit_name, fit_tag, fit_d, fit_md, truth_d, truth_md, truth_f, epsilon): + A_true = as_static_2d(truth_d["A_true"], "A_true").astype(np.float64) + E_true = as_static_2d(truth_d["E_true"], "E_true") != 0 + A_prune = as_static_2d(fit_d["A_prune"], "A_prune").astype(np.float64) + true_lab, est_lab, true_lab_2d, est_lab_2d, off_mask = graph_labels( + A_true, E_true, A_prune, epsilon + ) + pmet = presence_metrics(true_lab, est_lab) + smet = signed_metrics(true_lab, est_lab) + cm3 = edge_confusion_3x3(true_lab, est_lab) + smet["signed_mcc"] = multiclass_mcc(cm3) + + true_source_type = true_source_type_from_truth(truth_d, A_true, E_true) + src_met = source_type_metrics( + true_source_type, + fit_d["neuron_type"], + true_lab_2d, + est_lab_2d, + off_mask, + ) + + train_md = real_fit_metadata(fit_md)["train"] + stage_a = fit_md["bagsFDR_stageA"] + stage_b = fit_md["bagsFDR_stageB"] + + time_range_sec = np.asarray(train_md["time_range_sec"], dtype=np.float64) + duration_sec = float(time_range_sec[1] - time_range_sec[0]) + num_reco_edges = int(pmet["num_reco_edges"]) + fdr_bound = float(stage_b["stability_false_edge_bound"]) + + row = { + "fit_name": fit_name, + "fit_tag": fit_tag, + "fit_file": fit_file(base_path, fit_name), + "truth_file": truth_f, + "num_neurons": int(A_true.shape[0]), + "num_time_bins": int(train_md["num_time_bins"]), + "time_range_sec": time_range_sec.astype(np.float64), + "duration_sec": duration_sec, + "duration_min": duration_sec / 60.0, + "num_bags": int(stage_b["num_bags"]), + "bag_frac": float(stage_a["bag_frac"]), + "per_bag_quantile": float(stage_b["per_bag_quantile"]), + "stab_sel_thresh": float(stage_b["stab_sel_thresh"]), + "epsilon": float(epsilon), + "stability_false_edge_bound": fdr_bound, + "stability_bound_fdp": safe_div(fdr_bound, num_reco_edges), + "confusion3": cm3, + **pmet, + **smet, + **src_met, + } + return row + + +def assemble_metric_arrays(rows): + keys_float = ( + "duration_sec", "duration_min", "bag_frac", "per_bag_quantile", + "stab_sel_thresh", "epsilon", "stability_false_edge_bound", + "stability_bound_fdp", "precision", "recall", "f1", "mcc", "fdp", + "density_true", "density_reco", "signed_precision", "signed_recall", + "signed_f1", "signed_mcc", "sign_flip_rate", "frac_reco_und_neurons", + "frac_true_edges_from_reco_und_src", "frac_reco_edges_from_reco_und_src", + ) + keys_int = ( + "num_neurons", "num_time_bins", "num_bags", "TP", "FP", "FN", "TN", + "num_true_edges", "num_reco_edges", "num_candidates", "signed_TP", + "signed_FP", "signed_FN", "sign_flip_count", "num_reco_exc_neurons", + "num_reco_inh_neurons", "num_reco_und_neurons", + ) + + out = { + "fit_name": np.asarray([r["fit_name"] for r in rows], dtype=object), + "fit_tag": np.asarray([r["fit_tag"] for r in rows], dtype=object), + "fit_file": np.asarray([r["fit_file"] for r in rows], dtype=object), + "truth_file": np.asarray([r["truth_file"] for r in rows], dtype=object), + "time_range_sec": np.stack([r["time_range_sec"] for r in rows], axis=0).astype(np.float64), + "confusion3": np.stack([r["confusion3"] for r in rows], axis=0).astype(np.int64), + "source_type_confusion": np.stack([r["source_type_confusion"] for r in rows], axis=0).astype(np.int64), + "label_order": LABEL_ORDER.copy(), + "source_truth_order": np.asarray(["exc", "inh"], dtype=object), + "source_reco_order": SOURCE_RECO_NAMES.copy(), + } + out["confusion3_sum"] = np.sum(out["confusion3"], axis=0).astype(np.int64) + out["source_type_confusion_sum"] = np.sum(out["source_type_confusion"], axis=0).astype(np.int64) + for key in keys_float: + out[key] = np.asarray([r[key] for r in rows], dtype=np.float64) + for key in keys_int: + out[key] = np.asarray([r[key] for r in rows], dtype=np.int64) + return out + + +def main(): + args = parse_args() + if args.epsilon < 0.0: + raise ValueError("--epsilon must be non-negative") + + t0 = time.perf_counter() + fit_tags = flatten_words(args.fitTags) + if not fit_tags: + raise ValueError("--fitTags produced an empty list") + + out_name = args.outName + edge_dir = os.path.join(args.basePath, "edgeReco") + plot_dir = os.path.join(args.basePath, "plots") + os.makedirs(edge_dir, exist_ok=True) + os.makedirs(plot_dir, exist_ok=True) + + rows = [] + truth_d = None + truth_md = None + truth_f = None + truth_name_ref = None + fit_names = [full_fit_name(args.fitNameTrunk, tag) for tag in fit_tags] + + for fit_tag, fit_name in zip(fit_tags, fit_names): + fit_d, fit_md, inp_f = load_fit(args.basePath, fit_name, verb=args.verb) + truth_name = fit_md["provenance"]["state_model_file"] + if truth_name_ref is None: + truth_name_ref = truth_name + truth_d, truth_md, truth_f = load_truth_once(args.basePath, truth_name, verb=args.verb) + elif truth_name != truth_name_ref: + raise ValueError( + f"All fits must use the same truth. First={truth_name_ref}, {fit_name}={truth_name}" + ) + row = metric_row( + args.basePath, fit_name, fit_tag, + fit_d, fit_md, truth_d, truth_md, truth_f, + epsilon=args.epsilon, + ) + rows.append(row) + if args.verb > 0: + print( + f"metric {fit_tag}: TP={row['TP']} FP={row['FP']} " + f"FN={row['FN']} precision={row['precision']:.4f} " + f"recall={row['recall']:.4f} MCC={row['mcc']:.4f}" + ) + + out_d = assemble_metric_arrays(rows) + out_md = { + "program": "edgeMeterAccuracy3c.py", + "fitNameTrunk": args.fitNameTrunk, + "fitTags": fit_tags, + "fit_names": fit_names, + "truth_name": truth_name_ref, + "truth_file": truth_f, + "metric_target": "A_prune", + "candidate_universe": "all_off_diagonal_directed_pairs", + "edge_presence_rule": "A_prune != 0 using epsilon tolerance", + "truth_presence": "E_true != 0 on off-diagonal", + "truth_sign": "sign(A_true) on E_true edges", + "source_type_truth": "node_is_inhibitory from simTruth", + "source_type_reco": "neuron_type from aggregate fit; 0 is undecided abstention", + "epsilon": float(args.epsilon), + "label_order": LABEL_ORDER.tolist(), + "source_truth_order": ["exc", "inh"], + "source_reco_order": SOURCE_RECO_NAMES.tolist(), + "elapsed_sec": float(time.perf_counter() - t0), + } + + out_f = os.path.join(edge_dir, f"{out_name}.npz") + write_data_npz(out_d, out_f, metaD=out_md, verb=args.verb > 1) + if args.verb > 0: + print(f"\nSaved edge-meter metrics: {out_f}") + + plot_letters = normalize_plot_letters(args.showPlots) + if plot_letters: + from PlotterEdgeMeterAccuracy import Plotter + + unknown = sorted(set(plot_letters) - set("abc")) + if unknown: + raise ValueError(f"Unknown plot letters for edgeMeterAccuracy3c.py: {unknown}") + args.prjName = out_name + args.outPath = plot_dir + args.formatVenue = "prod" + plot = Plotter(args) + if "a" in plot_letters: + plot.recovery_summary(out_d, out_md, figId="a") + if "b" in plot_letters: + plot.count_summary(out_d, out_md, figId="b") + if "c" in plot_letters: + plot.confusion_summary(out_d, out_md, tagIdxL=[2, -1], figId="c") + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/edgeMeterFidelity3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/edgeMeterFidelity3c.py new file mode 100755 index 00000000..22e2de22 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/edgeMeterFidelity3c.py @@ -0,0 +1,553 @@ +#!/usr/bin/env python3 +"""Stability metrics for PRISM-FDR reconstructed connectivity graphs. + +Compares K >= 2 aggregate fit files (data subsets / time slices) without +requiring ground truth. Two comparison modes: + consecutive -- compare adjacent ordered pairs (k vs k+1) [default] + last -- compare each subset against the last subset as reference +""" + +import argparse +import os +import secrets +import time + +import numpy as np + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def parse_args(): + parser = argparse.ArgumentParser( + description="Stability metrics for PRISM-FDR aggregate fits (no ground truth)", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--basePath", required=True, + help="Run directory containing prismFit/, edgeFidelity/, and plots/") + parser.add_argument("--fitNameTrunk", required=True, + help="Common prefix of aggregate fit files") + parser.add_argument("--fitTags", nargs="+", required=True, + help="Ordered tags identifying each data subset") + parser.add_argument("--compare_mode", default="consecutive", + choices=["consecutive", "last"], + help="consecutive: compare adjacent pairs (x=midpoint duration); " + "last: compare each subset against the last subset (x=subset duration)") + parser.add_argument("--outName", default=argparse.SUPPRESS, + help="Output metric stem in edgeFidelity/; default is fitNameTrunk_efmHASH4") + parser.add_argument("-p", "--showPlots", nargs="+", default=["a"], + help="Plot groups: a=edge fidelity summary b=weight stats") + parser.add_argument("-X", "--noXterm", action="store_true", + help="Disable X terminal for plotting") + parser.add_argument("-v", "--verb", type=int, default=1, + help="Verbosity level") + return parser.parse_args() + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + +def flatten_words(items): + words = [] + for item in items: + words.extend(str(item).split()) + return words + + +def normalize_plot_letters(items): + letters = [] + for token in flatten_words(items): + letters.extend(list(token)) + return letters + + +def full_fit_name(trunk, tag): + tag = str(tag) + if tag.startswith(trunk): + return tag + return f"{trunk}_{tag}" + + +def fit_file(base_path, fit_name): + return os.path.join(base_path, "prismFit", f"{fit_name}.prismEM.npz") + + +def load_fit(base_path, fit_name, verb=1): + inp_f = fit_file(base_path, fit_name) + fit_d, fit_md = read_data_npz(inp_f, verb=verb > 1) + if "A_prune" not in fit_d: + raise KeyError(f"{inp_f} must contain A_prune") + if "A_sd_selected" not in fit_d: + raise KeyError(f"{inp_f} must contain A_sd_selected (bootstrap SD matrix)") + if "neuron_type" not in fit_d: + raise KeyError(f"{inp_f} must contain neuron_type") + return fit_d, fit_md, inp_f + + +def as_2d(arr, name): + arr = np.asarray(arr) + if arr.ndim != 2 or arr.shape[0] != arr.shape[1]: + raise ValueError(f"{name} must be square 2D, got {arr.shape}") + return arr + + +def safe_div(num, den): + den = float(den) + return float("nan") if den == 0.0 else float(num) / den + + +# --------------------------------------------------------------------------- +# build per-subset edge arrays +# --------------------------------------------------------------------------- + +def build_subsets(fit_records, verb=1): + """Extract A_prune, A_sd_selected, neuron_type, diagonal per subset. + + Returns a list of dicts, one per data subset. + """ + subsets = [] + for rec in fit_records: + fit_d, fit_md, inp_f, fit_name, fit_tag = rec + A_prune = as_2d(fit_d["A_prune"], "A_prune").astype(np.float64) + A_sd_selected = as_2d(fit_d["A_sd_selected"], "A_sd_selected").astype(np.float64) + # NaN means only 1 bag selected the edge (ddof=1 undefined) — treat as 0 uncertainty + np.nan_to_num(A_sd_selected, nan=0.0, copy=False) + neuron_type = np.asarray(fit_d["neuron_type"], dtype=np.int8).ravel() + single_rates = np.asarray(fit_d["single_rates"], dtype=np.float64).ravel() + n = A_prune.shape[0] + if neuron_type.shape != (n,): + raise ValueError(f"neuron_type shape {neuron_type.shape} != N={n}") + if single_rates.shape != (n,): + raise ValueError(f"single_rates shape {single_rates.shape} != N={n}") + + # parse duration from tag string (e.g. "30min" -> 30.0), fall back to metadata + duration_min = float("nan") + tag_str = str(fit_tag) + if tag_str.endswith("min"): + try: + duration_min = float(tag_str[:-3]) + except ValueError: + pass # tag doesn't encode duration; plotter will use index as x-axis + if not np.isfinite(duration_min): + if not fit_md: + raise ValueError(f"fit_tag '{fit_tag}' has no parseable duration and fit_md is empty") + if "bagsFDR_stageA" not in fit_md: + raise KeyError(f"fit_tag '{fit_tag}': fit_md missing 'bagsFDR_stageA' — re-run pipeline to regenerate files") + if "real_fit" not in fit_md["bagsFDR_stageA"]: + raise KeyError(f"fit_tag '{fit_tag}': bagsFDR_stageA missing 'real_fit' — re-run prism_FDR_Bags_train3c.py") + tr = fit_md["bagsFDR_stageA"]["real_fit"]["train"]["time_range_sec"] + tr = np.asarray(tr, dtype=np.float64) + duration_min = float(tr[1] - tr[0]) / 60.0 + + stage_b = fit_md.get("bagsFDR_stageB", {}) + min_posW = float(stage_b["min_posW"]) if "min_posW" in stage_b else float("nan") + max_negW = float(stage_b["max_negW"]) if "max_negW" in stage_b else float("nan") + + off = A_prune.copy() + np.fill_diagonal(off, 0.0) + source_edge_counts = (np.abs(off) > 1e-12).sum(axis=0).astype(np.float64) + source_edge_counts[neuron_type == 0] = 0.0 + weight_ranges = {} + edge_counts = {} + for cls, key in ((1, "exc"), (-1, "inh"), (0, "und")): + cols = neuron_type == cls + if np.any(cols): + vals = off[:, cols] + nz = vals[np.abs(vals) > 1e-12] + else: + nz = np.array([]) + weight_ranges[key] = ( + float(nz.min()) if nz.size else float("nan"), + float(nz.max()) if nz.size else float("nan"), + float(np.median(nz)) if nz.size else float("nan"), + ) + edge_counts[key] = int(nz.size) + + subsets.append({ + "fit_name": fit_name, + "fit_tag": fit_tag, + "fit_file": inp_f, + "A_prune": A_prune, + "A_sd_selected": A_sd_selected, + "neuron_type": neuron_type, + "single_rates": single_rates, + "source_edge_counts": source_edge_counts, + "diag": np.diag(A_prune).copy(), + "n": n, + "duration_min": duration_min, + "min_posW": min_posW, + "max_negW": max_negW, + "weight_ranges": weight_ranges, + "edge_counts": edge_counts, + }) + return subsets + + +# --------------------------------------------------------------------------- +# epsilon and E_global +# --------------------------------------------------------------------------- + +def compute_epsilon(subsets): + """Minimum non-zero absolute off-diagonal weight across all subsets.""" + vals = [] + for s in subsets: + A = s["A_prune"] + n = s["n"] + off = A.copy() + np.fill_diagonal(off, 0.0) + nz = np.abs(off[off != 0.0]) + if nz.size > 0: + vals.append(nz.min()) + if not vals: + raise ValueError("All A_prune matrices are entirely zero; cannot compute epsilon") + return float(np.min(vals)) + + +def compute_e_global(subsets, eps): + """Union of active off-diagonal edges across all subsets.""" + n = subsets[0]["n"] + union = np.zeros((n, n), dtype=bool) + off_mask = ~np.eye(n, dtype=bool) + for s in subsets: + union |= (np.abs(s["A_prune"]) >= eps) & off_mask + return union # (n, n) bool, diagonal always False + + + + +# --------------------------------------------------------------------------- +# off-diagonal pairwise metrics +# --------------------------------------------------------------------------- + +def edge_sets(A, e_global, eps): + """Return positive and negative edge sets as boolean arrays over e_global flat.""" + idx = np.where(e_global) + vals = A[idx] + E_pos = vals > eps + E_neg = vals < -eps + return E_pos, E_neg + + +def jaccard(a, b): + inter = np.count_nonzero(a & b) + union = np.count_nonzero(a | b) + return safe_div(inter, union) + + +def compute_topology_metrics(A1, A2, e_global, eps, include_smr=True): + """J+, J-, and optionally SMR for one pair of weight matrices.""" + E1p, E1n = edge_sets(A1, e_global, eps) + E2p, E2n = edge_sets(A2, e_global, eps) + + jp = jaccard(E1p, E2p) + jn = jaccard(E1n, E2n) + out = {"J_pos": jp, "J_neg": jn} + + if include_smr: + both_present = (E1p | E1n) & (E2p | E2n) + sign_agree = ((E1p & E2p) | (E1n & E2n)) + out["SMR"] = safe_div(np.count_nonzero(sign_agree), np.count_nonzero(both_present)) + + return out + + +def _spearman_weighted(X, Y, W=None): + """Weighted Spearman correlation of two 1D arrays. + + W=None gives the unweighted version. + """ + n = len(X) + if n < 2: + return float("nan") + # ranks (scipy-style average for ties) + from scipy.stats import rankdata + Rx = rankdata(X).astype(np.float64) + Ry = rankdata(Y).astype(np.float64) + + if W is None: + W = np.ones(n, dtype=np.float64) + W = np.asarray(W, dtype=np.float64) + + Ws = W.sum() + if Ws == 0.0: + return float("nan") + mRx = (W * Rx).sum() / Ws + mRy = (W * Ry).sum() / Ws + dx = Rx - mRx + dy = Ry - mRy + num = (W * dx * dy).sum() + den = np.sqrt((W * dx * dx).sum() * (W * dy * dy).sum()) + return safe_div(num, den) + + +def compute_magnitude_metrics(A1, sigma1, A2, sigma2, e_global, eps): + """r_s and r_s^w for one pair.""" + idx = np.where(e_global) + X = np.abs(A1[idx]) + Y = np.abs(A2[idx]) + S1 = np.nan_to_num(sigma1[idx], nan=0.0) + S2 = np.nan_to_num(sigma2[idx], nan=0.0) + W = 1.0 / ((S1 + S2) / 2.0 + eps) + + rs = _spearman_weighted(X, Y, W=None) + rsw = _spearman_weighted(X, Y, W=W) + return {"r_spearman": rs, "r_spearman_w": rsw} + + +def metrics_for_scope(A1, sigma1, A2, sigma2, e_global, eps, neuron_type1, neuron_type2, scope): + """Compute all off-diagonal metrics for a given source-type scope. + + scope: 'all', 'exc', 'inh' + """ + if scope == "all": + mask = e_global + else: + sign = 1 if scope == "exc" else -1 + n = A1.shape[0] + src_mask = np.zeros((n, n), dtype=bool) + for j in range(n): + if neuron_type1[j] == sign or neuron_type2[j] == sign: + src_mask[:, j] = True + mask = e_global & src_mask + + if not np.any(mask): + nan4 = {"J_pos": float("nan"), "J_neg": float("nan")} + if scope == "all": + nan4["SMR"] = float("nan") + return nan4 + nan2 = {"r_spearman": float("nan"), "r_spearman_w": float("nan")} + return {**nan4, **nan2} + + top = compute_topology_metrics(A1, A2, mask, eps, include_smr=(scope == "all")) + if scope == "all": + return top + mag = compute_magnitude_metrics(A1, sigma1, A2, sigma2, mask, eps) + return {**top, **mag} + + +# --------------------------------------------------------------------------- +# diagonal metrics +# --------------------------------------------------------------------------- + +def compute_diagonal_metrics_pair(d1, d2): + """Mean absolute change between two diagonal vectors.""" + d1 = np.asarray(d1, dtype=np.float64) + d2 = np.asarray(d2, dtype=np.float64) + if d1.size == 0 or d1.shape != d2.shape: + return float("nan") + return float(np.mean(np.abs(d1 - d2))) + + +# --------------------------------------------------------------------------- +# run comparisons +# --------------------------------------------------------------------------- + +SCOPES = ["all", "exc", "inh"] + + +def compare_pair(s1, s2, e_global, eps): + """All metrics for one ordered pair of subsets.""" + result = {} + for scope in SCOPES: + m = metrics_for_scope( + s1["A_prune"], s1["A_sd_selected"], + s2["A_prune"], s2["A_sd_selected"], + e_global, eps, + s1["neuron_type"], s2["neuron_type"], + scope, + ) + result[scope] = m + + result["diag_mae"] = compute_diagonal_metrics_pair(s1["diag"], s2["diag"]) + return result + + +def run_consecutive(subsets, e_global, eps): + """Strategy B: compare (k, k+1) pairs.""" + comparisons = [] + for i in range(len(subsets) - 1): + s1, s2 = subsets[i], subsets[i + 1] + res = compare_pair(s1, s2, e_global, eps) + res["label"] = f"{s1['fit_tag']} vs {s2['fit_tag']}" + res["x"] = 0.5 * (s1["duration_min"] + s2["duration_min"]) + comparisons.append(res) + return comparisons + + +def run_last(subsets, e_global, eps): + """Compare each subset (except the last) against the last subset as reference. + + x-axis is each subset's own duration, so caller can pass e.g. [20min,40min], [10min,50min], + [0min,60min] and study how shrinking the time window degrades results vs. the full window. + The last subset itself is the reference and is not compared against itself. + """ + ref = subsets[-1] + comparisons = [] + for s in subsets[:-1]: + res = compare_pair(s, ref, e_global, eps) + res["label"] = f"{s['fit_tag']} vs {ref['fit_tag']}" + res["x"] = s["duration_min"] + comparisons.append(res) + return comparisons + + +# --------------------------------------------------------------------------- +# assemble output arrays +# --------------------------------------------------------------------------- + +def assemble_output(comparisons, subsets, mode, eps): + """Pack all results into numpy arrays for the Plotter.""" + nc = len(comparisons) + ns = len(subsets) + all_rates = np.concatenate([s["single_rates"] for s in subsets]) + all_rates = all_rates[np.isfinite(all_rates)] + if all_rates.size: + n_rate_bins = min(30, int(all_rates.size)) + r0, r1 = float(np.min(all_rates)), float(np.max(all_rates)) + if r0 == r1: + pad = max(0.5, abs(r0) * 0.05) + r0, r1 = r0 - pad, r1 + pad + rate_edges = np.linspace(r0, r1, n_rate_bins + 1, dtype=np.float64) + edge_rate_hist = np.zeros((n_rate_bins, ns), dtype=np.float64) + for isub, s in enumerate(subsets): + hist, _ = np.histogram( + s["single_rates"], bins=rate_edges, weights=s["source_edge_counts"] + ) + edge_rate_hist[:, isub] = hist + else: + rate_edges = np.array([0.0, 1.0], dtype=np.float64) + edge_rate_hist = np.zeros((1, ns), dtype=np.float64) + + out = { + "compare_mode": np.array([mode], dtype=object), + "epsilon": np.array([eps]), + "x": np.array([c["x"] for c in comparisons], dtype=np.float64), + "labels": np.array([c["label"] for c in comparisons], dtype=object), + "subset_tags": np.array([s["fit_tag"] for s in subsets], dtype=object), + "subset_x": np.array([s["duration_min"] for s in subsets], dtype=np.float64), + "diag_mae": np.array([c["diag_mae"] for c in comparisons], dtype=np.float64), + "min_posW": np.array([s["min_posW"] for s in subsets], dtype=np.float64), + "max_negW": np.array([s["max_negW"] for s in subsets], dtype=np.float64), + "w_min_exc": np.array([s["weight_ranges"]["exc"][0] for s in subsets], dtype=np.float64), + "w_max_exc": np.array([s["weight_ranges"]["exc"][1] for s in subsets], dtype=np.float64), + "w_med_exc": np.array([s["weight_ranges"]["exc"][2] for s in subsets], dtype=np.float64), + "w_min_inh": np.array([s["weight_ranges"]["inh"][0] for s in subsets], dtype=np.float64), + "w_max_inh": np.array([s["weight_ranges"]["inh"][1] for s in subsets], dtype=np.float64), + "w_med_inh": np.array([s["weight_ranges"]["inh"][2] for s in subsets], dtype=np.float64), + "w_min_und": np.array([s["weight_ranges"]["und"][0] for s in subsets], dtype=np.float64), + "w_max_und": np.array([s["weight_ranges"]["und"][1] for s in subsets], dtype=np.float64), + "n_edges_exc": np.array([s["edge_counts"]["exc"] for s in subsets], dtype=np.int64), + "n_edges_inh": np.array([s["edge_counts"]["inh"] for s in subsets], dtype=np.int64), + "n_edges_und": np.array([s["edge_counts"]["und"] for s in subsets], dtype=np.int64), + "rate_bin_edges": rate_edges, + "edge_count_rate_hist": edge_rate_hist, + } + + # off-diagonal metrics per scope + for scope in SCOPES: + metric_keys = ["J_pos", "J_neg", "SMR"] if scope == "all" else [ + "J_pos", "J_neg", "r_spearman", "r_spearman_w" + ] + for key in metric_keys: + out[f"{scope}_{key}"] = np.array( + [c[scope][key] for c in comparisons], dtype=np.float64 + ) + + return out + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + +def main(): + args = parse_args() + t0 = time.perf_counter() + + fit_tags = flatten_words(args.fitTags) + if len(fit_tags) < 2: + raise ValueError("Need at least 2 --fitTags") + + if not hasattr(args, "outName"): + args.outName = f"{args.fitNameTrunk}_efm{secrets.token_hex(2)}" + + fidelity_dir = os.path.join(args.basePath, "edgeFidelity") + plot_dir = os.path.join(args.basePath, "plots") + os.makedirs(fidelity_dir, exist_ok=True) + os.makedirs(plot_dir, exist_ok=True) + + fit_names = [full_fit_name(args.fitNameTrunk, tag) for tag in fit_tags] + + fit_records = [] + for fit_tag, fit_name in zip(fit_tags, fit_names): + fit_d, fit_md, inp_f = load_fit(args.basePath, fit_name, verb=args.verb) + fit_records.append((fit_d, fit_md, inp_f, fit_name, fit_tag)) + if args.verb > 0: + print(f"loaded: {fit_name}") + + subsets = build_subsets(fit_records, verb=args.verb) + eps = compute_epsilon(subsets) + e_global = compute_e_global(subsets, eps) + + if args.verb > 0: + print(f"epsilon={eps:.3e} |E_global|={np.count_nonzero(e_global)}") + + if args.compare_mode == "consecutive": + comparisons = run_consecutive(subsets, e_global, eps) + else: + comparisons = run_last(subsets, e_global, eps) + + if args.verb > 0: + for i, c in enumerate(comparisons): + print( + f"comparison {c['label']}: " + f"J+={c['all']['J_pos']:.3f} " + f"J-={c['all']['J_neg']:.3f} " + f"SMR={c['all']['SMR']:.3f} " + f"rs_exc={c['exc']['r_spearman']:.3f} " + f"rs_inh={c['inh']['r_spearman']:.3f} " + f"rsw_exc={c['exc']['r_spearman_w']:.3f} " + f"rsw_inh={c['inh']['r_spearman_w']:.3f} " + f"diag_mae={c['diag_mae']:.3g}" + ) + + out_d = assemble_output(comparisons, subsets, args.compare_mode, eps) + out_md = { + "program": "edgeMeterFidelity3c.py", + "fitNameTrunk": args.fitNameTrunk, + "fitTags": fit_tags, + "compare_mode": args.compare_mode, + "epsilon": float(eps), + "elapsed_sec": float(time.perf_counter() - t0), + } + + out_f = os.path.join(fidelity_dir, f"{args.outName}.npz") + write_data_npz(out_d, out_f, metaD=out_md, verb=args.verb > 1) + if args.verb > 0: + print(f"\nSaved fidelity metrics: {out_f}") + + plot_letters = normalize_plot_letters(args.showPlots) + if plot_letters: + from PlotterEdgeMeterFidelity import Plotter + + unknown = sorted(set(plot_letters) - set("ab")) + if unknown: + raise ValueError(f"Unknown plot letters: {unknown}") + + args.prjName = args.outName + args.outPath = plot_dir + args.formatVenue = "prod" + plot = Plotter(args) + if "a" in plot_letters: + plot.plot_jaccard_smr(out_d, out_md, figId="a") + if "b" in plot_letters: + plot.plot_weight_thresholds(out_d, out_md, figId="b") + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/estEndTime.sh b/causal_net/nonStation_ver3c_states_EM_FDR/estEndTime.sh new file mode 100755 index 00000000..5560469a --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/estEndTime.sh @@ -0,0 +1,104 @@ +#!/bin/bash + +set -euo pipefail + +usage() { + echo "Usage: $0 LOGFILE" + echo "Example: $0 Lr5a20" +} + +if [[ "$#" -ne 1 ]]; then + usage + exit 2 +fi + +log="$1" +script_dir="$(cd "$(dirname "$0")" && pwd)" + +if [[ ! -f "$log" && -f "$script_dir/$log" ]]; then + log="$script_dir/$log" +fi + +if [[ ! -f "$log" ]]; then + echo "ERROR: logfile not found: $1" >&2 + exit 1 +fi + +awk ' +function die(msg) { + print "ERROR: " msg > "/dev/stderr" + exit 1 +} + +function wall_minutes(txt, h, m, s, rest, a) { + rest = txt + sub(/^real[ \t]+/, "", rest) + h = 0 + m = 0 + s = 0 + + if (index(rest, "h") > 0) { + split(rest, a, "h") + h = a[1] + 0 + rest = a[2] + } + if (index(rest, "m") > 0 && index(rest, "s") > 0) { + split(rest, a, "m") + m = a[1] + 0 + s = a[2] + sub(/s.*$/, "", s) + return h * 60 + m + s / 60.0 + } + if (rest ~ /^[0-9.]+s/) { + s = rest + sub(/s.*$/, "", s) + return h * 60 + s / 60.0 + } + die("cannot parse wall time line: " txt) +} + +BEGIN { + in_bags = 0 +} + +/^=== FDR bags:/ { + in_bags = 1 +} + +/^numBags=/ && num_bags == "" { + num_bags = $0 + sub(/^numBags=/, "", num_bags) + sub(/[^0-9].*$/, "", num_bags) +} + +/Stage \(b\): bags=[0-9]+/ && num_bags == "" { + num_bags = $0 + sub(/^.*bags=/, "", num_bags) + sub(/[^0-9].*$/, "", num_bags) +} + +/^--- bag[ \t]+[0-9][0-9][0-9]\/[0-9][0-9][0-9]/ && num_bags == "" { + num_bags = $0 + sub(/^.*\//, "", num_bags) + sub(/[^0-9].*$/, "", num_bags) + num_bags += 1 +} + +/^real[ \t]+[0-9]/ { + if (!in_bags && em_min == "") { + em_min = wall_minutes($0) + } else if (in_bags && bag_min == "") { + bag_min = wall_minutes($0) + } +} + +END { + if (em_min == "") die("missing Reference EM real time") + if (bag_min == "") die("missing first FDR bag real time") + if (num_bags == "") die("missing requested bag count") + + total_min = em_min + num_bags * bag_min + printf "expected_runtime_min=%.1f\n", total_min + printf "em_min=%.1f first_bag_min=%.1f num_bags=%d\n", em_min, bag_min, num_bags +} +' "$log" diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/fitPrismEM.sh b/causal_net/nonStation_ver3c_states_EM_FDR/fitPrismEM.sh new file mode 100755 index 00000000..c9449773 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/fitPrismEM.sh @@ -0,0 +1,86 @@ +#!/bin/bash +# salloc -q interactive -C gpu -t 4:00:00 -N 1 -A m2043 +# module load pytorch +# basePath=/pscratch/sd/b/balewski/2026_causalNet_tmp3/ +# Simplified wrapper script for distributed Lasso Poisson training +# +# ./fitPrismEM.sh --basePath $basePath --dataName daleN100_b6ce1e_2ce4d2 --num_states 2 --num_em_iters 2 --m_epochs 16 --fitName abc16 + + +# Capture variable arguments as text +varArgs="$*" + +# Require key variable arguments to be explicitly provided +for required_arg in --dataName --num_states --basePath ; do + if [[ " $varArgs " != *" $required_arg "* ]]; then + echo "ERROR: Missing required argument in varArgs: $required_arg" >&2 + echo "Provided varArgs: $varArgs" >&2 + exit 1 + fi +done + +# Validate that --basePath points to an existing directory +BASE_PATH="" +for ((i=1; i<=$#; i++)); do + arg="${!i}" + if [[ "$arg" == "--basePath" ]]; then + j=$((i + 1)) + if [ "$j" -le "$#" ]; then + BASE_PATH="${!j}" + fi + elif [[ "$arg" == --basePath=* ]]; then + BASE_PATH="${arg#--basePath=}" + fi +done + +if [ -z "$BASE_PATH" ]; then + echo "ERROR: --basePath is required and must include a directory value." >&2 + exit 1 +fi + +if [ ! -d "$BASE_PATH" ]; then + echo "ERROR: --basePath directory does not exist: $BASE_PATH" >&2 + exit 1 +fi + +# Default/fixed arguments +fixArgs=" --time_range_sec 0 80 " +# --dataPath /pscratch/sd/b/balewski/2025_causalNet_tmp/ + +# GPU configuration +NUM_GPUS=4 +CUDA_DEVICES="0,1,2,3" + +# Ensure enough GPUs are available before launching distributed training +HOST_NAME=$(hostname 2>/dev/null || echo "unknown-host") + +if ! command -v nvidia-smi >/dev/null 2>&1; then + echo "ERROR [$HOST_NAME]: nvidia-smi not found; cannot verify GPU availability." >&2 + exit 1 +fi + +AVAILABLE_GPUS=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if [ "$AVAILABLE_GPUS" -lt "$NUM_GPUS" ]; then + echo "ERROR [$HOST_NAME]: Need at least $NUM_GPUS GPUs, but only $AVAILABLE_GPUS detected." >&2 + exit 1 +fi + +# Print configuration +echo "=== Distributed Lasso Poisson Training ===" +echo "Variable args: $varArgs" +echo "Fixed args: $fixArgs" +echo "GPUs: $NUM_GPUS ($CUDA_DEVICES)" +echo "===========================================" +echo "" + +# Set environment variables and run the training +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export CUDA_VISIBLE_DEVICES=$CUDA_DEVICES + +# Execute the training command +echo "Running: ./prism_EM_train3c.py $fixArgs $varArgs" +echo "" + +time torchrun --standalone --nnodes=1 --nproc_per_node=4 ./prism_EM_train3c.py $fixArgs $varArgs diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/gen_daleMatrices3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/gen_daleMatrices3c.py new file mode 100755 index 00000000..aa6c8226 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/gen_daleMatrices3c.py @@ -0,0 +1,491 @@ +#!/usr/bin/env python3 +""" + ./gen_daleMatrices3c.py --num_neurons 100 --num_excite 70 --num_steps 10000 --dataName test1 + ./gen_daleMatrices3c.py --num_neurons 100 --num_excite 70 --spectral_radius 0.5 --Boffsets 0 10 20 --dataName test2 + +Primary purpose: generate the ground-truth dictionary (A_true, B_true) for use +by gen_nonStationarySpikes3c.py. The stationary spike generation performed here +is for evaluation only (firing-rate sanity check per B-vector). + +Dale's principle follows the paper/model convention + eta_t = A @ Y_{t-1} + B. +Thus A[i, j] is the effect from source neuron j to target neuron i, and a +source neuron's outgoing weights are stored in column j. Excitatory source +columns have positive off-diagonal weights; inhibitory source columns have +negative off-diagonal weights. Diagonal self-history terms are forced +non-positive for stability and are not used to define source type. + +Pipeline: +1. Sparse connectivity mask E_true (N, N). + Each neuron draws its own connection probability uniformly from + --edge_prob [lo, hi]. Self-loops (diagonal) are always included. + +2. Generate ONE weight matrix A_true for --spectral_radius R: + - Excitatory columns (0..N_E-1): weights ~ Uniform(1-v, 1+v), v=0.2. + - Inhibitory columns (N_E..N-1): weights ~ Uniform(-r-rv, -r+rv), + r = N_E/N_I (balance ratio). + - Mask with E_true, then rescale so rho(A_true) = R exactly. + +3. For each offset in --Boffsets generate one bias vector B_m (N,): + - idle rate range shifted by offset, converted to log scale, + with separate corrections for excitatory and inhibitory neurons. + - Stationary spikes Y are simulated (for evaluation) via + Y_t ~ Poisson(exp(clip(A @ Y_{t-1} + B_m, max=eta_clip)) * dt). + - Firing-rate statistics (rate, variance, Fano factor) computed. + +Output files saved to /truthDale/: + .simTruth.npz — A_true, B_true, E_true, node_positions + metadata + .spikes.npz — stationary spikes + rate statistics + +Output shapes (M = len(Boffsets), N = num_neurons, T = num_steps): + E_true (N, N) int — shared binary connectivity mask + A_true (N, N) float — single shared weight matrix + B_true (M, N) float — one bias vector per state/offset + node_positions (N, 2) float — static 2D node coordinates + node_is_inhibitory (N,) int — static Dale type label, 0=exc, 1=inh + node_distance_matrix (N, N) float — pairwise Euclidean distances + spikes (M, T, N) uint8 — stationary spikes per bias vector + single_rates (M, N) float — mean firing rates (Hz) per bias vector +""" + +import numpy as np +import time,hashlib +import scipy +import os +import sys + +import argparse +from pprint import pprint + +from toolbox.Util_NumpyIO import write_data_npz +from UtilDalePoisson import estimate_rates + +###### Matrix generation ################## + +def _build_placement_grid(L, H, d_min): + """Rectangular grid nodes on [0,L] x [0,H] with spacing d_min.""" + m_max = int(np.floor(L / d_min)) + n_max = int(np.floor(H / d_min)) + gx = (np.arange(0, m_max + 1, dtype=float) * d_min).reshape(-1, 1) + gy = (np.arange(0, n_max + 1, dtype=float) * d_min).reshape(1, -1) + xs = np.broadcast_to(gx, (gx.size, gy.size)).ravel() + ys = np.broadcast_to(gy, (gx.size, gy.size)).ravel() + G = np.column_stack([xs, ys]) + return G, (m_max + 1) * (n_max + 1) + + +def generate_node_locations(n_units, L, H, d_min, rng=None): + """Sample 2D node locations and return positions plus pairwise distances.""" + if rng is None: + rng = np.random.default_rng() + G, n_grid = _build_placement_grid(L, H, d_min) + if n_grid < n_units: + raise ValueError( + "grid has only %d nodes; need N <= N_G (reduce placement_min_dist or increase L/H)" + % n_grid + ) + idx = rng.choice(n_grid, size=n_units, replace=False) + P = G[idx].astype(np.float64) + + diff = P[:, np.newaxis, :] - P[np.newaxis, :, :] + D = np.sqrt(np.sum(diff * diff, axis=2)) + np.fill_diagonal(D, 0.0) + return P, D + + +def summarize_pairwise_distances(D, verb=1): + """Mean and median Euclidean distance over unique unordered pairs.""" + D = np.asarray(D, dtype=np.float64) + assert D.ndim == 2 and D.shape[0] == D.shape[1], "D must be square" + n_units = D.shape[0] + if n_units < 2: + return {"mean": float("nan"), "median": float("nan"), "n_pairs": 0} + iu = np.triu_indices(n_units, k=1) + distances = D[iu] + stats = { + "mean": float(np.mean(distances)), + "median": float(np.median(distances)), + "n_pairs": int(distances.size), + } + if verb > 0: + print( + "Pairwise distance (unique pairs): mean=%.6g median=%.6g (N=%d, pairs=%d)" + % (stats["mean"], stats["median"], n_units, stats["n_pairs"]) + ) + return stats + +def generate_sparse_mask(n_units, edge_prob): + """ + Creates a binary mask representing the network topology. + - n_units: Number of neurons + - edge_prob: [lo, hi] range; each neuron gets a random connectivity drawn uniformly from this range + - allows for self-loops (aka non-zero diagonal elements) + """ + prob_lo, prob_hi = edge_prob + conn_per_neuron = np.random.uniform(prob_lo, prob_hi, size=n_units) + E_true = (np.random.rand(n_units, n_units) < conn_per_neuron[:, None]) + np.fill_diagonal(E_true, True) # allow for self-loops + return E_true + +# Generate an initial network connectivity matrix +def init_W(n_units, n_excite, E_true, R, varyW, verb=1): + """ + Generalized weight initialization with source-column Dale convention. + + Model convention is eta = A @ y_prev + B, so A[i, j] is source j to + target i. Dale's law therefore constrains columns: one source neuron + can only excite or only inhibit all downstream targets. Diagonal + self-history terms are forced non-positive after masking. + """ + # 1. Assertions to ensure valid population counts + n_inhib = n_units - n_excite + assert n_excite > 0, "n_excite must be greater than 0" + assert n_inhib > 0, "n_inhib must be greater than 0 (n_units > n_excite)" + + # 2. Balance ratio: each target row receives N_E positive source columns + # and N_I negative source columns with approximately zero net mean. + ie_ratio = n_excite / n_inhib + + # 3. Initialize the weight container + W = np.zeros((n_units, n_units)) + + # 4. Assign excitatory source columns (0 to n_excite-1). + W[:, :n_excite] = np.random.uniform(1.0 - varyW, + 1.0 + varyW, + (n_units, n_excite)) + + # 5. Assign inhibitory source columns (n_excite to n_units-1). + i_center = -ie_ratio + i_vary = varyW * ie_ratio + W[:, n_excite:] = np.random.uniform(i_center - i_vary, + i_center + i_vary, + (n_units, n_inhib)) + + # 6. Apply the connectivity mask (Topology) + W = W * E_true + + # 7. Enforce stabilizing diagonal: flip sign of any positive self-loop. + d_idx = np.diag_indices(n_units) + W[d_idx] = np.where(W[d_idx] > 0.0, -W[d_idx], W[d_idx]) + + # 8. Spectral Normalization to ensure Stability + # This scales the entire cloud of eigenvalues to fit within radius R + current_rho = np.max(np.abs(np.linalg.eigvals(W))) + if verb > 0: + print('initW current_rho :',current_rho ) + if current_rho > 0: + W = W * (R / current_rho) + + return W + +# Generate the full set of matrices for use in subsequent synthetic experiments +def gen_dale_matrics(conf, E_true, verb=1): + """Generates one stable Dale matrix based on configuration.""" + num_neurons, num_excite, Rmax= conf['num_neurons'], conf['num_excite'],conf[ 'spect_radius'] + if verb > 0: + print(f"\n=== Generating Dale Matrix ===") + print(f"Parameters: N={num_neurons}, excit={num_excite}, Rmax={Rmax:.1f}") + + # additional configuration + varyW=0.2 # controls variation of excitatory weights + + A=init_W(num_neurons, num_excite, E_true, Rmax, varyW, verb=verb) + return A + +def spectral_radius_scaling(A, factors): + n = A.shape[0] + off_diag_mask = ~np.eye(n, dtype=bool) + + print(f"{'factor':>8s} {'ρ(all)':>10s} {'ρ(off-diag)':>12s}") + print("-" * 34) + + for c in factors: + # a) scale all elements + rho_all = np.max(np.abs(np.linalg.eigvals(c * A))) + + # b) scale only off-diagonal + A_off = A.copy() + A_off[off_diag_mask] *= c + rho_off = np.max(np.abs(np.linalg.eigvals(A_off))) + + print(f"{c:8.3f} {rho_all:10.4f} {rho_off:12.4f}") + + +#################### Simulation ################## + +def set_flat_selfSpiking(Nn, idleRate, spect_radius, num_excite, boffsets): + """Generate B_idle per B-offset; excitatory idle-rate range is scaled by sqrt(50/Nn).""" + assert 0 < num_excite <= Nn + sizeScale = np.sqrt(float(Nn)/100) + B_all = np.zeros((len(boffsets), Nn)) + R = float(spect_radius) + for ib, offset in enumerate(boffsets): + idle_eff = np.array(idleRate, dtype=float) + Ri_scaled = idle_eff * R + Bi = np.log(Ri_scaled) + B_all[ib] = np.random.uniform(Bi[0], Bi[1], size=(Nn,)) + #B_all+= + float(offset) + if 1: # for ver Mar 11 + B_all[ib, :num_excite] +=1+offset - R*1.5 - sizeScale # reduce excitatory rate + B_all[ib, num_excite:] += offset # reduce inhibitory rate + return B_all + +def gen_stationary_lag1_poisson(num_steps, dt, A, B_intercept, num_excite, eta_clip,verb=0): + """ + Generates a multivariate Poisson VAR(1) process: + Y_t ~ Poisson(exp(A @ Y_{t-1} + B_intercept)) + + Args: + num_steps (int): Number of time steps for the simulation. + dt (float): Integration time step in seconds. + A (np.ndarray): N x N autoregressive coefficient matrix (the Dale matrix). + B_intercept (np.ndarray): N-dimensional intercept vector (bias). + num_excite (int): Number of excitatory neurons. + verb (int): Verbosity level for printing progress. + + Returns: + (Y, A, B_intercept): Tuple containing: + Y (np.ndarray): num_steps x N time series of spike counts. + A (np.ndarray): The input connectivity matrix. + B_intercept (np.ndarray): The input intercept vector. + """ + if verb > 0: + print(f"\n=== Generating Poisson Process ===") + print(f"Simulation parameters: steps={num_steps}, dt={dt:.3f}, neurons={A.shape[0]}, excit={num_excite}") + print(f"Matrix A stats: min={np.min(A):.3f}, max={np.max(A):.3f}, mean={np.mean(A):.3f}") + print(f"Bias B stats: min={np.min(B_intercept):.3f}, max={np.max(B_intercept):.3f}, mean={np.mean(B_intercept):.3f}") + + if A is None: + raise ValueError("Connectivity matrix A cannot be None.") + d=A.shape[0] + + Y = np.zeros((num_steps, d), dtype=int) + Y[0] = np.random.poisson(np.exp(B_intercept)*dt) # initial state + if verb > 0: + print(f"Initial state: total spikes={np.sum(Y[0])}, excit spikes={np.sum(Y[0][:num_excite])}, inhib spikes={np.sum(Y[0][num_excite:])}") + + if verb>0: + print('t=0 Y[t] sum=%d, Excit(first 3):%s, Inhib(first 3):%s'%(np.sum(Y[0]), Y[0][:3], Y[0][num_excite:num_excite+3])) + + kk=5 + # Main simulation loop + if verb > 0: + print("Starting main simulation loop...") + for t in range(1, num_steps): + eta = A @ Y[t-1] + B_intercept + #eta = B_intercept # use it to see idle rate only + lambda_t = np.exp(np.clip(eta, max=eta_clip)) # avoid overflow + Y[t] = np.random.poisson(lambda_t*dt) + + if verb>0 and t<5: + print('t=%d Y[t] sum=%d, Excit:%s, Inhib:%s'%(t, np.sum(Y[t]), Y[t][:kk], Y[t][num_excite:num_excite+kk])) + + if verb > 0 and t % (num_steps // 4) == 0: + print(f" Progress: {t}/{num_steps} steps ({t/num_steps*100:.1f}%) - total spikes in this step: {np.sum(Y[t])}") + + if verb > 0: + print(f"Simulation complete. Final state: total spikes={np.sum(Y[-1])}") + print(f"Spike data stats: min={np.min(Y)}, max={np.max(Y)}, mean={np.mean(Y):.2f}, total spikes={np.sum(Y)}") + + return Y + +######################### +# MAIN +######################### + +def main(): + print("=" * 60) + print("DALE POISSON SIMULATION Lag=1") + print("=" * 60) + + parser = argparse.ArgumentParser(description="Simulate a recurrent neural network with Dale's principle.") + parser.add_argument("--num_neurons", type=int, default=50, help="Total number of neurons in the network.") + parser.add_argument("--num_excite", type=int, default=None, help="Number of excitatory neurons.") + parser.add_argument("--placement_H_L_delta", type=float, nargs=3, default=[1.0, 2.0, 2.0], + metavar=("placement_H", "placement_L", "placement_ker_delta"), + help="Node placement: [0,H] height, [0,L] width, distance-kernel exponent delta. Delta is recorded for topology provenance.") + parser.add_argument("--placement_min_dist", type=float, default=0.01, + help="Grid spacing d_min; minimum inter-neuron distance.") + parser.add_argument("--edge_prob", type=float, nargs=2, default=[0.05, 0.2], help="Range of edge probability [min, max]; mean is used as mask connectivity.") + parser.add_argument("--num_steps", type=int, default=10_001, help="Number of time steps for simulation.") + parser.add_argument("--step_size", type=float, default=0.01, help="Integration time step size (dt) in seconds.") + parser.add_argument("--spectral_radius", type=float, default=0.3, help="Target spectral radius value for the connectivity matrix.") + parser.add_argument("--idleRate", type=float, nargs=2, default=[15, 30.], help="Range of idle firing rates [min, max] in Hz.") + parser.add_argument("--Boffsets", type=float, nargs='+', default=[0.0], help="Per-state offsets added to idleRate range.") + parser.add_argument('-v',"--verb", type=int, default=1, help="Verbosity level (0=quiet, 1=normal).") + parser.add_argument("--dataName", type=str, default=None, help="Base name for output files (default: daleN_).") + parser.add_argument("--basePath", type=str, default='/pscratch/sd/b/balewski/2025_causalNet_tmp/', help="Output directory for all files.") + + np.set_printoptions(precision=3, suppress=True) + + args = parser.parse_args() + placement_H, placement_L, placement_ker_delta = ( + float(args.placement_H_L_delta[0]), + float(args.placement_H_L_delta[1]), + float(args.placement_H_L_delta[2]), + ) + if not (placement_H > 0 and placement_L > 0): + raise ValueError("placement_H_L_delta requires positive H and L") + if placement_ker_delta <= 0: + raise ValueError("placement_H_L_delta third value (placement_ker_delta) must be positive") + placement_min_dist = float(args.placement_min_dist) + if placement_min_dist <= 0: + raise ValueError("placement_min_dist must be positive") + + args.varTwindow=5 #(sec) + args.poisson_eta_clip=5 #~ [1e-3Hz , 1e3Hz] + if args.dataName is None: + args.dataName='daleN%d_'%args.num_neurons+hashlib.md5(os.urandom(32)).hexdigest()[:6] + + outPath=os.path.join(args.basePath, 'truthDale') + if args.num_excite==None : # per Roy & Kris wisdom + args.num_excite=int(0.8*args.num_neurons) + print('gen dale matrices args:', vars(args), '\n') + + # Validation checks + Nn = args.num_neurons + assert Nn >= 10 + assert args.num_excite >= 5 + assert args.num_excite < Nn + assert os.path.exists(args.basePath) + assert os.path.exists(outPath) + assert args.num_steps>=1000 + assert args.step_size>0.001 + assert args.idleRate[0]>=0.5 + assert args.idleRate[1]>args.idleRate[0] + assert len(args.Boffsets) >= 1 + + node_positions, node_distance_matrix = generate_node_locations( + Nn, placement_L, placement_H, placement_min_dist + ) + node_is_inhibitory = np.zeros(Nn, dtype=np.int32) + node_is_inhibitory[args.num_excite:] = 1 + distance_stats = summarize_pairwise_distances(node_distance_matrix, verb=args.verb) + + # Generate sparse connectivity mask (per-neuron random connectivity in edge_prob range) + E_true = generate_sparse_mask(Nn, args.edge_prob) + print(f"\n=== Generated sparse mask: edge_prob={args.edge_prob}, actual={np.mean(E_true):.3f}, non-zero={np.sum(E_true)} ===") + + dale_conf = { + 'num_neurons': args.num_neurons, + 'num_excite': args.num_excite, + 'spectral_radius': args.spectral_radius, + 'placement_L': placement_L, + 'placement_H': placement_H, + 'placement_min_dist': placement_min_dist, + 'placement_ker_delta': placement_ker_delta, + 'node_distance_stats': distance_stats, + 'edge_prob': args.edge_prob, + 'idleRate': args.idleRate, + 'Boffsets': args.Boffsets, + } + + print("Dale configuration:"); pprint(dale_conf) + + evol_conf={ + 'num_steps': args.num_steps, + 'step_size': args.step_size, + 'evol_time': args.num_steps*args.step_size, + 'poisson_eta_clip': args.poisson_eta_clip + } + + B_all = set_flat_selfSpiking(Nn, args.idleRate, args.spectral_radius, args.num_excite, boffsets=args.Boffsets) + + max_samples = 100_000 + + Y_list = [] + rates_list, rates_var_list, fano_list = [], [], [] + stats_list = [] + + verb_r = args.verb + print(f"\n{'='*60}") + print(f" Spectral radius: R={args.spectral_radius:.3f}") + print(f"{'='*60}") + + dale_conf_r = dale_conf.copy() + dale_conf_r['spect_radius'] = args.spectral_radius + + A_dale = gen_dale_matrics(dale_conf_r, E_true, verb=verb_r) + + if 0: # test scaling behavior of spectral radius when scaling A by different factors + factors = [0.2, 0.4, 0.6, 0.8, 0.9, ] + spectral_radius_scaling(A_dale, factors) + exit(1) + + if verb_r > 0: + total_connections = A_dale.size + zero_connections = np.sum(np.abs(A_dale) < 1e-10) + non_zero_connections = total_connections - zero_connections + sparsity = zero_connections / total_connections + print(f'Matrix sparsity: {sparsity*100:.1f}% ({zero_connections}/{total_connections} connections are zero)') + print(f'Non-zero connections: {non_zero_connections} ({(1-sparsity)*100:.1f}%)') + + for ib, offset in enumerate(args.Boffsets): + print(f"\n{'='*60}") + print(f" B offset [{ib+1}/{len(args.Boffsets)}]: offset={offset:.3f}") + print(f"{'='*60}") + + B_idle = B_all[ib] + start_time = time.time() + Y = gen_stationary_lag1_poisson(num_steps=args.num_steps, dt=args.step_size, A=A_dale, B_intercept=B_idle, num_excite=args.num_excite,eta_clip=args.poisson_eta_clip, verb=verb_r if ib == 0 else 0) + sim_time = time.time() - start_time + print("Spike generation completed in %.1f seconds" % sim_time) + + stats_dict, rates_dict, _ = estimate_rates(Y, dt=args.step_size, num_excite=args.num_excite, max_samples=max_samples, varTwindow=args.varTwindow, mxNn=5, verb=0, spect_radius=args.spectral_radius) + stats_dict['var_time_window_sec'] = float(args.varTwindow) + stats_dict['max_samples'] = int(max_samples) + + Y_list.append(np.clip(Y, 0, 255).astype(np.uint8)) + rates_list.append(rates_dict['single_rates']) + rates_var_list.append(rates_dict['sigle_rates_var']) + fano_list.append(rates_dict['single_fano_fact']) + stats_list.append(stats_dict) + + # Stack results along axis=0 (B-offset dimension) + trueD = { + 'A_true': A_dale, + 'B_true': B_all, + 'E_true': E_true, + 'node_positions': node_positions, + 'node_is_inhibitory': node_is_inhibitory, + 'node_distance_matrix': node_distance_matrix, + } + + trueMD = {'dale_conf': dale_conf, 'evol_conf': evol_conf, 'short_name': args.dataName, + 'provenance':{'state_model_file':args.dataName}} + # delete , 'dale_simu_stats': stats_list} + + spikeD = { + 'spikes': np.stack(Y_list, axis=0), + 'single_rates': np.stack(rates_list, axis=0), + 'sigle_rates_var': np.stack(rates_var_list, axis=0), + 'single_fano_fact': np.stack(fano_list, axis=0) + } + spikeMD={ 'short_name':args.dataName,'time_step_sec':args.step_size,'data_type':'simDaleStates', 'poisson_eta_clip': args.poisson_eta_clip } + + outFt = os.path.join(outPath,args.dataName + '.simTruth.npz') + write_data_npz(trueD, outFt, metaD=trueMD) + if args.verb>1: pprint(trueMD) + outFs = os.path.join(outPath, args.dataName + '.spikes.npz') + write_data_npz(spikeD, outFs, metaD=spikeMD) + if args.verb>1: pprint(spikeMD) + + print("\nSimulation completed successfully!") + print(f"\nRate Summary {args.dataName} N={args.num_neurons} exc={args.num_excite}, R={args.spectral_radius:.3f} ") + header = f"{'state':>5} {'B offset':>9} {'all rate (Hz)':>14} {'exc rate (Hz)':>14} {'inh rate (Hz)':>14}" + print(header) + print("-" * len(header)) + for ib, offset in enumerate(args.Boffsets): + s = stats_list[ib] + print(f"{ib:5d} {float(offset):9.1f} {s['avg_spike_rate_all']:14.1f} {s['avg_spike_rate_excit']:14.1f} {s['avg_spike_rate_inhib']:14.1f}") + print("\nNext step commands:") + print(" basePath="+args.basePath) + print(" ./view_daleMatrix3.py --basePath $basePath --dataName %s -p b a c d -m 0 -X " % args.dataName) + print(" ./view_spikesTrain3.py --basePath $basePath --dataName %s --time_range_sec 0 20 -p b -m 0 -X " % args.dataName) + print(" ./gen_nonStationarySpikes3c.py --basePath $basePath --inputStates %s --true_dwell_sec 1.0 " % args.dataName) + print(" ./fit_lassoPoisson.py --basePath $basePath --inpPath ${basePath}/truthDale --dataName %s --num_epochs 50 " % args.dataName) + + +if __name__ == '__main__': + main() + diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/gen_nonStationarySpikes3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/gen_nonStationarySpikes3c.py new file mode 100755 index 00000000..fc8065a4 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/gen_nonStationarySpikes3c.py @@ -0,0 +1,436 @@ +#!/usr/bin/env python3 +""" +Generate one non-stationary spike train using the ground-truth dictionary +(A_true, B_true) produced by gen_daleMatrices3.py. + +The connectivity matrix A is shared across all states. Each state m has its +own bias vector B_m. At every time bin the effective bias is the convex +combination B_eff = sum_m c_mt * B_m, where c_t tracks a slowly-moving +target that aims at the one-hot vector of the current target state S_true[t]. + +State-sequence schedules (--schedule): + mc — Markov chain with geometric dwell (mean = true_dwell_sec / dt bins), + minimum dwell = ceil(0.3 * true_dwell_sec / dt). + rr — Round-robin: states cycle 0,1,...,M-1 with exactly + true_dwell_sec / dt bins per visit. + +Smooth coefficient update at each bin (state_change_speed = nu): + c_t <- Simplex_project( c_{t-1} + clip(e_{S_t} - c_{t-1}, -nu, nu) ) + +Spike generation (Poisson GLM): + eta_t = A @ Y_{t-1} + B_eff + lambda_t = exp( clip(eta_t, max=eta_clip) ) * dt + Y_t ~ Poisson(lambda_t) + +An oracle state sequence S_oracle is also computed: at each t the state with +the highest Poisson log-likelihood given (A, {B_m}) and the observed spikes. + +Output files saved to /spikesData/: + .spikes.npz — spikes (T, N) int32, single_rates (N,) + .prismTruth.npz — S_true, C_true, S_oracle, state_transition, + A_true, B_true, rate variance, Fano factor + +Output shapes (T = num_steps, N = num_neurons, M = num_states): + spikes (T, N) int32 — non-stationary spike counts + single_rates (N,) float — mean firing rates (Hz) + S_true (T,) int32 — target state index per bin + C_true (T, M) float32 — smooth simplex coefficients + S_oracle (T,) int32 — oracle (max-likelihood) state per bin + state_transition (M, M) float32 — Markov transition matrix used/implied + A_true (N, N) float — shared connectivity matrix (copy) + B_true (M, N) float — per-state bias vectors (copy) +""" + +import os +import hashlib +import argparse +from pprint import pprint +import numpy as np + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from gen_daleMatrices3c import estimate_rates + + +def get_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v", "--verbosity", type=int, default=1, dest="verb", help="Verbosity level.") + parser.add_argument("--basePath", default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input data") + parser.add_argument("--inputStates", default=None, help="input simTruth base name") + parser.add_argument("--dataName", type=str, default=None, help="output spikes base name (default: _)") + + parser.add_argument("-t", "--num_steps", type=int, default=None, help="Number of time steps (default: from input evol_conf)") + parser.add_argument("--state_change_speed", type=float, default=0.33, help="Max coefficient change per step.") + parser.add_argument("--true_dwell_sec", type=float, default=1.0, + help="Mean dwell time in seconds to stay in a target state.") + parser.add_argument("--seed", type=int, default=42, help="Optional random seed.") + parser.add_argument("--schedule", choices=["mc", "rr"], default="mc", + help="State schedule: 'mc' = random Markov chain (default), " + "'rr' = deterministic round-robin (equal state coverage).") + + args = parser.parse_args() + args.inpPath = os.path.join(args.basePath, "truthDale") + args.outPath = os.path.join(args.basePath, "spikesData") + if args.dataName is None: + args.dataName = f"{args.inputStates}_{hashlib.md5(os.urandom(32)).hexdigest()[:6]}" + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + assert os.path.exists(args.basePath), f"missing basePath: {args.basePath}" + assert os.path.exists(args.inpPath), f"missing inpPath: {args.inpPath}" + os.makedirs(args.outPath, exist_ok=True) + return args + + +def ensure_state_atoms(A_in, B_in): + """Normalize A/B arrays to A(N,N), B(M,N) with shared A across states.""" + A = A_in + B = B_in + assert A.ndim == 2, f"A_true must be 2D, got shape={A.shape}" + if B.ndim == 1: + B = B[None, ...] + assert B.ndim == 2, f"B_true must be 1D or 2D, got shape={B.shape}" + assert A.shape[0] == A.shape[1], f"A must be square, got {A.shape}" + assert A.shape[0] == B.shape[1], f"neuron mismatch: A N={A.shape[0]}, B N={B.shape[1]}" + return A.astype(float), B.astype(float) + + +def _print_state_stats(S_true, n_states, n_steps): + """Print per-state counts and dwell statistics (shared by both schedules).""" + counts = np.bincount(S_true, minlength=n_states) + fracs = counts / n_steps + parts = ' '.join(f's{m}:{counts[m]}({fracs[m]:.1%})' for m in range(n_states)) + print(f"Target state distribution (T={n_steps}): {parts}") + + dwell_lens = {m: [] for m in range(n_states)} + run_state, run_len = S_true[0], 1 + for t in range(1, n_steps): + if S_true[t] == run_state: + run_len += 1 + else: + dwell_lens[run_state].append(run_len) + run_state, run_len = S_true[t], 1 + dwell_lens[run_state].append(run_len) + for m in range(n_states): + d = dwell_lens[m] + print(f" state {m}: {len(d)} episodes, dwell mean={np.mean(d):.1f} steps, " + f"total={np.sum(d)} steps") + + +def build_target_states_mc(n_steps, n_states, dwell_steps, rng): + """Random Markov chain target-state sequence with geometric dwell time. + + Each time step the chain either stays (prob = 1 - 1/dwell_steps) or + transitions uniformly to one of the other states. State coverage is + uncontrolled — a single unlucky run can give very few bins to one state. + """ + if n_states == 1: + return np.zeros(n_steps, dtype=np.int32), np.ones((1, 1), dtype=np.float32) + + minDwellFrac = 0.3 + min_dwell_steps = max(1, int(np.ceil(minDwellFrac * float(dwell_steps)))) + p_exit = 1.0 / float(dwell_steps) + p_stay = 1.0 - p_exit + trans = np.full((n_states, n_states), p_exit / float(n_states - 1), dtype=float) + np.fill_diagonal(trans, p_stay) + + S_true = np.zeros(n_steps, dtype=np.int32) + state = 0 + run_len = 0 + for t in range(n_steps): + if run_len < min_dwell_steps: + next_state = state + else: + next_state = rng.choice(n_states, p=trans[state]) + S_true[t] = next_state + if next_state == state: + run_len += 1 + else: + state = next_state + run_len = 1 + + _print_state_stats(S_true, n_states, n_steps) + return S_true, trans.astype(np.float32) + + +def build_target_states_rr(n_steps, n_states, dwell_steps, rng): + """Round-robin with fixed dwell_steps per visit. + + States cycle 0,1,2,0,1,2,... Every visit lasts exactly dwell_steps + steps (last visit of each state is trimmed to fit T). Total steps per + state is guaranteed to be T//M ± dwell_steps — no luck involved. + """ + if n_states == 1: + return np.zeros(n_steps, dtype=np.int32), np.ones((1, 1), dtype=np.float32) + + S_true = np.zeros(n_steps, dtype=np.int32) + t = 0 + visit = 0 + while t < n_steps: + state = visit % n_states # strict round-robin + t_end = min(t + dwell_steps, n_steps) # fixed dwell, trim at end + S_true[t:t_end] = state + t = t_end + visit += 1 + + # Transition matrix: same symmetric form as MC for metadata consistency + p_exit = 1.0 / float(dwell_steps) + p_stay = 1.0 - p_exit + trans = np.full((n_states, n_states), p_exit / float(n_states - 1), dtype=float) + np.fill_diagonal(trans, p_stay) + + _print_state_stats(S_true, n_states, n_steps) + return S_true, trans.astype(np.float32) + +def simulate_switching_poisson(n_steps, A, B_atoms, S_true, state_change_speed, dt, eta_clip, rng, verb=1): + """Switching Poisson generator: smooth c_t toward one-hot target state S_true[t].""" + n_states, n_neurons = B_atoms.shape + spikes = np.zeros((n_steps, n_neurons), dtype=np.int32) + C_true = np.zeros((n_steps, n_states), dtype=np.float32) + + c_curr = np.zeros(n_states, dtype=float) + c_curr[S_true[0]] = 1.0 + + for t in range(n_steps): + target = np.zeros(n_states, dtype=float) + target[S_true[t]] = 1.0 + + diff = target - c_curr + c_curr += np.clip(diff, -state_change_speed, state_change_speed) + c_sum = np.sum(c_curr) + if c_sum <= 0: + c_curr = target.copy() + c_sum = 1.0 + c_curr /= c_sum + C_true[t] = c_curr + + if verb > 1 and t < 12: + c_str = ", ".join([f"{x:.2f}" for x in c_curr]) + print(f"{t:<6} | {int(S_true[t]):<6} | [{c_str}]") + + A_eff = A + B_eff = np.einsum("m,mj->j", c_curr, B_atoms) + + prev_y = spikes[t - 1].astype(float) if t > 0 else np.zeros(n_neurons, dtype=float) + eta_t = A_eff @ prev_y + B_eff + lambda_t = np.exp(np.clip(eta_t, max=eta_clip)) + spikes[t] = rng.poisson(lambda_t * dt).astype(np.int32) + + return spikes, C_true + + +def compute_oracle_states_comA(spikes, A, B_atoms, dt, eta_clip): + """Oracle state by max Poisson log-likelihood using shared A and per-state B.""" + n_steps, n_neurons = spikes.shape + n_states = B_atoms.shape[0] + S_oracle = np.zeros((n_steps,), dtype=np.int32) + prev_y = np.zeros(n_neurons, dtype=np.float64) + for t in range(n_steps): + y_curr = spikes[t].astype(np.float64) + base = A @ prev_y + eta = base[None, :] + B_atoms + eta_c = np.clip(eta, max=eta_clip) + lam = np.exp(eta_c) * dt + scores = np.sum(y_curr * eta_c - lam, axis=1) + S_oracle[t] = int(np.argmax(scores)) + prev_y = y_curr + return S_oracle + + +def compute_oracle_score(S_true, S_oracle, n_states): + """Return overall and per-state agreement between oracle and target states.""" + if S_true is None or S_oracle is None: + return None, [] + nmin = min(S_true.shape[0], S_oracle.shape[0]) + if nmin <= 0: + return None, [] + avr_score = float(np.mean(S_true[:nmin] == S_oracle[:nmin])) + score_per_state = [] + for m in range(n_states): + mask = S_true[:nmin] == m + if mask.sum() > 0: + score_per_state.append(float((S_oracle[:nmin][mask] == m).mean())) + else: + score_per_state.append(float("nan")) + return avr_score, score_per_state + + +def main(): + args = get_parser() + print('gen non-stationary spikes args:', vars(args), '\n') + np.set_printoptions(precision=3, suppress=True) + + daleFF = os.path.join(args.inpPath, f"{args.inputStates}.simTruth.npz") + daleD, daleMD = read_data_npz(daleFF, verb=args.verb > 0) + if args.verb > 1: + print("\nInput simTruth metadata:") + pprint(daleMD) + + assert isinstance(daleMD, dict), "Expected dictionary metadata in simTruth file" + dale_conf_in = daleMD["dale_conf"] + evol_conf_in = daleMD["evol_conf"] + proven_in=daleMD['provenance'] + + step_size = evol_conf_in["step_size"] + var_time_window_sec = 5 #(sec) + max_samples = 100_000 # time steps + + if args.num_steps is None: + args.num_steps = int(evol_conf_in["num_steps"]) + + assert args.num_steps >= 100 + assert step_size > 0 + assert args.state_change_speed > 0 + assert args.true_dwell_sec > 0.0 + assert max_samples >= 100 + assert var_time_window_sec > 0 + + true_dwell_steps = max(1, int(np.ceil(float(args.true_dwell_sec) / float(step_size)))) + + A_atoms, B_atoms = ensure_state_atoms(daleD["A_true"], daleD["B_true"]) + n_states, n_neurons = B_atoms.shape + num_excite = int(dale_conf_in["num_excite"]) + num_excite = min(max(1, num_excite), n_neurons - 1) + + rng = np.random.default_rng(args.seed) + + schedule_fn = {"mc": build_target_states_mc, + "rr": build_target_states_rr}[args.schedule] + S_true, transition_matrix = schedule_fn(args.num_steps, n_states, true_dwell_steps, rng) + + spikes, C_true = simulate_switching_poisson( + n_steps=args.num_steps, + A=A_atoms, + B_atoms=B_atoms, + S_true=S_true, + state_change_speed=args.state_change_speed, + dt=step_size, + eta_clip=evol_conf_in['poisson_eta_clip'], + rng=rng, + verb=args.verb, + ) + S_oracle = compute_oracle_states_comA( + spikes=spikes, + A=A_atoms, + B_atoms=B_atoms, + dt=step_size, + eta_clip=evol_conf_in['poisson_eta_clip'], + ) + + stats_dict, rates_dict, _ = estimate_rates( + spikes, + dt=step_size, + num_excite=num_excite, + max_samples=max_samples, + varTwindow=var_time_window_sec, + mxNn=5, + verb=args.verb, + spect_radius=None, + ) + oracle_score, oracle_score_per_state = compute_oracle_score(S_true, S_oracle, n_states) + true_dwell_sec_req = float(args.true_dwell_sec) + true_dwell_sec_eff = float(true_dwell_steps * step_size) + + evol_conf = { + "num_steps": int(args.num_steps), + "step_size": float(step_size), + "evol_time": float(args.num_steps * step_size), + "state_change_speed": float(args.state_change_speed), + "true_dwell_sec": true_dwell_sec_req, + "true_dwell_sec_eff": true_dwell_sec_eff, + "true_dwell_steps": int(true_dwell_steps), + "true_dwell_time_sec": true_dwell_sec_eff, # legacy alias + "num_states": int(n_states), + "seed": args.seed, + "state_schedule": args.schedule, + "num_states": int(n_states), + "max_samples": int(max_samples), + "truth_input_name" : args.inputStates, + } + + dale_conf = dict(dale_conf_in) + dale_conf["num_neurons"] = int(dale_conf["num_neurons"]) + dale_conf["num_excite"] = int(dale_conf["num_excite"]) + + stats_meta = dict(stats_dict) + for key in ("num_excitatory", "num_inhibitory", "num_neurons", "num_steps"): + stats_meta.pop(key, None) + + spikesD = { + "spikes": spikes.astype(np.int32), + "single_rates": rates_dict["single_rates"], + } + spikesMD = { + "data_type": "simPrism", + "time_step_sec": evol_conf_in["step_size"], + 'poisson_eta_clip': evol_conf_in['poisson_eta_clip'], + 'provenance': proven_in + } + proven_in['state_transition_file']= args.dataName + + prismTruthD = { + "A_true": daleD["A_true"], + "B_true": daleD["B_true"], + "S_true": S_true.astype(np.int32), + "C_true": C_true.astype(np.float32), + "S_oracle": S_oracle.astype(np.int32), + "state_transition": transition_matrix.astype(np.float32), + "sigle_rates_var": rates_dict["sigle_rates_var"], + "single_fano_fact": rates_dict["single_fano_fact"], + } + + + prismTruthMD = { + "short_name": args.dataName, + "data_type": "simPrism", + "var_time_window_sec": float(var_time_window_sec), + "dale_conf": dale_conf, + "evol_conf": evol_conf, + } + prismTruthMD['oracle_eval'] = { + "avr_score": oracle_score, + "score_per_state": oracle_score_per_state, + } + + + # "dale_simu_stats": [stats_meta], + outFs = os.path.join(args.outPath, args.dataName + ".spikes.npz") + outFt = os.path.join(args.outPath, args.dataName + ".prismTruth.npz") + write_data_npz(spikesD, outFs, metaD=spikesMD) + write_data_npz(prismTruthD, outFt, metaD=prismTruthMD) + + if args.verb > 1: + print('\nspikes MD:'); pprint(spikesMD) + print('\nprismTruth MD:'); pprint(prismTruthMD) + + if oracle_score is not None: + target_bins_per_state = np.bincount(S_true, minlength=n_states) + switches_to_per_state = np.zeros(n_states, dtype=np.int64) + if S_true.shape[0] > 1: + switch_idx = np.where(S_true[1:] != S_true[:-1])[0] + 1 + if switch_idx.size > 0: + switches_to_per_state = np.bincount( + S_true[switch_idx], minlength=n_states + ).astype(np.int64) + print(f"gen, oracle avr score {oracle_score:.3f}, {args.dataName}") + print(f" {'state':>5s} {'score':>5s} {'bins':>7s} {'switches_to':>11s}") + print(f" {'-----':>5s} {'-----':>5s} {'-------':>7s} {'-----------':>11s}") + for m, sc in enumerate(oracle_score_per_state): + print( + f" {m:5d} {sc:5.3f} {int(target_bins_per_state[m]):7d} " + f"{int(switches_to_per_state[m]):11d}" + ) + + print("\n ./view_spikesTrain3.py --basePath $basePath --dataName %s --idxState -1 --time_range_sec 0 15 -p b " % args.dataName) + print("\n ./prism_Estep_train.py --basePath $basePath --dataName %s " % args.dataName) + print("\n ./prism_Mstep_train3.py --basePath $basePath --dataName %s " % args.dataName) + + print(" ./fit_lassoPoisson3.py --basePath $basePath --dataName %s --num_epochs 300 " % args.dataName) + + print(" ./fitPrismEM.sh --basePath $basePath --dataName %s --num_states %d --num_em_iters 4 --m_epochs 16 --time_range_sec 0 80 " % (args.dataName,n_states)) + print(" ./bigLassoBoots.sh --basePath $basePath --dataName %s --num_epochs 100 --dropDataFrac 0.33 --num_bootstraps 2 --bootsTag b2 --desyncTime " % args.dataName) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/prep_bioexp3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/prep_bioexp3c.py new file mode 100755 index 00000000..1c57178c --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/prep_bioexp3c.py @@ -0,0 +1,384 @@ +#!/usr/bin/env python3 +""" +Preprocessing pipeline for experimental neural data from Roy/Mandar laboratory. + +This script processes raw experimental neural recordings into standardized format +for connectivity analysis. The preprocessing pipeline includes: +- Raw data loading and format conversion +- Temporal binning and spike count extraction +- Data quality assessment and filtering +- Metadata extraction and session identification +- Metrics spreadsheet ingestion (metrics_curated.xlsx) +- Output formatting for downstream analysis tools + +Session naming convention: +- B6J: cell line name +- 250619: recording date (YYMMDD) +- M08020: chip identifier +- 000093: run number +- Well000: well number + +The script generates .spikes.npz files with standardized spike count matrices +and associated metadata for further analysis. + +Usage: + ./prep_bioexp.py --sessionName B6J_250619_M08020_000093_Well000 --inputPath /path/to/raw/data/ +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" +import sys, os, hashlib +import numpy as np +import pandas as pd +from pprint import pprint +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + +import argparse + + +#...!...!.................. +def commandline_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v", "--verb", type=int, help="increase debug verbosity", default=1) + parser.add_argument("--expPath", required=True, help="raw experimnetal data on CFS") + + parser.add_argument("--sessionName", default='celllinename/dateofrecording/chipID/Assaytype/runnumber/wellnumber', help='raw data session name') + parser.add_argument("--dataPath", default='/pscratch/sd/b/balewski/2025_causalNet_tmp/', help="head dir for any further data processing") + parser.add_argument("--shortName", default=None, help='(optional) output file name - Is it needed?') + parser.add_argument("--inpFormat", default=1, type=int, choices=[1, 2], + help="input format selector: 1=metrics_curated.xlsx, 2=quality_metrics.xlsx") + + # .... activity speciffic speciffic, + parser.add_argument('--samp_freq', default=100, type=int, help='sets binning of time axis') + + parser.add_argument('--freqRange', default=[1., 50], type=float, nargs=2, help='rebin of raw time axis') + + args = parser.parse_args() + + for arg in vars(args): + print('myArgs:', arg, getattr(args, arg)) + + return args + + +#...!...!.................... +def buildBioMeta(args): + pd = {} # payload + pd['raw_bioexp_path'] = args.expPath + pd['session_name'] = args.sessionName + pd['inp_format'] = args.inpFormat + txtL = args.sessionName.split('/') + #print('tt',txtL); aa + pd['cell_line_name'] = txtL[0] + pd['recording_date'] = txtL[1] + pd['chip_ID'] = txtL[2] + pd['run_num'] = txtL[4] + pd['well_num'] = txtL[5] + + sel = {'freq_range': args.freqRange} + md = {'bioexp': pd, 'data_selector': sel} + myHN = hashlib.md5(os.urandom(32)).hexdigest()[:6] + md['hash'] = myHN + if args.shortName == None: + md['short_name'] = '%s-%s' % (pd['recording_date'], md['hash']) + else: + md['short_name'] = args.shortName + + if args.verb > 1: + print('\nBMD:') + pprint(md) + return md + + +def read_spike_npy(md, args): + pmd = md['bioexp'] + inpF = os.path.join(args.expPath, args.sessionName, 'spike_times.npy') + print('inpF:', inpF) + assert os.path.exists(inpF) + # Load the dictionary from the .npy file + spike_dict = np.load(inpF, allow_pickle=True).item() + + pmd['sampling_freq'] = args.samp_freq + + # neuron ID MEA chip + meaIdL = np.array(sorted(spike_dict)) # here order of feature_id is settled + maxFeat = len(meaIdL) + + if args.verb > 1: + print('RSN: meaID list:', meaIdL) + pmd['num_feature'] = len(meaIdL) + + spikeT = {} # spike times + spikeCntL = np.zeros(pmd['num_feature'], dtype=int) # num spikes per neuron + maxTbin = 0 + + j = 0 + for k in meaIdL: + rec = np.array(spike_dict[k]) * args.samp_freq + #print(rec[:100],len(rec)) + spikeT[k] = rec.astype(int) # time-bin may repeat + spikeCntL[j] = len(rec) + j += 1 + if len(rec) == 0: + continue + mxTb = np.max(rec) + if maxTbin < mxTb: + maxTbin = mxTb + #exit(0) + pmd['num_time_bin'] = int(maxTbin) + 1 + pmd['max_time'] = pmd['num_time_bin'] / pmd['sampling_freq'] + chanFreq = spikeCntL / pmd['max_time'] + rawD = {'spikeT': spikeT, 'chanFreq': chanFreq, + 'spike_key': meaIdL, 'MEA_idx': meaIdL.copy()} + return rawD + + +def _mea_match_key(val): + """Canonical key for matching MEA_idx across npy dict keys and spreadsheet.""" + if isinstance(val, (bytes, np.bytes_)): + val = val.decode() + if isinstance(val, (int, np.integer)): + return str(int(val)) + if isinstance(val, (float, np.floating)): + if np.isnan(val): + raise ValueError("MEA_idx is NaN in metrics spreadsheet") + f = float(val) + return str(int(f)) if f == int(f) else str(f) + s = str(val).strip() + try: + f = float(s) + return str(int(f)) if f == int(f) else s + except ValueError: + return s + + +def metrics_xlsx_name(args): + return "quality_metrics.xlsx" if args.inpFormat == 2 else "metrics_curated.xlsx" + + +def _read_metrics_frame(args): + metrics_name = metrics_xlsx_name(args) + xlsxF = os.path.join(args.expPath, args.sessionName, metrics_name) + print("metrics xlsx:", xlsxF) + assert os.path.exists(xlsxF), f"missing metrics spreadsheet: {xlsxF}" + + df = pd.read_excel(xlsxF, engine="openpyxl") + cols = list(df.columns) + cols[0] = "MEA_idx" + df.columns = cols + + if args.inpFormat == 2: + rename_map = { + "location_X": "loc_x", + "location_Y": "loc_y", + "location_Z": "loc_z", + } + df = df.rename(columns={k: v for k, v in rename_map.items() if k in df.columns}) + else: + assert "MEA_idx" in df.columns, "first spreadsheet column must be MEA_idx" + + return df, metrics_name + + +def _is_integral_key(val): + try: + fval = float(val) + except (TypeError, ValueError): + return False + return np.isfinite(fval) and fval == int(fval) + + +def _integer_if_possible(vals): + arr = np.asarray(vals) + try: + farr = arr.astype(np.float64) + except (TypeError, ValueError): + return arr + if np.all(np.isfinite(farr)) and np.all(farr == farr.astype(np.int64)): + return farr.astype(np.int64) + return arr + + +def load_metrics_curated(args, mea_idx_order, spike_key_order): + """Load metrics spreadsheet; filter and reorder rows to match output neuron order.""" + df, metrics_name = _read_metrics_frame(args) + + mea_keys = [_mea_match_key(v) for v in df["MEA_idx"].values] + if len(mea_keys) != len(set(mea_keys)): + raise ValueError(f"duplicate MEA_idx rows in {metrics_name}") + + row_by_mea = {_k: i for i, _k in enumerate(mea_keys)} + order_idx = [] + for mea_id in np.asarray(mea_idx_order).ravel(): + key = _mea_match_key(mea_id) + assert key in row_by_mea, ( + f"MEA_idx {mea_id!r} (key={key!r}) not found in {metrics_name}" + ) + order_idx.append(row_by_mea[key]) + + df_ord = df.iloc[order_idx].reset_index(drop=True) + col_names = [str(c) for c in df_ord.columns] + metrics_2d = df_ord.to_numpy() + mea_idx_order = _integer_if_possible(df_ord["MEA_idx"].to_numpy()) + + n_match = len(order_idx) + n_sheet = len(df) + print( + "%s: kept %d/%d rows, shape=%s, cols=%d" + % (metrics_name, n_match, n_sheet, metrics_2d.shape, len(col_names)) + ) + if args.verb > 1: + print(" columns:", col_names) + return metrics_2d, col_names, mea_idx_order + + +#...!...!.................... +def unroll_bioexp(rawD, bioMD, args): + pmd = bioMD['bioexp'] + sel = bioMD['data_selector'] + frLo, frHi = sel['freq_range'] + print('frLo, frHi', frLo, frHi) + assert frLo < frHi + chanFreq = np.asarray(rawD['chanFreq'], dtype=float) + # vectorized boolean mask for channels within (frLo, frHi) range + freqMask = (chanFreq >= frLo) & (chanFreq <= frHi) + sel['num_drop_neur_lo_hi_freq'] = [int(np.sum(chanFreq < frLo)), int(np.sum(chanFreq > frHi))] + print('freqMask all=%d , passed=%d' % (freqMask.shape[0], np.sum(freqMask))) + #print(sel);aaa + # --- drop channles out of freq range + chanFreq = rawD['chanFreq'][freqMask] + MEA_idx = rawD['MEA_idx'][freqMask] + spike_key = rawD['spike_key'][freqMask] + + # .... REMAP MATRICES TO FREQUENCY-SORTED ORDER (PRIMARY INDEX) + neur_freqIdx = np.argsort(chanFreq) # indices that sort chanFreq by value + neur_revFreqIdx = np.empty(len(neur_freqIdx), dtype=int) # natural_index -> freq_sorted_position + neur_revFreqIdx[neur_freqIdx] = np.arange(len(neur_freqIdx)) + + #--- reorder channles by frequency + chanFreq = chanFreq[neur_freqIdx] + MEA_idx = MEA_idx[neur_freqIdx] + spike_key = spike_key[neur_freqIdx] + print('chanFreq', chanFreq[:5], '...', chanFreq[-5:], 'Hz') + + # create spike matrix: rows=time bins, cols=accepted channels + ntime = pmd['num_time_bin'] + nchan = spike_key.shape[0] + spikes2D = np.zeros((ntime, nchan), dtype=np.int32) + spikeT = rawD['spikeT'] + for ic, key in enumerate(spike_key): + tV = spikeT[key] + if len(tV) == 0: + continue + # tV holds time-bin indices where this channel fired one or more spikes + # bincount returns a length-ntime vector with spike counts per bin (zeros elsewhere) + # minlength=ntime guarantees the vector spans the full recording duration + cnt = np.bincount(tV, minlength=ntime) + spikes2D[:, ic] = cnt.astype(np.int32) + print('spikes2D shape', spikes2D.shape) + + # keep handy in meta for downstream + sel['num_chan'] = nchan + sel['max_spike_per_bin'] = int(np.max(spikes2D)) + + Y_uchar = np.clip(spikes2D, 0, 255).astype(np.uint8) + spikeD = {'spikes': Y_uchar, + 'single_rates': chanFreq + } + + bioD = {} + bioD['neur_freqIdx'] = neur_freqIdx + bioD['neur_revFreqIdx'] = neur_revFreqIdx + bioD['single_rates'] = np.asarray(chanFreq, dtype=np.float64) + + metrics_2d, metrics_cols, MEA_idx = load_metrics_curated(args, MEA_idx, spike_key) + bioD['MEA_idx'] = MEA_idx + bioD['spike_key'] = spike_key + bioD['metrics_curated'] = metrics_2d + bioMD['metrics_curated_columns'] = metrics_cols + metrics_col_map = {str(c): i for i, c in enumerate(metrics_cols)} + assert "loc_x" in metrics_col_map and "loc_y" in metrics_col_map, ( + f"metrics_curated must contain loc_x and loc_y columns; have {metrics_cols}" + ) + loc_x = metrics_2d[:, metrics_col_map["loc_x"]].astype(np.float64) + loc_y = metrics_2d[:, metrics_col_map["loc_y"]].astype(np.float64) + bioD['node_positions'] = np.column_stack([loc_x, loc_y]) + + #.... compute neural statistics for spikeMD + num_neurons = nchan + avg_rate = float(np.mean(chanFreq)) + std_rate = float(np.std(chanFreq)) + median_rate = float(np.median(chanFreq)) + min_rate = float(np.min(chanFreq)) + max_rate = float(np.max(chanFreq)) + + # Compute Fano factor (variance/mean) for each neuron + mean_counts_per_bin = np.mean(spikes2D, axis=0) + spike_variance = np.var(spikes2D, axis=0) + fano_factor = np.divide(spike_variance, mean_counts_per_bin, out=np.zeros_like(spike_variance), where=mean_counts_per_bin != 0) + avg_fano = float(np.mean(fano_factor)) + std_fano = float(np.std(fano_factor)) + + # Print summary statistics + print('Neural Statistics Summary:') + print('num neurons: %d, Avg Rate= %.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_neurons, avg_rate, std_rate, avg_fano, std_fano)) + print('Median rate %.2f Hz' % median_rate) + + #.... extract spikeMD for fitter + spikeMD = {'time_step_sec': 1. / pmd['sampling_freq'], + 'provenance': {'experiment_name': bioMD['short_name']}, + 'poisson_eta_clip': 5, # expected by fitter + 'data_type': 'bioExp', 'num_neurons': num_neurons} + + bioMD['rate_summary'] = { + 'avg_spike_rate': avg_rate, + 'std_spike_rate': std_rate, + 'avg_fano_factor': avg_fano, + 'std_fano_factor': std_fano, + 'median_spike_rate': median_rate, + 'min_spike_rate': min_rate, + 'max_spike_rate': max_rate + } + return bioD, spikeD, spikeMD + + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__ == "__main__": + + args = commandline_parser() + np.set_printoptions(precision=3) + bioMD = buildBioMeta(args) + + # read raw data + #rawD=read_spike_dict(bioMD,args) + rawD = read_spike_npy(bioMD, args) + + #.... filter & unroll data + bioD, spikeD, spikeMD = unroll_bioexp(rawD, bioMD, args) + + #...... WRITE OUTPUT ......... + outFt = os.path.join(args.dataPath, bioMD['short_name'] + '.bioExp.npz') + write_data_npz(bioD, outFt, metaD=bioMD) + if args.verb > 2: + print('\n bioD:', sorted(bioD)) + pprint(bioMD) + + outFs = outFt.replace('.bioExp.', '.spikes.') + write_data_npz(spikeD, outFs, metaD=spikeMD) + if args.verb > 2: + print('\nspikeD:', sorted(spikeD)) + pprint(spikeMD) + + print("\nNext step command:") + print(' ./view_bioexp.py --dataPath $dataPath --dataName %s -p a b c -T 0 3550 ' % (bioMD['short_name'])) + print(" ./fit_lassoPoisson.py --dataPath $dataPath --dataName %s --num_epochs 10 " % bioMD['short_name']) + print(" ./fitLasso4GPU.sh --dataPath $dataPath --dataName %s --num_epochs 200 " % bioMD['short_name']) + print(" ./bootsFit.sh --dataName %s --num_epochs 250 --dropDataFrac 0.5 --num_bootstraps 7 " % bioMD['short_name']) + + print(" ./selectEdges_FDR.py --dataName %s --num_bootstraps 6 10 -p a c d " % bioMD['short_name']) + + print(' --dataPath ' + args.dataPath) diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/prism_EM_FDR_Bags_aggregate3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/prism_EM_FDR_Bags_aggregate3c.py new file mode 100755 index 00000000..e719e7fe --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/prism_EM_FDR_Bags_aggregate3c.py @@ -0,0 +1,591 @@ +#!/usr/bin/env python3 +"""Aggregate EM-FDR-bagging Stage (a) outputs into one eval-compatible fit.""" + +import argparse +import itertools +import os +import secrets +import time + +import numpy as np + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + + +BASE_FIT_KEYS = ( + "A_init", + "A_hat", + "A_prune", + "neuron_type", + "neuron_Sedge", + "B_init", + "B_hat", + "freq_h1d", + "c_init", + "c_hat", + "S_init", + "S_hat", + "S_hat_CL", + "single_rates", + "e_nll_em", + "m_loss_epoch", + "m_nll_epoch", + "m_l1_epoch", + "rho_epoch", + "rho_correction_strength_epoch", + "nz_edges_epoch", + "learning_rates", +) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="PRISM-EM FDR bagging Stage (b): aggregate bag files", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--basePath", required=True, + help="Run directory containing prismFDR/ and prismFit/") + parser.add_argument("--dataName", required=True, + help="Stage (a) bag input stem in prismFDR/") + parser.add_argument("--outAgrName", default=None, + help="Output aggregate fit stem written to prismFit/; default is dataName plus random _hash4") + parser.add_argument("--num_bags", type=int, required=True, + help="Number of bag files to aggregate; e.g. 2 means bag000 and bag001") + parser.add_argument("--per_bag_quantile", type=float, default=0.99, + help="Per-source-column null magnitude quantile for bag selection") + parser.add_argument("--stab_sel_thresh", type=float, default=0.7, + help="Cross-bag stability selection frequency threshold") + parser.add_argument("--fdr_out_dir", default=None, + help="Directory containing Stage (a) bag files") + parser.add_argument("-v", "--verb", type=int, default=1) + return parser.parse_args() + + +def bag_file_name(data_name, bag_idx): + return f"{data_name}.bag{int(bag_idx):03d}.prismFDRbag.npz" + + +def bag_indices_from_count(num_bags): + return list(range(int(num_bags))) + + +def source_type_prune(A_hat, min_posW, max_negW): + """Compute source-neuron signs and Dale-style pruned A using columns.""" + A_hat = np.asarray(A_hat, dtype=np.float32) + n_neuron = A_hat.shape[0] + A_thr = A_hat.copy() + np.fill_diagonal(A_thr, 0.0) + + neuron_sedge = A_thr.sum(axis=0) + neuron_type = np.zeros((n_neuron,), dtype=np.int8) + neuron_type[neuron_sedge > float(min_posW)] = 1 + neuron_type[neuron_sedge < float(max_negW)] = -1 + + A_prune = A_hat.copy() + diag_A = np.diag(A_hat).copy() + exc_cols = neuron_type > 0 + inh_cols = neuron_type < 0 + A_prune[:, exc_cols] = np.where(A_prune[:, exc_cols] > 0, A_prune[:, exc_cols], 0.0) + A_prune[:, inh_cols] = np.where(A_prune[:, inh_cols] < 0, A_prune[:, inh_cols], 0.0) + np.fill_diagonal(A_prune, diag_A) + return A_prune.astype(np.float32), neuron_type, neuron_sedge.astype(np.float32) + + +def selected_weight_thresholds(A_hat, selected_mask): + A = np.asarray(A_hat, dtype=np.float64) + mask = np.asarray(selected_mask, dtype=bool) + pos = A[mask & (A > 0.0)] + neg = A[mask & (A < 0.0)] + min_posW = float(np.min(pos)) if pos.size else float("nan") + max_negW = float(np.max(neg)) if neg.size else float("nan") + return min_posW, max_negW + + +def load_bags(args): + inp_dir = args.fdr_out_dir + if inp_dir is None: + inp_dir = os.path.join(args.basePath, "prismFDR") + + bag_data = [] + bag_meta = [] + bag_files = [] + for bag_idx in bag_indices_from_count(args.num_bags): + inp_f = os.path.join(inp_dir, bag_file_name(args.dataName, bag_idx)) + if not os.path.exists(inp_f): + raise FileNotFoundError(f"Missing Stage (a) bag file: {inp_f}") + data_d, meta_d = read_data_npz(inp_f, verb=args.verb > 1) + if meta_d is None: + raise ValueError(f"Bag file has no metadata: {inp_f}") + bag_data.append(data_d) + bag_meta.append(meta_d) + bag_files.append(inp_f) + if args.verb > 0: + print(f"loaded bag {bag_idx:03d}: {inp_f}") + return bag_data, bag_meta, bag_files, inp_dir + + +def validate_bags(bag_data): + if not bag_data: + raise ValueError("No bags loaded") + n0 = np.asarray(bag_data[0]["A_hat"]).shape[0] + b0 = np.asarray(bag_data[0]["B_hat"]) + if b0.ndim == 1: + b0 = b0[None, :] + m0 = b0.shape[0] + for ib, data_d in enumerate(bag_data): + a_hat = np.asarray(data_d["A_hat"]) + a_null = np.asarray(data_d["A_null"]) + b_hat = np.asarray(data_d["B_hat"]) + if b_hat.ndim == 1: + b_hat = b_hat[None, :] + if a_hat.shape != (n0, n0): + raise ValueError(f"bag {ib}: A_hat shape {a_hat.shape} != {(n0, n0)}") + if a_null.ndim != 3 or a_null.shape[1:] != (n0, n0): + raise ValueError(f"bag {ib}: A_null shape {a_null.shape} incompatible with N={n0}") + if b_hat.shape != (m0, n0): + raise ValueError(f"bag {ib}: B_hat shape {b_hat.shape} != {(m0, n0)}") + return n0, m0 + + +def align_state_rows_to_reference(b_hat_stack): + """Align B state rows to bag 0 by minimum squared distance.""" + b_hat_stack = np.asarray(b_hat_stack, dtype=np.float64) + n_bag, n_state, _ = b_hat_stack.shape + ref = b_hat_stack[0] + aligned = np.empty_like(b_hat_stack) + perms = np.zeros((n_bag, n_state), dtype=np.int64) + aligned[0] = ref + perms[0] = np.arange(n_state, dtype=np.int64) + + all_perms = None + if n_state <= 8: + all_perms = list(itertools.permutations(range(n_state))) + + for ib in range(1, n_bag): + cur = b_hat_stack[ib] + if all_perms is not None: + best_perm = None + best_score = None + for perm in all_perms: + diff = ref - cur[np.asarray(perm)] + score = float(np.sum(diff * diff)) + if best_score is None or score < best_score: + best_score = score + best_perm = perm + perm_arr = np.asarray(best_perm, dtype=np.int64) + else: + remaining = set(range(n_state)) + perm = [] + for m in range(n_state): + best_j = min( + remaining, + key=lambda j: float(np.sum((ref[m] - cur[j]) ** 2)), + ) + perm.append(best_j) + remaining.remove(best_j) + perm_arr = np.asarray(perm, dtype=np.int64) + aligned[ib] = cur[perm_arr] + perms[ib] = perm_arr + return aligned.astype(np.float32), perms + + +def sd_or_nan(x, axis): + x = np.asarray(x, dtype=np.float64) + if x.shape[axis] < 2: + return np.full(x.shape[:axis] + x.shape[axis + 1:], np.nan, dtype=np.float64) + return np.std(x, axis=axis, ddof=1) + + +def aggregate_edges(a_bag, a_null, per_bag_quantile, stab_sel_thresh): + """Return Stage (b) matrices and diagnostics.""" + a_bag = np.asarray(a_bag, dtype=np.float64) + a_null = np.asarray(a_null, dtype=np.float64) + n_bag, n, n2 = a_bag.shape + if n != n2: + raise ValueError("A_hat stack must be square") + + off_mask = ~np.eye(n, dtype=bool) + tau = np.zeros((n_bag, n), dtype=np.float64) + null_mean_bag = np.zeros((n_bag, n), dtype=np.float64) + null_sd_bag = np.zeros((n_bag, n), dtype=np.float64) + + for ib in range(n_bag): + for j in range(n): + vals = a_null[ib, :, :, j][:, off_mask[:, j]].reshape(-1) + tau[ib, j] = float(np.quantile(np.abs(vals), per_bag_quantile)) + null_mean_bag[ib, j] = float(np.mean(vals)) + null_sd_bag[ib, j] = float(np.std(vals, ddof=1)) if vals.size > 1 else np.nan + + sel = (np.abs(a_bag) > tau[:, None, :]) & off_mask[None, :, :] + sel_count = sel.sum(axis=0).astype(np.int64) + sel_freq = sel_count.astype(np.float64) / float(n_bag) + final_sel = (sel_freq >= float(stab_sel_thresh)) & off_mask + + a_mean_all = np.mean(a_bag, axis=0) + a_sd_all = sd_or_nan(a_bag, axis=0) + + a_sum_sel = np.sum(np.where(sel, a_bag, 0.0), axis=0) + with np.errstate(invalid="ignore", divide="ignore"): + a_mean_sel = a_sum_sel / sel_count + a_mean_sel[sel_count == 0] = np.nan + + a_sd_sel = np.full((n, n), np.nan, dtype=np.float64) + for i in range(n): + for j in range(n): + if i == j: + continue + vals = a_bag[sel[:, i, j], i, j] + if vals.size > 1: + a_sd_sel[i, j] = float(np.std(vals, ddof=1)) + elif vals.size == 1: + a_sd_sel[i, j] = 0.0 + + src_null_mean = np.zeros((n,), dtype=np.float64) + src_null_sd = np.zeros((n,), dtype=np.float64) + for j in range(n): + vals = a_null[:, :, :, j][:, :, off_mask[:, j]].reshape(-1) + src_null_mean[j] = float(np.mean(vals)) + src_null_sd[j] = float(np.std(vals, ddof=1)) if vals.size > 1 else np.nan + + a_hat_final = np.zeros((n, n), dtype=np.float64) + a_hat_final[final_sel] = a_mean_sel[final_sel] + diag = np.diag(a_mean_all).copy() + np.fill_diagonal(a_hat_final, diag) + + z_null = np.full((n, n), np.nan, dtype=np.float64) + denom = src_null_sd[None, :] + with np.errstate(invalid="ignore", divide="ignore"): + z_null = (a_mean_sel - src_null_mean[None, :]) / denom + z_null[~off_mask] = np.nan + + selected_edges_per_bag_src = sel.sum(axis=1).astype(np.int64) + selected_edges_per_bag = sel.sum(axis=(1, 2)).astype(np.int64) + p_cand = int(n * (n - 1)) + q_lambda = float(np.mean(selected_edges_per_bag)) + if float(stab_sel_thresh) > 0.5: + stability_false_edge_bound = q_lambda * q_lambda / ( + float(p_cand) * (2.0 * float(stab_sel_thresh) - 1.0) + ) + else: + stability_false_edge_bound = np.inf + + return { + "A_hat_final": a_hat_final, + "selected_mask": final_sel, + "selected_in_bag": sel, + "selection_count": sel_count, + "selection_frequency": sel_freq, + "A_mean_selected": a_mean_sel, + "A_sd_selected": a_sd_sel, + "A_mean_all": a_mean_all, + "A_sd_all": a_sd_all, + "src_null_tau_bag": tau, + "src_null_tau_mean": np.mean(tau, axis=0), + "src_null_mean": src_null_mean, + "src_null_sd": src_null_sd, + "src_null_mean_bag": null_mean_bag, + "src_null_sd_bag": null_sd_bag, + "z_null": z_null, + "selected_edges_per_bag_src": selected_edges_per_bag_src, + "selected_edges_per_bag": selected_edges_per_bag, + "q_lambda": q_lambda, + "p_cand": p_cand, + "stability_false_edge_bound": float(stability_false_edge_bound), + } + + +def edge_table(final_sel, agg): + edge_i, edge_j = np.where(final_sel) + return { + "edge_i": edge_i.astype(np.int64), + "edge_j": edge_j.astype(np.int64), + "edge_sel_freq": agg["selection_frequency"][edge_i, edge_j].astype(np.float32), + "edge_A_mean": agg["A_mean_selected"][edge_i, edge_j].astype(np.float32), + "edge_A_sd_boot": agg["A_sd_selected"][edge_i, edge_j].astype(np.float32), + "edge_z_null": agg["z_null"][edge_i, edge_j].astype(np.float32), + "edge_A_mean_all": agg["A_mean_all"][edge_i, edge_j].astype(np.float32), + "edge_A_sd_all": agg["A_sd_all"][edge_i, edge_j].astype(np.float32), + "edge_src_null_mean": agg["src_null_mean"][edge_j].astype(np.float32), + "edge_src_null_sd": agg["src_null_sd"][edge_j].astype(np.float32), + "edge_src_tau_mean": agg["src_null_tau_mean"][edge_j].astype(np.float32), + } + + +def real_fit_metadata(md): + if "bagsFDR_stageA" not in md: + raise KeyError("metadata missing 'bagsFDR_stageA' — was this file produced by prism_FDR_Bags_train3c.py?") + if "real_fit" not in md["bagsFDR_stageA"]: + raise KeyError("bagsFDR_stageA missing 'real_fit' — re-run prism_FDR_Bags_train3c.py to regenerate bag files") + return md["bagsFDR_stageA"]["real_fit"] + + +def eval_loss_time(fitD, md, spikes): + """Precompute per-time fit-loss diagnostics for plotting.""" + trainMD = real_fit_metadata(md)["train"] + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + dt = float(trainMD["time_step_sec"]) + eta_clip = float(trainMD["eta_clip"]) + lambda2 = float(trainMD["lambda2"]) + + spikes_sub = np.asarray(spikes[t0_bin:t1_bin + 1], dtype=np.float64) + yp = spikes_sub[:-1] + yc = spikes_sub[1:] + + a_hat = np.asarray(fitD["A_hat"], dtype=np.float64) + b_hat = np.asarray(fitD["B_hat"], dtype=np.float64) + if b_hat.ndim == 1: + b_hat = b_hat[None, :] + c_hat = np.asarray(fitD["c_hat"], dtype=np.float64) + assert c_hat.shape[0] == spikes_sub.shape[0], "c_hat and spikes_sub must have matching time bins" + c_pairs = c_hat[1:] + c_prev = c_hat[:-1] + + rates = np.asarray(fitD["single_rates"], dtype=np.float64) + w = 1.0 / np.maximum(rates, 0.1) + w /= w.mean() + + eta = yp @ a_hat.T + c_pairs @ b_hat + eta_c = np.minimum(eta, eta_clip) + lam = np.exp(eta_c) * dt + log_dt = np.log(dt) + nll_t = np.sum(w[None, :] * (lam - yc * (eta + log_dt)), axis=1) + + dc = c_pairs - c_prev + l2_t = lambda2 * np.sum(dc * dc, axis=1) + + return { + "loss_nll_time": nll_t.astype(np.float32), + "loss_l2_time": l2_t.astype(np.float32), + } + + +def eval_true_state_recovery(fitD, md): + """Precompute state-recovery table for plotting.""" + trainMD = real_fit_metadata(md)["train"] + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + s_true = np.asarray(md["S_true"], dtype=np.int64)[t0_bin:t1_bin + 1] + s_hat = np.asarray(fitD["S_hat"], dtype=np.int64) + s_hat_cl = np.asarray(fitD["S_hat_CL"], dtype=np.float64) + n_state = int(trainMD["num_states"]) + + assert s_true.shape[0] == s_hat.shape[0] == s_hat_cl.shape[0], \ + "S_true/S_hat/S_hat_CL length mismatch on training window" + + state_acc_cl = [] + bins_hat = np.bincount(s_hat, minlength=n_state).astype(np.int64) + enter_hat = np.zeros(n_state, dtype=np.int64) + if s_hat.shape[0] > 1: + enter_idx = np.where(s_hat[1:] != s_hat[:-1])[0] + 1 + if enter_idx.size > 0: + enter_hat = np.bincount(s_hat[enter_idx], minlength=n_state).astype(np.int64) + + for m in range(n_state): + mask = (s_hat == m) + if int(mask.sum()) > 0: + acc_m = float((s_true[mask] == m).mean()) + cl_m = float(s_hat_cl[mask].mean()) + else: + acc_m = float("nan") + cl_m = float("nan") + state_acc_cl.append([acc_m, cl_m, int(bins_hat[m]), int(enter_hat[m])]) + + return { + "avg_acc": float((s_true == s_hat).mean()), + "state_acc_cl": state_acc_cl, + } + + +def source_spike_file(base_path, out_md): + prov = out_md["provenance"] + if out_md.get("data_type") == "bioExp": + data_name = prov["experiment_name"] + else: + data_name = prov["state_transition_file"] + return os.path.join(base_path, "spikesData", f"{data_name}.spikes.npz") + + +def add_fit_eval_metadata(out_d, out_md, base_path, verb=1): + """Attach display metrics to the aggregate output metadata.""" + spike_f = source_spike_file(base_path, out_md) + spike_d, _ = read_data_npz(spike_f, verb=verb > 1) + + eval_f = eval_loss_time(out_d, out_md, spike_d["spikes"]) + out_d["loss_nll_time"] = eval_f.pop("loss_nll_time") + out_d["loss_l2_time"] = eval_f.pop("loss_l2_time") + out_md["eval_f"] = { + **eval_f, + "loss_nll_time_key": "loss_nll_time", + "loss_l2_time_key": "loss_l2_time", + } + + if out_md.get("data_type") == "bioExp": + return out_md + + prov = out_md["provenance"] + st_name = prov["state_transition_file"] + truth_f = os.path.join(base_path, "spikesData", f"{st_name}.prismTruth.npz") + + truth_d, _ = read_data_npz(truth_f, verb=verb > 1) + + md_eval = dict(out_md) + md_eval["S_true"] = truth_d["S_true"] + md_eval["C_true"] = truth_d["C_true"] + + srec = real_fit_metadata(out_md)["states_recovery_eval"] + reco = eval_true_state_recovery(out_d, md_eval) + srec.update(reco) + out_md["eval_f"]["acc"] = reco["avg_acc"] + return out_md + + +def main(): + args = parse_args() + if args.outAgrName is None: + args.outAgrName = f"{args.dataName}_{secrets.token_hex(2)}" + if args.num_bags < 1: + raise ValueError("--num_bags must be at least 1") + if not (0.0 < float(args.per_bag_quantile) < 1.0): + raise ValueError("--per_bag_quantile must be between 0 and 1") + if not (0.0 < float(args.stab_sel_thresh) <= 1.0): + raise ValueError("--stab_sel_thresh must be in (0, 1]") + if args.verb > 0: + print("\nFDR-bag aggregate args:", vars(args), "\n") + + t0 = time.perf_counter() + bag_data, bag_meta, bag_files, inp_dir = load_bags(args) + n_neuron, n_state = validate_bags(bag_data) + + a_bag = np.stack([np.asarray(d["A_hat"], dtype=np.float32) for d in bag_data], axis=0) + a_null = np.stack([np.asarray(d["A_null"], dtype=np.float32) for d in bag_data], axis=0) + b_bag = np.stack([ + np.asarray(d["B_hat"], dtype=np.float32) + if np.asarray(d["B_hat"]).ndim == 2 + else np.asarray(d["B_hat"], dtype=np.float32)[None, :] + for d in bag_data + ], axis=0) + b_aligned = b_bag + state_perms = np.tile(np.arange(n_state, dtype=np.int64), (int(args.num_bags), 1)) + + agg = aggregate_edges( + a_bag, a_null, + per_bag_quantile=float(args.per_bag_quantile), + stab_sel_thresh=float(args.stab_sel_thresh), + ) + + a_hat = agg["A_hat_final"].astype(np.float32) + min_posW, max_negW = selected_weight_thresholds(a_hat, agg["selected_mask"]) + a_prune, neuron_type, neuron_sedge = source_type_prune(a_hat, min_posW, max_negW) + b_hat = np.mean(b_aligned, axis=0).astype(np.float32) + + ref_d = bag_data[0] + ref_md = bag_meta[0] + out_d = {} + for key in BASE_FIT_KEYS: + if key in ref_d: + out_d[key] = ref_d[key] + + out_d["A_hat"] = a_hat + out_d["A_prune"] = a_prune + out_d["neuron_type"] = neuron_type + out_d["neuron_Sedge"] = neuron_sedge + out_d["B_hat"] = b_hat + + out_d.update({ + "A_bag": a_bag.astype(np.float32), + "B_bag": b_aligned.astype(np.float32), + "state_permutation_to_ref": state_perms.astype(np.int64), + "selected_mask": agg["selected_mask"].astype(np.bool_), + "selected_in_bag": agg["selected_in_bag"].astype(np.bool_), + "selection_count": agg["selection_count"].astype(np.int64), + "selection_frequency": agg["selection_frequency"].astype(np.float32), + "A_mean_selected": agg["A_mean_selected"].astype(np.float32), + "A_sd_selected": agg["A_sd_selected"].astype(np.float32), + "A_mean_all": agg["A_mean_all"].astype(np.float32), + "A_sd_all": agg["A_sd_all"].astype(np.float32), + "A_diag_mean": np.diag(agg["A_mean_all"]).astype(np.float32), + "A_diag_std": np.diag(agg["A_sd_all"]).astype(np.float32), + "A_diag_stderr": (np.diag(agg["A_sd_all"]) / np.sqrt(float(args.num_bags))).astype(np.float32), + "src_null_tau_bag": agg["src_null_tau_bag"].astype(np.float32), + "src_null_tau_mean": agg["src_null_tau_mean"].astype(np.float32), + "src_null_mean": agg["src_null_mean"].astype(np.float32), + "src_null_sd": agg["src_null_sd"].astype(np.float32), + "src_null_mean_bag": agg["src_null_mean_bag"].astype(np.float32), + "src_null_sd_bag": agg["src_null_sd_bag"].astype(np.float32), + "z_null": agg["z_null"].astype(np.float32), + "selected_edges_per_bag_src": agg["selected_edges_per_bag_src"].astype(np.int64), + "selected_edges_per_bag": agg["selected_edges_per_bag"].astype(np.int64), + }) + out_d.update(edge_table(agg["selected_mask"], agg)) + + out_name = args.outAgrName + out_dir = os.path.join(args.basePath, "prismFit") + os.makedirs(out_dir, exist_ok=True) + out_f = os.path.join(out_dir, f"{out_name}.prismEM.npz") + + out_md = dict(ref_md) + out_md["fit_type"] = "prismEM_FDRbags_stageB" + out_md["bagsFDR_stageB"] = { + "program": "prism_EM_FDR_Bags_aggregate3c.py", + "dataName": args.dataName, + "input_dataName": args.dataName, + "outAgrName": out_name, + "output_name": out_name, + "num_bags": int(args.num_bags), + "expected_bag_indices": bag_indices_from_count(args.num_bags), + "per_bag_quantile": float(args.per_bag_quantile), + "stab_sel_thresh": float(args.stab_sel_thresh), + "fdr_out_dir": inp_dir, + "input_files": bag_files, + "reference_state_source": "reference_EM_fit_locked_in_stageA", + "reference_time_fields": [ + "A_init", "B_init", "freq_h1d", "c_init", "c_hat", + "S_init", "S_hat", "S_hat_CL", "EM histories", + ], + "null_threshold_axis": "source_column", + "A_hat_meaning": "selection_conditional_bagged_mean_for_stable_edges_zero_elsewhere", + "A_hat_diagonal": "all_bag_mean_diagonal", + "B_hat_meaning": "mean_across_bags_in_reference_state_order", + "state_permutation_to_ref": state_perms.tolist(), + "num_neurons": int(n_neuron), + "num_states": int(n_state), + "num_final_edges": int(np.sum(agg["selected_mask"])), + "min_posW": min_posW, + "max_negW": max_negW, + "q_lambda_mean_edges_selected_per_bag": float(agg["q_lambda"]), + "p_cand": int(agg["p_cand"]), + "stability_false_edge_bound": float(agg["stability_false_edge_bound"]), + "elapsed_sec": float(time.perf_counter() - t0), + } + + prov = dict(out_md.get("provenance", {})) + prov["EMtrain_file"] = out_name + prov["bagsFDR_stageB_file"] = out_name + prov["bagsFDR_stageB_dataName"] = args.dataName + prov["bagsFDR_stageB_outAgrName"] = out_name + out_md["provenance"] = prov + out_md = add_fit_eval_metadata(out_d, out_md, args.basePath, verb=args.verb) + + if args.verb > 0: + print( + f"Stage (b): bags={args.num_bags} N={n_neuron} " + f"q={args.per_bag_quantile:g} stab_sel_thresh={args.stab_sel_thresh:g}" + ) + print( + f" selected final edges: {int(np.sum(agg['selected_mask']))} / " + f"{int(agg['p_cand'])}" + ) + print( + f" q_lambda={agg['q_lambda']:.3f} " + f"stability_false_edge_bound={agg['stability_false_edge_bound']:.3f}" + ) + write_data_npz(out_d, out_f, metaD=out_md, verb=args.verb > 1) + if args.verb > 0: + print(f"\nSaved Stage (b) aggregate: {out_f}") + print( + f" ./prism_EM_eval3c.py --basePath $basePath " + f"--dataName {out_name} -p e f h i g" + ) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/prism_EM_eval3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/prism_EM_eval3c.py new file mode 100755 index 00000000..7db53071 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/prism_EM_eval3c.py @@ -0,0 +1,395 @@ +#!/usr/bin/env python3 +"""Evaluation and plotting for prism EM 3c results.""" + +import argparse +import os +import re +from pprint import pprint + +import numpy as np + +from toolbox.Util_NumpyIO import read_data_npz +from PlotterPrismEM3c import Plotter +from UtilBioExp import detect_spike_bursts + + +PLOT_FIG_ID = {chr(ord("a") + i): chr(ord("a") + i) for i in range(18)} +SIM_ONLY_PLOTS = set("mnopqr") +IMPLEMENTED_PLOTS = {"a", "b", "c", "d", "e", "f", "g", "h", "i", "m", "n", "o", "p", "r"} + + +def real_fit_metadata(md): + if is_stage_c_metadata(md): + return md + if "bagsFDR_stageA" in md and "real_fit" in md["bagsFDR_stageA"]: + return md["bagsFDR_stageA"]["real_fit"] + return md + + +def is_stage_c_metadata(md): + return md.get("fit_type") == "prismEM_deBias_stageC" or "deBias_stageC" in md + + +def display_fit_data(fitD, fitMD): + """Map Stage (c) output arrays onto legacy plotting keys without mutating input.""" + if not is_stage_c_metadata(fitMD): + return fitD + for key in ("A_debias", "B_debias"): + assert key in fitD, f"Stage (c) display requires {key}" + + disp = dict(fitD) + for key in ("A_hat", "A_prune", "A_init", "B_hat", "neuron_Sedge", + "A_diag_mean", "A_diag_stderr"): + if key in fitD: + disp["stageB_" + key] = fitD[key] + + A_debias = np.asarray(fitD["A_debias"]) + disp["A_hat"] = fitD["A_debias"] + disp["B_hat"] = fitD["B_debias"] + disp["A_prune"] = fitD["A_debias"] + disp["A_init"] = fitD["A_debias_init"] if "A_debias_init" in fitD else fitD["A_hat"] + if "neuron_Sedge_debias" in fitD: + disp["neuron_Sedge"] = fitD["neuron_Sedge_debias"] + else: + off_mask = ~np.eye(A_debias.shape[0], dtype=bool) + disp["neuron_Sedge"] = (A_debias * off_mask).sum(axis=0).astype(np.float32) + if "A_diag_debias" in fitD: + diag = np.asarray(fitD["A_diag_debias"], dtype=np.float32) + else: + diag = np.diag(A_debias).astype(np.float32) + disp["A_diag_mean"] = diag + disp["A_diag_stderr"] = np.zeros_like(diag, dtype=np.float32) + return disp + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Evaluate and plot prism EM 3c results", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--dataName", type=str, required=True, + help="Base name for prismEM file in basePath/prismFit") + parser.add_argument("--basePath", type=str, + default="/pscratch/sd/b/balewski/2026_causalNet_tmp3/", + help="Head dir for input/output data") + parser.add_argument("-p", "--showPlots", type=str, nargs="+", default="a", + help="Plot letters: a-l for experiments and simulations; m-r simulations only") + parser.add_argument("-T", "--time_range_sec", default=None, + nargs=2, type=float, + help="Time window [t0, t1] in seconds for time-axis plots") + parser.add_argument("-R", "--time_rebin2", default=20, type=int, + help="Burst panels: rebin factor on spike time axis") + parser.add_argument("--burst_freq_thres", default=5.0, type=float, + help="Burst panels: per-neuron rate threshold (Hz)") + parser.add_argument("--maxNeurons", default=24, type=int, + help="Spatial split panels: max source neurons per panel") + parser.add_argument("-X", "--noXterm", action="store_true", + help="Disable X terminal for plotting") + parser.add_argument("-v", "--verb", type=int, default=1, + help="Verbosity level") + return parser.parse_args() + + +def resolve_fit_file(base_path, data_name): + if re.search(r"\.bag\d{3}(?:$|\.)", data_name): + if data_name.endswith(".prismFDRbag.npz"): + fname = data_name + elif data_name.endswith(".prismFDRbag"): + fname = f"{data_name}.npz" + elif data_name.endswith(".npz"): + fname = data_name + else: + fname = f"{data_name}.prismFDRbag.npz" + fit_f = fname if os.path.isabs(fname) else os.path.join(base_path, "prismFDR", fname) + return fit_f, "prismFDR" + + fit_f = os.path.join(base_path, "prismFit", f"{data_name}.prismEM.npz") + return fit_f, "prismFit" + + +def load_fit(base_path, data_name, verb=1): + fit_f, fit_source = resolve_fit_file(base_path, data_name) + fitD, fitMD = read_data_npz(fit_f, verb=verb > 1) + assert isinstance(fitMD, dict), "Expected metadata dict in prismEM file" + return fitD, fitMD, fit_f, fit_source + + +def default_plot_time_range_sec(md, duration_sec=60.0): + trainMD = real_fit_metadata(md)["train"] + t0_sec = float(trainMD["time_range_sec"][0]) + t1_sec = min(float(trainMD["time_range_sec"][1]), t0_sec + float(duration_sec)) + return [t0_sec, t1_sec] + + +def load_simu_truth(base_path, fitD, fitMD, md, verb=1): + if "S_true" in fitD and "C_true" in fitD: + md["S_true"] = fitD["S_true"] + md["C_true"] = fitD["C_true"] + return md + + prov = fitMD["provenance"] + spikes_path = os.path.join(base_path, "spikesData") + + st_name = prov["state_transition_file"] + pt_f = os.path.join(spikes_path, f"{st_name}.prismTruth.npz") + trD, _ = read_data_npz(pt_f, verb=verb > 1) + md["S_true"] = trD["S_true"] + md["C_true"] = trD["C_true"] + return md + + +def load_simu_static_truth(base_path, fitMD, md, verb=1): + prov = fitMD["provenance"] + truth_name = prov["state_model_file"] + truth_f = os.path.join(base_path, "truthDale", f"{truth_name}.simTruth.npz") + trueD, trueMD = read_data_npz(truth_f, verb=verb > 1) + for key in ("A_true", "B_true", "E_true"): + assert key in trueD, f"{truth_f} must contain {key}" + md[key] = trueD[key] + if "dale_conf" in trueMD: + md["dale_conf"] = trueMD["dale_conf"] + return md + + +def source_spike_name(fitMD): + prov = fitMD["provenance"] + if fitMD.get("data_type") == "bioExp": + return prov["experiment_name"] + return prov["state_transition_file"] + + +def load_source_spikes(base_path, fitMD, verb=1): + spike_name = source_spike_name(fitMD) + spike_f = os.path.join(base_path, "spikesData", f"{spike_name}.spikes.npz") + spikeD, spikeMD = read_data_npz(spike_f, verb=verb > 1) + return spikeD, {**spikeMD, "short_name": spike_name}, spike_f + + +def load_node_metadata(base_path, fitMD, data_name, verb=1): + prov = fitMD["provenance"] + if fitMD.get("data_type") == "bioExp": + node_name = prov["experiment_name"] + node_f = os.path.join(base_path, "spikesData", f"{node_name}.bioExp.npz") + nodeD, nodeMD = read_data_npz(node_f, verb=verb > 1) + nodeMD = {**nodeMD, "short_name": data_name, "node_source_name": node_name} + else: + node_name = prov["state_model_file"] + node_f = os.path.join(base_path, "truthDale", f"{node_name}.simTruth.npz") + nodeD, nodeMD = read_data_npz(node_f, verb=verb > 1) + nodeMD = { + **nodeMD, + "short_name": data_name, + "data_type": fitMD.get("data_type", "simPrism"), + "node_source_name": node_name, + } + + assert "node_positions" in nodeD, ( + f"{node_f} must contain node_positions; regenerate the source metadata file" + ) + node_pos = np.asarray(nodeD["node_positions"], dtype=np.float64) + assert node_pos.ndim == 2 and node_pos.shape[1] == 2, ( + f"node_positions in {node_f} must have shape (N, 2), got {node_pos.shape}" + ) + return nodeD, nodeMD, node_f + + +def normalize_plot_letters(show_plots): + letters = "".join(show_plots) + unknown = sorted(set(letters) - set(PLOT_FIG_ID)) + assert not unknown, f"Unknown plot letters: {''.join(unknown)}" + missing = sorted(set(letters) - IMPLEMENTED_PLOTS) + assert not missing, f"Plots not implemented yet in prism_EM_eval3c.py: {''.join(missing)}" + return letters + + +def _binned_acceptance(x, accepted, bins): + total, _ = np.histogram(x, bins=bins) + passed, _ = np.histogram(x[accepted], bins=bins) + prob = np.full(total.shape, np.nan, dtype=np.float64) + m = total > 0 + prob[m] = passed[m].astype(np.float64) / total[m].astype(np.float64) + centers = 0.5 * (bins[:-1] + bins[1:]) + return centers, prob, total.astype(np.int64), passed.astype(np.int64) + + +def compute_fdr_acceptance_truth(fitD, md): + assert "bagsFDR_stageB" in md, "Plot o requires a Stage (b) FDR aggregate" + for key in ("A_prune", "single_rates"): + assert key in fitD, f"Plot o requires {key} in fit data" + for key in ("A_true", "E_true", "dale_conf"): + assert key in md, f"Plot o requires simulation truth metadata {key}" + + A_true = np.asarray(md["A_true"], dtype=np.float64) + if A_true.ndim == 3: + A_true = A_true[0] + E_true = np.asarray(md["E_true"]) != 0 + if E_true.ndim == 3: + E_true = E_true[0] != 0 + A_prune = np.asarray(fitD["A_prune"], dtype=np.float64) + assert A_prune.shape == A_true.shape and A_true.ndim == 2, ( + f"A_prune shape {A_prune.shape} must match 2D A_true shape {A_true.shape}" + ) + assert E_true.shape == A_true.shape, f"E_true shape {E_true.shape} must match A_true {A_true.shape}" + n_neur = A_true.shape[0] + rates = np.asarray(fitD["single_rates"], dtype=np.float64).reshape(-1) + assert rates.shape[0] == n_neur, "single_rates length must match A_true columns" + + num_exc = int(md["dale_conf"]["num_excite"]) + assert 0 < num_exc < n_neur, "dale_conf.num_excite must split excitatory/inhibitory source columns" + + off_mask = ~np.eye(n_neur, dtype=bool) + true_edge_mask = off_mask & E_true + assert np.any(true_edge_mask), "Plot o requires at least one true off-diagonal edge" + src_idx = np.broadcast_to(np.arange(n_neur, dtype=np.int64)[None, :], A_true.shape) + x_w = A_true[true_edge_mask] + accepted = (np.abs(A_prune) > 1e-12)[true_edge_mask] + src = src_idx[true_edge_mask] + src_rate = rates[src] + src_is_exc = src < num_exc + + w_step = 0.01 + w_lo = np.floor(float(np.min(x_w)) / w_step) * w_step + w_hi = np.ceil(float(np.max(x_w)) / w_step) * w_step + if np.isclose(w_lo, w_hi): + w_hi = w_lo + w_step + w_bins = np.arange(w_lo, w_hi + 1.5 * w_step, w_step, dtype=np.float64) + w_center, w_prob, w_total, w_pass = _binned_acceptance(x_w, accepted, w_bins) + + r_min = float(np.min(src_rate)) + r_max = float(np.max(src_rate)) + n_rate_bins = 10 + r_pos = src_rate[src_rate > 0.0] + if r_min > 0.0 and r_pos.size > 0 and r_max / float(np.min(r_pos)) > 20.0: + r_lo = float(np.min(r_pos)) + r_bins = np.geomspace(r_lo, r_max, n_rate_bins + 1) + r_scale = "log" + else: + if np.isclose(r_min, r_max): + r_max = r_min + 1.0 + r_bins = np.linspace(r_min, r_max, n_rate_bins + 1) + r_scale = "linear" + + rate_by_type = {} + for label, mask in (("exc", src_is_exc), ("inh", ~src_is_exc)): + c, p, t, a = _binned_acceptance(src_rate[mask], accepted[mask], r_bins) + rate_by_type[label] = { + "center": c, + "prob": p, + "total": t, + "passed": a, + } + + return { + "weight": { + "bin_edges": w_bins, + "center": w_center, + "prob": w_prob, + "total": w_total, + "passed": w_pass, + "bin_width": w_step, + }, + "rate": { + "bin_edges": r_bins, + "scale": r_scale, + "exc": rate_by_type["exc"], + "inh": rate_by_type["inh"], + }, + "summary": { + "num_candidates": int(x_w.size), + "num_accepted": int(np.sum(accepted)), + "num_exc": num_exc, + "num_inh": int(n_neur - num_exc), + }, + } + + +def main(): + args = parse_args() + args.showPlots = normalize_plot_letters(args.showPlots) + fit_f, fit_source = resolve_fit_file(args.basePath, args.dataName) + args.inpPath = os.path.dirname(fit_f) + args.outPath = os.path.join(args.basePath, "plots") + os.makedirs(args.outPath, exist_ok=True) + + print("EM-eval3c args:", vars(args), "\n") + + fitD, fitMD, fit_f, fit_source = load_fit(args.basePath, args.dataName, verb=args.verb) + if args.verb > 0: + print(f"Loaded fit ({fit_source}): {fit_f}") + if args.verb > 1: + pprint(fitMD) + plotD = display_fit_data(fitD, fitMD) + stage_c = is_stage_c_metadata(fitMD) + primary_A_label = "A_debias" if stage_c else "A_hat" + + is_sim = fitMD.get("data_type") != "bioExp" + requested_sim_only = sorted(set(args.showPlots) & SIM_ONLY_PLOTS) + assert is_sim or not requested_sim_only, ( + "Plots m-r require simulation truth; requested " + f"{''.join(requested_sim_only)} for data_type={fitMD.get('data_type')!r}" + ) + md = {**fitMD, "short_name": args.dataName} + if any(c in args.showPlots for c in "nopr"): + md = load_simu_static_truth(args.basePath, fitMD, md, verb=args.verb) + if "m" in args.showPlots: + md = load_simu_truth(args.basePath, fitD, fitMD, md, verb=args.verb) + if any(c in args.showPlots for c in "cdm") and args.time_range_sec is None: + args.time_range_sec = default_plot_time_range_sec(md) + if args.verb > 0: + print(f"default --time_range_sec from fitted data start: {args.time_range_sec}") + + args.prjName = args.dataName + plot = Plotter(args) + + if "a" in args.showPlots: + plot.summary(plotD, md, figId=PLOT_FIG_ID["a"]) + if "b" in args.showPlots: + plot.A_fitted(plotD, md, plotD["single_rates"], figId=PLOT_FIG_ID["b"]) + if "c" in args.showPlots: + plot.state_seq_fit(plotD, md, figId=PLOT_FIG_ID["c"], time_range_sec=args.time_range_sec) + if "d" in args.showPlots: + spikeD, spikeMD, spike_f = load_source_spikes(args.basePath, fitMD, verb=args.verb) + if args.verb > 0: + print(f"Loaded source spikes: {spike_f}") + rebD = detect_spike_bursts( + spikeD, spikeMD, args.time_rebin2, args.burst_freq_thres + ) + plot.spike_bursts(rebD, spikeMD, figId=PLOT_FIG_ID["d"], time_range_sec=args.time_range_sec) + if "e" in args.showPlots: + plot.node_outgoing_edge_stats( + plotD, md, plotD["single_rates"], plotD["neuron_type"], + figId=PLOT_FIG_ID["e"], est_key="A_hat", est_label=primary_A_label, + ) + if "f" in args.showPlots: + nodeD, nodeMD, node_f = load_node_metadata(args.basePath, fitMD, args.dataName, verb=args.verb) + if args.verb > 0: + print(f"Loaded node metadata: {node_f}") + plot.neuron_spatial_edges_split( + plotD, nodeD, nodeMD, plotD["neuron_type"], plotD["neuron_Sedge"], + maxNeurons=args.maxNeurons, figId=PLOT_FIG_ID["f"], + ) + if "g" in args.showPlots: + plot.fdr_selection_summary(fitD, md, figId=PLOT_FIG_ID["g"]) + if "h" in args.showPlots: + plot.final_weight_distributions(plotD, md, figId=PLOT_FIG_ID["h"]) + if "i" in args.showPlots: + plot.offdiag_weight_investigation(plotD, md, figId=PLOT_FIG_ID["i"]) + + #.... SIM only plots .... + if "m" in args.showPlots: + plot.state_seq_simu(plotD, md, figId=PLOT_FIG_ID["m"], time_range_sec=args.time_range_sec) + if "n" in args.showPlots: + plot.edge_stats_and_ABcorr(plotD, md, figId=PLOT_FIG_ID["n"]) + if "o" in args.showPlots: + plot.matrix_init(fitD, md, figId=PLOT_FIG_ID["o"]) + if "p" in args.showPlots: + plot.matrix_init(plotD, md, figId=PLOT_FIG_ID["p"], est_key="A_hat", est_label=primary_A_label) + if "r" in args.showPlots: + acceptD = compute_fdr_acceptance_truth(fitD, md) + plot.fdr_acceptance_truth(acceptD, md, figId=PLOT_FIG_ID["r"]) + + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/prism_EM_train3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/prism_EM_train3c.py new file mode 100755 index 00000000..abfc7fe3 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/prism_EM_train3c.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Standalone PRISM-EM training CLI. + +The reusable training implementation lives in PrismEM_Workhorse3c.py so the +ordinary trainer and the EM-FDR bagging driver run the same full-fit code. +""" + +import argparse +import math +import os +import secrets +from pprint import pprint + +import numpy as np + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from PrismEM_Workhorse3c import ( + add_prism_em_args, + barrier, + broadcast_array, + broadcast_object, + cleanup_distributed, + init_distributed, + is_rank0, + normalize_delay_args, + run_full_fit, + runtime_summary, + seed_everything, +) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="EM training: non-stationary Poisson GLM", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + add_prism_em_args(parser, include_time_range=True) + return parser.parse_args() + + +def slice_spikes_by_time(spikes, dt, time_range_sec): + t0_sec, t1_sec = [float(x) for x in time_range_sec] + if t1_sec < t0_sec: + t0_sec, t1_sec = t1_sec, t0_sec + T_raw = spikes.shape[0] + start_bin = max(0, int(math.floor(t0_sec / dt))) + end_bin = min(T_raw - 1, int(math.floor(t1_sec / dt))) + if end_bin <= start_bin: + raise ValueError("time_range_sec too small; need at least two bins") + return spikes[start_bin:end_bin + 1], [t0_sec, t1_sec], [start_bin, end_bin] + + +def main(): + args = normalize_delay_args(parse_args()) + ctx = init_distributed() + seed_everything(args.seed) + + try: + inpPath = os.path.join(args.basePath, "spikesData") + outPath = os.path.join(args.basePath, "prismFit") + if is_rank0(ctx): + os.makedirs(outPath, exist_ok=True) + if args.verb > 0: + print("Runtime:", runtime_summary(ctx)) + spikesFF = os.path.join(inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 0) + spikes = np.asarray(spikeD["spikes"]) + single_rates = np.asarray(spikeD["single_rates"]) + if args.verb > 1: + pprint(spikeMD) + dt = float(spikeMD["time_step_sec"]) + spikes, time_range_sec, time_range_bins = slice_spikes_by_time( + spikes, dt, args.time_range_sec + ) + if args.verb > 0: + print("\nEM-train args:", vars(args), "\n") + print( + f"N={spikes.shape[1]} EM={args.num_em_iters} M={args.num_states} " + f"T={spikes.shape[0]} dt={dt} " + f"time=[{time_range_sec[0]:.1f}, {time_range_sec[1]:.1f}]s " + f"bins={time_range_bins}" + ) + else: + spikeMD = None + spikes = None + single_rates = None + time_range_sec = None + time_range_bins = None + + spikeMD = broadcast_object(spikeMD, ctx) + spikes = broadcast_array(spikes, ctx) + single_rates = broadcast_array(single_rates, ctx) + time_range_sec = broadcast_object(time_range_sec, ctx) + time_range_bins = broadcast_object(time_range_bins, ctx) + + if args.fitName is None: + outF = f"{args.dataName}_{secrets.token_hex(2)}" if is_rank0(ctx) else None + outF = broadcast_object(outF, ctx) + else: + outF = args.fitName + + fitD, fitMD = run_full_fit( + spikes, spikeMD, single_rates, args, ctx, + time_range_sec=time_range_sec, + time_range_bins=time_range_bins, + fit_name=outF, + provenance_update={"dataName": args.dataName}, + ) + + if is_rank0(ctx): + outFF = os.path.join(outPath, f"{outF}.prismEM.npz") + write_data_npz(fitD, outFF, metaD=fitMD, verb=args.verb > 1) + print(f"\nSaved: {outFF}") + print(f" basePath={args.basePath}") + print( + f" ./prism_EM_eval3c.py --basePath $basePath " + f"--dataName {outF} -p a m n j i\n" + ) + barrier(ctx) + finally: + cleanup_distributed(ctx) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/prism_FDR_Bags_train3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/prism_FDR_Bags_train3c.py new file mode 100755 index 00000000..be5cf695 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/prism_FDR_Bags_train3c.py @@ -0,0 +1,514 @@ +#!/usr/bin/env python3 +"""Run one PRISM-EM FDR bag using reference-locked state labels. + +One invocation loads the reference full EM fit, samples a fraction of +reference-labeled lag pairs, runs one locked M-step for real A/B, runs +per-neuron time-shuffle locked-M-step null refits, and writes one bag file. +""" + +import argparse +import os +import re +import secrets +import time +from types import SimpleNamespace + +import numpy as np + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from Util_PrismEM import init_A_from_spikes, init_B_from_spikes +from PrismEM_Workhorse3c import ( + barrier, + broadcast_array, + broadcast_object, + broadcast_optional_array, + cleanup_distributed, + init_distributed, + init_null_params, + is_rank0, + run_locked_mstep, + runtime_summary, + seed_everything, + source_type_prune, +) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="PRISM-EM FDR bagging Stage (a): one reference-locked bag", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--emFitName", required=True, + help="Input EM-train fit stem in prismFit/") + parser.add_argument("--outFitName", default=None, + help="Output Stage (a) bag fit stem written to prismFDR/; default is emFitName plus random _hash4") + parser.add_argument("--basePath", required=True, + help="Run directory containing spikesData/, prismFit/, prismFDR/") + parser.add_argument("--bag_idx", type=int, required=True, + help="Bag index used in output filename and RNG stream") + parser.add_argument("--fdr_out_dir", type=str, default=None, + help="Output directory; default: /prismFDR") + parser.add_argument("--bag_frac", type=float, default=0.8, + help="Fraction of reference lag pairs sampled without replacement") + parser.add_argument("--num_scrambles", type=int, default=4, + help="Number of per-neuron time-shuffle null refits") + parser.add_argument("--min_roll_shift_sec", type=float, default=2.0, + help="Two-sided exclusion-zone half-width for circular shifts") + parser.add_argument("--init_A", choices=["ref", "data", "rand"], default="data", + help="A initialization for the real locked fit") + parser.add_argument("--init_B", choices=["ref", "data", "rand"], default="data", + help="B initialization for the real locked fit") + parser.add_argument("--null_A_init", choices=["scrambled_data", "rand"], + default="scrambled_data", + help="Fresh A initialization mode for each null refit") + parser.add_argument("--null_B_init", choices=["scrambled_data", "rand"], + default="scrambled_data", + help="Fresh B initialization mode for each null refit") + + g = parser.add_argument_group("locked M-step overrides") + g.add_argument("--epochs", type=int, default=None, + help="Locked M-step epochs; default from reference train.total_m_epochs") + g.add_argument("--lr_mstep", type=float, default=None) + g.add_argument("--lr_end_factor", type=float, default=None) + g.add_argument("--lambda3", type=float, default=None) + g.add_argument("--rho_max", type=float, default=None) + g.add_argument("--prescale_m_step_4_ArhoMax", type=int, default=None) + g.add_argument("--batch_size", type=int, default=None) + g.add_argument("--init_A_Tmax", type=int, default=None) + + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("-v", "--verb", type=int, default=1) + return parser.parse_args() + + +def validate_fit_stem(name, arg_name): + if not re.fullmatch(r"[A-Za-z0-9_.-]+", str(name)): + raise ValueError(f"{arg_name} must contain only letters, digits, underscores, dots, or hyphens: {name!r}") + if "/" in str(name): + raise ValueError(f"{arg_name} must be a file stem, not a path: {name!r}") + + +def ref_train_value(train_md, key, default=None): + if key in train_md: + return train_md[key] + if key == "epochs": + return int(train_md.get("total_m_epochs", int(train_md["num_em_iters"]) * int(train_md["m_epochs"]))) + return default + + +def source_spike_name(ref_md): + prov = ref_md["provenance"] + if ref_md.get("data_type") == "bioExp": + return prov["experiment_name"] + return prov["state_transition_file"] + + +def effective_locked_args(cli_args, ref_md, source_data_name): + train_md = ref_md["train"] + epochs = cli_args.epochs + if epochs is None: + epochs = ref_train_value(train_md, "epochs") + if int(epochs) < 1: + raise ValueError("--epochs must be >= 1") + + vals = { + "dataName": source_data_name, + "basePath": cli_args.basePath, + "num_states": int(train_md["num_states"]), + "num_em_iters": 1, + "m_epochs": int(epochs), + "epochs": int(epochs), + "lr_mstep": float(cli_args.lr_mstep if cli_args.lr_mstep is not None else train_md["lr_mstep"]), + "lr_end_factor": float(cli_args.lr_end_factor if cli_args.lr_end_factor is not None else train_md["lr_end_factor"]), + "lambda3": float(cli_args.lambda3 if cli_args.lambda3 is not None else train_md["lambda3"]), + "rho_max": float(cli_args.rho_max if cli_args.rho_max is not None else train_md["rho_max"]), + "prescale_m_step_4_ArhoMax": int( + cli_args.prescale_m_step_4_ArhoMax + if cli_args.prescale_m_step_4_ArhoMax is not None + else train_md["prescale_m_step_4_ArhoMax"] + ), + "batch_size": int(cli_args.batch_size if cli_args.batch_size is not None else train_md["batch_size"]), + "init_A": cli_args.init_A, + "init_B": cli_args.init_B, + "init_A_Tmax": int( + cli_args.init_A_Tmax + if cli_args.init_A_Tmax is not None + else train_md.get("init_A_Tmax", 50000) + ), + "seed": int(cli_args.seed), + "verb": int(cli_args.verb), + } + return SimpleNamespace(**vals) + + +def locked_hyperparam_dict(eff_args): + keys = ( + "epochs", "lr_mstep", "lr_end_factor", "lambda3", "rho_max", + "prescale_m_step_4_ArhoMax", "batch_size", "init_A", "init_B", + "init_A_Tmax", "num_states", "seed", + ) + return {k: getattr(eff_args, k) for k in keys} + + +def slice_reference_spikes(spikes, ref_md): + train_md = ref_md["train"] + b0, b1 = [int(x) for x in train_md["time_range_bins"]] + if b1 <= b0: + raise ValueError(f"Bad reference time_range_bins: {[b0, b1]}") + if b0 < 0 or b1 >= spikes.shape[0]: + raise ValueError( + f"Reference bins {[b0, b1]} outside spike array length {spikes.shape[0]}" + ) + return np.asarray(spikes[b0:b1 + 1]), [b0, b1] + + +def draw_pair_indices(rng, n_pair, bag_frac): + frac = float(bag_frac) + if not (0.0 < frac <= 1.0): + raise ValueError("--bag_frac must be in (0, 1]") + k = int(np.floor(frac * int(n_pair))) + k = max(1, min(int(n_pair), k)) + return np.sort(rng.choice(int(n_pair), size=k, replace=False)).astype(np.int64) + + +def circular_shift_by_neuron(spikes, shifts): + spikes = np.asarray(spikes) + shifts = np.asarray(shifts, dtype=np.int64) + out = np.empty_like(spikes) + for n, shift in enumerate(shifts): + out[:, n] = np.roll(spikes[:, n], int(shift)) + return out + + +def draw_roll_shifts(rng, T, N, min_shift_bins): + lo = int(min_shift_bins) + hi = int(T) - int(min_shift_bins) + if hi < lo: + raise ValueError( + f"Reference window length {T} too short for min_roll_shift_bins={min_shift_bins}" + ) + return rng.integers(lo, hi + 1, size=int(N), dtype=np.int64) + + +def stack_history(histories, key, dtype): + vals = [np.asarray(h[key], dtype=dtype) for h in histories] + return np.stack(vals, axis=0) if vals else np.zeros((0, 0), dtype=dtype) + + +def real_init_params(mode_a, mode_b, ref_d, ref_spikes, spike_md, eff_args, ctx): + if is_rank0(ctx): + if mode_a == "ref": + A_init = np.asarray(ref_d["A_hat"], dtype=np.float32) + A_md = {"method": "reference_fit", "source_key": "A_hat"} + elif mode_a == "data": + a_args = SimpleNamespace(**vars(eff_args)) + a_args.init_A = "data" + A_init, A_md = init_A_from_spikes(ref_spikes, a_args) + elif mode_a == "rand": + A_init, A_md = None, {"method": "rand"} + else: + raise ValueError(f"Bad init_A={mode_a}") + + if mode_b == "ref": + B_init = np.asarray(ref_d["B_hat"], dtype=np.float32) + B_md = {"method": "reference_fit", "source_key": "B_hat"} + elif mode_b == "data": + b_args = SimpleNamespace(**vars(eff_args)) + b_args.init_B = "data" + B_init, B_md = init_B_from_spikes(ref_spikes, float(spike_md["time_step_sec"]), b_args) + elif mode_b == "rand": + B_init, B_md = None, {"method": "rand"} + else: + raise ValueError(f"Bad init_B={mode_b}") + else: + A_init = B_init = None + A_md = B_md = None + + A_init = broadcast_optional_array(A_init, ctx) + B_init = broadcast_optional_array(B_init, ctx) + A_md = broadcast_object(A_md, ctx) + B_md = broadcast_object(B_md, ctx) + return A_init, B_init, A_md, B_md + + +def main(): + job_t0 = time.perf_counter() + args = parse_args() + if args.outFitName is None: + args.outFitName = f"{args.emFitName}_{secrets.token_hex(2)}" + validate_fit_stem(args.emFitName, "--emFitName") + validate_fit_stem(args.outFitName, "--outFitName") + ctx = init_distributed() + seed_base = int(args.seed) + int(args.bag_idx) * 1000003 + seed_everything(seed_base) + + try: + out_dir = args.fdr_out_dir or os.path.join(args.basePath, "prismFDR") + em_fit_name = args.emFitName + out_fit_name = args.outFitName + + if is_rank0(ctx): + os.makedirs(out_dir, exist_ok=True) + if args.verb > 0: + print("Runtime:", runtime_summary(ctx)) + + ref_f = os.path.join(args.basePath, "prismFit", f"{em_fit_name}.prismEM.npz") + ref_d, ref_md = read_data_npz(ref_f, verb=args.verb > 1) + source_data_name = source_spike_name(ref_md) + spike_f = os.path.join(args.basePath, "spikesData", f"{source_data_name}.spikes.npz") + spike_d, spike_md = read_data_npz(spike_f, verb=args.verb > 1) + source_spikes = np.asarray(spike_d["spikes"]) + ref_spikes, ref_bins = slice_reference_spikes(source_spikes, ref_md) + + s_hat = np.asarray(ref_d["S_hat"], dtype=np.int64) + if s_hat.shape[0] != ref_spikes.shape[0]: + raise ValueError( + f"Reference S_hat length {s_hat.shape[0]} != reference spikes length {ref_spikes.shape[0]}" + ) + if np.asarray(ref_d["B_hat"]).shape[0] != int(ref_md["train"]["num_states"]): + raise ValueError("Reference B_hat row count does not match train.num_states") + + eff_args = effective_locked_args(args, ref_md, source_data_name) + rng = np.random.default_rng(seed_base) + n_pair = ref_spikes.shape[0] - 1 + pair_idx = draw_pair_indices(rng, n_pair, args.bag_frac) + single_rates = ref_spikes.mean(axis=0).astype(np.float64) / float(spike_md["time_step_sec"]) + + if args.verb > 0: + print("\nFDR-bag args:", vars(args), "\n") + print( + f"reference={ref_f}\n" + f"source_spikes={spike_f}\n" + f"outFitName={out_fit_name} bag_idx={args.bag_idx} " + f"N={ref_spikes.shape[1]} M={eff_args.num_states} " + f"T_ref={ref_spikes.shape[0]} pairs={pair_idx.size}/{n_pair} " + f"epochs={eff_args.epochs} scrambles={args.num_scrambles}" + ) + else: + ref_d = ref_md = spike_md = ref_spikes = pair_idx = single_rates = eff_args = ref_bins = None + source_data_name = None + + ref_md = broadcast_object(ref_md, ctx) + source_data_name = broadcast_object(source_data_name, ctx) + em_fit_name = broadcast_object(em_fit_name, ctx) + out_fit_name = broadcast_object(out_fit_name, ctx) + spike_md = broadcast_object(spike_md, ctx) + eff_args = broadcast_object(eff_args, ctx) + ref_bins = broadcast_object(ref_bins, ctx) + ref_spikes = broadcast_array(ref_spikes, ctx) + pair_idx = broadcast_array(pair_idx, ctx) + single_rates = broadcast_array(single_rates, ctx) + + if is_rank0(ctx): + ref_A = np.asarray(ref_d["A_hat"], dtype=np.float32) + ref_B = np.asarray(ref_d["B_hat"], dtype=np.float32) + ref_S = np.asarray(ref_d["S_hat"], dtype=np.int64) + ref_c = np.asarray(ref_d["c_hat"], dtype=np.float32) + ref_cl = np.asarray(ref_d["S_hat_CL"], dtype=np.float32) + else: + ref_A = ref_B = ref_S = ref_c = ref_cl = None + ref_A = broadcast_array(ref_A, ctx) + ref_B = broadcast_array(ref_B, ctx) + ref_S = broadcast_array(ref_S, ctx) + ref_c = broadcast_array(ref_c, ctx) + ref_cl = broadcast_array(ref_cl, ctx) + ref_d_min = {"A_hat": ref_A, "B_hat": ref_B} + + yp = ref_spikes[:-1].astype(np.float32) + yc = ref_spikes[1:].astype(np.float32) + s_pair = ref_S[1:].astype(np.int64) + yp_bag = yp[pair_idx] + yc_bag = yc[pair_idx] + s_bag = s_pair[pair_idx] + + A_init, B_init, A_init_md, B_init_md = real_init_params( + eff_args.init_A, eff_args.init_B, ref_d_min, ref_spikes, spike_md, eff_args, ctx + ) + + real_t0 = time.perf_counter() + A_hat, B_hat, hist = run_locked_mstep( + yp_bag, yc_bag, s_bag, A_init, B_init, spike_md, eff_args, ctx + ) + if is_rank0(ctx) and args.verb > 0: + print(f"real locked fit finished: elapsed={time.perf_counter() - real_t0:.1f}s") + + if is_rank0(ctx): + rng = np.random.default_rng(seed_base + 17) + A_null = [] + B_null = [] + roll_shifts = [] + null_histories = [] + null_A_init_md = [] + null_B_init_md = [] + else: + rng = None + + dt = float(spike_md["time_step_sec"]) + min_roll_shift_bins = int(round(float(args.min_roll_shift_sec) / dt)) + T_ref, N = ref_spikes.shape + null_eff_args = SimpleNamespace(**vars(eff_args)) + if int(args.verb) == 1: + null_eff_args.verb = 0 + for p in range(int(args.num_scrambles)): + seed_everything(seed_base + 1009 * (p + 1)) + if is_rank0(ctx): + shifts = draw_roll_shifts(rng, T_ref, N, min_roll_shift_bins) + else: + shifts = None + shifts = broadcast_array(shifts, ctx) + scrambled = circular_shift_by_neuron(ref_spikes, shifts) + + A0, B0, A0_md, B0_md = init_null_params( + scrambled, spike_md, null_eff_args, ctx, args.null_A_init, args.null_B_init + ) + A_p, B_p, hist_p = run_locked_mstep( + scrambled[:-1].astype(np.float32)[pair_idx], + scrambled[1:].astype(np.float32)[pair_idx], + s_bag, + A0, + B0, + spike_md, + eff_args, + ctx, + ) + + if is_rank0(ctx): + A_null.append(A_p.astype(np.float32)) + B_null.append(B_p.astype(np.float32)) + roll_shifts.append(np.asarray(shifts, dtype=np.int64)) + null_histories.append(hist_p) + null_A_init_md.append(A0_md) + null_B_init_md.append(B0_md) + if args.verb > 0: + nz = int(hist_p["nz_edges_epoch"][-1]) if hist_p["nz_edges_epoch"].size else -1 + rho = float(hist_p["rho_epoch"][-1]) if hist_p["rho_epoch"].size else float("nan") + print( + f"null {p + 1}/{args.num_scrambles} finished: " + f"elapsed={time.perf_counter() - job_t0:.1f}s " + f"rho={rho:.4f} nz={nz}" + ) + + if is_rank0(ctx): + A_prune, neuron_type, neuron_sedge = source_type_prune(A_hat) + out_data_name = out_fit_name + out_stem = f"{out_fit_name}.bag{int(args.bag_idx):03d}" + ref_train = dict(ref_md["train"]) + train_md = dict(ref_train) + train_md.update({ + "fit_stage": "FDR_locked_mstep_bag", + "num_em_iters": 0, + "m_epochs": int(eff_args.epochs), + "total_m_epochs": int(eff_args.epochs), + "lr_mstep": float(eff_args.lr_mstep), + "lr_end_factor": float(eff_args.lr_end_factor), + "lambda3": float(eff_args.lambda3), + "rho_max": float(eff_args.rho_max), + "prescale_m_step_4_ArhoMax": int(eff_args.prescale_m_step_4_ArhoMax), + "batch_size": int(eff_args.batch_size), + "mstep_state_mode": "reference_viterbi_onehot_currbin", + "num_time_bins": int(T_ref), + "time_range_bins": [int(ref_bins[0]), int(ref_bins[1])], + "time_range_sec": list(ref_train["time_range_sec"]), + "seed": int(seed_base), + }) + + outD = { + "A_hat": A_hat.astype(np.float32), + "A_prune": A_prune.astype(np.float32), + "neuron_type": neuron_type.astype(np.int8), + "neuron_Sedge": neuron_sedge.astype(np.float32), + "B_hat": B_hat.astype(np.float32), + "single_rates": np.asarray(single_rates, dtype=np.float32), + "S_hat": ref_S.astype(np.int64), + "c_hat": ref_c.astype(np.float32), + "S_hat_CL": ref_cl.astype(np.float32), + "m_loss_epoch": hist["m_loss_epoch"], + "m_nll_epoch": hist["m_nll_epoch"], + "m_l1_epoch": hist["m_l1_epoch"], + "rho_epoch": hist["rho_epoch"], + "rho_correction_strength_epoch": hist["rho_correction_strength_epoch"], + "nz_edges_epoch": hist["nz_edges_epoch"], + "learning_rates": hist["learning_rates"], + "e_nll_em": np.zeros((0,), dtype=np.float64), + "A_null": np.stack(A_null, axis=0).astype(np.float32), + "B_null": np.stack(B_null, axis=0).astype(np.float32), + "roll_shifts_bin": np.stack(roll_shifts, axis=0).astype(np.int64), + "bag_pair_indices": pair_idx.astype(np.int64), + "bag_pair_dest_bins": (int(ref_bins[0]) + pair_idx + 1).astype(np.int64), + "bag_pair_state": s_bag.astype(np.int64), + "null_m_loss_epoch": stack_history(null_histories, "m_loss_epoch", np.float64), + "null_m_nll_epoch": stack_history(null_histories, "m_nll_epoch", np.float64), + "null_m_l1_epoch": stack_history(null_histories, "m_l1_epoch", np.float64), + "null_rho_epoch": stack_history(null_histories, "rho_epoch", np.float64), + "null_nz_edges_epoch": stack_history(null_histories, "nz_edges_epoch", np.int64), + "null_learning_rates": stack_history(null_histories, "learning_rates", np.float64), + } + if A_init is not None: + outD["A_init"] = np.asarray(A_init, dtype=np.float32) + if B_init is not None: + outD["B_init"] = np.asarray(B_init, dtype=np.float32) + + outMD = dict(spike_md) + outMD["fit_type"] = "prismEM_FDRbags_stageA" + outMD["train"] = train_md + outMD["states_recovery_eval"] = dict(ref_md.get("states_recovery_eval", {})) + outMD["init_A"] = A_init_md + outMD["init_B"] = B_init_md + outMD["init_state"] = { + "method": "reference_fit", + "source_key": "S_hat", + "em_fit_file": f"{em_fit_name}.prismEM.npz", + } + prov = dict(ref_md.get("provenance", spike_md.get("provenance", {}))) + prov.update({ + "dataName": source_data_name, + "EMtrain_file": out_stem, + "emFitName": em_fit_name, + "outFitName": out_fit_name, + "bagsFDR_stageA_dataName": out_data_name, + "bagsFDR_stageA_file": out_stem, + }) + outMD["provenance"] = prov + outMD["bagsFDR_stageA"] = { + "program": "prism_FDR_Bags_train3c.py", + "method": "reference_locked_pair_subsample", + "dataName": source_data_name, + "source_spike_name": source_data_name, + "emFitName": em_fit_name, + "outFitName": out_fit_name, + "output_dataName": out_data_name, + "output_name": out_stem, + "bag_idx": int(args.bag_idx), + "bag_frac": float(args.bag_frac), + "num_reference_pairs": int(T_ref - 1), + "num_selected_pairs": int(pair_idx.size), + "sample_without_replacement": True, + "em_fit_file": os.path.join(args.basePath, "prismFit", f"{em_fit_name}.prismEM.npz"), + "reference_time_range_bins": [int(ref_bins[0]), int(ref_bins[1])], + "reference_time_range_sec": list(ref_train["time_range_sec"]), + "locked_state_alignment": "S_hat[t] for pair (spikes[t-1], spikes[t])", + "locked_mstep": locked_hyperparam_dict(eff_args), + "real_init": {"A": A_init_md, "B": B_init_md}, + "null_refits": { + "num_scrambles": int(args.num_scrambles), + "A_init": args.null_A_init, + "B_init": args.null_B_init, + "A_init_md": null_A_init_md, + "B_init_md": null_B_init_md, + "min_roll_shift_sec": float(args.min_roll_shift_sec), + "min_roll_shift_bins": int(min_roll_shift_bins), + }, + "seed_base": int(seed_base), + "real_fit": {"train": train_md}, + } + + outF = os.path.join(out_dir, f"{out_stem}.prismFDRbag.npz") + write_data_npz(outD, outF, metaD=outMD, verb=args.verb > 1) + print(f"\nSaved FDR bag: {outF}") + barrier(ctx) + finally: + cleanup_distributed(ctx) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/prism_deBiasFit3c.py b/causal_net/nonStation_ver3c_states_EM_FDR/prism_deBiasFit3c.py new file mode 100755 index 00000000..3b761950 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/prism_deBiasFit3c.py @@ -0,0 +1,637 @@ +#!/usr/bin/env python3 +"""Stage (c) de-biased magnitude refit after EM-FDR bag aggregation.""" + +import argparse +import copy +import math +import os +import re +import secrets +import sys +import time +from types import SimpleNamespace + +import numpy as np + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from PrismEM_Workhorse3c import ( + barrier, + broadcast_array, + broadcast_object, + cleanup_distributed, + init_distributed, + is_rank0, + runtime_summary, + seed_everything, +) +from PrismDeBias_Workhorse3c import run_debias_fit + + +def print_argv_preflight(): + """Print raw invocation before argparse can exit on malformed CLI values.""" + argv = list(sys.argv) + rank = os.environ.get("RANK", "0") + local = os.environ.get("LOCAL_RANK", "0") + world = os.environ.get("WORLD_SIZE", "1") + if rank != "0": + return + missing_value_flags = [] + value_flags = { + "--basePath", + "--fdrFitName", + "--outFitName", + "--state_mode", + "--num_debias_iters", + "--warmup_locked_epochs", + "--m_epochs", + "--batch_size", + "--lr_mstep", + "--lr_end_factor", + "--eps_u", + "--pgd_iter", + "--lr_estep", + "--lambda2", + "--decode_dwell_sec", + "--rho_max", + "--prescale_m_step_4_ArhoMax", + "--delay_iter_4_ArhoMax", + "--target_iter_4_ArhoMax", + "--seed", + "--progress_every_batches", + "--progress_every_epochs", + "--progress_every_pairs", + "-v", + "--verb", + } + for i, tok in enumerate(argv[1:], start=1): + if tok in value_flags: + if i + 1 >= len(argv) or argv[i + 1].startswith("-"): + missing_value_flags.append(tok) + print( + "[prism_deBiasFit3c argv-preflight rank_env=%s local_env=%s/%s pid=%d] argv=%r" + % (rank, local, world, os.getpid(), argv), + flush=True, + ) + print( + "[prism_deBiasFit3c argv-preflight rank_env=%s] env basePath=%r BASEPATH=%r PWD=%r" + % ( + rank, + os.environ.get("basePath"), + os.environ.get("BASEPATH"), + os.environ.get("PWD"), + ), + flush=True, + ) + if missing_value_flags: + print( + "[prism_deBiasFit3c argv-preflight rank_env=%s] ERROR likely missing value after option(s): %s. " + "If you used --basePath $basePath, check that shell variable with: echo \"$basePath\"" + % (rank, ", ".join(missing_value_flags)), + flush=True, + ) + + +def parse_args(): + print_argv_preflight() + parser = argparse.ArgumentParser( + description="Stage (c) de-biased PRISM-EM refit on FDR-selected support", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--basePath", required=True, + help="Run directory containing spikesData/ and prismFit/") + parser.add_argument("--fdrFitName", required=True, + help="Stage (b) aggregate fit stem in prismFit/") + parser.add_argument("--outFitName", default=None, + help="Stage (c) output fit stem written to prismFit/") + + parser.add_argument("--state_mode", choices=["auto", "locked", "refit"], default="auto", + help="State treatment for Stage (c)") + parser.add_argument("--num_debias_iters", type=int, default=None, + help="Outer E/M iterations for state_mode=refit") + parser.add_argument("--warmup_locked_epochs", type=int, default=0, + help="Fixed-state M-step epochs before state refit") + + g = parser.add_argument_group("M-step") + g.add_argument("--m_epochs", type=int, default=None, + help="M-step epochs per de-bias iteration, or total epochs in locked mode") + g.add_argument("--batch_size", type=int, default=None) + g.add_argument("--lr_mstep", type=float, default=None) + g.add_argument("--lr_end_factor", type=float, default=None) + g.add_argument("--eps_u", type=float, default=1e-4, + help="Positive floor for Dale squared-parameter initialization") + + g = parser.add_argument_group("E-step") + g.add_argument("--pgd_iter", type=int, default=None) + g.add_argument("--lr_estep", type=float, default=None) + g.add_argument("--lambda2", type=float, default=None) + g.add_argument("--decode_dwell_sec", type=float, default=None) + + g = parser.add_argument_group("spectral radius") + g.add_argument("--rho_max", type=float, default=None) + g.add_argument("--prescale_m_step_4_ArhoMax", type=int, default=None) + g.add_argument("--delay_iter_4_ArhoMax", type=int, default=None) + g.add_argument("--target_iter_4_ArhoMax", type=int, default=None) + + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--progress_every_batches", type=int, default=50, + help="Print M-step progress every this many minibatches; 0 disables batch progress") + parser.add_argument("--progress_every_epochs", type=int, default=10, + help="Print M-step progress every this many M-epochs; 0 disables M-epoch progress") + parser.add_argument("--progress_every_pairs", type=int, default=20000, + help="Print E-step progress every this many local lag pairs; 0 disables pair progress") + parser.add_argument("-v", "--verb", type=int, default=1) + return parser.parse_args() + + +def validate_fit_stem(name, arg_name): + if not re.fullmatch(r"[A-Za-z0-9_.-]+", str(name)): + raise ValueError( + "%s must contain only letters, digits, underscores, dots, or hyphens: %r" + % (arg_name, name) + ) + if "/" in str(name): + raise ValueError("%s must be a file stem, not a path: %r" % (arg_name, name)) + + +def source_spike_name(md): + prov = md["provenance"] + if md.get("data_type") == "bioExp": + return prov["experiment_name"] + return prov["state_transition_file"] + + +def slice_spikes_by_bins(spikes, time_range_bins): + b0, b1 = [int(x) for x in time_range_bins] + if b1 <= b0: + raise ValueError("Bad time_range_bins: %r" % (time_range_bins,)) + if b0 < 0 or b1 >= spikes.shape[0]: + raise ValueError( + "time_range_bins %r outside spike array length %d" % (time_range_bins, spikes.shape[0]) + ) + return np.asarray(spikes[b0:b1 + 1]), [b0, b1] + + +def as_2d_B(B): + B = np.asarray(B, dtype=np.float32) + if B.ndim == 1: + B = B[None, :] + return B + + +def resolve_state_mode(mode, num_states): + if mode == "auto": + return "refit" if int(num_states) > 1 else "locked" + return str(mode) + + +def metadata_value(train_md, key, default): + return train_md[key] if key in train_md else default + + +def effective_args(cli_args, stage_md, train_md): + num_states = int(train_md["num_states"]) + state_mode = resolve_state_mode(cli_args.state_mode, num_states) + num_debias_iters = ( + int(cli_args.num_debias_iters) + if cli_args.num_debias_iters is not None + else 4 + ) + if state_mode == "locked": + num_debias_iters = 0 + if num_debias_iters < 0: + raise ValueError("--num_debias_iters must be non-negative") + + m_epochs = int( + cli_args.m_epochs + if cli_args.m_epochs is not None + else metadata_value(train_md, "m_epochs", 100) + ) + if m_epochs < 1: + raise ValueError("--m_epochs must be >= 1") + warmup = int(cli_args.warmup_locked_epochs) + if warmup < 0: + raise ValueError("--warmup_locked_epochs must be >= 0") + + target_rho_iter = ( + int(cli_args.target_iter_4_ArhoMax) + if cli_args.target_iter_4_ArhoMax is not None + else max(1, num_debias_iters) + ) + delay_rho_iter = ( + int(cli_args.delay_iter_4_ArhoMax) + if cli_args.delay_iter_4_ArhoMax is not None + else 0 + ) + if target_rho_iter <= delay_rho_iter: + target_rho_iter = delay_rho_iter + 1 + + decode_default = stage_md.get("states_recovery_eval", {}).get("decode_dwell_sec", 0.3) + + vals = { + "state_mode": state_mode, + "num_debias_iters": int(num_debias_iters), + "warmup_locked_epochs": int(warmup), + "m_epochs": int(m_epochs), + "batch_size": int( + cli_args.batch_size + if cli_args.batch_size is not None + else metadata_value(train_md, "batch_size", 4096) + ), + "lr_mstep": float( + cli_args.lr_mstep + if cli_args.lr_mstep is not None + else metadata_value(train_md, "lr_mstep", 0.003) + ), + "lr_end_factor": float( + cli_args.lr_end_factor + if cli_args.lr_end_factor is not None + else metadata_value(train_md, "lr_end_factor", 0.1) + ), + "pgd_iter": int( + cli_args.pgd_iter + if cli_args.pgd_iter is not None + else metadata_value(train_md, "pgd_iter", 5) + ), + "lr_estep": float( + cli_args.lr_estep + if cli_args.lr_estep is not None + else metadata_value(train_md, "lr_estep", 0.03) + ), + "lambda2": float( + cli_args.lambda2 + if cli_args.lambda2 is not None + else metadata_value(train_md, "lambda2", 2.0) + ), + "rho_max": float( + cli_args.rho_max + if cli_args.rho_max is not None + else metadata_value(train_md, "rho_max", 0.95) + ), + "prescale_m_step_4_ArhoMax": int( + cli_args.prescale_m_step_4_ArhoMax + if cli_args.prescale_m_step_4_ArhoMax is not None + else metadata_value(train_md, "prescale_m_step_4_ArhoMax", 120) + ), + "delay_iter_4_ArhoMax": int(delay_rho_iter), + "target_iter_4_ArhoMax": int(target_rho_iter), + "eps_u": float(cli_args.eps_u), + "decode_dwell_sec": float( + cli_args.decode_dwell_sec + if cli_args.decode_dwell_sec is not None + else decode_default + ), + "num_states": int(num_states), + "seed": int(cli_args.seed), + "progress_every_batches": int(cli_args.progress_every_batches), + "progress_every_epochs": int(cli_args.progress_every_epochs), + "progress_every_pairs": int(cli_args.progress_every_pairs), + "verb": int(cli_args.verb), + } + return SimpleNamespace(**vals) + + +def compute_loss_time(spikes, A, B, c_hat, train_md): + spikes = np.asarray(spikes, dtype=np.float64) + A = np.asarray(A, dtype=np.float64) + B = as_2d_B(B).astype(np.float64) + c_hat = np.asarray(c_hat, dtype=np.float64) + dt = float(train_md["time_step_sec"]) + eta_clip = float(train_md["eta_clip"]) + lambda2 = float(train_md["lambda2"]) + + yp = spikes[:-1] + yc = spikes[1:] + c_pairs = c_hat[1:] + c_prev = c_hat[:-1] + eta = yp @ A.T + c_pairs @ B + eta_c = np.minimum(eta, eta_clip) + lam = np.exp(eta_c) * dt + log_dt = math.log(dt) + nll_t = np.sum(lam - yc * (eta + log_dt), axis=1) + l2_t = lambda2 * np.sum((c_pairs - c_prev) ** 2, axis=1) + return nll_t.astype(np.float32), l2_t.astype(np.float32) + + +def train_metadata(eff_args, source_train, spike_md, n_neuron, n_time, time_range_bins, time_range_sec): + if str(eff_args.state_mode) == "locked": + num_em_iters = 0 + total_m_epochs = int(eff_args.warmup_locked_epochs) + int(eff_args.m_epochs) + else: + num_em_iters = int(eff_args.num_debias_iters) + total_m_epochs = int(eff_args.warmup_locked_epochs) + num_em_iters * int(eff_args.m_epochs) + return { + "fit_stage": "deBias_stageC", + "num_em_iters": int(num_em_iters), + "num_debias_iters": int(eff_args.num_debias_iters), + "warmup_locked_epochs": int(eff_args.warmup_locked_epochs), + "m_epochs": int(eff_args.m_epochs), + "total_m_epochs": int(total_m_epochs), + "pgd_iter": int(eff_args.pgd_iter), + "lr_estep": float(eff_args.lr_estep), + "lr_mstep": float(eff_args.lr_mstep), + "lr_end_factor": float(eff_args.lr_end_factor), + "lambda2": float(eff_args.lambda2), + "lambda3": 0.0, + "rho_max": float(eff_args.rho_max), + "rho_projection_mode": "active_uv_scale", + "prescale_m_step_4_ArhoMax": int(eff_args.prescale_m_step_4_ArhoMax), + "delay_em_iter_4_ArhoMax": int(eff_args.delay_iter_4_ArhoMax), + "target_em_iter_4_ArhoMax": int(eff_args.target_iter_4_ArhoMax), + "delay_iter_4_ArhoMax": int(eff_args.delay_iter_4_ArhoMax), + "target_iter_4_ArhoMax": int(eff_args.target_iter_4_ArhoMax), + "rho_sync_mode": "DDP_identical_active_params", + "delay_em_iter_4_lrDecay": 0, + "target_em_iter_4_lrDecay": max(1, int(eff_args.num_debias_iters)), + "delay_em_iter_4_Aprune": 0, + "mstep_state_mode": "destination_bin_onehot", + "state_mode": str(eff_args.state_mode), + "batch_size": int(eff_args.batch_size), + "num_states": int(eff_args.num_states), + "num_neurons": int(n_neuron), + "num_time_bins": int(n_time), + "time_step_sec": float(spike_md["time_step_sec"]), + "eta_clip": float(spike_md["poisson_eta_clip"]), + "time_range_sec": [float(time_range_sec[0]), float(time_range_sec[1])], + "time_range_bins": [int(time_range_bins[0]), int(time_range_bins[1])], + "seed": int(eff_args.seed), + "eps_u": float(eff_args.eps_u), + "source_stageB_mstep_state_mode": source_train.get("mstep_state_mode"), + } + + +def main(): + job_t0 = time.perf_counter() + args = parse_args() + env_rank = os.environ.get("RANK", "0") + env_local = os.environ.get("LOCAL_RANK", "0") + env_world = os.environ.get("WORLD_SIZE", "1") + if args.verb > 0 and env_rank == "0": + print( + "[prism_deBiasFit3c pre-init rank_env=%s local_env=%s/%s] raw args=%s" + % (env_rank, env_local, env_world, vars(args)), + flush=True, + ) + if not os.path.isdir(args.basePath): + raise FileNotFoundError( + "--basePath does not exist or is not a directory: %r. " + "If you used --basePath $basePath, verify it with: echo \"$basePath\"" + % args.basePath + ) + for subdir in ("spikesData", "prismFit"): + sub_path = os.path.join(args.basePath, subdir) + if not os.path.isdir(sub_path): + raise FileNotFoundError( + "--basePath %r is missing required subdirectory %r at %r" + % (args.basePath, subdir, sub_path) + ) + validate_fit_stem(args.fdrFitName, "--fdrFitName") + if args.outFitName is None: + args.outFitName = "%s_debias_%s" % (args.fdrFitName, secrets.token_hex(2)) + validate_fit_stem(args.outFitName, "--outFitName") + + ctx = init_distributed(verb=args.verb, program_name="prism_deBiasFit3c", rank0_only=True) + seed_everything(args.seed) + if args.verb > 0 and is_rank0(ctx): + print( + "[prism_deBiasFit3c rank=%d/%d] seed set to %d" + % (ctx.rank, ctx.world_size, int(args.seed)), + flush=True, + ) + + try: + if is_rank0(ctx): + if args.verb > 0: + print("Runtime:", runtime_summary(ctx), flush=True) + print("\nStage (c) deBias raw args:", vars(args), "\n", flush=True) + fit_f = os.path.join(args.basePath, "prismFit", "%s.prismEM.npz" % args.fdrFitName) + if args.verb > 0: + print("[rank 0] loading Stage (b) aggregate: %s" % fit_f, flush=True) + fdr_d, fdr_md = read_data_npz(fit_f, verb=args.verb > 1) + if "bagsFDR_stageB" not in fdr_md: + raise KeyError("Input must be a Stage (b) FDR aggregate with bagsFDR_stageB metadata") + required = ("A_hat", "B_hat", "selected_mask", "neuron_type", "S_hat", "c_hat", "S_hat_CL") + missing = [k for k in required if k not in fdr_d] + if missing: + raise KeyError("Stage (b) input missing required arrays: %s" % missing) + + train_b = dict(fdr_md["train"]) + eff_args = effective_args(args, fdr_md, train_b) + if args.verb > 0: + print("[rank 0] effective args:", vars(eff_args), flush=True) + source_name = source_spike_name(fdr_md) + spike_f = os.path.join(args.basePath, "spikesData", "%s.spikes.npz" % source_name) + if args.verb > 0: + print("[rank 0] loading source spikes: %s" % spike_f, flush=True) + spike_d, spike_md = read_data_npz(spike_f, verb=args.verb > 1) + source_spikes = np.asarray(spike_d["spikes"]) + spikes, time_bins = slice_spikes_by_bins(source_spikes, train_b["time_range_bins"]) + time_sec = list(train_b["time_range_sec"]) + + A_stage = np.asarray(fdr_d["A_hat"], dtype=np.float32) + B_stage = as_2d_B(fdr_d["B_hat"]) + selected_mask = np.asarray(fdr_d["selected_mask"], dtype=np.bool_) + neuron_type = np.asarray(fdr_d["neuron_type"], dtype=np.int8) + S_stageB = np.asarray(fdr_d["S_hat"], dtype=np.int64) + c_stageB = np.asarray(fdr_d["c_hat"], dtype=np.float32) + S_stageB_CL = np.asarray(fdr_d["S_hat_CL"], dtype=np.float32) + + if S_stageB.shape[0] != spikes.shape[0]: + raise ValueError("Stage (b) S_hat length does not match sliced spikes") + if c_stageB.shape[0] != spikes.shape[0]: + raise ValueError("Stage (b) c_hat length does not match sliced spikes") + if B_stage.shape[0] != int(eff_args.num_states): + raise ValueError("B_hat row count does not match train.num_states") + + single_rates = ( + np.asarray(fdr_d["single_rates"], dtype=np.float32) + if "single_rates" in fdr_d + else (spikes.mean(axis=0) / float(spike_md["time_step_sec"])).astype(np.float32) + ) + + if args.verb > 0: + print( + "Stage (c) input: %s N=%d M=%d T=%d state_mode=%s m_epochs=%d" + % ( + fit_f, + A_stage.shape[0], + int(eff_args.num_states), + spikes.shape[0], + eff_args.state_mode, + int(eff_args.m_epochs), + ), + flush=True, + ) + else: + fdr_d = fdr_md = spike_md = eff_args = None + spikes = A_stage = B_stage = selected_mask = neuron_type = None + S_stageB = c_stageB = S_stageB_CL = single_rates = None + source_name = time_bins = time_sec = None + + fdr_md = broadcast_object(fdr_md, ctx) + if args.verb > 0 and is_rank0(ctx): + print("[rank %d/%d] metadata broadcast complete" % (ctx.rank, ctx.world_size), flush=True) + spike_md = broadcast_object(spike_md, ctx) + eff_args = broadcast_object(eff_args, ctx) + source_name = broadcast_object(source_name, ctx) + time_bins = broadcast_object(time_bins, ctx) + time_sec = broadcast_object(time_sec, ctx) + spikes = broadcast_array(spikes, ctx) + A_stage = broadcast_array(A_stage, ctx) + B_stage = broadcast_array(B_stage, ctx) + selected_mask = broadcast_array(selected_mask, ctx) + neuron_type = broadcast_array(neuron_type, ctx) + S_stageB = broadcast_array(S_stageB, ctx) + c_stageB = broadcast_array(c_stageB, ctx) + S_stageB_CL = broadcast_array(S_stageB_CL, ctx) + single_rates = broadcast_array(single_rates, ctx) + if args.verb > 0 and is_rank0(ctx): + print( + "[rank %d/%d] data broadcast complete spikes=%s A=%s B=%s selected=%s" + % ( + ctx.rank, + ctx.world_size, + getattr(spikes, "shape", None), + getattr(A_stage, "shape", None), + getattr(B_stage, "shape", None), + getattr(selected_mask, "shape", None), + ), + flush=True, + ) + + result = run_debias_fit( + spikes, + spike_md, + A_stage, + B_stage, + selected_mask, + neuron_type, + c_stageB, + S_stageB, + eff_args, + ctx, + ) + if args.verb > 0 and is_rank0(ctx): + print("[rank %d/%d] run_debias_fit complete" % (ctx.rank, ctx.world_size), flush=True) + + if is_rank0(ctx): + out_d = dict(fdr_d) + for key in ("m_loss_epoch", "m_nll_epoch", "m_l1_epoch", "rho_epoch", + "rho_correction_strength_epoch", "nz_edges_epoch", + "learning_rates", "e_nll_em"): + if key in fdr_d: + out_d["stageB_" + key] = fdr_d[key] + + out_d["S_stageB"] = S_stageB.astype(np.int64) + out_d["c_stageB"] = c_stageB.astype(np.float32) + out_d["S_stageB_CL"] = S_stageB_CL.astype(np.float32) + out_d["A_debias"] = result["A_debias"].astype(np.float32) + out_d["B_debias"] = result["B_debias"].astype(np.float32) + out_d["A_debias_init"] = result["A_debias_init"].astype(np.float32) + out_d["B_debias_init"] = result["B_debias_init"].astype(np.float32) + out_d["S_hat"] = result["S_hat"].astype(np.int64) + out_d["c_hat"] = result["c_hat"].astype(np.float32) + out_d["S_hat_CL"] = result["S_hat_CL"].astype(np.float32) + out_d["single_rates"] = single_rates.astype(np.float32) + + hist = result["history"] + out_d.update(hist) + active = result["active"] + out_d["debias_dale_i"] = active["dale_i"].astype(np.int64) + out_d["debias_dale_j"] = active["dale_j"].astype(np.int64) + out_d["debias_dale_sign"] = active["dale_sign"].astype(np.float32) + out_d["debias_free_i"] = active["free_i"].astype(np.int64) + out_d["debias_free_j"] = active["free_j"].astype(np.int64) + out_d["A_diag_debias"] = np.diag(result["A_debias"]).astype(np.float32) + out_d["neuron_Sedge_debias"] = ( + result["A_debias"] * (~np.eye(result["A_debias"].shape[0], dtype=bool)) + ).sum(axis=0).astype(np.float32) + + train_md = train_metadata( + eff_args, + fdr_md["train"], + spike_md, + result["A_debias"].shape[0], + spikes.shape[0], + time_bins, + time_sec, + ) + loss_nll_t, loss_l2_t = compute_loss_time( + spikes, result["A_debias"], result["B_debias"], result["c_hat"], train_md + ) + out_d["loss_nll_time"] = loss_nll_t + out_d["loss_l2_time"] = loss_l2_t + + out_md = copy.deepcopy(fdr_md) + out_md["fit_type"] = "prismEM_deBias_stageC" + out_md["train"] = train_md + out_md["eval_f"] = { + "loss_nll_time_key": "loss_nll_time", + "loss_l2_time_key": "loss_l2_time", + } + out_md["deBias_stageC"] = { + "program": "prism_deBiasFit3c.py", + "workhorse": "PrismDeBias_Workhorse3c.py", + "input_fdrFitName": args.fdrFitName, + "input_stageB_file": os.path.join(args.basePath, "prismFit", "%s.prismEM.npz" % args.fdrFitName), + "outFitName": args.outFitName, + "source_spike_name": source_name, + "source_spike_file": os.path.join(args.basePath, "spikesData", "%s.spikes.npz" % source_name), + "reference_emFitName": out_md.get("provenance", {}).get("emFitName"), + "state_mode_requested": args.state_mode, + "state_mode": eff_args.state_mode, + "mstep_state_alignment": "destination_bin_S_t_for_pair_Y_tminus1_to_Y_t", + "support_source": "stageB_selected_mask_offdiag_plus_all_diagonal", + "dale_source": "stageB_neuron_type_source_columns", + "A_initialization": "stageB_A_hat_for_all_active_entries", + "B_initialization": "stageB_B_hat", + "A_stageB_preserved_keys": ["A_hat", "A_prune"], + "B_stageB_preserved_key": "B_hat", + "A_output_key": "A_debias", + "B_output_key": "B_debias", + "no_second_pruning": True, + "hyperparameters": vars(eff_args), + "num_dale_active": int(active["num_dale"]), + "num_free_active": int(active["num_free"]), + "num_diag_free": int(active["num_diag"]), + "num_selected_offdiag": int(active["num_selected_offdiag"]), + "num_trainable_A": int(active["num_dale"] + active["num_free"]), + "num_trainable_B": int(np.asarray(result["B_debias"]).size), + "elapsed_train_sec": float(result["elapsed_train_sec"]), + "elapsed_total_sec": float(time.perf_counter() - job_t0), + } + prov = dict(out_md.get("provenance", {})) + prov.update({ + "deBias_stageC_input": args.fdrFitName, + "deBias_stageC_file": args.outFitName, + "deBias_stageC_outFitName": args.outFitName, + "EMtrain_file": args.outFitName, + }) + out_md["provenance"] = prov + + out_dir = os.path.join(args.basePath, "prismFit") + os.makedirs(out_dir, exist_ok=True) + out_f = os.path.join(out_dir, "%s.prismEM.npz" % args.outFitName) + write_data_npz(out_d, out_f, metaD=out_md, verb=args.verb > 1) + if args.verb > 0: + print("\nSaved Stage (c) deBias fit: %s" % out_f, flush=True) + print( + " A_debias: %s B_debias: %s state_mode=%s" + % (out_d["A_debias"].shape, out_d["B_debias"].shape, eff_args.state_mode), + flush=True, + ) + print( + " ./prism_EM_eval3c.py --basePath %s --dataName %s -p a b c e h i" + % (args.basePath, args.outFitName), + flush=True, + ) + if args.verb > 0 and is_rank0(ctx): + print("[rank %d/%d] entering final barrier" % (ctx.rank, ctx.world_size), flush=True) + barrier(ctx) + if args.verb > 0 and is_rank0(ctx): + print("[rank %d/%d] finished total_elapsed=%.1fs" % (ctx.rank, ctx.world_size, time.perf_counter() - job_t0), flush=True) + finally: + cleanup_distributed(ctx) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/scan_FDR_BAG_hpar.sh b/causal_net/nonStation_ver3c_states_EM_FDR/scan_FDR_BAG_hpar.sh new file mode 100755 index 00000000..b21eac60 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/scan_FDR_BAG_hpar.sh @@ -0,0 +1,92 @@ +#!/bin/bash +# Scan Stage (b) FDR-bag aggregation hyperparameters for an existing bag set. +# +# This script does not rerun EM or bag fits. It reuses: +# $basePath/prismFDR/${fdrBagsName}.bagNNN.prismFDRbag.npz +# and writes aggregate fits to: +# $basePath/prismFit/${outAgrName}_thrX.prismEM.npz +# $basePath/prismFit/${outAgrName}_qX.prismEM.npz + +set -euo pipefail +export OMP_NUM_THREADS=1 +SECONDS=0 + +module load pytorch +cd "$(dirname "$0")" + +# ---------- input bag set ---------- +basePath=/pscratch/sd/b/balewski/2026_causalNet_exp_ver3c +fdrBagsName=daleN200_74e6d6_e2b7d7_take1_30min_embe98_fdrbe98 # N200 12/21 Hz +fdrBagsName=daleN200_55e5a6_ff089c_take1_60min_em4c42_fdr4c42 # N200 3/6 Hz +outAgrName=aaa24 +numBags=11 + +# ---------- baseline aggregation hyperparameters ---------- +base_per_bag_quantile=0.95 +base_stab_sel_thresh=0.7 + +# ---------- scan points ---------- +stabSelScan=(0.8 0.7 0.6 0.5 ) +perBagQuantileScan=(0.99 0.97 0.95 0.90 0.85) + +echo "basePath=$basePath" +echo "fdrBagsName=$fdrBagsName" +echo "outAgrName=$outAgrName" +echo "numBags=$numBags" +echo "base_per_bag_quantile=$base_per_bag_quantile" +echo "base_stab_sel_thresh=$base_stab_sel_thresh" +echo "stabSelScan=${stabSelScan[*]}" +echo "perBagQuantileScan=${perBagQuantileScan[*]}" + +echo +echo "=== Scan stability selection threshold ===" +for thr in "${stabSelScan[@]}"; do + scanOutAgrName="${outAgrName}_thr${thr}" + printf "\n--- stab_sel_thresh=%s, elapsed %.1f min ---\n" "$thr" "$(awk "BEGIN {print $SECONDS/60.0}")" + ./prism_EM_FDR_Bags_aggregate3c.py \ + --basePath "$basePath" \ + --dataName "$fdrBagsName" \ + --outAgrName "$scanOutAgrName" \ + --num_bags "$numBags" \ + --per_bag_quantile "$base_per_bag_quantile" \ + --stab_sel_thresh "$thr" +done + +echo +echo "=== Scan per-bag null quantile ===" +for q in "${perBagQuantileScan[@]}"; do + scanOutAgrName="${outAgrName}_q${q}" + printf "\n--- per_bag_quantile=%s, elapsed %.1f min ---\n" "$q" "$(awk "BEGIN {print $SECONDS/60.0}")" + ./prism_EM_FDR_Bags_aggregate3c.py \ + --basePath "$basePath" \ + --dataName "$fdrBagsName" \ + --outAgrName "$scanOutAgrName" \ + --num_bags "$numBags" \ + --per_bag_quantile "$q" \ + --stab_sel_thresh "$base_stab_sel_thresh" +done + +echo +echo "Done." +echo "Input bags: $basePath/prismFDR/${fdrBagsName}.bagNNN.prismFDRbag.npz" +echo "Aggregates:" +for thr in "${stabSelScan[@]}"; do + echo " $basePath/prismFit/${outAgrName}_thr${thr}.prismEM.npz" +done +for q in "${perBagQuantileScan[@]}"; do + echo " $basePath/prismFit/${outAgrName}_q${q}.prismEM.npz" +done + +thrTags=() +for thr in "${stabSelScan[@]}"; do + thrTags+=("thr${thr}") +done +qTags=() +for q in "${perBagQuantileScan[@]}"; do + qTags+=("q${q}") +done + +echo +echo "Metric commands:" +echo "./edgeMaterAbs3c.py --basePath \"$basePath\" --fitNameTrunk \"$outAgrName\" --fitTags ${thrTags[*]}" +echo "./edgeMaterAbs3c.py --basePath \"$basePath\" --fitNameTrunk \"$outAgrName\" --fitTags ${qTags[*]}" diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/toolbox b/causal_net/nonStation_ver3c_states_EM_FDR/toolbox new file mode 120000 index 00000000..c9203e9b --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/toolbox @@ -0,0 +1 @@ +../toolbox/ \ No newline at end of file diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/view_bioexp.py b/causal_net/nonStation_ver3c_states_EM_FDR/view_bioexp.py new file mode 100755 index 00000000..e4afe308 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/view_bioexp.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +""" +Visualization tool for biological experiment input features and data quality. + +This script provides comprehensive visualization and analysis of experimental +neural data, focusing on data quality assessment and feature exploration. +Main functionality includes: +- Time series visualization of neural recordings with customizable time ranges +- Data quality metrics and statistical summaries +- Cluster detection and activity pattern analysis +- Interactive plotting with configurable display options + +Used primarily for exploratory data analysis of biological neural recordings +before further processing and connectivity analysis. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os + +from pprint import pprint +import numpy as np +from PlotterBioExp import Plotter +from toolbox.Util_NumpyIO import read_data_npz +from UtilBioExp import detect_spike_bursts +import argparse + + +def _require_bioexp(bioD, bioMD): + for key in ("metrics_curated", "MEA_idx", "single_rates"): + assert key in bioD, f"bioExp.npz missing {key!r}; rerun prep_bioexp3b.py" + for key in ("metrics_curated_columns", "short_name", "data_selector", "rate_summary"): + assert key in bioMD, f"bioExp metadata missing {key!r}" + assert "freq_range" in bioMD["data_selector"] + + +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v", "--verbosity", type=int, + help="increase output verbosity", default=1, dest="verb") + parser.add_argument("-p", "--showPlots", default="a", nargs="+", + help="plots: a=freq histo, b=freq vs time, c=MEA layout, " + "d=metrics histograms, e=metrics correlations") + + parser.add_argument("-X", "--noXterm", action="store_true", + help="Disable X terminal for plotting") + + parser.add_argument("--dataPath", + default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", + help="head dir for any further data processing") + + parser.add_argument("-T", "--time_range", default=[0., 60], nargs=2, type=float, + help="display data time range in seconds") + parser.add_argument("--burst_freq_thres", default=5., type=float, + help="tags high freq channels for burst detection") + parser.add_argument("--dataName", default="HET_80k_1-fc62ef", + help="preprocessed session name") + + parser.add_argument("-R", "--time_rebin2", default=10, type=int, + help="rebin current time axis for burst panels") + + args = parser.parse_args() + args.outPath = "out/" + args.showPlots = "".join(args.showPlots) + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + assert args.time_range[0] < args.time_range[1] + assert os.path.exists(args.dataPath) + assert os.path.exists(args.outPath) + return args + + +#================================= +# M A I N +#================================= +if __name__ == "__main__": + args = get_parser() + np.set_printoptions(precision=3) + + bioFF = os.path.join(args.dataPath, f"{args.dataName}.bioExp.npz") + print("bioExp:", bioFF) + bioD, bioMD = read_data_npz(bioFF) + if args.verb > 1: pprint(bioMD) + _require_bioexp(bioD, bioMD) + + plotMD = dict(bioMD) + args.prjName = bioMD["short_name"] + + spikeD = None + rebD = None + if "a" in args.showPlots or "b" in args.showPlots: + spikesFF = os.path.join(args.dataPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF) + ''' + for key in ("spikes", "single_rates", "provenance", "data_type", + "time_step_sec", "num_neurons"): + assert key in spikeMD, f"spikes.npz metadata missing {key!r}" + ''' + assert "experiment_name" in spikeMD["provenance"] + plotMD = {**bioMD, **spikeMD} + args.prjName = spikeMD["provenance"]["experiment_name"] + if args.verb > 1: + pprint(spikeMD) + + if "b" in args.showPlots: + plotMD["plot"] = {"time_rangeLR": np.array(args.time_range)} + rebD = detect_spike_bursts( + spikeD, spikeMD, args.time_rebin2, + args.burst_freq_thres, + ) + + plot = Plotter(args) + + if "a" in args.showPlots: + plot.freq_histo(spikeD, plotMD, figId=1) + if "b" in args.showPlots: + plot.freq_vs_time(rebD, plotMD, figId=2) + if "c" in args.showPlots: + plot.neuron_spatial(bioD, plotMD, figId=3) + if "d" in args.showPlots: + plot.metrics_histograms(bioD, plotMD, figId=4) + if "e" in args.showPlots: + plot.metrics_correlations(bioD, plotMD, figId=5) + + plot.display_all() + print("M:done") + + diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/view_daleMatrix3.py b/causal_net/nonStation_ver3c_states_EM_FDR/view_daleMatrix3.py new file mode 100755 index 00000000..11c53c99 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/view_daleMatrix3.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +""" +Visualization tool for simulated Dale Poisson network data. + +Loads simulation output (simTruth.npz and spikes.npz) produced by +gen_daleMatrices3c.py and plots analysis figures for a selected state +index (--idxState). All arrays use natural neuron indexing +(first num_excite source columns are excitatory, remainder inhibitory). + +Available plots (-p flag): + a Dale connectivity matrix, eigenvalue scatter with spectral-radius + circle, and node-location map colored by signed firing rate + b Weight histogram, firing-rate histogram, outgoing-edge count + per neuron, and per-neuron firing-rate bar chart + d B_idle vs firing rate / SNR scatter, plus excitatory and + inhibitory rate histograms + e Pseudospectral contour plot with eigenvalue overlay + +Usage: + ./view_daleMatrix3.py --dataName daleN100_9fbe7f -m 0 -p a b d e +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os +from pprint import pprint +import numpy as np +from PlotterDaleMatrix import Plotter +from toolbox.Util_NumpyIO import read_data_npz +import argparse + +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize simulated Dale Poisson network data") + parser.add_argument("-v","--verbosity",type=int, help="increase output verbosity", default=1, dest='verb') + parser.add_argument("-p", "--showPlots", default='a b', nargs='+',help="abc-string listing shown plots: a=Dale_matrix_eigen_locations, b=histo_weights_rates, c=rates_study, d=pseudospectra") + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument("--basePath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for input data") + parser.add_argument("--dataName", default='daleN150_448b86',help='simulated Dale network base name') + #parser.add_argument("--outPath", type=str, default=None, help="Output path for plots (defaults to basePath)") + parser.add_argument('-m', '--idxState', type=int, default=0, help="Index into state list, selects which state to plot") + + args = parser.parse_args() + + # make arguments more flexible + args.inpPath = os.path.join(args.basePath,'truthDale') + args.outPath = os.path.join(args.basePath,'plots') + args.showPlots=''.join(args.showPlots) + + print( 'myArg-program:',parser.prog) + for arg in vars(args): print( 'myArg:',arg, getattr(args, arg)) + + assert os.path.exists(args.basePath) + return args + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__=="__main__": + args=get_parser() + np.set_printoptions(precision=3) + + # Load simulation truth data (Dale matrices, biases, etc.) + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb>0) + if args.verb>1: + print("\nSimulation Truth Metadata:"); pprint(trueMD) + + # Load spike data (generated spike counts and rates) + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb>0) + if args.verb>1: + print("\nSpike Data Metadata:"); pprint(spikeMD) + + # Merge metadata for plotting + trueMD['short_name'] = args.dataName + + # Select B-offset slice + ib = args.idxState + boffsets = trueMD['dale_conf']['Boffsets'] + assert ib < len(boffsets), f"idxState={ib} out of range, only {len(boffsets)} states available" + B_offset = boffsets[ib] + R_sel = trueMD['dale_conf']['spectral_radius'] + print(f"\nSelected B offset [{ib}]: offset={B_offset:.3f} (out of {boffsets})") + print(f"Spectral radius: R={R_sel:.3f}") + + # Slice stacked arrays along axis 0 for the selected B offset + trueD_r = { 'A_true': trueD['A_true'], 'B_true': trueD['B_true'][ib], 'E_true': trueD['E_true'] } + for key in ('node_positions', 'node_is_inhibitory', 'node_distance_matrix'): + if key in trueD: + trueD_r[key] = trueD[key] + spikeD_r = { k: spikeD[k][ib] for k in spikeD } + trueMD['sel_spect_radius'] = R_sel + trueMD['sel_Boffset'] = B_offset + trueMD['sel_state'] = ib + + #-------------------------------- + # .... plotting ........ + args.prjName=args.dataName+'_view%d'%ib + plot=Plotter(args) + + if 'a' in args.showPlots: + plot.Dale_matrix_and_eigen(trueD_r['A_true'], trueMD, trueD_r, spikeD_r, figId=1) + + if 'b' in args.showPlots: + plot.histo_weights_rates(trueD_r,spikeD_r,trueMD,figId=2) + + if 'c' in args.showPlots: + plot.rates_study(trueD_r,spikeD_r,trueMD,figId=3) + + if 'd' in args.showPlots: + plot.Dale_matrix_pseudospectra(trueD_r['A_true'],trueMD,trueD_r,figId=4) + + plot.display_all() + print('M:done - view_dalePoisson completed successfully!') diff --git a/causal_net/nonStation_ver3c_states_EM_FDR/view_spikesTrain3.py b/causal_net/nonStation_ver3c_states_EM_FDR/view_spikesTrain3.py new file mode 100755 index 00000000..31d6f5c9 --- /dev/null +++ b/causal_net/nonStation_ver3c_states_EM_FDR/view_spikesTrain3.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +""" +Visualization tool for biological experiment input features and data quality. + +This script provides comprehensive visualization and analysis of experimental +neural data, focusing on data quality assessment and feature exploration. +Main functionality includes: +- Time series visualization of neural recordings with customizable time ranges +- Data quality metrics and statistical summaries +- Interactive plotting with configurable display options + +Used primarily for exploratory data analysis of biological neural recordings +before further processing and connectivity analysis. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os + +from pprint import pprint +import numpy as np +from PlotterSpikesTrain import Plotter +from toolbox.Util_NumpyIO import read_data_npz +import argparse +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize spike-train data from simulated Dale network outputs") + parser.add_argument("-v","--verbosity",type=int, help="increase output verbosity", default=1, dest='verb') + parser.add_argument("-p", "--showPlots", default='a', nargs='+',help="abcd-string listing shown plots") + + + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + + parser.add_argument("--basePath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for input data") + parser.add_argument("--dataName", default=None,help='simulated Dale network base name') + + parser.add_argument('-m', '--idxState', type=int, default=0, help="Index into state list; if idxState<0 read spikes from spikesData/") + + parser.add_argument('-T','--time_range_sec' , default=[0., 50], nargs=2, type=float, help='display data time range in seconds') + parser.add_argument('-r','--time_rebin2', default=10, type=int, help='rebin current time axis') + + args = parser.parse_args() + # make arguments more flexible + if args.idxState >= 0: + args.inpPath = os.path.join(args.basePath, 'truthDale') + else: + args.inpPath = os.path.join(args.basePath, 'spikesData') + args.outPath = os.path.join(args.basePath,'plots') + args.showPlots=''.join(args.showPlots) + + print( 'myArg-program:',parser.prog) + for arg in vars(args): print( 'myArg:',arg, getattr(args, arg)) + + if args.time_range_sec!=None: assert args.time_range_sec[0] < args.time_range_sec[1] + assert os.path.exists(args.basePath) + assert os.path.exists(args.inpPath), f"missing inpPath: {args.inpPath}" + return args + + +#...!...!.................... +def rebin_spike_rates(spikeYield, md, tReb2): + """Rebin 2D spike train (time x neurons) and return rate summaries over rebinned time.""" + assert spikeYield.ndim == 2, f"Expected 2D spike train, got shape={spikeYield.shape}" + assert tReb2<101 # this would exceed 1 seconds + + #.... rebin time axis + ntime,nchan=spikeYield.shape + if ntime % tReb2 != 0: + ntime_c = ntime - (ntime % tReb2) + spikeYield = spikeYield[:ntime_c] # Clip the data + #print('ccl',ntime_c , ntime ,ntime % tReb2, tReb2) + #... sum yileds over tReb2 time bins + spikeYieldR= np.sum( spikeYield.reshape(-1, tReb2, nchan),axis=1) + time_step=md['time_step_sec'] + time_step2=time_step*tReb2 + #print('rr1',time_step2,spikeYieldR.shape,spikeYield.shape) + + rate2D=spikeYieldR/time_step2 + pop_spike_count = np.sum(spikeYieldR, axis=1) + pop_rate_hz = pop_spike_count / time_step2 + + rebD={'time_step2':time_step2} + rebD['rate2D']=rate2D + rebD['pop_spike_count'] = pop_spike_count + rebD['pop_rate_hz'] = pop_rate_hz + ntime//=tReb2 + rebD['timeV']= np.linspace(0, (ntime - 1) * time_step2, ntime) + + print('rebinned rates: nchan=%d dt=%.3f sec rebin=%d' % (nchan, time_step2, tReb2)) + return rebD + +#...!...!.................... +def XXXselect_radius_slice(spikeD, data_path, data_name, idxState, verb=1): + """Select one spectral-radius slice from stacked spikes arrays if present.""" + spikes = spikeD['spikes'] + if spikes.ndim == 2: + return spikeD, None + assert spikes.ndim == 3, f"Expected spikes ndim 2 or 3, got shape={spikes.shape}" + + nR = spikes.shape[0] + assert 0 <= idxState < nR, f"idxState={idxState} out of range for spikes with nR={nR}" + + truthFF = os.path.join(data_path, f"{data_name}.simTruth.npz") + assert os.path.exists(truthFF), f"missing simTruth file: {truthFF}" + _, trueMD = read_data_npz(truthFF, verb=verb>0) + spect_radii = trueMD['dale_conf']['spect_radius'] + assert idxState < len(spect_radii), f"idxState={idxState} out of range, only {len(spect_radii)} states available" + R_sel = spect_radii[idxState] + print(f"\nSelected spectral radius [{idxState}]: R={R_sel:.3f} (out of {spect_radii})") + + spikeD_r = {} + for key, arr in spikeD.items(): + if isinstance(arr, np.ndarray) and arr.ndim > 0 and arr.shape[0] == nR: + spikeD_r[key] = arr[idxState] + else: + spikeD_r[key] = arr + return spikeD_r, R_sel + + + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__=="__main__": + args=get_parser() + np.set_printoptions(precision=3) + + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb>0) + if args.verb>1: pprint(spikeMD) + + S_oracle = None + if args.idxState >=0: # per state information + # Load simulation truth data (Dale matrices, biases, etc.) + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb>0) + if args.verb>1: + print("\nSimulation Truth Metadata:"); pprint(trueMD) + + dataYield = spikeD['spikes'] + Mstate,Nt, Nn = dataYield.shape + m=args.idxState + assert m0) + S_true = prismD['S_true'] + S_oracle = prismD['S_oracle'] + if args.verb > 0: print('loaded prismTruth S_true:', S_true.shape); + if args.verb > 1: pprint(prismMD) + spikeMD['sel_spect_radius'] =-77 + oraE = prismMD['oracle_eval'] + spikeMD['oracle_score'] = oraE['avr_score'] + print(f"gen, oracle avr score {oraE['avr_score']:.3f}, {args.dataName}") + print(f" {'state':>5s} {'score':>5s}") + print(f" {'-----':>5s} {'-----':>5s}") + for m, sc in enumerate(oraE['score_per_state']): + print(f" {m:5d} {sc:5.3f}") + + #-------------------------------- + # .... plotting ........ + spikeMD['short_name']=f"{args.dataName}_state{args.idxState}" + spikeMD['sel_state'] = int(args.idxState) + args.prjName=spikeMD['short_name'] + spikeMD['plot']={} + spikeMD['plot']['time_rangeLR']=np.array(args.time_range_sec) + + plot=Plotter(args) + + if 'a' in args.showPlots: + plot.freq_histo(spikeD,spikeMD,figId=1) + if 'b' in args.showPlots: + rebD=rebin_spike_rates(spikeD['spikes'], spikeMD, args.time_rebin2) + if S_oracle is not None: + rebD['S_oracle'] = S_oracle + plot.freq_vs_time(rebD,spikeMD,figId=2, S_true=S_true, S_oracle=S_oracle) + + plot.display_all() + print('M:done') + #pprint(expMD) #tmp diff --git a/causal_net/nonStation_ver4_lagM_offDiag/PlotterDaleMatrix.py b/causal_net/nonStation_ver4_lagM_offDiag/PlotterDaleMatrix.py new file mode 100644 index 00000000..46827f1f --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/PlotterDaleMatrix.py @@ -0,0 +1,686 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for simulated Poisson process visualization. + +This module provides specialized plotting capabilities for analyzing +simulated Dale's principle neural networks. The Plotter class extends +PlotterBackbone to create visualizations including: +- Dale connectivity matrix plots with excitatory/inhibitory separation +- Eigenvalue analysis and network stability visualization +- Poisson process statistics and firing rate distributions +- Network dynamics and temporal evolution plots + +Designed specifically for validating and analyzing simulated neural +networks that follow Dale's principle with Poisson spiking dynamics. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec +import matplotlib.colors as colors +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.collections import LineCollection +from matplotlib.lines import Line2D +from scipy.optimize import curve_fit +from Util_pseudospectra import compute_pseudospectrum +from UtilDalePoisson4 import get_offdiag_triplets + + +def _placement_ker_delta(dmd): + return float(dmd["placement_ker_delta"]) + + +def _format_ker_delta_latex(delta): + """Matplotlib title fragment: δ_ker value (integers without .0).""" + x = float(delta) + if abs(x - round(x)) < 1e-9: + return str(int(round(x))) + return "%.4g" % x + + +def _kernel_figure_canvas_title(md): + """One-line caption: dataset, N, and (model B) Q, tau, placement H×L.""" + dmd = md["dale_conf"] + parts = [] + sn = md.get("short_name", "") + if sn: + parts.append("dataset: %s" % sn) + n = dmd.get("num_neurons") + if n is not None: + parts.append("N=%d" % int(n)) + if dmd.get("spike_model") == "B": + if "mem_Q" in dmd: + parts.append("Q=%.4g" % float(dmd["mem_Q"])) + if "mem_tau" in dmd: + parts.append("tau/sec=%.4g" % float(dmd["mem_tau"])) + ph = dmd.get("placement_H") + pl = dmd.get("placement_L") + if ph is not None and pl is not None: + parts.append("placement HxL=(%gx%g)" % (float(ph), float(pl))) + parts.append("δ_ker=%s" % _format_ker_delta_latex(_placement_ker_delta(dmd))) + return " ".join(parts) + + +def _dale_overview_title(md): + dmd = md["dale_conf"] + return r"True Dale, $N=%d$, $\delta_{\mathrm{ker}}=%s$, %s" % ( + int(dmd["num_neurons"]), + _format_ker_delta_latex(_placement_ker_delta(dmd)), + md["short_name"], + ) + + +def _realized_edge_lengths(trueD): + D = np.asarray(trueD["node_distance_matrix"], dtype=np.float64) + if D.ndim != 2 or D.shape[0] != D.shape[1]: + raise ValueError("node_distance_matrix must be square (N, N)") + N = D.shape[0] + edge_mask = np.asarray(trueD["E_true"]) != 0 + if edge_mask.shape != (N, N): + raise ValueError("edge mask shape mismatch with node_distance_matrix") + edge_mask = edge_mask.copy() + np.fill_diagonal(edge_mask, False) + d_off = D[edge_mask] + return d_off[d_off > 0] + + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + + def _plot_signed_offdiag_matrix(self, ax, trueD, md, title_txt=None): + dmd = md["dale_conf"] + numNeur = int(dmd["num_neurons"]) + placement_ker_delta = _placement_ker_delta(dmd) + + if "E_true" in trueD: + A_sign = np.asarray(trueD["E_true"], dtype=np.int8).copy() + else: + A_sign = np.sign(np.asarray(trueD["A_off_true"], dtype=np.float64)).astype(np.int8, copy=False) + if A_sign.shape != (numNeur, numNeur): + raise ValueError("signed off-diagonal topology shape mismatch") + np.fill_diagonal(A_sign, 0) + + sign_cmap = colors.ListedColormap(["blue", "white", "red"]) + sign_norm = colors.BoundaryNorm([-1.5, -0.5, 0.5, 1.5], sign_cmap.N) + ax.imshow( + A_sign.T, + aspect=1.0, + origin="lower", + cmap=sign_cmap, + norm=sign_norm, + interpolation="nearest", + ) + ax.set( + ylabel="neuron index (postsynaptic)", + xlabel="presyn. neuron index (target)", + ) + ax.set_aspect(1.0) + ax.grid(True, alpha=0.45) + if title_txt is None: + title_txt = _dale_overview_title(md) + ax.set_title(title_txt) + ax.plot([0, numNeur], [0, numNeur], "--", lw=0.8, color="magenta") + +#...!...!.................. + def Dale_matrix_and_eigen(self,A,md,trueD,figId=3): + + figId=self.smart_append(figId) + nrow,ncol=1,2 + fig=self.plt.figure(figId,facecolor='white', figsize=(12,5.)) + + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + placement_ker_delta = _placement_ker_delta(dmd) + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel if R_sel is not None else '' + + vmin = A.min() + vmax = A.max() + normMap = colors.TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax) + + #.... Position 1: Dale matrix (transpose: x = postsynaptic i, y = presynaptic j) ...... + ax = self.plt.subplot(nrow,ncol,1) + im = ax.imshow( + A.T, aspect=1., origin='lower', cmap='bwr', norm=normMap, interpolation='nearest' + ) + ax.set( + ylabel=' neuron index (postsynaptic)', + xlabel='presyn. neuron index (target)', + ) + ax.set_aspect(1.0) + ax.grid() + cbar = fig.colorbar(im, ax=ax, extend="both", shrink=0.7) + + tit_left = r"True Dale, $N=%d$, $\delta_{\mathrm{ker}}=%s$, %s" % ( + A.shape[0], + _format_ker_delta_latex(placement_ker_delta), + md["short_name"], + ) + ax.set(title=tit_left) + ax.plot([0,numNeur],[0,numNeur],'--',lw=0.8,color='magenta') + + #..... Position 2: Eigenvalues...... + ax = self.plt.subplot(nrow,ncol,2) + Eigen=np.linalg.eigvals(A) + real_parts = np.real(Eigen) + imag_parts = np.imag(Eigen) + ax.scatter(real_parts, imag_parts, color='blue', marker='o') + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + tit3='Eigenvalues, N%d%s, %s'%(A.shape[0], R_tag, md['short_name']) + ax.set_title(tit3) + ax.axhline(0, color='black', lw=0.5) + ax.axvline(0, color='black', lw=0.5) + ax.axvline(0,color='red', linestyle='--') + ax.set_aspect('equal') + if R_sel is not None: + theta = np.linspace(0, 2*np.pi, 200) + ax.plot(R_sel*np.cos(theta), R_sel*np.sin(theta), color='magenta', linestyle='--', lw=1.5, label='R=%.3f'%R_sel) + ax.legend() + ax.grid(True) + + def Dale_matrix_pseudospectra(self, A, md, trueD, figId=5): # p=e + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(8, 6)) + + npts = 80 + minY = -0.2 + epsMin = 0.03 + + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + ax = self.plt.subplot(1, 1, 1) + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel if R_sel is not None else '' + title = 'Pseudospectra, true N%d%s, %s' % (A.shape[0], R_tag, md['short_name']) + + levels = np.logspace(-2.5, -0.5, 10) + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + ax.text(0.05, 0.85, 'green: ε<%.3f'%epsMin, transform=ax.transAxes) + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) + + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(minY,1.1) + ax.set_xlim(-1.1,1.1) + ax.set_aspect(1.) + ax.axhline(1,color='m',linestyle='--') + ax.axvline(1,color='m',linestyle='--') + ax.axvline(-1,color='m',linestyle='--') + +#...!...!.................. + def histo_weights_rates(self,trueD,spikeD,md,figId=3): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,3.5)) + data_name=md['short_name'] + dmd=md['dale_conf'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.2f' % R_sel + + A=trueD['A_true'] + single_rates = spikeD['single_rates'] + + # output: np.column_stack([i_indices, j_indices, values]) + EposT=get_offdiag_triplets(A,isPos=True) + EnegT=get_offdiag_triplets(A,isPos=False) + print('True0 num edges pos=%d neg=%d'%(EposT.shape[0],EnegT.shape[0])) + + A_diag = np.diag(A) + wMin = min(np.min(EnegT[:,2]), np.min(A_diag)) + wMax = max(np.max(EposT[:,2]), np.max(A_diag)) + + def count_elements(E, Nn=numNeur): + i_indices = E[:, 0].astype(int) + counts = np.bincount(i_indices, minlength=Nn) + return counts + + edgeCount=count_elements(EnegT) + count_elements(EposT) + + #.... weights histogram + ax = self.plt.subplot(nrow,ncol,1) + binX= np.linspace(wMin, wMax, 100) + ax.hist(EnegT[:,2], bins=binX, color='blue', alpha=0.7, edgecolor=None,label='neg:%d'%EnegT.shape[0]) + ax.hist(EposT[:,2], bins=binX, color='red', alpha=0.7, edgecolor=None,label='pos:%d'%EposT.shape[0]) + ax.hist(A_diag, bins=binX, color='cyan', alpha=0.65, edgecolor=None, + label='diag:%d'%A_diag.shape[0]) + + ax.legend(loc='upper left') + tit='True Dale, N%d%s, %s'%(A.shape[0], R_tag, md['short_name']) + ax.set(title=tit, xlabel='True weight value',ylabel='num edges') + ax.axvline(0,color='k',ls='--') + ax.grid(True, alpha=0.3) + + #.... : histogram of rates + ax = self.plt.subplot(nrow,ncol,3) + ax.hist(single_rates, bins=20) + ax.set_xlabel('Firing rate (Hz)') + ax.set_ylabel('num neurons') + ax.grid(True, alpha=0.3) + ax.set_title('Single rates spectrum') + ax.set_xlim(0,) + median_val = np.median(single_rates) + ax.axvline(median_val, color='r', linestyle='--', linewidth=1.5) + y_max = ax.get_ylim()[1] + median_text = f"median: {median_val:.2f} (Hz)\n N={single_rates.shape[0]}" + ax.text( x=median_val * 1.1, y=y_max * 0.7, s=median_text, color='red') + + tau = np.asarray(trueD["node_is_inhibitory"]).reshape(-1) + if tau.shape[0] != numNeur: + raise ValueError("node_is_inhibitory length must match num_neurons") + if tau.dtype.kind not in "iu": + tau = tau.astype(np.int64) + if np.any((tau != 0) & (tau != 1)): + raise ValueError("node_is_inhibitory must be 0 (exc) or 1 (inh)") + exc_mask = tau == 0 + inh_mask = tau == 1 + neurXlab = 'neuron index' + + x_vals = np.arange(numNeur) + #.... edge count + ax = self.plt.subplot(nrow,ncol,2) + ax.fill_between(x_vals, edgeCount, step='mid', color='salmon', alpha=0.7) + ax.set_xlabel(neurXlab) + ax.set_ylabel('num true edges') + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + probLo, probHi = dmd['edge_prob'] + rho_title = f'outgoing edges, true, prob=[{probLo:.2f}, {probHi:.2f}]' + ax.set_title(rho_title) + + + #.... firing rates ..... + ax = self.plt.subplot(nrow,ncol,4) + chanW=0.9 + + ax.bar(x_vals[exc_mask], single_rates[exc_mask], width=chanW, color='red', align='center', alpha=0.7, label='Excitatory') + ax.bar(x_vals[inh_mask], single_rates[inh_mask], width=chanW, color='blue', align='center', alpha=0.7, label='Inhibitory') + + ax.set_xlabel(neurXlab) + ax.set_ylabel('Firing rate (Hz)') + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + state_tag = md.get('sel_state', None) + if state_tag is None: + ax.set_title('Single Neurons') + else: + ax.set_title(f'Single Neurons, state={state_tag}') + ax.legend() + + +#............................ +#............................ +#............................ + def rates_study(self,trueD,spikeD,md,figId=4): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,4)) + data_name=md['short_name'] + dmd=md['dale_conf'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel if R_sel is not None else '' + + B_idle = np.asarray(trueD["B_true"]).reshape(-1) + single_rates = np.asarray(spikeD["single_rates"]).reshape(-1) + if B_idle.shape[0] != numNeur: + raise ValueError("B_true length must match dale_conf num_neurons") + if single_rates.shape[0] != numNeur: + raise ValueError("single_rates length must match num_neurons (same order as B_true, node_is_inhibitory)") + + tau = np.asarray(trueD["node_is_inhibitory"]).reshape(-1) + if tau.shape[0] != numNeur: + raise ValueError("node_is_inhibitory length must match num_neurons") + if tau.dtype.kind not in "iu": + tau = tau.astype(np.int64) + if np.any((tau != 0) & (tau != 1)): + raise ValueError("node_is_inhibitory must be 0 (exc) or 1 (inh)") + exc_mask = tau == 0 + inh_mask = tau == 1 + n_exc = int(np.sum(exc_mask)) + n_inh = int(np.sum(inh_mask)) + + # 1) Scatter: x=B_idle, y=log(single_rates) + ax = self.plt.subplot(nrow,ncol,1) + ax.scatter(B_idle[exc_mask], single_rates[exc_mask], s=14, alpha=0.6, marker='^', facecolors='none', edgecolors='red', label='Excitatory') + ax.scatter(B_idle[inh_mask], single_rates[inh_mask], s=14, alpha=0.6, marker='^', facecolors='none', edgecolors='blue', label='Inhibitory') + ax.set_xlabel('true B_idle ') + ax.set_ylabel('single_rates (Hz)') + ax.set_yscale('log') + ax.grid(True, alpha=0.3) + tit='True Dale, N%d%s, %s'%(numNeur, R_tag, data_name) + ax.set_title(tit) + + # reference line: y = exp(B), clipped to central 80% of B range + bmin = float(np.min(B_idle)) + bmax = float(np.max(B_idle)) + bLo = bmin + 0.1 * (bmax - bmin) + bHi = bmax - 0.1 * (bmax - bmin) + bx = np.linspace(bLo, bHi, 100) + by = np.exp(bx) + ax.plot(bx, by, linestyle='--', color='black', linewidth=0.8, label='y=exp(x)') + ax.legend() + + # 2) Histogram: true B_idle (all neurons) + ax2 = self.plt.subplot(nrow,ncol,2) + b_bins = min(60, max(20, int(np.sqrt(numNeur) * 3))) + ax2.hist(B_idle, bins=b_bins, color='dimgray', alpha=0.8, edgecolor=None) + ax2.set_xlabel('true B_idle') + ax2.set_ylabel('num neurons') + ax2.grid(True, alpha=0.3) + state_tag = md.get('sel_state', None) + if state_tag is None: + ax2.set_title('true B_idle') + else: + ax2.set_title(f'true B_idle, state={state_tag}') + + # 3) Histogram: single_rates for excitatory + ax3 = self.plt.subplot(nrow,ncol,3) + exc_vals = single_rates[exc_mask] + inh_vals = single_rates[inh_mask] + # compute common bins and range (start at 0) + exc_max = np.max(exc_vals) if exc_vals.size > 0 else 0.0 + inh_max = np.max(inh_vals) if inh_vals.size > 0 else 0.0 + x_max = max(exc_max, inh_max) + if x_max <= 0: x_max = 1.0 + #common_bins = np.linspace(0, x_max, 21) + common_bins = np.linspace(0, x_max, int(2*x_max)) + ax3.hist(exc_vals, bins=common_bins, color='red', alpha=0.7, edgecolor=None) + ax3.set_xlabel('single_rates (Hz)') + ax3.set_ylabel('num neurons') + ax3.grid(True, alpha=0.3) + ax3.set_title('Rates: Excitatory (N=%d)' % (n_exc)) + + # 4) Histogram: single_rates for inhibitory + ax4 = self.plt.subplot(nrow,ncol,4) + ax4.hist(inh_vals, bins=common_bins, color='blue', alpha=0.7, edgecolor=None) + ax4.set_xlabel('single_rates (Hz)') + ax4.set_ylabel('num neurons') + ax4.grid(True, alpha=0.3) + ax4.set_title('Rates: Inhibitory (N=%d)' % (n_inh)) + # unify x-range starting at 0 for both histograms + ax3.set_xlim(0, x_max) + ax4.set_xlim(0, x_max) + + def _plot_placement_topology_ax(self, ax, trueD, md, title_txt=None, arch_radius=None): + dmd = md["dale_conf"] + L = float(dmd["placement_L"]) + H = float(dmd["placement_H"]) + if not (L > 0 and H > 0): + raise ValueError("placement_L and placement_H must be positive") + placement_ker_delta = _placement_ker_delta(dmd) + + P = np.asarray(trueD["node_positions"], dtype=float) + tau = np.asarray(trueD["node_is_inhibitory"]).reshape(-1) + if tau.dtype.kind not in "iu": + tau = tau.astype(np.int64) + N = P.shape[0] + if tau.shape[0] != N: + raise ValueError("node_positions and node_is_inhibitory length mismatch") + exc_m = tau == 0 + inh_m = tau == 1 + if not np.any(exc_m) or not np.any(inh_m): + raise ValueError("node_is_inhibitory must contain both excitatory (0) and inhibitory (1) labels") + + M = np.asarray(trueD["E_true"]) + mask = M != 0 + if M.shape != (N, N): + raise ValueError("E_true shape mismatch") + + red_segs = [] + blue_segs = [] + for j in range(N): + for i in range(N): + if i == j or not mask[i, j]: + continue + seg = np.array([P[j], P[i]]) + if exc_m[j]: + red_segs.append(seg) + elif inh_m[j]: + blue_segs.append(seg) + + if red_segs: + ax.add_collection(LineCollection(red_segs, colors="red", linewidths=0.9, alpha=0.5, zorder=1)) + if blue_segs: + ax.add_collection(LineCollection(blue_segs, colors="blue", linewidths=0.9, alpha=0.5, zorder=1)) + + ax.scatter( + P[exc_m, 0], + P[exc_m, 1], + s=65, + marker="^", + facecolors="none", + edgecolors="darkred", + linewidths=1.2, + zorder=3, + ) + ax.scatter( + P[inh_m, 0], + P[inh_m, 1], + s=55, + marker="s", + facecolors="none", + edgecolors="darkblue", + linewidths=1.2, + zorder=3, + ) + + label_color = "black" + for k in range(0, N, 10): + ax.text( + P[k, 0], + P[k, 1], + str(k), + color=label_color, + fontsize=8, + ha="left", + va="bottom", + bbox=dict(boxstyle="round,pad=0.12", facecolor="white", edgecolor="none", alpha=0.9), + zorder=4, + ) + + leg_handles = [] + if red_segs: + leg_handles.append(Line2D([0], [0], color="red", alpha=0.6, lw=2, label="Exc outgoing")) + if blue_segs: + leg_handles.append(Line2D([0], [0], color="blue", alpha=0.6, lw=2, label="Inh outgoing")) + leg_handles.append( + Line2D( + [0], + [0], + marker="^", + color="w", + markerfacecolor="none", + markeredgecolor="darkred", + markersize=9, + label="Excitatory", + ) + ) + leg_handles.append( + Line2D( + [0], + [0], + marker="s", + color="w", + markerfacecolor="none", + markeredgecolor="darkblue", + markersize=8, + label="Inhibitory", + ) + ) + leg_handles.append(Line2D([0], [0], color="lime", alpha=0.95, lw=2, label="1,2,3 hops")) + + if title_txt is None: + title_txt = r"Neuron placement, $N=%d$, $\delta_{\mathrm{ker}}=%s$, %s" % ( + N, + _format_ker_delta_latex(placement_ker_delta), + md["short_name"], + ) + ax.set_title(title_txt, pad=36) + ax.set_xlabel("x (placement)") + ax.set_ylabel("y (placement)") + ax.grid(True, alpha=0.3) + ax.legend( + handles=leg_handles, + loc="lower center", + bbox_to_anchor=(0.5, 1.0), + ncol=len(leg_handles), + fontsize=9, + frameon=True, + ) + + if arch_radius is not None and np.isfinite(arch_radius) and arch_radius > 0: + thetaV = np.linspace(-0.5 * np.pi, 0.5 * np.pi, 300) + for x_center in (0.0, arch_radius, 2.0 * arch_radius): + xV = x_center + arch_radius * np.cos(thetaV) + yV = 0.5 + arch_radius * np.sin(thetaV) + keep = np.abs(yV) <= 1.0 + if not np.any(keep): + continue + xP = np.where(keep, xV, np.nan) + yP = np.where(keep, yV, np.nan) + ax.plot(xP, yP, color="lime", linewidth=2.0, alpha=0.95, zorder=2.2, clip_on=True) + + mx, my = 0.2, 0.05 + ax.set_xlim(0.0 - mx, L + mx) + ax.set_ylim(0.0 - my, H + my) + + ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True)) + ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) + ax.xaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: "%d" % int(round(x)))) + ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: "%d" % int(round(x)))) + + def plot_placement_topology(self, trueD, md, figId=6): + """ + 2D placement: triangles = excitatory, squares = inhibitory; + outgoing edges from j→i: red if j is excitatory, blue if inhibitory (presynaptic type). + """ + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(10, 5)) + ax = fig.add_subplot(1, 1, 1) + d_off = _realized_edge_lengths(trueD) + arch_radius = float(np.percentile(d_off, 50)) if d_off.size > 0 else None + self._plot_placement_topology_ax(ax, trueD, md, arch_radius=arch_radius) + + def _plot_offdiag_distance_hist_ax(self, ax, trueD, md, title_txt=None): + d_off = _realized_edge_lengths(trueD) + dmd = md["dale_conf"] + d_ker = _format_ker_delta_latex(_placement_ker_delta(dmd)) + pctD = None + if d_off.size > 0: + ax.hist( + d_off, + bins=min(60, max(10, d_off.size // 5)), + color="steelblue", + edgecolor="white", + alpha=0.9, + ) + p16, p50, p84 = np.percentile(d_off, [16, 50, 84]) + pctD = {"p16": float(p16), "p50": float(p50), "p84": float(p84)} + ax.axvline(p16, color="darkorange", linestyle="--", linewidth=1.3, zorder=4) + ax.axvline(p50, color="darkgreen", linestyle="--", linewidth=1.3, zorder=4) + ax.axvline(p84, color="darkviolet", linestyle="--", linewidth=1.3, zorder=4) + pct_txt = "p16 = %.5g\np50 = %.5g\np84 = %.5g" % (p16, p50, p84) + ax.text( + 0.97, + 0.97, + pct_txt, + transform=ax.transAxes, + ha="right", + va="top", + fontsize=8, + family="monospace", + bbox=dict(boxstyle="round,pad=0.35", facecolor="white", edgecolor="0.6", alpha=0.92), + zorder=5, + ) + if title_txt is None: + title_txt = r"Realized edge lengths ($\delta_{\mathrm{ker}}=%s$)" % d_ker + ax.set_title(title_txt) + ax.set_xlabel("distance") + ax.set_ylabel("count") + ax.grid(True, alpha=0.3) + return pctD + + def offdiag_distance_kernel_hist(self, trueD, md, figId=6): + """ + One row, three panels: (1) histogram of off-diagonal distances from + node_distance_matrix; (2) κ vs lag (model B); (3) empty. + """ + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(14, 4)) + gs = gridspec.GridSpec(1, 3, figure=fig, wspace=0.35) + + ax1 = fig.add_subplot(gs[0, 0]) + self._plot_offdiag_distance_hist_ax(ax1, trueD, md) + + ax2 = fig.add_subplot(gs[0, 1]) + kappa = np.asarray(trueD["offdiag_kernel"], dtype=np.float64).ravel() + if kappa.size > 0: + lags = np.arange(1, kappa.size + 1, dtype=np.float64) + ax2.plot(lags, kappa, color="darkorange", marker="o", markersize=4, linewidth=1.0) + ax2.axhline(0.0, color="k", lw=0.5) + ax2.set_title(r"offdiag_kernel $\kappa_\ell$ vs lag $\ell$") + ax2.set_xlabel(r"lag $\ell$") + ax2.set_ylabel(r"$\kappa$") + ax2.grid(True, alpha=0.3) + else: + ax2.text( + 0.5, + 0.5, + r"model A: lag-1 (empty $\kappa$, $M=0$)", + ha="center", + va="center", + transform=ax2.transAxes, + ) + ax2.set_axis_off() + + ax3 = fig.add_subplot(gs[0, 2]) + ax3.set_axis_off() + + title_txt = _kernel_figure_canvas_title(md) + fig.suptitle(title_txt, fontsize=10, y=0.93) + fig.subplots_adjust(top=0.74) + + def topo_overview(self, trueD, md, figId=7): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(13, 10)) + gs = gridspec.GridSpec(2, 2, figure=fig, height_ratios=[1.0, 1.35], hspace=0.42, wspace=0.28) + + ax1 = fig.add_subplot(gs[0, 0]) + self._plot_signed_offdiag_matrix(ax1, trueD, md, title_txt="Signed off-diagonal topology") + + ax2 = fig.add_subplot(gs[0, 1]) + pctD = self._plot_offdiag_distance_hist_ax(ax2, trueD, md, title_txt="Realized edge lengths") + + ax3 = fig.add_subplot(gs[1, :]) + arch_radius = pctD["p50"] + self._plot_placement_topology_ax(ax3, trueD, md, title_txt="Neuron placement", arch_radius=arch_radius) + fig.suptitle(_dale_overview_title(md), fontsize=16, y=0.985) + fig.subplots_adjust(top=0.90) diff --git a/causal_net/nonStation_ver4_lagM_offDiag/PlotterMemKernEM.py b/causal_net/nonStation_ver4_lagM_offDiag/PlotterMemKernEM.py new file mode 100644 index 00000000..126ece92 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/PlotterMemKernEM.py @@ -0,0 +1,494 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for prism EM evaluation. +""" + +from toolbox.PlotterBackbone import PlotterBackbone +import numpy as np +from matplotlib.ticker import MaxNLocator +import matplotlib.colors as colors + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def _add_aoff_hist(self, ax, a_off_hat): + """Plot non-zero off-diagonal edge weights on a log-count histogram.""" + a_fit = np.asarray(a_off_hat) + nn = a_fit.shape[0] + diag_mask = np.eye(nn, dtype=bool) + a_off = a_fit[~diag_mask] + valid = np.abs(a_off) > 1e-10 + a_off_nz = a_off[valid] + n_edges = int(a_off_nz.size) + if n_edges > 0: + ax.hist(a_off_nz, bins=100, color='g', alpha=0.8) + ax.set_yscale('log') + else: + ax.text(0.5, 0.5, "No non-zero off-diagonal edges", + ha='center', va='center', transform=ax.transAxes) + ax.set(title=f"A off-diagonal, {n_edges} edges", + xlabel="edge value", ylabel="edges") + ax.grid(True, alpha=0.3) + + + def summary_memKerEM(self, fitD, md, figId=1): + """EM convergence overview: 2 rows × 3 columns. + + Row 1: E-step NLL vs EM iter | M-step NLL+L1 vs M-epoch | spectral radius vs M-epoch + Row 2: nz off-diag edges vs M-epoch | mean occupancy | A off-diag weight histogram + """ + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(12, 6)) + fig.subplots_adjust(hspace=0.45, wspace=0.35) + + trainMD = md["train"] + short_name = md["short_name"] + + e_nll = fitD["e_nll_em"] + m_nll = fitD["m_nll_epoch"] + m_l1 = fitD["m_l1_epoch"] + m_loss = fitD["m_loss_epoch"] + rho = fitD["rho_epoch"] + nz = fitD["nz_edges_epoch"] + lr = fitD["learning_rates"] + + n_em = int(trainMD["num_em_iters"]) + m_per_em = int(trainMD["m_epochs"]) + n_m_total = len(m_nll) + m_epochs = np.arange(1, n_m_total + 1) + em_iters = np.arange(1, len(e_nll) + 1) + + prune_em = int(trainMD["delay_em_iter_4_Aprune"]) + rho_start_em = int(trainMD["delay_em_iter_4_ArhoMax"]) + lr_drop_em = int(trainMD["delay_em_iter_4_lrDecay"]) + prune_m_epoch = prune_em * m_per_em + rho_start_m_epoch = rho_start_em * m_per_em + lr_drop_m_epoch = lr_drop_em * m_per_em + + def draw_threshold_marker(ax, x_pos, x_max, txt, color,yFac=0.02): + if not (0 < x_pos <= x_max): + return + ax.axvline(x_pos, color=color, ls='--', lw=0.9, alpha=0.95) + y0, y1 = ax.get_ylim() + x0, x1 = ax.get_xlim() + x_off = 0.01 * max(1e-9, x1 - x0) + y_txt = y1 - yFac * (y1 - y0) + ax.text( + x_pos + x_off, y_txt, txt, rotation=90, color=color, fontsize=7, + ha='left', va='top', + bbox=dict(facecolor='white', alpha=0.55, edgecolor='none', pad=0.2) + ) + + jSkipEM = 0 + jSkipM = int(jSkipEM * m_per_em) + # ── Row 1, Col 1: E-step NLL vs EM iteration ──────────────── + ax = self.plt.subplot(2, 3, 1) + ax.plot(em_iters[jSkipEM:], e_nll[jSkipEM:], 'o-', color='tab:blue', markersize=3, + linewidth=1.2) + ax.set(title="E-step weighted NLL", xlabel="EM iteration", + ylabel="NLL / bin") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_em, len(e_nll), "start Aprune", "k") + draw_threshold_marker(ax, rho_start_em, len(e_nll), "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_em, len(e_nll), "start_lrDrop", "tab:gray", yFac=0.6) + pgd_it = trainMD["block1_iter"] + lr_e = trainMD["lr_kappa"] + txt = (f"block1_iter={pgd_it}\n" + f"lr_E={lr_e}\n" + f"λ₂={trainMD['lambda2']}") + ax.text(0.97, 0.97, txt, transform=ax.transAxes, + va="top", ha="right", fontsize=8, + bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')) + + # ── Row 1, Col 2: M-step NLL + L1 vs M-epoch ──────────────── + ax = self.plt.subplot(2, 3, 2) + ax.plot(m_epochs[jSkipM:], m_nll[jSkipM:], color='tab:blue', linewidth=1, label='NLL') + ax.set_ylabel('NLL', color='tab:blue') + ax.tick_params(axis='y', labelcolor='tab:blue') + ax.set(title="M-step loss", xlabel="M-epoch (global)") + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + ax2.plot(m_epochs[jSkipM:], m_l1[jSkipM:] , color='tab:red', linewidth=1, + linestyle='--', label='L1') + ax2.set_ylabel('L1', color='tab:red') + ax2.tick_params(axis='y', labelcolor='tab:red') + + lines1, lab1 = ax.get_legend_handles_labels() + lines2, lab2 = ax2.get_legend_handles_labels() + ax.legend(lines1 + lines2, lab1 + lab2, fontsize=7, loc='upper right') + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + + lr_m = trainMD["lr_net"] + txt = (f"lr_M={lr_m}\n" + f"L1 λ3={trainMD['lambda3']}\n" + f"batch={trainMD['batch_size']}") + ax.text(0.50, 0.50, txt, transform=ax.transAxes, + va="center", ha="center", fontsize=8, + bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')) + + # ── Row 1, Col 3: spectral radius vs M-epoch ──────────────── + ax = self.plt.subplot(2, 3, 3) + rho_max = float(trainMD["rho_max"]) + ax.plot(m_epochs[jSkipM:], rho[jSkipM:] , color='tab:green', linewidth=1, + label='ρ(A)') + ax.axhline(rho_max, color='red', ls='--', lw=1, + label=f'ρ_max={rho_max}') + ax.set(title="Spectral radius ρ(A)", xlabel="M-epoch (global)", + ylabel="ρ") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + ax.legend(fontsize=8) + + # ── Row 2, Col 1: non-zero off-diag edges vs M-epoch ──────── + ax = self.plt.subplot(2, 3, 4) + ax.plot(m_epochs[jSkipM:], nz[jSkipM:], color='tab:purple', linewidth=1) + ax.set(title=f"Non-zero off-diag edges (|A|>{trainMD['minW']})", + xlabel="M-epoch (global)", ylabel="count") + ax.grid(True, alpha=0.3) + draw_threshold_marker(ax, prune_m_epoch, n_m_total, "start Aprune", "k") + draw_threshold_marker(ax, rho_start_m_epoch, n_m_total, "start rhoMax", "tab:brown") + draw_threshold_marker(ax, lr_drop_m_epoch, n_m_total, "start_lrDrop", "tab:gray", yFac=0.4) + + # learning rate on twin axis + ax2 = ax.twinx() + ax2.plot(m_epochs[jSkipM:], lr[jSkipM:], color='tab:orange', linewidth=0.8, + linestyle='--', alpha=0.6) + ax2.set_ylabel('learning rate', color='tab:orange') + ax2.tick_params(axis='y', labelcolor='tab:orange') + + # ── Row 2, Col 2: temporal kernel kappa_hat ────────────────── + ax = self.plt.subplot(2, 3, 5) + kappa_hat = np.asarray(fitD["kappa_hat"]) + lags = np.arange(1, len(kappa_hat) + 1) + ax.step(lags, kappa_hat, where='post', color='tab:blue', label='fit') + ax.plot(lags, kappa_hat, 'o', markersize=3, alpha=0.5, color='tab:blue') + kappa_true_arr = np.asarray(md["offdiag_kernel"]).ravel() + if kappa_true_arr.size > 0: + kappa_true = kappa_true_arr + lags_true = np.arange(1, len(kappa_true) + 1) + ax.step(lags_true, kappa_true, where='post', color='tab:red', linestyle='--', alpha=0.7, label='truth') + kappa_init = np.asarray(fitD["kappa_init"]) + ax.step(lags, kappa_init, where='post', color='black', linestyle='--', alpha=0.6, label='init') + ax.legend(fontsize=8) + ax.set(title="Temporal kernel $\kappa(\ell)$", xlabel="lag $\ell$", ylabel="weight") + ax.grid(True, alpha=0.3) + ax.axhline(0, color='k', lw=0.8, alpha=0.5) + + # ── Row 2, Col 3: A off-diagonal weight histogram ──────────── + ax = self.plt.subplot(2, 3, 6) + self._add_aoff_hist(ax, fitD["A_off_hat"]) + + fig.suptitle(f"MemKern EM: {short_name}, " + f"K_EM={n_em}×K_M={m_per_em}", + fontsize=12) + fig.tight_layout(rect=[0, 0, 1, 0.95]) + + def correl_fit_truth(self, fitD, md, figId=2): + """Four-panel correlation plot (top) + Residual Histograms (bottom).""" + figId = self.smart_append(figId) + # 2x8 grid: Top row 4 plots (2cols each), Bottom row 6 plots (1col each, centered) + # Use height_ratios to make diagnostic histograms 1/2 of correlation plots + fig = self.plt.figure(figId, facecolor='white', figsize=(20.0, 7.5)) + gs = fig.add_gridspec(2, 8, height_ratios=[2.0, 1.0], hspace=0.35, wspace=0.45, + left=0.05, right=0.95, top=0.9, bottom=0.08) + + short_name = md["short_name"] + A_off_true = md["A_off_true"] + A_diag_true = md["A_diag_true"] + + # Prepare connectivity data + if A_off_true.ndim == 2: + diag_mask = np.eye(A_off_true.shape[0], dtype=bool) + A_off_true = A_off_true[~diag_mask] + + A_off_hat_full = np.asarray(fitD["A_off_hat"]) + diag_mask = np.eye(A_off_hat_full.shape[0], dtype=bool) + A_off_hat = A_off_hat_full[~diag_mask] + A_diag_hat = np.asarray(fitD["A_diag_hat"]) + B_true = md["B_true"] + B_hat = np.asarray(fitD["B_hat"]) + + minW = float(md["train"]["minW"]) + B_div = 2.0 # Common dividing line for B groups + + def get_corr_stats(x, y): + if len(x) < 2: + return 0.0, len(x) + if np.allclose(x, x[0]) or np.allclose(y, y[0]): + return 0.0, len(x) + r = np.corrcoef(x, y)[0, 1] + if not np.isfinite(r): + r = 0.0 + return float(r), len(x) + + def add_corr_ax(ax, x, y, labels, title, x_split=0, min_w=None, show_split=True, dot_color='blue'): + x = np.asarray(x).ravel() + y = np.asarray(y).ravel() + if x.size == 0 or y.size == 0: + ax.axhline(0, color='k', lw=0.6, alpha=0.3) + if show_split: + ax.axvline(x_split, color='red', lw=1.0, ls='--') + ax.text(0.5, 0.5, "No points to plot", transform=ax.transAxes, + va='center', ha='center', fontsize=11, + bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')) + ax.set(title=title, xlabel=labels[0], ylabel=labels[1]) + ax.grid(True, alpha=0.2) + if min_w is not None: + for side in [-1, 1]: + ax.axvline(side * min_w, color='red', ls='--', lw=0.8, alpha=0.6) + return + ax.scatter(x, y, s=10, alpha=0.3, color=dot_color, edgecolors='none') + ax.axhline(0, color='k', lw=0.6, alpha=0.3) + if show_split: + ax.axvline(x_split, color='red', lw=1.0, ls='--') + + low, high = min(x.min(), y.min()), max(x.max(), y.max()) + ax.plot([low, high], [low, high], 'k--', lw=0.8, alpha=0.7) + + maskL, maskR = (x < x_split), (x > x_split) + rL, nL = get_corr_stats(x[maskL], y[maskL]) + rR, nR = get_corr_stats(x[maskR], y[maskR]) + + if nL > 0: ax.plot(np.mean(x[maskL]), np.mean(y[maskL]), 'k+', ms=12, mew=2) + if nR > 0: ax.plot(np.mean(x[maskR]), np.mean(y[maskR]), 'k+', ms=12, mew=2) + + ax.text(0.05, 0.5, f"rL={rL:.3f}\nnL={nL}", transform=ax.transAxes, va='center', ha='left', fontsize=11) + ax.text(0.95, 0.4, f"rR={rR:.3f}\nnR={nR}", transform=ax.transAxes, va='center', ha='right', fontsize=11) + ax.set(title=title, xlabel=labels[0], ylabel=labels[1]) + ax.grid(True, alpha=0.2) + + if min_w is not None: + for side in [-1, 1]: ax.axvline(side*min_w, color='red', ls='--', lw=0.8, alpha=0.6) + n0 = np.sum(np.abs(x) <= min_w) + ax.text(-4*min_w, 0.1, f"n0={n0}", color='red', transform=ax.get_xaxis_transform(), + rotation='vertical', va='bottom', ha='center', fontsize=11) + + def add_resid_hist(ax, resid, title, color): + if len(resid) == 0: + ax.text(0.5, 0.5, "Empty", ha='center', va='center', transform=ax.transAxes) + return + med, std = np.median(resid), np.std(resid) + q25, q75 = np.percentile(resid, [25, 75]) + ax.hist(resid, bins=30, color=color, alpha=0.5, density=True) + ax.axvline(0, color='k', ls='--', lw=1.0, alpha=0.8) + ax.axvline(q25, color='k', ls=':', lw=0.8, alpha=0.6) + ax.axvline(q75, color='k', ls=':', lw=0.8, alpha=0.6) + y_lim = ax.get_ylim() + y0 = (y_lim[0] + y_lim[1]) * 0.5 + ax.errorbar(med, y0, xerr=std, fmt='ko', ms=8, elinewidth=2, capsize=4) + ax.text(0.05, 0.95, f"med={med:.3f}\nstd={std:.3f}", transform=ax.transAxes, fontsize=8, + va='top', bbox=dict(facecolor='white', alpha=0.7)) + ax.set(title=title, xlabel="Residual (fit-true)") + ax.grid(True, alpha=0.2) + + # ── Row 1: Correlations (span 2 columns each) ───────────────── + ax1 = fig.add_subplot(gs[0, 0:2]) + valid = np.abs(A_off_hat) >= minW + add_corr_ax(ax1, A_off_true[valid], A_off_hat[valid], ["A_off_true", "A_off_hat"], + f"$A_{{off}}$ fit, minW={minW:.2f}", min_w=minW, show_split=False, dot_color='green') + + ax2 = fig.add_subplot(gs[0, 2:4]) + ax2.scatter(A_diag_true, A_diag_hat, s=12, alpha=0.6, color='salmon', edgecolors='none') + ax2.plot(np.mean(A_diag_true), np.mean(A_diag_hat), 'k+', ms=12, mew=2) + r, n = get_corr_stats(A_diag_true, A_diag_hat) + l, h = A_diag_true.min(), A_diag_true.max() + ax2.plot([l, h], [l, h], 'k--', lw=0.8, alpha=0.7) + ax2.set(title=f"A_diag fit, r={r:.3f}, n={n}", xlabel="A_diag_true", ylabel="A_diag_hat") + ax2.grid(True, alpha=0.2) + + ax3 = fig.add_subplot(gs[0, 4:6]) + add_corr_ax(ax3, B_true, B_hat, ["B_true", "B_hat"], "B fit", x_split=B_div, show_split=True, dot_color='blue') + + ax4 = fig.add_subplot(gs[0, 6:8]) + kappa_hat = np.asarray(fitD["kappa_hat"]) + lags = np.arange(1, len(kappa_hat) + 1) + ax4.step(lags, kappa_hat, where='post', color='tab:blue', label='kappa_fit', lw=2) + kappa_true_arr = np.asarray(md["offdiag_kernel"]).ravel() + if kappa_true_arr.size > 0: + kappa_true = kappa_true_arr + lags_true = np.arange(1, len(kappa_true) + 1) + ax4.step(lags_true, kappa_true, where='post', color='tab:red', linestyle='--', alpha=0.8, label='kappa_true', lw=1.5) + n_overlap = min(len(kappa_hat), len(kappa_true)) + diff = np.abs(kappa_hat[:n_overlap] - kappa_true[:n_overlap]) + ax4.text(0.5, 0.98, f"sum abs err={np.sum(diff):.3f}\nmax abs err={np.max(diff):.3f}", + transform=ax4.transAxes, fontsize=9, fontweight='bold', color='tab:red', ha='center', va='top') + kappa_init = np.asarray(fitD["kappa_init"]) + ax4.step(lags, kappa_init, where='post', color='black', linestyle=':', alpha=0.5, label='kappa_init') + ax4.legend(fontsize=8, loc='best') + ax4.set(title=r"Temporal kernel $\kappa(\ell)$", xlabel=r"lag $\ell$", ylabel="weight") + ax4.grid(True, alpha=0.2) + ax4.axhline(0, color='k', lw=0.8, alpha=0.5) + + # ── Row 2: Diagnostic Row (8 cells total) ────────────────────── + # 1-3: Connectivity Residuals (A_off x2, A_diag) shifted left to col 0-2 + resids_left = [ + (A_off_hat[A_off_true < -minW] - A_off_true[A_off_true < -minW], "A_off (Inhib) resid", 'tab:blue'), + (A_off_hat[A_off_true > minW] - A_off_true[A_off_true > minW], "A_off (Excit) resid", 'tab:red'), + (A_diag_hat - A_diag_true, "A_diag residual", 'salmon'), + ] + for i, (res, tit, col) in enumerate(resids_left): + add_resid_hist(fig.add_subplot(gs[1, i]), res, tit, col) + + # 4: Correctness plot (TP/FP/FN) for A_off + ax_acc = fig.add_subplot(gs[1, 3]) + # Compute TP, FP, FN based on threshold minW for A_off + truth_mask = np.abs(A_off_true) > 1e-6 # True edges + fit_mask = np.abs(A_off_hat) >= minW # Detected edges + TP = np.sum(truth_mask & fit_mask) + FP = np.sum(~truth_mask & fit_mask) + FN = np.sum(truth_mask & ~fit_mask) + + bars = ax_acc.bar(['TP', 'FP', 'FN'], [TP, FP, FN], color=['tab:green', 'tab:red', 'tab:gray'], alpha=0.7) + # Place counts at exactly half-height of the plot area, independent of bar height + ax_acc.set_ylim(bottom=0) + y_mid = 0.5 * (ax_acc.get_ylim()[0] + ax_acc.get_ylim()[1]) + for i, val in enumerate([TP, FP, FN]): + ax_acc.text(i, y_mid, f"{val:d}", ha='center', va='center', fontsize=11, fontweight='bold', color='black') + ax_acc.set(title="A_off Edge Accuracy", ylabel="count") + ax_acc.grid(True, axis='y', alpha=0.2) + + # 5-7: Remaining residuals (B x2, Kappa) shifted to col 4-6 + resids_right = [ + (B_hat[B_true < B_div] - B_true[B_true < B_div], f"B (Low < {B_div}) resid", 'tab:blue'), + (B_hat[B_true > B_div] - B_true[B_true > B_div], f"B (High > {B_div}) resid", 'tab:red'), + ] + for i, (res, tit, col) in enumerate(resids_right): + add_resid_hist(fig.add_subplot(gs[1, i+4]), res, tit, col) + + kappa_true_r = np.asarray(md["offdiag_kernel"]).ravel() + if kappa_true_r.size > 0: + n_overlap = min(len(kappa_hat), len(kappa_true_r)) + k_res = kappa_hat[:n_overlap] - kappa_true_r[:n_overlap] + add_resid_hist(fig.add_subplot(gs[1, 6]), k_res, "Kappa residual", 'tab:green') + + self._add_aoff_hist(fig.add_subplot(gs[1, 7]), fitD["A_off_hat"]) + + fig.suptitle(f"Prism EM Diagnostics: {short_name}", fontsize=14) + + + def state_init_prismEM(self, fitD, md, figId=2, time_reb=20): + """Two-panel plot: initialization vs truth state trajectories.""" + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(13.5, 8.0)) + gs = fig.add_gridspec(3, 3, height_ratios=[1.0, 1.0, 0.85], hspace=0.75, wspace=0.35) + + trainMD = md["train"] + t0_bin, t1_bin = trainMD["time_range_bins"] + dt = float(trainMD["time_step_sec"]) + + S_init = np.asarray(fitD["S_init"]) + C_init = np.asarray(fitD["c_init"]) + S_true = np.asarray(md["S_true"])[t0_bin : t1_bin + 1] + C_true = np.asarray(md["C_true"])[t0_bin : t1_bin + 1] + + assert S_init.shape[0] == S_true.shape[0] == C_init.shape[0] == C_true.shape[0], \ + "S_init/S_true/C_init/C_true must have matching lengths" + n_cmp = S_init.shape[0] + + # Use original EM time bins (dt), not coarse rate-bin or plotting rebinning. + S_init_cl = 1.0 - np.max(C_init, axis=1) + t = (t0_bin + np.arange(n_cmp)) * dt + n_states = C_init.shape[1] + + ax = fig.add_subplot(gs[0, :]) + ax.step(t, S_init, where="post", color="k", linewidth=1.0, label="S_init") + ax.plot(t, S_init, linestyle="none", marker=".", markersize=1.5, color="k", alpha=0.6) + lo = np.clip(S_init - S_init_cl, -1.0, float(n_states - 1)) + hi = np.clip(S_init + S_init_cl, -1.0, float(n_states - 1)) + ax.fill_between(t, lo, hi, color="gray", alpha=0.3, label="S_init_CL") + ax.set(title="Initialization", xlabel="time (s)", ylabel="state") + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + for m in range(C_init.shape[1]): + ax2.step(t, C_init[:, m], where="post", linewidth=0.8, alpha=0.85, label=f"C_init[{m}]") + ax2.plot(t, C_init[:, m], linestyle="none", marker=".", markersize=1.4, alpha=0.55) + ax2.set_ylabel("C_init") + ax2.set_ylim(0.0, 1.0) + + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + ax = fig.add_subplot(gs[1, :]) + ax.step(t, S_true, where="post", color="k", linewidth=1.0, label="S_true") + ax.plot(t, S_true, linestyle="none", marker=".", markersize=1.5, color="k", alpha=0.6) + ax.set(title="Truth", xlabel="time (s)", ylabel="state") + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax2 = ax.twinx() + for m in range(C_true.shape[1]): + ax2.step(t, C_true[:, m], where="post", linewidth=0.8, alpha=0.85, label=f"C_true[{m}]") + ax2.plot(t, C_true[:, m], linestyle="none", marker=".", markersize=1.4, alpha=0.55) + ax2.set_ylabel("C_true") + + ax.legend(loc="lower left", bbox_to_anchor=(0.0, 1.02), fontsize=8) + ax2.legend(loc="lower right", bbox_to_anchor=(1.0, 1.02), ncol=4, fontsize=8) + + # Bottom row, left: histogram of classification frequency with thresholds. + ax = fig.add_subplot(gs[2, 0]) + init_state_md = md["init_state"] + + freq_h1d = np.asarray(fitD["freq_h1d"], dtype=np.float64).ravel() + bins = min(40, max(10, int(np.sqrt(max(1, freq_h1d.size))))) + ax.hist(freq_h1d, bins=bins, color='tab:blue', alpha=0.75) + rate_bin_sec = float(init_state_md["rate_bin_sec"]) + ax.set(title=f"Rate (time bin {rate_bin_sec:g} s)", xlabel="rate", ylabel="count") + ax.grid(True, alpha=0.3) + + thres = [float(x) for x in init_state_md["rate_thres"]] + for thr in thres: + ax.axvline(float(thr), color='k', linestyle='--', linewidth=1.0, alpha=0.9) + + # Label state regions 0,1,... inside histogram intervals. + xlo, xhi = ax.get_xlim() + ylo, yhi = ax.get_ylim() + edges = [xlo] + thres + [xhi] + for i in range(max(0, len(edges) - 1)): + xc = 0.5 * (edges[i] + edges[i + 1]) + ax.text( + xc, ylo + 0.88 * (yhi - ylo), f"{i}", + ha="center", va="center", fontsize=10, color="k", + bbox=dict(facecolor="white", alpha=0.6, edgecolor="none"), + ) + + # Bottom row, middle: B_init correlation vs log(single_rates). + ax_mid = fig.add_subplot(gs[2, 1]) + b_init = np.asarray(fitD["B_init"], dtype=np.float64).ravel() + single_rates = np.asarray(fitD["single_rates"], dtype=np.float64).ravel() + assert b_init.ndim == 1 and b_init.size > 0, "Expected B_init to be 1D and non-empty" + assert b_init.size == single_rates.size, "B_init and single_rates must have matching lengths" + + n = b_init.size + assert n > 1, "Need at least 2 points for correlation plot" + x = np.log(np.clip(single_rates, 1e-12, None)) + y = b_init + ax_mid.scatter(x, y, s=10, alpha=0.75, color='tab:green', edgecolors='none') + corr = float(np.corrcoef(x, y)[0, 1]) + lo = float(min(np.min(x), np.min(y))) + hi = float(max(np.max(x), np.max(y))) + ax_mid.plot([lo, hi], [lo, hi], color='k', linestyle='--', linewidth=0.8, alpha=0.6) + ax_mid.set(title=f"B_init r={corr:.3f}", + xlabel="log(single_rates)", ylabel="B_init") + ax_mid.set_aspect("equal", adjustable="box") + ax_mid.grid(True, alpha=0.3) + + # Bottom row, right: reserved cell for future diagnostics. + ax_right = fig.add_subplot(gs[2, 2]) + ax_right.axis("off") + + t0s, t1s = trainMD["time_range_sec"] + fig.suptitle( + f"Prism EM Init: {md['short_name']} T=[{t0s:.1f}, {t1s:.1f}] s", + fontsize=12, + ) + + diff --git a/causal_net/nonStation_ver4_lagM_offDiag/PlotterSpikesTrain.py b/causal_net/nonStation_ver4_lagM_offDiag/PlotterSpikesTrain.py new file mode 100644 index 00000000..97645699 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/PlotterSpikesTrain.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for biological experiment data visualization. + +This module provides specialized plotting capabilities for experimental +neural data analysis. The Plotter class extends PlotterBackbone to create +visualizations tailored for biological neural recordings including: +- Time series plots of neural activity patterns +- Statistical summaries and distribution analysis +- Data quality assessment plots +- Comparative analysis between experimental conditions + +Designed specifically for processing and visualizing data from biological +neural experiments, with automatic adaptation to experimental metadata. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +import numpy as np +import matplotlib.gridspec as gridspec + +#...!...!.................... +def summary_column(md): + return 'fix-me 32678' + #print(sorted(md)) + rs=md['rate_summary'] + ds=md['data_selector'] + #pprint(md) + txt='dataset: '+md['short_name'] + txt += '\nnum acc neurons: %d' % md['num_neurons'] + txt += '\ndrop neur: %d <%.1f Hz, %d >%.1f Hz' % (ds['drop_neur_by_freq_range'][0],ds['freq_range'][0],ds['drop_neur_by_freq_range'][1],ds['freq_range'][1]) + txt+='\ndata type: %s\ntime_step=%.2f sec'%(md['data_type'],md['time_step_sec']) + txt += '\nMedian rate: %.1f Hz' % rs['median_spike_rate'] + txt += '\nrate range [%.1f, %.1f Hz]' % (rs['min_spike_rate'], rs['max_spike_rate']) + txt += '\nAvg Rate: %.1f±%.1f Hz' % (rs['avg_spike_rate'], rs['std_spike_rate']) + txt += '\nAvg Fano: %.2f±%.2f' % (rs['avg_fano_factor'], rs['std_fano_factor']) + + return txt + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def freq_histo(self,spikeD,md,figId=1): + + tit='dataset '+md['short_name'] + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,4)) + + # Create gridspec with 7:3 width ratio + gs = gridspec.GridSpec(1, 2, width_ratios=[7, 3]) + + dataYield, dataRates = spikeD['spikes'], spikeD['single_rates'] + + # Compute median + median_val = np.median(dataRates) + txtM= f'median : {median_val:.2f} Hz' + + #.... freq per channel + ax = self.plt.subplot(gs[0, 0]) + chanV=np.arange(dataRates.shape[0]) + ax.bar(chanV, dataRates , width=0.8, color='g', align='center', alpha=0.7) + + ax.axhline(median_val, color='red', linestyle='--', linewidth=1) + if md['data_type']=='simDale': + ax.set_yscale('log') + ax.grid() + ax.set(xlabel='input neuron index', ylabel='avr frequency (Hz)',title=tit) + ax.text( ax.get_xlim()[1]*0.1,median_val,txtM, + color='red', ha='center', va='bottom', fontsize=10) + + txt=summary_column(md) + ax.text(0.05, 0.3, txt, transform=ax.transAxes, fontsize=12,color='b') + + #........ freq histo ..... + ax = self.plt.subplot(gs[0, 1]) + ax.hist(dataRates,bins=20) + #ax.set_yscale('log') + ax.set(ylabel='num channels', xlabel='avr frequency (Hz)',title=tit) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + # Add median text annotation + ax.text(median_val+1, ax.get_ylim()[1]*0.5,txtM, + color='red', ha='left', va='bottom', fontsize=10,rotation=90) + ax.text(0.25, 0.3, txt, transform=ax.transAxes, fontsize=11,color='m') + + ax.grid() + +#...!...!.................. + def freq_vs_time(self,rebD,md,figId=2,S_true=None,S_oracle=None): + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16, 7.7)) + + time_step=rebD['time_step2'] + R_sel = md['sel_spect_radius'] + + # clip time data for display + itL, itR = (md['plot']['time_rangeLR'] / time_step).astype(int) + ntime_tot = rebD['rate2D'].shape[0] + itL = max(0, min(itL, ntime_tot - 1)) + itR = max(itL + 1, min(itR, ntime_tot)) + print('iTL,R', itL,itR) + + # unpack data and clip the time range + rate2D=rebD['rate2D'][itL:itR] + timeV=rebD['timeV'][itL:itR] + + _,nchan=rate2D.shape + Tbin = time_step + xL = float(timeV[0]) + xR = float(timeV[-1] + Tbin) + if isinstance(rebD, dict) and 'pop_rate_hz' in rebD: + pop_rate_hz = np.asarray(rebD['pop_rate_hz'][itL:itR], dtype=np.float64) + else: + pop_rate_hz = np.sum(rate2D, axis=1) + # Median of per-neuron time-averaged rate (global median of all T×N bins is ~0 when activity is sparse). + mean_per_neuron_hz = np.mean(rate2D, axis=0) + medRateDisp = float(np.median(mean_per_neuron_hz)) + cntAboveMed = np.sum(rate2D > medRateDisp, axis=1) + med_sync_hz = float(np.median(pop_rate_hz)) + + tit0 = 'dataset: %s R=%.3f nchan=%d Tbin=%.2f sec' % ( + md['short_name'], + R_sel, + nchan, + time_step, + ) + if ( + md.get('spike_model') == 'B' + and 'mem_Q' in md + and 'mem_tau' in md + ): + tit0 += ' Q=%.4g tau/sec=%.4g' % ( + float(md['mem_Q']), + float(md['mem_tau']), + ) + tit0 += ' placement HxL=(%gx%g)' % ( + float(md['placement_H']), + float(md['placement_L']), + ) + + # Layout: top-count trace, heatmap, synchronicity trace. + gs = fig.add_gridspec(3, 1, height_ratios=[0.15, 0.55, 0.22], hspace=0.36) + + # ..... top plot + ax = fig.add_subplot(gs[0, 0]) + ax.bar(timeV+Tbin*.5, cntAboveMed, width=Tbin, color='teal', align='center', alpha=0.7) + ax.set( + ylabel='neurons > median', + title=f'neurons above median mean rate ({medRateDisp:.2f} Hz)', + ) + ax.set_xlim(xL,xR) + ax.set_xlabel('Time (s)') + ax.tick_params(axis='x', labelbottom=False) + ax.grid() + + # ....... main heatmap + ax = fig.add_subplot(gs[1, 0]) + #cmap='tab20c', 'Oranges' + # - - - Get the colormap and modify it --- + original_cmap = self.plt.get_cmap('tab20c') + custom_cmap = original_cmap.copy() + custom_cmap.set_under('white') + + cax = ax.imshow(rate2D.T, + aspect='auto', + cmap=custom_cmap, # Use our modified colormap + origin='lower', + extent=[xL, xR, 0, nchan], + interpolation='nearest', + vmin=1, # Set the lower bound of the colormap + vmax=rate2D.max()) # Optional: ensure the upper bound is set + + ax.set_ylabel('Neuron index') + ax.set_title(tit0 ) + ax.set_xlabel('Time (s)') + ax.tick_params(axis='x', labelbottom=False) + cbar = fig.colorbar(cax, ax=ax, orientation='horizontal', pad=0.02, label=f'Instantanous freq (Hz) per neuron over Tbin={time_step:.2f} sec', shrink=0.55) + cbar.ax.tick_params(labelsize=10) + ax.grid() + + # ..... synchronicity (yellow) below heatmap + ax = fig.add_subplot(gs[2, 0]) + ax.bar(timeV+Tbin*.5, pop_rate_hz, width=Tbin, color='gold', align='center', alpha=0.8) + ax.axhline(med_sync_hz, color='red', linestyle='--', linewidth=1.0, zorder=3) + y_rng = float(np.max(pop_rate_hz) - np.min(pop_rate_hz)) + y_pad = max(0.02 * y_rng, 0.01 * max(float(np.max(pop_rate_hz)), 1.0)) + ax.text( + xL + 0.01 * (xR - xL), + med_sync_hz + y_pad, + 'median', + va='bottom', + ha='left', + color='red', + fontsize=9, + zorder=4, + ) + ax.set( + ylabel='synchronicity (Hz)', + title=( + f'Network synchronicity from {nchan} neurons, Tbin={Tbin:.2f} sec, ' + f'median rate={med_sync_hz:.2f} Hz' + ), + ) + ax.set_xlim(xL,xR) + ax.tick_params(axis='x', labelbottom=True) + ax.grid() + ax.set_xlabel('Time (s)') diff --git a/causal_net/nonStation_ver4_lagM_offDiag/UtilDalePoisson4.py b/causal_net/nonStation_ver4_lagM_offDiag/UtilDalePoisson4.py new file mode 100644 index 00000000..a49dc15d --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/UtilDalePoisson4.py @@ -0,0 +1,276 @@ +""" +Utility functions for Dale Poisson simulation data processing + +""" + +import numpy as np +import os +import time +from pprint import pprint + +def compute_consecutive_coincidence_rate(Y,time_evol): + """ + Compute the frequency of coincidences for 2 consecutive time bins for 2 different channels. + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + + Returns: + float: conc_rate_all - the overall coincidence rate + """ + num_steps, num_neurons = Y.shape + + # We need at least 2 time steps for consecutive bins + if num_steps < 2: + return 0.0 + + # Get consecutive time slices using vectorized operations + Y_t = Y[:-1, :] # Y[0:T-1, :] - current time bins + Y_t_plus_1 = Y[1:, :] # Y[1:T, :] - next time bins + + # Create boolean masks for non-zero values + mask_t = (Y_t != 0) # Shape: (T-1, N) + mask_t_plus_1 = (Y_t_plus_1 != 0) # Shape: (T-1, N) + + # Use broadcasting to compute all channel pairs at once + # mask_t[:, :, None] has shape (T-1, N, 1) + # mask_t_plus_1[:, None, :] has shape (T-1, 1, N) + # Broadcasting gives shape (T-1, N, N) for all pairs + coincidences = mask_t[:, :, None] & mask_t_plus_1[:, None, :] + + # Remove diagonal (same channel pairs) using boolean indexing + diagonal_mask = np.eye(num_neurons, dtype=bool) + coincidences[:, diagonal_mask] = False + + # Count total coincidences across all time steps + coincidence_count = np.sum(coincidences) + + conc_rate = coincidence_count / time_evol/num_neurons + return conc_rate # Hz, per neuron + +def estimate_rates(Y, dt, tau, max_samples, spect_radius, varTwindow=5, mxNn=5, verb=1): + """ + Evaluates spike statistics and estimates firing rates and Fano factors. + + tau[i] == 0 excitatory, tau[i] == 1 inhibitory; must match Y columns (same order as A_true / B). + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + dt (float): Time bin size in seconds. + tau (np.ndarray): shape (N,), int; 0 = excitatory, 1 = inhibitory. + max_samples (int): Max time samples to use for rate computation. + spect_radius (float): Reported in summary (required). + varTwindow (float): Window length (seconds) for variance / Fano. + mxNn (int): Max neurons per type in detailed printout. + verb (int): Verbose diagnostics. + """ + if verb > 0: + print("\n=== Estimating Rates & Stats ===") + if Y.shape[0] > max_samples: + if verb > 0: + print("Using %d samples (out of %d) for rate computation" % (max_samples, Y.shape[0])) + Y = Y[:max_samples] + + num_steps_sim, Nn_sim = Y.shape + tau = np.asarray(tau).reshape(-1) + if tau.shape[0] != Nn_sim: + raise ValueError("tau length %d != Y.shape[1] %d" % (tau.shape[0], Nn_sim)) + if tau.dtype.kind not in "iu": + raise ValueError("tau must have integer dtype (0=exc, 1=inh)") + tau_i = tau.astype(np.int64, copy=False) + if np.any((tau_i != 0) & (tau_i != 1)): + raise ValueError("tau entries must be 0 (excitatory) or 1 (inhibitory)") + exc_mask = tau_i == 0 + inh_mask = tau_i == 1 + if not np.any(exc_mask) or not np.any(inh_mask): + raise ValueError("tau must contain at least one excitatory and one inhibitory neuron") + num_excite = int(np.sum(exc_mask)) + num_inhib = int(np.sum(inh_mask)) + + time_evol = num_steps_sim * dt + if verb > 0: + print('steps num_steps=%d, time_evol=%.1f sec, Nn=%d (%d Excit, %d Inhib)' % (num_steps_sim, time_evol, Nn_sim, num_excite, num_inhib)) + + # Compute windowed spike counts and statistics over non-overlapping time windows of length varTwindow (sec), per neuron + window_size = max(1, int(round(varTwindow / dt))) + num_windows = num_steps_sim // window_size + if verb > 0: + print('variance computation: num_windows=%d, window_size=%d' % (num_windows, window_size)) + if num_windows <= 1: + raise ValueError("varTwindow=%s is too large for data (num_windows=%d). Need at least 2 windows for variance." % (varTwindow, num_windows)) + Y_trim = Y[:num_windows * window_size] + Y_win = Y_trim.reshape(num_windows, window_size, Nn_sim) + + # Sum spike counts over each window: shape (num_windows, n_neurons) + window_spike_counts = np.sum(Y_win, axis=1) + + # Fano Factor: Var[spike count] / Mean[spike count] over windows + mean_window_spike_counts = np.mean(window_spike_counts, axis=0) + var_window_spike_counts = np.var(window_spike_counts, axis=0) + fano_factor = np.divide(var_window_spike_counts, mean_window_spike_counts, + out=np.zeros_like(var_window_spike_counts), + where=mean_window_spike_counts != 0) + + # Convert window spike counts to rates (Hz) for rate variance + window_rates = window_spike_counts / (window_size * dt) # shape: (num_windows, n_neurons), Hz + single_rates_var = np.var(window_rates, axis=0) # variance of rate across windows, per neuron (Hz^2) + + # Compute raw arrays (full time series) + spike_counts = np.sum(Y, axis=0) + spike_rates = spike_counts / time_evol + mean_counts_per_bin = np.mean(Y, axis=0) + + # Compute consecutive coincidence rate + start_time = time.time() + conc_rate_per_neuron = compute_consecutive_coincidence_rate(Y,time_evol) + elapsed_time = time.time() - start_time + if verb > 0: + print('Coincidence rate %.2g Hz, Y.shape=%s elaT %.3f sec' % (conc_rate_per_neuron, Y.shape, elapsed_time)) + + # Compute all population statistics once + med_rate_all = np.median(spike_rates) + avg_rate_all = float(np.mean(spike_rates)) + std_rate_all = float(np.std(spike_rates)) + avg_fano_all = float(np.mean(fano_factor)) + std_fano_all = float(np.std(fano_factor)) + avg_rate_excit = float(np.mean(spike_rates[exc_mask])) + std_rate_excit = float(np.std(spike_rates[exc_mask])) + avg_fano_excit = float(np.mean(fano_factor[exc_mask])) + std_fano_excit = float(np.std(fano_factor[exc_mask])) + avg_rate_inhib = float(np.mean(spike_rates[inh_mask])) + std_rate_inhib = float(np.std(spike_rates[inh_mask])) + avg_fano_inhib = float(np.mean(fano_factor[inh_mask])) + std_fano_inhib = float(np.std(fano_factor[inh_mask])) + + # Build stats dictionary with computed values + stats_dict = { + 'num_steps': num_steps_sim, + 'time_evol_sec': time_evol, + 'time_step_sec': dt, + 'num_neurons': Nn_sim, + 'num_excitatory': num_excite, + 'num_inhibitory': num_inhib, + 'avg_spike_rate_all': avg_rate_all, + 'std_spike_rate_all': std_rate_all, + 'avg_fano_factor_all': avg_fano_all, + 'std_fano_factor_all': std_fano_all, + 'avg_spike_rate_excit': avg_rate_excit, + 'std_spike_rate_excit': std_rate_excit, + 'avg_fano_factor_excit': avg_fano_excit, + 'std_fano_factor_excit': std_fano_excit, + 'avg_spike_rate_inhib': avg_rate_inhib, + 'std_spike_rate_inhib': std_rate_inhib, + 'avg_fano_factor_inhib': avg_fano_inhib, + 'std_fano_factor_inhib': std_fano_inhib, + 'conc_rate_per_neuron': float(conc_rate_per_neuron), + 'median_spike_rate_all': med_rate_all + } + + # Printing detailed stats for individual neurons + mxE = min(mxNn, num_excite) + mxI = min(mxNn, num_inhib) + + exc_show = np.where(exc_mask)[0][:mxE] + inh_show = np.where(inh_mask)[0][:mxI] + + if verb > 0: + print('\n--- Stats for %d Excitatory Neurons (by index order) ---' % len(exc_show)) + np.set_printoptions(precision=2) + print('Total Spike Counts: %s' % spike_counts[exc_show]) + print('Mean Firing Rate (Hz): %s' % spike_rates[exc_show]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[exc_show])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[exc_show])) + print('Fano Factor (Var/Mean): %s' % fano_factor[exc_show]) + + print('\n--- Stats for %d Inhibitory Neurons (by index order) ---' % len(inh_show)) + np.set_printoptions(precision=2) + print('Total Spike Counts: %s' % spike_counts[inh_show]) + print('Mean Firing Rate (Hz): %s' % spike_rates[inh_show]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[inh_show])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[inh_show])) + print('Fano Factor (Var/Mean): %s' % fano_factor[inh_show]) + + R_tag = ', R=%.3f' % float(spect_radius) + print('\n--- Population Summary Statistics%s ---' % R_tag) + print('All (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (Nn_sim, stats_dict['avg_spike_rate_all'], stats_dict['std_spike_rate_all'], stats_dict['avg_fano_factor_all'], stats_dict['std_fano_factor_all'])) + print('Excit (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_excite, stats_dict['avg_spike_rate_excit'], stats_dict['std_spike_rate_excit'], stats_dict['avg_fano_factor_excit'], stats_dict['std_fano_factor_excit'])) + if num_inhib > 0: + print('Inhib (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_inhib, stats_dict['avg_spike_rate_inhib'], stats_dict['std_spike_rate_inhib'], stats_dict['avg_fano_factor_inhib'], stats_dict['std_fano_factor_inhib'])) + + # Print key summary values using dictionary + summary_keys = ['conc_rate_per_neuron', 'median_spike_rate_all'] + for key in summary_keys: + if key == 'conc_rate_per_neuron': + print('Coincidence rate per neuron %.2g Hz' % stats_dict[key]) + elif key == 'median_spike_rate_all': + print('Median rate %.2f Hz\n' % stats_dict[key]) + + # Part 2: Compute frequency sorting + neur_freq_index = np.argsort(spike_rates) + + rates_dict = { + 'single_rates': spike_rates, + 'neur_freqIdx':neur_freq_index, + 'sigle_rates_var': single_rates_var, + 'single_fano_fact': fano_factor + } + + return stats_dict, rates_dict, neur_freq_index + +def get_offdiag_triplets(A, isPos=True): + """ + Returns positive/negative, off-diagonal elements of 2D array A as array of [i, j, value] triplets. + + Parameters: + A : 2D numpy array + isPos : bool, if True only return positive values (>0), if False return all negative values + + Returns: + Array of shape (n, 3) where n is the number of valid off-diagonal elements + Each row is [row_index, col_index, value] + """ + # Create mask for off-diagonal elements + offdiag_mask = ~np.eye(A.shape[0], A.shape[1], dtype=bool) + + # Add value condition based on isPos + if isPos: + value_mask = A > 0 + else: + value_mask = A < 0 + + # Combine masks + mask = value_mask & offdiag_mask + + i_indices, j_indices = np.where(mask) + values = A[i_indices, j_indices] + + return np.column_stack([i_indices, j_indices, values]) + +def get_offdiag_triplets(A, isPos=True): + """ + Returns positive/negative, off-diagonal elements of 2D array A as array of [i, j, value] triplets. + + Parameters: + A : 2D numpy array + isPos : bool, if True only return positive values (>0), if False return all negative values + + Returns: + Array of shape (n, 3) where n is the number of valid off-diagonal elements + Each row is [row_index, col_index, value] + """ + # Create mask for off-diagonal elements + offdiag_mask = ~np.eye(A.shape[0], A.shape[1], dtype=bool) + + # Add value condition based on isPos + if isPos: + value_mask = A > 0 + else: + value_mask = A < 0 + + # Combine masks + mask = value_mask & offdiag_mask + + i_indices, j_indices = np.where(mask) + values = A[i_indices, j_indices] + + return np.column_stack([i_indices, j_indices, values]) diff --git a/causal_net/nonStation_ver4_lagM_offDiag/Util_pseudospectra.py b/causal_net/nonStation_ver4_lagM_offDiag/Util_pseudospectra.py new file mode 100755 index 00000000..23fc3e1f --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/Util_pseudospectra.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +import numpy as np +import time + +def create_example_matrices(size=100): + A_base = np.diag(np.linspace(-5, -0.5, size)) + np.diag(np.ones(size - 2) * 4, k=2) + sparsity = 0.15 + random_connections = np.random.randn(size, size) + mask = np.random.rand(size, size) > sparsity + random_connections[mask] = 0 + A_true = A_base + random_connections * 0.1 + + noise = np.random.randn(size, size) * 0.10 + noise_mask = np.random.rand(size, size) > sparsity + noise[noise_mask] = 0 + A_fit = A_true + noise + + return A_true, A_fit + +def compute_pseudospectrum(A, npts, minY): + print(f'computing pseudospectrum A:{A.shape} .... ') + eigs = np.linalg.eigvals(A) + + # Calculate bounds + real_min, real_max = np.real(eigs).min(), np.real(eigs).max() + imag_min, imag_max = np.imag(eigs).min(), np.imag(eigs).max() + + real_pad = (real_max - real_min) * 0.2 + imag_pad = (imag_max - imag_min) * 0.2 + + bbox = [ + real_min - real_pad, + min(real_max + real_pad, 1.0), + imag_min, + imag_max + imag_pad + ] + + x_coords = np.linspace(bbox[0], bbox[1], npts) + y_coords = np.linspace(minY, bbox[3], npts) + #y_coords = np.linspace(bbox[2], bbox[3], npts) + X, Y = np.meshgrid(x_coords, y_coords) + Z = X + 1j * Y + + sigma_grid = np.zeros_like(Z, dtype=float) + I = np.eye(A.shape[0]) + + for i in range(npts): + for j in range(npts): + z = Z[i, j] + s = np.linalg.svd(z * I - A, compute_uv=False) + sigma_grid[i, j] = s[-1] + + return X, Y, sigma_grid, eigs + +def plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin): + levels = np.logspace(-2.5, -0.5, 10) # Use 10 contour levels + + # Plot the normal contour lines + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + + # Create a mask for the area greater than epsMin + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) # Masking areas greater than epsMin + + # Fill the area below epsMin + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + # Scatter plot for eigenvalues + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + # Vertical line at Re=0 + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + # Set titles and labels + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + # Additional formatting + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(bottom=minY) # Clip at minY + + +def main(size=100, minY=-0.9, epsMin=0.032): + A_true, A_fit = create_example_matrices(size) + + fig, axes = plt.subplots(1, 2, figsize=(16, 7)) + + matrices_to_plot = [ + ('True Matrix ($A_{true}$)', A_true), + ('Distorted Matrix ($A_{fit}$) ', A_fit) + ] + + print("Starting pseudospectrum calculations...", minY) + start_time = time.time() + + for i, (title, A) in enumerate(matrices_to_plot): + ax = axes[i] + print("Processing '{}'...".format(title)) + + npts = 80 # Grid resolution + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin) # Pass epsMin + + total_time = time.time() - start_time + print(f"Calculation finished in {total_time:.2f} seconds.") + + fig.suptitle('Pseudospectrum Comparison (Clipped Imaginary Part)', fontsize=16) + plt.tight_layout(rect=[0, 0.03, 1, 0.95]) + + outFile = 'out/pseudospectra_comparison.png' + plt.savefig(outFile) + print(f"Plot saved as '{outFile}'") + plt.close(fig) + +# --- Main execution --- +if __name__ == '__main__': + + import matplotlib.pyplot as plt + + main(size=50, minY=-0.5, epsMin=0.024) diff --git a/causal_net/nonStation_ver4_lagM_offDiag/docs/buildFit.sh b/causal_net/nonStation_ver4_lagM_offDiag/docs/buildFit.sh new file mode 100755 index 00000000..7ae78c18 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/docs/buildFit.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +set -e + +doc_name=nonStationary-ver4_oscKer_EMfit + +is_ubuntu=0 +if [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + is_ubuntu=1 +fi + +if [ "$is_ubuntu" -eq 1 ]; then + pdflatex "$doc_name" + pdflatex "$doc_name" + open "${doc_name}.pdf" +else + if ! command -v latex >/dev/null 2>&1; then + module load texlive + fi + latex "$doc_name" + latex "$doc_name" + pdflatex "$doc_name" + xdvi -s 5 "${doc_name}.dvi" +fi diff --git a/causal_net/nonStation_ver4_lagM_offDiag/docs/buildGen.sh b/causal_net/nonStation_ver4_lagM_offDiag/docs/buildGen.sh new file mode 100755 index 00000000..eb0304a0 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/docs/buildGen.sh @@ -0,0 +1,25 @@ +#!/bin/bash + +set -e + +doc_name=nonStationary-ver4_topoDaleGen + +is_ubuntu=0 +if [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + is_ubuntu=1 +fi + +if [ "$is_ubuntu" -eq 1 ]; then + pdflatex "$doc_name" + pdflatex "$doc_name" + open "${doc_name}.pdf" +else + if ! command -v latex >/dev/null 2>&1; then + module load texlive + fi + latex "$doc_name" + latex "$doc_name" + pdflatex "$doc_name" # to produce .pdf for github + + xdvi -s 5 "${doc_name}.dvi" +fi diff --git a/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_oscKer_EMfit.pdf b/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_oscKer_EMfit.pdf new file mode 100644 index 00000000..ef8d1856 Binary files /dev/null and b/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_oscKer_EMfit.pdf differ diff --git a/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_oscKer_EMfit.tex b/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_oscKer_EMfit.tex new file mode 100644 index 00000000..e6545fc5 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_oscKer_EMfit.tex @@ -0,0 +1,163 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} + +\geometry{margin=0.8in} + +\title{Sparse Poisson GLM with Off-Diagonal Memory Kernel:\\Model, Inference, and Training} +\author{Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +Two-block inference scheme for a Poisson GLM with a lagged off-diagonal memory kernel on a non-stationary neuronal population. The model uses a diagonal lag-1 drive plus a shared history kernel on off-diagonal inputs. Alternating optimization yields tractable convex subproblems for the kernel (\(\lambda_2\) smoothness) and network parameters (\(\lambda_3\) off-diagonal sparsity). +\end{abstract} + +\tableofcontents +\newpage + +% =============================================================== +\section{Generative Model} +\label{sec:gen} +% =============================================================== + +Population of \(N\) neurons, \(T\) time bins of width \(\Delta t\). Spike counts \(Y_t \in \mathbb{N}_0^{N}\), \(t=0,\dots,T-1\). + +\paragraph{Parameters.} +Diagonal: \(A^{\mathrm{diag}} \in \mathbb{R}^{N}\) with \((A^{\mathrm{diag}})_i = (A_{\text{true}})_{ii}\). +Off-diagonal: \(A^{\mathrm{off}} \in \mathbb{R}^{N\times N}\), zeros on diagonal. +Kernel: \(\kappa \in \mathbb{R}^{M}\), 0-based indexing, anchor \(\kappa_0=1\), learnable entries \(\kappa_1,\dots,\kappa_{M-1}\). +Bias: \(B \in \mathbb{R}^{N}\). +Clipping threshold: \(\eta_{\mathrm{clip}}=5\). + +\paragraph{Off-diagonal history.} +\[ +S_t \;=\; \sum_{\ell=0}^{M-1} \kappa_\ell \, Y_{t-1-\ell} \;\in\; \mathbb{R}^{N}, +\qquad t-1-\ell \ge 0. +\] + +\paragraph{State equation.} +\[ +\eta_{t,i} \;=\; (A^{\mathrm{diag}})_i \, Y_{t-1,i} +\;+\; \big(A^{\mathrm{off}} S_t\big)_i +\;+\; B_i, +\] +\[ +\mu_{t,i} \;=\; \exp\!\big(\tilde{\eta}_{t,i}\big)\,\Delta t,\quad +\tilde{\eta}_{t,i} \;=\; \min\big(\eta_{t,i},\,\eta_{\mathrm{clip}}\big). +\] + +\paragraph{Observation model.} +\[ +Y_{t,i} \sim \mathrm{Poisson}(\mu_{t,i}), \quad i=1,\dots,N,\; t=0,\dots,T-1. +\] + +% =============================================================== +\section{Objective and Identifiability} +\label{sec:objective} +% =============================================================== + +\paragraph{Log-likelihood.} +Dropping the constant \(\log(Y_{t,i}!)\) terms: +\[ +\mathcal{L} = \sum_{i=1}^N \sum_{t=0}^{T-1} +\left[ Y_{t,i}\,\tilde{\eta}_{t,i} - \exp(\tilde{\eta}_{t,i})\,\Delta t \right]. +\] + +\paragraph{Objective function.} +\begin{equation} +\mathcal{J} +\;=\; +\underbrace{\sum_{t=2}^{T}\sum_{i=1}^{N} +\left[\mu_{t,i}-Y_{t,i}\left(\tilde{\eta}_{t,i}+\log\Delta t\right)\right]}_{\text{Poisson NLL}} +\;+\; +\underbrace{\lambda_2 \sum_{\ell=2}^{M-1} \big(\kappa_\ell-\kappa_{\ell-1}\big)^2}_{\text{kernel smoothness (Block~1)}} +\;+\; +\underbrace{\lambda_3\,\overline{|A_{\mathrm{off}}|}}_{\text{off-diagonal sparsity (Block~2)}}, +\label{eq:objective} +\end{equation} +where +\begin{equation} +\overline{|A_{\mathrm{off}}|} +\;=\; +\frac{1}{N(N-1)} +\sum_{\substack{i,j=1\\i\neq j}}^{N}|A_{ij}^{\mathrm{off}}|. +\label{eq:l1_offdiag} +\end{equation} + +\paragraph{Identifiability.} +The product \(A^{\mathrm{off}}_{ij}\,\kappa_\ell\) has a scaling degeneracy, removed by the anchor \(\kappa_0=1\). +If clipping activates frequently, a censored Poisson likelihood should replace the standard one. + +% =============================================================== +\section{Inference: Two-Block Alternating Scheme} +\label{sec:inference} +% =============================================================== + +The bilinear structure \(A^{\mathrm{off}}_{ij}\,\kappa_\ell\) precludes joint convexity. We use block coordinate descent: each block is convex given the other fixed, guaranteeing a non-increasing objective sequence converging to a stationary point. + +\subsection{Initialization} +Initialize \(A^{\mathrm{diag}}\) and \(B\) from mean firing rates and lag-1 covariances. +Initialize \(\kappa\) with a damped-oscillator profile or simply \(\kappa_0=1\) with small subsequent entries. + +\subsection{Block~1 (kernel update)} +Fix \((A^{\mathrm{diag}}, A^{\mathrm{off}}, B)\). Update \(\kappa_1,\dots,\kappa_{M-1}\) by minimizing \(\mathcal{J}\) with the \(\lambda_2\) smoothness penalty. The problem is convex in \(\kappa\) since \(S_t\) is linear in \(\kappa\). + +\subsection{Block~2 (network update)} +Fix \(\kappa\). Update \((A^{\mathrm{diag}}, A^{\mathrm{off}}, B)\) by minimizing the Poisson NLL plus the \(\lambda_3\) L1 penalty~\eqref{eq:l1_offdiag}. Dale's law sign constraints on columns of \(A^{\mathrm{off}}\) can be added as linear inequalities. + +% =============================================================== +\section{Training Samples} +\label{sec:training} +% =============================================================== + +Input: spike counts \(Y_t \in \mathbb{N}_0^{N}\), \(t=0,\dots,T-1\). +Let \(M_{\mathrm{cut}} \ge 2\) be the history cutoff. + +\paragraph{Sample construction.} +Each sample at output time \(t_{\mathrm{out}} \in [M_{\mathrm{cut}},\, T-1]\): +\[ +X^{(t_{\mathrm{out}})} = \big\{Y_{t_{\mathrm{out}}-1},\, Y_{t_{\mathrm{out}}-2},\, \dots,\, Y_{t_{\mathrm{out}}-M_{\mathrm{cut}}}\big\} \in \mathbb{N}_0^{N \times M_{\mathrm{cut}}}, +\qquad +y^{(t_{\mathrm{out}})} = Y_{t_{\mathrm{out}}} \in \mathbb{N}_0^{N}. +\] + +\paragraph{Sampling strategy.} +Draw \(t_{\mathrm{out}}\) densely from \([M_{\mathrm{cut}},\, T-1]\) with stride \(s\); optionally enforce a minimum separation to reduce autocorrelation. Total samples: +\[ +N_{\mathrm{samples}} \approx \left\lfloor \frac{T - M_{\mathrm{cut}}}{s} \right\rfloor. +\] +For validation, sample windows within random non-overlapping blocks of length \(L \gg M_{\mathrm{cut}}\) to obtain a temporally separated hold-out set. + +% =============================================================== +\section{Kernel Normalization} +\label{sec:lags} +% =============================================================== + +We use 0-based lag indexing \(\ell \in \{0,\dots,M-1\}\) with anchor \(\kappa_0=1\). An alternative \(\|\kappa\|_1=1\) constraint was tested but significantly reduced non-stationarity in simulations, so the ground truth \(\kappa_{\mathrm{true}}\) does not obey it. + +The \(\lambda_2\) smoothness penalty imposes no parametric form on \(\kappa\). If prior knowledge constrains the kernel shape, one can reparametrize \(\kappa\) with fewer parameters while keeping \(\kappa_0=1\). + +% =============================================================== +\section{Implementation Notes} +\label{sec:practical} +% =============================================================== + +\paragraph{Sparsity and sign constraints.} +The \(\lambda_3\) penalty~\eqref{eq:l1_offdiag} promotes sparsity; normalization by \(N(N-1)\) makes the penalty strength network-size invariant. Column sign constraints (Dale's law) can be imposed if neuron types are known. + +\paragraph{Regularization tuning.} +Start \(\lambda_2\) small; increase if the kernel estimate oscillates. Tune \(\lambda_3\) to match expected connectivity density (e.g., 5--20\%). + + + +\end{document} diff --git a/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_topoDaleGen.pdf b/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_topoDaleGen.pdf new file mode 100644 index 00000000..e42d127d Binary files /dev/null and b/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_topoDaleGen.pdf differ diff --git a/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_topoDaleGen.tex b/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_topoDaleGen.tex new file mode 100644 index 00000000..9dee8596 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/docs/nonStationary-ver4_topoDaleGen.tex @@ -0,0 +1,719 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} +\usepackage{algorithm} +\usepackage{algorithmic} + + +\geometry{margin=0.8in} + +\title{Synthetic Dale Network Topology and Spike-Count Generation:\\ +Distance-Dependent Sparse Connectivity with a Multivariate Poisson GLM for Count Time Series} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +We describe the pipeline implemented in \texttt{gen\_daleMatrices4.py}: +(1)~build a spatially embedded, Dale-law-compliant sparse network +$A_{\text{true}}$ with target spectral radius $R<1$ and a bias vector $B$; +(2)~simulate spike counts with a multivariate Poisson GLM for count time series, +either +\textbf{Model~A} (lag-1, no off-diagonal memory) or +\textbf{Model~B} (lag-$M$: diagonal lag-1 drive plus an off-diagonal +history filtered by a damped-oscillator kernel $\kappa_\ell$). +A third variant (Model~C: additional diagonal adaptation kernel) is +discussed separately but is \emph{not} implemented in this script. +\end{abstract} + +\tableofcontents +\newpage + +% =============================================================== +\section{Overview of the Method} +\label{sec:overview} +% =============================================================== + +The objective is to produce synthetic spike-count time series +$Y_t\in\mathbb{Z}_{\geq 0}^{N}$ ($t=1,\dots,T$) with configurable +spatial connectivity and dynamics given by a multivariate Poisson GLM +for count time series. + +\paragraph{Stage 1 --- Network construction (Section~\ref{sec:generation_AB}).} +$N$ neurons are sampled from a regular grid on $[0,L]\times[0,H]$ and +connected by distance-dependent out-degree sampling with Dale's law. +Outputs: $A_{\text{true}}$ with $\rho(A_{\text{true}})=R$ and bias $B$. + +\paragraph{Stage 2 --- Spike generation (Section~\ref{sec:spike_gen}).} +Spike counts are generated by a multivariate Poisson GLM for count time +series with autoregressive lag-$M$ history. +\textbf{Model~A:} lag-1 recurrence on all couplings. +\textbf{Model~B:} diagonal term uses $Y_{t-1}$ only; off-diagonal +block uses lags $1,\dots,M$ with kernel $\kappa_\ell$ +(Section~\ref{sec:lagM_single}). The number of lags $M$ is computed from +the oscillation period (not a separate CLI flag). +Model~C (Section~\ref{sec:lagM_two}) is not implemented in code. + +\paragraph{Notation conventions.} +$A_{\text{true}}$ denotes the connectivity matrix throughout.\\ +$A_{\text{true}}^{\mathrm{diag}}=\operatorname{diag} +((A_{\text{true}})_{11},\dots,(A_{\text{true}})_{NN})$ is its diagonal +and $A_{\text{true}}^{\mathrm{off}}=A_{\text{true}}- +A_{\text{true}}^{\mathrm{diag}}$ its off-diagonal part. +In Model~B (Section~\ref{sec:lagM_single}), $A^{\mathrm{diag}}\in\mathbb{R}^N$ +denotes the \emph{vector} of diagonal entries $(A_{\text{true}})_{ii}$, +parallel to the matrix shorthand $A^{\mathrm{off}}$. +$Y_t\in\mathbb{Z}_{\geq 0}^{N}$ is the vector of spike counts in +time bin $t$, where each bin has width $\Delta t$ seconds. +$B\in\mathbb{R}^{N}$ is the bias (log-baseline-rate) vector. +$\eta_{\mathrm{clip}}=5$ is a saturation threshold that prevents +numerical overflow in the exponential link function. +Column convention: $({A_{\text{true}}})_{ij}$ is the weight from +presynaptic neuron $j$ to postsynaptic neuron $i$. + + +% =============================================================== +\section{Generation of $A_{\text{true}}$ and $B$} +\label{sec:generation_AB} +% =============================================================== + +\subsection{Parameters} + +The script takes (among others): $N$ and $N_{\text{exc}}$ (default +$0.8\,N$); rectangular domain $[0,L]\times[0,H]$; grid spacing +$d_{\text{min}}>0$ (\texttt{--placement\_min\_dist}); distance-kernel +exponent $\delta_{\mathrm{ker}}\in\{1,2\}$ (third entry of +\texttt{--placement\_H\_L\_delta}); out-degree range from two fractions +$p_{\text{lo}}0$. +Rescaling preserves all signs and the zero pattern, so Dale's law, +E-I ratio, and negative self-loops are all maintained. +For $R<1$, contractivity of $A_{\text{true}}$ is a sufficient +condition for stationarity of the Poisson GLM. + +\subsubsection{Phase 5: Bias Vector Generation} + +Let $f_{\mathrm{idle}}\in[f_{\min},f_{\max}]$ be a \emph{target idle +firing-rate band} in Hz (argument \texttt{idleRate}, default +$[15,30]$). Independent draws +\begin{equation} + B_i \sim \mathrm{Uniform}\bigl(\ln(R f_{\min}),\,\ln(R f_{\max})\bigr), + \label{eq:bias} +\end{equation} +then all excitatory indices receive a fixed downward shift +$B_i\leftarrow B_i-2$ (hard-coded in the script). Inhibitory +components are unchanged. This matches \texttt{set\_flat\_selfSpiking} +in code. + +\subsection{Output Summary} + +\begin{table}[h] +\centering +\begin{tabular}{llll} +\toprule +Array & Shape & Type & Description \\ +\midrule +$\mathbf{P}$ & $N\times 2$ & float + & Neuron positions $(x_i,y_i)$, sorted by $(x,y)$ \\ +$\mathbf{D}$ & $N\times N$ & float + & Pairwise distances; diagonal $0$ in output \\ +$\tau$ & $N$ & int + & Type labels: $0$=exc, $1$=inh \\ +$\sigma$ & $N$ & $\pm1$ + & Synaptic signs \\ +$\mathbf{E}$ & $N\times N$ & $\{-1,0,+1\}$ + & Signed topology (Eq.~\ref{eq:topology}) \\ +$k$ & $N$ & int + & Out-degree, $k_j\in[k_{\text{min}},k_{\text{max}}]$ \\ +$A_{\text{true}}$ & $N\times N$ & float + & Weighted connectivity, $\rho(A_{\text{true}})=R$ \\ +$B$ & $N$ & float + & Bias vector (Eq.~\ref{eq:bias}) \\ +\bottomrule +\end{tabular} +\caption{Output arrays of Stage~1. +$A_{\text{true}}$ and $B$ are passed to the spike generator +(Section~\ref{sec:spike_gen}).} +\label{tab:output} +\end{table} + +% =============================================================== +\section{Generation of Spike Trains} +\label{sec:spike_gen} +% =============================================================== + +Given $A_{\text{true}}$, $B$, time-bin width $\Delta t$, horizon +$T$, and $\eta_{\mathrm{clip}}=5$ (hard-coded), Models~A and~B are +implemented as a multivariate Poisson GLM for count time series with +autoregressive lag-$M$ history. Bin $t$ indexes rows $Y_t$ (code uses +$t=0,\dots,T-1$). +\textbf{Model~A:} $Y_0\sim\mathrm{Poisson}(\exp(B)\Delta t)$, then +Eq.~\eqref{eq:lag1} for $t\ge 1$. +\textbf{Model~B:} the first $M$ rows are set to zero (pre-history), +with $M=\texttt{mem\_lag\_steps}$ computed from $\tau_{\mathrm{mem}}$ +(Section~\ref{sec:lagM_single}); dynamics follow that section for $t\ge 1$. + +% --------------------------------------------------------------- +\subsection{Model A: Lag-1, No Memory, Not Bursting} +\label{sec:lag1} +% --------------------------------------------------------------- + +The simplest model conditions only on the immediately preceding bin: +\begin{equation} + \eta_t = A_{\text{true}}\,Y_{t-1}+B, + \qquad + \tilde{\eta}_t = \min(\eta_t,\;\eta_{\mathrm{clip}}), + \qquad + \mu_t = \exp(\tilde{\eta}_t)\,\Delta t, + \qquad + Y_t \sim \mathrm{Poisson}(\mu_t). + \label{eq:lag1} +\end{equation} +Here $\min(\cdot)$ and $\exp(\cdot)$ act element-wise. +Because $\rho(A_{\text{true}})<1$, the process admits a stationary +distribution. + +% --------------------------------------------------------------- +\subsection{Model B: Multivariate Poisson GLM with Autoregressive Lag-$M$ History} +\label{sec:lagM_single} +% --------------------------------------------------------------- + +\subsubsection{Off-diagonal kernel $\kappa_\ell$ (damped oscillator)} + + +Let $\tau_{\mathrm{mem}}>0$ be the oscillation period in \textbf{seconds} +(required; the code asserts $\tau_{\mathrm{mem}}>0$). +Superscript ${}^{(s)}$ marks quantities expressed in \textbf{discrete steps} +(time bins). +Define $\tau^{(s)}=\mathrm{round}(\tau_{\mathrm{mem}}/\Delta t)\ge 1$, the oscillation period in bins. +Let $Q>0$ (quality factor); the CLI passes \texttt{mem\_Q} and \texttt{mem\_tau} +(\texttt{--time\_kernel\_q\_tau}). +Set $\omega=2\pi/\tau^{(s)}$ and $\gamma=\pi/(Q\,\tau^{(s)})$. +The off-diagonal history uses only lags $\ell=1,\dots,M$ with +\begin{equation} + M = \max\!\left(1,\;\bigl\lfloor 3\tau^{(s)}/4\bigr\rfloor\right) +\end{equation} +This choice sets the \emph{truncation depth} $M\approx \tfrac{3}{4}\tau^{(s)}$, +so $\omega M\approx \tfrac{3\pi}{2}$ (about $1.5\pi$ radians of oscillation phase +across the retained lags). +For $\ell=1,\dots,M$, define raw weights +\begin{equation} + \tilde{\kappa}_\ell = \exp(-\gamma \ell)\,\cos(\omega \ell). + \label{eq:kappa_offdiag_gen4} +\end{equation} +There is \emph{no} $\ell=0$ term. +The implementation then fixes the scale by normalizing to the first lag: +$\kappa_\ell = \tilde{\kappa}_\ell / \tilde{\kappa}_1=1.0$ + + +\subsubsection{Evolution} + +Write $A^{\mathrm{diag}}\in\mathbb{R}^N$ with +$(A^{\mathrm{diag}})_i=(A_{\text{true}})_{ii}$, and let +$A^{\mathrm{off}}=A_{\text{true}}$ with zero diagonal. +The off-diagonal spike-history vector is +$S_t=\sum_{\ell=1}^{M}\kappa_\ell\,Y_{t-\ell}$ (vector sum). +For $t\ge 1$, +\begin{equation} + \eta_t + = A^{\mathrm{diag}}\odot Y_{t-1} + A^{\mathrm{off}} S_t + B, + \label{eq:modelB_drive} +\end{equation} +then $\tilde\eta_t=\min(\eta_t,\eta_{\mathrm{clip}})$, +$\mu_t=\exp(\tilde\eta_t)\,\Delta t$, and +$Y_t\sim\mathrm{Poisson}(\mu_t)$ element-wise as in +Eq.~\eqref{eq:lag1}. +The first $M$ rows $Y_0,\dots,Y_{M-1}$ are held at~$0$ (pre-history). + +\paragraph{Relation to Model~A.} +Model~A is a separate code path (\texttt{--spike\_model A}); it is +\textbf{not} recovered by formally setting $M=1$ in Model~B. + +\subsubsection{Stability} + +\texttt{gen\_daleMatrices4.py} does not impose an extra analytic +stability test for Model~B beyond rescaling $A_{\text{true}}$ to +spectral radius $R<1$ (Eq.~\ref{eq:spectral}). A sufficient +condition for a two-kernel extension appears in +Section~\ref{sec:lagM_two} (Model~C). + + +% --------------------------------------------------------------- +\subsection{Model C: Lag-$M$, Two Kernels, Bursting - not implemented} +\label{sec:lagM_two} +% --------------------------------------------------------------- + +\subsubsection{Adaptation kernel for the diagonal} + +Spike-frequency adaptation is modeled as a purely inhibitory, +exponentially decaying self-history kernel with time constant +$\tau_{\mathrm{adapt}}$ (in time bins) and amplitude +$\alpha_{\mathrm{adapt}}>0$. +Let $M_{\mathrm{d}}$ denote the number of lags retained for the +diagonal kernel (which may differ from $M$). +The unnormalized weights are: +\begin{equation} + \tilde{h}^{\mathrm{diag}}_\ell + = \exp\!\left(-\frac{\ell}{\tau_{\mathrm{adapt}}}\right), + \qquad \ell = 1,\dots,M_{\mathrm{d}}. + \label{eq:adapt_kernel_raw} +\end{equation} +Normalizing and attaching the inhibitory sign: +\begin{equation} + h^{\mathrm{diag}}_\ell + = -\,\alpha_{\mathrm{adapt}}\; + \frac{\tilde{h}^{\mathrm{diag}}_\ell} + {\sum_{k=1}^{M_{\mathrm{d}}} + \tilde{h}^{\mathrm{diag}}_k}, + \qquad \ell = 1,\dots,M_{\mathrm{d}}. + \label{eq:adapt_kernel} +\end{equation} +All weights are negative, ensuring that a neuron's own past +activity is always suppressive. +The dimensionless amplitude $\alpha_{\mathrm{adapt}}$ controls the +overall strength of adaptation relative to the diagonal entries of +$A_{\text{true}}$. + +\subsubsection{Evolution equations} + +The history drive now has two kernel-filtered components: +\begin{equation} + H_t + = A_{\text{true}}^{\mathrm{diag}} + \!\left(\sum_{\ell=1}^{M_{\mathrm{d}}} + h^{\mathrm{diag}}_\ell\,Y_{t-\ell}\right) + + A_{\text{true}}^{\mathrm{off}} + \!\left(\sum_{\ell=1}^{M} + h^{\mathrm{off}}_\ell\,Y_{t-\ell}\right). + \label{eq:histC} +\end{equation} +Note that, unlike Model~B, the diagonal term now integrates over +$M_{\mathrm{d}}$ lags rather than using only $Y_{t-1}$. +Spike generation proceeds identically: +\begin{equation} + \eta_t = H_t + B, + \qquad + \tilde{\eta}_t = \min(\eta_t,\;\eta_{\mathrm{clip}}), + \qquad + \mu_t = \exp(\tilde{\eta}_t)\,\Delta t, + \qquad + Y_t \sim \mathrm{Poisson}(\mu_t). + \label{eq:lagM_gen_C} +\end{equation} + +\subsubsection{Reduction to Model B} + +Setting $M_{\mathrm{d}}=1$ and $\alpha_{\mathrm{adapt}}=1$ gives +$h^{\mathrm{diag}}_1=-1$, so the diagonal term becomes +$-A_{\text{true}}^{\mathrm{diag}}Y_{t-1}$. +Since $(A_{\text{true}})_{ii}<0$ by construction +(negative self-loops), the product +$-A_{\text{true}}^{\mathrm{diag}}Y_{t-1} + =|A_{\text{true}}^{\mathrm{diag}}|Y_{t-1}$ +has the wrong sign. +To recover Model~B exactly, one must instead set +$\alpha_{\mathrm{adapt}}=-1$ (i.e.\ disable the adaptation sign +flip) or, equivalently, simply bypass the diagonal kernel and use +$A_{\text{true}}^{\mathrm{diag}}Y_{t-1}$ directly. +In practice, Model~C is not intended to reduce to Model~B; rather, +Model~B is the appropriate choice when adaptation is not modeled. + +\subsubsection{Stability} + +A sufficient condition for stability of Model~C is: +\begin{equation} + \rho\!\bigl(A_{\text{true}}^{\mathrm{off}}\bigr) + \cdot\sum_{\ell=1}^{M}|h^{\mathrm{off}}_\ell| + \;+\; + \max_i\bigl|(A_{\text{true}})_{ii}\bigr| + \cdot\sum_{\ell=1}^{M_{\mathrm{d}}} + |h^{\mathrm{diag}}_\ell| + \;<\; 1. + \label{eq:stabilityC} +\end{equation} +The first term equals +$\rho(A_{\text{true}}^{\mathrm{off}})$ by the off-diagonal kernel +normalization. +The second term equals +$\alpha_{\mathrm{adapt}}\max_i|(A_{\text{true}})_{ii}|$ by the +diagonal kernel normalization. +Both contributions must be kept small enough that their sum remains +below unity. + +% --------------------------------------------------------------- +\subsection{Parameter Choices} +\label{sec:params} +% --------------------------------------------------------------- + +The script defaults include $\Delta t=0.01$\,s and +\texttt{--time\_kernel\_q\_tau} $=(Q,\tau_{\mathrm{mem}})=(3,\,0.4)$ +in code (two floats: \texttt{mem\_Q}, \texttt{mem\_tau} in seconds). +The lag count $M=\texttt{mem\_lag\_steps}$ equals the length of +\texttt{\_offdiag\_time\_kernel} (computed from $\tau^{(s)}$ as above), not passed +separately. For typical simulations choose $T=\texttt{num\_steps}$ well above +$M$ so the trajectory spends most bins outside the initial zero pre-history. + +Model~C-specific parameter guidance remains in +Section~\ref{sec:lagM_two}. + +\newpage +% --------------------------------------------------------------- +\subsection{Notation Summary} +\label{sec:notation} +% --------------------------------------------------------------- + +For reference, all symbols used in the spike generation models are +collected here. + +\begin{table}[h] +\centering +\begingroup +\renewcommand{\arraystretch}{1.22} +\begin{tabular}{lll} +\toprule +Symbol & Domain & Description \\ +\midrule +$N$ & $\mathbb{Z}_{>0}$ & Number of neurons \\ +$T$ & $\mathbb{Z}_{>0}$ & Number of time bins \\ +$\Delta t$ & $\mathbb{R}_{>0}$ & + Time-bin width (seconds); default $0.01$\,s \\ +$Y_t$ & $\mathbb{Z}_{\geq 0}^{N}$ & + Spike-count vector in bin $t$ \\ +$A_{\text{true}}$ & $\mathbb{R}^{N\times N}$ & + Connectivity matrix from Section~\ref{sec:generation_AB} \\ +$A_{\text{true}}^{\mathrm{diag}}$ & $\mathbb{R}^{N\times N}$ & + Diagonal of $A_{\text{true}}$ (self-coupling) \\ +$A_{\text{true}}^{\mathrm{off}}$ & $\mathbb{R}^{N\times N}$ & + Off-diagonal of $A_{\text{true}}$ (cross-neuron coupling) \\ +$A^{\mathrm{diag}}$ & $\mathbb{R}^{N}$ & + Vector $(A_{\text{true}})_{ii}$; Model~B shorthand (Eq.~\ref{eq:modelB_drive}) \\ +$B$ & $\mathbb{R}^{N}$ & + Bias vector (log-baseline rate) \\ +$\eta_{\mathrm{clip}}$ & $\mathbb{R}_{>0}$ & + Clipping threshold; default $5$ \\ +\midrule +\multicolumn{3}{l}{\textit{Off-diagonal kernel (Model~B; Model~C off-part)}} \\ +$\tau_{\mathrm{mem}}$ & $\mathbb{R}_{>0}$ & + Oscillation period (\textbf{seconds}); stored as \texttt{mem\_tau}; $\tau^{(s)}=\mathrm{round}(\tau_{\mathrm{mem}}/\Delta t)$ \\ +$Q$ & $\mathbb{R}_{>0}$ & + Quality factor; stored as \texttt{mem\_Q} \\ +$\omega$ & $\mathbb{R}_{>0}$ & + $2\pi/\tau^{(s)}$ \\ +$\gamma$ & $\mathbb{R}_{>0}$ & + $\pi/(Q\,\tau^{(s)})$ \\ +$M$ & $\mathbb{Z}_{>0}$ & + \texttt{mem\_lag\_steps}$=\max(1,\lfloor 3\tau^{(s)}/4\rfloor)$; history lags $\ell=1,\dots,M$ \\ +$\kappa_\ell$ & $\mathbb{R}$ & + Eq.~\ref{eq:kappa_offdiag_gen4}, $\kappa_\ell=\tilde{\kappa}_\ell/\tilde{\kappa}_1$ \\ +\midrule +\multicolumn{3}{l}{\textit{Diagonal adaptation kernel + (Model~C only)}} \\ +$\tau_{\mathrm{adapt}}$ & $\mathbb{R}_{>0}$ & + Adaptation time constant (time bins); default $10$ \\ +$\alpha_{\mathrm{adapt}}$ & $\mathbb{R}_{>0}$ & + Adaptation amplitude; default $0.5$ \\ +$M_{\mathrm{d}}$ & $\mathbb{Z}_{>0}$ & + Number of diagonal history lags \\ +$h^{\mathrm{diag}}_\ell$ & $\mathbb{R}_{\leq 0}$ & + Diagonal kernel weight at lag $\ell$ + (Eq.~\ref{eq:adapt_kernel}) \\ +\bottomrule +\end{tabular} +\endgroup +\caption{Complete notation for the spike generation models + (Section~\ref{sec:spike_gen}). + All symbols from Section~\ref{sec:generation_AB} + ($\mathbf{P}$, $\mathbf{E}$, $\sigma$, $\tau_j$, $r$, $R$, + $\delta_{\mathrm{ker}}$, etc.)\ retain their definitions + from that section.} +\label{tab:notation} +\end{table} + +% = = = = = = = = = ==== +% = = = = = = = = = ==== +\appendix +\section*{\Huge Appendix} +\addcontentsline{toc}{section}{Appendix} +\vspace{1cm} + +\subsubsection{Excitable memory kernel} + +Experimental organoid recordings show brief synchronous bursts +separated by long, irregular quiescent intervals. This pattern is +characteristic of an excitable system rather than an oscillator: +recent activity first promotes further firing (burst build-up) +and then strongly suppresses it (post-burst refractoriness). +We model this with a difference-of-exponentials kernel: +\begin{equation} + h(\tau) = e^{-\tau/\tau_{\mathrm{exc}}} + - \beta\,e^{-\tau/\tau_{\mathrm{inh}}}, + \qquad \tau>0, +\end{equation} +where $\tau_{\mathrm{exc}}\ll\tau_{\mathrm{inh}}$ and +$\beta\in(0,1)$. +The fast excitatory lobe ($\tau_{\mathrm{exc}}$, a few tens of +milliseconds) creates the burst; the slow inhibitory lobe +($\tau_{\mathrm{inh}}$, several hundred milliseconds) enforces a +prolonged refractory period. +The kernel crosses zero at +$\tau^{*}=\frac{\tau_{\mathrm{exc}}\,\tau_{\mathrm{inh}}} + {\tau_{\mathrm{inh}}-\tau_{\mathrm{exc}}} + \ln(1/\beta)$. +Because the inhibitory tail decays slowly, the network remains +quiescent until the suppression wears off and a noise fluctuation +triggers the next burst, producing irregular inter-burst intervals +without a fixed oscillation period. + +Sampling at integer lags and normalizing gives the discrete +weights: +\begin{equation} + \tilde{h}^{\mathrm{off}}_\ell + = \exp(-\ell/\tau_{\mathrm{exc}}) + - \beta\,\exp(-\ell/\tau_{\mathrm{inh}}), + \qquad + h^{\mathrm{off}}_\ell + = \frac{\tilde{h}^{\mathrm{off}}_\ell} + {\sum_{k=1}^{M}|\tilde{h}^{\mathrm{off}}_k|}, + \qquad \ell=1,\dots,M. + \label{eq:mem_kernel} +\end{equation} +A kernel support of $M\geq 3\tau_{\mathrm{inh}}$ suffices to +capture the exponential tail. + + +\subsubsection{Asymmetric damped-oscillator memory kernel} + +Experimental organoid recordings show brief synchronous bursts +separated by long, irregular quiescent intervals. The standard +damped-cosine kernel $h(\tau)=e^{-\gamma\tau}\cos(\omega\tau)$ +produces symmetric positive and negative lobes, so the excitatory +and suppressive phases of each cycle have equal duration. To +obtain short bursts followed by prolonged silences we introduce a +smooth, single-parameter frequency modulation that compresses the +positive (excitatory) lobe and stretches the negative +(suppressive) lobe. + +Define the phase-modulated kernel +\begin{equation} + h(\tau) = e^{-\gamma\tau} + \cos\!\bigl(\omega\tau + \delta\sin(\omega\tau)\bigr), + \qquad \tau>0, + \label{eq:asym_kernel} +\end{equation} +with +\begin{equation} + \omega = \frac{2\pi}{\tau_{\mathrm{mem}}}, + \qquad + \gamma = \frac{\pi}{Q\,\tau_{\mathrm{mem}}}, + \label{eq:asym_params} +\end{equation} +exactly as in the symmetric case. The new dimensionless parameter +$\delta\in[0,1)$ controls the asymmetry between the two lobes. + +The instantaneous frequency of the argument is +$\dot\Phi(\tau)=\omega\bigl(1+\delta\cos(\omega\tau)\bigr)$. +Near the positive peak of the cosine +($\cos(\omega\tau)\approx+1$) the phase advances at rate +$\omega(1+\delta)$, compressing the positive lobe; near the +negative trough ($\cos(\omega\tau)\approx-1$) it advances at rate +$\omega(1-\delta)$, stretching the negative lobe. The ratio of +time spent in the negative versus positive half-cycle is +approximately $(1+\delta)/(1-\delta)$. Setting $\delta=0$ +recovers the original symmetric kernel; $\delta=0.6$ makes the +negative lobe roughly four times longer than the positive one. + +Sampling at integer lags and normalizing by the sum of absolute +values gives the discrete weights: +\begin{equation} + \tilde{h}^{\mathrm{off}}_\ell + = \exp(-\gamma\,\ell)\, + \cos\!\bigl(\omega\,\ell+\delta\sin(\omega\,\ell)\bigr), + \qquad + h^{\mathrm{off}}_\ell + = \frac{\tilde{h}^{\mathrm{off}}_\ell} + {\sum_{k=1}^{M}|\tilde{h}^{\mathrm{off}}_k|}, + \qquad \ell=1,\dots,M. + \label{eq:asym_mem_kernel} +\end{equation} +A kernel support of $M\geq 2\tau_{\mathrm{mem}}$ suffices to +capture the full tail. Typical parameter choices are +$\tau_{\mathrm{mem}}=40$--$60$ bins, $Q=2$--$3$, and +$\delta=0.5$--$0.7$. +\subsubsection{Evolution equations} + +Decompose $A_{\text{true}}$ into diagonal and off-diagonal parts: +\begin{equation} + A_{\text{true}}^{\mathrm{diag}} + = \operatorname{diag}\!\bigl((A_{\text{true}})_{11},\dots, + (A_{\text{true}})_{NN}\bigr), + \qquad + A_{\text{true}}^{\mathrm{off}} + = A_{\text{true}} - A_{\text{true}}^{\mathrm{diag}}. +\end{equation} +The history drive depending on the past $M$ time bins is: +\begin{equation} + H_t + = A_{\text{true}}^{\mathrm{diag}}\,Y_{t-1} + + A_{\text{true}}^{\mathrm{off}} + \!\left(\sum_{\ell=1}^{M} + h^{\mathrm{off}}_\ell\,Y_{t-\ell}\right). + \label{eq:histB} +\end{equation} +The diagonal term uses only the most recent time bin (lag-1 self-coupling), +while the off-diagonal term integrates kernel-weighted spike history +over $M$ lags. +Spikes are then generated as in Model~A: +\begin{equation} + \eta_t = H_t + B, + \qquad + \tilde{\eta}_t = \min(\eta_t,\;\eta_{\mathrm{clip}}), + \qquad + \mu_t = \exp(\tilde{\eta}_t)\,\Delta t, + \qquad + Y_t \sim \mathrm{Poisson}(\mu_t). + \label{eq:lagM_gen} +\end{equation} + +\end{document} diff --git a/causal_net/nonStation_ver4_lagM_offDiag/gen_daleMatrices4.py b/causal_net/nonStation_ver4_lagM_offDiag/gen_daleMatrices4.py new file mode 100755 index 00000000..5f485080 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/gen_daleMatrices4.py @@ -0,0 +1,738 @@ +#!/usr/bin/env python3 +""" + ./gen_daleMatrices4.py --num_neurons 100 --num_excite 80 --placement_H_L_delta 1 2 2 --placement_min_dist 0.01 --dataName test1 + +Primary purpose: generate the ground-truth dictionary (A_true, B_true) for use +by gen_nonStationarySpikes3.py. The stationary spike generation performed here +is for evaluation only (firing-rate sanity check per B-vector). + +Topology (writeup): grid placement on [0,L]×[0,H], distance-kernel affinities, +fixed random out-degree per neuron, Dale's law (column sign = presynaptic type), +negative self-loops, E–I weight scaling, spectral radius normalization. + +CLI placement: --placement_H_L_delta (H, L, placement_ker_delta), --placement_min_dist (grid spacing); out-degree +fraction range: --edge_prob lo hi (maps to k_min, k_max). + +Pipeline: +1. Build rectangular grid (spacing placement_min_dist); sample N nodes without replacement; + sort positions by x. Assign N_exc excitatory / rest inhibitory at random. +2. Affinity V_ij = D_ij^{-δ} with placement_ker_delta from --placement_H_L_delta; sample k_j targets per column j without + replacement with probabilities proportional to V_ij. +3. Signed topology E (diagonal −1); raw weights W from Uniform(1−v,1+v) with + inhibitory scale −r U; rescale so rho(A_true) = R. + +Neuron index order follows placement (sorted by x, then y). Type labels are in +tau (0=exc, 1=inh); B vectors and rate summaries use tau, not index blocks. + +Output files saved to /truthDale/: + .simTruth.npz — arrays below + metadata (dale_conf, evol_conf, …) + .spikes.npz — spikes + rate statistics (+ spike-side metadata) + +simTruth.npz arrays (N = num_neurons, T = num_steps): + A_off_true (N, N) float — off-diagonal weights (diag forced 0) + A_diag_true (N,) float — diagonal of A_true + B_true (N,) float — bias / log-baseline vector + E_true (N, N) int8 — signed Dale topology {−1,0,+1} + node_positions (N, 2) float — (x,y) placement, sorted by x then y + node_is_inhibitory (N,) int32 — 0 = excitatory, 1 = inhibitory + node_distance_matrix (N, N) float — pairwise Euclidean distances; diagonal 0 + offdiag_kernel (M,) float — κ_ℓ memory kernel; M=0 for model A (lag-1), M>0 for B + +simTruth metadata dale_conf (among others): num_neurons, num_excite, spectral_radius, + placement_H, placement_L, placement_min_dist, placement_ker_delta, edge_prob, + spike_model, weight_var, idleRate; mem_Q, mem_tau, mem_tau_steps, mem_lag_steps (zeros for A; + for B, mem_lag_steps = len(offdiag_kernel)). + +spikes.npz arrays: + spikes (T, N) uint8 — spike counts per bin + single_rates, sigle_rates_var, single_fano_fact — per-neuron summaries + +spikes.npz metadata repeats spike_model, placement_*, mem_lag_steps (0 for A, same as dale_conf for B). +evol_conf includes mem_lag_steps mirroring dale_conf. + +Spike GLM (--spike_model): A = lag-1 (baseline Y[0] Poisson); B = lag-M +off-diagonal kernel (--time_kernel_q_tau); mem_tau in seconds; mem_tau_steps = round(mem_tau/dt); +mem_lag_steps = max(1, (3·mem_tau_steps)//4) so active lags reach ~1.5π phase. +Require mem_tau>0 (seconds). First mem_lag_steps bins zero for B. +""" + +import numpy as np +import time +import hashlib +import os +import sys + +import argparse +from pprint import pprint + +from toolbox.Util_NumpyIO import write_data_npz +from UtilDalePoisson4 import estimate_rates + +if sys.version_info < (3, 0): + sys.stderr.write("ERROR: gen_daleMatrices4.py requires Python 3.0 or newer.\n") + sys.exit(1) + +###### Matrix generation (spatial + Dale, see writeup) ################## + + +def _build_placement_grid(L, H, d_min): + """Rectangular grid nodes on [0,L]×[0,H] with spacing d_min.""" + m_max = int(np.floor(L / d_min)) + n_max = int(np.floor(H / d_min)) + gx = (np.arange(0, m_max + 1, dtype=float) * d_min).reshape(-1, 1) + gy = (np.arange(0, n_max + 1, dtype=float) * d_min).reshape(1, -1) + xs = np.broadcast_to(gx, (gx.size, gy.size)).ravel() + ys = np.broadcast_to(gy, (gx.size, gy.size)).ravel() + G = np.column_stack([xs, ys]) + return G, (m_max + 1) * (n_max + 1) + + +def _sample_targets_without_replacement(j, V_col, k_j, rng): + """Draw k_j distinct row indices i≠j with P(i) ∝ V_ij.""" + n = V_col.shape[0] + p = np.asarray(V_col, dtype=float).copy() + p[j] = 0.0 + s = p.sum() + if s <= 0: + raise RuntimeError("zero affinity sum for presynaptic neuron %d" % j) + p /= s + return rng.choice(n, size=k_j, replace=False, p=p) + + +def generate_spatial_dale_network( + n_units, + n_excite, + L, + H, + d_min, + k_min, + k_max, + placement_ker_delta, + weight_var, + R, + rng=None, + verb=1, +): + """ + Returns A_true, E_true, P, tau, sigma, k_out, rho0, D_save in placement order + (x-sorted). D_save has pairwise Euclidean distances and zeros on the diagonal. + """ + if rng is None: + rng = np.random.default_rng() + n_inhib = n_units - n_excite + assert n_excite > 0 and n_inhib > 0 + assert k_min >= 1 and k_max <= n_units - 1 and k_min <= k_max + assert float(placement_ker_delta) > 0 + assert 0 < weight_var < 1 + assert 0 < R < 1 + + r_balance = n_excite / float(n_inhib) + + G, n_grid = _build_placement_grid(L, H, d_min) + if n_grid < n_units: + raise ValueError( + "grid has only %d nodes; need N <= N_G (reduce d_min or increase L/H)" + % n_grid + ) + + idx = rng.choice(n_grid, size=n_units, replace=False) + P = G[idx].astype(np.float64) + order_x = np.lexsort((P[:, 1], P[:, 0])) + P = P[order_x] + + exc_mask = np.zeros(n_units, dtype=bool) + exc_mask[rng.choice(n_units, size=n_excite, replace=False)] = True + tau = np.where(exc_mask, 0, 1).astype(np.int32) + sigma = np.where(exc_mask, 1, -1).astype(np.int8) + + # D_ij = distance row i to column j (postsynaptic i, presynaptic j) + diff = P[:, np.newaxis, :] - P[np.newaxis, :, :] + D = np.sqrt(np.sum(diff * diff, axis=2)) + np.fill_diagonal(D, np.inf) + with np.errstate(divide="ignore"): + V = np.power(D, -float(placement_ker_delta)) + np.fill_diagonal(V, 0.0) + + E = np.zeros((n_units, n_units), dtype=np.int8) + k_out = np.zeros(n_units, dtype=np.int32) + + for j in range(n_units): + kj = int(rng.integers(k_min, k_max + 1)) + k_out[j] = kj + targets = _sample_targets_without_replacement(j, V[:, j], kj, rng) + E[targets, j] = sigma[j] + #E[j, j] = -1 # to tage + + W = np.zeros((n_units, n_units), dtype=np.float64) + for j in range(n_units): + for i in range(n_units): + if i == j: + W[i, j] = float(rng.uniform( -5* weight_var, - 2*weight_var)) + continue # done with diagonal self-feedback + eij = E[i, j] + if eij == 0: continue # edge not existing + W[i, j] = float(rng.uniform(1.0 - weight_var, 1.0 + weight_var)) + if eij == -1: # inhibitory + W[i, j] *= -r_balance + + ev = np.linalg.eigvals(W) + rho0 = float(np.max(np.abs(ev))) + if verb > 0: + print("initW current_rho :", rho0) + if rho0 <= 0: + raise RuntimeError("spectral radius of W is zero (degenerate)") + A_true = W * (R / rho0) + + D_save = D.copy() + np.fill_diagonal(D_save, 0.0) + + return A_true, E, P, tau, sigma, k_out, rho0, D_save + + +def spectral_radius_scaling(A, factors): + n = A.shape[0] + off_diag_mask = ~np.eye(n, dtype=bool) + + print(f"{'factor':>8s} {'ρ(all)':>10s} {'ρ(off-diag)':>12s}") + print("-" * 34) + + for c in factors: + rho_all = np.max(np.abs(np.linalg.eigvals(c * A))) + + A_off = A.copy() + A_off[off_diag_mask] *= c + rho_off = np.max(np.abs(np.linalg.eigvals(A_off))) + + print(f"{c:8.3f} {rho_all:10.4f} {rho_off:12.4f}") + + +def summarize_pairwise_distances(D, verb=1): + """ + Mean and median Euclidean distance over unique unordered pairs (i < j). + D: (N, N) symmetric nonnegative; diagonal entries are ignored. + Returns dict with mean, median, n_nodes, n_pairs. + """ + D = np.asarray(D, dtype=float) + if D.ndim != 2 or D.shape[0] != D.shape[1]: + raise ValueError("D must be a square matrix") + n = D.shape[0] + if n < 2: + out = {"mean": float("nan"), "median": float("nan"), "n_nodes": n, "n_pairs": 0} + if verb > 0: + print("Pairwise distance: N<2 — no pairs (mean/median undefined).") + return out + iu = np.triu_indices(n, k=1) + d = D[iu] + m = float(np.mean(d)) + med = float(np.median(d)) + out = {"mean": m, "median": med, "n_nodes": n, "n_pairs": int(d.size)} + if verb > 0: + print( + "Pairwise distance (unique pairs): mean=%.6g median=%.6g (N=%d, pairs=%d)" + % (m, med, n, out["n_pairs"]) + ) + return out + + +#################### Simulation ################## + + +def set_flat_selfSpiking(Nn, idleRate, spect_radius, tau): + """Bias vector (N,); per-neuron E/I from tau (0=exc, 1=inh).""" + tau = np.asarray(tau).reshape(-1) + if tau.shape[0] != Nn: + raise ValueError("tau length %d != Nn %d" % (tau.shape[0], Nn)) + if tau.dtype.kind not in "iu": + raise ValueError("tau must be integer dtype") + exc = tau.astype(np.int64) == 0 + inh = tau.astype(np.int64) == 1 + if not np.any(exc) or not np.any(inh): + raise ValueError("tau must label at least one excitatory and one inhibitory neuron") + sizeScale = np.sqrt(float(Nn) / 100) + R = float(spect_radius) + idle_eff = np.array(idleRate, dtype=float) + Ri_scaled = idle_eff * R + Bi = np.log(Ri_scaled) + B = np.random.uniform(Bi[0], Bi[1], size=(Nn,)) + #B[exc] += 1.0 - R * 1.5 # - sizeScale + B[exc] += -2.0 + return B + + +def gen_stationary_lag1_poisson(num_steps, dt, A, B_intercept, tau, eta_clip, verb=0): + """ + Generates a multivariate Poisson VAR(1) process: + Y_t ~ Poisson(exp(A @ Y_{t-1} + B_intercept)) + tau: (N,) int, 0=excitatory, 1=inhibitory (same order as A columns / Y). + """ + if A is None: + raise ValueError("Connectivity matrix A cannot be None.") + d = A.shape[0] + tau = np.asarray(tau).reshape(-1) + if tau.shape[0] != d: + raise ValueError("tau length %d != A.shape[0] %d" % (tau.shape[0], d)) + if tau.dtype.kind not in "iu": + raise ValueError("tau must be integer dtype") + tau_i = tau.astype(np.int64) + if np.any((tau_i != 0) & (tau_i != 1)): + raise ValueError("tau entries must be 0 or 1") + exc_idx = np.where(tau_i == 0)[0] + inh_idx = np.where(tau_i == 1)[0] + n_exc = len(exc_idx) + if verb > 0: + print(f"\n=== Generating Poisson Process ===") + print(f"Simulation parameters: steps={num_steps}, dt={dt:.3f}, neurons={d}, excit={n_exc}") + print(f"Matrix A stats: min={np.min(A):.3f}, max={np.max(A):.3f}, mean={np.mean(A):.3f}") + print(f"Bias B stats: min={np.min(B_intercept):.3f}, max={np.max(B_intercept):.3f}, mean={np.mean(B_intercept):.3f}") + + Y = np.zeros((num_steps, d), dtype=int) + Y[0] = np.random.poisson(np.exp(B_intercept) * dt) # initial state + if verb > 0: + print( + f"Initial state: total spikes={np.sum(Y[0])}, excit spikes={np.sum(Y[0, exc_idx])}, inhib spikes={np.sum(Y[0, inh_idx])}" + ) + + kk = min(5, n_exc, len(inh_idx)) + if verb > 0 and kk > 0: + eix = exc_idx[:kk] + iix = inh_idx[:kk] + print( + "t=0 Y[t] sum=%d, Excit(sample idx %s):%s, Inhib(sample idx %s):%s" + % (np.sum(Y[0]), eix, Y[0, eix], iix, Y[0, iix]) + ) + + if verb > 0: + print("Starting main simulation loop...") + for t in range(1, num_steps): + eta = A @ Y[t - 1] + B_intercept + lambda_t = np.exp(np.minimum(eta, eta_clip)) # avoid overflow + Y[t] = np.random.poisson(lambda_t * dt) + + if verb > 0 and t < 5 and kk > 0: + eix = exc_idx[:kk] + iix = inh_idx[:kk] + print( + "t=%d Y[t] sum=%d, Excit(sample):%s, Inhib(sample):%s" + % (t, np.sum(Y[t]), Y[t, eix], Y[t, iix]) + ) + + if verb > 0 and t % (num_steps // 4) == 0: + print(f" Progress: {t}/{num_steps} steps ({t/num_steps*100:.1f}%) - total spikes in this step: {np.sum(Y[t])}") + + if verb > 0: + print(f"Simulation complete. Final state: total spikes={np.sum(Y[-1])}") + print(f"Spike data stats: min={np.min(Y)}, max={np.max(Y)}, mean={np.mean(Y):.2f}, total spikes={np.sum(Y)}") + + return Y + + +def _offdiag_time_kernel( mem_tau_steps, mem_Q): + """ + Off-diagonal damped-oscillator: κ_ℓ = ~ exp(-γℓ) cos(ωℓ) for ℓ ≤ (3/4)*mem_tau_steps; else κ_ℓ = 0. + mem_tau_steps: period in bins (int). + Returns kappa shape (mem_lag_steps,). + """ + mem_lag_steps = max(1, (3 * mem_tau_steps) // 4) + if mem_lag_steps < 1: + raise ValueError("mem_lag_steps must be >= 1") + mem_tau_steps = int(mem_tau_steps) + if mem_tau_steps < 1: + raise ValueError("mem_tau_steps must be >= 1") + mem_Q = float(mem_Q) + if mem_Q <= 0: + raise ValueError("mem_Q must be positive") + tau_f = float(mem_tau_steps) + omega = 2.0 * np.pi / tau_f + gamma = np.pi / (mem_Q * tau_f) + + ell = np.arange(1, mem_lag_steps+1 , dtype=np.float64) + kappa = np.exp(-gamma * ell) * np.cos(omega * ell) + kappa/=kappa[0] # sets 1st amplitude to 1 + m = kappa.shape[0] + print(f" κ (len={m}, oscillator lags ℓ≤(3/4) mem_lag_steps : mem_tau_steps: {mem_tau_steps}): \n {kappa}") + sum_abs = float(np.sum(np.abs(kappa))) + print(f" sum_k |κ_k| = {sum_abs:.12g}") + return kappa + + + +def gen_nonstationary_lagM_modelB_poisson( + num_steps, + dt, + A, + B_intercept, + tau, + mem_tau_steps, + mem_Q, + eta_clip, + h_off, + verb=0, +): + """ + Model B: H_t = diag(A) Y_{t-1} + A_off (sum_l h_l Y_{t-l}), then Poisson as Model A. + First mem_lag_steps bins are zero (pre-history). + mem_tau_steps: oscillation period in bins (int). Uses h_off (κ) from _offdiag_time_kernel; κ_ℓ = 0 for lag ℓ > (3/4)*mem_tau_steps. + """ + if A is None: + raise ValueError("Connectivity matrix A cannot be None.") + d = A.shape[0] + mem_lag_steps=h_off.shape[0] + tau = np.asarray(tau).reshape(-1) + if tau.shape[0] != d: + raise ValueError("tau length %d != A.shape[0] %d" % (tau.shape[0], d)) + if tau.dtype.kind not in "iu": + raise ValueError("tau must be integer dtype") + tau_i = tau.astype(np.int64) + if np.any((tau_i != 0) & (tau_i != 1)): + raise ValueError("tau entries must be 0 or 1") + exc_idx = np.where(tau_i == 0)[0] + inh_idx = np.where(tau_i == 1)[0] + n_exc = len(exc_idx) + + A_off = A.copy() + np.fill_diagonal(A_off, 0.0) + a_diag = np.diag(A).astype(np.float64, copy=False) + + if verb > 0: + print(f"\n=== Generating Model B (lag-M) Poisson ===") + print( + f"steps={num_steps}, dt={dt:.3f}, neurons={d}, excit={n_exc}, mem_lag_steps={mem_lag_steps}, mem_tau_steps={mem_tau_steps}, mem_Q={mem_Q}" + ) + print(f"Matrix A stats: min={np.min(A):.3f}, max={np.max(A):.3f}, mean={np.mean(A):.3f}") + print(f"Bias B stats: min={np.min(B_intercept):.3f}, max={np.max(B_intercept):.3f}, mean={np.mean(B_intercept):.3f}") + + Y = np.zeros((num_steps, d), dtype=int) + + kk = min(5, n_exc, len(inh_idx)) + if verb > 0: + print("Initial Y[0:%d] set to zero (pre-history)." % min(mem_lag_steps, num_steps)) + + if verb > 0: + print("Starting main simulation loop...") + t_start = time.time() + for t in range(1, num_steps): + S = np.zeros(d, dtype=np.float64) + for ell in range(1, mem_lag_steps + 1): + tt = t - ell + if tt >= 0: + S += h_off[ell - 1] * Y[tt] + eta = a_diag * Y[t - 1] + A_off @ S + B_intercept + lambda_t = np.exp(np.minimum(eta, eta_clip)) + Y[t] = np.random.poisson(lambda_t * dt) + + if verb > 0 and t < 5 and kk > 0: + eix = exc_idx[:kk] + iix = inh_idx[:kk] + print( + "t=%d Y[t] sum=%d, Excit(sample):%s, Inhib(sample):%s" + % (t, np.sum(Y[t]), Y[t, eix], Y[t, iix]) + ) + + if verb > 0 and t % (num_steps // 4) == 0: + ela_t = time.time() - t_start + print(f" Progress: {t}/{num_steps} steps ({t/num_steps*100:.1f}%) - total spikes in this step: {np.sum(Y[t])} elaT={ela_t:.1f}s") + + if verb > 0: + print(f"Simulation complete. Final state: total spikes={np.sum(Y[-1])}") + print(f"Spike data stats: min={np.min(Y)}, max={np.max(Y)}, mean={np.mean(Y):.2f}, total spikes={np.sum(Y)}") + + return Y + + +######################### +# MAIN +######################### + + +def main(): + print("=" * 60) + print("DALE POISSON SIMULATION (Model A or B)") + print("=" * 60) + + parser = argparse.ArgumentParser(description="Simulate a recurrent neural network with Dale's principle.") + parser.add_argument("--spike_model", type=str, default="A", choices=["A", "B"], + help="Spike GLM: A = lag-1 stationary; B = lag-M off-diagonal damped-oscillator kernel.") + parser.add_argument("--time_kernel_q_tau", type=float, nargs=2, default=[3.0, 0.4], metavar=("mem_Q", "mem_tau"), + help="Model B: [mem_Q, mem_tau] — quality factor, period mem_tau>0 (s); mem_lag_steps=(3*mem_tau_steps)//4 (~1.5π phase).") + parser.add_argument("--num_neurons", type=int, default=50, help="Total number of neurons in the network.") + parser.add_argument("--num_excite", type=int, default=None, help="Number of excitatory neurons.") + parser.add_argument("--placement_H_L_delta", type=float, nargs=3, default=[1.0, 2.0, 2.0], + metavar=("placement_H", "placement_L", "placement_ker_delta"), help="Placement: [0,H] height, [0,L] width, distance-kernel exponent δ>0 in V_ij ∝ D_ij^{-δ} (e.g. 1, 2, 0.2).") + parser.add_argument("--placement_min_dist", type=float, default=0.01, + help="Grid spacing d_min; minimum inter-neuron distance.") + parser.add_argument("--edge_prob", type=float, nargs=2, default=[0.05, 0.2], + help="Out-degree range as fractions of N: k_min=max(1,floor(lo*N)), k_max=min(N-1,floor(hi*N)).") + parser.add_argument("--init_weight_var", type=float, default=0.4, help="Fractional weight variation v for Uniform(1-v,1+v).") + parser.add_argument("--num_steps", type=int, default=10_001, help="Number of time steps for simulation.") + parser.add_argument("--step_size", type=float, default=0.01, help="Integration time step size (dt) in seconds.") + parser.add_argument("--spectral_radius", type=float, default=0.90, help="Target spectral radius value for the connectivity matrix.") + parser.add_argument("--idleRate", type=float, nargs=2, default=[10, 20], help="Range of idle firing rates [min, max] in Hz.") + parser.add_argument("-v", "--verb", type=int, default=1, help="Verbosity level (0=quiet, 1=normal).") + parser.add_argument("--dataName", type=str, default=None, help="Base name for output files (default: daleN_).") + parser.add_argument("--basePath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="Output directory for all files.") + + np.set_printoptions(precision=3, suppress=True) + + args = parser.parse_args() + placement_H, placement_L, placement_ker_delta_raw = ( + float(args.placement_H_L_delta[0]), + float(args.placement_H_L_delta[1]), + float(args.placement_H_L_delta[2]), + ) + if not (placement_H > 0 and placement_L > 0): + raise ValueError("placement_H_L_delta requires positive H and L") + placement_ker_delta = float(placement_ker_delta_raw) + if placement_ker_delta <= 0: + raise ValueError("placement_H_L_delta third value (placement_ker_delta, δ) must be positive") + placement_min_dist = float(args.placement_min_dist) + if placement_min_dist <= 0: + raise ValueError("placement_min_dist must be positive") + mem_Q, mem_tau_arg = args.time_kernel_q_tau + mem_tau = None + mem_tau_steps = None + + if args.spike_model == "B": + mem_tau = float(mem_tau_arg) + if mem_Q <= 0: + raise ValueError("mem_Q must be positive for spike_model B") + assert mem_tau > 0, "mem_tau must be positive for spike_model B" + mem_tau_steps = max(1, int(round(float(mem_tau) / float(args.step_size)))) + + args.varTwindow = 5 # (sec) + args.poisson_eta_clip = 5 # ~ [1e-3Hz , 1e+3Hz] + if args.dataName is None: + args.dataName = "daleN%d_" % args.num_neurons + hashlib.md5(os.urandom(32)).hexdigest()[:6] + + outPath = os.path.join(args.basePath, "truthDale") + if args.num_excite is None: + args.num_excite = int(0.8 * args.num_neurons) + + Nn = args.num_neurons + prob_lo, prob_hi = float(args.edge_prob[0]), float(args.edge_prob[1]) + if not (0 < prob_lo < prob_hi <= 1.0): + raise ValueError("edge_prob must satisfy 0 < lo < hi <= 1") + k_out_min = max(1, int(np.floor(prob_lo * Nn))) + k_out_max = int(np.floor(prob_hi * Nn)) + k_out_max = max(k_out_min, k_out_max) + k_out_max = min(k_out_max, Nn - 1) + + rng = np.random.default_rng() + + print("gen dale matrices args:", vars(args), "\n") + + assert Nn >= 10 + assert args.num_excite >= 5 + assert args.num_excite < Nn + assert os.path.exists(args.basePath) + assert os.path.exists(outPath) + assert args.num_steps >= 1000 + assert args.step_size > 0.001 + assert args.idleRate[0] >= 0.5 + assert args.idleRate[1] > args.idleRate[0] + + print(f"\n{'='*60}") + print(f" Spectral radius: R={args.spectral_radius:.3f}") + print(f"{'='*60}") + + A_dale, E_true, P_pos, tau, sigma, k_out, rho0, D_mat = generate_spatial_dale_network( + n_units=Nn, + n_excite=args.num_excite, + L=placement_L, + H=placement_H, + d_min=placement_min_dist, + k_min=k_out_min, + k_max=k_out_max, + placement_ker_delta=placement_ker_delta, + weight_var=args.init_weight_var, + R=args.spectral_radius, + rng=rng, + verb=args.verb, + ) + + if args.verb > 0: + summarize_pairwise_distances(D_mat, verb=1) + + if args.verb > 0: + total_connections = A_dale.size + zero_connections = np.sum(np.abs(A_dale) < 1e-10) + non_zero_connections = total_connections - zero_connections + sparsity = zero_connections / total_connections + print(f"Matrix sparsity: {sparsity*100:.1f}% ({zero_connections}/{total_connections} connections are zero)") + print(f"Non-zero connections: {non_zero_connections} ({(1-sparsity)*100:.1f}%)") + + dale_conf = { + "num_neurons": args.num_neurons, + "num_excite": args.num_excite, + "spectral_radius": args.spectral_radius, + "placement_L": placement_L, + "placement_H": placement_H, + "placement_min_dist": placement_min_dist, + "edge_prob": [prob_lo, prob_hi], + "k_out_min": k_out_min, + "k_out_max": k_out_max, + "placement_ker_delta": placement_ker_delta, + "init_weight_var": args.init_weight_var, + "rho0_unscaled": rho0, + "idleRate": args.idleRate, + "spike_model": args.spike_model, + } + if args.spike_model == "B": + dale_conf["mem_tau"] = float(mem_tau) + dale_conf["mem_tau_steps"] = int(mem_tau_steps) + dale_conf["mem_Q"] = float(mem_Q) + else: + dale_conf["mem_tau"] = 0.0 + dale_conf["mem_tau_steps"] = 0 + dale_conf["mem_Q"] = 0.0 + evol_conf = { + "num_steps": args.num_steps, + "step_size": args.step_size, + "evol_time": args.num_steps * args.step_size, + "poisson_eta_clip": args.poisson_eta_clip, + "spike_model": args.spike_model, + } + + + B_true = set_flat_selfSpiking(Nn, args.idleRate, args.spectral_radius, tau) + + max_samples = 100_000 + + print(f"\n{'='*60}") + print(" Poisson evaluation spike generation (model %s)" % args.spike_model) + print(f"{'='*60}") + + if args.spike_model == "B": + offdiag_kernel = _offdiag_time_kernel(mem_tau_steps, mem_Q) + dale_conf["mem_lag_steps"] = int(offdiag_kernel.shape[0]) + else: + offdiag_kernel = np.zeros((0,), dtype=np.float64) + dale_conf["mem_lag_steps"] = 0 + + evol_conf["mem_lag_steps"] = dale_conf["mem_lag_steps"] + + print("Dale configuration:") + pprint(dale_conf) + + start_time = time.time() + if args.spike_model == "A": + Y = gen_stationary_lag1_poisson( + num_steps=args.num_steps, + dt=args.step_size, + A=A_dale, + B_intercept=B_true, + tau=tau, + eta_clip=args.poisson_eta_clip, + verb=args.verb, + ) + else: + Y = gen_nonstationary_lagM_modelB_poisson( + num_steps=args.num_steps, + dt=args.step_size, + A=A_dale, + B_intercept=B_true, + tau=tau, + mem_tau_steps=mem_tau_steps, + mem_Q=mem_Q, + eta_clip=args.poisson_eta_clip, + h_off=offdiag_kernel, + verb=args.verb, + + ) + sim_time = time.time() - start_time + print("Spike generation completed in %.1f seconds" % sim_time) + + stats_dict, rates_dict, _ = estimate_rates( + Y, + args.step_size, + tau, + max_samples, + args.spectral_radius, + varTwindow=args.varTwindow, + mxNn=5, + verb=0, + ) + stats_dict["var_time_window_sec"] = float(args.varTwindow) + stats_dict["max_samples"] = int(max_samples) + + Y_u8 = np.clip(Y, 0, 255).astype(np.uint8) + A_off_true = A_dale.copy() + np.fill_diagonal(A_off_true, 0.0) + A_diag_true = np.diag(A_dale).copy() + + trueD = { + "A_diag_true": A_diag_true, + "A_off_true": A_off_true, + "B_true": B_true, + "E_true": E_true, + "node_positions": P_pos, + "node_is_inhibitory": tau, + "node_distance_matrix": D_mat, + "offdiag_kernel": offdiag_kernel, + } + + # Compute max spikes per neuron across all bins + max_spikes_per_neuron = Y.max(axis=0) + sp_min = float(max_spikes_per_neuron.min()) + sp_avg = float(max_spikes_per_neuron.mean()) + sp_max = float(max_spikes_per_neuron.max()) + sp_pc = np.percentile(max_spikes_per_neuron, [25, 50, 75]).tolist() + + if args.verb > 0: + print(f"\nSpike count statistics (max per neuron over all bins):") + print(f" min: {sp_min:.1f}, avg: {sp_avg:.1f}, max: {sp_max:.1f}") + print(f" percentiles [25, 50, 75]: {sp_pc}") + + trueMD = { + "dale_conf": dale_conf, + "evol_conf": evol_conf, + "short_name": args.dataName, + "provenance": {"state_model_file": args.dataName}, + "max_spike_stats": { + "min": sp_min, "avg": sp_avg, "max": sp_max, "percentiles": sp_pc + } + } + + spikeD = { + "spikes": Y_u8, + "single_rates": rates_dict["single_rates"], + "sigle_rates_var": rates_dict["sigle_rates_var"], + "single_fano_fact": rates_dict["single_fano_fact"], + } + spikeMD = { + "short_name": args.dataName, + "time_step_sec": args.step_size, + "data_type": "simDaleMemKer", + "poisson_eta_clip": args.poisson_eta_clip, + "spike_model": args.spike_model, + "placement_L": float(placement_L), + "placement_H": float(placement_H), + "placement_min_dist": float(placement_min_dist), + "mem_lag_steps": int(dale_conf["mem_lag_steps"]), + } + + outFt = os.path.join(outPath, args.dataName + ".simTruth.npz") + write_data_npz(trueD, outFt, metaD=trueMD) + if args.verb > 1: + pprint(trueMD) + outFs = os.path.join(outPath, args.dataName + ".spikes.npz") + write_data_npz(spikeD, outFs, metaD=spikeMD) + if args.verb > 1: + pprint(spikeMD) + + print("\nSimulation completed successfully!") + print(f"\nRate Summary {args.dataName} N={args.num_neurons} exc={args.num_excite}, R={args.spectral_radius:.3f} ") + s = stats_dict + print( + f" rates: all={s['avg_spike_rate_all']:14.1f} exc={s['avg_spike_rate_excit']:14.1f} inh={s['avg_spike_rate_inhib']:14.1f}" + ) + print("\nNext step commands:") + print(" basePath=" + args.basePath) + print(" ./view_daleMatrix4.py --basePath $basePath --dataName %s -p a e b c d g f -X " % args.dataName) + print(" ./view_spikesTrain4.py --basePath $basePath --dataName %s --time_range_sec 1 8 -p b --time_rebin2 2 -X " % args.dataName) + print(" ./movie_spikesTrain4.py --basePath $basePath --dataName %s --time_range_sec 1 8 --flushSize .2 " % args.dataName) + print(" ./view_spikesTrain4.py --basePath $basePath --dataName %s --time_range_sec 0 20 -p b -X " % args.dataName) + print(" ./memKern_EM_train4.py --basePath $basePath --dataName %s --time_range_sec 1 20 " % args.dataName) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver4_lagM_offDiag/memKern_EM_eval4.py b/causal_net/nonStation_ver4_lagM_offDiag/memKern_EM_eval4.py new file mode 100755 index 00000000..b76ae33c --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/memKern_EM_eval4.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +""" +Evaluation and plotting for prism EM results. +""" + +import os +import argparse +import numpy as np +from pprint import pprint +from toolbox.Util_NumpyIO import read_data_npz +from PlotterMemKernEM import Plotter + + +def compute_ll_gap(spikes_sub, A_true, B_true, dt, eta_clip): + """Per-time LL gap between best and 2nd-best true-state likelihood.""" + y_prev = np.asarray(spikes_sub[:-1], dtype=np.float64) + y_curr = np.asarray(spikes_sub[1:], dtype=np.float64) + a_true = np.asarray(A_true, dtype=np.float64) + b_true = np.asarray(B_true, dtype=np.float64) + + t_pairs = y_prev.shape[0] + m_states = b_true.shape[0] + ll_gap = np.zeros((t_pairs,), dtype=np.float64) + if m_states <= 1: + return ll_gap + + for t in range(t_pairs): + yp = y_prev[t] + yc = y_curr[t] + base = a_true @ yp + eta = base[None, :] + b_true + eta_c = np.minimum(eta, float(eta_clip)) + lam = np.exp(eta_c) * float(dt) + scores = np.sum(yc[None, :] * eta_c - lam, axis=1) + top2 = np.partition(scores, -2)[-2:] + ll_gap[t] = float(top2[-1] - top2[-2]) + return ll_gap + + +def eval_em_metrics_time(fitD, md, spikes): + """Compute per-time metrics for -p f canvas.""" + trainMD = md["train"] + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + dt = float(trainMD["time_step_sec"]) + eta_clip = float(trainMD["eta_clip"]) + lambda2 = float(trainMD["lambda2"]) + + spikes_sub = np.asarray(spikes[t0_bin : t1_bin + 1], dtype=np.float64) + yp = spikes_sub[:-1] + yc = spikes_sub[1:] + + if "c_hat" not in fitD or "S_true" not in md: + return {} + + a_hat = np.asarray(fitD["A_off_hat"], dtype=np.float64) + b_hat = np.asarray(fitD["B_hat"], dtype=np.float64) + c_hat = np.asarray(fitD["c_hat"], dtype=np.float64) + + n_pairs = spikes_sub.shape[0] - 1 + assert c_hat.shape[0] == spikes_sub.shape[0], "c_hat and spikes_sub must have matching time bins" + c_pairs = c_hat[1:] + c_prev = c_hat[:-1] + + rates = np.asarray(fitD["single_rates"], dtype=np.float64) + w = 1.0 / np.maximum(rates, 0.1) + w /= w.mean() + + eta = yp @ a_hat.T + c_pairs @ b_hat + eta_c = np.minimum(eta, eta_clip) + lam = np.exp(eta_c) * dt + log_dt = np.log(dt) + nll_t = np.sum(w[None, :] * (lam - yc * (eta + log_dt)), axis=1) + + dc = c_pairs - c_prev + l2_t = lambda2 * np.sum(dc * dc, axis=1) + + ll_gap = compute_ll_gap(spikes_sub, md["A_true"], md["B_true"], dt, eta_clip) + + s_true = np.asarray(md["S_true"])[t0_bin : t1_bin + 1] + s_hat = np.asarray(fitD["S_hat"]) + assert s_true.shape[0] == s_hat.shape[0], "S_true and S_hat must have matching length" + acc = float((s_true == s_hat).mean()) + + return { + "acc": acc, + "loss_nll_time": nll_t, + "loss_l2_time": l2_t, + "ll_gap": ll_gap, + "ll_gap_mean": float(np.mean(ll_gap)), + } + + +def eval_state_recovery(fitD, md): + """Compute state recovery table on full training window.""" + trainMD = md["train"] + t0_bin, t1_bin = [int(x) for x in trainMD["time_range_bins"]] + s_true = np.asarray(md["S_true"], dtype=np.int64)[t0_bin : t1_bin + 1] + s_hat = np.asarray(fitD["S_hat"], dtype=np.int64) + s_hat_cl = np.asarray(fitD["S_hat_CL"], dtype=np.float64) + m_states = int(trainMD["num_states"]) + + assert s_true.shape[0] == s_hat.shape[0] == s_hat_cl.shape[0], \ + "S_true/S_hat/S_hat_CL length mismatch on training window" + + avg_acc = float((s_true == s_hat).mean()) + state_acc_cl = [] + bins_hat = np.bincount(s_hat, minlength=m_states).astype(np.int64) + enter_hat = np.zeros(m_states, dtype=np.int64) + if s_hat.shape[0] > 1: + enter_idx = np.where(s_hat[1:] != s_hat[:-1])[0] + 1 + if enter_idx.size > 0: + enter_hat = np.bincount(s_hat[enter_idx], minlength=m_states).astype(np.int64) + + trans_hat = np.zeros((m_states, m_states), dtype=np.int64) + if s_hat.shape[0] > 1: + np.add.at(trans_hat, (s_hat[:-1], s_hat[1:]), 1) + + for m in range(m_states): + mask = (s_hat == m) + cnt = int(mask.sum()) + if cnt > 0: + acc_m = float((s_true[mask] == m).mean()) + cl_m = float(s_hat_cl[mask].mean()) + else: + acc_m = float("nan") + cl_m = float("nan") + state_acc_cl.append([acc_m, cl_m, int(bins_hat[m]), int(enter_hat[m])]) # [acc, CL, bins_hat, enter_hat] + + return { + "avg_acc": avg_acc, + "state_acc_cl": state_acc_cl, + "trans_hat": trans_hat, + } + + +def main(): + parser = argparse.ArgumentParser( + description="Evaluate and plot prism EM results") + parser.add_argument("--dataName", type=str, required=True, + help="Base name for prismEM file") + parser.add_argument("--basePath", type=str, + default="/pscratch/sd/b/balewski/2026_causalNet_tmp3/", + help="Head dir for input/output data") + parser.add_argument("-p", "--showPlots", type=str, nargs='+', + default="a", + help="Plot types: a=EM convergence summary, b=fit-vs-truth diagnostics (A_off, A_diag, B, κ vs simTruth .npz)") + parser.add_argument("--minW", type=float, default=0.02, + help="Threshold for A-matrix edge eval") + parser.add_argument("--timeReb", type=int, default=20, + help="Time rebin factor for time-axis plots") + g = parser.add_argument_group("data") + g.add_argument("-T", "--time_range_sec", default=[0.0, 15.0], + nargs=2, type=float, + help="Time window [t0, t1] in seconds") + parser.add_argument("-X", "--noXterm", action="store_true", + help="Disable X terminal for plotting") + parser.add_argument("-v", "--verb", type=int, default=1, + help="Verbosity level") + args = parser.parse_args() + + args.inpPath = os.path.join(args.basePath, "memKernFit") + args.outPath = os.path.join(args.basePath, "plots") + os.makedirs(args.outPath, exist_ok=True) + + show_req = ''.join(args.showPlots) + args.showPlots = "" + for c in show_req: + if c in "abcdefgh" and c not in args.showPlots: + args.showPlots += c + + print("EM-eval args:", vars(args), "\n") + + # ── load EM fit ────────────────────────────────────────────────── + fitFF = os.path.join(args.inpPath, f"{args.dataName}.memKernEM.npz") + fitD, fitMD = read_data_npz(fitFF) + assert isinstance(fitMD, dict), "Expected metadata dict in prismEM file" + + if args.verb > 1: pprint(fitMD) + + MD = {**fitMD} #, "short_name": args.dataName} + + #--- load spikes ---- + prov = fitMD["provenance"] + spikesF=prov['spiksData_file'] + spikesFF = os.path.join(args.basePath, "truthDale", f"{spikesF}.spikes.npz") # should be; "spikesData" + spikeD, _ = read_data_npz(spikesFF, verb=args.verb > 1) + spikes = spikeD["spikes"] + + # ── load ground truth if available ─────────────────────────────── + prov = fitMD["provenance"] + trueF=spikesF + trueFF = os.path.join(args.basePath, "truthDale", f"{trueF}.simTruth.npz") + trueD, trueMD = read_data_npz(trueFF, verb=args.verb > 1) + if args.verb > 1: pprint(trueMD) + + has_truth = False + for xx in [ 'dale_conf', 'evol_conf']: + MD[xx]=trueMD[xx] + + + MD["A_off_true"] = trueD["A_off_true"] + MD["A_diag_true"] = trueD["A_diag_true"] + MD["B_true"] = trueD["B_true"] + MD["E_true"] = trueD["E_true"] + MD["offdiag_kernel"] = trueD["offdiag_kernel"] + MD["A_true"] = MD["A_off_true"] + np.diag(MD["A_diag_true"]) + MD["eval_f"] = eval_em_metrics_time(fitD, MD, spikes) + + #MD["short_name"] = args.dataName + + # ── plot ────────────────────────────────────────────────────────── + args.prjName = args.dataName + plot = Plotter(args) + + if 'a' in args.showPlots: + plot.summary_memKerEM(fitD, MD, figId=1) + + if 'b' in args.showPlots: + plot.correl_fit_truth(fitD, MD, figId=2) + + + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver4_lagM_offDiag/memKern_EM_train4.py b/causal_net/nonStation_ver4_lagM_offDiag/memKern_EM_train4.py new file mode 100755 index 00000000..560ff9af --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/memKern_EM_train4.py @@ -0,0 +1,875 @@ +#!/usr/bin/env python3 +""" +memKern_EM_train.py — Two-block EM training for Poisson GLM +with off-diagonal memory kernel. + +Multi-GPU version: uses all visible GPUs on the node via +torch.distributed (NCCL). Each rank gets non-overlapping samples +(offset by rank * stride). Only rank 0 reads/writes files and prints. + +Jointly infers from observed spikes: + - diagonal drive A^diag (N,) + - off-diagonal matrix A^off (N, N), zeros on diagonal + - bias vector B (N,) + - shared lag kernel kappa (M_cut,), kappa[0]=1 fixed + +Algorithm (from writeup): + Block 1 (E-like): update kappa with lambda2 smoothness, A/B fixed + Block 2 (M-like): update A^diag, A^off, B with lambda3 L1 sparsity, kappa fixed + +Usage: + torchrun --standalone --nproc_per_node=gpu memKern_EM_train4.py --dataName --basePath + # Or single-GPU: + python memKern_EM_train4.py --dataName --basePath +""" + +import os +import time +import math +import secrets +import argparse +import logging # Added by user instruction + +# Suppress annoying torch.compile profiler warnings +logging.getLogger('torch._logging').setLevel(logging.ERROR) # Added by user instruction + +from pprint import pprint + +import numpy as np +import torch +import torch.optim as optim +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + + +# ═══════════════════════════════════════════════════════════════════════ +# Distributed helpers +# ═══════════════════════════════════════════════════════════════════════ + +def setup_distributed(): + """Initialize distributed backend. Returns (rank, world_size, device). + Falls back to single-GPU if env vars are not set.""" + if "RANK" in os.environ: + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + else: + rank = 0 + world_size = 1 + device = torch.device("cuda:0") + torch.cuda.set_device(device) + return rank, world_size, device + + +def cleanup_distributed(): + if dist.is_initialized(): + dist.destroy_process_group() + + +def is_rank0(rank): + return rank == 0 + + +def broadcast_array(arr, rank, device, dtype=torch.float32): + """Broadcast a numpy array from rank 0 to all ranks. + Returns a numpy array on CPU.""" + if rank == 0: + t = torch.tensor(arr, device=device) + else: + t = torch.empty(0, device=device) + + # First broadcast shape + if rank == 0: + shape = torch.tensor(list(arr.shape), dtype=torch.long, device=device) + else: + shape = torch.empty(0, dtype=torch.long, device=device) + + ndim = torch.tensor(arr.ndim if rank == 0 else 0, dtype=torch.long, device=device) + if dist.is_initialized(): + dist.broadcast(ndim, src=0) + nd = int(ndim.item()) + + if rank != 0: + shape = torch.empty(nd, dtype=torch.long, device=device) + if dist.is_initialized(): + dist.broadcast(shape, src=0) + + if rank != 0: + t = torch.empty(*shape.tolist(), dtype=dtype, device=device) + if dist.is_initialized(): + dist.broadcast(t, src=0) + return t.cpu().numpy() + + +# ═══════════════════════════════════════════════════════════════════════ +# Generative model forward pass +# ═══════════════════════════════════════════════════════════════════════ + +def compute_S(Y_hist, kappa): + """Compute off-diagonal history vector S_t for all samples. + + Args: + Y_hist: (batch, M_cut, N) — spike history, Y_hist[:,0,:] = Y_{t-1}, etc. + kappa: (M_cut,) — kernel weights, kappa[0]=1 + + Returns: + S: (batch, N) — weighted sum of history + """ + # S = sum_{ell=0}^{M-1} kappa[ell] * Y_hist[:, ell, :] + # kappa: (M_cut,) -> (1, M_cut, 1) for broadcasting + return (Y_hist * kappa[None, :, None]).sum(dim=1) # (batch, N) + + +def compute_eta(Y_prev, S, A_diag, A_off, B): + """Compute log-rate eta for all samples. + + Args: + Y_prev: (batch, N) — most recent spike counts (= Y_hist[:,0,:]) + S: (batch, N) — off-diagonal history + A_diag: (N,) + A_off: (N, N) — zeros on diagonal + B: (N,) + + Returns: + eta: (batch, N) + """ + # diagonal contribution: A_diag_i * Y_{t-1,i} + diag_term = Y_prev * A_diag[None, :] # (batch, N) + # off-diagonal contribution: (A_off @ S^T)^T = S @ A_off^T + off_term = S @ A_off.t() # (batch, N) + return diag_term + off_term + B[None, :] + + +def compute_mu(eta, eta_clip, dt): + """Clipped intensity: mu = exp(clamp(eta)) * dt.""" + eta_c = torch.clamp(eta, max=eta_clip) + return torch.exp(eta_c) * dt, eta_c + + +def poisson_nll(mu, eta_c, Y_curr, log_dt): + """Poisson NLL: sum of [mu - Y * (eta_c + log_dt)].""" + return (mu - Y_curr * (eta_c + log_dt)).sum() + + +# ═══════════════════════════════════════════════════════════════════════ +# Dataset: (Y_hist, Y_curr) pairs +# ═══════════════════════════════════════════════════════════════════════ + +def build_spike_history(spikes_np, M_cut, stride=1, offset=0): + """Build (Y_hist, Y_curr) arrays from spike data — fully vectorized. + + Args: + spikes_np: (T, N) int array of spike counts + M_cut: history depth + stride: sliding window step + offset: starting offset for this rank's samples + + Returns: + Y_hist: (N_samples, M_cut, N) float32 + Y_curr: (N_samples, N) float32 + """ + T, N = spikes_np.shape + # indices of the 'current' time bin for each sample + t_indices = np.arange(M_cut + offset, T, stride) + n_samples = len(t_indices) + + # Vectorized: build lag indices (n_samples, M_cut) + # lag_indices[i, ell] = t_indices[i] - 1 - ell + lag_offsets = np.arange(M_cut) # [0, 1, ..., M_cut-1] + lag_indices = t_indices[:, None] - 1 - lag_offsets[None, :] # (n_samples, M_cut) + + # Advanced indexing: gather all history windows at once + Y_hist = spikes_np[lag_indices].astype(np.float32) # (n_samples, M_cut, N) + Y_curr = spikes_np[t_indices].astype(np.float32) # (n_samples, N) + return Y_hist, Y_curr + + +# ═══════════════════════════════════════════════════════════════════════ +# Block 1: kernel update (kappa) +# ═══════════════════════════════════════════════════════════════════════ + +def run_block1_kernel_update( + Y_hist_gpu, Y_curr_gpu, A_diag, A_off, B, + kappa, dt, eta_clip, lambda2, lr_kappa, n_iter, M_cut, + rank, world_size +): + """Update kappa[1:] by gradient descent on Poisson NLL + smoothness penalty. + + kappa[0] = 1 is fixed (anchor). + The problem is convex in kappa since S_t is linear in kappa. + Gradients are all-reduced across ranks. + + Args: + Y_hist_gpu: (N_samples_local, M_cut, N) on GPU + Y_curr_gpu: (N_samples_local, N) on GPU + A_diag, A_off, B: fixed parameters on GPU + kappa: (M_cut,) tensor on GPU, kappa[0]=1 + dt, eta_clip: scalars + lambda2: smoothness penalty weight + lr_kappa: learning rate for kappa update + n_iter: number of gradient steps + M_cut: kernel length + rank, world_size: for gradient all-reduce + + Returns: + kappa: updated (M_cut,) tensor + nll_val: final NLL value + """ + log_dt = math.log(dt) + N_samples_local = Y_hist_gpu.shape[0] + + # Get total N_samples across all ranks for proper averaging + n_total = torch.tensor(float(N_samples_local), device=kappa.device) + if dist.is_initialized(): + dist.all_reduce(n_total, op=dist.ReduceOp.SUM) + n_total_val = n_total.item() + + # We only optimize kappa[1:] + kappa_free = kappa[1:].clone().detach().requires_grad_(True) + + for iteration in range(n_iter): + # Reconstruct full kappa with anchor + kappa_full = torch.cat([torch.ones(1, device=kappa.device), kappa_free]) + + # Forward pass + S = compute_S(Y_hist_gpu, kappa_full) + Y_prev = Y_hist_gpu[:, 0, :] # most recent frame + eta = compute_eta(Y_prev, S, A_diag, A_off, B) + mu, eta_c = compute_mu(eta, eta_clip, dt) + + # Poisson NLL (local, unnormalized) + nll = poisson_nll(mu, eta_c, Y_curr_gpu, log_dt) + + # Smoothness penalty on kappa: sum_{ell=2}^{M-1} (kappa[ell] - kappa[ell-1])^2 + if M_cut >= 3: + diffs = kappa_full[2:] - kappa_full[1:-1] + smooth_pen = lambda2 * (diffs ** 2).sum() + else: + smooth_pen = torch.tensor(0.0, device=kappa.device) + + loss = nll / n_total_val + smooth_pen + + # Backward + if kappa_free.grad is not None: + kappa_free.grad.zero_() + loss.backward() + + # All-reduce the gradient across ranks + if dist.is_initialized(): + dist.all_reduce(kappa_free.grad, op=dist.ReduceOp.SUM) + + # Gradient step + with torch.no_grad(): + kappa_free -= lr_kappa * kappa_free.grad + + # Reconstruct final kappa + with torch.no_grad(): + kappa_out = torch.cat([torch.ones(1, device=kappa.device), kappa_free.detach()]) + + # For reporting, get global NLL + nll_report = nll.detach() + if dist.is_initialized(): + dist.all_reduce(nll_report, op=dist.ReduceOp.SUM) + + return kappa_out, float((nll_report / n_total_val).item()) + + +# ═══════════════════════════════════════════════════════════════════════ +# Block 2: network update (A_diag, A_off, B) +# ═══════════════════════════════════════════════════════════════════════ + +class NetworkModel(torch.nn.Module): + """Wraps A_diag, A_off, B as learnable parameters for Adam.""" + + def __init__(self, N, A_diag_init, A_off_init, B_init): + super().__init__() + self.N = N + self.A_diag = torch.nn.Parameter(A_diag_init.clone()) + self.A_off = torch.nn.Parameter(A_off_init.clone()) + self.B = torch.nn.Parameter(B_init.clone()) + + def forward(self, Y_prev, S, dt, eta_clip): + """ + Args: + Y_prev: (batch, N) + S: (batch, N) — precomputed with fixed kappa + dt: scalar + eta_clip: scalar + Returns: + mu: (batch, N) + eta_c: (batch, N) + """ + # Zero out diagonal of A_off during forward + A_off_masked = self.A_off * (1.0 - torch.eye(self.N, device=self.A_off.device)) + eta = compute_eta(Y_prev, S, self.A_diag, A_off_masked, self.B) + mu, eta_c = compute_mu(eta, eta_clip, dt) + return mu, eta_c + + +def offdiag_l1_mean(A_off, N): + """Mean absolute value of off-diagonal entries: eq (2) from writeup.""" + mask = ~torch.eye(N, dtype=torch.bool, device=A_off.device) + return A_off[mask].abs().mean() + + +def enforce_zero_diagonal_(A_off, N): + """Zero out diagonal of A_off in-place.""" + with torch.no_grad(): + idx = torch.arange(N, device=A_off.device) + A_off[idx, idx] = 0.0 + + +def enforce_spectral_radius_(A_off, rho_max, N): + """Hard spectral projection on A_off. Returns rho before projection.""" + with torch.no_grad(): + rho = torch.linalg.eigvals(A_off).abs().max().item() + if rho > rho_max: + A_off.mul_(rho_max / rho) + return rho + + +def offdiag_soft_threshold_(A_off, lr, lam, N): + """Proximal L1 on off-diagonal elements of A_off.""" + if lam <= 0: + return + with torch.no_grad(): + mask = ~torch.eye(N, dtype=torch.bool, device=A_off.device) + thresh = lr * lam + v = A_off[mask] + A_off[mask] = v.sign() * (v.abs() - thresh).clamp(min=0.0) + + +# ═══════════════════════════════════════════════════════════════════════ +# Initialization helpers +# ═══════════════════════════════════════════════════════════════════════ + +def init_A_diag_from_spikes(N, dt): + """Initialize A_diag with random values in range [-0.5, -0.15].""" + return np.random.uniform(-0.5, -0.15, size=N).astype(np.float32) + + +def init_A_off_from_spikes(Y_hist_gpu, Y_curr_gpu): + """Initialize A_off from lag-1 cross-correlations on GPU, zeros on diagonal.""" + # Y_hist_gpu[:, 0, :] is Y_{t-1}, Y_curr_gpu is Y_t + Y0 = Y_hist_gpu[:, 0, :].double() # (n_samples, N) + Y1 = Y_curr_gpu.double() + n = Y0.shape[0] + m0 = Y0.mean(dim=0, keepdim=True) + s0 = Y0.std(dim=0, keepdim=True) + 1e-8 + m1 = Y1.mean(dim=0, keepdim=True) + s1 = Y1.std(dim=0, keepdim=True) + 1e-8 + Z0 = (Y0 - m0) / s0 + Z1 = (Y1 - m1) / s1 + C = (Z0.t() @ Z1) / n # (N, N) on GPU + A_off = (C * 0.1).float().cpu().numpy() + np.fill_diagonal(A_off, 0.0) + return A_off + + +def init_B_from_spikes(Y_curr_gpu, dt): + """Initialize B from mean log firing rates using GPU data.""" + mean_rate = Y_curr_gpu.double().mean(dim=0) / dt + mean_rate = mean_rate.clamp(min=1e-6) + B = mean_rate.log().float().cpu().numpy() + return B + + +def init_kappa_damped_oscillator(M_cut): + """Initialize kappa as a damped oscillator, kappa[0]=1. + Functional form: kappa(ell) = exp(-alpha*ell) * cos(omega*ell) + where: + omega = 1.5 * pi / M_cut (1.5 pi phase over M_cut steps) + alpha = omega / (2 * Q) + Q = 1.5 (quality factor) + """ + Q = 1.5 + omega = 1.5 * np.pi / M_cut + alpha = omega / (2.0 * Q) + + kappa = np.zeros(M_cut, dtype=np.float32) + for ell in range(M_cut): + kappa[ell] = np.exp(-alpha * ell) * np.cos(omega * ell) + kappa[0] = 1.0 # enforce anchor + return kappa + + +# ═══════════════════════════════════════════════════════════════════════ +# Args +# ═══════════════════════════════════════════════════════════════════════ + +def parse_args(): + p = argparse.ArgumentParser( + description="Two-block EM for Poisson GLM with memory kernel") + p.add_argument("--dataName", required=True) + p.add_argument("--basePath", + default="/pscratch/sd/b/balewski/2025_causalNet_tmp/") + + g = p.add_argument_group("model") + g.add_argument("--M_cut", type=int, default=10, + help="History depth (kernel length)") + g.add_argument("--eta_clip", type=float, default=5.0, + help="Log-rate clipping threshold") + + g = p.add_argument_group("EM structure") + g.add_argument("--num_em_iters", type=int, default=20, + help="Outer EM iterations") + g.add_argument("--block1_iter", type=int, default=50, + help="Gradient steps per Block 1 (kernel update)") + g.add_argument("--block2_epochs", type=int, default=100, + help="Adam epochs per Block 2 (network update)") + + g = p.add_argument_group("Block 1 (kernel)") + g.add_argument("--lr_kappa", type=float, default=0.01, + help="Learning rate for kappa update") + g.add_argument("--lambda2", type=float, default=1.0, + help="Kernel smoothness penalty weight") + + g = p.add_argument_group("Block 2 (network)") + g.add_argument("--lr_net", type=float, default=0.003, + help="Adam learning rate for A_diag, A_off, B") + g.add_argument("--lr_end_factor", type=float, default=0.1, + help="LR decays to lr_net * lr_end_factor") + g.add_argument("--lambda3", type=float, default=0.02, + help="L1 penalty on off-diagonal A_off") + g.add_argument("--rho_max", type=float, default=0.95, + help="Hard spectral radius ceiling for A_off") + g.add_argument("--rho_every", type=int, default=50, + help="Spectral projection frequency (batches)") + g.add_argument("--delay_em_iter_4_ArhoMax", type=int, default=3, + help="EM iter to start rho_max enforcement") + g.add_argument("--delay_em_iter_4_lrDecay", type=int, default=5, + help="EM iter to start LR decay") + g.add_argument("--delay_em_iter_4_Aprune", type=int, default=3, + help="EM iter to start L1 pruning") + g.add_argument("--batch_size", type=int, default=2048) + g.add_argument("--minW", type=float, default=0.01, + help="Threshold for edge counting (reporting)") + + g = p.add_argument_group("data") + g.add_argument("-T", "--time_range_sec", default=[0.0, 60.0], + nargs=2, type=float, + help="Time window [t0, t1] in seconds") + g.add_argument("--sample_stride", type=int, default=1, + help="Stride for sliding window in data preparation") + + g = p.add_argument_group("init") + + g = p.add_argument_group("misc") + g.add_argument("--seed", type=int, default=42) + g.add_argument("--fitName", type=str, default=None, + help="Output name; if omitted a random name is generated") + g.add_argument("-v", "--verb", type=int, default=1) + + return p.parse_args() + + +# ═══════════════════════════════════════════════════════════════════════ +# Main +# ═══════════════════════════════════════════════════════════════════════ + +def main(): + args = parse_args() + + rank, world_size, device = setup_distributed() + + inpPath = os.path.join(args.basePath, "truthDale") + outPath = os.path.join(args.basePath, "memKernFit") + + if is_rank0(rank): + assert os.path.exists(inpPath), f"missing inpPath: {inpPath}" + os.makedirs(outPath, exist_ok=True) + print(f"\nmemKern_EM_train args (world_size={world_size}):") + for arg in vars(args): + print(f" {arg}: {getattr(args, arg)}") + + torch.manual_seed(args.seed + rank) + np.random.seed(args.seed) # keep same on all ranks for init + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision("high") + + # ── Load data (rank 0 reads, then broadcasts) ──────────────── + if is_rank0(rank): + spikesFF = os.path.join(inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 0) + if args.verb > 1: + pprint(spikeMD) + spikes = np.asarray(spikeD["spikes"]) + assert spikes.ndim == 2, f"spikes must be (T, N); got {spikes.shape}" + single_rates = np.asarray(spikeD["single_rates"]) + + truthFF = os.path.join(inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 0) + else: + spikes = np.empty(0, dtype=np.uint8) + single_rates = np.empty(0, dtype=np.float64) + spikeMD = None + + # Broadcast spikes and single_rates to all ranks + if world_size > 1: + spikes = broadcast_array(spikes, rank, device, dtype=torch.uint8) + single_rates = broadcast_array(single_rates, rank, device, dtype=torch.float64) + # Broadcast spikeMD keys needed by all ranks + dt_t = torch.tensor(float(spikeMD["time_step_sec"]) if rank == 0 else 0.0, device=device) + if dist.is_initialized(): + dist.broadcast(dt_t, src=0) + dt = float(dt_t.item()) + else: + dt = float(spikeMD["time_step_sec"]) + + T_raw, N = spikes.shape + M_cut = args.M_cut + eta_clip = args.eta_clip + + # ── Time range selection ───────────────────────────────────── + t0_sec, t1_sec = float(args.time_range_sec[0]), float(args.time_range_sec[1]) + assert t0_sec < t1_sec, "time_range_sec[0] must be < time_range_sec[1]" + start_bin = max(0, int(math.floor(t0_sec / dt))) + end_bin = min(T_raw - 1, int(math.floor(t1_sec / dt))) + assert end_bin - start_bin >= M_cut + 1, \ + f"Time range too short: need at least {M_cut + 2} bins, got {end_bin - start_bin + 1}" + spikes = spikes[start_bin: end_bin + 1] + T_full = spikes.shape[0] + + if is_rank0(rank): + print(f"\nN={N} M_cut={M_cut} T={T_full} " + f"dt={dt} eta_clip={eta_clip}") + print(f"time=[{t0_sec:.1f}, {t1_sec:.1f}]s bins=[{start_bin}, {end_bin}]") + + # ── Build dataset (vectorized, GPU-resident) ───────────────── + # Each rank gets non-overlapping samples: + # rank r sees samples at t_indices starting at M_cut + r*stride, + # stepping by stride*world_size + effective_stride = args.sample_stride * world_size + rank_offset = rank * args.sample_stride + t_build = time.time() + Y_hist_np, Y_curr_np = build_spike_history( + spikes, M_cut, stride=effective_stride, offset=rank_offset + ) + N_samples_local = Y_hist_np.shape[0] + + # All ranks compute total + n_total_t = torch.tensor(float(N_samples_local), device=device) + if dist.is_initialized(): + dist.all_reduce(n_total_t, op=dist.ReduceOp.SUM) + N_samples_total = int(n_total_t.item()) + + if is_rank0(rank): + print(f"N_samples_total={N_samples_total} per_rank={N_samples_local} " + f"stride={args.sample_stride} world_size={world_size} " + f"effective_stride={effective_stride} build={time.time()-t_build:.1f}s") + + # Move entire dataset to GPU — no DataLoader needed + Y_hist_gpu = torch.tensor(Y_hist_np, dtype=torch.float32, device=device) + Y_curr_gpu = torch.tensor(Y_curr_np, dtype=torch.float32, device=device) + del Y_hist_np, Y_curr_np # free CPU memory + + # ── Initialize parameters (same on all ranks via same np seed) ─ + del spikes # free CPU memory, no longer needed + A_diag_init = init_A_diag_from_spikes(N, dt) + A_off_init = init_A_off_from_spikes(Y_hist_gpu, Y_curr_gpu) + B_init = init_B_from_spikes(Y_curr_gpu, dt) + kappa_init = init_kappa_damped_oscillator(M_cut) + + if is_rank0(rank): + print(f"\nInitialization:") + print(f" A_diag range: [{A_diag_init.min():.4f}, {A_diag_init.max():.4f}]") + print(f" A_off range: [{A_off_init.min():.4f}, {A_off_init.max():.4f}]") + print(f" B range: [{B_init.min():.4f}, {B_init.max():.4f}]") + print(f" kappa: {kappa_init}") + + # Move to GPU + A_diag = torch.tensor(A_diag_init, dtype=torch.float32, device=device) + A_off = torch.tensor(A_off_init, dtype=torch.float32, device=device) + B = torch.tensor(B_init, dtype=torch.float32, device=device) + kappa = torch.tensor(kappa_init, dtype=torch.float32, device=device) + + # Save copies of initial state + A_diag_init_save = A_diag_init.copy() + A_off_init_save = A_off_init.copy() + B_init_save = B_init.copy() + kappa_init_save = kappa_init.copy() + + # ── Build NetworkModel for Block 2 ─────────────────────────── + net_model = NetworkModel(N, A_diag, A_off, B).to(device) + + # Wrap in DDP for automatic gradient synchronization in Block 2 + if world_size > 1: + ddp_model = DDP(net_model, device_ids=[rank]) + else: + ddp_model = net_model + # Access underlying model for parameter projections + raw_model = ddp_model.module if world_size > 1 else ddp_model + + # Compile for kernel fusion (reduces launch overhead for small ops) + ddp_model = torch.compile(ddp_model) + + # Optimizer and scheduler for Block 2 + total_b2_epochs = args.num_em_iters * args.block2_epochs + decay_start = args.delay_em_iter_4_lrDecay * args.block2_epochs + decay_epochs = max(1, total_b2_epochs - decay_start) + optimizer = optim.Adam(ddp_model.parameters(), lr=args.lr_net, fused=True) + scheduler = optim.lr_scheduler.LinearLR( + optimizer, start_factor=1.0, end_factor=args.lr_end_factor, + total_iters=decay_epochs, last_epoch=-1 + ) + scheduler._step_count = 1 + + # ── History tracking ───────────────────────────────────────── + h_b1_nll = [] + h_kappa = [] + + # Monitoring histories (per global M-epoch) + h_m_loss = [] + h_m_nll = [] + h_m_l1 = [] + h_rho = [] + h_nz = [] + h_lr = [] + + h_kappa.append(kappa_init.copy()) + + off_mask = ~torch.eye(N, dtype=torch.bool, device=device) + + t_start = time.time() + b2_epoch_global = 0 + curr_lr_kappa = args.lr_kappa + + # ══════════════════════════════════════════════════════════════ + # EM loop + # ══════════════════════════════════════════════════════════════ + for em in range(1, args.num_em_iters + 1): + apply_prune = em > args.delay_em_iter_4_Aprune + apply_rho = em > args.delay_em_iter_4_ArhoMax + + # Update lr_kappa using same LinearLR schedule logic + if em > args.delay_em_iter_4_lrDecay: + # Fraction into the decay period (0 to 1) + total_decay_steps = args.num_em_iters - args.delay_em_iter_4_lrDecay + step_in_decay = em - args.delay_em_iter_4_lrDecay + factor = 1.0 - (1.0 - args.lr_end_factor) * (step_in_decay / total_decay_steps) + curr_lr_kappa = args.lr_kappa * factor + + if is_rank0(rank): + print(f"\nEM {em:3d}/{args.num_em_iters:d} " + f"lr_kappa={curr_lr_kappa:.2e} " + f"lr_net={optimizer.param_groups[0]['lr']:.2e} " + f"prune={apply_prune} rho={apply_rho}") + + # ── Block 1: kernel update ─────────────────────────────── + tb1_start = time.time() + # Use current A_diag, A_off, B from net_model (detached) + with torch.no_grad(): + A_diag_fixed = raw_model.A_diag.data.clone() + A_off_fixed = raw_model.A_off.data.clone() + # Zero diagonal of A_off + idx = torch.arange(N, device=device) + A_off_fixed[idx, idx] = 0.0 + B_fixed = raw_model.B.data.clone() + + kappa, b1_nll = run_block1_kernel_update( + Y_hist_gpu, Y_curr_gpu, + A_diag_fixed, A_off_fixed, B_fixed, + kappa, dt, eta_clip, args.lambda2, + curr_lr_kappa, args.block1_iter, M_cut, + rank, world_size + ) + tb1 = time.time() - tb1_start + h_b1_nll.append(float(b1_nll)) + h_kappa.append(kappa.cpu().numpy().copy()) + + # ── Block 2: network update ────────────────────────────── + tb2_start = time.time() + + # Fix kappa for this block — precompute S and Y_prev ONCE + kappa_fixed = kappa.detach().clone() + with torch.no_grad(): + S_full = compute_S(Y_hist_gpu, kappa_fixed) # (N_local, N) + Y_prev_full = Y_hist_gpu[:, 0, :].contiguous() # (N_local, N) + + BS = args.batch_size + n_batches = (N_samples_local + BS - 1) // BS + log_dt = math.log(dt) + + for ep in range(args.block2_epochs): + b2_epoch_global += 1 + # Accumulate losses on GPU to avoid CPU-GPU sync per batch + acc_nll = torch.tensor(0.0, device=device) + acc_l1 = torch.tensor(0.0, device=device) + acc_loss = torch.tensor(0.0, device=device) + + # Shuffle indices on GPU + perm = torch.randperm(N_samples_local, device=device) + + for bi in range(n_batches): + i0 = bi * BS + i1 = min(i0 + BS, N_samples_local) + idx = perm[i0:i1] + + Y_prev_b = Y_prev_full[idx] + S_b = S_full[idx] + Y_curr_b = Y_curr_gpu[idx] + + optimizer.zero_grad(set_to_none=True) + mu, eta_c = ddp_model(Y_prev_b, S_b, dt, eta_clip) + nll = poisson_nll(mu, eta_c, Y_curr_b, log_dt) / Y_curr_b.shape[0] + + if args.lambda3 > 0: + l1_term = args.lambda3 * offdiag_l1_mean(raw_model.A_off, N) + else: + l1_term = torch.tensor(0.0, device=device) + + loss = nll + l1_term + loss.backward() + optimizer.step() + + # Post-step projections (on raw model params) + enforce_zero_diagonal_(raw_model.A_off.data, N) + if apply_prune and args.lambda3 > 0: + offdiag_soft_threshold_( + raw_model.A_off.data, + optimizer.param_groups[0]["lr"], + args.lambda3, N + ) + if apply_rho and (bi % args.rho_every == 0): + enforce_spectral_radius_(raw_model.A_off.data, args.rho_max, N) + + # Accumulate on GPU — no .item() sync + with torch.no_grad(): + acc_nll += nll.detach() + acc_l1 += l1_term.detach() + acc_loss += loss.detach() + + if em > args.delay_em_iter_4_lrDecay: + scheduler.step() + + # Record metrics ONCE per EM iteration (not per epoch) — single sync + nb = max(1, n_batches) + with torch.no_grad(): + rho = float(torch.linalg.eigvals(raw_model.A_off.data).abs().max().item()) + nz = int((raw_model.A_off.abs() > args.minW).sum().item()) + last_nll = float(acc_nll.item()) / nb + last_l1 = float(acc_l1.item()) / nb + last_loss = float(acc_loss.item()) / nb + + h_m_loss.append(last_loss) + h_m_nll.append(last_nll) + h_m_l1.append(last_l1) + h_rho.append(rho) + h_nz.append(nz) + h_lr.append(float(optimizer.param_groups[0]["lr"])) + + tb2 = time.time() - tb2_start + met = dict(nll=h_m_nll[-1], l1=h_m_l1[-1], rho=h_rho[-1], nz=h_nz[-1]) + + # ── Re-synchronize parameters from rank 0 ──────────────── + # Prevents drift from non-gradient ops (pruning, spectral projection) + if dist.is_initialized(): + for p in raw_model.parameters(): + dist.broadcast(p.data, src=0) + + if is_rank0(rank) and args.verb > 0: + n_off = int(off_mask.sum().item()) + sp = 1.0 - met["nz"] / max(1, n_off) + print( + f"EM {em:3d}/{args.num_em_iters} " + f"B1_nll={b1_nll:.4e}({tb1:.1f}s) " + f"B2_nll={met['nll']:.4e} l1={met['l1']:.4e} " + f"rho={met['rho']:.4f} sp={sp:.3f} nz={met['nz']} " + f"lr={h_lr[-1]:.2e} ({tb2:.1f}s) " + f"kappa_norm={float(kappa.abs().sum()):.3f} " + f"tot={time.time() - t_start:.0f}s" + ) + + # ══════════════════════════════════════════════════════════════ + # Save results (rank 0 only) + # ══════════════════════════════════════════════════════════════ + if is_rank0(rank): + print(f"\nTotal EM time: {time.time() - t_start:.1f}s") + + A_diag_hat = raw_model.A_diag.detach().cpu().numpy() + A_off_hat = raw_model.A_off.detach().cpu().numpy() + np.fill_diagonal(A_off_hat, 0.0) + B_hat = raw_model.B.detach().cpu().numpy() + kappa_hat = kappa.detach().cpu().numpy() + + if args.fitName is None: + h6 = secrets.token_hex(3) + outF = f"{args.dataName}-memKern-{h6}" + else: + outF = args.fitName + outFF = os.path.join(outPath, f"{outF}.memKernEM.npz") + + outD = { + "A_diag_init": A_diag_init_save, + "A_diag_hat": A_diag_hat.astype(np.float32), + "A_off_init": A_off_init_save, + "A_off_hat": A_off_hat.astype(np.float32), + "B_init": B_init_save, + "B_hat": B_hat.astype(np.float32), + "kappa_init": kappa_init_save, + "kappa_hat": kappa_hat.astype(np.float32), + "kappa_history": np.array(h_kappa, dtype=np.float32), + "single_rates": single_rates.astype(np.float32), + "e_nll_em": np.array(h_b1_nll, dtype=np.float64), + "m_loss_epoch": np.array(h_m_loss, dtype=np.float64), + "m_nll_epoch": np.array(h_m_nll, dtype=np.float64), + "m_l1_epoch": np.array(h_m_l1, dtype=np.float64), + "rho_epoch": np.array(h_rho, dtype=np.float64), + "nz_edges_epoch": np.array(h_nz, dtype=np.int64), + "learning_rates": np.array(h_lr, dtype=np.float64), + } + + spikeMD.pop('short_name') + outMD = {'spike_gen':spikeMD} + outMD["provenance"]={"memKernEM_file": outF, "spiksData_file": args.dataName} + outMD['short_name']=outF + + outMD["fit_type"] = "memKernEM" + outMD["train"] = { + "num_em_iters": args.num_em_iters, + "sample_stride": args.sample_stride, + "m_epochs": args.block2_epochs, + "block1_iter": args.block1_iter, + "block2_epochs": args.block2_epochs, + "lr_kappa": args.lr_kappa, + "lr_net": args.lr_net, + "lr_end_factor": args.lr_end_factor, + "lambda2": args.lambda2, + "lambda3": args.lambda3, + "rho_max": args.rho_max, + "rho_every": args.rho_every, + "delay_em_iter_4_ArhoMax": args.delay_em_iter_4_ArhoMax, + "delay_em_iter_4_lrDecay": args.delay_em_iter_4_lrDecay, + "delay_em_iter_4_Aprune": args.delay_em_iter_4_Aprune, + "batch_size": args.batch_size, + "minW": args.minW, + "M_cut": M_cut, + "num_states": 1, + "num_neurons": N, + "num_time_bins": T_full, + "num_samples": N_samples_total, + "num_samples_per_rank": N_samples_local, + "world_size": world_size, + "time_step_sec": dt, + "eta_clip": eta_clip, + "time_range_sec": [t0_sec, t1_sec], + "time_range_bins": [start_bin, end_bin], + "seed": args.seed, + "training_time_sec": round(time.time() - t_start, 1), + } + + if args.verb > 1: pprint(outMD) + write_data_npz(outD, outFF, metaD=outMD) + print(f"\nSaved: {outFF}") + print(f" basePath={args.basePath}") + print(f" ./memKern_EM_eval4.py --basePath {args.basePath} --dataName {outF} -p b a \n") + + cleanup_distributed() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver4_lagM_offDiag/movie_spikesTrain4.py b/causal_net/nonStation_ver4_lagM_offDiag/movie_spikesTrain4.py new file mode 100755 index 00000000..85ec071a --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/movie_spikesTrain4.py @@ -0,0 +1,547 @@ +#!/usr/bin/env python3 +""" +Spike train movie generator — uses the same ``--basePath`` / ``--dataName`` data layout as view_spikesTrain4.py. + +Loads: + /truthDale/.simTruth.npz → node_positions, node_is_inhibitory + /truthDale/.spikes.npz → spikes (T, N) uint8 + +via toolbox.Util_NumpyIO.read_data_npz (same as view_spikesTrain4.py). + +MP4 requires ``imageio-ffmpeg`` (see ``container/ubu24-python.dockerfile``). +Also writes ``_frame0.png`` (first movie frame, same DPI as MP4). +With ``--stillFrames N`` (default 20), also writes ``_still.html``: +first N frames as embedded PNGs; **click the image** to advance (no autoplay). +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import argparse +import base64 +import io +import json +import os +import sys +import numpy as np +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.animation as animation +from matplotlib.collections import LineCollection +from matplotlib.patches import Ellipse +from matplotlib.ticker import MaxNLocator +import imageio_ffmpeg + +MP4_QUALITY_PRESETS = { + "low": {"dpi": 140, "bitrate": 3000, "crf": 26}, + "medium": {"dpi": 220, "bitrate": 8000, "crf": 18}, + "high": {"dpi": 260, "bitrate": 16000, "crf": 15}, +} +MP4_QUALITY_TAG = {"low": "l", "medium": "m", "high": "h"} + + +def get_parser(): + parser = argparse.ArgumentParser(description="Generate spike movie from Dale simTruth + spikes (same paths as view_spikesTrain4)") + parser.add_argument("-v", "--verbosity", type=int, help="increase output verbosity", default=1, dest="verb") + parser.add_argument("--basePath", default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input data") + parser.add_argument("--dataName", default=None, help="simulated Dale network base name") + parser.add_argument("-T", "--time_range_sec", default=[0.0, 10.0], nargs=2, type=float, help="time range (seconds) to render") + parser.add_argument("-r", "--time_rebin2", default=1, type=int, help="rebin factor on time axis (same meaning as view_spikesTrain4)") + # Movie-only + parser.add_argument("--output", type=str, default=None, help="Output path; must end with .mp4 (default: /_spike_movie.mp4)") + parser.add_argument("--tau", type=int, default=20, help="Frames for spike glow decay (default: 20)") + parser.add_argument("--fps", type=int, default=60, help="Frames per second (default: 60)") + parser.add_argument("--speed", type=float, default=1.0, help="Playback speed factor in [0.2, 1.0]; lower is slower (default: 1.0)") + parser.add_argument("--flushSize", type=float, default=0.05, help="Max spike-flash circle radius at g=1 (plot units); 0=auto 0.45×min inter-neuron distance (default: 0)") + parser.add_argument("--marker_size", type=float, default=0.02, help="Neuron marker size scale; 0=auto (default: 0)") + parser.add_argument("--mp4quality", choices=sorted(MP4_QUALITY_PRESETS), default="medium", help="MP4 quality preset controlling dpi, bitrate, and crf (default: medium)") + parser.add_argument("--stillFrames", type=int, default=40, help="Also write _still.html: first N frames, click image to advance (0=skip)") + parser.add_argument("--edgePrescale", type=int, default=10, help="Show a neuron's outgoing edges on every Nth spike event for that neuron (default: 10)") + + args = parser.parse_args() + preset = MP4_QUALITY_PRESETS[args.mp4quality] + args.dpi = preset["dpi"] + args.mp4Bitrate = preset["bitrate"] + args.mp4Crf = preset["crf"] + args.inpPath = os.path.join(args.basePath, "truthDale") + args.outPath = os.path.join(args.basePath, "plots") + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + if args.time_range_sec is not None: + assert args.time_range_sec[0] < args.time_range_sec[1] + assert os.path.exists(args.basePath), args.basePath + assert os.path.exists(args.inpPath), f"missing inpPath: {args.inpPath}" + if args.dataName is None: + print("ERROR: --dataName is required", file=sys.stderr) + sys.exit(1) + if args.flushSize < 0: + print("ERROR: --flushSize must be >= 0 (0 = auto radius)", file=sys.stderr) + sys.exit(1) + if args.stillFrames < 0: + print("ERROR: --stillFrames must be >= 0", file=sys.stderr) + sys.exit(1) + if args.edgePrescale < 1: + print("ERROR: --edgePrescale must be >= 1", file=sys.stderr) + sys.exit(1) + if not (0.2 <= args.speed <= 1.0): + print("ERROR: --speed must be in [0.2, 1.0]", file=sys.stderr) + sys.exit(1) + _out = args.output if args.output else os.path.join( + args.outPath, f"{args.dataName}_spike_movie_q-{MP4_QUALITY_TAG[args.mp4quality]}.mp4" + ) + if not _out.endswith(".mp4"): + print("ERROR: --output must be a .mp4 path", file=sys.stderr) + sys.exit(1) + args.movie_output_path = _out + return args + + +def rebin_spike_rates(spikeYield, md, tReb2): + """Rebin 2D spike train (time x neurons); copied from view_spikesTrain4.py.""" + assert spikeYield.ndim == 2, f"Expected 2D spike train, got shape={spikeYield.shape}" + assert tReb2 < 101 + + ntime, nchan = spikeYield.shape + if ntime % tReb2 != 0: + ntime_c = ntime - (ntime % tReb2) + spikeYield = spikeYield[:ntime_c] + spikeYieldR = np.sum(spikeYield.reshape(-1, tReb2, nchan), axis=1) + time_step = md["time_step_sec"] + time_step2 = time_step * tReb2 + rate2D = spikeYieldR / time_step2 + pop_spike_count = np.sum(spikeYieldR, axis=1) + pop_rate_hz = pop_spike_count / time_step2 + rebD = {"time_step2": time_step2} + rebD["rate2D"] = rate2D + rebD["pop_spike_count"] = pop_spike_count + rebD["pop_rate_hz"] = pop_rate_hz + ntime //= tReb2 + rebD["timeV"] = np.linspace(0, (ntime - 1) * time_step2, ntime) + rebD["spike_counts"] = spikeYieldR + print("rebinned: nchan=%d dt=%.3f sec rebin=%d" % (nchan, time_step2, tReb2)) + return rebD + + +def movie_title_header(spikeMD, trueMD, nchan, tbin_sec): + """Metadata line matching PlotterSpikesTrain.freq_vs_time (dataset, R, nchan, Tbin, …).""" + dale = trueMD.get("dale_conf", {}) + + def md_get_float(key, *, required=False): + if key in spikeMD: + return float(spikeMD[key]) + if key in dale: + return float(dale[key]) + if required: + raise KeyError(f"Missing metadata field: {key}") + return None + + R_sel = md_get_float("spectral_radius", required=True) + sn = spikeMD.get("short_name", trueMD.get("short_name", "unknown")) + tit0 = "dataset: %s R=%.3f nchan=%d Tbin=%.2f sec" % (sn, R_sel, nchan, tbin_sec) + mem_q = md_get_float("mem_Q") + mem_tau = md_get_float("mem_tau") + if spikeMD.get("spike_model") == "B" and mem_q is not None and mem_tau is not None: + tit0 += " Q=%.4g tau/sec=%.4g" % ( + mem_q, + mem_tau, + ) + ph = md_get_float("placement_H", required=True) + pl = md_get_float("placement_L", required=True) + tit0 += " placement HxL=(%gx%g)" % ( + ph, + pl, + ) + return tit0 + + +def write_still_click_html(fig, update, n_frames, dpi, out_path): + """Write self-contained HTML: embedded PNGs, click image to cycle frames (no autoplay).""" + chunks = [] + for i in range(n_frames): + update(i) + fig.canvas.draw() + buf = io.BytesIO() + fig.savefig(buf, format="png", dpi=dpi, bbox_inches=None) + chunks.append(base64.standard_b64encode(buf.getvalue()).decode("ascii")) + b64_json = json.dumps(chunks) + html = ( + "\nstill frames\n" + "\n" + "

\n" + "\"frame\"\n" + "\n\n" + ) + with open(out_path, "w", encoding="utf-8") as f: + f.write(html) + + +def main(): + from toolbox.Util_NumpyIO import read_data_npz + + args = get_parser() + np.set_printoptions(precision=3) + + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + assert os.path.isfile(spikesFF), f"missing {spikesFF}" + assert os.path.isfile(truthFF), f"missing {truthFF}" + + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 0) + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 0) + if args.verb > 1: + from pprint import pprint + + print("\nspikeMD:"); pprint(spikeMD) + print("\ntrueMD:"); pprint(trueMD) + + positions = np.asarray(trueD["node_positions"]) + is_inhib = np.asarray(trueD["node_is_inhibitory"]) + edge_sign = np.asarray(trueD["E_true"]) if "E_true" in trueD else None + spikes = np.asarray(spikeD["spikes"]) + if spikes.ndim != 2: + raise ValueError( + "spikes must have shape (T, N); got %s" % (spikes.shape,) + ) + + dt = float(spikeMD["time_step_sec"]) + ntime, nchan = spikes.shape + if positions.shape[0] != nchan or is_inhib.shape[0] != nchan: + raise ValueError( + "N mismatch: spikes N=%d vs positions %s vs inhib %s" + % (nchan, positions.shape, is_inhib.shape) + ) + if edge_sign is not None and edge_sign.shape != (nchan, nchan): + raise ValueError("E_true must have shape (%d, %d)" % (nchan, nchan)) + + t0s, t1s = float(args.time_range_sec[0]), float(args.time_range_sec[1]) + i0 = max(0, int(np.floor(t0s / dt))) + i1 = min(ntime, int(np.ceil(t1s / dt))) + if i0 >= i1: + raise ValueError("empty time window: indices [%d,%d) from time_range_sec" % (i0, i1)) + spikes = spikes[i0:i1, :].astype(np.float64, copy=False) + t_off_sec = float(i0) * dt + + t_reb = int(args.time_rebin2) + if t_reb < 1: + raise ValueError("time_rebin2 must be >= 1") + if t_reb > 1: + rebD = rebin_spike_rates(spikes, spikeMD, t_reb) + spikes = rebD["spike_counts"].astype(np.float64, copy=False) + frame_dt_sec = float(rebD["time_step2"]) + else: + frame_dt_sec = dt + + title_header = movie_title_header(spikeMD, trueMD, nchan, frame_dt_sec) + + # ---- Load data (movie core) ---- + N = positions.shape[0] + T_total = spikes.shape[0] + + t_start = 0 + t_end = T_total + T_render = t_end - t_start + + tau = args.tau + export_fps = float(args.fps) + output_frame_idx = np.floor(np.arange(int(np.ceil(T_render / args.speed))) * args.speed).astype(np.int32) + output_frame_idx = np.clip(output_frame_idx, 0, T_render - 1) + T_export = int(output_frame_idx.shape[0]) + print(f"Neurons: {N}, frames: {T_total} (time slice + rebin), frame_dt={frame_dt_sec:.4g}s") + print(f"Rendering frames {t_start} to {t_end-1} ({T_render} frames)") + print(f"Tau (decay frames): {tau}, FPS: {args.fps}, speed={args.speed:.3g}, export_fps={export_fps:.3f}, export_frames={T_export}") + + x = positions[:, 0] + y = positions[:, 1] + + if N > 1: + from scipy.spatial import distance as sp_dist + + dists = sp_dist.pdist(positions) + min_dist = float(np.min(dists)) + else: + min_dist = 1.0 + + if float(args.flushSize) > 0: + max_circle_r = float(args.flushSize) + else: + max_circle_r = min_dist * 0.45 + + # Axis span is L+2*flush and H+2*flush (one flush radius per side around [0,L]×[0,H]). + dale = trueMD.get("dale_conf", {}) + pl = float(spikeMD.get("placement_L", dale["placement_L"])) + ph = float(spikeMD.get("placement_H", dale["placement_H"])) + pad = max_circle_r + xlim_lo, xlim_hi = -pad, pl + pad + ylim_lo, ylim_hi = -pad, ph + pad + w_data = xlim_hi - xlim_lo + h_data = ylim_hi - ylim_lo + if args.verb > 0: + print( + "Axis limits: placement [−r,L+r]×[−r,H+r] (window %.4g × %.4g, pad=flushSize=%.4g)" + % (w_data, h_data, pad) + ) + + marker_size = args.marker_size if args.marker_size > 0 else min_dist * 0.18 + + print(f"Min neuron distance: {min_dist:.4f}") + print( + f"Marker size: {marker_size:.4f}, max flash radius (flushSize): {max_circle_r:.4f}" + + ( + f" [explicit --flushSize {args.flushSize}]" + if float(args.flushSize) > 0 + else " [auto from min_dist]" + ) + ) + print(f"Edge display prescale: every {args.edgePrescale}th spike event per neuron") + + print("Precomputing glow intensities...") + glow = np.zeros((T_render, N), dtype=np.float32) + time_since_spike = np.full(N, tau + 1, dtype=np.int32) + + pre_start = max(0, t_start - tau) + for t in range(pre_start, t_start): + firing = spikes[t] > 0 + time_since_spike[firing] = 0 + time_since_spike[~firing] += 1 + + for frame in range(T_render): + t = t_start + frame + firing = spikes[t] > 0 + time_since_spike[firing] = 0 + active = time_since_spike <= tau + glow[frame, active] = 1.0 - time_since_spike[active] / tau + time_since_spike[~firing] += 1 + + fig, ax = plt.subplots(1, 1, figsize=(16.0, 5.0)) + ax.set_xlim(xlim_lo, xlim_hi) + ax.set_ylim(ylim_lo, ylim_hi) + ax.set_facecolor("white") + ax.set_xlabel("x", labelpad=12) + ax.set_ylabel("y", labelpad=12) + ax.tick_params(axis="x", which="major", pad=8) + ax.tick_params(axis="y", which="major", pad=6) + ax.xaxis.set_major_locator(MaxNLocator(integer=True)) + ax.yaxis.set_major_locator(MaxNLocator(integer=True)) + + title_text = ax.set_title( + title_header + f" t = {t_off_sec:.3f}s (frame 0)", + fontsize=12, + pad=10, + ) + + exc_mask = is_inhib == 0 + inh_mask = is_inhib == 1 + exc_fill = "#f4a3a3" + exc_edge = "#8b3a3a" + inh_fill = "#a9c8ff" + inh_edge = "#355c99" + exc_flush = "#ff5c5c" + inh_flush = "#4d8dff" + + ax.scatter( + x[exc_mask], + y[exc_mask], + marker="^", + s=marker_size * 800, + facecolors=exc_fill, + edgecolors=exc_edge, + linewidths=0.8, + zorder=3, + ) + ax.scatter( + x[inh_mask], + y[inh_mask], + marker="s", + s=marker_size * 600, + facecolors=inh_fill, + edgecolors=inh_edge, + linewidths=0.8, + zorder=3, + ) + + exc_segments_by_src = [[] for _ in range(N)] + inh_segments_by_src = [[] for _ in range(N)] + if edge_sign is not None: + edge_mask = edge_sign != 0 + for j in np.flatnonzero(exc_mask): + for i in np.flatnonzero(edge_mask[:, j]): + if i == j: + continue + exc_segments_by_src[j].append(np.array([positions[j], positions[i]], dtype=np.float64)) + for j in np.flatnonzero(inh_mask): + for i in np.flatnonzero(edge_mask[:, j]): + if i == j: + continue + inh_segments_by_src[j].append(np.array([positions[j], positions[i]], dtype=np.float64)) + exc_edge_collection = LineCollection( + [], + colors=exc_edge, + linewidths=1.0, + alpha=0.45, + zorder=2, + ) + inh_edge_collection = LineCollection( + [], + colors=inh_edge, + linewidths=1.0, + alpha=0.45, + zorder=2, + ) + ax.add_collection(exc_edge_collection) + ax.add_collection(inh_edge_collection) + + # Shrink default side margins so the axes use more of the frame (wide L×H plots were ~70% width). + fig.subplots_adjust(left=0.065, right=0.995, bottom=0.13, top=0.86) + + # Ellipse in data coords with height/width = sx/sy so the patch is a circle in pixels (non-equal data aspect). + fig.set_dpi(args.dpi) + fig.canvas.draw() + _bbox = ax.get_window_extent() + _dx = xlim_hi - xlim_lo + _dy = ylim_hi - ylim_lo + flush_sx_over_sy = (_bbox.width / _dx) / (_bbox.height / _dy) + + circles = [] + for i in range(N): + circle_face = inh_flush if inh_mask[i] else exc_flush + c = Ellipse( + (x[i], y[i]), + 0.0, + 0.0, + angle=0.0, + facecolor=circle_face, + edgecolor="none", + linewidth=0, + alpha=0.0, + zorder=5, + ) + ax.add_patch(c) + circles.append(c) + + if args.verb > 0: + fw, fh = fig.get_size_inches() + print( + "figure: %.2f×%.2f in (data window %.4g×%.4g)" + % (fw, fh, w_data, h_data) + ) + + _progress_seen = set() + edge_spike_counts = np.zeros(N, dtype=np.int32) + last_src_frame = None + current_exc_segments = [] + current_inh_segments = [] + + def update(out_frame): + nonlocal last_src_frame, current_exc_segments, current_inh_segments + src_frame = int(output_frame_idx[out_frame]) + if out_frame % 200 == 0 and out_frame not in _progress_seen: + _progress_seen.add(out_frame) + print(f" Rendering output frame {out_frame}/{T_export} from source frame {src_frame}/{T_render}...") + t_sec = t_off_sec + (t_start + src_frame) * frame_dt_sec + title_text.set_text(title_header + f" t = {t_sec:.3f}s (frame {src_frame})") + if src_frame != last_src_frame: + firing = spikes[t_start + src_frame] > 0 + current_exc_segments = [] + current_inh_segments = [] + firing_idx = np.flatnonzero(firing) + if firing_idx.size > 0: + edge_spike_counts[firing_idx] += 1 + for j in np.flatnonzero(firing & exc_mask): + if edge_spike_counts[j] % args.edgePrescale == 0: + current_exc_segments.extend(exc_segments_by_src[j]) + for j in np.flatnonzero(firing & inh_mask): + if edge_spike_counts[j] % args.edgePrescale == 0: + current_inh_segments.extend(inh_segments_by_src[j]) + last_src_frame = src_frame + exc_edge_collection.set_segments(current_exc_segments) + inh_edge_collection.set_segments(current_inh_segments) + for i in range(N): + g = glow[src_frame, i] + if g > 0: + r = max_circle_r * g + circles[i].set_width(2.0 * r) + circles[i].set_height(2.0 * r * flush_sx_over_sy) + circles[i].set_alpha(0.7 * g) + circles[i].set_edgecolor("none") + circles[i].set_linewidth(0) + circles[i].set_visible(True) + else: + circles[i].set_visible(False) + return circles + [exc_edge_collection, inh_edge_collection] + + output = args.movie_output_path + _out_dir = os.path.dirname(os.path.abspath(output)) + assert os.path.isdir(_out_dir), "output directory must exist: %r" % _out_dir + + print(f"Creating animation with {T_render} frames -> {output}") + # blit=True breaks growing flush Ellipse patches (bbox stays ~0); green flashes vanish. + anim = animation.FuncAnimation( + fig, update, frames=T_export, interval=1000 / export_fps, blit=False + ) + + png_frame0 = os.path.splitext(output)[0] + "_frame0.png" + update(0) + fig.canvas.draw() + fig.savefig(png_frame0, format="png", dpi=args.dpi, bbox_inches=None) + print(f"Saved first frame to {png_frame0}") + + ff = imageio_ffmpeg.get_ffmpeg_exe() + if not ff or not os.path.isfile(ff) or not os.access(ff, os.X_OK): + raise RuntimeError( + "imageio_ffmpeg.get_ffmpeg_exe() is missing or not executable: %r" % (ff,) + ) + plt.rcParams["animation.ffmpeg_path"] = ff + if args.verb > 0: + print(f"Using ffmpeg: {ff}") + print(f"MP4 quality: dpi={args.dpi}, bitrate={args.mp4Bitrate} kbps, crf={args.mp4Crf}, speed={args.speed:.3g}, export_fps={export_fps:.3f}, export_frames={T_export}") + + ffmpeg_args = [ + "-vcodec", "libx264", + "-crf", str(args.mp4Crf), + "-preset", "slow", + "-pix_fmt", "yuv420p", + ] + writer = animation.FFMpegWriter( + fps=export_fps, bitrate=args.mp4Bitrate, extra_args=ffmpeg_args + ) + anim.save(output, writer=writer, dpi=args.dpi) + print(f"Saved movie to {output}") + + if args.stillFrames > 0: + n_still = min(int(args.stillFrames), T_render) + html_out = os.path.splitext(output)[0] + "_still.html" + print(f"Creating still viewer ({n_still} frames, click to advance) -> {html_out}") + write_still_click_html(fig, update, n_still, args.dpi, html_out) + print(f"Saved still HTML to {html_out}") + + plt.close(fig) + print("M:done") + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver4_lagM_offDiag/toolbox b/causal_net/nonStation_ver4_lagM_offDiag/toolbox new file mode 120000 index 00000000..ac05dc12 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/toolbox @@ -0,0 +1 @@ +../toolbox \ No newline at end of file diff --git a/causal_net/nonStation_ver4_lagM_offDiag/topoAna_daleMatrix4.py b/causal_net/nonStation_ver4_lagM_offDiag/topoAna_daleMatrix4.py new file mode 100644 index 00000000..3947d9fd --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/topoAna_daleMatrix4.py @@ -0,0 +1,304 @@ +#!/usr/bin/env python3 +""" +Topology analysis for simulated Dale Poisson network data. + +Loads simTruth.npz from gen_daleMatrices4.py. From the weighted off-diagonal +matrix A_off_true, builds the binary adjacency G (doc/network_ver4_topoDiscovery.tex, +Eq. adjacency): G_ij = 1 iff i != j and A_off_ij != 0 (directed edge +presynaptic j -> postsynaptic i). + +Method 1 — degree assortativity (Newman, directed): Pearson r between +k_i^in at the target and k_j^out at the source over all edges (Eq. assortativity). + +Usage: + ./topoAna_daleMatrix4.py --basePath $basePath --dataName daleN200_fa3733 +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os +from pprint import pprint +import numpy as np +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +import argparse + + +def adjacency_from_A_off(A_off, atol=0.0): + """ + Binary directed adjacency G: G_ij = 1 if |A_off_ij| > atol and i != j. + Rows index postsynaptic i, columns presynaptic j (matches gen_daleMatrices4 / TeX). + """ + A = np.asarray(A_off, dtype=np.float64) + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError("A_off must be square (N, N)") + G = (np.abs(A) > atol).astype(np.float64) + np.fill_diagonal(G, 0.0) + return G + + +def shuffle_neuron_order(A_off, perm=None): + """ + Reindex neurons with a shared permutation on rows and columns. + Rows are shuffled first; columns are reordered to the same new neuron order. + Returns (A_off_shuffled, perm). + """ + A = np.asarray(A_off, dtype=np.float64) + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError("A_off must be square (N, N)") + + n = A.shape[0] + if perm is None: + perm = np.random.permutation(n) + else: + perm = np.asarray(perm) + if perm.shape != (n,): + raise ValueError("perm must have shape (N,)") + return A[perm][:, perm], perm + + +def directed_degree_assortativity(G): + """ + Directed degree assortativity r per network_ver4_topoDiscovery.tex, Sec. 1. + For each edge j -> i (G_ij = 1): pair (k_i^in, k_j^out). + Uses |E|^{-1} moments (population cov / product of population std devs). + Returns (r, |E|, extra dict). r is nan if |E| < 2 or a std dev is 0. + """ + G = np.asarray(G, dtype=np.float64) + n = G.shape[0] + k_in = G.sum(axis=1) + k_out = G.sum(axis=0) + ei, ej = np.nonzero(G > 0) + m = int(ei.size) + if m < 2: + return float("nan"), m, {"n_nodes": n, "n_edges": m, "reason": "|E| < 2"} + + x = k_in[ei].astype(np.float64) # Target in-degrees + y = k_in[ej].astype(np.float64) # Source in-degrees (In-In correlation) + mx = x.mean() + my = y.mean() + cov = (x * y).mean() - mx * my + vx = ((x - mx) ** 2).mean() + vy = ((y - my) ** 2).mean() + sig_in_target = float(np.sqrt(vx)) + sig_in_source = float(np.sqrt(vy)) + if sig_in_target <= 0.0 or sig_in_source <= 0.0: + return float("nan"), m, { + "n_nodes": n, + "n_edges": m, + "sigma_in_target": sig_in_target, + "sigma_in_source": sig_in_source, + "reason": "zero std on edge multiset", + } + r = cov / (sig_in_target * sig_in_source) + return float(r), m, { + "n_nodes": n, + "n_edges": m, + "r_assortativity": float(r), + "mean_k_in_target": float(mx), + "mean_k_in_source": float(my), + "sigma_in_target": sig_in_target, + "sigma_in_source": sig_in_source, + "cov_edge": float(cov), + } + + +def run_method1(G, dmd, verb=1): + r, m_edges, info = directed_degree_assortativity(G) + d_ker = dmd["placement_ker_delta"] + + print("\n--- Method 1: directed degree assortativity ---") + print(" N=%d |E|=%d r=%s" % (info["n_nodes"], m_edges, repr(r) if np.isnan(r) else "%.6f" % r)) + if verb > 1: + print(" mean target k_in on edges: %.6f" % info["mean_k_in_target"]) + print(" mean source k_in on edges: %.6f" % info["mean_k_in_source"]) + print(" sigma target in: %.6f" % info["sigma_in_target"]) + print(" sigma source in: %.6f" % info["sigma_in_source"]) + print(" (truth metadata) placement_ker_delta = %s" % (d_ker,)) + + return info + + +def run_method2(G, dmd, verb=1): + G = np.asarray(G, dtype=np.float64) + G2 = G @ G + k_in = G.sum(axis=1) # Target nodes (rows) + k_out = G.sum(axis=0) # Source nodes (columns) + + ei, ej = np.nonzero(G > 0) + m = int(ei.size) + if m == 0: + return {"mean_jaccard": float("nan"), "n_edges": 0, "reason": "no edges"} + + num = G2[ei, ej] + den = k_in[ei] + k_out[ej] - num + # den is always >= 2 because i!=j, j in N_in(i), i in N_out(j) + j_idx = num / den + j_mean = float(np.mean(j_idx)) + + print("\n--- Method 2: common-neighbor overlap (Jaccard Index) ---") + print(" mean Jaccard index: %.6f" % j_mean) + + info = { + "mean_jaccard": j_mean, + "n_edges": m + } + return info + + +def run_method3(G, dmd, verb=1): + G = np.asarray(G, dtype=np.float64) + K = 6 + k_vals = np.arange(1, K + 1) # k = 1..K + + # Tk = trace(G^(k+1)) / sum(G^k) + transitivity = [] + Gk = G.copy() + for k in k_vals: + Gnext = Gk @ G + num = float(np.trace(Gnext)) + den = float(Gk.sum()) + if den > 0: + transitivity.append(num / den) + else: + transitivity.append(0.0) + Gk = Gnext + + transitivity = np.array(transitivity) + + # xi = -d log Tk / dk for k in 2..K + # Handle log compatibility + fit_mask = (k_vals >= 2) & (transitivity > 1e-16) + xi = float("nan") + if np.sum(fit_mask) >= 2: + slope, _ = np.polyfit(k_vals[fit_mask], np.log(transitivity[fit_mask]), 1) + xi = -float(slope) + + print("\n--- Method 4: cycle closure decay rate ---") + print(" decay rate xi: %.6f" % xi) + + return { + "cycle_decay": xi, + "transitivity_k": transitivity.tolist() + } + + +def run_method4(A_off, dmd, verb=1): + try: + import gudhi + except ImportError as exc: + raise ImportError( + "Method 4 requires GUDHI. Install it in this environment before running topology analysis." + ) from exc + A = np.abs(A_off) + # Symmetrized interaction strength (LaTeX Method 5, weighted variant) + W = (A + A.T) / 2.0 + N = W.shape[0] + + st = gudhi.SimplexTree() + # Add nodes at filtration 0 + for i in range(N): + st.insert([i], filtration=0.0) + + # Add edges with filtration -W_ij (strongest edges enter first) + ei, ej = np.nonzero(W > 1e-12) + for i, j in zip(ei, ej): + if i < j: + st.insert([i, j], filtration=-float(W[i, j])) + + # Expand to Rips complex of dimension 2 to capture H1 (loops) + st.expansion(2) + + # Compute persistence + st.persistence() + + # Total persistence Pi_1 (dim 1) + diag1 = st.persistence_intervals_in_dimension(1) + pi1 = 0.0 + if len(diag1) > 0: + # Filter out infinite death + valid1 = diag1[np.isfinite(diag1[:, 1])] + pi1 = float(np.sum(valid1[:, 1] - valid1[:, 0])) + + print("\n--- Method 4: persistent homology (H1) ---") + print(" total persistence Pi_1: %.6f" % pi1) + + return { + "homology_k1": pi1, + "n_features_k1": len(diag1) + } + + +def get_parser(): + parser = argparse.ArgumentParser(description="Topology analysis on Dale simTruth.npz (method 1: assortativity)") + parser.add_argument("-v", "--verbosity", type=int, help="increase output verbosity", default=1, dest="verb") + parser.add_argument("--basePath", default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input data") + parser.add_argument("--dataName", default="daleN150_448b86", help="simulated Dale network base name") + parser.add_argument("--shuffle", action="store_true", help="randomize neuron order of A_off before computing G") + + args = parser.parse_args() + + args.inpPath = os.path.join(args.basePath, "truthDale") + args.outPath = os.path.join(args.basePath, "topoAna") + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + assert os.path.exists(args.basePath) + os.makedirs(args.outPath, exist_ok=True) + return args + + +if __name__ == "__main__": + args = get_parser() + np.set_printoptions(precision=3) + + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 0) + if args.verb > 1: + print("\nSimulation Truth Metadata:") + pprint(trueMD) + + trueMD["short_name"] = args.dataName + + A_off = np.asarray(trueD["A_off_true"], dtype=np.float64) + print("read obj: A_off_true %s %s" % (A_off.shape, A_off.dtype)) + neuron_perm = np.arange(A_off.shape[0], dtype=np.int64) + if args.shuffle: + A_off, neuron_perm = shuffle_neuron_order(A_off) + print("applied neuron shuffle to A_off, first 10 perm entries:", neuron_perm[:10]) + + G = adjacency_from_A_off(A_off) + dmd = trueMD["dale_conf"] + + info1 = run_method1(G, dmd, verb=args.verb) + info2 = run_method2(G, dmd, verb=args.verb) + info3 = run_method3(G, dmd, verb=args.verb) + info4 = run_method4(A_off, dmd, verb=args.verb) + + print("\nM:done topoAna_daleMatrix4") + + + # ══════════════════════════════════════════════════════════════ + # Save results + # ══════════════════════════════════════════════════════════════ + outMD ={"dale_conf":dmd} + outMD["shuffle"] = bool(args.shuffle) + outMD['methods']={ + 'assortativity': info1, + 'jaccard_index': info2, + 'cycle_decay': info3, + 'homology_k1': info4 + } + + outD={} #'G': G,'A_off': A_off, 'neuron_perm': neuron_perm} + if args.verb>1: pprint(outMD) + outF = args.dataName + outFF = os.path.join(args.outPath, f"{outF}.topoAna.npz") + + write_data_npz(outD, outFF, metaD=outMD) + print(f"\nSaved: {outFF}") + print("\n#Summary: dataName, ker_delta, r_assort, jaccard, cycle_decay, hom_k1") + print("#Values: %s %.1f %.6f %.6f %.6f %.6f" % (args.dataName, dmd["placement_ker_delta"], info1["r_assortativity"], info2["mean_jaccard"], info3["cycle_decay"], info4["homology_k1"])) + diff --git a/causal_net/nonStation_ver4_lagM_offDiag/view_daleMatrix4.py b/causal_net/nonStation_ver4_lagM_offDiag/view_daleMatrix4.py new file mode 100755 index 00000000..a9198073 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/view_daleMatrix4.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +""" +Visualization tool for simulated Dale Poisson network data. + +Loads simulation output (simTruth.npz and spikes.npz) produced by +gen_daleMatrices4.py. Neuron order follows placement; use node_is_inhibitory +in simTruth for E/I (not index blocks). + +Available plots (-p flag): + a Dale connectivity matrix (color-coded) + eigenvalue scatter + with spectral-radius circle + b Weight histogram, firing-rate histogram, outgoing-edge count + per neuron, and per-neuron firing-rate bar chart + c B_idle vs firing rate / SNR scatter, plus excitatory and + inhibitory rate histograms (rates_study) + d 2D neuron placement: triangles=excitatory, squares=inhibitory; + red=outgoing excitatory edges, blue=outgoing inhibitory edges (presynaptic) + e Off-diagonal distance histogram, offdiag_kernel histogram, empty panel + f Pseudospectral contour plot with eigenvalue overlay + g Topology overview: signed off-diagonal matrix, distance histogram, + and placement topology on one canvas + +Usage: + ./view_daleMatrix4.py --dataName daleN100_9fbe7f -p a b c d e f g +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os +from pprint import pprint +import numpy as np +from PlotterDaleMatrix import Plotter +from toolbox.Util_NumpyIO import read_data_npz +import argparse + + +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize simulated Dale Poisson network data") + parser.add_argument("-v", "--verbosity", type=int, help="increase output verbosity", default=1, dest="verb") + parser.add_argument( + "-p", + "--showPlots", + default="a b", + nargs="+", + help="plot letters: a=Dale+eigen, b=histograms, c=rates_study, d=placement, e=dist/kernel hist, f=pseudospectra, g=topo overview", + ) + parser.add_argument("-X", "--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument("--basePath", default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input data") + parser.add_argument("--dataName", default="daleN150_448b86", help="simulated Dale network base name") + + args = parser.parse_args() + + args.inpPath = os.path.join(args.basePath, "truthDale") + args.outPath = os.path.join(args.basePath, "plots") + args.showPlots = "".join(args.showPlots) + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + assert os.path.exists(args.basePath) + return args + + +if __name__ == "__main__": + args = get_parser() + np.set_printoptions(precision=3) + + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 0) + assert trueD["B_true"].ndim == 1, "B_true must be shape (N,) from gen_daleMatrices4" + if args.verb > 1: + print("\nSimulation Truth Metadata:") + pprint(trueMD) + + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 0) + if args.verb > 1: + print("\nSpike Data Metadata:") + pprint(spikeMD) + + trueMD["short_name"] = args.dataName + + R_sel = trueMD["dale_conf"]["spectral_radius"] + print(f"\nSpectral radius: R={R_sel:.3f}") + + trueD_r = { + "A_true": trueD["A_off_true"] + np.diag(trueD["A_diag_true"]), + "B_true": trueD["B_true"], + "E_true": trueD["E_true"], + "node_is_inhibitory": trueD["node_is_inhibitory"], + } + spikeD_r = dict(spikeD) + + trueMD["sel_spect_radius"] = R_sel + + args.prjName = args.dataName + "_view" + plot = Plotter(args) + + if "a" in args.showPlots: + plot.Dale_matrix_and_eigen(trueD_r["A_true"], trueMD, trueD_r, figId=1) + + if "b" in args.showPlots: + plot.histo_weights_rates(trueD_r, spikeD_r, trueMD, figId=2) + + if "c" in args.showPlots: + plot.rates_study(trueD_r, spikeD_r, trueMD, figId=3) + + if "d" in args.showPlots: + plot.plot_placement_topology(trueD, trueMD, figId=4) + + if "e" in args.showPlots: + plot.offdiag_distance_kernel_hist(trueD, trueMD, figId=5) + + if "f" in args.showPlots: + plot.Dale_matrix_pseudospectra(trueD_r["A_true"], trueMD, trueD_r, figId=6) + + if "g" in args.showPlots: + plot.topo_overview(trueD, trueMD, figId=7) + + + plot.display_all() + print("M:done - view_dalePoisson completed successfully!") diff --git a/causal_net/nonStation_ver4_lagM_offDiag/view_spikesTrain4.py b/causal_net/nonStation_ver4_lagM_offDiag/view_spikesTrain4.py new file mode 100755 index 00000000..0a420320 --- /dev/null +++ b/causal_net/nonStation_ver4_lagM_offDiag/view_spikesTrain4.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +""" +Visualization tool for biological experiment input features and data quality. + +This script provides comprehensive visualization and analysis of experimental +neural data, focusing on data quality assessment and feature exploration. +Main functionality includes: +- Time series visualization of neural recordings with customizable time ranges +- Data quality metrics and statistical summaries +- Interactive plotting with configurable display options + +Used primarily for exploratory data analysis of biological neural recordings +before further processing and connectivity analysis. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os + +from pprint import pprint +import numpy as np +from PlotterSpikesTrain import Plotter +from toolbox.Util_NumpyIO import read_data_npz +import argparse +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize spike-train data from simulated Dale network outputs") + parser.add_argument("-v","--verbosity",type=int, help="increase output verbosity", default=1, dest='verb') + parser.add_argument("-p", "--showPlots", default='a', nargs='+',help="abcd-string listing shown plots") + + + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + + parser.add_argument("--basePath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for input data") + parser.add_argument("--dataName", default=None,help='simulated Dale network base name') + + + parser.add_argument('-T','--time_range_sec' , default=[0., 50], nargs=2, type=float, help='display data time range in seconds') + parser.add_argument('-r','--time_rebin2', default=10, type=int, help='rebin current time axis') + + args = parser.parse_args() + # make arguments more flexible + args.inpPath = os.path.join(args.basePath, 'truthDale') + + # args.inpPath = os.path.join(args.basePath, 'spikesData') + args.outPath = os.path.join(args.basePath,'plots') + args.showPlots=''.join(args.showPlots) + + print( 'myArg-program:',parser.prog) + for arg in vars(args): print( 'myArg:',arg, getattr(args, arg)) + + if args.time_range_sec!=None: assert args.time_range_sec[0] < args.time_range_sec[1] + assert os.path.exists(args.basePath) + assert os.path.exists(args.inpPath), f"missing inpPath: {args.inpPath}" + return args + + +#...!...!.................... +def rebin_spike_rates(spikeYield, md, tReb2): + """Rebin 2D spike train (time x neurons) and return rate summaries over rebinned time.""" + assert spikeYield.ndim == 2, f"Expected 2D spike train, got shape={spikeYield.shape}" + assert tReb2<101 # this would exceed 1 seconds + + #.... rebin time axis + ntime,nchan=spikeYield.shape + if ntime % tReb2 != 0: + ntime_c = ntime - (ntime % tReb2) + spikeYield = spikeYield[:ntime_c] # Clip the data + #print('ccl',ntime_c , ntime ,ntime % tReb2, tReb2) + #... sum yileds over tReb2 time bins + spikeYieldR= np.sum( spikeYield.reshape(-1, tReb2, nchan),axis=1) + time_step=md['time_step_sec'] + time_step2=time_step*tReb2 + #print('rr1',time_step2,spikeYieldR.shape,spikeYield.shape) + + rate2D=spikeYieldR/time_step2 + pop_spike_count = np.sum(spikeYieldR, axis=1) + pop_rate_hz = pop_spike_count / time_step2 + + rebD={'time_step2':time_step2} + rebD['rate2D']=rate2D + rebD['pop_spike_count'] = pop_spike_count + rebD['pop_rate_hz'] = pop_rate_hz + ntime//=tReb2 + rebD['timeV']= np.linspace(0, (ntime - 1) * time_step2, ntime) + + print('rebinned rates: nchan=%d dt=%.3f sec rebin=%d' % (nchan, time_step2, tReb2)) + return rebD + + + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__=="__main__": + args=get_parser() + np.set_printoptions(precision=3) + + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb>0) + if args.verb>1: pprint(spikeMD) + + S_oracle = None + if 1: # per state information + # Load simulation truth data (Dale matrices, biases, etc.) + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb>0) + if args.verb>1: + print("\nSimulation Truth Metadata:"); pprint(trueMD) + + dataYield = np.asarray(spikeD["spikes"]) + if dataYield.ndim != 2: + raise ValueError( + "spikes must have shape (T, N) from gen_daleMatrices4; got %s. Multi-state (M, T, N) is not supported." + % (dataYield.shape,) + ) + sr = np.asarray(spikeD["single_rates"]) + if sr.ndim != 1: + raise ValueError("single_rates must be (N,); got shape %s" % (sr.shape,)) + if sr.shape[0] != dataYield.shape[1]: + raise ValueError("single_rates length must match spikes.shape[1]") + spikeD["spikes"] = dataYield + spikeD["single_rates"] = sr + spikeMD["sel_spect_radius"] = trueMD["dale_conf"]["spectral_radius"] + S_true = None + + + #-------------------------------- + # .... plotting ........ + spikeMD["short_name"] = args.dataName + args.prjName = spikeMD["short_name"] + spikeMD['plot']={} + spikeMD['plot']['time_rangeLR']=np.array(args.time_range_sec) + + plot=Plotter(args) + + if 'a' in args.showPlots: + plot.freq_histo(spikeD,spikeMD,figId=1) + if 'b' in args.showPlots: + rebD=rebin_spike_rates(spikeD['spikes'], spikeMD, args.time_rebin2) + if S_oracle is not None: + rebD['S_oracle'] = S_oracle + plot.freq_vs_time(rebD,spikeMD,figId=2, S_true=S_true, S_oracle=S_oracle) + + plot.display_all() + print('M:done') + #pprint(expMD) #tmp diff --git a/causal_net/nonStation_ver5_synapDepres/PlotterBSSM_fit.py b/causal_net/nonStation_ver5_synapDepres/PlotterBSSM_fit.py new file mode 100644 index 00000000..25289bd9 --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/PlotterBSSM_fit.py @@ -0,0 +1,408 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for BSSM-STD block-coordinate fit evaluation. +""" + +import os +import warnings +import numpy as np + +os.environ.setdefault("MPLCONFIGDIR", "/tmp") + +import matplotlib.colors as mcolors +from matplotlib.ticker import MaxNLocator + +from toolbox.PlotterBackbone import PlotterBackbone + + +def _scalar(arr): + return float(np.asarray(arr).reshape(-1)[0]) + + +def _offdiag(W): + W = np.asarray(W) + assert W.ndim == 2 and W.shape[0] == W.shape[1], "W must be square" + return W[~np.eye(W.shape[0], dtype=bool)] + + +def _corr(x, y): + x = np.asarray(x, dtype=np.float64).reshape(-1) + y = np.asarray(y, dtype=np.float64).reshape(-1) + assert x.shape == y.shape and x.size >= 2, "correlation arrays must match and be nontrivial" + if np.std(x) <= 0.0 or np.std(y) <= 0.0: + return float("nan") + return float(np.corrcoef(x, y)[0, 1]) + + +def _sym_limit(*arrs): + vmax = 0.0 + for arr in arrs: + if arr is None: + continue + val = float(np.nanmax(np.abs(np.asarray(arr)))) + vmax = max(vmax, val) + return max(vmax, 1e-12) + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def _heat(self, ax, mat, title, cmap="coolwarm", vlim=None): + mat = np.asarray(mat) + assert mat.ndim == 2, "heatmap input must be 2D" + if vlim is None: + im = ax.imshow(mat, origin="lower", aspect="auto", cmap=cmap) + else: + im = ax.imshow( + mat, + origin="lower", + aspect="auto", + cmap=cmap, + norm=mcolors.TwoSlopeNorm(vcenter=0.0, vmin=-vlim, vmax=vlim), + ) + ax.set(title=title, xlabel="pre neuron j", ylabel="post neuron i") + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + return im + + def two_block_summary(self, fitD, diag, figId=1): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(14, 8)) + gs = fig.add_gridspec(2, 3, hspace=0.38, wspace=0.34) + + outer = np.asarray(fitD["outer_idx"], dtype=np.int32) + 1 + U = np.asarray(fitD["U_outer"], dtype=np.float64) + tau_rec = np.asarray(fitD["tau_rec_outer"], dtype=np.float64) + nll_A = np.asarray(fitD["nll_after_blockA"], dtype=np.float64) + nll_B = np.asarray(fitD["nll_after_std_grid"], dtype=np.float64) + nll = np.asarray(fitD["nll_outer"], dtype=np.float64) + blockB_executed = np.asarray(fitD["blockB_executed"], dtype=np.int32).astype(bool) + rho = np.asarray(fitD["spectral_radius_outer"], dtype=np.float64) + nz = np.asarray(fitD["nz_weight_outer"], dtype=np.int64) + block_outer = np.asarray(fitD["blockA_outer"], dtype=np.int32) + 1 + block_lr = np.asarray(fitD["blockA_lr"], dtype=np.float64) + + ax = fig.add_subplot(gs[0, 0]) + ax.plot(outer, nll_A, "o-", label="after Block A", color="tab:blue") + ax.plot(outer, nll, "s-", label="outer result", color="tab:green") + ax.plot(outer[blockB_executed], nll_B[blockB_executed], "^-", label="after Block B", color="tab:orange") + ax.set(title="Outer mean Bernoulli NLL", xlabel="outer iteration", ylabel="mean NLL") + ax.xaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + lr_outer = np.empty_like(outer, dtype=np.float64) + for i, o in enumerate(outer): + m = block_outer == o + assert np.any(m), "missing Block A LR values for outer %d" % int(o) + lr_outer[i] = block_lr[np.where(m)[0][-1]] + ax_lr = ax.twinx() + ax_lr.plot(outer, lr_outer, "d--", color="tab:red", lw=1.0, ms=4, label="Block A LR") + ax_lr.set_ylabel("Block A LR", color="tab:red") + ax_lr.tick_params(axis="y", labelcolor="tab:red") + lines, labels = ax.get_legend_handles_labels() + lines_lr, labels_lr = ax_lr.get_legend_handles_labels() + ax.legend(lines + lines_lr, labels + labels_lr, fontsize=8) + + ax = fig.add_subplot(gs[0, 1]) + ax.plot(outer, U, "o-", color="tab:purple", label="U") + ax.set(title=r"Block B STD parameter path", xlabel="outer iteration", ylabel=r"$U$") + ax.tick_params(axis="y", labelcolor="tab:purple") + ax.yaxis.label.set_color("tab:purple") + ax.grid(True, alpha=0.3) + ax.xaxis.set_major_locator(MaxNLocator(integer=True)) + ax2 = ax.twinx() + ax2.plot(outer, tau_rec, "s-", color="tab:brown", label=r"$\tau_{rec}$") + ax2.set_ylabel(r"$\tau_{rec}$ (s)", color="tab:brown") + ax2.tick_params(axis="y", labelcolor="tab:brown") + + ax = fig.add_subplot(gs[0, 2]) + ep = np.arange(1, block_outer.size + 1) + bce = np.asarray(fitD["blockA_bce_loss"], dtype=np.float64) + l1 = np.asarray(fitD["blockA_l1_loss"], dtype=np.float64) + total = np.asarray(fitD["blockA_loss"], dtype=np.float64) + ax.plot(ep, bce, color="tab:blue", lw=1.2, label="BCE") + ax.plot(ep, total, color="k", lw=1.0, alpha=0.75, label="BCE+L1") + ax.set(title="Block A objective", xlabel="global Block A epoch", ylabel=r"$\mathcal{L}_A$") + ax.grid(True, alpha=0.3) + ax_l1 = ax.twinx() + ax_l1.plot(ep, l1, color="tab:red", lw=1.0, ls="--", label="L1") + ax_l1.set_ylabel("L1", color="tab:red") + ax_l1.tick_params(axis="y", labelcolor="tab:red") + for o in np.unique(block_outer)[1:]: + x = int(np.where(block_outer == o)[0][0]) + 1 + ax.axvline(x, color="gray", lw=0.6, alpha=0.35) + ax.legend(fontsize=8, loc="upper right") + + ax = fig.add_subplot(gs[1, 0]) + rho_max = _scalar(fitD["rho_max"]) + ax.plot(outer, rho, "o-", color="tab:green") + ax.axhline(rho_max, color="tab:red", lw=1.0, ls="--", label=r"$\rho_{max}$") + ax.set(title="Spectral radius", xlabel="outer iteration", ylabel=r"$\rho(W)$") + ax.xaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + + ax = fig.add_subplot(gs[1, 1]) + ax.plot(outer, nz, "o-", color="tab:cyan") + ax.set( + title="Reported nonzero W count", + xlabel="outer iteration", + ylabel="count", + ) + ax.xaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + + ax = fig.add_subplot(gs[1, 2]) + ax.axis("off") + text = ( + f"N={int(_scalar(fitD['num_neurons']))}\n" + f"eval bins={int(_scalar(fitD['num_eval_bins']))}\n" + f"dt={_scalar(fitD['time_step_sec']):.4g} s\n" + f"burn bins={int(_scalar(fitD['burn_bins']))}\n" + f"init bins={int(_scalar(fitD['num_init_bins']))}\n" + f"mean |W_init|={_scalar(fitD['W_init_abs_mean']):.3g}\n" + f"M={len(fitD['kappa_fit'])}, alpha={_scalar(fitD['alpha_fit']):.5f}\n" + f"lambda_l1={_scalar(fitD['lambda_l1']):.3g}\n" + f"eta_clip={_scalar(fitD['eta_clip']):.3g}\n" + f"final U={U[-1]:.5f}\n" + f"final tau_rec={tau_rec[-1]:.5f} s\n" + f"final NLL={nll[-1]:.6f}" + ) + ax.text(0.02, 0.98, text, transform=ax.transAxes, va="top", ha="left", fontsize=10) + fig.suptitle(f"BSSM-STD Two-Block Fit Summary: {diag['short_name']}", fontsize=14) + + def parameter_init_vs_fit(self, fitD, diag, figId=2): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(15, 8.5)) + gs = fig.add_gridspec(2, 3, hspace=0.38, wspace=0.32) + + W0 = np.asarray(fitD["W_init"], dtype=np.float64) + W = np.asarray(fitD["W_fit"], dtype=np.float64) + B0 = np.asarray(fitD["B_init"], dtype=np.float64) + B = np.asarray(fitD["B_fit"], dtype=np.float64) + assert W0.shape == W.shape, "W_init/W_fit shape mismatch" + assert B0.shape == B.shape == (W.shape[0],), "B_init/B_fit shape mismatch" + + vlim = _sym_limit(W0, W) + self._heat(fig.add_subplot(gs[0, 0]), W0, "W_init", vlim=vlim) + self._heat(fig.add_subplot(gs[0, 1]), W, "W_fit", vlim=vlim) + self._heat(fig.add_subplot(gs[0, 2]), W - W0, "W_fit - W_init", vlim=_sym_limit(W - W0)) + + ax = fig.add_subplot(gs[1, 0]) + rB = _corr(B0, B) + ax.scatter(B0, B, s=14, alpha=0.75, color="tab:blue", edgecolors="none") + lo = min(float(B0.min()), float(B.min())) + hi = max(float(B0.max()), float(B.max())) + ax.plot([lo, hi], [lo, hi], "k--", lw=0.8) + ax.set(title=f"$B_i$ init vs fit, r={rB:.3f}", xlabel="B_init", ylabel="B_fit") + ax.grid(True, alpha=0.3) + + ax = fig.add_subplot(gs[1, 1]) + w0 = _offdiag(W0) + wf = _offdiag(W) + rW = _corr(w0, wf) + ax.scatter(w0, wf, s=5, alpha=0.25, color="tab:green", edgecolors="none") + xspan = max(float(np.max(np.abs(w0))), 1e-12) + yspan = max(float(np.max(np.abs(wf))), 1e-12) + ax.set_xlim(float(w0.min()) - 0.05 * xspan, float(w0.max()) + 0.05 * xspan) + ax.set_ylim(float(wf.min()) - 0.05 * yspan, float(wf.max()) + 0.05 * yspan) + ax.axvline(0.0, color="k", lw=0.7, alpha=0.45) + ax.axhline(0.0, color="k", lw=0.7, alpha=0.45) + ax.set( + title=f"Eq.17 offdiag $W_{{ij}}^{{(0)}}$ vs fit, r={rW:.3f}", + xlabel="raw lag-1 covariance W_init", + ylabel="W_fit", + ) + ax.text( + 0.03, + 0.97, + "mean |W_init|=%.3g\nmean |W_fit|=%.3g" % (np.mean(np.abs(w0)), np.mean(np.abs(wf))), + transform=ax.transAxes, + va="top", + ha="left", + fontsize=8, + bbox=dict(facecolor="white", alpha=0.75, edgecolor="none"), + ) + ax.grid(True, alpha=0.3) + + ax = fig.add_subplot(gs[1, 2]) + ax.hist(w0, bins=80, histtype="step", lw=1.3, color="gray", label="W_init") + ax.hist(wf, bins=80, histtype="step", lw=1.3, color="tab:green", label="W_fit") + ax.set_yscale("log") + ax.set(title=r"Off-diagonal $W_{ij}$ distribution", xlabel="weight", ylabel="count") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + fig.suptitle(f"BSSM-STD $B_i,W_{{ij}}$ Initialization and Fit: {diag['short_name']}", fontsize=14) + + def blockB_joint_grid(self, fitD, diag, figId=3): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(15, 8)) + gs = fig.add_gridspec(2, 3, hspace=0.4, wspace=0.35) + + outer = np.asarray(fitD["joint_search_outer"], dtype=np.int32) + 1 + refine = np.asarray(fitD["joint_search_refine"], dtype=np.int32) + eval_id = np.asarray(fitD["joint_search_eval"], dtype=np.int32) + eval_total = np.asarray(fitD["joint_search_eval_total"], dtype=np.int32) + U = np.asarray(fitD["joint_search_U"], dtype=np.float64) + tau_rec = np.asarray(fitD["joint_search_tau_rec"], dtype=np.float64) + nll = np.asarray(fitD["joint_search_nll"], dtype=np.float64) + accepted_outer = np.asarray(fitD["outer_idx"], dtype=np.int32) + 1 + accepted_U = np.asarray(fitD["U_outer"], dtype=np.float64) + accepted_tau = np.asarray(fitD["tau_rec_outer"], dtype=np.float64) + assert outer.size > 0, "joint search arrays are empty" + assert outer.shape == refine.shape == eval_id.shape == eval_total.shape == U.shape == tau_rec.shape == nll.shape + assert accepted_outer.shape == accepted_U.shape == accepted_tau.shape, "outer STD trajectory shape mismatch" + + rows = [] + for o in np.unique(outer): + for r in np.unique(refine[outer == o]): + m = (outer == o) & (refine == r) + idx = np.where(m)[0][np.argmin(nll[m])] + rows.append((o, r, U[idx], tau_rec[idx], nll[idx], int(np.sum(m)))) + rows = np.asarray(rows, dtype=np.float64) + + ax = fig.add_subplot(gs[0, 0]) + for r in np.unique(rows[:, 1].astype(np.int32)): + m = rows[:, 1] == r + ax.plot(rows[m, 0], rows[m, 2], "o-", label=f"refine {r}") + ax.set(title="Block B best U per grid round", xlabel="outer", ylabel=r"$U$") + ax.xaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + + ax = fig.add_subplot(gs[0, 1]) + for r in np.unique(rows[:, 1].astype(np.int32)): + m = rows[:, 1] == r + ax.plot(rows[m, 0], rows[m, 3], "o-", label=f"refine {r}") + ax.set(title=r"Block B best $\tau_{rec}$ per grid round", xlabel="outer", ylabel=r"$\tau_{rec}$ (s)") + ax.xaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + + ax = fig.add_subplot(gs[0, 2]) + for r in np.unique(rows[:, 1].astype(np.int32)): + m = rows[:, 1] == r + ax.plot(rows[m, 0], rows[m, 4], "o-", label=f"refine {r}") + ax.set(title="Best Block B grid NLL", xlabel="outer", ylabel="mean NLL") + ax.xaxis.set_major_locator(MaxNLocator(integer=True)) + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + + last_outer = int(np.max(outer)) + last_refine = int(np.max(refine[outer == last_outer])) + m = (outer == last_outer) & (refine == last_refine) + ax = fig.add_subplot(gs[1, 0]) + sc = ax.scatter(U[m], tau_rec[m], c=nll[m], s=70, cmap="viridis", edgecolors="k", linewidths=0.2) + ax.plot( + accepted_U, + accepted_tau, + "o-", + color="black", + lw=1.2, + ms=4, + label="outer path", + zorder=5, + ) + for oo, uu, tt in zip(accepted_outer, accepted_U, accepted_tau): + ax.text(uu, tt, str(int(oo)), fontsize=7, color="black", ha="left", va="bottom") + best = np.where(m)[0][np.argmin(nll[m])] + ax.plot(U[best], tau_rec[best], "r*", ms=16, label="best") + ax.set( + title=f"Block B candidates outer {last_outer}, refine {last_refine}", + xlabel=r"$U$", + ylabel=r"$\tau_{rec}$ (s)", + ) + self.plt.colorbar(sc, ax=ax, fraction=0.046, pad=0.04, label="mean NLL") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + + ax = fig.add_subplot(gs[1, 1]) + u_vals = np.unique(U[m]) + tau_vals = np.unique(tau_rec[m]) + Z = np.full((tau_vals.size, u_vals.size), np.nan, dtype=np.float64) + for uu, tt, nn in zip(U[m], tau_rec[m], nll[m]): + iu = int(np.where(u_vals == uu)[0][0]) + it = int(np.where(tau_vals == tt)[0][0]) + if not np.isfinite(Z[it, iu]) or nn < Z[it, iu]: + Z[it, iu] = nn + im = ax.imshow( + Z, + origin="lower", + aspect="auto", + extent=[u_vals.min(), u_vals.max(), tau_vals.min(), tau_vals.max()], + cmap="viridis", + ) + ax.set(title="Final Block B refine NLL surface", xlabel=r"$U$", ylabel=r"$\tau_{rec}$ (s)") + self.plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + ax = fig.add_subplot(gs[1, 2]) + ax.hist(nll - np.min(nll), bins=60, color="tab:purple", alpha=0.75) + ax.set(title="All Block B candidates above best", xlabel=r"$\Delta$NLL", ylabel="count") + ax.grid(True, alpha=0.3) + + fig.suptitle(f"BSSM-STD Block B Joint $U,\\tau_{{rec}}$ Grid: {diag['short_name']}", fontsize=14) + + def bernoulli_observation_diagnostics(self, fitD, diag, figId=4): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(15, 8)) + gs = fig.add_gridspec(2, 3, hspace=0.4, wspace=0.35) + + obs = np.asarray(diag["obs_rate_hz"], dtype=np.float64) + model = np.asarray(diag["model_rate_hz"], dtype=np.float64) + nll = np.asarray(diag["nll_neuron"], dtype=np.float64) + single_rates = np.asarray(diag["single_rates_hz"], dtype=np.float64) + resid = model - obs + + ax = fig.add_subplot(gs[0, 0]) + r = _corr(obs, model) + ax.scatter(obs, model, s=16, alpha=0.75, color="tab:blue", edgecolors="none") + lo = min(float(obs.min()), float(model.min())) + hi = max(float(obs.max()), float(model.max())) + ax.plot([lo, hi], [lo, hi], "k--", lw=0.8) + ax.set(title=f"Model vs observed rates, r={r:.3f}", xlabel="observed rate (Hz)", ylabel="model rate (Hz)") + ax.grid(True, alpha=0.3) + + ax = fig.add_subplot(gs[0, 1]) + ax.hist(resid, bins=40, color="tab:orange", alpha=0.65, label="model - fit window obs") + ax.hist(obs - single_rates, bins=40, histtype="step", color="tab:gray", lw=1.3, label="window obs - full data") + ax.axvline(0, color="k", lw=0.8, ls="--") + ax.set(title="Rate residuals", xlabel="rate difference (Hz)", ylabel="neurons") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + + ax = fig.add_subplot(gs[0, 2]) + ax.plot(np.arange(nll.size), nll, color="tab:green") + ax.set(title="Per-neuron mean Bernoulli NLL", xlabel="neuron index", ylabel="NLL") + ax.grid(True, alpha=0.3) + + ax = fig.add_subplot(gs[1, 0]) + calib_p = np.asarray(diag["calib_p"], dtype=np.float64) + calib_obs = np.asarray(diag["calib_obs"], dtype=np.float64) + calib_count = np.asarray(diag["calib_count"], dtype=np.int64) + valid = calib_count > 0 + ax.plot([0, 1], [0, 1], "k--", lw=0.8) + ax.scatter(calib_p[valid], calib_obs[valid], s=np.clip(calib_count[valid] / calib_count[valid].max() * 120, 20, 120), + color="tab:red", alpha=0.75, edgecolors="none") + ax.set(title="Bernoulli probability calibration", xlabel=r"mean model $\tilde{p}_{i,t}$", ylabel="observed spike fraction") + ax.grid(True, alpha=0.3) + + ax = fig.add_subplot(gs[1, 1]) + edges = np.asarray(diag["p_hist_edges"], dtype=np.float64) + cnt = np.asarray(diag["p_hist_count"], dtype=np.int64) + centers = 0.5 * (edges[:-1] + edges[1:]) + ax.bar(centers, cnt, width=np.diff(edges), align="center", color="tab:blue", alpha=0.75) + ax.set_yscale("log") + ax.set(title=r"Model probability distribution", xlabel=r"$\tilde{p}_{i,t}$", ylabel="count") + ax.grid(True, alpha=0.3) + + ax = fig.add_subplot(gs[1, 2]) + B = np.asarray(fitD["B_fit"], dtype=np.float64) + x = np.log(np.clip(obs * _scalar(fitD["time_step_sec"]), 1e-9, 1.0 - 1e-9)) + ax.scatter(x, B, s=16, alpha=0.75, color="tab:purple", edgecolors="none") + rr = _corr(x, B) + ax.set(title=f"B_fit vs log observed spike prob, r={rr:.3f}", xlabel="log(obs p/bin)", ylabel="B_fit") + ax.grid(True, alpha=0.3) + + fig.suptitle( + f"BSSM-STD Bernoulli Observation Diagnostics: {diag['short_name']} recomputed NLL={diag['nll_eval']:.6f}", + fontsize=14, + ) diff --git a/causal_net/nonStation_ver5_synapDepres/PlotterDaleMatrix.py b/causal_net/nonStation_ver5_synapDepres/PlotterDaleMatrix.py new file mode 100644 index 00000000..e15b770a --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/PlotterDaleMatrix.py @@ -0,0 +1,656 @@ +#!/usr/bin/env python3 +"""Plotting utilities for the BSSM-STD Bernoulli network generator.""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +import matplotlib.ticker as ticker +import numpy as np +import matplotlib.gridspec as gridspec +import matplotlib.colors as colors +from matplotlib.collections import LineCollection +from matplotlib.lines import Line2D +from Util_pseudospectra import compute_pseudospectrum + + +def _placement_ker_delta(dmd): + return float(dmd["placement_ker_delta"]) + + +def _format_ker_delta_latex(delta): + """Matplotlib title fragment: δ_ker value (integers without .0).""" + x = float(delta) + if abs(x - round(x)) < 1e-9: + return str(int(round(x))) + return "%.4g" % x + + +def _kernel_figure_canvas_title(md): + """One-line caption for the kernel / latent summary canvas.""" + dmd = md["dale_conf"] + parts = [] + parts.append("dataset: %s" % md["short_name"]) + parts.append("N=%d" % int(dmd["num_neurons"])) + parts.append(r"$\tau_s$=%.4g s" % float(dmd["synaptic_tau"])) + parts.append(r"$\tau_{rec}$=%.4g s" % float(dmd["std_tau_rec"])) + parts.append("U=%.3g" % float(dmd["std_u"])) + parts.append("M=%d" % int(dmd["mem_lag_steps"])) + parts.append("placement HxL=(%gx%g)" % (float(dmd["placement_H"]), float(dmd["placement_L"]))) + parts.append("δ_ker=%s" % _format_ker_delta_latex(_placement_ker_delta(dmd))) + return " ".join(parts) + + +def _dale_overview_title(md): + dmd = md["dale_conf"] + return r"BSSM-STD network, $N=%d$, $\delta_{\mathrm{ker}}=%s$, %s" % ( + int(dmd["num_neurons"]), + _format_ker_delta_latex(_placement_ker_delta(dmd)), + md["short_name"], + ) + + +def _baseline_rate_hz(trueD, md): + b_true = np.asarray(trueD["b_true"], dtype=np.float64).reshape(-1) + dt = float(md["evol_conf"]["step_size"]) + p0 = 1.0 / (1.0 + np.exp(-b_true)) + return p0 / dt + + +def _realized_edge_lengths(trueD): + D = np.asarray(trueD["node_distance_matrix"], dtype=np.float64) + if D.ndim != 2 or D.shape[0] != D.shape[1]: + raise ValueError("node_distance_matrix must be square (N, N)") + N = D.shape[0] + edge_mask = np.asarray(trueD["sign_true"]) != 0 + if edge_mask.shape != (N, N): + raise ValueError("edge mask shape mismatch with node_distance_matrix") + edge_mask = edge_mask.copy() + np.fill_diagonal(edge_mask, False) + d_off = D[edge_mask] + return d_off[d_off > 0] + + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + + def _plot_signed_connectivity_matrix(self, ax, trueD, md, title_txt=None): + dmd = md["dale_conf"] + numNeur = int(dmd["num_neurons"]) + placement_ker_delta = _placement_ker_delta(dmd) + + A_sign = np.asarray(trueD["sign_true"], dtype=np.int8).copy() + if A_sign.shape != (numNeur, numNeur): + raise ValueError("signed connectivity shape mismatch") + np.fill_diagonal(A_sign, 0) + + sign_cmap = colors.ListedColormap(["blue", "white", "red"]) + sign_norm = colors.BoundaryNorm([-1.5, -0.5, 0.5, 1.5], sign_cmap.N) + ax.imshow( + A_sign.T, + aspect=1.0, + origin="lower", + cmap=sign_cmap, + norm=sign_norm, + interpolation="nearest", + ) + ax.set( + xlabel="postsynaptic neuron index $i$", + ylabel="presynaptic neuron index $j$", + ) + ax.set_aspect(1.0) + ax.grid(True, alpha=0.45) + if title_txt is None: + title_txt = _dale_overview_title(md) + ax.set_title(title_txt) + ax.plot([0, numNeur], [0, numNeur], "--", lw=0.8, color="magenta") + +#...!...!.................. + def Dale_matrix_and_eigen(self,A,md,trueD,figId=3): + + figId=self.smart_append(figId) + nrow,ncol=1,2 + fig=self.plt.figure(figId,facecolor='white', figsize=(12,5.)) + + dmd=md['dale_conf'] + numNeur=dmd['num_neurons'] + placement_ker_delta = _placement_ker_delta(dmd) + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel + + vmin = A.min() + vmax = A.max() + normMap = colors.TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax) + + #.... Position 1: Dale matrix (transpose: x = postsynaptic i, y = presynaptic j) ...... + ax = self.plt.subplot(nrow,ncol,1) + im = ax.imshow( + A.T, aspect=1., origin='lower', cmap='bwr', norm=normMap, interpolation='nearest' + ) + ax.set( + xlabel='postsynaptic neuron index $i$', + ylabel='presynaptic neuron index $j$', + ) + ax.set_aspect(1.0) + ax.grid() + cbar = fig.colorbar(im, ax=ax, extend="both", shrink=0.7) + cbar.set_label('weight $W_{ij}$') + + tit_left = r"BSSM-STD weights $W$, $N=%d$, $\delta_{\mathrm{ker}}=%s$, %s" % ( + A.shape[0], + _format_ker_delta_latex(placement_ker_delta), + md["short_name"], + ) + ax.set(title=tit_left) + ax.plot([0,numNeur],[0,numNeur],'--',lw=0.8,color='magenta') + + #..... Position 2: Eigenvalues...... + ax = self.plt.subplot(nrow,ncol,2) + Eigen=np.linalg.eigvals(A) + real_parts = np.real(Eigen) + imag_parts = np.imag(Eigen) + ax.scatter(real_parts, imag_parts, color='blue', marker='o') + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + tit3='Eigenvalues of recurrent $W$, N%d%s, %s'%(A.shape[0], R_tag, md['short_name']) + ax.set_title(tit3) + ax.axhline(0, color='black', lw=0.5) + ax.axvline(0, color='black', lw=0.5) + ax.axvline(0,color='red', linestyle='--') + ax.set_aspect('equal') + theta = np.linspace(0, 2*np.pi, 200) + ax.plot(R_sel*np.cos(theta), R_sel*np.sin(theta), color='magenta', linestyle='--', lw=1.5, label='R=%.3f'%R_sel) + ax.legend() + ax.grid(True) + + def Dale_matrix_pseudospectra(self, A, md, trueD, figId=5): # p=e + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(8, 6)) + + npts = 80 + minY = -0.2 + epsMin = 0.03 + + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + ax = self.plt.subplot(1, 1, 1) + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel + title = 'Pseudospectra of recurrent $W$, N%d%s, %s' % (A.shape[0], R_tag, md['short_name']) + + levels = np.logspace(-2.5, -0.5, 10) + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + ax.text(0.05, 0.85, 'green: ε<%.3f'%epsMin, transform=ax.transAxes) + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) + + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(minY,1.1) + ax.set_xlim(-1.1,1.1) + ax.set_aspect(1.) + ax.axhline(1,color='m',linestyle='--') + ax.axvline(1,color='m',linestyle='--') + ax.axvline(-1,color='m',linestyle='--') + +#...!...!.................. + def histo_weights_rates(self,trueD,spikeD,md,figId=3): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,3.5)) + dmd=md['dale_conf'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.2f' % R_sel + + W = np.asarray(trueD["W_true"], dtype=np.float64) + off_mask = ~np.eye(numNeur, dtype=bool) + pos_weights = W[(W > 0) & off_mask] + neg_weights = W[(W < 0) & off_mask] + single_rates = np.asarray(spikeD['single_rates']).reshape(-1) + tau = np.asarray(trueD["node_is_inhibitory"]).reshape(-1) + if tau.shape[0] != numNeur: + raise ValueError("node_is_inhibitory length must match num_neurons") + if tau.dtype.kind not in "iu": + tau = tau.astype(np.int64) + if np.any((tau != 0) & (tau != 1)): + raise ValueError("node_is_inhibitory must be 0 (exc) or 1 (inh)") + exc_mask = tau == 0 + inh_mask = tau == 1 + baseline_rates = _baseline_rate_hz(trueD, md) + edge_mask = np.asarray(trueD["sign_true"]) != 0 + in_degree = np.sum(edge_mask, axis=1) + out_degree = np.sum(edge_mask, axis=0) + x_vals = np.arange(numNeur) + + #.... weights histogram + ax = self.plt.subplot(nrow,ncol,1) + weight_all = np.concatenate([neg_weights, pos_weights]) if (neg_weights.size + pos_weights.size) > 0 else np.array([-1.0, 1.0]) + binX= np.linspace(float(np.min(weight_all)), float(np.max(weight_all)), 80) + if neg_weights.size > 0: + ax.hist(neg_weights, bins=binX, color='blue', alpha=0.7, edgecolor=None, label='inhib:%d'%neg_weights.size) + if pos_weights.size > 0: + ax.hist(pos_weights, bins=binX, color='red', alpha=0.7, edgecolor=None, label='excit:%d'%pos_weights.size) + ax.legend(loc='upper left') + tit='$W_{ij}$, N%d%s, %s'%(W.shape[0], R_tag, md['short_name']) + ax.set(title=tit, xlabel='off-diagonal weight value',ylabel='num edges') + ax.axvline(0,color='k',ls='--') + ax.grid(True, alpha=0.3) + + #.... in/out degree + ax = self.plt.subplot(nrow,ncol,2) + ax.plot(x_vals, out_degree, color='darkred', lw=1.2, alpha=0.85, label='out-degree') + ax.plot(x_vals, in_degree, color='darkblue', lw=1.2, alpha=0.85, label='in-degree') + ax.set_xlabel('neuron index') + ax.set_ylabel('num edges') + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + probLo, probHi = dmd['edge_prob'] + rho_title = f'realized degree counts, prob=[{probLo:.2f}, {probHi:.2f}]' + ax.set_title(rho_title) + ax.legend() + + #.... firing-rate histogram + ax = self.plt.subplot(nrow,ncol,3) + common_bins = np.linspace(0, max(20.0, float(np.max(single_rates)) * 1.05), 24) + ax.hist(single_rates[exc_mask], bins=common_bins, color='red', alpha=0.6, edgecolor=None, label='exc') + ax.hist(single_rates[inh_mask], bins=common_bins, color='blue', alpha=0.6, edgecolor=None, label='inh') + ax.set_xlabel('firing rate (Hz)') + ax.set_ylabel('num neurons') + ax.grid(True, alpha=0.3) + ax.set_title('Empirical rate distribution') + ax.legend() + + #.... baseline-vs-empirical rates + ax = self.plt.subplot(nrow,ncol,4) + ax.scatter(baseline_rates[exc_mask], single_rates[exc_mask], s=18, alpha=0.65, + marker='^', facecolors='none', edgecolors='red', label='Excitatory') + ax.scatter(baseline_rates[inh_mask], single_rates[inh_mask], s=18, alpha=0.65, + marker='s', facecolors='none', edgecolors='blue', label='Inhibitory') + lim_hi = max(float(np.max(baseline_rates)), float(np.max(single_rates)), 1.0) + ax.plot([0, lim_hi], [0, lim_hi], '--', color='black', lw=0.8, label='identity') + ax.set_xlabel('baseline Bernoulli rate from $b_i$ (Hz)') + ax.set_ylabel('empirical firing rate (Hz)') + ax.grid(True, alpha=0.3) + ax.set_title('Baseline rate vs empirical rate') + ax.legend() + + +#............................ +#............................ +#............................ + def rates_study(self,trueD,spikeD,md,figId=4): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,4)) + dmd=md['dale_conf'] + numNeur=dmd['num_neurons'] + R_sel = md['sel_spect_radius'] + R_tag = ', R=%.3f' % R_sel + + single_rates = np.asarray(spikeD["single_rates"]).reshape(-1) + x_true = np.asarray(trueD["x_true"], dtype=np.float64) + u_true = np.asarray(trueD["u_true"], dtype=np.float64) + h_true = np.asarray(trueD["h_true"], dtype=np.float64) + p_true = np.asarray(trueD["p_true"], dtype=np.float64) + + tau = np.asarray(trueD["node_is_inhibitory"]).reshape(-1) + if tau.shape[0] != numNeur: + raise ValueError("node_is_inhibitory length must match num_neurons") + if tau.dtype.kind not in "iu": + tau = tau.astype(np.int64) + if np.any((tau != 0) & (tau != 1)): + raise ValueError("node_is_inhibitory must be 0 (exc) or 1 (inh)") + exc_mask = tau == 0 + inh_mask = tau == 1 + dt = float(md["evol_conf"]["step_size"]) + T = x_true.shape[0] + stride = max(1, T // 2000) + t_sec = dt * np.arange(0, T, stride, dtype=np.float64) + mean_x_t = np.mean(x_true[::stride], axis=1) + p10_x_t = np.percentile(x_true[::stride], 10, axis=1) + p90_x_t = np.percentile(x_true[::stride], 90, axis=1) + mean_u_t = np.mean(u_true[::stride], axis=1) + mean_h_t = np.mean(h_true[::stride], axis=1) + mean_p_rate = np.mean(p_true, axis=0) / dt + mean_x_per_neuron = np.mean(x_true, axis=0) + + # 1) population-average STD resource trace + ax = self.plt.subplot(nrow,ncol,1) + ax.fill_between(t_sec, p10_x_t, p90_x_t, color='lightsteelblue', alpha=0.45, label='10-90 pct') + ax.plot(t_sec, mean_x_t, color='navy', lw=1.4, label=r'mean $x_t$') + ax.set_xlabel('time (s)') + ax.set_ylabel(r'STD resource $x$') + ax.set_ylim(0.0, 1.05) + ax.grid(True, alpha=0.3) + tit='BSSM-STD latent state, N%d%s, %s'%(numNeur, R_tag, md['short_name']) + ax.set_title(tit) + ax.legend() + + # 2) population-average release and filtered drive + ax2 = self.plt.subplot(nrow,ncol,2) + ax2.plot(t_sec, mean_u_t, color='darkorange', lw=1.2, label=r'mean $u_t$') + ax2.plot(t_sec, mean_h_t, color='darkgreen', lw=1.2, label=r'mean $h_t$') + ax2.set_xlabel('time (s)') + ax2.set_ylabel('population mean') + ax2.grid(True, alpha=0.3) + ax2.set_title(r'STD release $u_t$ and filtered drive $h_t$') + ax2.legend() + + # 3) mean Bernoulli probability vs empirical rate + ax3 = self.plt.subplot(nrow,ncol,3) + ax3.scatter(mean_p_rate[exc_mask], single_rates[exc_mask], s=18, alpha=0.65, + marker='^', facecolors='none', edgecolors='red', label='Excitatory') + ax3.scatter(mean_p_rate[inh_mask], single_rates[inh_mask], s=18, alpha=0.65, + marker='s', facecolors='none', edgecolors='blue', label='Inhibitory') + lim_hi = max(float(np.max(mean_p_rate)), float(np.max(single_rates)), 1.0) + ax3.plot([0, lim_hi], [0, lim_hi], '--', color='black', lw=0.8) + ax3.set_xlabel('model-based mean rate from $p_{i,t}$ (Hz)') + ax3.set_ylabel('observed spike rate from $s_{i,t}$ (Hz)') + ax3.grid(True, alpha=0.3) + ax3.set_title('Model-based vs observed firing rate') + ax3.text( + 0.03, 0.55, r"$\hat r_i^{\mathrm{obs}}=\frac{1}{T\Delta t}\sum_{t=0}^{T-1}s_{i,t}$", + transform=ax3.transAxes, ha="left", va="center", fontsize=9, + bbox=dict(boxstyle="round,pad=0.30", facecolor="white", edgecolor="0.6", alpha=0.92), + ) + ax3.text( + 0.97, 0.03, r"$\bar r_i^{\mathrm{model}}=\frac{1}{T\Delta t}\sum_{t=0}^{T-1}p_{i,t}$", + transform=ax3.transAxes, ha="right", va="bottom", fontsize=9, + bbox=dict(boxstyle="round,pad=0.30", facecolor="white", edgecolor="0.6", alpha=0.92), + ) + ax3.legend() + + # 4) mean STD resource vs empirical rate + ax4 = self.plt.subplot(nrow,ncol,4) + ax4.scatter(mean_x_per_neuron[exc_mask], single_rates[exc_mask], s=18, alpha=0.65, + marker='^', facecolors='none', edgecolors='red', label='Excitatory') + ax4.scatter(mean_x_per_neuron[inh_mask], single_rates[inh_mask], s=18, alpha=0.65, + marker='s', facecolors='none', edgecolors='blue', label='Inhibitory') + ax4.set_xlabel(r'time-avg resource $\bar{x}_i$') + ax4.set_ylabel('empirical firing rate (Hz)') + ax4.grid(True, alpha=0.3) + ax4.set_title(r'Resource depletion vs firing rate') + ax4.legend() + + def _plot_placement_topology_ax(self, ax, trueD, md, title_txt=None): + dmd = md["dale_conf"] + L = float(dmd["placement_L"]) + H = float(dmd["placement_H"]) + if not (L > 0 and H > 0): + raise ValueError("placement_L and placement_H must be positive") + placement_ker_delta = _placement_ker_delta(dmd) + + P = np.asarray(trueD["node_positions"], dtype=float) + tau = np.asarray(trueD["node_is_inhibitory"]).reshape(-1) + if tau.dtype.kind not in "iu": + tau = tau.astype(np.int64) + N = P.shape[0] + if tau.shape[0] != N: + raise ValueError("node_positions and node_is_inhibitory length mismatch") + exc_m = tau == 0 + inh_m = tau == 1 + if not np.any(exc_m) or not np.any(inh_m): + raise ValueError("node_is_inhibitory must contain both excitatory (0) and inhibitory (1) labels") + + M = np.asarray(trueD["sign_true"]) + mask = M != 0 + + red_segs = [] + blue_segs = [] + for j in range(N): + for i in range(N): + if i == j or not mask[i, j]: + continue + seg = np.array([P[j], P[i]]) + if exc_m[j]: + red_segs.append(seg) + elif inh_m[j]: + blue_segs.append(seg) + + if red_segs: + ax.add_collection(LineCollection(red_segs, colors="red", linewidths=0.9, alpha=0.5, zorder=1)) + if blue_segs: + ax.add_collection(LineCollection(blue_segs, colors="blue", linewidths=0.9, alpha=0.5, zorder=1)) + + ax.scatter( + P[exc_m, 0], + P[exc_m, 1], + s=65, + marker="^", + facecolors="none", + edgecolors="darkred", + linewidths=1.2, + zorder=3, + ) + ax.scatter( + P[inh_m, 0], + P[inh_m, 1], + s=55, + marker="s", + facecolors="none", + edgecolors="darkblue", + linewidths=1.2, + zorder=3, + ) + + label_color = "black" + for k in range(0, N, 10): + ax.text( + P[k, 0], + P[k, 1], + str(k), + color=label_color, + fontsize=8, + ha="left", + va="bottom", + bbox=dict(boxstyle="round,pad=0.12", facecolor="white", edgecolor="none", alpha=0.9), + zorder=4, + ) + + leg_handles = [] + if red_segs: + leg_handles.append(Line2D([0], [0], color="red", alpha=0.6, lw=2, label="Exc outgoing")) + if blue_segs: + leg_handles.append(Line2D([0], [0], color="blue", alpha=0.6, lw=2, label="Inh outgoing")) + leg_handles.append( + Line2D( + [0], + [0], + marker="^", + color="w", + markerfacecolor="none", + markeredgecolor="darkred", + markersize=9, + label="Excitatory", + ) + ) + leg_handles.append( + Line2D( + [0], + [0], + marker="s", + color="w", + markerfacecolor="none", + markeredgecolor="darkblue", + markersize=8, + label="Inhibitory", + ) + ) + if title_txt is None: + title_txt = r"Spatial connectivity, $N=%d$, $\delta_{\mathrm{ker}}=%s$, %s" % ( + N, + _format_ker_delta_latex(placement_ker_delta), + md["short_name"], + ) + ax.set_title(title_txt, pad=36) + ax.set_xlabel("x (placement)") + ax.set_ylabel("y (placement)") + ax.grid(True, alpha=0.3) + ax.legend( + handles=leg_handles, + loc="lower center", + bbox_to_anchor=(0.5, 1.0), + ncol=len(leg_handles), + fontsize=9, + frameon=True, + ) + + mx, my = 0.2, 0.05 + ax.set_xlim(0.0 - mx, L + mx) + ax.set_ylim(0.0 - my, H + my) + + ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True)) + ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) + ax.xaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: "%d" % int(round(x)))) + ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: "%d" % int(round(x)))) + + def plot_placement_topology(self, trueD, md, figId=6): + """ + 2D placement: triangles = excitatory, squares = inhibitory; + outgoing edges from j→i: red if j is excitatory, blue if inhibitory (presynaptic type). + """ + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(10, 5)) + ax = fig.add_subplot(1, 1, 1) + self._plot_placement_topology_ax(ax, trueD, md) + + def _plot_connection_length_hist_ax(self, ax, trueD, md, title_txt=None): + d_off = _realized_edge_lengths(trueD) + dmd = md["dale_conf"] + d_ker = _format_ker_delta_latex(_placement_ker_delta(dmd)) + pctD = None + if d_off.size > 0: + ax.hist( + d_off, + bins=min(60, max(10, d_off.size // 5)), + color="steelblue", + edgecolor="white", + alpha=0.9, + ) + p16, p50, p84 = np.percentile(d_off, [16, 50, 84]) + pctD = {"p16": float(p16), "p50": float(p50), "p84": float(p84)} + ax.axvline(p16, color="darkorange", linestyle="--", linewidth=1.3, zorder=4) + ax.axvline(p50, color="darkgreen", linestyle="--", linewidth=1.3, zorder=4) + ax.axvline(p84, color="darkviolet", linestyle="--", linewidth=1.3, zorder=4) + pct_txt = "p16 = %.5g\np50 = %.5g\np84 = %.5g" % (p16, p50, p84) + ax.text( + 0.97, + 0.97, + pct_txt, + transform=ax.transAxes, + ha="right", + va="top", + fontsize=8, + family="monospace", + bbox=dict(boxstyle="round,pad=0.35", facecolor="white", edgecolor="0.6", alpha=0.92), + zorder=5, + ) + if title_txt is None: + title_txt = r"Realized edge lengths ($\delta_{\mathrm{ker}}=%s$)" % d_ker + ax.set_title(title_txt) + ax.set_xlabel("distance") + ax.set_ylabel("count") + ax.grid(True, alpha=0.3) + return pctD + + def connection_length_kernel_hist(self, trueD, md, figId=6): + """ + One row, three panels: edge lengths, lag-M synaptic kernel, STD recovery envelope. + """ + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(14, 4)) + gs = gridspec.GridSpec(1, 3, figure=fig, wspace=0.35) + + ax1 = fig.add_subplot(gs[0, 0]) + self._plot_connection_length_hist_ax(ax1, trueD, md) + + ax2 = fig.add_subplot(gs[0, 1]) + kappa = np.asarray(trueD["kappa_true"], dtype=np.float64).ravel() + kappa_plot = np.maximum(kappa, np.finfo(np.float64).tiny) + dt = float(md["evol_conf"]["step_size"]) + tau_s = float(md["dale_conf"]["synaptic_tau"]) + lags = dt * np.arange(1, kappa.size + 1, dtype=np.float64) + kernel_xmax_sec = 0.015 + lags_ms = lags * 1000.0 + kernel_xmax_ms = kernel_xmax_sec * 1000.0 + kernel_visible = lags <= kernel_xmax_sec + if not np.any(kernel_visible): + kernel_visible = lags <= lags[0] + kappa_sum_in_window = float(np.sum(kappa[kernel_visible])) + ax2.plot(lags_ms, kappa_plot, color="darkorange", marker="o", markersize=3.5, linewidth=1.0) + ax2.set_title(r"Fast synaptic kernel $\kappa_\ell$ ($\tau_s$)") + ax2.set_xlabel("lag time (ms)") + ax2.set_ylabel(r"kernel entry $\kappa_\ell$") + ax2.yaxis.set_major_formatter(ticker.FuncFormatter(lambda v, p: "%.3f" % v)) + ax2.set_xlim(0.0, kernel_xmax_ms) + ax2.xaxis.set_major_formatter(ticker.FuncFormatter(lambda v, p: "%d" % int(round(v)))) + ymax = float(np.max(kappa_plot[kernel_visible]) * 1.05) + ax2.set_ylim(0.0, ymax) + ax2.grid(True, alpha=0.3) + txt = "true $\\tau_s$=%.3f s\n$\\sum \\kappa$=%.4f" % (tau_s, kappa_sum_in_window) + ax2.text( + 0.97, 0.97, txt, transform=ax2.transAxes, ha="right", va="top", + fontsize=8, family="monospace", + bbox=dict(boxstyle="round,pad=0.35", facecolor="white", edgecolor="0.6", alpha=0.92), + ) + + ax3 = fig.add_subplot(gs[0, 2]) + tau_rec = float(md["dale_conf"]["std_tau_rec"]) + recovery_xmax = 1.0 + t_recovery = np.linspace(0.0, recovery_xmax, 400) + recovery_curve = 1.0 - np.exp(-t_recovery / tau_rec) + ax3.plot(t_recovery, recovery_curve, color="darkgreen", linewidth=1.3) + ax3.axhline(1.0 - np.exp(-1.0), color="black", linestyle="--", lw=0.8) + ax3.axvline(tau_rec, color="magenta", linestyle="--", lw=1.0) + ax3.set_ylim(-0.02, 1.02) + ax3.set_title(r"STD recovery set by $\tau_{rec}$") + ax3.set_xlabel("lag time (s)") + ax3.set_ylabel("recovered fraction") + ax3.set_xlim(0.0, recovery_xmax) + ax3.grid(True, alpha=0.3) + txt = r"synaptic $M\Delta t$=%.4g s" "\n" r"$\tau_{rec}$=%.4g s" "\n" r"$x(\tau_{rec})$=%.3f" % ( + float(lags[-1]), tau_rec, float(1.0 - np.exp(-1.0)) + ) + ax3.text( + 0.97, 0.08, txt, transform=ax3.transAxes, ha="right", va="bottom", + fontsize=8, family="monospace", + bbox=dict(boxstyle="round,pad=0.35", facecolor="white", edgecolor="0.6", alpha=0.92), + ) + + title_txt = _kernel_figure_canvas_title(md) + fig.suptitle(title_txt, fontsize=10, y=0.93) + fig.subplots_adjust(top=0.74) + + def topo_overview(self, trueD, md, figId=7): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor="white", figsize=(13, 10)) + gs = gridspec.GridSpec(2, 2, figure=fig, height_ratios=[1.0, 1.35], hspace=0.42, wspace=0.28) + + ax1 = fig.add_subplot(gs[0, 0]) + self._plot_signed_connectivity_matrix(ax1, trueD, md, title_txt="Signed connectivity") + + ax2 = fig.add_subplot(gs[0, 1]) + self._plot_connection_length_hist_ax(ax2, trueD, md, title_txt="Realized connection lengths") + + ax3 = fig.add_subplot(gs[1, :]) + self._plot_placement_topology_ax(ax3, trueD, md, title_txt="Spatial connectivity") + fig.suptitle(_dale_overview_title(md), fontsize=16, y=0.985) + fig.subplots_adjust(top=0.90) diff --git a/causal_net/nonStation_ver5_synapDepres/PlotterSpikesTrain.py b/causal_net/nonStation_ver5_synapDepres/PlotterSpikesTrain.py new file mode 100644 index 00000000..d13539a7 --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/PlotterSpikesTrain.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for biological experiment data visualization. + +This module provides specialized plotting capabilities for experimental +neural data analysis. The Plotter class extends PlotterBackbone to create +visualizations tailored for biological neural recordings including: +- Time series plots of neural activity patterns +- Statistical summaries and distribution analysis +- Data quality assessment plots +- Comparative analysis between experimental conditions + +Designed specifically for processing and visualizing data from biological +neural experiments, with automatic adaptation to experimental metadata. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +import numpy as np +import matplotlib.gridspec as gridspec + +#...!...!.................... +def summary_column(md): + return 'fix-me 32678' + #print(sorted(md)) + rs=md['rate_summary'] + ds=md['data_selector'] + #pprint(md) + txt='dataset: '+md['short_name'] + txt += '\nnum acc neurons: %d' % md['num_neurons'] + txt += '\ndrop neur: %d <%.1f Hz, %d >%.1f Hz' % (ds['drop_neur_by_freq_range'][0],ds['freq_range'][0],ds['drop_neur_by_freq_range'][1],ds['freq_range'][1]) + txt+='\ndata type: %s\ntime_step=%.2f sec'%(md['data_type'],md['time_step_sec']) + txt += '\nMedian rate: %.1f Hz' % rs['median_spike_rate'] + txt += '\nrate range [%.1f, %.1f Hz]' % (rs['min_spike_rate'], rs['max_spike_rate']) + txt += '\nAvg Rate: %.1f±%.1f Hz' % (rs['avg_spike_rate'], rs['std_spike_rate']) + txt += '\nAvg Fano: %.2f±%.2f' % (rs['avg_fano_factor'], rs['std_fano_factor']) + + return txt + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def freq_histo(self,spikeD,md,figId=1): + + tit='dataset '+md['short_name'] + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,4)) + + # Create gridspec with 7:3 width ratio + gs = gridspec.GridSpec(1, 2, width_ratios=[7, 3]) + + dataYield, dataRates = spikeD['spikes'], spikeD['single_rates'] + + # Compute median + median_val = np.median(dataRates) + txtM= f'median : {median_val:.2f} Hz' + + #.... freq per channel + ax = self.plt.subplot(gs[0, 0]) + chanV=np.arange(dataRates.shape[0]) + ax.bar(chanV, dataRates , width=0.8, color='g', align='center', alpha=0.7) + + ax.axhline(median_val, color='red', linestyle='--', linewidth=1) + if md['data_type']=='simDale': + ax.set_yscale('log') + ax.grid() + ax.set(xlabel='input neuron index', ylabel='avr frequency (Hz)',title=tit) + ax.text( ax.get_xlim()[1]*0.1,median_val,txtM, + color='red', ha='center', va='bottom', fontsize=10) + + txt=summary_column(md) + ax.text(0.05, 0.3, txt, transform=ax.transAxes, fontsize=12,color='b') + + #........ freq histo ..... + ax = self.plt.subplot(gs[0, 1]) + ax.hist(dataRates,bins=20) + #ax.set_yscale('log') + ax.set(ylabel='num channels', xlabel='avr frequency (Hz)',title=tit) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + # Add median text annotation + ax.text(median_val+1, ax.get_ylim()[1]*0.5,txtM, + color='red', ha='left', va='bottom', fontsize=10,rotation=90) + ax.text(0.25, 0.3, txt, transform=ax.transAxes, fontsize=11,color='m') + + ax.grid() + +#...!...!.................. + def freq_vs_time(self,rebD,md,figId=2): + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16, 7.7)) + + time_step=rebD['time_step2'] + R_sel = md['sel_spect_radius'] + + # clip time data for display + itL, itR = (md['plot']['time_rangeLR'] / time_step).astype(int) + ntime_tot = rebD['rate2D'].shape[0] + itL = max(0, min(itL, ntime_tot - 1)) + itR = max(itL + 1, min(itR, ntime_tot)) + print('iTL,R', itL,itR) + + # unpack data and clip the time range + rate2D=rebD['rate2D'][itL:itR] + timeV=rebD['timeV'][itL:itR] + + _,nchan=rate2D.shape + Tbin = time_step + xL = float(timeV[0]) + xR = float(timeV[-1] + Tbin) + if isinstance(rebD, dict) and 'pop_rate_hz' in rebD: + pop_rate_hz = np.asarray(rebD['pop_rate_hz'][itL:itR], dtype=np.float64) + else: + pop_rate_hz = np.sum(rate2D, axis=1) + # Median of per-neuron time-averaged rate (global median of all T×N bins is ~0 when activity is sparse). + mean_per_neuron_hz = np.mean(rate2D, axis=0) + medRateDisp = float(np.median(mean_per_neuron_hz)) + cntAboveMed = np.sum(rate2D > medRateDisp, axis=1) + med_sync_hz = float(np.median(pop_rate_hz)) + + tit0 = 'dataset: %s R=%.3f nchan=%d Tbin=%.2f sec' % ( + md['short_name'], + R_sel, + nchan, + time_step, + ) + if ( + md.get('spike_model') == 'B' + and 'mem_Q' in md + and 'mem_tau' in md + ): + tit0 += ' Q=%.4g tau/sec=%.4g' % ( + float(md['mem_Q']), + float(md['mem_tau']), + ) + tit0 += ' placement HxL=(%gx%g)' % ( + float(md['placement_H']), + float(md['placement_L']), + ) + + # Layout: top-count trace, heatmap, synchronicity trace. + gs = fig.add_gridspec(3, 1, height_ratios=[0.15, 0.55, 0.22], hspace=0.36) + + # ..... top plot + ax = fig.add_subplot(gs[0, 0]) + ax.bar(timeV+Tbin*.5, cntAboveMed, width=Tbin, color='teal', align='center', alpha=0.7) + ax.set( + ylabel='neurons > median', + title=f'neurons above median mean rate ({medRateDisp:.2f} Hz)', + ) + ax.set_xlim(xL,xR) + ax.set_xlabel('Time (s)') + ax.tick_params(axis='x', labelbottom=False) + ax.grid() + + # ....... main heatmap + ax = fig.add_subplot(gs[1, 0]) + #cmap='tab20c', 'Oranges' + # - - - Get the colormap and modify it --- + original_cmap = self.plt.get_cmap('tab20c') + custom_cmap = original_cmap.copy() + custom_cmap.set_under('white') + + cax = ax.imshow(rate2D.T, + aspect='auto', + cmap=custom_cmap, # Use our modified colormap + origin='lower', + extent=[xL, xR, 0, nchan], + interpolation='nearest', + vmin=1, # Set the lower bound of the colormap + vmax=rate2D.max()) # Optional: ensure the upper bound is set + + ax.set_ylabel('Neuron index') + ax.set_title(tit0 ) + ax.set_xlabel('Time (s)') + ax.tick_params(axis='x', labelbottom=False) + cbar = fig.colorbar( + cax, ax=ax, orientation='horizontal', pad=0.08, shrink=0.55, + aspect=40, anchor=(0.0, 0.5) + ) + cbar.ax.text( + 1.02, 0.5, f'Instantaneous freq (Hz) per neuron over Tbin={time_step:.2f} sec', + transform=cbar.ax.transAxes, va='center', ha='left', fontsize=10 + ) + cbar.ax.tick_params(labelsize=10) + ax.grid() + + # ..... synchronicity (yellow) below heatmap + ax = fig.add_subplot(gs[2, 0]) + ax.bar(timeV+Tbin*.5, pop_rate_hz, width=Tbin, color='gold', align='center', alpha=0.8) + ax.axhline(med_sync_hz, color='red', linestyle='--', linewidth=1.0, zorder=3) + y_rng = float(np.max(pop_rate_hz) - np.min(pop_rate_hz)) + y_pad = max(0.02 * y_rng, 0.01 * max(float(np.max(pop_rate_hz)), 1.0)) + ax.text( + xL + 0.01 * (xR - xL), + med_sync_hz + y_pad, + 'median', + va='bottom', + ha='left', + color='red', + fontsize=9, + zorder=4, + ) + ax.set( + ylabel='synchronicity (Hz)', + title=( + f'Network synchronicity from {nchan} neurons, Tbin={Tbin:.2f} sec, ' + f'median rate={med_sync_hz:.2f} Hz' + ), + ) + ax.set_xlim(xL,xR) + ax.tick_params(axis='x', labelbottom=True) + ax.grid() + ax.set_xlabel('Time (s)') + +#...!...!.................. + def state_x_vs_time(self, rebD, md, figId=3): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(16, 7.7)) + + time_step = rebD['time_step2'] + state2D_all = np.asarray(rebD['state2D'], dtype=np.float64) + + # clip time data for display, matching freq_vs_time() + itL, itR = (md['plot']['time_rangeLR'] / time_step).astype(int) + ntime_tot = state2D_all.shape[0] + itL = max(0, min(itL, ntime_tot - 1)) + itR = max(itL + 1, min(itR, ntime_tot)) + print('x-state iTL,R', itL, itR) + + x2D = state2D_all[itL:itR] + timeV = rebD['timeV'][itL:itR] + + _, nchan = x2D.shape + Tbin = time_step + xL = float(timeV[0]) + xR = float(timeV[-1] + Tbin) + + mean_x = np.mean(x2D, axis=1) + median_x = np.median(x2D, axis=1) + pct_lo_x = np.percentile(x2D, 10, axis=1) + pct_hi_x = np.percentile(x2D, 90, axis=1) + med_x = float(np.median(x2D)) + cnt_above_med = np.sum(x2D > med_x, axis=1) + + R_sel = md['sel_spect_radius'] + tit0 = 'dataset: %s R=%.3f nchan=%d Tbin=%.2f sec' % ( + md['short_name'], float(R_sel), nchan, time_step + ) + tit0 += r' STD resource $x_j(t)$' + + gs = fig.add_gridspec(3, 1, height_ratios=[0.15, 0.55, 0.22], hspace=0.36) + + # ..... top plot: count above global displayed median + ax = fig.add_subplot(gs[0, 0]) + ax.bar(timeV + Tbin * .5, cnt_above_med, width=Tbin, color='lightskyblue', + align='center', alpha=0.85) + ax.set( + ylabel='neurons > median', + title=r'neurons with $x_j$ above displayed median %.3f' % med_x, + ) + ax.set_xlim(xL, xR) + ax.set_xlabel('Time (s)') + ax.tick_params(axis='x', labelbottom=False) + ax.grid() + + # ....... main heatmap + ax = fig.add_subplot(gs[1, 0]) + cax = ax.imshow( + x2D.T, + aspect='auto', + cmap='viridis', + origin='lower', + extent=[xL, xR, 0, nchan], + interpolation='nearest', + vmin=0.0, + vmax=1.0, + ) + ax.set_ylabel('Neuron index') + ax.set_title(tit0) + ax.set_xlabel('Time (s)') + ax.tick_params(axis='x', labelbottom=False) + cbar = fig.colorbar( + cax, ax=ax, orientation='horizontal', pad=0.08, shrink=0.55, + aspect=40, anchor=(0.0, 0.5) + ) + cbar.ax.text( + 1.02, 0.5, r'rebinned STD resource $x_j$', + transform=cbar.ax.transAxes, va='center', ha='left', fontsize=10 + ) + cbar.ax.tick_params(labelsize=10) + ax.grid() + + # ..... bottom plot: population median with central 40 percentile envelope + ax = fig.add_subplot(gs[2, 0]) + ax.bar(timeV + Tbin * .5, median_x, width=Tbin, color='salmon', + align='center', alpha=0.65, label=r'median $x$') + ax.plot(timeV + Tbin * .5, pct_hi_x, color='firebrick', + linestyle='-', linewidth=0.8, label=r'10-90 percentile') + ax.plot(timeV + Tbin * .5, pct_lo_x, color='firebrick', + linestyle='-', linewidth=0.8) + ax.axhline(med_x, color='black', linestyle='--', linewidth=1.0, + label='median %.3f' % med_x) + ax.set( + ylabel=r'population $x$', + title=r'Population STD resource, median %.3f' % ( + med_x + ), + ) + ax.set_xlim(xL, xR) + ax.set_ylim(0.0, 1.02) + ax.tick_params(axis='x', labelbottom=True) + ax.grid() + ax.legend(loc='best', fontsize=9) + ax.set_xlabel('Time (s)') diff --git a/causal_net/nonStation_ver5_synapDepres/UtilDalePoisson5.py b/causal_net/nonStation_ver5_synapDepres/UtilDalePoisson5.py new file mode 100644 index 00000000..33352134 --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/UtilDalePoisson5.py @@ -0,0 +1,249 @@ +""" +Utility functions for Dale Poisson simulation data processing + +""" + +import numpy as np +import os +import time +from pprint import pprint + +def compute_consecutive_coincidence_rate(Y,time_evol): + """ + Compute the frequency of coincidences for 2 consecutive time bins for 2 different channels. + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + + Returns: + float: conc_rate_all - the overall coincidence rate + """ + num_steps, num_neurons = Y.shape + + # We need at least 2 time steps for consecutive bins + if num_steps < 2: + return 0.0 + + # Get consecutive time slices using vectorized operations + Y_t = Y[:-1, :] # Y[0:T-1, :] - current time bins + Y_t_plus_1 = Y[1:, :] # Y[1:T, :] - next time bins + + # Create boolean masks for non-zero values + mask_t = (Y_t != 0) # Shape: (T-1, N) + mask_t_plus_1 = (Y_t_plus_1 != 0) # Shape: (T-1, N) + + # Use broadcasting to compute all channel pairs at once + # mask_t[:, :, None] has shape (T-1, N, 1) + # mask_t_plus_1[:, None, :] has shape (T-1, 1, N) + # Broadcasting gives shape (T-1, N, N) for all pairs + coincidences = mask_t[:, :, None] & mask_t_plus_1[:, None, :] + + # Remove diagonal (same channel pairs) using boolean indexing + diagonal_mask = np.eye(num_neurons, dtype=bool) + coincidences[:, diagonal_mask] = False + + # Count total coincidences across all time steps + coincidence_count = np.sum(coincidences) + + conc_rate = coincidence_count / time_evol/num_neurons + return conc_rate # Hz, per neuron + +def estimate_rates(Y, dt, tau, max_samples, spect_radius, varTwindow=5, mxNn=5, verb=1): + """ + Evaluates spike statistics and estimates firing rates and Fano factors. + + tau[i] == 0 excitatory, tau[i] == 1 inhibitory; must match Y columns (same order as A_true / B). + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + dt (float): Time bin size in seconds. + tau (np.ndarray): shape (N,), int; 0 = excitatory, 1 = inhibitory. + max_samples (int): Max time samples to use for rate computation. + spect_radius (float): Reported in summary (required). + varTwindow (float): Window length (seconds) for variance / Fano. + mxNn (int): Max neurons per type in detailed printout. + verb (int): Verbose diagnostics. + """ + if verb > 0: + print("\n=== Estimating Rates & Stats ===") + if Y.shape[0] > max_samples: + if verb > 0: + print("Using %d samples (out of %d) for rate computation" % (max_samples, Y.shape[0])) + Y = Y[:max_samples] + + num_steps_sim, Nn_sim = Y.shape + tau = np.asarray(tau).reshape(-1) + if tau.shape[0] != Nn_sim: + raise ValueError("tau length %d != Y.shape[1] %d" % (tau.shape[0], Nn_sim)) + if tau.dtype.kind not in "iu": + raise ValueError("tau must have integer dtype (0=exc, 1=inh)") + tau_i = tau.astype(np.int64, copy=False) + if np.any((tau_i != 0) & (tau_i != 1)): + raise ValueError("tau entries must be 0 (excitatory) or 1 (inhibitory)") + exc_mask = tau_i == 0 + inh_mask = tau_i == 1 + if not np.any(exc_mask) or not np.any(inh_mask): + raise ValueError("tau must contain at least one excitatory and one inhibitory neuron") + num_excite = int(np.sum(exc_mask)) + num_inhib = int(np.sum(inh_mask)) + + time_evol = num_steps_sim * dt + if verb > 0: + print('steps num_steps=%d, time_evol=%.1f sec, Nn=%d (%d Excit, %d Inhib)' % (num_steps_sim, time_evol, Nn_sim, num_excite, num_inhib)) + + # Compute windowed spike counts and statistics over non-overlapping time windows of length varTwindow (sec), per neuron + window_size = max(1, int(round(varTwindow / dt))) + num_windows = num_steps_sim // window_size + if verb > 0: + print('variance computation: num_windows=%d, window_size=%d' % (num_windows, window_size)) + if num_windows <= 1: + raise ValueError("varTwindow=%s is too large for data (num_windows=%d). Need at least 2 windows for variance." % (varTwindow, num_windows)) + Y_trim = Y[:num_windows * window_size] + Y_win = Y_trim.reshape(num_windows, window_size, Nn_sim) + + # Sum spike counts over each window: shape (num_windows, n_neurons) + window_spike_counts = np.sum(Y_win, axis=1) + + # Fano Factor: Var[spike count] / Mean[spike count] over windows + mean_window_spike_counts = np.mean(window_spike_counts, axis=0) + var_window_spike_counts = np.var(window_spike_counts, axis=0) + fano_factor = np.divide(var_window_spike_counts, mean_window_spike_counts, + out=np.zeros_like(var_window_spike_counts), + where=mean_window_spike_counts != 0) + + # Convert window spike counts to rates (Hz) for rate variance + window_rates = window_spike_counts / (window_size * dt) # shape: (num_windows, n_neurons), Hz + single_rates_var = np.var(window_rates, axis=0) # variance of rate across windows, per neuron (Hz^2) + + # Compute raw arrays (full time series) + spike_counts = np.sum(Y, axis=0) + spike_rates = spike_counts / time_evol + mean_counts_per_bin = np.mean(Y, axis=0) + + # Compute consecutive coincidence rate + start_time = time.time() + conc_rate_per_neuron = compute_consecutive_coincidence_rate(Y,time_evol) + elapsed_time = time.time() - start_time + if verb > 0: + print('Coincidence rate %.2g Hz, Y.shape=%s elaT %.3f sec' % (conc_rate_per_neuron, Y.shape, elapsed_time)) + + # Compute all population statistics once + med_rate_all = np.median(spike_rates) + avg_rate_all = float(np.mean(spike_rates)) + std_rate_all = float(np.std(spike_rates)) + avg_fano_all = float(np.mean(fano_factor)) + std_fano_all = float(np.std(fano_factor)) + avg_rate_excit = float(np.mean(spike_rates[exc_mask])) + std_rate_excit = float(np.std(spike_rates[exc_mask])) + avg_fano_excit = float(np.mean(fano_factor[exc_mask])) + std_fano_excit = float(np.std(fano_factor[exc_mask])) + avg_rate_inhib = float(np.mean(spike_rates[inh_mask])) + std_rate_inhib = float(np.std(spike_rates[inh_mask])) + avg_fano_inhib = float(np.mean(fano_factor[inh_mask])) + std_fano_inhib = float(np.std(fano_factor[inh_mask])) + + # Build stats dictionary with computed values + stats_dict = { + 'num_steps': num_steps_sim, + 'time_evol_sec': time_evol, + 'time_step_sec': dt, + 'num_neurons': Nn_sim, + 'num_excitatory': num_excite, + 'num_inhibitory': num_inhib, + 'avg_spike_rate_all': avg_rate_all, + 'std_spike_rate_all': std_rate_all, + 'avg_fano_factor_all': avg_fano_all, + 'std_fano_factor_all': std_fano_all, + 'avg_spike_rate_excit': avg_rate_excit, + 'std_spike_rate_excit': std_rate_excit, + 'avg_fano_factor_excit': avg_fano_excit, + 'std_fano_factor_excit': std_fano_excit, + 'avg_spike_rate_inhib': avg_rate_inhib, + 'std_spike_rate_inhib': std_rate_inhib, + 'avg_fano_factor_inhib': avg_fano_inhib, + 'std_fano_factor_inhib': std_fano_inhib, + 'conc_rate_per_neuron': float(conc_rate_per_neuron), + 'median_spike_rate_all': med_rate_all + } + + # Printing detailed stats for individual neurons + mxE = min(mxNn, num_excite) + mxI = min(mxNn, num_inhib) + + exc_show = np.where(exc_mask)[0][:mxE] + inh_show = np.where(inh_mask)[0][:mxI] + + if verb > 0: + print('\n--- Stats for %d Excitatory Neurons (by index order) ---' % len(exc_show)) + np.set_printoptions(precision=2) + print('Total Spike Counts: %s' % spike_counts[exc_show]) + print('Mean Firing Rate (Hz): %s' % spike_rates[exc_show]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[exc_show])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[exc_show])) + print('Fano Factor (Var/Mean): %s' % fano_factor[exc_show]) + + print('\n--- Stats for %d Inhibitory Neurons (by index order) ---' % len(inh_show)) + np.set_printoptions(precision=2) + print('Total Spike Counts: %s' % spike_counts[inh_show]) + print('Mean Firing Rate (Hz): %s' % spike_rates[inh_show]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[inh_show])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[inh_show])) + print('Fano Factor (Var/Mean): %s' % fano_factor[inh_show]) + + R_tag = ', R=%.3f' % float(spect_radius) + print('\n--- Population Summary Statistics%s ---' % R_tag) + print('All (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (Nn_sim, stats_dict['avg_spike_rate_all'], stats_dict['std_spike_rate_all'], stats_dict['avg_fano_factor_all'], stats_dict['std_fano_factor_all'])) + print('Excit (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_excite, stats_dict['avg_spike_rate_excit'], stats_dict['std_spike_rate_excit'], stats_dict['avg_fano_factor_excit'], stats_dict['std_fano_factor_excit'])) + if num_inhib > 0: + print('Inhib (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_inhib, stats_dict['avg_spike_rate_inhib'], stats_dict['std_spike_rate_inhib'], stats_dict['avg_fano_factor_inhib'], stats_dict['std_fano_factor_inhib'])) + + # Print key summary values using dictionary + summary_keys = ['conc_rate_per_neuron', 'median_spike_rate_all'] + for key in summary_keys: + if key == 'conc_rate_per_neuron': + print('Coincidence rate per neuron %.2g Hz' % stats_dict[key]) + elif key == 'median_spike_rate_all': + print('Median rate %.2f Hz\n' % stats_dict[key]) + + # Part 2: Compute frequency sorting + neur_freq_index = np.argsort(spike_rates) + + rates_dict = { + 'single_rates': spike_rates, + 'neur_freqIdx':neur_freq_index, + 'single_rates_var': single_rates_var, + 'single_fano_fact': fano_factor + } + + return stats_dict, rates_dict, neur_freq_index + + + +def get_offdiag_triplets(A, isPos=True): + """ + Returns positive/negative, off-diagonal elements of 2D array A as array of [i, j, value] triplets. + + Parameters: + A : 2D numpy array + isPos : bool, if True only return positive values (>0), if False return all negative values + + Returns: + Array of shape (n, 3) where n is the number of valid off-diagonal elements + Each row is [row_index, col_index, value] + """ + # Create mask for off-diagonal elements + offdiag_mask = ~np.eye(A.shape[0], A.shape[1], dtype=bool) + + # Add value condition based on isPos + if isPos: + value_mask = A > 0 + else: + value_mask = A < 0 + + # Combine masks + mask = value_mask & offdiag_mask + + i_indices, j_indices = np.where(mask) + values = A[i_indices, j_indices] + + return np.column_stack([i_indices, j_indices, values]) diff --git a/causal_net/nonStation_ver5_synapDepres/Util_pseudospectra.py b/causal_net/nonStation_ver5_synapDepres/Util_pseudospectra.py new file mode 100755 index 00000000..23fc3e1f --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/Util_pseudospectra.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +import numpy as np +import time + +def create_example_matrices(size=100): + A_base = np.diag(np.linspace(-5, -0.5, size)) + np.diag(np.ones(size - 2) * 4, k=2) + sparsity = 0.15 + random_connections = np.random.randn(size, size) + mask = np.random.rand(size, size) > sparsity + random_connections[mask] = 0 + A_true = A_base + random_connections * 0.1 + + noise = np.random.randn(size, size) * 0.10 + noise_mask = np.random.rand(size, size) > sparsity + noise[noise_mask] = 0 + A_fit = A_true + noise + + return A_true, A_fit + +def compute_pseudospectrum(A, npts, minY): + print(f'computing pseudospectrum A:{A.shape} .... ') + eigs = np.linalg.eigvals(A) + + # Calculate bounds + real_min, real_max = np.real(eigs).min(), np.real(eigs).max() + imag_min, imag_max = np.imag(eigs).min(), np.imag(eigs).max() + + real_pad = (real_max - real_min) * 0.2 + imag_pad = (imag_max - imag_min) * 0.2 + + bbox = [ + real_min - real_pad, + min(real_max + real_pad, 1.0), + imag_min, + imag_max + imag_pad + ] + + x_coords = np.linspace(bbox[0], bbox[1], npts) + y_coords = np.linspace(minY, bbox[3], npts) + #y_coords = np.linspace(bbox[2], bbox[3], npts) + X, Y = np.meshgrid(x_coords, y_coords) + Z = X + 1j * Y + + sigma_grid = np.zeros_like(Z, dtype=float) + I = np.eye(A.shape[0]) + + for i in range(npts): + for j in range(npts): + z = Z[i, j] + s = np.linalg.svd(z * I - A, compute_uv=False) + sigma_grid[i, j] = s[-1] + + return X, Y, sigma_grid, eigs + +def plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin): + levels = np.logspace(-2.5, -0.5, 10) # Use 10 contour levels + + # Plot the normal contour lines + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + + # Create a mask for the area greater than epsMin + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) # Masking areas greater than epsMin + + # Fill the area below epsMin + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + # Scatter plot for eigenvalues + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + # Vertical line at Re=0 + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + # Set titles and labels + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + # Additional formatting + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(bottom=minY) # Clip at minY + + +def main(size=100, minY=-0.9, epsMin=0.032): + A_true, A_fit = create_example_matrices(size) + + fig, axes = plt.subplots(1, 2, figsize=(16, 7)) + + matrices_to_plot = [ + ('True Matrix ($A_{true}$)', A_true), + ('Distorted Matrix ($A_{fit}$) ', A_fit) + ] + + print("Starting pseudospectrum calculations...", minY) + start_time = time.time() + + for i, (title, A) in enumerate(matrices_to_plot): + ax = axes[i] + print("Processing '{}'...".format(title)) + + npts = 80 # Grid resolution + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin) # Pass epsMin + + total_time = time.time() - start_time + print(f"Calculation finished in {total_time:.2f} seconds.") + + fig.suptitle('Pseudospectrum Comparison (Clipped Imaginary Part)', fontsize=16) + plt.tight_layout(rect=[0, 0.03, 1, 0.95]) + + outFile = 'out/pseudospectra_comparison.png' + plt.savefig(outFile) + print(f"Plot saved as '{outFile}'") + plt.close(fig) + +# --- Main execution --- +if __name__ == '__main__': + + import matplotlib.pyplot as plt + + main(size=50, minY=-0.5, epsMin=0.024) diff --git a/causal_net/nonStation_ver5_synapDepres/docs/bssm_std_model_cartoon.png b/causal_net/nonStation_ver5_synapDepres/docs/bssm_std_model_cartoon.png new file mode 100644 index 00000000..960c5374 Binary files /dev/null and b/causal_net/nonStation_ver5_synapDepres/docs/bssm_std_model_cartoon.png differ diff --git a/causal_net/nonStation_ver5_synapDepres/docs/buildGen.sh b/causal_net/nonStation_ver5_synapDepres/docs/buildGen.sh new file mode 100755 index 00000000..decc943b --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/docs/buildGen.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +set -e + +doc_name=nonStation-ver5_synDeprGen + +is_ubuntu_or_macos=0 +if [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + is_ubuntu_or_macos=1 +elif [ "$(uname -s)" = "Darwin" ]; then + is_ubuntu_or_macos=1 +fi + +if ! command -v pdflatex >/dev/null 2>&1; then + module load texlive +fi + +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" +pdflatex -interaction=nonstopmode -halt-on-error "$doc_name" + +pdf_file="${doc_name}.pdf" +echo "Built $pdf_file" + +if [ "$is_ubuntu_or_macos" -eq 0 ]; then + gv "$pdf_file" +fi diff --git a/causal_net/nonStation_ver5_synapDepres/docs/nonStation-ver5_synDeprGen.pdf b/causal_net/nonStation_ver5_synapDepres/docs/nonStation-ver5_synDeprGen.pdf new file mode 100644 index 00000000..da699219 Binary files /dev/null and b/causal_net/nonStation_ver5_synapDepres/docs/nonStation-ver5_synDeprGen.pdf differ diff --git a/causal_net/nonStation_ver5_synapDepres/docs/nonStation-ver5_synDeprGen.tex b/causal_net/nonStation_ver5_synapDepres/docs/nonStation-ver5_synDeprGen.tex new file mode 100644 index 00000000..305f10b6 --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/docs/nonStation-ver5_synDeprGen.tex @@ -0,0 +1,683 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{graphicx} +\usepackage{xcolor} +\usepackage{booktabs} +\usepackage{algorithm} +\usepackage{algorithmic} + +\geometry{margin=0.8in} + +\title{Synaptic Depression Modeling \\ Poisson spike generation and model fitting} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +We describe a real-time adaptive Poisson state-space model (PSSM) for a network of +interacting neurons with short-term synaptic depression (STD) and a lag-$M$ synaptic +memory kernel. The STD mechanism is represented as a latent dynamical resource variable +that is updated at each discrete time step, coupled to observed spike counts +through a Poisson likelihood. For context, we also briefly review the Tsodyks--Markram +(TM) resource model, which provides the continuous-time biological foundation for the +discrete STD dynamics used in the PSSM. +\end{abstract} + +%========= +\section{Short-Term Synaptic Depression (STD) in Cortical Synapses} +\label{sec:std_intro} + +Short-term synaptic depression (STD) is a fast, reversible decrease in synaptic strength +that occurs when a presynaptic cortical neuron fires repeatedly over short time windows +(typically from milliseconds to seconds). + +\paragraph{Activity-decaying phenomenon.} +When presynaptic firing is high, the postsynaptic response (EPSCs or IPSCs) becomes +progressively smaller for successive spikes, producing a decaying envelope of synaptic +efficacy during bursts. After presynaptic activity pauses, synaptic strength gradually +recovers to baseline levels. + +\paragraph{Main biological mechanism.} +STD is commonly modeled as depletion of synaptic ``resources'' required for effective +neurotransmitter release. At each presynaptic spike, release depends on +(i) the availability of readily releasable vesicles (or release-ready sites) and +(ii) an effective probability of release (often related to intracellular calcium dynamics). +Sustained activity depletes resources and reduces future release, while periods of silence +allow recovery. + +\paragraph{Relevant timescales.} +Recovery is typically on intermediate timescales ($\sim 100$\,ms to a few seconds), +depending on synapse type and experimental protocol. + +\paragraph{Functional consequences.} +STD acts as a rate-dependent synaptic filter: it attenuates sustained high-rate input, +emphasizes transient changes in firing, and contributes to network stabilization by +preventing runaway excitation. + +%========= +\section{Computational Models with STD} +\label{sec:std_models} + +We describe two STD frameworks. The TM model (Section~\ref{sec:tm_std}) is included +for conceptual context only. The PSSM (Section~\ref{sec:bssm_std}) is the target +implementation model. + +The following notation is shared across both models: +$U\in(0,1]$ is the utilization parameter (fraction of available resources used per spike), +and $\tau_{\mathrm{rec}}$ is the STD recovery time constant. + +%========= +\subsection{Tsodyks--Markram (TM) Resource Model: Conceptual Background} +\label{sec:tm_std} + +The \emph{Tsodyks--Markram (TM)} model~\cite{Tsodyks1997} is a widely used continuous-time +resource model for \emph{short-term plasticity (STP)}. In the STD-only version, the synapse +has a single \emph{available-resource} variable $x(t)\in[0,1]$. Between presynaptic spikes, +resources recover exponentially: + +\begin{figure}[t] +\centering +\includegraphics[width=\textwidth]{tm_std_cartoon.png} +\caption{Cartoon of the Tsodyks--Markram STD resource model. Presynaptic spikes +consume a fraction $U$ of the available resource $x$, making successive release +events smaller; during silence, the resource recovers toward $x=1$ with time +constant $\tau_{\mathrm{rec}}$.} +\label{fig:tm_std_cartoon} +\end{figure} + +\noindent +The cartoon in Fig.~\ref{fig:tm_std_cartoon} summarizes the STD-only TM mechanism: +spikes multiply the available resource by $(1-U)$, while inter-spike intervals allow +the resource to recover toward one. + +\noindent +Between presynaptic spikes, +\begin{equation} +\frac{dx}{dt}=\frac{1-x}{\tau_{\mathrm{rec}}}. +\label{eq:tm_recovery} +\end{equation} +At a presynaptic spike time $t=t_k$, a fraction $U$ of available resources is consumed: +\begin{equation} +x(t_k^+)=x(t_k^-)\,(1-U), +\label{eq:tm_spike_update} +\end{equation} +and the released neurotransmitter amplitude at that spike is +\begin{equation} +r(t_k)=U\,x(t_k^-). +\label{eq:tm_release} +\end{equation} +The TM model provides the biological motivation for the discrete STD resource update +used in the PSSM below (Eq.~\ref{eq:std_ssm_state}), which is the exact discretization +of Eq.~\eqref{eq:tm_recovery} combined with Eq.~\eqref{eq:tm_spike_update}. + +\paragraph{Interpretation of $U$.} +The utilization parameter $U$ is a fraction of the resources available at the +moment of a spike, not a fixed fraction of the original full pool. For example, +if $U=0.5$ and several spikes arrive with negligible recovery between them, a +fully recovered synapse releases $0.5$ of the full pool on the first spike, +then $0.25$ on the second, then $0.125$ on the third, and so on. Thus the axon +can still generate more than two closely spaced spikes; STD only makes the +effective synaptic release decrease geometrically during the burst. + +\subsection{Real-Time Adaptive Poisson State-Space Network Model (PSSM) with STD + and Lag-$M$ Kernel} +\label{sec:bssm_std} + +The \emph{Poisson state-space model (PSSM)} represents a network of $N$ neurons using +(i) latent dynamical states encoding STD resources and (ii) a Poisson observation model +for per-bin spike counts conditioned on those states. This is the primary model of +interest. At physiological firing rates ($30$--$50$\,Hz) and bin width +($\Delta t = 1$\,ms), the expected spike count per bin is +$\lambda \approx r\Delta t = 0.03$--$0.05$. In that regime, most bins still contain +zero or one spike, but the Poisson model also allows occasional bins with two or more +spikes instead of enforcing a hard binary cap. + +\begin{figure}[t] +\centering +\includegraphics[width=\textwidth]{bssm_std_model_cartoon.png} +\caption{Cartoon of the full PSSM with STD and a lag-$M$ synaptic kernel. +Presynaptic spike counts $s_{j,t}$ modulate the latent resource $x_{j,t}$, +produce STD-weighted release $u_{j,t}$, pass through a causal synaptic kernel to +form filtered drive $h_{j,t}$, and enter the postsynaptic Poisson observation +model through signed recurrent weights $W_{ij}$ and baseline $B_i$ to generate +the output spike count $s_{i,t}$.} +\label{fig:bssm_std_cartoon} +\end{figure} + +\noindent +Figure~\ref{fig:bssm_std_cartoon} summarizes the sequential data-flow view of +the model: the recurrent weighted input determines the postsynaptic spike +count intensity, sampled presynaptic spikes generate release from the current STD +resource, and that release updates both the next STD state and the filtered +drive available in the next bin. + +\paragraph{Discrete time and notation.} +Let $t\in\{0,1,\dots,T-1\}$ index time bins of width $\Delta t$. +For each neuron $j$, the primary observable is the nonnegative integer \emph{spike count} +\begin{equation} +s_{j,t}\in\{0,1,2,\dots\}, +\end{equation} +which records how many spikes neuron $j$ emitted in bin $t$. The downstream efficacy of +those spikes is determined separately by the latent STD resource variable +$x_{j,t}\in[0,1]$, which encodes the fraction of neurotransmitter available for release +at the start of bin $t$ (see Eq.~\ref{eq:std_ssm_state}). Spikes that occur when +resources are depleted ($x_{j,t}\approx 0$) contribute little to postsynaptic drive, +whereas spikes after full recovery ($x_{j,t}\approx 1$) contribute maximally. + +\paragraph{STD latent state dynamics.} +Each presynaptic neuron $j$ has its own latent resource variable $x_{j,t}$, +tracking the STD state of that neuron's outgoing synapses independently of all +other neurons. The full network state is the collection $\{x_{j,t}\}_{j=1}^{N}$, +with each variable evolving as the exact discretization of the TM continuous-time +recovery (Eq.~\ref{eq:tm_recovery}) combined with spike-triggered consumption +(Eq.~\ref{eq:tm_spike_update}): +\begin{equation} +\begin{aligned} +x^{\mathrm{dep}}_{j,t} +&= +\underbrace{x_{j,t}(1-U)^{s_{j,t}}}_{\text{spike-triggered depletion}}, +\\ +x_{j,t+1} +&= +\underbrace{1-\big(1-x^{\mathrm{dep}}_{j,t}\big) +e^{-\Delta t/\tau_{\mathrm{rec}}}}_{\text{exponential recovery after depletion}}, +\qquad j=1,\dots,N . +\end{aligned} +\label{eq:std_ssm_state} +\end{equation} +Release during bin $t$ uses the pre-bin resource $x_{j,t}$. If $s_{j,t}>0$ spikes occur +within the bin, each spike consumes a fraction $U$ of the resource remaining at that +moment, so the net depletion factor is $(1-U)^{s_{j,t}}$. Because +$\tau_{\mathrm{rec}}\gg \Delta t$ in the intended regime, recovery within a bin is +neglected and the depleted resource then recovers toward $1$ over the interval to +bin $t+1$. Because the update depends only on neuron $j$'s own spike count, resource +variables of different neurons are +conditionally independent given their respective spike histories, allowing all +$N$ updates to proceed in parallel. + +\paragraph{STD-modulated release.} +The effective synaptic release from neuron $j$ in bin $t$ is: +\begin{equation} +u_{j,t}=x_{j,t}\Bigl[1-(1-U)^{s_{j,t}}\Bigr]. +\label{eq:std_release} +\end{equation} +This is the total released resource from all spikes in bin $t$: it is zero when +$s_{j,t}=0$, small when resources are depleted by recent activity, and approaches +$x_{j,t}$ when many spikes arrive in the same bin. + +\paragraph{Lag-$M$ synaptic kernel.} +Synaptic currents have temporal extent beyond a single bin. We convolve the +STD-modulated release with a causal lag-$M$ kernel +$\bm{\kappa}=(\kappa_0,\dots,\kappa_{M-1})^\top$, $\kappa_\ell\geq 0$, encoding how +a release event at lag $\ell+1$ bins contributes to postsynaptic drive at time $t$. +The effective STD-filtered synaptic input from neuron $j$ at time $t$ is: +\begin{equation} +h_{j,t} = \sum_{\ell=0}^{M-1}\kappa_\ell\,u_{j,t-1-\ell}. +\label{eq:lag_kernel} +\end{equation} +Here $\ell=0$ is one bin back and $M\Delta t$ is the finite lag-history span for the fast synaptic current. +For an exponential synaptic kernel, practical choices take $M\Delta t$ to be +several synaptic time constants $\tau_s$ so that the current tail is small at +the truncation boundary. This lag window is distinct from the STD recovery +memory: the slower recovery governed by $\tau_{\mathrm{rec}}$ is carried by the +latent state $x_{j,t}$ rather than by the finite synaptic kernel. Practical +kernel choices include exponential decay, raised cosine basis functions, or a +free learnable filter (see Section~\ref{sec:gen_kernel} for the generator +choice). + +\paragraph{Poisson observation model.} +For each postsynaptic neuron $i$, the conditional expected spike count is: +\begin{equation} +\lambda_{i,t} += +\exp\!\left( +B_i ++ +\sum_{\substack{j=1\\j\neq i}}^{N} +W_{ij}\,h_{j,t} +\right), +\label{eq:bernoulli_obs} +\end{equation} +where $B_i$ is a baseline log-mean-count for neuron $i$, and $W_{ij}$ is the synaptic weight from neuron $j$ +to neuron $i$ (positive for excitatory, negative for inhibitory connections). +Self-connections are excluded by the constraint $j\neq i$. +The spike count is then drawn as: +\begin{equation} +s_{i,t} \sim \mathrm{Poisson}(\lambda_{i,t}). +\label{eq:bernoulli_likelihood} +\end{equation} +For $t-1-\ell<0$, pre-history values $u_{j,t-1-\ell}$ are set to zero, corresponding +to a fully recovered synapse with no prior spiking. +Because all kernel coefficients $\kappa_\ell\geq0$ and all release values +$u_{j,t}\geq0$, the filtered drive satisfies $h_{j,t}\geq0$ at all times; +the implementation enforces this with a floor clamp after each recursive update. + +\paragraph{Summary of model parameters.} +The full parameter set of the PSSM is $\theta=\{B_i, W_{ij}, \kappa_\ell, U, +\tau_{\mathrm{rec}}\}$, where $B_i$ and $W_{ij}$ are neuron- and synapse-specific, +while this model assumes a single shared utilization parameter $U$ for all +presynaptic neurons and outgoing synapses. The recovery time +$\tau_{\mathrm{rec}}$ is also treated as shared in the generator and estimation +procedure below. In the generator +(Section~\ref{sec:generator}), all parameters are fixed to known ground-truth +values; in inference, they constitute the quantities to be estimated. + +\paragraph{Real-time and adaptive aspects.} +The term \emph{real-time} refers to the fact that both the latent STD state +$x_{j,t}$ and parameter estimates can be updated sequentially as new spike +counts $s_{i,t}$ arrive, without reprocessing the full data history. The term +\emph{adaptive} refers to the property that $x_{j,t}$ automatically tracks recent +presynaptic spiking through Eq.~\eqref{eq:std_ssm_state}: resources deplete during +bursts and recover during silence, so the effective synaptic gain $u_{j,t}$ adapts +online to the presynaptic firing history. Together, these properties make the PSSM +suitable for closed-loop neural decoding or real-time inference scenarios. + +%========= +\section{Data Generator: Cortical Network with STD} +\label{sec:generator} + +To validate the PSSM and test inference procedures, we generate synthetic spike +trains from a network that mimics key properties of mammalian cortical dynamics. +The generator uses the same PSSM structure as the inference model +(Section~\ref{sec:bssm_std}), but with all parameters fixed to known ground-truth +values, enabling direct comparison between true and estimated quantities. + +%========= +\subsection{Network Architecture} +\label{sec:gen_architecture} + +We consider a network of $N\in[50, 500]$ neurons partitioned following Dale's law +into $N_E = 0.8\,N$ excitatory and $N_I = 0.2\,N$ inhibitory neurons. +The code constructs the initial synaptic matrix $W\in\mathbb{R}^{N\times N}$ +with a spatial connections Dale generator. Neurons are +placed on a two-dimensional grid over $[0,L]\times[0,H]$ with minimum spacing +$d_{\min}$, and their order is sorted by position. For each presynaptic neuron +$j$, an out-degree $k_j$ is drawn from the configured range and $k_j$ distinct +postsynaptic targets are sampled without replacement with probability +proportional to $D_{ij}^{-\delta}$, excluding $i=j$. This gives nearby neurons +higher connection probability while keeping the diagonal zero. + +The nonzero weights obey Dale's law column-by-column: +\begin{equation} +W^{(0)}_{ij} = +\begin{cases} ++a_{ij} & \text{if neuron } j \text{ is excitatory},\\ +-\frac{N_E}{N_I}\,a_{ij} & \text{if neuron } j \text{ is inhibitory}, +\end{cases} +\qquad +a_{ij}\sim\mathrm{Uniform}(1-v,1+v), +\label{eq:weight_draw} +\end{equation} +for sampled edges only, with $W^{(0)}_{ii}=0$. Finally, the whole matrix is +rescaled to the requested spectral radius $R$, +\begin{equation} +W = W^{(0)}\,\frac{R}{\rho(W^{(0)})}, +\label{eq:weight_matrix} +\end{equation} +which preserves sparsity, signs, Dale's law, and the zero diagonal. + + +%========= +\subsection{Synaptic Kernel} +\label{sec:gen_kernel} + +We use an \emph{exponential decay kernel} as the ground-truth synaptic filter: +\begin{equation} +\kappa_\ell = (1-\alpha)\,\alpha^{\ell}, \qquad \ell = 0,1,\dots,M-1, +\label{eq:exp_kernel} +\end{equation} +where $\alpha = e^{-\Delta t/\tau_s}$ and $\tau_s$ is the synaptic time constant. +The index $\ell=0$ corresponds to the most recent lag (one bin back), and +$\ell=M-1$ to the oldest lag retained. +This kernel corresponds to a standard single-compartment synaptic current model +and is the exact discrete-time analog of the continuous-time shot-noise synaptic +model. It is normalized so that $\sum_{\ell=0}^{M-1}\kappa_\ell \approx 1$ for +$M\Delta t \gg \tau_s$. + +A key computational advantage of this choice is that the lag-$M$ convolution +in Eq.~\eqref{eq:lag_kernel} can be evaluated by a tail-corrected recursive +update: +\begin{equation} +h_{j,t+1} += +\alpha\,h_{j,t} + (1-\alpha)\,u_{j,t} +- +(1-\alpha)\alpha^{M}\,u_{j,t-M}, +\label{eq:recursive_kernel} +\end{equation} +with $u_{j,t-M}=0$ when $t-M<0$. This recurrence is exactly equivalent to the +finite lag-$M$ convolution for the exponential kernel and avoids explicitly +summing over all $M$ lags at every time step. The decay rate of $\kappa_\ell$ +is set by the fast synaptic time constant $\tau_s$, whereas the slower STD time +constant $\tau_{\mathrm{rec}}$ enters through the latent state update in +Eq.~\eqref{eq:std_ssm_state}. +In the fitting code, this exponential kernel is fixed by command-line arguments: +\texttt{--synaptic\_tau} sets $\tau_s$, and \texttt{--kernel\_len\_steps} sets +the retained lag length $M$. + +%========= +\subsection{Ground-Truth Parameter Values} +\label{sec:gen_params} + +Table~\ref{tab:gen_params} lists the ground-truth parameter values used in the +generator, motivated by mammalian cortical physiology. + +\begin{table}[h] +\centering +\small +\setlength{\tabcolsep}{3pt} +\begin{tabular}{@{}llll@{}} +\toprule +Parameter & Symbol & Value & Rationale \\ +\midrule +Time bin width & $\Delta t$ & $1$\,ms & Resolves spikes; keeps $\lambda$ small \\ +Synaptic time constant & $\tau_s$ & $3$--$5$\,ms & AMPA-like fast excitation \\ +STD recovery time & $\tau_{\mathrm{rec}}$ & $200$--$500$\,ms & Typical cortical STD \\ +Utilization parameter & $U$ & $0.3$--$0.5$ & Shared moderate depression \\ +Weight scale & $v, R$ & $v=0.4,\ R=0.90$ & Spectral-radius scaling \\ +Baseline log-count & $B_i$ & sampled idle-rate range & User-adjusted baseline drive \\ +Kernel length & $M$ & $\gtrsim 3$--$5\,\tau_s/\Delta t$ & Fast synaptic lag window \\ +\bottomrule +\end{tabular} +\caption{Ground-truth generator parameters. Ranges reflect variability across +synapse types and network sizes. The kernel length $M$ controls the finite lag +history for the fast synaptic filter and is separate from the STD recovery time +$\tau_{\mathrm{rec}}$. The baseline $B_i$ is expressed as a log-mean-count consistent +with the Poisson observation model (Eq.~\ref{eq:bernoulli_obs}). +Each neuron's idle rate is drawn independently and uniformly from +$[\text{idleRate}_{\min}, \text{idleRate}_{\max}]$, and the corresponding +log-mean-count is $B_i = \log(r_i\Delta t)$.} +\label{tab:gen_params} +\end{table} + +%========= +\subsection{Generation Procedure} +\label{sec:gen_procedure} + +Given the fixed parameters, synthetic spike counts are generated sequentially +for $t = 0, 1, \dots, T-1$ as follows. + +\begin{algorithm} +\caption{Synthetic network spike train generator} +\label{alg:generator} +\begin{algorithmic}[1] +\STATE \textbf{Initialize:} $x_{j,0} \leftarrow 1$ for all $j$ (fully recovered); + $h_{j,0} \leftarrow 0$ and $u_{j,t}\leftarrow0$ for $t<0$; draw $W$, $\{B_i\}$ as in + Eqs.~\eqref{eq:weight_draw}--\eqref{eq:weight_matrix}. +\FOR{$t = 0, 1, \dots, T-1$} + \FOR{each neuron $i = 1,\dots,N$} + \STATE Compute expected spike count: + $\lambda_{i,t} \leftarrow \exp\!\left(B_i + + \sum_{j \neq i} W_{ij}\,h_{j,t}\right)$ + \hfill(Eq.~\ref{eq:bernoulli_obs}) + \STATE Draw spike count: + $s_{i,t} \sim \mathrm{Poisson}(\lambda_{i,t})$ + \hfill(Eq.~\ref{eq:bernoulli_likelihood}) + \ENDFOR + \FOR{each neuron $j = 1,\dots,N$} + \STATE Compute STD-modulated release from pre-spike resources: + $u_{j,t} \leftarrow x_{j,t}\bigl[1-(1-U)^{s_{j,t}}\bigr]$ + \hfill(Eq.~\ref{eq:std_release}) + \STATE Update STD resource: + $x_{j,t+1} \leftarrow + 1-\big(1-x_{j,t}(1-U)^{s_{j,t}}\big)e^{-\Delta t/\tau_{\mathrm{rec}}}$ + \hfill(Eq.~\ref{eq:std_ssm_state}) + \STATE Update recursive synaptic filter: + $h_{j,t+1} \leftarrow + \alpha h_{j,t} + (1-\alpha)u_{j,t} + -(1-\alpha)\alpha^M u_{j,t-M}$ + \hfill(Eq.~\ref{eq:recursive_kernel}) + \ENDFOR +\ENDFOR +\STATE \textbf{Output:} spike counts $\{s_{j,t}\}$, ground-truth STD states + $\{x_{j,t}\}$, weight matrix $W$. +\end{algorithmic} +\end{algorithm} + +%========= +\subsection{Network Stability Check} +\label{sec:gen_stability} + +With recurrent excitatory connections, the network can exhibit runaway activity +if synaptic weights are too large. Before generating long recordings, we recommend +a short pilot run ($\sim 1$\,s) and verification that mean firing rates across +neurons remain in the target range of $30$--$50$\,Hz. If rates exceed this +range, one should reduce the spectral-radius target $R$ or increase $U$: stronger +STD provides a natural activity-dependent stabilizing mechanism by reducing +synaptic efficacy during high-rate episodes. + +%========= +\section{Parameter Estimation: Two-Block Coordinate Descent} +\label{sec:param_estimation} + +The true BSSM parameter set is +$\theta = \{B_i, W_{ij}, \tau_\text{rec}, U, \tau_s\}$. The estimated set is +$\tilde{\theta} = \{\tilde{B}_i, \tilde{W}_{ij}, \tilde{\tau}_\text{rec}, \tilde{U}\}$, +with the exponential synaptic kernel fixed rather than fitted. In +\texttt{fit5\_BSSM\_STD\_blocks.py}, \texttt{--synaptic\_tau} sets $\tau_s$ and +\texttt{--kernel\_len\_steps} sets $M$. The latent state $x_{j,t}$ is not a free +parameter: given observed spikes and a candidate +$(\tilde{U},\tilde{\tau}_\text{rec})$, it is fully determined by the forward +recursion~\eqref{eq:std_ssm_state}. Consequently $\{x,u,h\}$ are recomputed from +scratch whenever $\tilde{U}$ or $\tilde{\tau}_\text{rec}$ changes. + +%========= +\subsection{Initialization} +\label{sec:init} + +\textbf{Step 1 --- $\tilde{\tau}_\text{rec}^{(0)}$ and $\tilde{U}^{(0)}$.} +Do not pre-estimate these from raw ISIs or spike-pair statistics. +Set physiological starting values: +\begin{equation} + \tilde{\tau}_\text{rec}^{(0)} = 0.4~\text{s}, \qquad \tilde{U}^{(0)} = 0.4. +\end{equation} +These initialize the forward STD recursion only; both are updated later in +Block~B. + +\textbf{Step 2 --- Burn-in.} +Set $x_{j,0}=1$ for all $j$ and exclude the first +\begin{equation} + T_\text{burn} = \lceil \texttt{--burn\_sec}/\Delta t \rceil . +\end{equation} +bins from all likelihood evaluations (retain them for the forward state pass). +The \texttt{--burn\_sec} argument is mandatory and must be positive. +This removes sensitivity to the initial condition, which decays on timescale +$\tau_\text{rec}$. +The fitter treats the user-specified \texttt{-T} interval as the requested +post-burn fitting duration and requires the input file to contain both the +burn-in bins and the full requested fitting window. If the input is too short, +the run aborts rather than shortening the requested fit. + +\textbf{Step 3 --- $\{B_i^{(0)}, W_{ij}^{(0)}\}$.} +Initialize both $\{B_i^{(0)},W_{ij}^{(0)}\}$ from the first +$J=\texttt{--init\_samples}$ post-burn bins. +The \texttt{--init\_samples} argument is mandatory and must be greater than +$1000$; the fitter aborts if fewer than $J$ post-burn bins are available. +The baseline log-odds use the mean firing rates in this initialization window: +\begin{equation} + B_i^{(0)} = \log\frac{\bar{p}_i}{1-\bar{p}_i}, \qquad + \bar{p}_i = \frac{1}{J}\sum_{r=0}^{J-1} s_{i,T_\text{burn}+r}. +\end{equation} +Weights are initialized from raw lag-1 spike covariances over the same +post-burn window: +\begin{equation} + W_{ij}^{(0)} = + \frac{1}{J-1}\sum_{r=1}^{J-1} + \left(s_{i,T_\text{burn}+r}-\bar{p}^{+}_i\right) + \left(s_{j,T_\text{burn}+r-1}-\bar{p}^{-}_j\right), + \qquad i\neq j, +\end{equation} +with $W_{ii}^{(0)}=0$. Here $\bar{p}^{+}_i$ and $\bar{p}^{-}_j$ are the +means of the post- and pre-lag samples used in the covariance. Thus +$W_{ij}^{(0)}$ is positive when spikes in neuron $j$ one bin earlier are +associated with increased spiking in neuron $i$. +The output NPZ stores these initialization arrays as \texttt{B\_init} and +\texttt{W\_init}; it also stores the Eq.~(17) means as +\texttt{init\_p\_mean}, \texttt{init\_p\_plus}, and +\texttt{init\_p\_minus}. + +%========= +\subsection{Block A: Update $\{B_i, W_{ij}\}$} +\label{sec:block_a} + +\paragraph{Objective.} +With $(\tilde{U},\tilde{\tau}_\text{rec})$ fixed and $\{h_{j,t}\}$ precomputed, +the per-bin clipped logit and predicted probability are +\begin{equation} + \eta_{i,t} = \mathrm{clip}\!\left( + B_i + \sum_{j\neq i} \tilde{W}_{ij}\,h_{j,t},\; -\eta_c,\; \eta_c + \right), + \qquad + \tilde{p}_{i,t} = \sigma(\eta_{i,t}), + \label{eq:blockA_logit} +\end{equation} +where $\eta_c$ is a logit-clip threshold (default $\eta_c=10$). +For reuse across the coordinate-descent blocks, define the mean Bernoulli +negative log-likelihood over all post-burn bins and all neurons: +\begin{equation} +\mathcal{N}(\tilde{p}) += +-\frac{1}{(T-T_\text{burn})N} +\sum_{i=1}^{N}\sum_{t=T_\text{burn}}^{T-1} +\Bigl[s_{i,t}\log\tilde{p}_{i,t}+(1-s_{i,t})\log(1-\tilde{p}_{i,t})\Bigr]. +\label{eq:mean_bernoulli_nll} +\end{equation} +Block~A minimizes this NLL plus L1 regularization: +\begin{equation} +\mathcal{L}_A(\{B_i,\tilde{W}_{ij}\}) += +\mathcal{N}(\tilde{p}) ++\lambda_1\,\overline{|\tilde{W}|}, +\label{eq:blockA_loss} +\end{equation} +where $\overline{|\tilde{W}|}$ is the mean absolute value of the off-diagonal +weights. The L1 term ($\lambda_1>0$, default $10^{-4}$) promotes sparsity. +The fitter saves the epoch-wise BCE term, L1 term, and their sum separately as +\texttt{blockA\_bce\_loss}, \texttt{blockA\_l1\_loss}, and +\texttt{blockA\_loss}, respectively. The final fitted arrays are saved as +\texttt{B\_fit} and \texttt{W\_fit}. + +\paragraph{Optimizer.} +$\mathcal{L}_A$ is minimized by the Adam optimizer with learning rate +$\eta_W$ (default $10^{-2}$), run for a fixed number of epochs. +Each epoch draws a random permutation of the $T-T_\text{burn}$ post-burn bins +and processes them in mini-batches of size $B$ (default $8192$). + +\paragraph{Constraint applied after each mini-batch step.} +\begin{enumerate} + \item \textbf{Zero diagonal:} $\tilde{W}_{ii}\leftarrow 0$ for all $i$ + (no self-connections). +\end{enumerate} + +\paragraph{Constraint applied after each full epoch.} +\begin{enumerate} + \setcounter{enumi}{1} + \item \textbf{Spectral-radius cap:} if $\rho(\tilde{W})>\rho_\text{max}$ + (default $0.99$), rescale + $\tilde{W}\leftarrow\tilde{W}\cdot\rho_\text{max}/\rho(\tilde{W})$. + The cap is delayed by \texttt{--delay\_epoch\_4\_rhoMax}: if this + value is $D$, the first $D$ Block~A epochs of each outer iteration + skip the cap, and rescaling starts at epoch $D+1$. + This scalar rescaling preserves the zero diagonal already enforced + after each mini-batch update. +\end{enumerate} + +%========= +\subsection{Block B: Joint Update of $\tilde{U}$ and $\tilde{\tau}_\text{rec}$} +\label{sec:block_b} + +With $\{B_i,\tilde{W}_{ij}\}$ fixed, the STD parameters are updated jointly by +minimizing the mean Bernoulli NLL over a two-dimensional grid: +\begin{equation} +\mathcal{L}_B(\tilde{U},\tilde{\tau}_\text{rec}) += +\mathcal{N}(\tilde{p}(\tilde{U},\tilde{\tau}_\text{rec})). +\label{eq:blockB_loss} +\end{equation} +Here $\tilde{p}_{i,t}$ is evaluated by Eq.~\eqref{eq:blockA_logit}, and the +drive $\{h_{j,t}\}$ is recomputed from scratch for each candidate pair +$(\tilde{U},\tilde{\tau}_\text{rec})$. No true values and no quadratic anchor +terms are used. + +\paragraph{Search procedure.} +A bounded two-dimensional grid-refinement search is used over +$\tilde{U}\in[U_\text{lo},U_\text{hi}]$ and +$\tilde{\tau}_\text{rec}\in[\tau_\text{lo},\tau_\text{hi}]$, supplied by the +mandatory \texttt{--u\_bounds} and \texttt{--tau\_bounds} arguments. +Each refinement round evaluates $\mathcal{L}_B$ on a Cartesian grid with +\texttt{--u\_grid\_points} values in the $U$ direction and +\texttt{--tau\_grid\_points} values in the $\tau_\text{rec}$ direction +(both default to $9$). The best pair is retained, then the rectangular search +box is narrowed by \texttt{--grid\_shrink} (default $0.35$) around that pair +for \texttt{--grid\_refine} refinement rounds (default $1$). +No gradients of $\mathcal{L}_B$ with respect to the STD parameters are used. + +%========= +\subsection{Outer Loop and Scheduling} +\label{sec:outer_loop} + +The outer loop cycles through Blocks A$\to$B for a user-specified +number of iterations $K_\text{outer}$ (default $6$). +To avoid early-stage bias from poor STD parameters corrupting the +$\tilde{W}_{ij}$ fit, the first $K_\text{freeze}$ outer iterations +(default $2$) run Block~A only and skip Block~B. +The number of Block~A epochs per outer iteration can differ between the +frozen phase ($E_\text{freeze}$, default $40$) and the live phase +($E_\text{live}$, default: same as $E_\text{freeze}$). +The number of outer iterations and epochs are set by the user based on +observed loss curves; no automatic stopping criterion is applied. + +%========= +\subsection{Computational Scaling} +\label{sec:scaling} + +Each $\mathcal{O}(NT)$ forward pass (Eq.~\ref{eq:std_ssm_state}--\ref{eq:recursive_kernel}) +is a sequential scan over $T$ bins, vectorized across the $N$ neurons. +Block~A processes $(T-T_\text{burn})$ time steps in mini-batches; +the dominant cost is the $(T-T_\text{burn})\times N$ matrix operations for +$\eta$ and the backward pass through Adam. +Block~B requires a forward pass for each pair in the joint +$U\times\tau_\text{rec}$ grid. With current-best insertion, this is at most +$(G_U+1)(G_\tau+1)(R+1)$ forward passes per outer iteration; at the default +$G_U=G_\tau=9$ and $R=1$, this is at most $200$ forward passes. + +%========= +\subsection{Identifiability} +\label{sec:identifiability} + +The strongest coupling is between $\tilde{U}$ and $\tilde{W}_{ij}$: larger +$\tilde{U}$ inflates the release amplitude $u_{j,t}$, which can be partially +compensated by smaller weights. The bounded search domain $0<\tilde{U}\leq1$ +prevents nonphysical values but does not remove this coupling by itself. +A second coupling exists between $\tilde{\tau}_\text{rec}$ and $\tilde{W}_{ij}$: +longer recovery produces more depleted $x_{j,t}$ trajectories, which Block~A +can partially offset by inflating $\tilde{W}_{ij}$. Neither coupling is an +exact degeneracy because the model must match the temporal dynamics of recovery +after bursts, not only the mean coupling level. The joint Block~B search avoids +an additional sequential-coordinate artifact between $\tilde{U}$ and +$\tilde{\tau}_\text{rec}$, but it does not by itself eliminate the coupling +between STD amplitude and recurrent weight scale. The frozen phase exploits +this structure by stabilizing $\tilde{W}_{ij}$ before releasing the STD +parameters. +\begin{thebibliography}{1} + +\bibitem{Tsodyks1997} +M.~V. Tsodyks and H.~Markram. +\newblock The neural code between neocortical pyramidal neurons depends on + neurotransmitter release probability. +\newblock \emph{Proceedings of the National Academy of Sciences}, 94(2):719--723, + 1997. +\newblock doi:{10.1073/pnas.94.2.719}. + +\end{thebibliography} + +\end{document} + diff --git a/causal_net/nonStation_ver5_synapDepres/docs/tm_std_cartoon.png b/causal_net/nonStation_ver5_synapDepres/docs/tm_std_cartoon.png new file mode 100644 index 00000000..ec814164 Binary files /dev/null and b/causal_net/nonStation_ver5_synapDepres/docs/tm_std_cartoon.png differ diff --git a/causal_net/nonStation_ver5_synapDepres/eval_BSSM_fit.py b/causal_net/nonStation_ver5_synapDepres/eval_BSSM_fit.py new file mode 100755 index 00000000..8f577907 --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/eval_BSSM_fit.py @@ -0,0 +1,297 @@ +#!/usr/bin/env python3 +""" +Strict evaluation and plotting for fit5_BSSM_STD_blocks.py outputs. + +The evaluator intentionally requires the current BSSM-STD NPZ schema. Missing +records are errors, not silently inferred defaults. +""" + +import argparse +import os +from pprint import pprint + +import numpy as np + +from PlotterBSSM_fit import Plotter +from toolbox.Util_NumpyIO import read_data_npz + + +FIT_REQUIRED = [ + "W_fit", "B_fit", "W_init", "B_init", + "init_p_mean", "init_p_plus", "init_p_minus", "W_init_abs_mean", + "initW_random", "W_init_rebin_steps", "W_init_rebin_bins", + "outer_idx", "U_outer", "tau_rec_outer", + "nll_after_blockA", "nll_after_std_grid", "nll_outer", + "spectral_radius_outer", "nz_weight_outer", "blockB_executed", + "blockA_outer", "blockA_epoch", "blockA_bce_loss", "blockA_l1_loss", "blockA_loss", "blockA_lr", + "joint_search_outer", "joint_search_U", "joint_search_tau_rec", "joint_search_nll", + "joint_search_refine", "joint_search_eval", "joint_search_eval_total", + "kappa_fit", "alpha_fit", "burn_bins", "time_range_bins", "requested_time_range_bins", + "fit_time_range_bins", "init_time_range_bins", "time_step_sec", + "num_neurons", "num_input_bins", "num_eval_bins", "num_init_bins", + "U_init", "tau_rec_init", "u_bounds", "tau_bounds", + "u_grid_points", "tau_grid_points", "grid_refine", "grid_shrink", + "delay_epoch_4_blockB", "eta_clip", "lr_w", "lr_end_factor", + "lambda_l1", "rho_max", "weight_threshold", "h_chunk_steps", +] + +SPIKE_REQUIRED = ["spikes", "single_rates", "single_rates_var", "single_fano_fact"] + + +def abort(msg): + raise SystemExit("ERROR: %s" % msg) + + +def require_records(data, keys, label): + missing = [key for key in keys if key not in data] + if missing: + abort("%s missing required records: %s" % (label, ", ".join(missing))) + + +def scalar(arr, dtype=float): + val = np.asarray(arr).reshape(-1)[0] + return dtype(val) + + +def stable_sigmoid(x): + x = np.asarray(x, dtype=np.float64) + out = np.empty_like(x) + pos = x >= 0 + out[pos] = 1.0 / (1.0 + np.exp(-x[pos])) + exp_x = np.exp(x[~pos]) + out[~pos] = exp_x / (1.0 + exp_x) + return out + + +def std_h_forward_numpy(spikes, U, tau_rec, dt, alpha, mem_lag_steps, burn_bins): + spikes = np.asarray(spikes, dtype=np.float64) + assert spikes.ndim == 2, "spikes window must be 2D" + assert 0.0 < float(U) < 0.999, "U must satisfy 0 < U < 0.999" + assert float(tau_rec) > 0.0, "tau_rec must be positive" + assert int(mem_lag_steps) >= 1, "mem_lag_steps must be positive" + assert 0 <= int(burn_bins) < spikes.shape[0], "bad burn_bins" + + T, N = spikes.shape + M = int(mem_lag_steps) + H = np.empty((T - int(burn_bins), N), dtype=np.float32) + x = np.ones((N,), dtype=np.float64) + h = np.zeros((N,), dtype=np.float64) + u_hist = np.zeros((M, N), dtype=np.float64) + + rec = float(np.exp(-float(dt) / float(tau_rec))) + gain = 1.0 - float(alpha) + tail_gain = gain * (float(alpha) ** M) + + for t in range(T): + hist_idx = t % M + old_u = u_hist[hist_idx].copy() + if t >= int(burn_bins): + H[t - int(burn_bins)] = h.astype(np.float32) + + s_t = spikes[t] + u_t = float(U) * x * s_t + x_depleted = x * (1.0 - float(U) * s_t) + x = 1.0 - (1.0 - x_depleted) * rec + x = np.clip(x, 0.0, 1.0) + + h = float(alpha) * h + gain * u_t - tail_gain * old_u + h = np.maximum(h, 0.0) + u_hist[hist_idx] = u_t + + return H + + +def validate_fit_shapes(fitD): + W = np.asarray(fitD["W_fit"]) + B = np.asarray(fitD["B_fit"]) + W0 = np.asarray(fitD["W_init"]) + B0 = np.asarray(fitD["B_init"]) + assert W.ndim == 2 and W.shape[0] == W.shape[1], "W_fit must be square" + assert W0.shape == W.shape, "W_init/W_fit shape mismatch" + assert B.shape == (W.shape[0],), "B_fit shape mismatch" + assert B0.shape == B.shape, "B_init/B_fit shape mismatch" + assert int(scalar(fitD["num_neurons"], int)) == W.shape[0], "num_neurons disagrees with W shape" + assert np.allclose(np.diag(W0), 0.0), "Eq.17 requires zero diagonal in W_init" + for key in ["init_p_mean", "init_p_plus", "init_p_minus"]: + assert np.asarray(fitD[key]).shape == B.shape, "%s shape mismatch" % key + + n_outer = fitD["outer_idx"].shape[0] + for key in ["U_outer", "tau_rec_outer", "nll_after_blockA", "nll_after_std_grid", + "nll_outer", "spectral_radius_outer", "nz_weight_outer", "blockB_executed"]: + assert fitD[key].shape[0] == n_outer, "%s length mismatch" % key + + n_block = fitD["blockA_outer"].shape[0] + for key in ["blockA_epoch", "blockA_bce_loss", "blockA_l1_loss", "blockA_loss", "blockA_lr"]: + assert fitD[key].shape[0] == n_block, "%s length mismatch" % key + + n_search = fitD["joint_search_outer"].shape[0] + for key in ["joint_search_U", "joint_search_tau_rec", "joint_search_nll", + "joint_search_refine", "joint_search_eval", "joint_search_eval_total"]: + assert fitD[key].shape[0] == n_search, "%s length mismatch" % key + + +def build_diag_base(fitD, args, fit_file, spikes_file): + validate_fit_shapes(fitD) + raw_bins = np.asarray(fitD["time_range_bins"], dtype=np.int64) + fit_bins = np.asarray(fitD["fit_time_range_bins"], dtype=np.int64) + init_bins = np.asarray(fitD["init_time_range_bins"], dtype=np.int64) + assert raw_bins.shape == fit_bins.shape == init_bins.shape == (2,), "time-bin records must be length 2" + assert raw_bins[0] <= fit_bins[0] <= fit_bins[1] <= raw_bins[1], "fit bins must lie inside raw bins" + assert fit_bins[0] <= init_bins[0] <= init_bins[1] <= fit_bins[1], "init bins must lie inside fit bins" + + return { + "short_name": args.dataName, + "data_name": args.dataName, + "fit_file": fit_file, + "spikes_file": spikes_file, + } + + +def compute_observation_diagnostics(fitD, spikeD, obs_batch): + spikes = np.asarray(spikeD["spikes"]) + assert spikes.ndim == 2, "spikes record must be 2D" + assert np.max(spikes) <= 1, "spikes must be binary" + + raw0, raw1 = np.asarray(fitD["time_range_bins"], dtype=np.int64) + assert 0 <= raw0 <= raw1 < spikes.shape[0], "fit raw bins exceed spikes record" + spikes_w = spikes[raw0:raw1 + 1] + + burn_bins = scalar(fitD["burn_bins"], int) + dt = scalar(fitD["time_step_sec"], float) + alpha = scalar(fitD["alpha_fit"], float) + M = int(np.asarray(fitD["kappa_fit"]).size) + U = float(np.asarray(fitD["U_outer"], dtype=np.float64)[-1]) + tau_rec = float(np.asarray(fitD["tau_rec_outer"], dtype=np.float64)[-1]) + W = np.asarray(fitD["W_fit"], dtype=np.float64) + B = np.asarray(fitD["B_fit"], dtype=np.float64) + eta_clip = scalar(fitD["eta_clip"], float) + n_eval = scalar(fitD["num_eval_bins"], int) + + H = std_h_forward_numpy(spikes_w, U, tau_rec, dt, alpha, M, burn_bins) + Y = np.asarray(spikes_w[burn_bins:], dtype=np.float64) + assert H.shape == Y.shape, "H/Y shape mismatch" + assert H.shape[0] == n_eval, "computed H length disagrees with num_eval_bins" + assert H.shape[1] == W.shape[0], "H/W neuron count mismatch" + + N = W.shape[0] + sum_nll_neuron = np.zeros((N,), dtype=np.float64) + sum_prob_neuron = np.zeros((N,), dtype=np.float64) + total_nll = 0.0 + total_count = 0 + + p_edges = np.linspace(0.0, 1.0, 41) + p_hist_count = np.zeros((p_edges.size - 1,), dtype=np.int64) + calib_sum_p = np.zeros_like(p_hist_count, dtype=np.float64) + calib_sum_y = np.zeros_like(p_hist_count, dtype=np.float64) + + batch = int(obs_batch) + assert batch > 0, "--obs_batch must be positive" + for i0 in range(0, H.shape[0], batch): + i1 = min(H.shape[0], i0 + batch) + eta = np.asarray(H[i0:i1], dtype=np.float64) @ W.T + B + eta = np.clip(eta, -eta_clip, eta_clip) + p = stable_sigmoid(eta) + y = Y[i0:i1] + p_safe = np.clip(p, 1e-12, 1.0 - 1e-12) + nll = -(y * np.log(p_safe) + (1.0 - y) * np.log1p(-p_safe)) + sum_nll_neuron += np.sum(nll, axis=0) + sum_prob_neuron += np.sum(p, axis=0) + total_nll += float(np.sum(nll)) + total_count += int(nll.size) + + flat_p = p.reshape(-1) + flat_y = y.reshape(-1) + bins = np.searchsorted(p_edges, flat_p, side="right") - 1 + bins = np.clip(bins, 0, p_hist_count.size - 1) + p_hist_count += np.bincount(bins, minlength=p_hist_count.size) + calib_sum_p += np.bincount(bins, weights=flat_p, minlength=p_hist_count.size) + calib_sum_y += np.bincount(bins, weights=flat_y, minlength=p_hist_count.size) + + obs_rate_hz = np.mean(Y, axis=0) / dt + model_rate_hz = sum_prob_neuron / float(Y.shape[0]) / dt + nll_neuron = sum_nll_neuron / float(Y.shape[0]) + valid = p_hist_count > 0 + calib_p = np.zeros_like(calib_sum_p) + calib_obs = np.zeros_like(calib_sum_y) + calib_p[valid] = calib_sum_p[valid] / p_hist_count[valid] + calib_obs[valid] = calib_sum_y[valid] / p_hist_count[valid] + + single_rates = np.asarray(spikeD["single_rates"], dtype=np.float64) + assert single_rates.shape == obs_rate_hz.shape, "single_rates shape mismatch" + + return { + "obs_rate_hz": obs_rate_hz, + "model_rate_hz": model_rate_hz, + "nll_neuron": nll_neuron, + "calib_p": calib_p, + "calib_obs": calib_obs, + "calib_count": p_hist_count, + "p_hist_edges": p_edges, + "p_hist_count": p_hist_count, + "single_rates_hz": single_rates, + "nll_eval": total_nll / float(total_count), + } + + +def main(): + parser = argparse.ArgumentParser(description="Evaluate and plot BSSM-STD fit results") + prs=parser.add_argument + prs("--dataName", required=True, help="Spike data base name, e.g. daleN100_861b17") + prs("--basePath", default="/pscratch/sd/b/balewski/2026_causalNet_tmp", + help="Root containing truthDale/ and plots/") + prs("-p", "--showPlots", nargs="+", default=["a", "b", "c", "d"], + help="Plots: a=two-block summary, b=B/W initialization vs fit, c=Block B grid, d=Bernoulli observation diagnostics") + prs("--obs_batch", type=int, default=65536, + help="Time-bin batch size for Bernoulli observation diagnostics") + prs("-X", "--noXterm", action="store_true", help="Disable X display and save PNG only") + prs("-v", "--verb", type=int, default=1) + args = parser.parse_args() + if args.verb > 0: + print("BSSM-fit eval args:"); pprint(vars(args)) + + inpPath = os.path.join(args.basePath, "fitBssmStd") + args.outPath = os.path.join(args.basePath, "plots") + + # ── load BSSM-STD fit ──────────────────────────────────────────── + fitFF = os.path.join(inpPath, f"{args.dataName}.fitBSSMSTD.npz") + fitD, fitMD = read_data_npz(fitFF) + require_records(fitD, FIT_REQUIRED, "fit NPZ") + + if args.verb > 1: pprint(fitMD) + + #--- load spikes ---- + if "provenance" not in fitMD or "spikesData_file" not in fitMD["provenance"]: + abort("fit metadata missing provenance.spikesData_file") + prov = fitMD["provenance"] + spikesF=prov['spikesData_file'] + spikesFF = os.path.join(args.basePath, "truthDale", f"{spikesF}.spikes.npz") + spikeD, _ = read_data_npz(spikesFF, verb=args.verb > 0) + require_records(spikeD, SPIKE_REQUIRED, "spikes NPZ") + + args.showPlots =''.join(args.showPlots) + + # ── plot ─────────────────────────────── + args.prjName = args.dataName + plot = Plotter(args) + diag = build_diag_base(fitD, args, fitFF, spikesFF) + + if "d" in args.showPlots: + obsD = compute_observation_diagnostics(fitD, spikeD, args.obs_batch) + diag.update(obsD) + if args.verb > 0: + print(" recomputed Bernoulli observation NLL=%.9f" % diag["nll_eval"]) + + plot = Plotter(args) + if "a" in args.showPlots: + plot.two_block_summary(fitD, diag, figId=1) + if "b" in args.showPlots: + plot.parameter_init_vs_fit(fitD, diag, figId=2) + if "c" in args.showPlots: + plot.blockB_joint_grid(fitD, diag, figId=3) + if "d" in args.showPlots: + plot.bernoulli_observation_diagnostics(fitD, diag, figId=4) + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver5_synapDepres/fit5_BSSM_STD_blocks.py b/causal_net/nonStation_ver5_synapDepres/fit5_BSSM_STD_blocks.py new file mode 100755 index 00000000..789067c7 --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/fit5_BSSM_STD_blocks.py @@ -0,0 +1,1310 @@ +#!/usr/bin/env python3 +""" +Block-coordinate fit for the BSSM-STD Bernoulli model. + +The script fits + + p_t = sigmoid(B + W h_t) + +from observed spikes. The fast synaptic kernel is fixed from command-line +arguments, while the STD release trajectory h_t is recomputed from +observed spikes for each candidate (U, tau_rec). + +Blocks: + A. fit B and W with U, tau_rec fixed, using mini-batch Bernoulli NLL; + B. jointly update U and tau_rec by bounded 2D grid-refinement search with B,W fixed. + +Run with a PyTorch module on NERSC, for example: + + module load pytorch + ./fit5_BSSM_STD_blocks.py --basePath $basePath --dataName daleN100_xxxxxx +""" + +import argparse +import json +import math +import os +import time +import zipfile +from pprint import pprint + +import numpy as np + +try: + import torch + import torch.nn.functional as F +except ModuleNotFoundError as exc: + raise SystemExit("PyTorch is required. On NERSC run: module load pytorch") from exc + +from toolbox.Util_NumpyIO import write_data_npz + + +def parse_args(): + p = argparse.ArgumentParser( + description="Fit BSSM-STD Bernoulli model by block coordinate descent." + ) + p.add_argument("--dataName", required=True) + p.add_argument( + "--basePath", + default="/pscratch/sd/b/balewski/2026_causalNet_tmp/", + help="Input root; reads /truthDale/.*.npz.", + ) + p.add_argument( + "--outPath", + default=None, + help="Output directory. Default: /fitBssmStd.", + ) + p.add_argument( + "-T", + "--time_range_sec", + nargs=2, + type=float, + default=[0.0, 100.0], + help="Requested post-burn fitting window [t0, t1] seconds. " + "The input must contain this window plus burn-in.", + ) + p.add_argument("--dtype", default="float32", choices=["float32", "float64"]) + p.add_argument( + "--synaptic_tau", + type=float, + default=0.005, + help="Fast synaptic tau_s in seconds for the fixed exponential kernel.", + ) + p.add_argument( + "--kernel_len_steps", + type=int, + default=25, + help="Number of lag bins M retained in the fixed exponential kernel.", + ) + p.add_argument( + "--h_chunk_steps", + type=int, + default=2048, + help="CUDA chunk length for the STD recurrence. Lower this only if an underflow abort occurs.", + ) + p.add_argument("--u0", type=float, default=0.4, help="Initial STD utilization U.") + p.add_argument("--tau_rec0", type=float, default=0.4, help="Initial tau_rec in seconds.") + p.add_argument( + "--init_samples", + type=int, + required=True, + help="Mandatory number of post-burn bins, >1000, used to initialize B from rates " + "and W from lag-1 covariances.", + ) + p.add_argument( + "--initW_random", + type=float, + default=0.0, + help="If >0, initialize off-diagonal W_ij uniformly in [-value,value]; " + "otherwise initialize W from post-burn lag-1 covariances.", + ) + p.add_argument( + "--u_bounds", + type=float, + nargs=2, + required=True, + help="Mandatory U scan range [U_lo U_hi] for the joint STD grid search.", + ) + p.add_argument( + "--tau_bounds", + type=float, + nargs=2, + required=True, + help="Mandatory tau_rec scan range [tau_lo tau_hi] in seconds for the joint STD grid search.", + ) + p.add_argument( + "--burn_sec", + type=float, + required=True, + help="Mandatory positive burn-in seconds excluded from likelihood.", + ) + p.add_argument("--num_outer", type=int, default=6) + p.add_argument( + "--freeze_std_outer", + type=int, + default=2, + help="Number of initial outer iterations that update only B,W.", + ) + p.add_argument( + "--delay_epoch_4_blockB", + type=int, + default=5, + help="Number of initial outer iterations to skip before running Block B grid search.", + ) + p.add_argument("--blockA_epochs", type=int, default=40, + help="Block A epochs for frozen-STD outer iterations.") + p.add_argument("--blockA_epochs_live", type=int, default=-1, + help="Block A epochs once Block B is live. If <0, reuse --blockA_epochs.") + p.add_argument( + "--progress_every", + type=int, + default=5, + help="Print Block A progress every this many epochs when --verb > 0; use 0 to disable.", + ) + p.add_argument("--batch_size", type=int, default=8192) + p.add_argument("--lr_w", type=float, default=1e-2) + p.add_argument( + "--lr_end_factor", + type=float, + default=0.1, + help="LR decays to lr_w * lr_end_factor", + ) + p.add_argument("--lambda_l1", type=float, default=1e-4) + p.add_argument( + "--rho_max", + type=float, + default=0.99, + help="If >0, rescale W after each Block A epoch when spectral radius exceeds this value.", + ) + p.add_argument( + "--delay_epoch_4_rhoMax", + type=int, + default=99, + help="Number of initial Block A epochs per outer iteration to skip before applying --rho_max.", + ) + p.add_argument("--eta_clip", type=float, default=10.0) + p.add_argument("--u_grid_points", type=int, default=9) + p.add_argument("--tau_grid_points", type=int, default=9) + p.add_argument("--grid_refine", type=int, default=1) + p.add_argument( + "--grid_shrink", + type=float, + default=0.35, + help="Refinement half-width as a fraction of the current U/tau_rec search intervals.", + ) + p.add_argument( + "--weight_threshold", + type=float, + default=0.03, + help="Absolute off-diagonal weight threshold used only for reporting nz_weight_outer.", + ) + p.add_argument("--seed", type=int, default=12345) + p.add_argument("-v", "--verb", type=int, default=1) + return p.parse_args() + + +def require_keys(obj, keys, obj_name): + if obj is None: + abort_fit("missing required object %s" % obj_name) + for key in keys: + if key not in obj: + abort_fit("missing required key %s[%r]" % (obj_name, key)) + + +def _extract_json_object(text): + start = text.find('{"') + if start < 0: + abort_fit("could not locate JSON object in meta.JSON.npy") + depth = 0 + in_str = False + escape = False + for pos in range(start, len(text)): + ch = text[pos] + if in_str: + if escape: + escape = False + elif ch == "\\": + escape = True + elif ch == '"': + in_str = False + continue + if ch == '"': + in_str = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return text[start : pos + 1] + abort_fit("unterminated JSON object in meta.JSON.npy") + + +def read_meta_json(npz_path): + with zipfile.ZipFile(npz_path, "r") as zf: + if "meta.JSON.npy" not in zf.namelist(): + abort_fit("missing meta.JSON.npy in %s" % npz_path) + raw = zf.read("meta.JSON.npy") + text = raw.decode("utf-8", errors="ignore") + return json.loads(_extract_json_object(text)) + + +def read_npz_records(npz_path, keys, verb=0): + """Read selected non-object arrays plus metadata without loading all records.""" + out = {} + with np.load(npz_path, allow_pickle=False) as data: + for key in keys: + if key not in data.files: + abort_fit("missing required key %s in %s" % (key, npz_path)) + out[key] = data[key] + meta = read_meta_json(npz_path) + if verb > 0: + print("read selected records from:", npz_path) + for key, val in out.items(): + print(" %s %s %s" % (key, val.shape, val.dtype)) + return out, meta + + +def choose_device(): + if not torch.cuda.is_available(): + raise SystemExit("ERROR: CUDA GPU is required but torch.cuda.is_available() is false") + if torch.cuda.device_count() < 1: + raise SystemExit("ERROR: no CUDA GPU devices are visible") + return torch.device("cuda:0") + + +def abort_fit(msg): + raise SystemExit("ERROR: %s" % msg) + + +def time_window_with_burn(spikes, dt, time_range_sec, burn_bins): + t0_sec, t1_sec = float(time_range_sec[0]), float(time_range_sec[1]) + if not t0_sec < t1_sec: + abort_fit("time_range_sec[0] must be < time_range_sec[1]") + if burn_bins < 0: + abort_fit("burn_bins must be nonnegative") + + requested_start_bin = int(math.floor(t0_sec / dt)) + requested_end_bin = int(math.floor(t1_sec / dt)) + if requested_start_bin < 0: + abort_fit("requested start_bin=%d is negative" % requested_start_bin) + + requested_bins = requested_end_bin - requested_start_bin + 1 + raw_start_bin = requested_start_bin + raw_end_bin = requested_end_bin + int(burn_bins) + input_bins = int(spikes.shape[0]) + if raw_end_bin >= input_bins: + abort_fit( + "not enough input time bins for requested window plus burn-in: " + "requested_bins=%d burn_bins=%d requires raw bins [%d,%d], " + "but input has bins [0,%d]" + % ( + requested_bins, + int(burn_bins), + raw_start_bin, + raw_end_bin, + input_bins - 1, + ) + ) + return ( + raw_start_bin, + raw_end_bin, + requested_start_bin, + requested_end_bin, + requested_bins, + spikes[raw_start_bin : raw_end_bin + 1], + ) + + +def logit_np(p): + p = np.clip(np.asarray(p, dtype=np.float64), 1e-6, 1.0 - 1e-6) + return np.log(p) - np.log1p(-p) + + +def init_B_W_from_postburn_spikes( + spikes_w, + burn_bins, + init_samples, + kernel_len_steps, + initW_random, + seed, +): + init_samples = int(init_samples) + if init_samples < 2: + abort_fit("--init_samples must be at least 2") + kernel_len_steps = int(kernel_len_steps) + if kernel_len_steps < 1: + abort_fit("--kernel_len_steps must be at least 1") + postburn = np.asarray(spikes_w[burn_bins:], dtype=np.float64) + init_bins = int(postburn.shape[0]) + if init_bins < init_samples: + abort_fit( + "not enough post-burn bins for initialization: " + "available=%d required --init_samples=%d" + % (init_bins, init_samples) + ) + + init_spikes = postburn[:init_samples] + p_mean = np.mean(init_spikes, axis=0) + B0 = logit_np(p_mean) + + if float(initW_random) > 0.0: + pre = init_spikes[:-1] + post = init_spikes[1:] + p_minus = np.mean(pre, axis=0) + p_plus = np.mean(post, axis=0) + rng = np.random.default_rng(int(seed)) + W0 = rng.uniform( + -float(initW_random), + float(initW_random), + size=(init_spikes.shape[1], init_spikes.shape[1]), + ) + np.fill_diagonal(W0, 0.0) + rebin_bins = 0 + return B0, W0, init_samples, p_mean, p_plus, p_minus, rebin_bins + + if init_samples % kernel_len_steps != 0: + abort_fit( + "--init_samples must be an integer multiple of --kernel_len_steps " + "for data-based W initialization: init_samples=%d kernel_len_steps=%d" + % (init_samples, kernel_len_steps) + ) + rebin_bins = init_samples // kernel_len_steps + if rebin_bins < 2: + abort_fit( + "data-based W initialization needs at least 2 rebinned samples: " + "init_samples=%d kernel_len_steps=%d gives %d" + % (init_samples, kernel_len_steps, rebin_bins) + ) + init_rebinned = init_spikes.reshape(rebin_bins, kernel_len_steps, init_spikes.shape[1]).sum(axis=1) + pre = init_rebinned[:-1] + post = init_rebinned[1:] + p_minus = np.mean(pre, axis=0) + p_plus = np.mean(post, axis=0) + pre_centered = pre - p_minus.reshape(1, -1) + post_centered = post - p_plus.reshape(1, -1) + W0 = (post_centered.T @ pre_centered) / float(rebin_bins - 1) + np.fill_diagonal(W0, 0.0) + return B0, W0, init_samples, p_mean, p_plus, p_minus, rebin_bins + + +def infer_alpha(kappa): + kappa = np.asarray(kappa, dtype=np.float64).reshape(-1) + if kappa.size < 1: + abort_fit("kernel must contain at least one entry") + if kappa.size == 1: + return 0.0 + if not kappa[0] > 0: + abort_fit("kernel first entry must be positive") + alpha = float(kappa[1] / kappa[0]) + if not (0.0 <= alpha < 1.0): + abort_fit("kernel ratio alpha must satisfy 0 <= alpha < 1") + return alpha + + +def build_exponential_kernel(dt, synaptic_tau, mem_lag_steps): + synaptic_tau = float(synaptic_tau) + mem_lag_steps = int(mem_lag_steps) + if not synaptic_tau > 0.0: + abort_fit("--synaptic_tau must be positive") + if mem_lag_steps < 1: + abort_fit("--kernel_len_steps must be at least 1") + alpha = float(math.exp(-float(dt) / synaptic_tau)) + ell = np.arange(mem_lag_steps, dtype=np.float64) + return (1.0 - alpha) * np.power(alpha, ell) + + +def project_w_(W): + with torch.no_grad(): + W.diagonal().zero_() + + +def enforce_spectral_radius_(W, rho_max): + if float(rho_max) <= 0.0: + abort_fit("--rho_max must be positive") + with torch.no_grad(): + rho = torch.linalg.eigvals(W).abs().max() + if not torch.isfinite(rho): + abort_fit("spectral radius is not finite") + if rho > float(rho_max): + W.mul_(float(rho_max) / float(rho.item())) + + +@torch.no_grad() +def spectral_radius_torch(W): + if W.device.type != "cuda": + abort_fit("spectral radius diagnostic requires CUDA tensor") + if W.numel() == 0: + abort_fit("spectral radius diagnostic requires nonempty W") + vals = torch.linalg.eigvals(W) + return float(torch.max(torch.abs(vals)).item()) + + +_H_FORWARD_SEC_PER_BIN = None + + +def _record_h_timing(elapsed, n_bins): + global _H_FORWARD_SEC_PER_BIN + sec_per_bin = float(elapsed) / float(max(1, int(n_bins))) + if _H_FORWARD_SEC_PER_BIN is None: + _H_FORWARD_SEC_PER_BIN = sec_per_bin + else: + _H_FORWARD_SEC_PER_BIN = 0.5 * float(_H_FORWARD_SEC_PER_BIN) + 0.5 * sec_per_bin + + +def std_h_forward_vectorized_cuda( + spikes, + U, + tau_rec, + dt, + alpha, + mem_lag_steps, + burn_bins, + label="H", + verb=1, + t0=None, + chunk_steps=2048, +): + """GPU-friendly H computation. Abort if the CUDA path cannot be used safely.""" + if spikes.device.type != "cuda": + abort_fit("H forward requires CUDA tensors; no CPU fallback is available") + + T = int(spikes.shape[0]) + N = int(spikes.shape[1]) + dtype = spikes.dtype + work_dtype = torch.float64 if dtype == torch.float32 else dtype + device = spikes.device + rec = float(math.exp(-float(dt) / float(tau_rec))) + beta = 1.0 - rec + tiny = 1e-250 if work_dtype == torch.float64 else 1e-35 + + x0 = torch.ones((N,), dtype=work_dtype, device=device) + u_all = torch.empty((T, N), dtype=dtype, device=device) + last_print = time.time() + t_start = time.time() if t0 is None else float(t0) + + for start in range(0, T, int(chunk_steps)): + end = min(T, start + int(chunk_steps)) + s = spikes[start:end].to(dtype=work_dtype) + a = rec * (1.0 - float(U) * s) + p_after = torch.cumprod(a, dim=0) + min_abs = float(torch.min(torch.abs(p_after)).item()) + if (not math.isfinite(min_abs)) or min_abs < tiny: + abort_fit( + "CUDA H recurrence underflow in %s at bins [%d,%d) with --h_chunk_steps=%d; " + "rerun with a smaller --h_chunk_steps" + % (label, start, end, int(chunk_steps)) + ) + + inv_cumsum = torch.cumsum(torch.reciprocal(p_after), dim=0) + p_before = torch.empty_like(p_after) + sum_before = torch.empty_like(inv_cumsum) + p_before[0].fill_(1.0) + sum_before[0].zero_() + if end - start > 1: + p_before[1:] = p_after[:-1] + sum_before[1:] = inv_cumsum[:-1] + + x_vals = p_before * (x0.unsqueeze(0) + beta * sum_before) + x_vals = torch.clamp(x_vals, 0.0, 1.0) + u_all[start:end] = (float(U) * x_vals * s).to(dtype=dtype) + x0 = p_after[-1] * (x0 + beta * inv_cumsum[-1]) + x0 = torch.clamp(x0, 0.0, 1.0) + + now = time.time() + if verb > 0 and now - last_print >= 5.0: + done_bins = end + frac = min(0.999, float(done_bins) / float(max(1, T))) + elapsed = now - t_start + eta = elapsed * (1.0 / max(1e-9, frac) - 1.0) + print( + " %s release scan: %d/%d bins %.0f%% elapsed=%.0fs eta=%.0fs" + % (label, done_bins, T, 100.0 * frac, elapsed, eta), + flush=True, + ) + last_print = now + + if verb > 0: + print(" %s synaptic convolution on GPU..." % label, flush=True) + ell = torch.arange(int(mem_lag_steps), dtype=dtype, device=device) + kappa = (1.0 - float(alpha)) * torch.pow( + torch.tensor(float(alpha), dtype=dtype, device=device), + ell, + ) + weight = torch.flip(kappa, dims=[0]).view(1, 1, int(mem_lag_steps)).repeat(N, 1, 1) + u_ch = u_all.transpose(0, 1).unsqueeze(0).contiguous() + u_pad = F.pad(u_ch, (int(mem_lag_steps), 0)) + h_ch = F.conv1d(u_pad, weight, groups=N) + H = h_ch[0, :, :T].transpose(0, 1).contiguous() + H = torch.clamp(H[int(burn_bins):], min=0.0) + return H + + +def make_h( + spikes_t, + U, + tau_rec, + dt, + alpha, + mem_lag_steps, + burn_bins, + h_chunk_steps, + label="H", + verb=1, +): + global _H_FORWARD_SEC_PER_BIN + if not (0.0 < float(U) < 0.999): + abort_fit("U must satisfy 0 < U < 0.999") + if not tau_rec > 0.0: + abort_fit("tau_rec must be positive") + + n_bins = int(spikes_t.shape[0]) + t0 = time.time() + H = std_h_forward_vectorized_cuda( + spikes_t, + U, + tau_rec, + dt, + alpha, + mem_lag_steps, + burn_bins, + label=label, + verb=verb, + t0=t0, + chunk_steps=int(h_chunk_steps), + ) + torch.cuda.synchronize() + elapsed = time.time() - t0 + _record_h_timing(elapsed, n_bins) + + if verb > 0: + print(" %s forward pass done in %.1fs shape=%s mean=%.6g max=%.6g" + % (label, elapsed, tuple(H.shape), + float(torch.mean(H).item()), float(torch.max(H).item())), flush=True) + return H + + +@torch.no_grad() +def bernoulli_nll_mean(H, Y_eval, W, B, eta_clip, batch_size): + n = H.shape[0] + N = H.shape[1] + total = 0.0 + for i0 in range(0, n, batch_size): + i1 = min(n, i0 + batch_size) + eta = H[i0:i1] @ W.t() + B + eta = torch.clamp(eta, -float(eta_clip), float(eta_clip)) + loss = F.binary_cross_entropy_with_logits( + eta, Y_eval[i0:i1].to(dtype=H.dtype), reduction="sum" + ) + total += float(loss.item()) + return total / float(max(1, n * N)) + + +def progress_line(args, t_start, msg): + if args.verb > 0: + print(" %s elapsed=%.1fs" % (msg, time.time() - t_start), flush=True) + + +def blockA_epochs_for_outer(args, outer): + live_std = ( + int(outer) >= int(args.freeze_std_outer) + and int(outer) >= int(args.delay_epoch_4_blockB) + ) + if live_std and int(args.blockA_epochs_live) >= 0: + return int(args.blockA_epochs_live) + return int(args.blockA_epochs) + + +def train_block_a( + H, + Y_eval, + W, + B, + args, + offdiag_mask, + generator, + outer, + t_start, + epochs=None, + global_epoch_offset=0, + global_epoch_total=None, +): + params = [W, B] + optimizer = torch.optim.Adam(params, lr=float(args.lr_w)) + n = int(H.shape[0]) + N = int(H.shape[1]) + bce_losses = [] + l1_losses = [] + total_losses = [] + lr_values = [] + progress_every = int(args.progress_every) + rho_delay = int(args.delay_epoch_4_rhoMax) + if epochs is None: + epochs = int(args.blockA_epochs) + if global_epoch_total is None: + global_epoch_total = int(epochs) + + for epoch in range(epochs): + global_epoch = int(global_epoch_offset) + int(epoch) + if int(global_epoch_total) > 1: + lr_frac = float(global_epoch) / float(int(global_epoch_total) - 1) + else: + lr_frac = 0.0 + lr_epoch = float(args.lr_w) * ( + 1.0 - lr_frac * (1.0 - float(args.lr_end_factor)) + ) + lr_values.append(lr_epoch) + for group in optimizer.param_groups: + group["lr"] = lr_epoch + + perm = torch.randperm(n, device=H.device, generator=generator) + epoch_bce_loss = 0.0 + epoch_l1_loss = 0.0 + epoch_total_loss = 0.0 + epoch_count = 0 + for i0 in range(0, n, int(args.batch_size)): + idx = perm[i0 : min(n, i0 + int(args.batch_size))] + Hb = H.index_select(0, idx) + Yb = Y_eval.index_select(0, idx).to(dtype=H.dtype) + + optimizer.zero_grad(set_to_none=True) + eta = Hb @ W.t() + B + eta = torch.clamp(eta, -float(args.eta_clip), float(args.eta_clip)) + base_loss = F.binary_cross_entropy_with_logits(eta, Yb, reduction="mean") + l1_loss = torch.zeros((), dtype=base_loss.dtype, device=base_loss.device) + if args.lambda_l1 > 0: + l1_loss = float(args.lambda_l1) * torch.mean(torch.abs(W[offdiag_mask])) + loss = base_loss + l1_loss + loss.backward() + optimizer.step() + project_w_(W) + + nb = int(idx.numel()) * N + epoch_bce_loss += float(base_loss.item()) * nb + epoch_l1_loss += float(l1_loss.item()) * nb + epoch_total_loss += float(loss.item()) * nb + epoch_count += nb + + epoch_bce = epoch_bce_loss / float(max(1, epoch_count)) + epoch_l1 = epoch_l1_loss / float(max(1, epoch_count)) + epoch_total = epoch_total_loss / float(max(1, epoch_count)) + bce_losses.append(epoch_bce) + l1_losses.append(epoch_l1) + total_losses.append(epoch_total) + if epoch + 1 > rho_delay: + enforce_spectral_radius_(W, float(args.rho_max)) + if ( + args.verb > 0 + and progress_every > 0 + and ( + epoch == 0 + or (epoch + 1) % progress_every == 0 + or epoch + 1 == epochs + ) + ): + print( + " progress outer=%d/%d epoch=%d/%d lr=%.3g blockA_bce=%.6f blockA_l1=%.3g blockA_total=%.6f elapsed=%.1fs" + % ( + outer + 1, + int(args.num_outer), + epoch + 1, + epochs, + lr_epoch, + epoch_bce, + epoch_l1, + epoch_total, + time.time() - t_start, + ), + flush=True, + ) + + return ( + np.asarray(bce_losses, dtype=np.float64), + np.asarray(l1_losses, dtype=np.float64), + np.asarray(total_losses, dtype=np.float64), + np.asarray(lr_values, dtype=np.float64), + ) + + +def blockB_joint_grid_search( + current_u, + current_tau, + u_bounds, + tau_bounds, + u_grid_points, + tau_grid_points, + refine_steps, + shrink, + objective_nll_fn, + t_start=None, + outer=None, + num_outer=None, + verb=1, + eval_progress_every=5, +): + u_lo = float(u_bounds[0]) + u_hi = float(u_bounds[1]) + tau_lo = float(tau_bounds[0]) + tau_hi = float(tau_bounds[1]) + if not u_lo < u_hi: + abort_fit("joint grid U bounds must satisfy U_lo < U_hi") + if not tau_lo < tau_hi: + abort_fit("joint grid tau bounds must satisfy tau_lo < tau_hi") + best_u = float(np.clip(current_u, u_lo, u_hi)) + best_tau = float(np.clip(current_tau, tau_lo, tau_hi)) + best_nll = float("inf") + u_all = [] + tau_all = [] + nll_all = [] + refine_all = [] + eval_all = [] + eval_total_all = [] + + cur_u_lo, cur_u_hi = u_lo, u_hi + cur_tau_lo, cur_tau_hi = tau_lo, tau_hi + for ref in range(int(refine_steps) + 1): + u_grid = np.linspace(cur_u_lo, cur_u_hi, int(u_grid_points), dtype=np.float64) + tau_grid = np.linspace(cur_tau_lo, cur_tau_hi, int(tau_grid_points), dtype=np.float64) + if best_u > cur_u_lo and best_u < cur_u_hi: + u_grid = np.unique(np.sort(np.append(u_grid, best_u))) + if best_tau > cur_tau_lo and best_tau < cur_tau_hi: + tau_grid = np.unique(np.sort(np.append(tau_grid, best_tau))) + + nlls = [] + pairs = [] + n_eval = int(u_grid.size * tau_grid.size) + i_eval = 0 + for u_val in u_grid: + for tau_val in tau_grid: + i_eval += 1 + nll = float(objective_nll_fn(float(u_val), float(tau_val))) + u_all.append(float(u_val)) + tau_all.append(float(tau_val)) + nll_all.append(nll) + refine_all.append(int(ref)) + eval_all.append(int(i_eval)) + eval_total_all.append(int(n_eval)) + nlls.append(nll) + pairs.append((float(u_val), float(tau_val))) + print_eval = ( + int(eval_progress_every) > 0 + and (i_eval % int(eval_progress_every) == 0 or i_eval == n_eval) + ) + if verb > 0 and print_eval: + prefix = ( + " search U_tau refine=%d eval=%d/%d U=%.6g tau_rec=%.6g nll=%.6f" + % (ref, i_eval, n_eval, float(u_val), float(tau_val), nll) + ) + if outer is not None and num_outer is not None: + prefix = " outer=%d/%d %s" % (int(outer) + 1, int(num_outer), prefix.strip()) + if t_start is not None: + prefix += " elapsed=%.1fs" % (time.time() - t_start) + print(prefix, flush=True) + nlls = np.asarray(nlls, dtype=np.float64) + idx = int(np.argmin(nlls)) + best_u, best_tau = pairs[idx] + best_nll = float(nlls[idx]) + + u_width = (cur_u_hi - cur_u_lo) * float(shrink) + tau_width = (cur_tau_hi - cur_tau_lo) * float(shrink) + cur_u_lo = max(u_lo, best_u - u_width) + cur_u_hi = min(u_hi, best_u + u_width) + cur_tau_lo = max(tau_lo, best_tau - tau_width) + cur_tau_hi = min(tau_hi, best_tau + tau_width) + if verb > 1: + print( + " U_tau refine=%d best_U=%.6g best_tau=%.6g nll=%.6f" + % (ref, best_u, best_tau, best_nll), + flush=True, + ) + + return ( + best_u, + best_tau, + best_nll, + np.asarray(u_all, dtype=np.float64), + np.asarray(tau_all, dtype=np.float64), + np.asarray(nll_all, dtype=np.float64), + np.asarray(refine_all, dtype=np.int32), + np.asarray(eval_all, dtype=np.int32), + np.asarray(eval_total_all, dtype=np.int32), + ) + + +def main(): + args = parse_args() + print("BSSM-fit args:"); pprint(vars(args)) + t_start = time.time() + device = choose_device() + dtype = torch.float32 if args.dtype == "float32" else torch.float64 + torch.cuda.set_device(device) + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision("high") + torch.manual_seed(int(args.seed)) + torch.cuda.manual_seed_all(int(args.seed)) + np.random.seed(int(args.seed)) + + if not (0.0 < args.u0 < 0.999): + abort_fit("--u0 must satisfy 0 < U < 0.999") + if not args.tau_rec0 > 0.0: + abort_fit("--tau_rec0 must be positive") + if not args.synaptic_tau > 0.0: + abort_fit("--synaptic_tau must be positive") + if not args.kernel_len_steps >= 1: + abort_fit("--kernel_len_steps must be at least 1") + if args.h_chunk_steps < 1: + abort_fit("--h_chunk_steps must be at least 1") + if args.burn_sec <= 0.0: + abort_fit("--burn_sec must be positive") + if args.init_samples <= 1000: + abort_fit("--init_samples must be above 1000") + if not (args.u_bounds[0] > 0.0 and args.u_bounds[0] < args.u_bounds[1] < 0.999): + abort_fit("--u_bounds must satisfy 0 < U_lo < U_hi < 0.999") + if not (args.tau_bounds[0] > 0.0 and args.tau_bounds[0] < args.tau_bounds[1]): + abort_fit("--tau_bounds must satisfy 0 < tau_lo < tau_hi") + if args.num_outer < 1: + abort_fit("--num_outer must be at least 1") + if args.blockA_epochs < 1: + abort_fit("--blockA_epochs must be at least 1") + if args.progress_every < 0: + abort_fit("--progress_every must be nonnegative") + if args.delay_epoch_4_rhoMax < 0: + abort_fit("--delay_epoch_4_rhoMax must be nonnegative") + if args.delay_epoch_4_blockB < 0: + abort_fit("--delay_epoch_4_blockB must be nonnegative") + if args.batch_size < 1: + abort_fit("--batch_size must be at least 1") + if not (0.0 < args.lr_end_factor <= 1.0): + abort_fit("--lr_end_factor must satisfy 0 < lr_end_factor <= 1") + if args.rho_max <= 0.0: + abort_fit("--rho_max must be positive") + if args.u_grid_points < 3 or args.tau_grid_points < 3: + abort_fit("--u_grid_points and --tau_grid_points must each be at least 3") + if args.grid_refine < 0: + abort_fit("--grid_refine must be nonnegative") + if not (0.0 < args.grid_shrink <= 1.0): + abort_fit("--grid_shrink must satisfy 0 < grid_shrink <= 1") + + inp_path = os.path.join(args.basePath, "truthDale") + out_path = args.outPath if args.outPath is not None else os.path.join(args.basePath, "fitBssmStd") + if not os.path.exists(inp_path): + abort_fit("missing input path: %s" % inp_path) + + spikes_ff = os.path.join(inp_path, args.dataName + ".spikes.npz") + if not os.path.exists(spikes_ff): + abort_fit("missing spikes file: %s" % spikes_ff) + + spike_d, spike_md = read_npz_records(spikes_ff, ["spikes"], verb=args.verb > 1) + require_keys(spike_d, ["spikes"], "spike_d") + require_keys(spike_md, ["time_step_sec"], "spike_md") + + spikes_np = np.asarray(spike_d["spikes"]) + if spikes_np.ndim != 2: + abort_fit("spikes array must be 2D") + if np.max(spikes_np) > 1: + abort_fit("BSSM-STD fit expects binary Bernoulli spikes") + dt = float(spike_md["time_step_sec"]) + burn_sec = float(args.burn_sec) + burn_bins = int(math.ceil(burn_sec / dt)) + ( + start_bin, + end_bin, + requested_start_bin, + requested_end_bin, + requested_bins, + spikes_w, + ) = time_window_with_burn(spikes_np, dt, args.time_range_sec, burn_bins) + T_w, N = spikes_w.shape + if T_w != burn_bins + requested_bins: + abort_fit( + "internal time-window mismatch: loaded=%d burn_bins=%d requested_bins=%d" + % (T_w, burn_bins, requested_bins) + ) + if requested_bins < int(args.init_samples): + abort_fit( + "not enough requested post-burn bins for initialization: " + "requested_bins=%d required --init_samples=%d" + % (requested_bins, int(args.init_samples)) + ) + os.makedirs(out_path, exist_ok=True) + + kappa = build_exponential_kernel( + dt=dt, + synaptic_tau=float(args.synaptic_tau), + mem_lag_steps=int(args.kernel_len_steps), + ) + mem_lag_steps = int(kappa.size) + alpha = infer_alpha(kappa) + if mem_lag_steps < 1: + abort_fit("kernel length must be at least 1") + + spikes_t = torch.as_tensor(spikes_w.astype(np.float32), dtype=dtype, device=device) + Y_eval = spikes_t[burn_bins:] + ( + B0, + W0, + init_bins, + init_p_mean, + init_p_plus, + init_p_minus, + W_init_rebin_bins, + ) = init_B_W_from_postburn_spikes( + spikes_w, + burn_bins, + args.init_samples, + args.kernel_len_steps, + args.initW_random, + args.seed, + ) + np_dtype = np.float32 if dtype == torch.float32 else np.float64 + B0 = B0.astype(np_dtype) + W0 = W0.astype(np_dtype) + offdiag_np = ~np.eye(N, dtype=bool) + W0_abs_mean = float(np.mean(np.abs(W0[offdiag_np]))) if np.any(offdiag_np) else 0.0 + + W = torch.tensor(W0, dtype=dtype, device=device, requires_grad=True) + B = torch.tensor(B0, dtype=dtype, device=device, requires_grad=True) + + offdiag_mask = ~torch.eye(N, dtype=torch.bool, device=device) + project_w_(W) + + gen = torch.Generator(device=device) + gen.manual_seed(int(args.seed)) + + U = float(args.u0) + tau_rec = float(args.tau_rec0) + + outer_rows = [] + blockA_outer = [] + blockA_epoch = [] + blockA_bce_loss = [] + blockA_l1_loss = [] + blockA_loss = [] + blockA_lr = [] + joint_search_outer = [] + joint_search_U = [] + joint_search_tau_rec = [] + joint_search_nll = [] + joint_search_refine = [] + joint_search_eval = [] + joint_search_eval_total = [] + blockB_executed = [] + + if args.verb > 0: + print("\nfit5_BSSM_STD_blocks.py", flush=True) + print(" dataName=%s" % args.dataName, flush=True) + print(" device=%s dtype=%s torch=%s" % (device, args.dtype, torch.__version__), flush=True) + print(" gpu=%s" % torch.cuda.get_device_name(device), flush=True) + print( + " bins raw=[%d,%d] requested=[%d,%d] T_raw=%d N=%d dt=%.6g burn_bins=%d eval_bins=%d" + % ( + start_bin, + end_bin, + requested_start_bin, + requested_end_bin, + T_w, + N, + dt, + burn_bins, + int(Y_eval.shape[0]), + ), + flush=True, + ) + print( + " init B from first %d post-burn bins: B_mean=%.6g" + % (init_bins, float(np.mean(B0))), + flush=True, + ) + if float(args.initW_random) > 0.0: + print( + " init W random uniform offdiag in [-%.6g,%.6g]: W_abs_mean=%.6g" + % (float(args.initW_random), float(args.initW_random), W0_abs_mean), + flush=True, + ) + else: + print( + " init W from %d rebinned post-burn samples of %d bins: W_lag1_cov_abs_mean=%.6g" + % (W_init_rebin_bins, int(args.kernel_len_steps), W0_abs_mean), + flush=True, + ) + print( + " kernel: tau_s=%.6g M=%d alpha=%.6g init U=%.6g tau_rec=%.6g" + % (float(args.synaptic_tau), mem_lag_steps, alpha, U, tau_rec), + flush=True, + ) + print(" CUDA H recurrence chunk_steps=%d" % int(args.h_chunk_steps), flush=True) + + total_blockA_epochs = sum( + blockA_epochs_for_outer(args, outer) + for outer in range(int(args.num_outer)) + ) + blockA_epoch_offset = 0 + + for outer in range(int(args.num_outer)): + if args.verb > 0: + print( + "\nOuter %d/%d U=%.6f tau_rec=%.6f" + % (outer + 1, args.num_outer, U, tau_rec), + flush=True, + ) + + torch.cuda.synchronize() + torch.cuda.empty_cache() + progress_line(args, t_start, "start computing H for outer=%d/%d" % (outer + 1, int(args.num_outer))) + H = make_h( + spikes_t, + U, + tau_rec, + dt, + alpha, + mem_lag_steps, + burn_bins, + args.h_chunk_steps, + label="H outer=%d/%d" % (outer + 1, int(args.num_outer)), + verb=args.verb, + ) + + live_std = ( + outer >= int(args.freeze_std_outer) + and outer >= int(args.delay_epoch_4_blockB) + ) + epochs_this_outer = blockA_epochs_for_outer(args, outer) + progress_line( + args, + t_start, + "start Block A outer=%d/%d epochs=%d batch_size=%d" + % (outer + 1, int(args.num_outer), epochs_this_outer, int(args.batch_size)), + ) + bce_losses, l1_losses, losses, lr_values = train_block_a( + H, Y_eval, W, B, args, offdiag_mask, gen, outer, t_start, + epochs=epochs_this_outer, + global_epoch_offset=blockA_epoch_offset, + global_epoch_total=total_blockA_epochs, + ) + blockA_epoch_offset += epochs_this_outer + progress_line(args, t_start, "finished Block A outer=%d/%d" % (outer + 1, int(args.num_outer))) + for ep, val in enumerate(losses): + blockA_outer.append(outer) + blockA_epoch.append(ep) + blockA_bce_loss.append(float(bce_losses[ep])) + blockA_l1_loss.append(float(l1_losses[ep])) + blockA_loss.append(float(val)) + blockA_lr.append(float(lr_values[ep])) + + progress_line(args, t_start, "start evaluating nll_after_blockA outer=%d/%d" % (outer + 1, int(args.num_outer))) + nll_after_A = bernoulli_nll_mean(H, Y_eval, W, B, args.eta_clip, args.batch_size) + progress_line( + args, + t_start, + "finished nll_after_blockA outer=%d/%d nll_A=%.9f" + % (outer + 1, int(args.num_outer), nll_after_A), + ) + + if live_std: + progress_line( + args, + t_start, + "start joint U/tau_rec grid search outer=%d/%d current_U=%.6f current_tau=%.6f" + % (outer + 1, int(args.num_outer), U, tau_rec), + ) + + def obj_std_nll(u_val, tau_val): + H_std = make_h( + spikes_t, + float(u_val), + float(tau_val), + dt, + alpha, + mem_lag_steps, + burn_bins, + args.h_chunk_steps, + verb=0, + ) + nll_std = bernoulli_nll_mean(H_std, Y_eval, W, B, args.eta_clip, args.batch_size) + del H_std + return nll_std + + ( + U, + tau_rec, + nll_std, + vals_u, + vals_tau, + vals_nll, + vals_refine, + vals_eval, + vals_eval_total, + ) = blockB_joint_grid_search( + U, + tau_rec, + args.u_bounds, + args.tau_bounds, + args.u_grid_points, + args.tau_grid_points, + args.grid_refine, + args.grid_shrink, + obj_std_nll, + t_start=t_start, + outer=outer, + num_outer=args.num_outer, + verb=args.verb, + ) + progress_line( + args, + t_start, + "finished joint U/tau_rec grid search outer=%d/%d best_U=%.6f best_tau=%.6f nll=%.9f" + % (outer + 1, int(args.num_outer), U, tau_rec, nll_std), + ) + joint_search_outer.extend([outer] * vals_u.size) + joint_search_U.extend(vals_u.tolist()) + joint_search_tau_rec.extend(vals_tau.tolist()) + joint_search_nll.extend(vals_nll.tolist()) + joint_search_refine.extend(vals_refine.tolist()) + joint_search_eval.extend(vals_eval.tolist()) + joint_search_eval_total.extend(vals_eval_total.tolist()) + blockB_executed.append(1) + nll_outer = float(nll_std) + else: + progress_line( + args, + t_start, + "skip joint U/tau_rec search outer=%d/%d because freeze_std_outer=%d delay_epoch_4_blockB=%d" + % ( + outer + 1, + int(args.num_outer), + int(args.freeze_std_outer), + int(args.delay_epoch_4_blockB), + ), + ) + nll_std = float(nll_after_A) + blockB_executed.append(0) + nll_outer = float(nll_after_A) + + rho_w = spectral_radius_torch(W.detach()) + nz = int(torch.sum(torch.abs(W.detach()[offdiag_mask]) >= float(args.weight_threshold)).item()) + outer_rows.append( + [ + outer, + U, + tau_rec, + nll_after_A, + nll_std, + nll_outer, + rho_w, + nz, + ] + ) + + if args.verb > 0: + print( + " nll_A=%.9f nll=%.9f rhoW=%.4f nz=%d elapsed=%.1fs" + % (nll_after_A, nll_outer, rho_w, nz, time.time() - t_start), + flush=True, + ) + print(" now U=%.6f tau_rec=%.6f" % (U, tau_rec), flush=True) + + del H + torch.cuda.empty_cache() + + outer_arr = np.asarray(outer_rows, dtype=np.float64) + elapsed = time.time() - t_start + + out_stem = args.dataName + out_ff = os.path.join(out_path, out_stem + ".fitBSSMSTD.npz") + + W_fit = W.detach().cpu().numpy().astype(np.float32) + B_fit = B.detach().cpu().numpy().astype(np.float32) + W_init = W0.astype(np.float32) + B_init = B0.astype(np.float32) + out_d = { + "W_fit": W_fit, + "B_fit": B_fit, + "W_init": W_init, + "B_init": B_init, + "init_p_mean": init_p_mean.astype(np.float32), + "init_p_plus": init_p_plus.astype(np.float32), + "init_p_minus": init_p_minus.astype(np.float32), + "W_init_abs_mean": np.asarray([W0_abs_mean], dtype=np.float64), + "initW_random": np.asarray([float(args.initW_random)], dtype=np.float64), + "W_init_rebin_steps": np.asarray([int(args.kernel_len_steps)], dtype=np.int32), + "W_init_rebin_bins": np.asarray([int(W_init_rebin_bins)], dtype=np.int64), + "outer_idx": outer_arr[:, 0].astype(np.int32), + "U_outer": outer_arr[:, 1], + "tau_rec_outer": outer_arr[:, 2], + "nll_after_blockA": outer_arr[:, 3], + "nll_after_std_grid": outer_arr[:, 4], + "nll_outer": outer_arr[:, 5], + "spectral_radius_outer": outer_arr[:, 6], + "nz_weight_outer": outer_arr[:, 7].astype(np.int32), + "blockB_executed": np.asarray(blockB_executed, dtype=np.int32), + "blockA_outer": np.asarray(blockA_outer, dtype=np.int32), + "blockA_epoch": np.asarray(blockA_epoch, dtype=np.int32), + "blockA_bce_loss": np.asarray(blockA_bce_loss, dtype=np.float64), + "blockA_l1_loss": np.asarray(blockA_l1_loss, dtype=np.float64), + "blockA_loss": np.asarray(blockA_loss, dtype=np.float64), + "blockA_lr": np.asarray(blockA_lr, dtype=np.float64), + "joint_search_outer": np.asarray(joint_search_outer, dtype=np.int32), + "joint_search_U": np.asarray(joint_search_U, dtype=np.float64), + "joint_search_tau_rec": np.asarray(joint_search_tau_rec, dtype=np.float64), + "joint_search_nll": np.asarray(joint_search_nll, dtype=np.float64), + "joint_search_refine": np.asarray(joint_search_refine, dtype=np.int32), + "joint_search_eval": np.asarray(joint_search_eval, dtype=np.int32), + "joint_search_eval_total": np.asarray(joint_search_eval_total, dtype=np.int32), + "kappa_fit": kappa.astype(np.float32), + "alpha_fit": np.asarray([alpha], dtype=np.float64), + "burn_bins": np.asarray([burn_bins], dtype=np.int32), + "time_range_bins": np.asarray([start_bin, end_bin], dtype=np.int64), + "requested_time_range_bins": np.asarray([requested_start_bin, requested_end_bin], dtype=np.int64), + "fit_time_range_bins": np.asarray([start_bin + burn_bins, end_bin], dtype=np.int64), + "init_time_range_bins": np.asarray( + [start_bin + burn_bins, start_bin + burn_bins + init_bins - 1], + dtype=np.int64, + ), + "time_step_sec": np.asarray([dt], dtype=np.float64), + "num_neurons": np.asarray([N], dtype=np.int32), + "num_input_bins": np.asarray([T_w], dtype=np.int64), + "num_eval_bins": np.asarray([Y_eval.shape[0]], dtype=np.int64), + "num_init_bins": np.asarray([init_bins], dtype=np.int64), + "U_init": np.asarray([float(args.u0)], dtype=np.float64), + "tau_rec_init": np.asarray([float(args.tau_rec0)], dtype=np.float64), + "u_bounds": np.asarray(args.u_bounds, dtype=np.float64), + "tau_bounds": np.asarray(args.tau_bounds, dtype=np.float64), + "u_grid_points": np.asarray([int(args.u_grid_points)], dtype=np.int32), + "tau_grid_points": np.asarray([int(args.tau_grid_points)], dtype=np.int32), + "grid_refine": np.asarray([int(args.grid_refine)], dtype=np.int32), + "grid_shrink": np.asarray([float(args.grid_shrink)], dtype=np.float64), + "delay_epoch_4_blockB": np.asarray([int(args.delay_epoch_4_blockB)], dtype=np.int32), + "eta_clip": np.asarray([float(args.eta_clip)], dtype=np.float64), + "lr_w": np.asarray([float(args.lr_w)], dtype=np.float64), + "lr_end_factor": np.asarray([float(args.lr_end_factor)], dtype=np.float64), + "lambda_l1": np.asarray([float(args.lambda_l1)], dtype=np.float64), + "rho_max": np.asarray([float(args.rho_max)], dtype=np.float64), + "weight_threshold": np.asarray([float(args.weight_threshold)], dtype=np.float64), + "h_chunk_steps": np.asarray([int(args.h_chunk_steps)], dtype=np.int32), + } + + out_md = { + "short_name": out_stem, + "fit_type": "BSSM_STD_blockFit", + "provenance": { + "script": os.path.basename(__file__), + "spikesData_file": args.dataName, + }, + "config": vars(args), + "model": { + "num_neurons": int(N), + "num_input_bins": int(T_w), + "num_eval_bins": int(Y_eval.shape[0]), + "num_requested_bins": int(requested_bins), + "num_init_bins": int(init_bins), + "init_start_bin": int(start_bin + burn_bins), + "init_end_bin": int(start_bin + burn_bins + init_bins - 1), + "time_step_sec": float(dt), + "kernel_synaptic_tau": float(args.synaptic_tau), + "kernel_len_steps": int(mem_lag_steps), + "kernel_alpha": float(alpha), + }, + "result": { + "U_final": float(out_d["U_outer"][-1]), + "tau_rec_final": float(out_d["tau_rec_outer"][-1]), + "nll_final": float(out_d["nll_outer"][-1]), + "elapsed_sec": round(elapsed, 3), + }, + } + + if args.verb > 1: + pprint(out_md) + write_data_npz(out_d, out_ff, metaD=out_md, verb=max(1, int(args.verb))) + + if args.verb > 0: + print("\nSaved: %s" % out_ff) + print( + " final_U=%.6f final_tau=%.6f elapsed=%.1fs" + % (out_d["U_outer"][-1], out_d["tau_rec_outer"][-1], elapsed) + ) + print(f" basePath={args.basePath}") + print(f" ./eval_BSSM_fit.py --basePath {args.basePath} --dataName {args.dataName} -p a b\n") + + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver5_synapDepres/gen5_BSSM_STD_spikes.py b/causal_net/nonStation_ver5_synapDepres/gen5_BSSM_STD_spikes.py new file mode 100755 index 00000000..51d7792a --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/gen5_BSSM_STD_spikes.py @@ -0,0 +1,716 @@ +#!/usr/bin/env python3 +""" +Generate synthetic spike trains for the Real-Time Adaptive Bernoulli +State-Space Network Model (BSSM) with short-term synaptic depression (STD) +and a lag-M synaptic kernel. + +Output files: + + /truthDale/.simTruth.npz + /truthDale/.spikes.npz + +The one architectural exception requested by the user is preserved: +the connectivity matrix is generated from 2D spatial placement exactly in the +style of the current Dale generator, but the diagonal is forced to zero so the +network has no self-loops. +""" + +import argparse +import os +import sys +import time +from pprint import pprint + +import numpy as np + + +from toolbox.Util_NumpyIO import write_data_npz +from UtilDalePoisson5 import estimate_rates + + +if sys.version_info < (3, 0): + sys.stderr.write("ERROR: gen5_BSSM_STD_spikes.py requires Python 3.0 or newer.\n") + sys.exit(1) + + +def _build_placement_grid(length_x, height_y, d_min): + """Rectangular grid nodes on [0,L] x [0,H] with spacing d_min.""" + m_max = int(np.floor(length_x / d_min)) + n_max = int(np.floor(height_y / d_min)) + gx = (np.arange(0, m_max + 1, dtype=float) * d_min).reshape(-1, 1) + gy = (np.arange(0, n_max + 1, dtype=float) * d_min).reshape(1, -1) + xs = np.broadcast_to(gx, (gx.size, gy.size)).ravel() + ys = np.broadcast_to(gy, (gx.size, gy.size)).ravel() + grid = np.column_stack([xs, ys]) + return grid, (m_max + 1) * (n_max + 1) + + +def _sample_targets_without_replacement(j, affinity_column, k_j, rng): + """Draw k_j distinct postsynaptic targets i != j with P(i) proportional to affinity.""" + n_units = affinity_column.shape[0] + p = np.asarray(affinity_column, dtype=float).copy() + p[j] = 0.0 + norm = p.sum() + if norm <= 0: + raise RuntimeError("zero affinity sum for presynaptic neuron %d" % j) + p /= norm + return rng.choice(n_units, size=k_j, replace=False, p=p) + + +def generate_spatial_dale_network( + n_units, + n_excite, + length_x, + height_y, + d_min, + k_min, + k_max, + placement_ker_delta, + weight_var, + spectral_radius_target, + rng=None, + verb=1, +): + """ + Generate the spatial/Dale recurrent matrix used by the old generator, + but with a zero diagonal (no self-loops). + + Returns + ------- + W_true : (N, N) + Spectrally scaled recurrent weight matrix in postsynaptic-row, + presynaptic-column convention. + sign_true : (N, N) + Signed binary topology: +1 excitatory edge, -1 inhibitory edge, 0 absent. + positions : (N, 2) + Neuron positions sorted by x then y. + tau : (N,) + 0 excitatory, 1 inhibitory. + rho0 : float + Spectral radius before rescaling. + distance_matrix : (N, N) + Pairwise Euclidean distances with zero diagonal. + """ + if rng is None: + rng = np.random.default_rng() + + n_inhib = n_units - n_excite + assert n_excite > 0 and n_inhib > 0 + assert k_min >= 1 and k_max <= n_units - 1 and k_min <= k_max + assert float(placement_ker_delta) > 0 + assert 0 < weight_var < 1 + assert 0 < spectral_radius_target < 1 + + balance_scale = n_excite / float(n_inhib) + + grid, n_grid = _build_placement_grid(length_x, height_y, d_min) + if n_grid < n_units: + raise ValueError( + "grid has only %d nodes; need N <= N_G (reduce d_min or increase L/H)" + % n_grid + ) + + sample_idx = rng.choice(n_grid, size=n_units, replace=False) + positions = grid[sample_idx].astype(np.float64) + order_x = np.lexsort((positions[:, 1], positions[:, 0])) + positions = positions[order_x] + + exc_mask = np.zeros(n_units, dtype=bool) + exc_mask[rng.choice(n_units, size=n_excite, replace=False)] = True + tau = np.where(exc_mask, 0, 1).astype(np.int32) + sigma = np.where(exc_mask, 1, -1).astype(np.int8) + + diff = positions[:, np.newaxis, :] - positions[np.newaxis, :, :] + distance_matrix = np.sqrt(np.sum(diff * diff, axis=2)) + np.fill_diagonal(distance_matrix, np.inf) + with np.errstate(divide="ignore"): + affinity = np.power(distance_matrix, -float(placement_ker_delta)) + np.fill_diagonal(affinity, 0.0) + + sign_true = np.zeros((n_units, n_units), dtype=np.int8) + for j in range(n_units): + kj = int(rng.integers(k_min, k_max + 1)) + targets = _sample_targets_without_replacement(j, affinity[:, j], kj, rng) + sign_true[targets, j] = sigma[j] + + raw_W = np.zeros((n_units, n_units), dtype=np.float64) + for j in range(n_units): + edge_mask = sign_true[:, j] != 0 + if not np.any(edge_mask): + continue + weights = rng.uniform(1.0 - weight_var, 1.0 + weight_var, size=int(np.sum(edge_mask))) + if sigma[j] < 0: + weights *= -balance_scale + raw_W[edge_mask, j] = weights + + np.fill_diagonal(raw_W, 0.0) + + eigvals = np.linalg.eigvals(raw_W) + rho0 = float(np.max(np.abs(eigvals))) + if verb > 0: + print("initW current_rho :", rho0) + if rho0 <= 0: + raise RuntimeError("spectral radius of raw_W is zero (degenerate)") + + W_true = raw_W * (spectral_radius_target / rho0) + np.fill_diagonal(W_true, 0.0) + + distance_save = distance_matrix.copy() + np.fill_diagonal(distance_save, 0.0) + return W_true, sign_true, positions, tau, rho0, distance_save + + +def summarize_pairwise_distances(distance_matrix, verb=1): + """Mean and median Euclidean distance over unique unordered pairs (i < j).""" + distance_matrix = np.asarray(distance_matrix, dtype=float) + if distance_matrix.ndim != 2 or distance_matrix.shape[0] != distance_matrix.shape[1]: + raise ValueError("distance_matrix must be a square matrix") + n_units = distance_matrix.shape[0] + if n_units < 2: + raise ValueError("at least 2 neurons are required for pairwise distances") + iu = np.triu_indices(n_units, k=1) + distances = distance_matrix[iu] + out = { + "mean": float(np.mean(distances)), + "median": float(np.median(distances)), + "n_nodes": int(n_units), + "n_pairs": int(distances.size), + } + if verb > 0: + print( + "Pairwise distance (unique pairs): mean=%.6g median=%.6g (N=%d, pairs=%d)" + % (out["mean"], out["median"], out["n_nodes"], out["n_pairs"]) + ) + return out + + +def _sigmoid(logits): + logits = np.asarray(logits, dtype=np.float64) + return 1.0 / (1.0 + np.exp(-logits)) + + +def _logit(prob): + prob = np.asarray(prob, dtype=np.float64) + return np.log(prob) - np.log1p(-prob) + + +def set_flat_selfSpiking(n_units, idle_rate_hz, dt, tau, rng): + """ + Bernoulli baseline logits derived from target idle firing rates in Hz. + + The returned B_true is a log-odds vector. If recurrent input were zero, + neuron i would spike with probability p_i ~ rate_i * dt in each bin. + """ + tau = np.asarray(tau).reshape(-1) + if tau.shape[0] != n_units: + raise ValueError("tau length %d != n_units %d" % (tau.shape[0], n_units)) + if tau.dtype.kind not in "iu": + raise ValueError("tau must be integer dtype") + exc = tau.astype(np.int64) == 0 + inh = tau.astype(np.int64) == 1 + if not np.any(exc) or not np.any(inh): + raise ValueError("tau must label at least one excitatory and one inhibitory neuron") + + idle_rate_hz = np.asarray(idle_rate_hz, dtype=float) + if idle_rate_hz.shape != (2,): + raise ValueError("idleRate must be a length-2 range [min_hz, max_hz]") + if idle_rate_hz[0] <= 0 or idle_rate_hz[1] <= idle_rate_hz[0]: + raise ValueError("idleRate must satisfy 0 < min_hz < max_hz") + + per_neuron_rate_hz = rng.uniform(idle_rate_hz[0], idle_rate_hz[1], size=n_units) + p_bin = np.clip(per_neuron_rate_hz * dt, 1e-6, 1.0 - 1e-6) + return _logit(p_bin) + + +def build_exponential_kernel(dt, synaptic_tau, mem_lag_steps, verb=0): + """ + Exponential lag-M kernel: + kappa_l = (1 - alpha) alpha^(l-1), alpha = exp(-dt / tau_s) + + This finite kernel is implemented with the equivalent tail-corrected + recursion + h_{t+1} = alpha h_t + (1 - alpha) u_t + - (1 - alpha) alpha^M u_{t-M}. + """ + synaptic_tau = float(synaptic_tau) + if synaptic_tau <= 0: + raise ValueError("synaptic_tau must be positive") + mem_lag_steps = int(mem_lag_steps) + if mem_lag_steps < 1: + raise ValueError("mem_lag_steps must be >= 1") + + alpha = float(np.exp(-dt / synaptic_tau)) + ell = np.arange(mem_lag_steps, dtype=np.float64) + kernel = (1.0 - alpha) * np.power(alpha, ell) + if verb > 0: + print( + " exponential kernel: len=%d, span=%.6g sec, tau_s=%.6g sec, alpha=%.6g, sum=%.12g" + % (mem_lag_steps, mem_lag_steps * float(dt), synaptic_tau, alpha, float(np.sum(kernel))) + ) + return kernel, alpha + + +def gen_bssm_std_bernoulli( + num_steps, + dt, + W_true, + B_true, + tau, + std_u, + std_tau_rec, + synaptic_alpha, + mem_lag_steps, + logit_clip, + rng, + verb=0, +): + """ + Simulate the BSSM-STD Bernoulli model. + + h_t is the filtered presynaptic drive available at the start of bin t. + x_t is the available resource fraction at the start of bin t. + """ + if W_true is None: + raise ValueError("W_true cannot be None") + n_units = W_true.shape[0] + if W_true.shape != (n_units, n_units): + raise ValueError("W_true must be square") + tau = np.asarray(tau).reshape(-1) + if tau.shape[0] != n_units: + raise ValueError("tau length %d != W_true.shape[0] %d" % (tau.shape[0], n_units)) + if tau.dtype.kind not in "iu": + raise ValueError("tau must be integer dtype") + tau_i = tau.astype(np.int64) + if np.any((tau_i != 0) & (tau_i != 1)): + raise ValueError("tau entries must be 0 or 1") + if B_true.shape != (n_units,): + raise ValueError("B_true shape %s does not match (%d,)" % (B_true.shape, n_units)) + if not (0 < std_u <= 1): + raise ValueError("std_u must satisfy 0 < std_u <= 1") + if std_tau_rec <= 0: + raise ValueError("std_tau_rec must be positive") + if not (0 <= synaptic_alpha < 1): + raise ValueError("synaptic_alpha must satisfy 0 <= alpha < 1") + mem_lag_steps = int(mem_lag_steps) + if mem_lag_steps < 1: + raise ValueError("mem_lag_steps must be >= 1") + + x = np.ones(n_units, dtype=np.float64) + h = np.zeros(n_units, dtype=np.float64) + u_history = np.zeros((mem_lag_steps, n_units), dtype=np.float64) + recovery_decay = float(np.exp(-dt / std_tau_rec)) + synaptic_gain = 1.0 - synaptic_alpha + tail_gain = synaptic_gain * float(np.power(synaptic_alpha, mem_lag_steps)) + + spikes = np.zeros((num_steps, n_units), dtype=np.uint8) + x_true = np.zeros((num_steps, n_units), dtype=np.float32) + u_true = np.zeros((num_steps, n_units), dtype=np.float32) + h_true = np.zeros((num_steps, n_units), dtype=np.float32) + p_true = np.zeros((num_steps, n_units), dtype=np.float32) + n_exc = int(np.sum(tau_i == 0)) + progress_stride = max(1, num_steps // 4) + + if verb > 0: + print("\n=== Generating BSSM-STD Bernoulli spikes ===") + print( + "steps=%d, dt=%.6g, neurons=%d, excit=%d, std_u=%.6g, tau_rec=%.6g, alpha=%.6g, gain=%.6g, M=%d" + % ( + num_steps, + dt, + n_units, + n_exc, + std_u, + std_tau_rec, + synaptic_alpha, + synaptic_gain, + mem_lag_steps, + ) + ) + print( + "W stats: min=%.3f, max=%.3f, mean=%.3f" + % (float(np.min(W_true)), float(np.max(W_true)), float(np.mean(W_true))) + ) + print( + "b stats: min=%.3f, max=%.3f, mean=%.3f" + % (float(np.min(B_true)), float(np.max(B_true)), float(np.mean(B_true))) + ) + + t_start = time.time() + for t in range(num_steps): + old_u = u_history[t % mem_lag_steps].copy() + x_true[t] = x + h_true[t] = h + + logits_t = np.clip(B_true + W_true @ h, -logit_clip, logit_clip) + p_t = _sigmoid(logits_t) + s_t = (rng.random(n_units) < p_t).astype(np.uint8) + u_t = std_u * x * s_t + + spikes[t] = s_t + p_true[t] = p_t + u_true[t] = u_t + + if verb > 0 and t < 5: + print( + "t=%d spikes=%d mean_p=%.4f mean_x=%.4f mean_h=%.4f" + % (t, int(np.sum(s_t)), float(np.mean(p_t)), float(np.mean(x)), float(np.mean(h))) + ) + + x_depleted = x * (1.0 - std_u * s_t) + x = 1.0 - (1.0 - x_depleted) * recovery_decay + x = np.clip(x, 0.0, 1.0) + h = synaptic_alpha * h + synaptic_gain * u_t - tail_gain * old_u + h = np.maximum(h, 0.0) + u_history[t % mem_lag_steps] = u_t + + if verb > 0 and t > 0 and t % progress_stride == 0: + elapsed = time.time() - t_start + print( + " Progress: %d/%d steps (%.1f%%) - total spikes this bin: %d elaT=%.1fs" + % (t, num_steps, 100.0 * t / num_steps, int(np.sum(s_t)), elapsed) + ) + + if verb > 0: + print("Simulation complete. Final total spikes=%d" % int(np.sum(spikes[-1]))) + print( + "Spike data stats: min=%d, max=%d, mean=%.4f, total spikes=%d" + % ( + int(np.min(spikes)), + int(np.max(spikes)), + float(np.mean(spikes)), + int(np.sum(spikes)), + ) + ) + + out = { + "spikes": spikes, + "x_true": x_true, + "u_true": u_true, + "h_true": h_true, + "p_true": p_true, + "x_final": x.astype(np.float32), + "h_final": h.astype(np.float32), + } + return out + + +def main(): + print("=" * 60) + print("BSSM-STD BERNOULLI SIMULATION") + print("=" * 60) + + parser = argparse.ArgumentParser( + description="Simulate the BSSM-STD Bernoulli network with spatial/Dale connectivity." + ) + p = parser.add_argument + p("--synaptic_tau", type=float, default=0.005, help="Synaptic decay constant tau_s in seconds.") + p("--std_recovery_tau", type=float, default=0.300, help="STD recovery constant tau_rec in seconds.") + p("--std_u", type=float, default=0.4, help="STD utilization U in (0, 1].") + p("--kernel_len_steps", type=int, default=25, + help="Lag-M synaptic-kernel length in bins.") + p("--num_neurons", type=int, default=50, help="Total number of neurons in the network.") + p("--num_excite", type=int, required=True, help="Number of excitatory neurons.") + p("--placement_H_L_delta", type=float, nargs=3, default=[1.0, 2.0, 2.0], + metavar=("placement_H", "placement_L", "placement_ker_delta"), + help="Placement: [0,H] height, [0,L] width, and distance-kernel exponent delta > 0.") + p("--placement_min_dist", type=float, default=0.01, help="Grid spacing d_min; minimum inter-neuron distance.") + p("--edge_prob", type=float, nargs=2, default=[0.05, 0.2], + help="Out-degree range as fractions of N: k_min=max(1,floor(lo*N)), k_max=min(N-1,floor(hi*N)).") + p("--init_weight_var", type=float, default=0.4, help="Fractional weight variation v for Uniform(1-v,1+v).") + p("--num_steps", type=int, default=5_001, help="Number of time steps for simulation.") + p("--step_size", type=float, default=0.001, help="Time bin width dt in seconds.") + p("--spectral_radius", type=float, default=0.90, help="Target spectral radius for the off-diagonal recurrent matrix.") + p("--idleRate", type=float, nargs=2, default=[30.0, 50.0], help="Range of baseline firing rates [min, max] in Hz.") + p("--logit_clip", type=float, default=20.0, help="Clip logits to [-logit_clip, +logit_clip].") + p("-v", "--verb", type=int, default=1, help="Verbosity level (0=quiet, 1=normal).") + p("--dataName", type=str, required=True, help="Base name for output files.") + p("--basePath", type=str, default="/pscratch/sd/b/balewski/2026_causalNet_tmp/", + help="Output directory root; files are written under /truthDale/.") + + np.set_printoptions(precision=3, suppress=True) + args = parser.parse_args() + print("gen args:", vars(args), "\n") + placement_H = float(args.placement_H_L_delta[0]) + placement_L = float(args.placement_H_L_delta[1]) + placement_ker_delta = float(args.placement_H_L_delta[2]) + if not (placement_H > 0 and placement_L > 0): + raise ValueError("placement_H_L_delta requires positive H and L") + if placement_ker_delta <= 0: + raise ValueError("placement_H_L_delta third value (placement_ker_delta) must be positive") + + placement_min_dist = float(args.placement_min_dist) + if placement_min_dist <= 0: + raise ValueError("placement_min_dist must be positive") + + synaptic_tau = float(args.synaptic_tau) + std_tau_rec = float(args.std_recovery_tau) + if synaptic_tau <= 0: + raise ValueError("synaptic_tau must be positive") + if std_tau_rec <= 0: + raise ValueError("std_recovery_tau must be positive") + if not (0 < args.std_u <= 1): + raise ValueError("std_u must satisfy 0 < U <= 1") + if args.logit_clip <= 0: + raise ValueError("logit_clip must be positive") + + n_units = int(args.num_neurons) + if n_units < 10: + raise ValueError("num_neurons must be >= 10") + if args.num_excite < 5 or args.num_excite >= n_units: + raise ValueError("num_excite must satisfy 5 <= num_excite < num_neurons") + if args.num_steps < 1000: + raise ValueError("num_steps must be >= 1000") + if args.step_size <= 0: + raise ValueError("step_size must be positive") + if args.idleRate[0] <= 0 or args.idleRate[1] <= args.idleRate[0]: + raise ValueError("idleRate must satisfy 0 < min < max") + if args.idleRate[1] * args.step_size >= 0.5: + raise ValueError("idleRate[1] * step_size is too large for a Bernoulli approximation") + + total_time_sec = float(args.num_steps) * float(args.step_size) + var_t_window = min(5.0, max(0.5, total_time_sec / 4.0)) + + if args.kernel_len_steps < 1: + raise ValueError("kernel_len_steps must be >= 1") + mem_lag_steps = int(args.kernel_len_steps) + kernel_span_sec = mem_lag_steps * float(args.step_size) + kernel_capture_fraction = float(1.0 - np.exp(-kernel_span_sec / synaptic_tau)) + if args.verb > 0 and kernel_span_sec < 3.0 * synaptic_tau: + print( + "WARNING: kernel lag span M*dt=%.6g sec is less than 3*tau_s=%.6g sec; " + "the fast synaptic tail will be strongly truncated." + % (kernel_span_sec, 3.0 * synaptic_tau) + ) + + prob_lo = float(args.edge_prob[0]) + prob_hi = float(args.edge_prob[1]) + if not (0 < prob_lo < prob_hi <= 1.0): + raise ValueError("edge_prob must satisfy 0 < lo < hi <= 1") + k_out_min = max(1, int(np.floor(prob_lo * n_units))) + k_out_max = int(np.floor(prob_hi * n_units)) + k_out_max = max(k_out_min, k_out_max) + k_out_max = min(k_out_max, n_units - 1) + + print("gen BSSM-STD args:"); pprint( vars(args)) + + rng = np.random.default_rng() + out_path = os.path.join(args.basePath, "truthDale") + if not os.path.exists(args.basePath): + raise FileNotFoundError("missing basePath: %s" % args.basePath) + if not os.path.exists(out_path): + raise FileNotFoundError("missing output path: %s" % out_path) + + print("\n%s" % ("=" * 60)) + print(" Spectral radius target: R=%.3f" % float(args.spectral_radius)) + print("%s" % ("=" * 60)) + + W_true, sign_true, positions, tau, rho0, distance_matrix = generate_spatial_dale_network( + n_units=n_units, + n_excite=args.num_excite, + length_x=placement_L, + height_y=placement_H, + d_min=placement_min_dist, + k_min=k_out_min, + k_max=k_out_max, + placement_ker_delta=placement_ker_delta, + weight_var=args.init_weight_var, + spectral_radius_target=args.spectral_radius, + rng=rng, + verb=args.verb, + ) + + if args.verb > 0: + summarize_pairwise_distances(distance_matrix, verb=1) + total_connections = W_true.size + zero_connections = int(np.sum(np.abs(W_true) < 1e-10)) + non_zero_connections = total_connections - zero_connections + sparsity = zero_connections / float(total_connections) + print( + "Matrix sparsity: %.1f%% (%d/%d connections are zero)" + % (100.0 * sparsity, zero_connections, total_connections) + ) + print("Non-zero connections: %d (%.1f%%)" % (non_zero_connections, 100.0 * (1.0 - sparsity))) + + kappa_true, synaptic_alpha = build_exponential_kernel( + dt=float(args.step_size), + synaptic_tau=synaptic_tau, + mem_lag_steps=mem_lag_steps, + verb=args.verb, + ) + + B_true = set_flat_selfSpiking( + n_units=n_units, + idle_rate_hz=args.idleRate, + dt=float(args.step_size), + tau=tau, + rng=rng, + ) + + dale_conf = { + "num_neurons": args.num_neurons, + "num_excite": args.num_excite, + "spectral_radius": args.spectral_radius, + "placement_L": placement_L, + "placement_H": placement_H, + "placement_min_dist": placement_min_dist, + "edge_prob": [prob_lo, prob_hi], + "k_out_min": k_out_min, + "k_out_max": k_out_max, + "placement_ker_delta": placement_ker_delta, + "init_weight_var": args.init_weight_var, + "rho0_unscaled": rho0, + "idleRate": args.idleRate, + "model_name": "BSSM_STD", + "mem_lag_steps": int(mem_lag_steps), + "kernel_span_sec": float(kernel_span_sec), + "kernel_capture_fraction": float(kernel_capture_fraction), + "synaptic_tau": float(synaptic_tau), + "std_tau_rec": float(std_tau_rec), + "std_u": float(args.std_u), + "self_loops": False, + } + evol_conf = { + "num_steps": args.num_steps, + "step_size": args.step_size, + "evol_time": args.num_steps * args.step_size, + "logit_clip": args.logit_clip, + "model_name": "BSSM_STD", + "mem_lag_steps": int(mem_lag_steps), + "kernel_span_sec": float(kernel_span_sec), + } + + print("\n%s" % ("=" * 60)) + print(" BSSM-STD Bernoulli spike generation") + print("%s" % ("=" * 60)) + print("Dale configuration:") + pprint(dale_conf) + + start_time = time.time() + sim_out = gen_bssm_std_bernoulli( + num_steps=args.num_steps, + dt=float(args.step_size), + W_true=W_true, + B_true=B_true, + tau=tau, + std_u=float(args.std_u), + std_tau_rec=float(std_tau_rec), + synaptic_alpha=synaptic_alpha, + mem_lag_steps=mem_lag_steps, + logit_clip=float(args.logit_clip), + rng=rng, + verb=args.verb, + ) + sim_time = time.time() - start_time + print("Spike generation completed in %.1f seconds" % sim_time) + + spikes = sim_out["spikes"] + max_samples = 100_000 + stats_dict, rates_dict, _ = estimate_rates( + spikes, + args.step_size, + tau, + max_samples, + args.spectral_radius, + varTwindow=var_t_window, + mxNn=5, + verb=0, + ) + stats_dict["var_time_window_sec"] = float(var_t_window) + stats_dict["max_samples"] = int(max_samples) + + trueD = { + "W_true": W_true, + "B_true": B_true, + "sign_true": sign_true, + "node_positions": positions, + "node_is_inhibitory": tau, + "node_distance_matrix": distance_matrix, + "kappa_true": kappa_true.astype(np.float32), + "x_final": sim_out["x_final"], + "h_final": sim_out["h_final"], + "x_true": sim_out["x_true"], + "u_true": sim_out["u_true"], + "h_true": sim_out["h_true"], + "p_true": sim_out["p_true"], + } + + max_spikes_per_neuron = spikes.max(axis=0) + sp_min = float(max_spikes_per_neuron.min()) + sp_avg = float(max_spikes_per_neuron.mean()) + sp_max = float(max_spikes_per_neuron.max()) + sp_pc = np.percentile(max_spikes_per_neuron, [25, 50, 75]).tolist() + + if args.verb > 0: + print("\nSpike count statistics (max per neuron over all bins):") + print(" min: %.1f, avg: %.1f, max: %.1f" % (sp_min, sp_avg, sp_max)) + print(" percentiles [25, 50, 75]: %s" % sp_pc) + + trueMD = { + "dale_conf": dale_conf, + "evol_conf": evol_conf, + "short_name": args.dataName, + "provenance": { + "state_model_file": args.dataName, + "generator_script": os.path.basename(__file__), + }, + "max_spike_stats": { + "min": sp_min, + "avg": sp_avg, + "max": sp_max, + "percentiles": sp_pc, + }, + } + + spikeD = { + "spikes": spikes.astype(np.uint8), + "single_rates": rates_dict["single_rates"], + "single_rates_var": rates_dict["single_rates_var"], + "single_fano_fact": rates_dict["single_fano_fact"], + } + spikeMD = { + "short_name": args.dataName, + "time_step_sec": args.step_size, + "data_type": "simBSSM_STD", + "logit_clip": args.logit_clip, + "model_name": "BSSM_STD", + "placement_L": float(placement_L), + "placement_H": float(placement_H), + "placement_min_dist": float(placement_min_dist), + "mem_lag_steps": int(mem_lag_steps), + "kernel_span_sec": float(kernel_span_sec), + } + + out_truth = os.path.join(out_path, args.dataName + ".simTruth.npz") + write_data_npz(trueD, out_truth, metaD=trueMD) + if args.verb > 1: + pprint(trueMD) + + out_spikes = os.path.join(out_path, args.dataName + ".spikes.npz") + write_data_npz(spikeD, out_spikes, metaD=spikeMD) + if args.verb > 1: + pprint(spikeMD) + + print("\nSimulation completed successfully!") + print( + "\nRate Summary %s N=%d exc=%d, R=%.3f" + % (args.dataName, args.num_neurons, args.num_excite, float(args.spectral_radius)) + ) + s = stats_dict + print( + " rates: all=%14.1f exc=%14.1f inh=%14.1f" + % (s["avg_spike_rate_all"], s["avg_spike_rate_excit"], s["avg_spike_rate_inhib"]) + ) + print("\nNext step commands:") + print(" basePath=" + args.basePath) + print( + " ./fit5_BSSM_STD_blocks.py --basePath $basePath --dataName %s --burn_sec 2 --init_samples 50000 --u_bounds 0.2 0.8 --tau_bounds 0.1 0.8" + % args.dataName + ) + print( + " ./eval_BSSM_fit.py --basePath $basePath --dataName %s -X" + % args.dataName + ) + + +if __name__ == "__main__": + main() diff --git a/causal_net/nonStation_ver5_synapDepres/toolbox b/causal_net/nonStation_ver5_synapDepres/toolbox new file mode 120000 index 00000000..ac05dc12 --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/toolbox @@ -0,0 +1 @@ +../toolbox \ No newline at end of file diff --git a/causal_net/nonStation_ver5_synapDepres/view_daleMatrix5.py b/causal_net/nonStation_ver5_synapDepres/view_daleMatrix5.py new file mode 100755 index 00000000..2cb73a59 --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/view_daleMatrix5.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +""" +Visualization tool for BSSM-STD network data. + +Loads simulation output (simTruth.npz and spikes.npz) produced by +gen5_BSSM_STD_spikes.py. Neuron order follows placement; use +node_is_inhibitory in simTruth for E/I (not index blocks). + +Available plots (-p flag): + a Recurrent weight matrix W (color-coded) + eigenvalue scatter + with spectral-radius circle + b Weight histogram, in/out-degree counts, empirical firing-rate histogram, + and baseline-vs-empirical rate scatter + c Latent BSSM-STD summaries: population resource trace, release/filter trace, + mean spike-probability vs empirical rate, and mean resource vs rate + d 2D neuron placement: triangles=excitatory, squares=inhibitory; + red=outgoing excitatory edges, blue=outgoing inhibitory edges (presynaptic) + e Connection-length histogram, lag-M synaptic kernel, and STD recovery curve + f Pseudospectral contour plot with eigenvalue overlay + g Topology overview: signed off-diagonal matrix, distance histogram, + and placement topology on one canvas + +Usage: + ./view_daleMatrix5.py --dataName daleN100_9fbe7f -p a b c d e f g +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os +from pprint import pprint +import numpy as np +from PlotterDaleMatrix import Plotter +from toolbox.Util_NumpyIO import read_data_npz +import argparse + + +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize simulated BSSM-STD network data") + parser.add_argument("-v", "--verbosity", type=int, help="increase output verbosity", default=1, dest="verb") + parser.add_argument( + "-p", + "--showPlots", + default="a b", + nargs="+", + help="plot letters: a=W+eigen, b=weights/rates, c=latent summaries, d=placement, e=dist+kernel+STD, f=pseudospectra, g=topo overview", + ) + parser.add_argument("-X", "--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument("--basePath", default="/private/tmp/2025_causalNet_tmp/", help="head dir for input data") + parser.add_argument("--dataName", default="daleN150_448b86", help="simulated Dale network base name") + + args = parser.parse_args() + + args.inpPath = os.path.join(args.basePath, "truthDale") + args.outPath = os.path.join(args.basePath, "plots") + args.showPlots = "".join(args.showPlots) + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + assert os.path.exists(args.basePath) + return args + + +if __name__ == "__main__": + args = get_parser() + np.set_printoptions(precision=3) + + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 0) + assert trueD["b_true"].ndim == 1, "b_true must be shape (N,) from gen5_BSSM_STD_spikes" + assert trueMD["dale_conf"]["model_name"] == "BSSM_STD" + if args.verb > 1: + print("\nSimulation Truth Metadata:") + pprint(trueMD) + + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb > 0) + if args.verb > 1: + print("\nSpike Data Metadata:") + pprint(spikeMD) + + trueMD["short_name"] = args.dataName + + R_sel = trueMD["dale_conf"]["spectral_radius"] + print(f"\nSpectral radius: R={R_sel:.3f}") + + trueMD["sel_spect_radius"] = R_sel + + args.prjName = args.dataName + "_view" + plot = Plotter(args) + + if "a" in args.showPlots: + plot.Dale_matrix_and_eigen(trueD["W_true"], trueMD, trueD, figId=1) + + if "b" in args.showPlots: + plot.histo_weights_rates(trueD, spikeD, trueMD, figId=2) + + if "c" in args.showPlots: + plot.rates_study(trueD, spikeD, trueMD, figId=3) + + if "d" in args.showPlots: + plot.plot_placement_topology(trueD, trueMD, figId=4) + + if "e" in args.showPlots: + plot.connection_length_kernel_hist(trueD, trueMD, figId=5) + + if "f" in args.showPlots: + plot.Dale_matrix_pseudospectra(trueD["W_true"], trueMD, trueD, figId=6) + + if "g" in args.showPlots: + plot.topo_overview(trueD, trueMD, figId=7) + + + plot.display_all() + print("M:done - view_daleMatrix5 completed successfully!") diff --git a/causal_net/nonStation_ver5_synapDepres/view_spikesTrain5.py b/causal_net/nonStation_ver5_synapDepres/view_spikesTrain5.py new file mode 100755 index 00000000..00d8bbc8 --- /dev/null +++ b/causal_net/nonStation_ver5_synapDepres/view_spikesTrain5.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +""" +Visualization tool for biological experiment input features and data quality. + +This script provides comprehensive visualization and analysis of experimental +neural data, focusing on data quality assessment and feature exploration. +Main functionality includes: +- Time series visualization of neural recordings with customizable time ranges +- Data quality metrics and statistical summaries +- Interactive plotting with configurable display options + +Used primarily for exploratory data analysis of biological neural recordings +before further processing and connectivity analysis. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os + +from pprint import pprint +import numpy as np +from PlotterSpikesTrain import Plotter +from toolbox.Util_NumpyIO import read_data_npz +import argparse +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize spike-train data from simulated Dale network outputs") + parser.add_argument("-v","--verbosity",type=int, help="increase output verbosity", default=1, dest='verb') + parser.add_argument("-p", "--showPlots", default='a', nargs='+',help="abcd-string listing shown plots") + + + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + + parser.add_argument("--basePath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for input data") + parser.add_argument("--dataName", default=None,help='simulated Dale network base name') + + + parser.add_argument('-T','--time_range_sec' , default=[0., 50], nargs=2, type=float, help='display data time range in seconds') + parser.add_argument('-r','--time_rebin2', default=10, type=int, help='rebin current time axis') + + args = parser.parse_args() + # make arguments more flexible + args.inpPath = os.path.join(args.basePath, 'truthDale') + + # args.inpPath = os.path.join(args.basePath, 'spikesData') + args.outPath = os.path.join(args.basePath,'plots') + args.showPlots=''.join(args.showPlots) + + print( 'myArg-program:',parser.prog) + for arg in vars(args): print( 'myArg:',arg, getattr(args, arg)) + + if args.time_range_sec!=None: assert args.time_range_sec[0] < args.time_range_sec[1] + assert os.path.exists(args.basePath) + assert os.path.exists(args.inpPath), f"missing inpPath: {args.inpPath}" + return args + + +#...!...!.................... +def rebin_spike_rates(spikeYield, md, tReb2): + """Rebin 2D spike train (time x neurons) and return rate summaries over rebinned time.""" + assert spikeYield.ndim == 2, f"Expected 2D spike train, got shape={spikeYield.shape}" + assert tReb2<101 # this would exceed 1 seconds + + #.... rebin time axis + ntime,nchan=spikeYield.shape + if ntime % tReb2 != 0: + ntime_c = ntime - (ntime % tReb2) + spikeYield = spikeYield[:ntime_c] # Clip the data + #print('ccl',ntime_c , ntime ,ntime % tReb2, tReb2) + #... sum yileds over tReb2 time bins + spikeYieldR= np.sum( spikeYield.reshape(-1, tReb2, nchan),axis=1) + time_step=md['time_step_sec'] + time_step2=time_step*tReb2 + #print('rr1',time_step2,spikeYieldR.shape,spikeYield.shape) + + rate2D=spikeYieldR/time_step2 + pop_spike_count = np.sum(spikeYieldR, axis=1) + pop_rate_hz = pop_spike_count / time_step2 + + rebD={'time_step2':time_step2} + rebD['rate2D']=rate2D + rebD['pop_spike_count'] = pop_spike_count + rebD['pop_rate_hz'] = pop_rate_hz + ntime//=tReb2 + rebD['timeV']= np.linspace(0, (ntime - 1) * time_step2, ntime) + + print('rebinned rates: nchan=%d dt=%.3f sec rebin=%d' % (nchan, time_step2, tReb2)) + return rebD + + +#...!...!.................... +def rebin_network_state(state2D, md, tReb2, state_name='state'): + """Rebin a continuous network state (time x neurons) by averaging over time.""" + state2D = np.asarray(state2D) + assert state2D.ndim == 2, f"Expected 2D {state_name}, got shape={state2D.shape}" + assert tReb2 < 101 # keep consistent with spike-rate rebinning guard + + ntime, nchan = state2D.shape + if ntime % tReb2 != 0: + ntime_c = ntime - (ntime % tReb2) + state2D = state2D[:ntime_c] + ntime = ntime_c + + state2DR = np.mean(state2D.reshape(-1, tReb2, nchan), axis=1) + time_step = md['time_step_sec'] + time_step2 = time_step * tReb2 + ntime2 = state2DR.shape[0] + + rebD = { + 'time_step2': time_step2, + 'state2D': state2DR, + 'timeV': np.linspace(0, (ntime2 - 1) * time_step2, ntime2), + 'state_name': state_name, + } + + print('rebinned %s: nchan=%d dt=%.3f sec rebin=%d' % (state_name, nchan, time_step2, tReb2)) + return rebD + + + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__=="__main__": + args=get_parser() + np.set_printoptions(precision=3) + + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb>0) + if args.verb>1: pprint(spikeMD) + + if 1: # per state information + # Load simulation truth data (Dale matrices, biases, etc.) + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb>0) + if args.verb>1: + print("\nSimulation Truth Metadata:"); pprint(trueMD) + + dataYield = np.asarray(spikeD["spikes"]) + if dataYield.ndim != 2: + raise ValueError( + "spikes must have shape (T, N) from gen_daleMatrices4; got %s. Multi-state (M, T, N) is not supported." + % (dataYield.shape,) + ) + sr = np.asarray(spikeD["single_rates"]) + if sr.ndim != 1: + raise ValueError("single_rates must be (N,); got shape %s" % (sr.shape,)) + if sr.shape[0] != dataYield.shape[1]: + raise ValueError("single_rates length must match spikes.shape[1]") + spikeD["spikes"] = dataYield + spikeD["single_rates"] = sr + spikeMD["sel_spect_radius"] = trueMD["dale_conf"]["spectral_radius"] + + + #-------------------------------- + # .... plotting ........ + spikeMD["short_name"] = args.dataName + args.prjName = spikeMD["short_name"] + spikeMD['plot']={} + spikeMD['plot']['time_rangeLR']=np.array(args.time_range_sec) + + plot=Plotter(args) + + if 'a' in args.showPlots: + plot.freq_histo(spikeD,spikeMD,figId=1) + if 'b' in args.showPlots: + rebD=rebin_spike_rates(spikeD['spikes'], spikeMD, args.time_rebin2) + plot.freq_vs_time(rebD,spikeMD,figId=2) + if 'c' in args.showPlots: + rebX = rebin_network_state(trueD['x_true'], spikeMD, args.time_rebin2, state_name='x') + plot.state_x_vs_time(rebX, spikeMD, figId=3) + + plot.display_all() + print('M:done') + #pprint(expMD) #tmp diff --git a/causal_net/quantum/doc/explore_Amatrix.pdf b/causal_net/quantum/doc/explore_Amatrix.pdf new file mode 100644 index 00000000..7c88bce8 Binary files /dev/null and b/causal_net/quantum/doc/explore_Amatrix.pdf differ diff --git a/causal_net/quantum/doc/explore_Amatrix.tex b/causal_net/quantum/doc/explore_Amatrix.tex new file mode 100644 index 00000000..9df7bfab --- /dev/null +++ b/causal_net/quantum/doc/explore_Amatrix.tex @@ -0,0 +1,47 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} + +\geometry{margin=0.8in} + +\title{Explore Quantum walk in Neuronal Network Modeling} +\author{Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\section{Block Matrix Formulation of Time Evolution} + +Let the original time evolution be given by $\eta_t = A Y_{t-1} + B$, where the elements of $B$ are bounded in $[0, b_{\max}]$ and the elements of $Y$ are binary ($0$ or $1$). We construct the extended block matrices $A'$ and $B'$ as: +\begin{equation} + A' = \begin{bmatrix} 0 & A \\ A^T & 0 \end{bmatrix}, \quad B' = \begin{bmatrix} B \\ 0 \end{bmatrix} +\end{equation} + +To properly reconstruct $\eta_t$ in the upper block, we must define $Y'_{t-1}$ with the variables in the lower block: +\begin{equation} + Y'_{t-1} = \begin{bmatrix} 0 \\ Y_{t-1} \end{bmatrix} +\end{equation} + +The extended time evolution of $\eta'_t$ then evaluates correctly: +\begin{align} + \eta'_t &= A' Y'_{t-1} + B' \nonumber \\ + &= \begin{bmatrix} 0 & A \\ A^T & 0 \end{bmatrix} \begin{bmatrix} 0 \\ Y_{t-1} \end{bmatrix} + \begin{bmatrix} B \\ 0 \end{bmatrix} \nonumber \\ + &= \begin{bmatrix} A Y_{t-1} \\ 0 \end{bmatrix} + \begin{bmatrix} B \\ 0 \end{bmatrix} \nonumber \\ + &= \begin{bmatrix} A Y_{t-1} + B \\ 0 \end{bmatrix} +\end{align} + +Substituting the original equation into the upper block, we arrive at the final compact form: +\begin{equation} + \eta'_t = \begin{bmatrix} \eta_t \\ 0 \end{bmatrix} +\end{equation} + +Finally we measure the upper half of qubits to retrieve $\eta_t$. Assume $b_{max}=0.3$ and draw B,Y randomly. +\end{document} diff --git a/causal_net/stationaryLag1_ver6_lasso/PlotterBioExp.py b/causal_net/stationaryLag1_ver6_lasso/PlotterBioExp.py new file mode 100644 index 00000000..c8ff9d9e --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/PlotterBioExp.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for biological experiment data visualization. + +This module provides specialized plotting capabilities for experimental +neural data analysis. The Plotter class extends PlotterBackbone to create +visualizations tailored for biological neural recordings including: +- Time series plots of neural activity patterns +- Statistical summaries and distribution analysis +- Data quality assessment plots +- Comparative analysis between experimental conditions + +Designed specifically for processing and visualizing data from biological +neural experiments, with automatic adaptation to experimental metadata. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec + +from matplotlib.colors import LinearSegmentedColormap + +#...!...!.................... +def summary_column(md): + #print(sorted(md)) + rs=md['rate_summary'] + ds=md['data_selector'] + #pprint(md) + txt='dataset: '+md['short_name'] + txt += '\nnum acc neurons: %d' % md['num_neurons'] + txt += '\ndrop neur: %d <%.1f Hz, %d >%.1f Hz' % (ds['drop_neur_by_freq_range'][0],ds['freq_range'][0],ds['drop_neur_by_freq_range'][1],ds['freq_range'][1]) + txt+='\ndata type: %s\ntime_step=%.2f sec'%(md['data_type'],md['time_step_sec']) + txt += '\nMedian rate: %.1f Hz' % rs['median_spike_rate'] + txt += '\nrate range [%.1f, %.1f Hz]' % (rs['min_spike_rate'], rs['max_spike_rate']) + txt += '\nAvg Rate: %.1f±%.1f Hz' % (rs['avg_spike_rate'], rs['std_spike_rate']) + txt += '\nAvg Fano: %.2f±%.2f' % (rs['avg_fano_factor'], rs['std_fano_factor']) + + return txt + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def freq_histo(self,spikeD,md,figId=1): + + tit='dataset '+md['short_name'] + + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,4)) + + # Create gridspec with 7:3 width ratio + gs = gridspec.GridSpec(1, 2, width_ratios=[7, 3]) + + dataYield, dataRates = spikeD['spikes'], spikeD['single_rates'] + + # Compute median + median_val = np.median(dataRates) + txtM= f'median : {median_val:.2f} Hz' + + #.... freq per channel + ax = self.plt.subplot(gs[0, 0]) + chanV=np.arange(dataRates.shape[0]) + ax.bar(chanV, dataRates , width=0.8, color='g', align='center', alpha=0.7) + + ax.axhline(median_val, color='red', linestyle='--', linewidth=1) + if md['data_type']=='simDale': + ax.set_yscale('log') + ax.grid() + ax.set(xlabel='input neuron index', ylabel='avr frequency (Hz)',title=tit) + ax.text( ax.get_xlim()[1]*0.1,median_val,txtM, + color='red', ha='center', va='bottom', fontsize=10) + + txt=summary_column(md) + ax.text(0.05, 0.3, txt, transform=ax.transAxes, fontsize=12,color='b') + + #........ freq histo ..... + ax = self.plt.subplot(gs[0, 1]) + ax.hist(dataRates,bins=20) + #ax.set_yscale('log') + ax.set(ylabel='num channels', xlabel='avr frequency (Hz)',title=tit) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + # Add median text annotation + ax.text(median_val+1, ax.get_ylim()[1]*0.5,txtM, + color='red', ha='left', va='bottom', fontsize=10,rotation=90) + ax.text(0.25, 0.3, txt, transform=ax.transAxes, fontsize=11,color='m') + + ax.grid() + +#...!...!.................. + def freq_vs_time(self,rebD,md,figId=2): + figId=self.smart_append(figId) + fig=self.plt.figure(figId,facecolor='white', figsize=(16,11)) + + tit='dataset '+md['short_name'] + time_step=rebD['time_step2'] + rateThr2=rebD['rate_thres2'] + highCntThr=rebD['high_cnt_thres'] + + # clip time data for display + tL,tR=md['plot']['time_rangeLR'] + itL,itR=(md['plot']['time_rangeLR']/time_step).astype(int) + print('iTL,R', itL,itR) + + # unpack data and clip the time range + rate2D=rebD['rate2D'][itL:itR] + highChanCnt=rebD['highChanCnt'][itL:itR] + highChanSm=rebD['highChanSmooth'][itL:itR] + highChanMask=rebD['highChanMask'][itL:itR] + + rate1D=rebD['rate1D'][itL:itR] + timeV=rebD['timeV'][itL:itR] + + _,nchan=rate2D.shape + Tbin = (timeV[1] - timeV[0]) + + tit0='dataset: %s nchan=%d Tbin=%.2f sec'%(md['short_name'],nchan,time_step) + + # Create gridspec with 2 rows in a 0.6:0.4 ratio + gs = fig.add_gridspec(5, 1, height_ratios=[0.15,0.15,0.15, 0.54,0.01]) + + # ..... top plot + ax = fig.add_subplot(gs[1, 0]) + ax.bar(timeV+Tbin*.5, highChanSm, width=Tbin, color='slateblue', align='center', alpha=0.7) + ax.bar(timeV+Tbin*.5, highChanSm*highChanMask , width=Tbin, color='red', align='center', alpha=0.7) + tit=f'num neurons w/ smooth kernel {rebD["smooth_kernel"]} , rate thres> {rateThr2:.0f} (Hz), usable time frac:{rebD["usable_time_fract"]:.3f} nchan={nchan}' + ax.set(ylabel='num neurons', title=tit) + ax.axhline( highCntThr,c='g',lw=1,ls='--') + ax.set_xlim(tL,tR) + ax.grid() + + # ..... middle1 plot + ax = fig.add_subplot(gs[2, 0]) + + ax.bar(timeV+Tbin*.5, highChanCnt, width=Tbin, color='forestgreen', align='center', alpha=0.7) + + ax.bar(timeV+Tbin*.5, highChanCnt*highChanMask , width=Tbin, color='red', align='center', alpha=0.7) + ax.set(xlabel='Time (s)', ylabel='num neurons', title=f'num neurons with instantanous rate > thres= {rateThr2:.0f} (Hz)') + ax.set_xlim(tL,tR) + ax.grid() + + # ..... middle2 plot + ax = fig.add_subplot(gs[0, 0]) + ax.bar(timeV+Tbin*.5, rate1D, width=Tbin, color='orange', align='center', alpha=0.7) + ax.bar(timeV+Tbin*.5, rate1D*highChanMask , width=Tbin, color='red', align='center', alpha=0.7) + ax.set( ylabel='sum rate (Hz)', title=f'sum rate from all {nchan} neurons') + ax.set_xlim(tL,tR) + ax.grid() + + # ....... bottom row: 2D histogram + ax = fig.add_subplot(gs[3, 0]) + #cmap='tab20c', 'Oranges' + # - - - Get the colormap and modify it --- + original_cmap = self.plt.get_cmap('tab20c') + custom_cmap = original_cmap.copy() + custom_cmap.set_under('white') + + cax = ax.imshow(rate2D.T, + aspect='auto', + cmap=custom_cmap, # Use our modified colormap + origin='lower', + extent=[tL, tR, 0, nchan], + interpolation='nearest', + vmin=1, # Set the lower bound of the colormap + vmax=rate2D.max()) # Optional: ensure the upper bound is set + + ax.set_ylabel('Neuron index') + ax.set_xlabel('Time (sec)') + ax.set_title(tit0 ) + cbar = fig.colorbar(cax, ax=ax, orientation='horizontal', pad=0.17, label=f'Instantanous freq (Hz) per neuron over Tbin={time_step:.1f} sec', shrink=0.5) + cbar.ax.tick_params(labelsize=10) + ax.grid() + + diff --git a/causal_net/stationaryLag1_ver6_lasso/PlotterFitEval.py b/causal_net/stationaryLag1_ver6_lasso/PlotterFitEval.py new file mode 100644 index 00000000..e76b7b41 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/PlotterFitEval.py @@ -0,0 +1,765 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for LASSO fit evaluation and visualization. + +This module provides comprehensive plotting capabilities for analyzing +LASSO Poisson model fitting results. The Plotter class extends PlotterBackbone +to create specialized visualizations including: +- Connectivity matrix heatmaps with frequency-sorted neurons +- Weight distribution histograms and statistical summaries +- Network structure plots showing excitatory/inhibitory connections +- Reconstruction quality plots comparing fitted vs observed rates +- Comparative analysis plots for simulation validation + +The plots support both simulated and experimental data with automatic +adaptation based on available metadata. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec +import matplotlib.colors as colors +from matplotlib.colors import TwoSlopeNorm +from scipy.stats import gennorm +from matplotlib.colors import LogNorm + + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def summary_fitLasso(self,fitD, md,byFreq=False, figId=1): # p=a + figId=self.smart_append(figId) + nrow,ncol=2,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(19,6)) + + #pprint(md) + fmd=md['fit_lasso'] + esmd=md['edge_selector'] + Nn=fmd['num_neurons'] + + # Create diagonal mask + diagM = np.eye(Nn, dtype=bool) + #!off_diagM = ~diagM + + fitType=md['fit_type'] + eselType=esmd['selector_type'] + isExp = md.get('data_type') == 'bioExp' # Automatically detect experimental data + # Unpack arrays from bigD + A_fit = fitD['A_'+fitType] + B = fitD['B_'+fitType] + Freq=fitD['single_rates'] + + + #...... Training curves + ax = self.plt.subplot(nrow,ncol,1) + plot_trainingCurves(ax,fitD,md) + + #.... : histogram of rates + ax = self.plt.subplot(nrow,ncol,1+ncol) + ax.hist(Freq, bins=20) + x_vals = np.arange(Nn) + frLab='Firing rate (Hz)' + ax.set(ylabel='num neurons',xlabel=frLab,title='Single rates, %d neurons'%Nn) + ax.grid(True, alpha=0.4) + ax.set_xlim(0,) + # Compute median + median_val = np.median(Freq) + txtM= f'median rate: {median_val:.2f} Hz' + ax.text( median_val,ax.get_ylim()[1]*0.8,txtM, + color='red', ha='left', va='bottom', fontsize=10) + ax.axvline(median_val, color='red', linestyle='--', linewidth=1) + + # ... A off-diagonal + xLab='edge value' + frLab='log10( Firing rate/ Hz )' + Aedg=A_fit[~diagM] + i_indices, j_indices = np.where(~diagM) + Freq_expanded = Freq[i_indices] + # Filter out 0 values + valid_mask =np.abs(Aedg)>1e-10 + Aedg_clean = Aedg[valid_mask] + Freq_expanded_clean = Freq_expanded[valid_mask] + # Count non-zero edges (non-NaN values) + n_edges = len(Aedg_clean) + + ax = self.plt.subplot(nrow,ncol,2) + ax.hist(Aedg_clean, bins=100,color='g',alpha=0.8) + ax.set_yscale('log') + ax.grid(True, alpha=0.4) + ax.set(ylabel='edges',xlabel=xLab,title='A off-diagonal') + txt='edge sel meth: %s \n acc frac=%.2f' %( eselType,n_edges/Nn/(Nn-1)) + if eselType=='FDR': txt+='\n alpha=%.3f '%(esmd['alpha']) + ax.text(0.05, 0.75, txt, transform=ax.transAxes, fontsize=10) + # Plot 2D histogram + ax = self.plt.subplot(nrow,ncol,2+ncol) + h = ax.hist2d(Aedg_clean, np.log10(Freq_expanded_clean), bins=30, cmap='viridis', norm=LogNorm()) + cbar = fig.colorbar(h[3],ax=ax) + ax.set(ylabel=frLab,xlabel=xLab,title='accept %d of %d edges '%(n_edges,Nn*(Nn-1))) + ax.grid(True, alpha=0.4) + + + # ... A diagonal ..... + xLab='A-diag value' + Adia=A_fit[diagM] + i_indices, j_indices = np.where(diagM) + Freq_expanded = Freq[i_indices] + + ax = self.plt.subplot(nrow,ncol,3) + ax.hist(Adia, bins=30,color='salmon') + ax.grid(True, alpha=0.4) + ax.set(ylabel='neurons',xlabel=xLab,title='A diagonal') + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + # Plot 2D histogram + ax = self.plt.subplot(nrow,ncol,3+ncol) + h = ax.hist2d(Adia,np.log10(Freq_expanded), bins=30, cmap='Greys',vmax=2.1) + cbar = fig.colorbar(h[3],ax=ax) + ax.set(ylabel=frLab,xlabel=xLab,title='num neurons') + ax.grid(True, alpha=0.4) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + # ... B-term ..... + xLab='B-term value' + ax = self.plt.subplot(nrow,ncol,4) + ax.hist(B, bins=30,color='darkviolet') + ax.grid(True, alpha=0.4) + ax.set(ylabel='neurons',xlabel=xLab,title='B-term') + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + # Plot 2D histogram + ax = self.plt.subplot(nrow,ncol,4+ncol) + h = ax.hist2d(B,np.log10(Freq), bins=30, cmap='Grays',vmax=2.1) + cbar = fig.colorbar(h[3],ax=ax) + ax.set(ylabel=frLab,xlabel=xLab,title='num neurons') + ax.grid(True, alpha=0.4) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + #...!...!.................. + def residuals(self, evalD,md, figId=1): + #pprint(md) + fitType=md['fit_type'] + fmd=md['fit_'+fitType] + esmd=md['edge_selector'] + + figId=self.smart_append(figId) + nrow,ncol=2,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(12,6)) + + # Unpack arrays from bigD + for j,etype in enumerate(['neg','pos']): + ax = self.plt.subplot(nrow,ncol,1+j) + + title = 'edges: %s' % etype + valT,valF=plot_correl_offdiag(fig,ax,evalD[etype]) + ax.set_title(title) + + txt='alpha=%.3f '%(esmd['alpha']) + ax.text(0.40, 0.15, txt, transform=ax.transAxes, fontsize=10) + + ax = self.plt.subplot(nrow,ncol,1+j+ncol) + plot_1D_residuals(ax, valT,valF,lab='TP off-diag '+etype,col='green') + + if etype=='pos': + txt=md["short_name"] + else: + txt='fit: '+fitType + ax.text(0.05, 0.2, txt, transform=ax.transAxes, fontsize=8) + + # ... diagonal + title = 'fit (diagonal)' + ax = self.plt.subplot(nrow,ncol,3) + V=evalD['diag'] + dCol='salmon' + ax.scatter(V[:,1], V[:,0], alpha=0.6, color=dCol,marker='.',s=5) + ax.set(title=title,xlabel='truth',ylabel='fitted') + add_x45_lins(ax, only45=True) + ax.grid(True, alpha=0.5) + + ax = self.plt.subplot(nrow,ncol,3+ncol) + plot_1D_residuals(ax,V[:,1],V[:,0],lab='diag',col=dCol) + + + # ... B-term + title = 'fit (B-term)' + ax = self.plt.subplot(nrow,ncol,4) + dCol='darkviolet' + V=evalD['bterm'] + ax.scatter(V[:,1], V[:,0], alpha=0.6, color=dCol,marker='.',s=5) + ax.set(title=title,xlabel='truth',ylabel='fitted') + add_x45_lins(ax, only45=True) + ax.grid(True, alpha=0.5) + + ax = self.plt.subplot(nrow,ncol,4+ncol) + plot_1D_residuals(ax,V[:,1],V[:,0],lab='diag',col=dCol) + + +#...!...!.................. + def freqSortA_histos(self, fitD, md, spikeD, figId=1, k=6): + #pprint(md); aa67 + fitType=md['fit_type'] + + fmd=md['fit_'+fitType] + + #1isExp = md.get('data_type') == 'bioExp' # Automatically detect experimental data + figId=self.smart_append(figId) + nrow,ncol=k,2 # Add 1 extra row for the neuron stats plot + fig=self.plt.figure(figId,facecolor='white', figsize=(16,12)) + + # Unpack arrays from bigD + A_fit = fitD['A_'+fitType].copy() + Freq=fitD['single_rates'] + + # Multiply each row of A_fit by the corresponding element of F + A_fit=A_fit * np.sqrt(Freq[:,None]) + + # convert 0's to Nan + A_fit[ A_fit==0.]=np.nan + + num_neurons = A_fit.shape[0] + wzoomMx=0.05 + # Remove diagonal elements by setting them to NaN + A_fit_no_diag = A_fit.copy() + np.fill_diagonal(A_fit_no_diag, np.nan) + A_flat = A_fit_no_diag.flatten() + xLab=r' weights * $\sqrt{ rate}$' + + #------- top left plot w/ 1D histo fo weights + ax = self.plt.subplot(nrow,ncol,1) + ax.hist(A_flat, bins=200, color='green', alpha=0.7, edgecolor=None) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + ax.grid() + method_name = md.get('edge_selection_method', 'unknown') + ax.set(xlabel=xLab,ylabel='count',title=f'All non-diag weights, Fit {fitType}, Method: {method_name}') + ax.set_yscale('log') + + + # Left column: 2D histogram of A-matrix (mutiple rows) + ax = self.plt.subplot2grid((nrow, ncol), (1, 0), rowspan=5) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + # Create 2D histogram: x-axis is value, y-axis is row index + row_indices = np.repeat(np.arange(num_neurons), num_neurons) + + # Remove NaN values (diagonal elements) + valid_mask = ~np.isnan(A_flat) + A_flat = A_flat[valid_mask] + row_indices = row_indices[valid_mask] + + # Use ax.hist2d() directly with log scale + H, xedges, yedges, im = ax.hist2d(A_flat, row_indices, bins=[50, num_neurons], cmap='Greys', vmax=6) + + cbar = fig.colorbar(im, ax=ax, + orientation='horizontal', # Make it horizontal + location='bottom', # Place it below + shrink=0.7, # Make it 70% width + pad=0.1) # Add some padding from plot + cbar.set_label('num edges') + + ax.set_xlabel(xLab) + ax.set_ylabel('freq-sorted neuron index') + + lasso_name = md['fit_lasso']['lassoFit_output_name'] + title_text = f'{lasso_name}, Fit {fitType}, epochs={fmd["num_epochs"]}, sampl/k={fmd["num_samples_used"]/1000}' + + ax.set_title(title_text) + + # Compute row indices and collect data for global binning + single_rates = spikeD['single_rates'] + idxOff = (num_neurons // k) // 2 + 2 + row_indices = [min(i * (num_neurons // k) + idxOff, num_neurons - 1) for i in range(k)] + selected_rows_data = [A_fit_no_diag[idx, :][~np.isnan(A_fit_no_diag[idx, :])] for idx in row_indices] + all_valid_data = np.concatenate([data for data in selected_rows_data if len(data) > 0]) + global_bin_edges = np.linspace(np.min(all_valid_data), np.max(all_valid_data), 51) if len(all_valid_data) > 0 else None + + # Add frequency annotations and plot histograms in single loop + xlim_offset = 0.02 * (ax.get_xlim()[1] - ax.get_xlim()[0]) + for i in range(k): + rowIdx = row_indices[i] + ax.axhline(rowIdx, color='yellow', linewidth=2, alpha=0.8) + ax.text(ax.get_xlim()[0] + xlim_offset, rowIdx, f'{single_rates[rowIdx]:.1f} Hz', verticalalignment='center', horizontalalignment='left', fontsize=12) + ax_hist = self.plt.subplot(nrow, ncol, 2 + i * ncol) + reversed_rowIdx = row_indices[k - 1 - i] + #xLab1 = xLab if (i == k - 1) else None + plot_row_histogram(ax_hist, A_fit_no_diag, reversed_rowIdx, single_rates, global_bin_edges, xLab, isExp=True) + + # Minimize whitespace between 1D plots + self.plt.subplots_adjust(hspace=0.05, wspace=0.1) + +#...!...!.................. + def summary_network(self, fitD, edgeD, md, procFrac=0.8,figId=5): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(16,3.5)) + + fitType=md['fit_type'] + + # Unpack arrays from bigD + A_fit = fitD['A_'+fitType] + Neu_sum=edgeD['edge_sum'] + Neu_edg=edgeD['edge_vals'] + + # .... A-matrix .... + ax = self.plt.subplot(nrow,ncol,1) + plot_A2D(fig,ax,A_fit) + + # .... A-eigen .... + ax = self.plt.subplot(nrow,ncol,2) + ax.set_title(f'{fitType}: {md["short_name"]}') + eigF=np.linalg.eigvals(A_fit) + reF = np.real(eigF) + imF = np.imag(eigF) + + ax.scatter(reF, imF, color='red', marker='o', facecolors='none',label='fit',s=20) + ax.set_ylim(-0.1,) + #1ax.set_xlim(right=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.grid(True, alpha=0.3) + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + ax.legend() + + # .... pos vs. neg count .... + ax = self.plt.subplot(nrow,ncol,3) + # Extract positive and negative counts + n_pos = Neu_sum[:, 1] + n_neg = Neu_sum[:, 2] + # Create 2D histogram + h = ax.hist2d(n_pos, n_neg, bins=20, cmap='Greys', cmin=1,cmax=5) + # Add colorbar + self.plt.colorbar(h[3], ax=ax, label='Neurons') + ax.set_xlim(0,) + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + + # Add diagonal line (y=x) + lims = [ + np.min([ax.get_xlim(), ax.get_ylim()]), + np.max([ax.get_xlim(), ax.get_ylim()]), + ] + ax.plot(lims, lims, 'b--', alpha=0.5, linewidth=1.5, label='y=x') + ax.set(title='Edge type correlation, %d neurons'%(A_fit.shape[0]),xlabel='num pos',ylabel='num neg') + + # .... edge std vs. value .... + ax = self.plt.subplot(nrow,ncol,4) + ax.grid(True, alpha=0.3) + ax.set(title='Edge value accuracy',xlabel='edge val',ylabel='edge std') + + # Extract average weights and standard deviations + avg_weights = Neu_edg[:, 2] + std_devs = Neu_edg[:, 3] + + # Compute percentile range to contain procFrac of data by magnitude + abs_weights = np.abs(avg_weights) + percentile_cutoff = procFrac * 100 + threshold = np.percentile(abs_weights, percentile_cutoff) + + # Filter data within range + mask = abs_weights <= threshold + avg_weights_accepted = avg_weights[mask] + std_devs_accepted = std_devs[mask] + + # Count accepted and rejected + n_total = len(avg_weights) + n_accepted = len(avg_weights_accepted) + n_rejected = n_total - n_accepted + + # Create 2D histogram + h = ax.hist2d(avg_weights_accepted, std_devs_accepted, bins=30, cmap='Greys', cmin=1) + + # Add colorbar + self.plt.colorbar(h[3], ax=ax, label='edges') + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + # Add text box with statistics + stats_text = (f'display fract = {procFrac:.2f}\n' + f'Accepted: {n_accepted:,} ({n_accepted/n_total*100:.1f}%)\n' + f'Rejected: {n_rejected:,} ({n_rejected/n_total*100:.1f}%)\n' + f'abs(x) cutoff: ±{threshold:.3f}') + + ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, + fontsize=10, verticalalignment='top') + + + +#............................ +#............................ +#............................ + +def plot_trainingCurves(ax,fitD,md,title='aa3'): + fitType=md['fit_type'] + fmd=md['fit_'+fitType] + + epoch0=10 # skip intial loss values + lossTot = fitD['losses_total'][epoch0:] + epochsT=fitD['losses_epochs'][epoch0:] + + # Plot total loss on left y-axis + ax.plot(epochsT, lossTot, label='total '+fitType, color='blue', linestyle='-') + + titl='%dk samp'%(fmd["num_samples_used"]/1000) + + if fitType=='lasso': # Create second y-axis for L1 loss + lossL1= lossTot - fitD['losses_wo_L1'][epoch0:] + ax2 = ax.twinx() + ax2.plot(epochsT, lossL1, label='L1 loss', color='red', linestyle='--') + ax2.set_ylabel('L1 loss', color='red') + ax2.tick_params(axis='y', labelcolor='red') + # Move y-axis ticks and labels inside + ax2.tick_params(axis='y', direction='in', pad=-60, labelsize=10) + ax2.yaxis.set_label_position('right') + + # Format second y-axis ticks in scientific notation + from matplotlib.ticker import FuncFormatter + ax2.yaxis.set_major_formatter(FuncFormatter(lambda x, p: f'{x:.1e}')) + # Combine legends from both axes + lines1, labels1 = ax.get_legend_handles_labels() + lines2, labels2 = ax2.get_legend_handles_labels() + ax.legend(lines1 + lines2, labels1 + labels2, title=titl) + else: + ax.legend( title=titl) + + title = 'Loss: '+md["short_name"] + + ax.set(xlabel='Epoch', title=title) + ax.set_ylabel('Loss', color='blue') + + # Color the y-axis labels to match the lines + ax.tick_params(axis='y', labelcolor='blue') + ax.grid(True, alpha=0.3) + +def plot_A2D(fig,ax,A,title='aa',byFreq=True,trueD=None): + Am=A.copy() + + + # If byFreq=False, reorder matrix to natural neuron indexing + if not byFreq and trueD is not None and 'neur_revFreqIdx' in trueD: + neur_revFreqIdx = trueD['neur_revFreqIdx'] # freq_sorted_position → natural_index + # Reorder both rows and columns from frequency-sorted to natural order + Am = Am[np.ix_(neur_revFreqIdx, neur_revFreqIdx)] + mask = mask[np.ix_(neur_revFreqIdx, neur_revFreqIdx)] + Am[~mask]=0 # Re-apply mask after reordering + + vmin = Am.min()-0.3 + vmax = Am.max()+0.3 + + norm = TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax) + im1=ax.imshow(Am, cmap='bwr', norm=norm, origin='lower') # 'RdBu_r' + + #masked_values = Am[mask] + #vmin = masked_values.min() + #vmax = masked_values.max() + + nval=np.sum(Am!=0.) + title='%s n=%d'%(title,nval) + # Simplified axis labels based on display order + if byFreq: + ylabel_text = 'From neuron (freq-sorted)' + xlabel_text = 'To neuron (freq-sorted)' + else: + ylabel_text = 'From neuron (natural idx)' + xlabel_text = 'To neuron (natural idx)' + ax.set(title=title, ylabel=ylabel_text, xlabel=xlabel_text) + fig.colorbar(im1, ax=ax, shrink=0.7) + ax.grid(True, alpha=0.5) + #add_x45_lins(ax, only45=True) + + +def add_x45_lins(ax, only45=False): + # Add dashed lines through (0,0) + lims = [ + np.min([ax.get_xlim(), ax.get_ylim()]), # min of both axes + np.max([ax.get_xlim(), ax.get_ylim()]), # max of both axes + ] + # 45 degree line y=x + ax.plot(lims, lims, '--', color='k', linewidth=0.8) + if only45: return + + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + + +def plot_correl_offdiag(fig,ax,tripV): + TP,FP,FN=tripV + #print('ss',TP.shape) + ax.scatter(TP[:,3], TP[:,2], alpha=0.6, color='green',label='TP: %d'%TP.shape[0],marker='.',s=5) + n=FN.shape[0] + ax.scatter(FN[:,2], [0]*n, alpha=0.6, color='red',s=5,label='FN: %d'%FN.shape[0]) + + n=FP.shape[0] + ax.scatter([0]*n,FP[:,2], alpha=0.6, color='blue',s=5,label='FP: %d'%FP.shape[0]) + + ax.set(aspect=1. ,xlabel='true weight',ylabel='fitted') + ax.grid(True, alpha=0.5) + + add_x45_lins(ax) + if np.any(TP): # Only plot if there are TP points + x_cg = np.mean(TP[:,3]) + y_cg = np.mean(TP[:,2]) + ax.scatter(x_cg, y_cg, marker='+', s=200, color='k', linewidths=2) #, label='TP avr') + ax.legend() + + return TP[:,3], TP[:,2] + + + +#...!...!.................. +def plot_1d_weight_histo_with_stats(ax, group_values, start_row, end_row, group_idx, color, ggd_rangeL, isExp=False): + """Plot 1D histogram with statistics for a single group""" + + ggd_range=(ggd_rangeL[0]+ggd_rangeL[1])/2. # TMP + # Create histogram for this group with log scale + counts, bins, _ = ax.hist(group_values, bins=100, alpha=0.7, color=color) + + ax.axvline(x=ggd_range, color='black', linestyle='--', linewidth=1, alpha=0.8) + ax.axvline(x=-ggd_range, color='black', linestyle='--', linewidth=1, alpha=0.8) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + # Compute mean and RMSE in the range ±ggd_range, excluding values too close to zero + mask_range = (np.abs(group_values) < ggd_range) & (np.abs(group_values) > ggd_range/5) + values_in_range = group_values[mask_range] + + if len(values_in_range) > 0: + mean_val = np.mean(values_in_range) + rmse_val = np.sqrt(np.mean(values_in_range**2)) + + # Fit Generalized Gaussian Distribution (GGD) + try: + # Fit GGD without any constraints to allow all parameters to be estimated freely + beta, loc, scale = gennorm.fit(values_in_range) + + # Format with scientific notation for very small values + if abs(loc) < 1e-3: + loc_str = f'{loc:.2e}' + else: + loc_str = f'{loc:.4f}' + + if abs(scale) < 1e-3: + scale_str = f'{scale:.2e}' + else: + scale_str = f'{scale:.4f}' + + ggd_text = f'β={beta:.4f}\nx₀={loc_str}\nμ={scale_str}' + + # Plot GGD fit curve + ggd_plot_range=2*ggd_range + x_fit = np.linspace(-ggd_plot_range, ggd_plot_range, 100) + # Scale the PDF to match histogram counts + ggd_pdf = gennorm.pdf(x_fit, beta, loc, scale) + + # Scale to match histogram by finding the maximum bin count in range + bin_centers = (bins[:-1] + bins[1:]) / 2 + mask_fit_bins = (bin_centers >= -ggd_plot_range) & (bin_centers <= ggd_plot_range) + if np.any(mask_fit_bins): + max_count_in_range = np.max(counts[mask_fit_bins]) + max_pdf = np.max(ggd_pdf) + if max_pdf > 0: + scaled_pdf = ggd_pdf * max_count_in_range / max_pdf + ax.plot(x_fit, scaled_pdf, '--', color='magenta', linewidth=2, + label='GGD fit', zorder=8) + + except Exception as e: + print(f"GGD fit failed for group {group_idx+1}: {e}") + beta, loc, scale = np.nan, np.nan, np.nan + ggd_text = f'β=failed\nx₀=failed\nμ=failed' + + # Add text with statistics + ax.text(0.25, 0.85, f'mean={mean_val:.4f}\nRMSE={rmse_val:.4f}\n{ggd_text}', + transform=ax.transAxes, fontsize=10, + bbox=dict(boxstyle='round', facecolor='white', alpha=0.8), + verticalalignment='top') + + # Add circle for mean with horizontal error bar for RMSE at half height + max_count = np.max(counts) + half_height = max_count / 200 + + # Draw horizontal error bar for RMSE + ax.errorbar(mean_val, half_height, xerr=rmse_val, fmt='o', + color='black', capsize=8, capthick=1, linewidth=1, zorder=9) + + ax.set_xlabel('off-diagonal weight value') + ax.set_ylabel('count') + ax.set_title(f'Fitted weights, group {group_idx+1}: neurons {start_row}-{end_row-1} (n={len(group_values)})') + ax.grid(True, alpha=0.3) + + # Limit y-range from 0.5 to max count + max_count = np.max(counts) + ax.set_ylim(0.5, max_count * 1.1) # Add 10% margin above max + + +#...!...!.................. +def add_k_block_lines(ax, num_neurons, k): + """Add horizontal lines to mark K block boundaries""" + rows_per_group = num_neurons // k + for i in range(1, k): # Don't draw line at the very top or bottom + y_line = i * rows_per_group - 0.5 # Position between blocks + ax.axhline(y=y_line, color='black', linestyle='--', linewidth=0.5, alpha=1.0) + +#...!...!.................. +def plot_neuron_stats(ax, A_flat_narrow, row_indices_narrow, num_neurons): + """Plot mean and RMSE for each neuron index""" + + means = [] + rmses = [] + neuron_indices = [] + + # Compute statistics for each neuron + for neuron_idx in range(num_neurons): + # Get values for this specific neuron + mask = row_indices_narrow == neuron_idx + values = A_flat_narrow[mask] + + if len(values) > 0: + mean_val = np.mean(values) + rmse_val = np.sqrt(np.mean(values**2)) + + means.append(mean_val) + rmses.append(rmse_val) + neuron_indices.append(neuron_idx) + + # Convert to numpy arrays + means = np.array(means) + rmses = np.array(rmses) + neuron_indices = np.array(neuron_indices) + + # Plot mean values with RMSE as error bars + ax.errorbar(neuron_indices, means, yerr=rmses, fmt='o', + color='blue', capsize=3, capthick=1, linewidth=1, markersize=3) + + ax.set_xlabel('Neuron freq-sorted index') + ax.set_ylabel('Mean ± RMSE') + ax.set_title('Per-neuron zero-values stats') + ax.grid(True, alpha=0.3) + ax.axhline(y=0, color='black', linestyle='-', linewidth=0.5, alpha=0.5) + + +#...!...!.................. +def plot_1D_residuals(ax, valT,valF,lab,col,first=True): + fac=100 + res=fac*(valT-valF) + + # Plot histogram of residuals + n, bins, patches = ax.hist(res, bins=30, color=col, edgecolor='none', alpha=0.7, density=True ) + + # Compute statistics: mean, standard deviation, and RMSE + mean_val = np.mean(res) + std_val = np.std(res) + n_entries = len(valT) + + # Determine half of the maximum bin height to position the error bar + half_height = np.max(n) / 2.0 + + # Draw horizontal error bar for RMSE at the computed mean and half-height + ax.errorbar(mean_val, half_height, xerr=std_val, fmt='o', + color='black', capsize=8, capthick=1, linewidth=1, zorder=9) + + # Add text to the figure displaying the mean, standard deviation, and RMSE + textstr = f"{lab}\nN = {n_entries}\nMean = {mean_val:.2f}\nStd = {std_val:.2f}" + + if first: + x0=0.95 + else: + x0=0.4 + ax.text(x0, 0.95, textstr, transform=ax.transAxes, fontsize=8, + verticalalignment='top', horizontalalignment='right') + + # Add a black vertical line at x = 0 + ax.axvline(0, color='black', linestyle='--', linewidth=1) + ax.set(xlabel='residuals x %d'%fac, ylabel='density/bin') + + +#...!...!.................. +def plot_both_eigen(ax, eigT,eigF): + reT = np.real(eigT) + imT = np.imag(eigT) + + ax.scatter(reT, imT, color='blue', marker='o',label='True',s=10) + + reF = np.real(eigF) + imF = np.imag(eigF) + + ax.scatter(reF, imF, color='red', marker='o', facecolors='none',label='fit',s=20) + ax.set_ylim(-0.1,) + ax.set_xlim(right=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.grid(True, alpha=0.3) + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + ax.legend() + + +#...!...!.................. +def plot_row_histogram(ax, A_fit_no_diag, rowIdx, single_rates=None, global_bin_edges=None, xLab=None, isExp=False): + """Plot histogram of a specific row from the 2D matrix with statistics""" + + # Extract row data (exclude diagonal element which is NaN) + row_data = A_fit_no_diag[rowIdx, :] + valid_data = row_data[~np.isnan(row_data)] + + if len(valid_data) == 0: + ax.text(0.5, 0.5, 'No valid data', transform=ax.transAxes, ha='center', va='center') + return + + # Plot histogram with consistent binning + if global_bin_edges is not None: + counts, bin_edges, _ = ax.hist(valid_data, bins=global_bin_edges, color='chocolate') + else: + neve_happens + n_bins = min(60, len(valid_data)//3) # Adaptive bin count fallback + counts, bin_edges, _ = ax.hist(valid_data, bins=n_bins, alpha=0.7, color='skyblue') + + # Calculate median value + xMedian = np.median(valid_data) + + # Define slice range around median + xDel = 0.1 + slice_mask = (valid_data >= (xMedian - xDel)) & (valid_data <= (xMedian + xDel)) + sliced_data = valid_data[slice_mask] + + # Compute standard deviation of sliced data + if len(sliced_data) > 1: + std_val = np.std(sliced_data) + x0 = np.mean(sliced_data) + else: + std_val = 0.0 + x0 = xMedian + + # Draw x=0 reference line (always shown) + ax.axvline(0, linestyle='--', color='lime', linewidth=1) + + # Draw median line (always shown) + ax.axvline(xMedian, linestyle='-', color='k', linewidth=2, label='Median') + + # Draw dashed lines for the ±0.1 range (only for simulation data) + if not isExp: + ax.axvline(xMedian - xDel, linestyle='--', color='k', alpha=0.7) + ax.axvline(xMedian + xDel, linestyle='--', color='k', alpha=0.7) + else: + ax.set_yscale('log') + + # Add merged text with statistics and row/frequency info inside the plot + freq_text = f', {single_rates[rowIdx]:.1f}Hz' if single_rates is not None else '' + info_text = f'Row {rowIdx}{freq_text}\nmedian={xMedian:.3f}, std={std_val:.3f}' + + ax.text(0.98, 0.95, info_text, transform=ax.transAxes, fontsize=8, + verticalalignment='top', horizontalalignment='right', + bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.8)) + + # Set consistent x-axis limits when using global binning + if global_bin_edges is not None: + ax.set_xlim(global_bin_edges[0], global_bin_edges[-1]) + + if xLab!=None: ax.set_xlabel(xLab) + + ax.set_ylabel('Count') # Restore y-axis label + ax.grid(True, alpha=0.3) diff --git a/causal_net/stationaryLag1_ver6_lasso/PlotterSimPoisson.py b/causal_net/stationaryLag1_ver6_lasso/PlotterSimPoisson.py new file mode 100644 index 00000000..66c766f5 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/PlotterSimPoisson.py @@ -0,0 +1,336 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for simulated Poisson process visualization. + +This module provides specialized plotting capabilities for analyzing +simulated Dale's principle neural networks. The Plotter class extends +PlotterBackbone to create visualizations including: +- Dale connectivity matrix plots with excitatory/inhibitory separation +- Eigenvalue analysis and network stability visualization +- Poisson process statistics and firing rate distributions +- Network dynamics and temporal evolution plots + +Designed specifically for validating and analyzing simulated neural +networks that follow Dale's principle with Poisson spiking dynamics. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec +import matplotlib.colors as colors +import numpy as np +import matplotlib.pyplot as plt +from scipy.optimize import curve_fit +from Util_pseudospectra import compute_pseudospectrum + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def Dale_matrix_and_eigen(self,W_freq,md,trueD,figId=3): + + figId=self.smart_append(figId) + nrow,ncol=1,3 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,4.)) + + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + + # Reconstruct natural-order matrix from frequency-sorted matrix + neur_freqIdx = trueD['neur_freqIdx'] # natural_index → freq_sorted_position + W_natural = W_freq[np.ix_(neur_freqIdx, neur_freqIdx)] + + # Common normalization for both plots + vmin = min(W_freq.min(), W_natural.min()) + vmax = max(W_freq.max(), W_natural.max()) + normMap = colors.TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax) + + #.... Position 1: Natural order Dale matrix ...... + ax = self.plt.subplot(nrow,ncol,1) + im=ax.imshow(W_natural, aspect=1., origin='lower', cmap='bwr', norm=normMap, interpolation='nearest') + ax.set( xlabel='presyn. neuron index (natural idx)', ylabel='postsyn. neuron index (natural idx)') + ax.set_aspect(1.0) + ax.grid() + cbar = fig.colorbar(im, ax=ax, extend="both", shrink=0.5) + + tit_natural='True Dale (natural), M%d, %s'%(W_natural.shape[0],md['short_name']) + ax.set(title=tit_natural) + ax.axhline(numExc-0.5,color='k',ls='--', label='E/I boundary') + ax.axvline(numExc-0.5,color='k',ls='--') + + #.... Position 2: Frequency-sorted Dale matrix ...... + ax = self.plt.subplot(nrow,ncol,2) + im=ax.imshow(W_freq, aspect=1., origin='lower', cmap='bwr', norm=normMap, interpolation='nearest') + ax.set( xlabel='presyn. neuron index (freq-sorted)', ylabel='postsyn. neuron index (freq-sorted)') + ax.set_aspect(1.0) + ax.grid() + cbar = fig.colorbar(im, ax=ax, extend="both", shrink=0.7) + + tit_freq='True Dale (freq-sorted), M%d, %s'%(W_freq.shape[0],md['short_name']) + ax.set(title=tit_freq) + # Note: In frequency-sorted order, excitatory/inhibitory neurons are mixed, so no simple boundary line + + #..... Position 3: Eigenvalues...... + ax = self.plt.subplot(nrow,ncol,3) + Eigen=np.linalg.eigvals(W_freq) # Use freq-sorted matrix for eigenvalues + real_parts = np.real(Eigen) + imag_parts = np.imag(Eigen) + ax.scatter(real_parts, imag_parts, color='blue', marker='o') + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + tit3='Eigenvalues, M%d, %s'%(W_freq.shape[0],md['short_name']) + ax.set_title(tit3) + ax.axhline(0, color='black', lw=0.5) + ax.axvline(0, color='black', lw=0.5) + ax.axvline(0,color='red', linestyle='--') + ax.grid(True) + + def Dale_matrix_pseudospectra(self, A, md, trueD, figId=5): # p=e + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(8, 6)) + + npts = 80 + minY = -0.2 + epsMin = 0.03 + + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + ax = self.plt.subplot(1, 1, 1) + title = 'Pseudospectra, true M%d, %s' % (A.shape[0], md['short_name']) + + levels = np.logspace(-2.5, -0.5, 10) + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + ax.text(0.05, 0.85, 'green: ε<%.3f'%epsMin, transform=ax.transAxes) + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) + + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(bottom=minY) + ax.set_xlim(-3,) +#...!...!.................. + def histo_weights_rates(self,trueD,spikeD,md,byFreq=False,figId=3): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,4)) + data_name=md['short_name'] + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + step_size=md['evol_conf']['step_size'] + + A=trueD['A_true'] + + # Data is now stored in frequency-sorted order by default + single_rates = spikeD['single_rates'] # Already in freq-sorted order + + from UtilSelectFDR import get_offdiag_triplets + # output: np.column_stack([i_indices, j_indices, values]) + EposT=get_offdiag_triplets(A,isPos=True) + EnegT=get_offdiag_triplets(A,isPos=False) + print('True0 num edges pos=%d neg=%d'%(EposT.shape[0],EnegT.shape[0])) + + wMin=np.min(EnegT[:,2]) + wMax=np.max(EposT[:,2]) + + def count_elements(E, Nn=numNeur): + i_indices = E[:, 0].astype(int) + counts = np.bincount(i_indices, minlength=Nn) + return counts + + edgeCount=count_elements(EnegT) + count_elements(EposT) + m_inh1d=count_elements(EnegT)>0 + #print('num inh neur=',np.sum(m_inh1d), ' num inh edges:',np.sum(edgeCount)) + + #.... weights histogram + ax = self.plt.subplot(nrow,ncol,1) + binX= np.linspace(wMin, wMax, 100) + ax.hist(EnegT[:,2], bins=binX, color='blue', alpha=0.7, edgecolor=None,label='neg:%d'%EnegT.shape[0]) + ax.hist(EposT[:,2], bins=binX, color='red', alpha=0.7, edgecolor=None,label='pos:%d'%EposT.shape[0]) + + ax.legend(loc='upper left') + tit='True Dale, M%d, %s'%(A.shape[0],md['short_name']) + ax.set(title=tit, xlabel='True weight value',ylabel='num edges') + ax.axvline(0,color='k',ls='--') + ax.grid(True, alpha=0.3) + + #.... : histogram of rates + ax = self.plt.subplot(nrow,ncol,2) + yLog= md['evol_conf']['expRate'] + ax.hist(single_rates, bins=20)#, log=yLog) + x_vals = np.arange(numNeur) + ax.set_xlabel('num neurons') + ax.set_xlabel('Firing rate (Hz)') + ax.set_ylabel('num neurons') + ax.grid(True, alpha=0.3) + ax.set_title('Single rates spectrum') + ax.set_xlim(0,) + median_val = np.median(single_rates) + ax.axvline(median_val, color='r', linestyle='--', linewidth=1.5) + y_max = ax.get_ylim()[1] + median_text = f"median: {median_val:.2f} (Hz)\n N={single_rates.shape[0]}" + ax.text( x=median_val * 1.1, y=y_max * 0.7, s=median_text, color='red') + + # Prepare data for neuron-indexed plots based on display order + if byFreq: + # Display data in frequency-sorted order (as stored) + display_single_rates = single_rates + display_edgeTV = edgeCount + display_inh_mask = m_inh1d + neurXlab = 'freq sorted neurons index' + else: + # Convert to natural neuron order for display + neur_revFreqIdx = trueD['neur_revFreqIdx'] # freq_sorted_position → natural_index + + # Create arrays in natural order + display_single_rates = np.zeros_like(single_rates) + display_edgeTV = np.zeros_like(edgeCount) + + # Map firing rates and edge counts from freq-sorted back to natural order + display_single_rates[neur_revFreqIdx] = single_rates + display_edgeTV[neur_revFreqIdx] = edgeCount + + # For natural order, create inhibitory mask based on original neuron types + # First num_excite neurons are excitatory, rest are inhibitory + display_inh_mask = np.zeros(numNeur, dtype=bool) + display_inh_mask[numExc:] = True # Inhibitory neurons start at index numExc + neurXlab = 'natural indexed neurons' + + x_vals = np.arange(numNeur) + #.... edge count + ax = self.plt.subplot(nrow,ncol,3) + ax.fill_between(x_vals, display_edgeTV, step='mid', color='salmon', alpha=0.7) + ax.set_xlabel(neurXlab) + ax.set_ylabel('num true edges') + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + probLo, probHi = dmd['edge_prob'] + rho_title = f'outgoing edges, true, prob=[{probLo:.2f}, {probHi:.2f}]' + ax.set_title(rho_title) + + + #.... firing rates ..... + ax = self.plt.subplot(nrow,ncol,4) + chanW=0.9 + # Create masks for inhibitory and excitatory neurons + inh_mask_display = display_inh_mask + exc_mask_display = ~display_inh_mask + + ax.bar(x_vals[exc_mask_display], display_single_rates[exc_mask_display], width=chanW, color='red', align='center', alpha=0.7, label='Excitatory') + ax.bar(x_vals[inh_mask_display], display_single_rates[inh_mask_display], width=chanW, color='blue', align='center', alpha=0.7, label='Inhibitory') + + ax.set_xlabel(neurXlab) + ax.set_ylabel('Firing rate (Hz)') + ax.set_ylim(0,) + ax.grid(True, alpha=0.3) + ax.set_title('Single Neuron Firing Rates') + ax.legend() + + +#............................ +#............................ +#............................ + def rates_study(self,trueD,spikeD,md,figId=4): + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(15,4)) + data_name=md['short_name'] + dmd=md['dale_conf'] + numExc=dmd['num_excite'] + numNeur=dmd['num_neurons'] + + # Data are frequency-sorted in primary storage + B_idle = trueD['B_true'] # freq-sorted bias + single_rates = spikeD['single_rates'] # freq-sorted rates (Hz) + single_rates_snr = spikeD['sigle_rates_snr'] # freq-sorted SNR (dimensionless) + neur_revFreqIdx = trueD['neur_revFreqIdx'] # freq_sorted_position → natural_index + + # Build excitatory/inhibitory masks in freq-sorted order using natural index split + nat_index = neur_revFreqIdx + exc_mask = nat_index < numExc + inh_mask = ~exc_mask + + # 1) Scatter: x=B_idle, y=log(single_rates) + ax = self.plt.subplot(nrow,ncol,1) + ax.scatter(B_idle[exc_mask], single_rates[exc_mask], s=14, alpha=0.6, marker='^', facecolors='none', edgecolors='red', label='Excitatory') + ax.scatter(B_idle[inh_mask], single_rates[inh_mask], s=14, alpha=0.6, marker='^', facecolors='none', edgecolors='blue', label='Inhibitory') + ax.set_xlabel('true B_idle ') + ax.set_ylabel('single_rates (Hz)') + ax.set_yscale('log') + ax.grid(True, alpha=0.3) + tit='True Dale, M%d, %s'%(numNeur, data_name) + ax.set_title(tit) + + # reference line: y = exp(B), clipped to central 80% of B range + bmin = float(np.min(B_idle)) + bmax = float(np.max(B_idle)) + bLo = bmin + 0.1 * (bmax - bmin) + bHi = bmax - 0.1 * (bmax - bmin) + bx = np.linspace(bLo, bHi, 100) + by = np.exp(bx) + ax.plot(bx, by, linestyle='--', color='black', linewidth=0.8, label='y=exp(x)') + ax.legend() + + # 2) Scatter: x=B_idle, y=single SNR (from spikeD) + ax = self.plt.subplot(nrow,ncol,2) + ax.scatter(B_idle[exc_mask], single_rates_snr[exc_mask], s=14, alpha=0.6, facecolors='none', edgecolors='red', label='Excitatory') + ax.scatter(B_idle[inh_mask], single_rates_snr[inh_mask], s=14, alpha=0.6, facecolors='none', edgecolors='blue', label='Inhibitory') + ax.set_xlabel('true B_idle') + ax.set_ylabel('single SNR (rate^2/var)') + ax.grid(True, alpha=0.3) + ax.set_yscale('log') + ax.set_title('single SNR vs B_idle') + ax.legend() + + # 3) Histogram: single_rates for excitatory + ax3 = self.plt.subplot(nrow,ncol,3) + exc_vals = single_rates[exc_mask] + inh_vals = single_rates[inh_mask] + # compute common bins and range (start at 0) + exc_max = np.max(exc_vals) if exc_vals.size > 0 else 0.0 + inh_max = np.max(inh_vals) if inh_vals.size > 0 else 0.0 + x_max = max(exc_max, inh_max) + if x_max <= 0: x_max = 1.0 + #common_bins = np.linspace(0, x_max, 21) + common_bins = np.linspace(0, x_max, int(2*x_max)) + ax3.hist(exc_vals, bins=common_bins, color='red', alpha=0.7, edgecolor=None) + ax3.set_xlabel('single_rates (Hz)') + ax3.set_ylabel('num neurons') + ax3.grid(True, alpha=0.3) + ax3.set_title('Rates: Excitatory (N=%d)' % (numExc)) + + # 4) Histogram: single_rates for inhibitory + ax4 = self.plt.subplot(nrow,ncol,4) + ax4.hist(inh_vals, bins=common_bins, color='blue', alpha=0.7, edgecolor=None) + ax4.set_xlabel('single_rates (Hz)') + ax4.set_ylabel('num neurons') + ax4.grid(True, alpha=0.3) + ax4.set_title('Rates: Inhibitory (N=%d)' % (numNeur-numExc)) + # unify x-range starting at 0 for both histograms + ax3.set_xlim(0, x_max) + ax4.set_xlim(0, x_max) + diff --git a/causal_net/stationaryLag1_ver6_lasso/PoissonGLModel.py b/causal_net/stationaryLag1_ver6_lasso/PoissonGLModel.py new file mode 100644 index 00000000..d9c2352f --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/PoissonGLModel.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +""" +PyTorch implementation of Poisson Generalized Linear Model for neural spike data. + +This module defines the PoissonGLModel class, a neural network that models +multivariate Poisson spike counts using the GLM framework: + rate_t = exp(A @ Y_prev + B) * dt + +The model supports two initialization modes: +- Stage A: Standard parameterization with full A and B matrices +- Stage B: Efficient parameterization using only trainable weights with masks + +Key features: +- Recurrent connectivity matrix A (neuron-to-neuron weights) +- Bias term B (baseline firing rates) +- Forward pass computing Poisson rates from previous spike counts +- Custom negative log-likelihood loss function for Poisson distributions +- Support for sparse parameterization and noise injection +""" + +# PoissonGML refers to a Poisson Generalized Linear Model (GLM) +import torch +import torch.nn as nn + + +class PoissonGLModel(nn.Module): + """ + Generalized Linear Model for multivariate Poisson spike counts. + rate_t = exp(A @ Y_prev + B) * dt + + Constructor: + PoissonGLM(n_neurons) + - Stage A: random initialization of A and B + PoissonGLM(n_neurons, A_init, B_init, trainable_mask) + - Stage B: efficient parameterization with only trainable weights in a single tensor + """ + def __init__(self, n_neurons, A_init=None, B_init=None, trainable_mask=None, noise_scale=0.0): + super().__init__() + self.n_neurons = n_neurons + + if A_init is None: # Stage A: standard parameterization + self.A = nn.Parameter(torch.randn(n_neurons, n_neurons) * 0.1) + self.B = nn.Parameter(torch.randn(n_neurons) * 0.1) + self.is_stage_b = False + else: # Stage B: efficient parameterization + self.is_stage_b = True + + # Store fixed components as buffers + A_t = torch.tensor(A_init, dtype=torch.float32) + B_t = torch.tensor(B_init, dtype=torch.float32) + mask = torch.tensor(trainable_mask, dtype=torch.bool) + + self.register_buffer('A_fixed', A_t) # Full A matrix for initialization + self.register_buffer('B_fixed', B_t) # Fixed B values (will be overridden) + self.register_buffer('trainable_mask', mask) + + # Find trainable positions + trainable_A_indices = torch.nonzero(mask, as_tuple=True) + self.register_buffer('trainable_A_row_idx', trainable_A_indices[0]) + self.register_buffer('trainable_A_col_idx', trainable_A_indices[1]) + + # Count trainable parameters + n_trainable_A = mask.sum().item() + n_trainable_B = n_neurons # All B elements are trainable + total_trainable = n_trainable_A + n_trainable_B + + # Single parameter tensor for all trainable weights + self.C = nn.Parameter(torch.zeros(total_trainable)) + + # Initialize C with current values + optional noise + with torch.no_grad(): + # First part: trainable A elements + trainable_A_values = A_t[trainable_A_indices] + self.C[:n_trainable_A] = trainable_A_values + + # Second part: B elements + self.C[n_trainable_A:] = B_t + # Add random noise if requested + if noise_scale > 0: + noise = torch.randn_like(self.C) * noise_scale + self.C += noise + #print(f"Added random noise with scale {noise_scale} to Stage B initialization") + + self.n_trainable_A = n_trainable_A + + @property + def A(self): + if not self.is_stage_b: + return self._parameters['A'] + else: + # Reconstruct A matrix from C and fixed values + A_reconstructed = self.A_fixed.clone() + A_reconstructed[self.trainable_A_row_idx, self.trainable_A_col_idx] = self.C[:self.n_trainable_A] + return A_reconstructed + + @property + def B(self): + if not self.is_stage_b: + return self._parameters['B'] + else: + # B elements are stored in the second part of C + return self.C[self.n_trainable_A:] + + def forward(self, Y_prev, dt=0.01): + linear = torch.addmm(self.B.unsqueeze(0), Y_prev, self.A.t()).clamp(min=-10, max=10) + return torch.exp(linear) * dt + + +def poisson_nll_loss(spikes, targets, firing_rates): + eps = 1e-8 + firing_rates_safe = torch.maximum(firing_rates, torch.tensor(0.1, device=firing_rates.device)) + weights = 1.0 / firing_rates_safe + weights /= torch.mean(weights) + weights = weights.unsqueeze(0) + loss = -weights * targets * torch.log(spikes + eps) + weights * spikes + return loss.mean() diff --git a/causal_net/stationaryLag1_ver6_lasso/README.md b/causal_net/stationaryLag1_ver6_lasso/README.md new file mode 100644 index 00000000..f5df8cbc --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/README.md @@ -0,0 +1,235 @@ +# Neural Connectivity Inference with LASSO Poisson GLM + +This package provides a complete pipeline for inferring neural connectivity from spike count data using LASSO-regularized Poisson Generalized Linear Models (GLM) with statistical significance testing via False Discovery Rate (FDR) control. + +## Overview + +The pipeline implements a state-of-the-art approach for neural connectivity analysis that combines: +- **Poisson GLM modeling** for biologically realistic spike count data +- **LASSO regularization** for sparse connectivity estimation +- **Bootstrap resampling** for statistical robustness +- **FDR-controlled edge selection** for multiple comparison correction +- **Multi-GPU distributed training** for computational efficiency + +### Key Features +- Support for both simulated and experimental neural data +- Distributed training across multiple GPUs using PyTorch DDP +- Bootstrap-based statistical inference with configurable sample sizes +- Row-wise FDR control for conservative edge selection +- Comprehensive visualization and evaluation tools +- Dale's principle simulation for method validation + +## Workflow Pipeline + +``` +Data Preparation → Bootstrap Training → Edge Selection → Evaluation + ↓ ↓ ↓ ↓ +[prep_bioexp.py] [bootsFit.sh] [selectEdges_FDR.py] [eval_fitLasso.py] +[sim_dalePoisson.py] ↓ [view_bioexp.py] + [fitLasso4GPU.sh] + ↓ + [fit_lassoPoisson.py] +``` + +## Installation & Requirements + +- Python 3.8+ +- PyTorch with CUDA support +- NumPy, SciPy, Matplotlib +- statsmodels (for FDR correction) +- Custom toolbox modules (PlotterBackbone, Util_NumpyIO) + +## Executables Reference + +### Data Preparation Scripts + +#### `prep_bioexp.py` +**Purpose**: Preprocessing pipeline for experimental neural data from Roy/Mandar laboratory +**Input**: Raw experimental neural recordings +**Output**: Standardized `.spikes.npz` files with spike count matrices and metadata +**Usage**: +```bash +./prep_bioexp.py --sessionName B6J_250619_M08020_000093_Well000 --inputPath /path/to/raw/data/ +``` + +#### `sim_dalePoisson.py` +**Purpose**: Simulates Dale's principle neural networks with Poisson spiking dynamics +**Input**: Network parameters (num_neurons, connectivity strength, etc.) +**Output**: `.spikes.npz` and `.simTruth.npz` files with simulated data and ground truth +**Usage**: +```bash +./sim_dalePoisson.py --num_neurons 50 --num_excite 35 --num_steps 10000 --dataName test_dale +``` + +#### `view_bioexp.py` +**Purpose**: Visualization tool for biological experiment data quality assessment +**Input**: `.spikes.npz` files +**Output**: Interactive plots for exploratory data analysis +**Usage**: +```bash +./view_bioexp.py --dataPath /path/to/data/ --sessionName mydata -p a +``` + +### Model Training Scripts + +#### `bootsFit.sh` +**Purpose**: Bootstrap wrapper for running multiple LASSO training iterations +**Function**: Orchestrates multiple training runs with different random seeds/data splits +**Key Features**: +- Configurable number of bootstrap iterations (--num_bootstraps) +- Support for time-shuffled control data (--shuffleTime) +- Automatic naming convention (dataName-boot0, dataName-boot1, etc.) +- Progress tracking and error handling +- Pass-through of training parameters to underlying scripts + +**Usage**: +```bash +./bootsFit.sh --dataName mydata --num_bootstraps 10 --num_epochs 100 --dropDataFrac 0.5 +``` + +#### `fitLasso4GPU.sh` +**Purpose**: Multi-GPU wrapper for distributed LASSO Poisson training +**Function**: Sets up GPU environment and executes fit_lassoPoisson.py +**Key Features**: +- Multi-GPU configuration (4 GPUs default) +- Optimized environment variables (OMP_NUM_THREADS=1, etc.) +- Default hyperparameters (batch_size=2048, lr=1e-3, L1_alpha=1e-3) +- Execution timing and logging + +**Usage**: Called automatically by bootsFit.sh, not typically run directly + +#### `fit_lassoPoisson.py` +**Purpose**: Core distributed training of Poisson GLM with LASSO regularization +**Input**: `.spikes.npz` files with neural spike count data +**Output**: `.lassoFit.npz` files with fitted connectivity matrices and metadata +**Key Features**: +- PyTorch DistributedDataParallel for multi-GPU training +- Custom Poisson negative log-likelihood loss +- L1 regularization for sparse connectivity +- Support for time decorrelation and data shuffling +- Automatic model checkpointing + +**Usage**: +```bash +# Multi-GPU (via torchrun) +OMP_NUM_THREADS=1 CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 ./fit_lassoPoisson.py --dataName mydata --num_epochs 100 + +# Single GPU +./fit_lassoPoisson.py --dataName mydata --num_epochs 100 +``` + +### Statistical Analysis Scripts + +#### `selectEdges_FDR.py` +**Purpose**: Row-wise FDR edge selection from LASSO bootstrap results +**Input**: Multiple `.lassoFit.npz` files from bootstrap training +**Output**: `.selFdr.lassoFit.npz` with FDR-selected edges and statistics +**Key Features**: +- Row-wise Benjamini-Hochberg FDR correction +- Configurable number of real vs shuffled bootstraps +- Empirical p-value computation using pooled null distribution +- Comprehensive statistical summaries and edge selection metrics + +**Usage**: +```bash +# Single value (duplicated for real/shuffled) +./selectEdges_FDR.py --dataName mydata --num_bootstraps 10 --alpha 0.01 + +# Different numbers of real vs shuffled +./selectEdges_FDR.py --dataName mydata --num_bootstraps 8 12 --alpha 0.01 +``` + +### Evaluation Scripts + +#### `eval_fitLasso.py` +**Purpose**: Comprehensive evaluation and visualization of LASSO fitting results +**Input**: `.lassoFit.npz` files (individual fits or FDR-selected) +**Output**: Statistical summaries and visualization plots +**Key Features**: +- Connectivity matrix visualization with frequency sorting +- Weight distribution analysis and sparsity metrics +- Network structure plots for excitatory/inhibitory connections +- Ground truth comparison for simulated data + +**Usage**: +```bash +./eval_fitLasso.py --dataName mydata-selFdr --dataPath /path/to/data/ -p ab +``` + +#### `eval_fitRegress.py` +**Purpose**: Evaluation of regression-based connectivity analysis +**Input**: Regression fit results +**Output**: Alternative connectivity analysis using regression methods +**Usage**: +```bash +./eval_fitRegress.py --dataName mydata --dataPath /path/to/data/ -p ab +``` + +## Utility Modules + +### Core Utilities +- **`PoissonGLModel.py`**: PyTorch implementation of Poisson GLM with dual parameterization modes +- **`UtilTorch.py`**: GPU setup, data preprocessing, and distributed training utilities +- **`UtileSelectFDR.py`**: Statistical functions for FDR edge selection and bootstrap analysis +- **`UtilDalePoisson.py`**: Dale's principle simulation utilities and connectivity analysis +- **`UtilBioExp.py`**: Biological experiment data processing and cluster detection +- **`UtilFreqGen.py`**: Realistic neural firing frequency distribution generation + +### Plotting Modules +- **`PlotterFitEval.py`**: LASSO fit evaluation and connectivity visualization +- **`PlotterBioExp.py`**: Biological experiment data visualization +- **`PlotterSimPoisson.py`**: Simulated neural network and Dale's principle plots + +## Typical Analysis Workflow + +### 1. Data Preparation +```bash +# For experimental data +./prep_bioexp.py --sessionName myexperiment --inputPath /raw/data/ + +# For simulated data +./sim_dalePoisson.py --num_neurons 50 --num_excite 35 --dataName mysim +``` + +### 2. Bootstrap Training +```bash +# Train multiple bootstrap iterations +./bootsFit.sh --dataName mydata --num_bootstraps 10 --num_epochs 200 --dropDataFrac 0.5 + +# Optional: Train time-shuffled controls +./bootsFit.sh --dataName mydata --num_bootstraps 10 --shuffleTime --num_epochs 200 +``` + +### 3. Statistical Edge Selection +```bash +# Apply FDR correction with real and shuffled data +./selectEdges_FDR.py --dataName mydata --num_bootstraps 10 --alpha 0.01 -p af +``` + +### 4. Results Evaluation +```bash +# Visualize selected connectivity +./eval_fitLasso.py --dataName mydata-selFdr --dataPath /path/to/results/ -p abc + +# Compare with ground truth (for simulations) +./eval_fitLasso.py --dataName mysim-selFdr --dataPath /path/to/results/ -p abcd +``` + +## Output File Conventions + +- **`.spikes.npz`**: Preprocessed spike count matrices +- **`.simTruth.npz`**: Ground truth connectivity (simulations only) +- **`dataName-bootX.lassoFit.npz`**: Individual bootstrap training results +- **`dataName-shufX.lassoFit.npz`**: Time-shuffled control results +- **`dataName-selFdr.lassoFit.npz`**: FDR-selected final connectivity + +## Performance Considerations + +- **GPU Memory**: Ensure sufficient GPU memory for large networks (>100 neurons) +- **Bootstrap Size**: More bootstraps improve statistical power but increase computation time +- **Batch Size**: Adjust based on GPU memory and network size +- **Distributed Training**: Use multiple GPUs for networks >50 neurons or long recordings + +## Citation + +If you use this code, please cite the appropriate papers describing the methods implemented here. diff --git a/causal_net/stationaryLag1_ver6_lasso/UtilBioExp.py b/causal_net/stationaryLag1_ver6_lasso/UtilBioExp.py new file mode 100644 index 00000000..7aefdd8d --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/UtilBioExp.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +""" +Utility functions for biological experiment data processing. + +This module provides specialized functions for analyzing biological neural data, +particularly focused on identifying and processing clusters of neural activity. +Main functionality includes: +- Cluster detection in time series data using connectivity analysis +- Threshold-based cluster validation and masking +- Signal processing utilities for experimental neural recordings + +Used primarily in conjunction with experimental data preprocessing and +analysis pipelines for biological neural network studies. +""" + +import numpy as np +from scipy.ndimage import label + +def create_clusters_mask(X, th): + """ + Creates a boolean mask for a 1D NumPy array. + + A contiguous cluster of non-zero values is marked as True if any value + within that cluster is greater than or equal to the threshold 'th'. + Clusters are assumed to be separated by zero values. + + Parameters: + X (np.ndarray): Input 1D NumPy array. + th (float): The threshold value. + + Returns: + np.ndarray: A boolean mask with the same shape as X. + """ + # Step 1: Find all contiguous clusters of non-zero values. + # The `label` function assigns a unique integer to each cluster. + # e.g., [1, 1, 1, 0, 2, 2, 2, 0, 3, 3, 3] + labeled_array, num_clusters = label(X > 0) + + # Step 2: Initialize the final mask to all False. + final_mask = np.zeros_like(X, dtype=bool) + + # Step 3: Iterate through each cluster found by the label function. + # We start from 1 because 0 is the background (the zeros in X). + for i in range(1, num_clusters + 1): + # Create a boolean mask for the current cluster only. + current_cluster_mask = (labeled_array == i) + + # Step 4: Check if ANY value within this specific cluster meets the threshold. + if np.any(X[current_cluster_mask] >= th): + # Step 5: If the condition is met, mark this entire cluster as True in our final mask. + final_mask[current_cluster_mask] = True + + return final_mask + +if __name__=="__main__": + + # Example usage + X = np.array([1, 2, 3, 0, 2, 6, 2, 0, 1, 2, 8, 0, 0]) + threshold = 5 + mask = create_clusters_mask(X, threshold) + + # Print the output in two columns + print(f"Input | Mask , thr:{threshold}") + print("-" * 25) + for x_val, mask_val in zip(X, mask): + print(f"{x_val:<4} | {mask_val}") diff --git a/causal_net/stationaryLag1_ver6_lasso/UtilDalePoisson.py b/causal_net/stationaryLag1_ver6_lasso/UtilDalePoisson.py new file mode 100644 index 00000000..d5e5f34b --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/UtilDalePoisson.py @@ -0,0 +1,246 @@ +""" +Utility functions for Dale Poisson simulation data processing + +""" + +import numpy as np +import os +import time +from pprint import pprint + +def compute_consecutive_coincidence_rate(Y,time_evol): + """ + Compute the frequency of coincidences for 2 consecutive time bins for 2 different channels. + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + + Returns: + float: conc_rate_all - the overall coincidence rate + """ + num_steps, num_neurons = Y.shape + + # We need at least 2 time steps for consecutive bins + if num_steps < 2: + return 0.0 + + # Get consecutive time slices using vectorized operations + Y_t = Y[:-1, :] # Y[0:T-1, :] - current time bins + Y_t_plus_1 = Y[1:, :] # Y[1:T, :] - next time bins + + # Create boolean masks for non-zero values + mask_t = (Y_t != 0) # Shape: (T-1, N) + mask_t_plus_1 = (Y_t_plus_1 != 0) # Shape: (T-1, N) + + # Use broadcasting to compute all channel pairs at once + # mask_t[:, :, None] has shape (T-1, N, 1) + # mask_t_plus_1[:, None, :] has shape (T-1, 1, N) + # Broadcasting gives shape (T-1, N, N) for all pairs + coincidences = mask_t[:, :, None] & mask_t_plus_1[:, None, :] + + # Remove diagonal (same channel pairs) using boolean indexing + diagonal_mask = np.eye(num_neurons, dtype=bool) + coincidences[:, diagonal_mask] = False + + # Count total coincidences across all time steps + coincidence_count = np.sum(coincidences) + + conc_rate = coincidence_count / time_evol/num_neurons + return conc_rate # Hz, per neuron + +def estimate_rates(Y, dt, num_excite, max_samples, varTwindow=5, mxNn=5): + """ + Evaluates spike statistics and estimates firing rates, Fano factors, and SNR. + + Computes statistics over non-overlapping time windows of length varTwindow: + - Fano Factor: Var[spike count] / Mean[spike count] per neuron (dimensionless) + - Rate variance: Var[rate] per neuron (Hz^2) + - SNR: (mean rate)^2 / Var[rate] per neuron (dimensionless) + + Args: + Y (np.ndarray): Spike data array (time_steps x n_neurons). + dt (float): Time bin size in seconds. + num_excite (int): Number of excitatory neurons. + max_samples (int): The maximum number of time samples to use for calculation. + varTwindow (float, optional): Time window length in seconds for variance computation. Defaults to 5. + mxNn (int, optional): Max number of neurons to show in detailed stats. Defaults to 5. + + Returns: + tuple: A tuple containing (stats_dict, rates_dict, neur_freq_index): + - stats_dict (dict): Contains population-level spike statistics. + - rates_dict (dict): Contains per-neuron firing rates, Fano factors, rate variance, and SNR. + - neur_freq_index (np.ndarray): Neuron indices sorted by firing rate (low to high). + """ + # 1. Clip data to max_samples_for_rates + print("\n=== Estimating Rates & Stats ===") + if Y.shape[0] > max_samples: + print("Using %d samples (out of %d) for rate computation" % (max_samples, Y.shape[0])) + Y = Y[:max_samples] + + # Part 1: Compute all basic statistics once + num_steps_sim, Nn_sim = Y.shape + num_inhib = Nn_sim - num_excite + time_evol = num_steps_sim * dt + print('steps num_steps=%d, time_evol=%.1f sec, Nn=%d (%d Excit, %d Inhib)' % (num_steps_sim, time_evol, Nn_sim, num_excite, num_inhib)) + + # Compute windowed spike counts and statistics over non-overlapping time windows of length varTwindow (sec), per neuron + window_size = max(1, int(round(varTwindow / dt))) + num_windows = num_steps_sim // window_size + print('variance computation: num_windows=%d, window_size=%d' % (num_windows, window_size)) + if num_windows <= 1: + raise ValueError("varTwindow=%s is too large for data (num_windows=%d). Need at least 2 windows for variance." % (varTwindow, num_windows)) + Y_trim = Y[:num_windows * window_size] + Y_win = Y_trim.reshape(num_windows, window_size, Nn_sim) + + # Sum spike counts over each window: shape (num_windows, n_neurons) + window_spike_counts = np.sum(Y_win, axis=1) + + # Fano Factor: Var[spike count] / Mean[spike count] over windows + mean_window_spike_counts = np.mean(window_spike_counts, axis=0) + var_window_spike_counts = np.var(window_spike_counts, axis=0) + fano_factor = np.divide(var_window_spike_counts, mean_window_spike_counts, + out=np.zeros_like(var_window_spike_counts), + where=mean_window_spike_counts != 0) + + # Convert window spike counts to rates (Hz) for rate variance and SNR + window_rates = window_spike_counts / (window_size * dt) # shape: (num_windows, n_neurons), Hz + mean_window_rates = np.mean(window_rates, axis=0) # mean rate per neuron (Hz) + single_rates_var = np.var(window_rates, axis=0) # variance of rate across windows, per neuron (Hz^2) + + # SNR: (mean rate)^2 / var(rate) - dimensionless + single_rates_snr = np.divide(mean_window_rates**2, single_rates_var, + out=np.zeros_like(single_rates_var), + where=single_rates_var != 0) + + # Compute raw arrays (full time series) + spike_counts = np.sum(Y, axis=0) + spike_rates = spike_counts / time_evol + mean_counts_per_bin = np.mean(Y, axis=0) + + # Compute consecutive coincidence rate + start_time = time.time() + conc_rate_per_neuron = compute_consecutive_coincidence_rate(Y,time_evol) + elapsed_time = time.time() - start_time + print('Coincidence rate %.2g Hz, Y.shape=%s elaT %.3f sec' % (conc_rate_per_neuron, Y.shape, elapsed_time)) + + # Compute all population statistics once + med_rate_all = np.median(spike_rates) + avg_rate_all = float(np.mean(spike_rates)) + std_rate_all = float(np.std(spike_rates)) + avg_fano_all = float(np.mean(fano_factor)) + std_fano_all = float(np.std(fano_factor)) + avg_snr_all = float(np.mean(single_rates_snr)) + std_snr_all = float(np.std(single_rates_snr)) + avg_rate_excit = float(np.mean(spike_rates[:num_excite])) + std_rate_excit = float(np.std(spike_rates[:num_excite])) + avg_fano_excit = float(np.mean(fano_factor[:num_excite])) + std_fano_excit = float(np.std(fano_factor[:num_excite])) + avg_snr_excit = float(np.mean(single_rates_snr[:num_excite])) + std_snr_excit = float(np.std(single_rates_snr[:num_excite])) + avg_rate_inhib = float(np.mean(spike_rates[num_excite:])) + std_rate_inhib = float(np.std(spike_rates[num_excite:])) + avg_fano_inhib = float(np.mean(fano_factor[num_excite:])) + std_fano_inhib = float(np.std(fano_factor[num_excite:])) + avg_snr_inhib = float(np.mean(single_rates_snr[num_excite:])) + std_snr_inhib = float(np.std(single_rates_snr[num_excite:])) + + # Build stats dictionary with computed values + stats_dict = { + 'num_steps': num_steps_sim, + 'time_evol_sec': time_evol, + 'time_step_sec': dt, + 'num_neurons': Nn_sim, + 'num_excitatory': num_excite, + 'num_inhibitory': num_inhib, + 'avg_spike_rate_all': avg_rate_all, + 'std_spike_rate_all': std_rate_all, + 'avg_fano_factor_all': avg_fano_all, + 'std_fano_factor_all': std_fano_all, + 'avg_snr_all': avg_snr_all, + 'std_snr_all': std_snr_all, + 'avg_spike_rate_excit': avg_rate_excit, + 'std_spike_rate_excit': std_rate_excit, + 'avg_fano_factor_excit': avg_fano_excit, + 'std_fano_factor_excit': std_fano_excit, + 'avg_snr_excit': avg_snr_excit, + 'std_snr_excit': std_snr_excit, + 'avg_spike_rate_inhib': avg_rate_inhib, + 'std_spike_rate_inhib': std_rate_inhib, + 'avg_fano_factor_inhib': avg_fano_inhib, + 'std_fano_factor_inhib': std_fano_inhib, + 'avg_snr_inhib': avg_snr_inhib, + 'std_snr_inhib': std_snr_inhib, + 'conc_rate_per_neuron': float(conc_rate_per_neuron), + 'median_spike_rate_all': med_rate_all + } + + # Printing detailed stats for individual neurons + mxE = min(mxNn, num_excite) + mxI = min(mxNn, num_inhib) + + print('\n--- Stats for first %d Excitatory Neurons ---' % mxE) + np.set_printoptions(precision=2) + print('Total Spike Counts: %s' % spike_counts[:mxE]) + print('Mean Firing Rate (Hz): %s' % spike_rates[:mxE]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[:mxE])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[:mxE])) + print('Fano Factor (Var/Mean): %s' % fano_factor[:mxE]) + print('SNR (rate^2/var): %s' % single_rates_snr[:mxE]) + + print('\n--- Stats for first %d Inhibitory Neurons ---' % mxI) + np.set_printoptions(precision=2) + inhib_slice = slice(num_excite, num_excite + mxI) + print('Total Spike Counts: %s' % spike_counts[inhib_slice]) + print('Mean Firing Rate (Hz): %s' % spike_rates[inhib_slice]) + print('Mean Spike Count per bin (dt=%.3fs): %s' % (dt, mean_counts_per_bin[inhib_slice])) + print('Rate Variance (window=%.1fs): %s' % (varTwindow, single_rates_var[inhib_slice])) + print('Fano Factor (Var/Mean): %s' % fano_factor[inhib_slice]) + print('SNR (rate^2/var): %s' % single_rates_snr[inhib_slice]) + + # Print population summary using dictionary values + print('\n--- Population Summary Statistics ---') + print('All (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f, Avg SNR=%.2f±%.2f' % (Nn_sim, stats_dict['avg_spike_rate_all'], stats_dict['std_spike_rate_all'], stats_dict['avg_fano_factor_all'], stats_dict['std_fano_factor_all'], stats_dict['avg_snr_all'], stats_dict['std_snr_all'])) + print('Excit (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f, Avg SNR=%.2f±%.2f' % (num_excite, stats_dict['avg_spike_rate_excit'], stats_dict['std_spike_rate_excit'], stats_dict['avg_fano_factor_excit'], stats_dict['std_fano_factor_excit'], stats_dict['avg_snr_excit'], stats_dict['std_snr_excit'])) + if num_inhib > 0: + print('Inhib (%d neurons): Avg Rate=%.2f±%.2f Hz, Avg Fano=%.2f±%.2f, Avg SNR=%.2f±%.2f' % (num_inhib, stats_dict['avg_spike_rate_inhib'], stats_dict['std_spike_rate_inhib'], stats_dict['avg_fano_factor_inhib'], stats_dict['std_fano_factor_inhib'], stats_dict['avg_snr_inhib'], stats_dict['std_snr_inhib'])) + + # Print key summary values using dictionary + summary_keys = ['conc_rate_per_neuron', 'median_spike_rate_all'] + for key in summary_keys: + if key == 'conc_rate_per_neuron': + print('Coincidence rate per neuron %.2g Hz' % stats_dict[key]) + elif key == 'median_spike_rate_all': + print('Median rate %.2f Hz\n' % stats_dict[key]) + + # Part 2: Compute frequency sorting + neur_freq_index = np.argsort(spike_rates) + + rates_dict = { + 'single_rates': spike_rates, + 'neur_freqIdx':neur_freq_index, + 'sigle_rates_var': single_rates_var, + 'sigle_rates_snr': single_rates_snr, + 'single_fano_fact': fano_factor + } + + return stats_dict, rates_dict, neur_freq_index + + + +def do_neuron_classifier(A): # ??? + # + thrMaj=0.9 # edge count for: pure | majority + thrMix=0.5 # edge count for: majority | mix + + # Create mask for off-diagonal elements + offdiag_mask = ~np.eye(A.shape[0], A.shape[1], dtype=bool) + posEdge_mask= A>0 + negEdge_mask= A<0 + + # combine masks + + mask = posEdge_mask & offdiag_mask + i_indices, j_indices = np.where(mask) + values = A[i_indices, j_indices] + not_used_yet + return np.column_stack([i_indices, j_indices, values]) diff --git a/causal_net/stationaryLag1_ver6_lasso/UtilGenExpFreqs.py b/causal_net/stationaryLag1_ver6_lasso/UtilGenExpFreqs.py new file mode 100755 index 00000000..0b9ee315 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/UtilGenExpFreqs.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 + +import numpy as np +import matplotlib.pyplot as plt + +def gen_realistic_freqs(min_freq, max_freq, trapezoid_height, trapezoid_rmin, num_samples, sigma): + """ + Generate samples using a combination of Gaussian and trapezoidal distributions. + + Parameters: + - min_freq: Minimum frequency value. + - max_freq: Maximum frequency value. + - trapezoid_height: Height of the trapezoidal distribution at min_freq. + - trapezoid_rmin: Relative reduction at max_freq (0..1). + - num_samples: Total number of samples to generate. + - sigma: Standard deviation of the Gaussian distribution. + + Returns: + - samples: Array of generated samples. + """ + + # Define the Gaussian PDF (centered at 0, evaluated for f>=min_freq) + def gaussian_pdf(f): + return np.exp(- (f / sigma) ** 2 / 2) + + # Define the trapezoidal PDF using vectorized operations + def trapezoidal_pdf(f): + inside = (f >= min_freq) & (f <= max_freq) + val = np.zeros_like(f, dtype=float) + # Linearly decreasing from trapezoid_height at min_freq + val[inside] = trapezoid_height * (1 - (f[inside] - min_freq) / (max_freq - min_freq) * trapezoid_rmin) + return val + + # Create a combined PDF + def combined_pdf(f): + return gaussian_pdf(f) + trapezoidal_pdf(f) + + # Rejection sampling over continuous interval [min_freq, max_freq] to avoid quantization + x_dense = np.linspace(min_freq, max_freq, 5000) + pdf_dense = combined_pdf(x_dense) + M = float(np.max(pdf_dense)) if np.all(np.isfinite(pdf_dense)) else 1.0 + if M <= 0: + # Fallback: uniform if pathological configuration + return np.random.uniform(min_freq, max_freq, size=num_samples) + + samples = np.empty(num_samples, dtype=float) + filled = 0 + # Choose a reasonable batch size for efficiency + batch_size = max(1000, num_samples * 2) + while filled < num_samples: + f_try = np.random.uniform(min_freq, max_freq, size=batch_size) + u = np.random.uniform(0.0, M, size=batch_size) + accept_mask = u < combined_pdf(f_try) + accepted = f_try[accept_mask] + take = min(accepted.size, num_samples - filled) + if take > 0: + samples[filled:filled+take] = accepted[:take] + filled += take + + return samples + +if __name__ == '__main__': + + # Parameters + min_freq = 0.5 # Set minimum frequency + max_freq = 90 #(Hz) + trapezoid_height = 0.15 # Height of the trapezoidal distribution + trapezoid_rmin=0.3 # fraction of trapzoid at the max freq + sigma = 5 # (Hz) Standard deviation of the Gaussian + + num_samples = 150 + + # Generate samples + samples = gen_realistic_freqs(min_freq, max_freq, trapezoid_height, trapezoid_rmin,num_samples, sigma) + + print('samples:',samples[samples<5]) + # Calculate the median of the samples + median_value = np.median(samples) + + # Plotting the generated samples + fig, ax = plt.subplots(figsize=(8, 3)) + + # Create the histogram without normalization + ax.hist(samples, bins=np.arange(min_freq, max_freq + 2, 2), color='steelblue', edgecolor='black', alpha=0.7) + + # Draw the median line + ax.axvline(median_value, color='red', linestyle='--', linewidth=2, label=f'Median: {median_value:.2f}') + + # Formatting + ax.set_xlabel('Frequency (Hz)', fontsize=12) + ax.set_ylabel('Count', fontsize=12) + ax.set_title('Combined Distribution of Gaussian and Trapezoidal', fontsize=14) + ax.grid(True, alpha=0.3) + ax.legend() + + # Print parameters inside the canvas + params_text = (f"Min Frequency: {min_freq} Hz\n" + f"Max Frequency: {max_freq} Hz\n" + f"Trapezoid Height: {trapezoid_height}\n" + f"Number of Samples: {num_samples}\n" + f"Gaussian Sigma: {sigma:.2f}\n" + f"Median: {median_value:.2f}") + + # Adding text to the plot + ax.text(0.05, 0.95, params_text, transform=ax.transAxes, fontsize=10, + verticalalignment='top', bbox=dict(facecolor='white', alpha=0.5)) + + plt.tight_layout() + plt.show() diff --git a/causal_net/stationaryLag1_ver6_lasso/UtilSelectFDR.py b/causal_net/stationaryLag1_ver6_lasso/UtilSelectFDR.py new file mode 100644 index 00000000..f470b8e3 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/UtilSelectFDR.py @@ -0,0 +1,375 @@ +#!/usr/bin/env python3 +""" +Utility functions for FDR edge selection from LASSO bootstraps. +""" +import os +import numpy as np +from statsmodels.stats.multitest import fdrcorrection +from toolbox.Util_NumpyIO import read_data_npz + +def summary_reco_neuronNet(A_avr, A_std): + """ + Analyze neuron connectivity patterns. + + Parameters: + A_avr: 2D square array of average connection weights + A_std: 2D square array of standard deviations + + Returns: + Neu_sum: array with [neuron_id, n_positive, n_negative] for each neuron + Neu_edg: array with [neuron_id, target_id, avg_weight, std_dev] for non-zero connections + """ + n_neurons = A_avr.shape[0] + + # Part A: Count positive and negative values per row + Neu_sum = np.zeros((n_neurons, 4), dtype=int) + Neu_sum[:, 0] = np.arange(n_neurons) # neuron indices + Neu_sum[:, 1] = (A_avr > 0).sum(axis=1) # count positives per row + Neu_sum[:, 2] = (A_avr < 0).sum(axis=1) # count negatives per row + Neu_sum[:, 3] = Neu_sum[:, 1] + Neu_sum[:, 2] # cound any non-zero edges + + # Part B: Extract non-zero edges with their statistics + # Find all non-zero positions + i_neur, j_target = np.nonzero(A_avr) + + # Create edge array + n_edges = len(i_neur) + Neu_edg = np.zeros((n_edges, 4)) + Neu_edg[:, 0] = i_neur # source neuron + Neu_edg[:, 1] = j_target # target neuron + Neu_edg[:, 2] = A_avr[i_neur, j_target] # average weights + Neu_edg[:, 3] = A_std[i_neur, j_target] # standard deviations + + edgeD={} + edgeD['edge_sum']=Neu_sum + edgeD['edge_vals']=Neu_edg + + #pprint(Neu_sum) + return edgeD + +def compare_triplets(Rec, trueD): + """ + Compares two triplet arrays and returns True Positives, False Positives, and False Negatives. + + Parameters: + Rec : array of shape (n, 3) with [i, j, value] triplets (Reconstructed/Predicted) + trueD : array of shape (m, 3) with [i, j, value] triplets (Measured/Ground Truth) + + Returns: + TP : array of shape (k, 4) with [i, j, rec_value, meas_value] for matching indices + FP : array of shape (p, 3) with [i, j, value] for indices only in Rec + FN : array of shape (q, 3) with [i, j, value] for indices only in Meas + """ + # Convert to dictionaries for efficient lookup + rec_dict = {(int(row[0]), int(row[1])): row[2] for row in Rec} + true_dict = {(int(row[0]), int(row[1])): row[2] for row in trueD} + + # Get sets of indices + rec_indices = set(rec_dict.keys()) + true_indices = set(true_dict.keys()) + + # Find TP, FP, FN indices + tp_indices = rec_indices & true_indices # Intersection + fp_indices = rec_indices - true_indices # In Rec but not in Meas + fn_indices = true_indices - rec_indices # In Meas but not in Rec + + # Build output arrays + TP = np.array([[i, j, rec_dict[(i,j)], true_dict[(i,j)]] + for i, j in sorted(tp_indices)]) + FP = np.array([[i, j, rec_dict[(i,j)]] + for i, j in sorted(fp_indices)]) + FN = np.array([[i, j, true_dict[(i,j)]] + for i, j in sorted(fn_indices)]) + + # Handle empty arrays + if len(TP) == 0: + TP = np.empty((0, 4)) # i,j,vr,vt + if len(FP) == 0: + FP = np.empty((0, 3)) # i,j,vr + if len(FN) == 0: + FN = np.empty((0, 3)) # i,j,vr,vt + + print('comp tripl TP=%d, FP=%d, FN=%d'%(TP.shape[0],FP.shape[0],FN.shape[0])) + #print('shapes tripl TP=%s, FP=%s, FN=%s'%(TP.shape,FP.shape,FN.shape)) + return [TP, FP, FN] + +def eval_tagged_edges_4_simu(fitD, trueD): + #print(sorted(trueD)) + At=trueD['A_true'] + Bt=trueD['B_true'] + EposT=get_offdiag_triplets(At,True) + EnegT=get_offdiag_triplets(At,False) + print('True num edges pos=%d neg=%d'%(EposT.shape[0],EnegT.shape[0])) + + Ar=fitD['A_avr'] # reco + Br=fitD['B_avr'] # reco + EposR=get_offdiag_triplets(Ar,True) + EnegR=get_offdiag_triplets(Ar,False) + print('Reco num edges pos=%d neg=%d'%(EposR.shape[0],EnegR.shape[0])) + + #... zip diagonal + diag_At = np.diag(At) + diag_Ar = np.diag(Ar) + + evalD={} + evalD['pos']=compare_triplets(EposR, EposT) + evalD['neg']=compare_triplets(EnegR, EnegT) + evalD['diag']=np.column_stack([diag_Ar, diag_At]) + evalD['bterm']=np.column_stack([Br, Bt]) + + return evalD + + +def get_offdiag_triplets(A, isPos=True): + """ + Returns positive/negative, off-diagonal elements of 2D array A as array of [i, j, value] triplets. + + Parameters: + A : 2D numpy array + isPos : bool, if True only return positive values (>0), if False return all negative values + + Returns: + Array of shape (n, 3) where n is the number of valid off-diagonal elements + Each row is [row_index, col_index, value] + """ + # Create mask for off-diagonal elements + offdiag_mask = ~np.eye(A.shape[0], A.shape[1], dtype=bool) + + # Add value condition based on isPos + if isPos: + value_mask = A > 0 + else: + value_mask = A < 0 + + # Combine masks + mask = value_mask & offdiag_mask + + i_indices, j_indices = np.where(mask) + values = A[i_indices, j_indices] + + return np.column_stack([i_indices, j_indices, values]) + +def edge_selector_fdr(A_edges_real_list, A_edges_shuf_list, alpha=0.01): + """ + Apply row-wise FDR to select significant edges. + + Parameters + ---------- + A_edges_real_list : list of np.ndarray + List of Kreal real-data edge matrices (off-diagonals only). + A_edges_shuf_list : list of np.ndarray + List of Kshuf shuffled-data edge matrices (off-diagonals only). + alpha : float + FDR control level. + + Returns + ------- + E_edges : np.ndarray + Binary mask of significant edges (same shape as A_edges_real). + E_pval : np.ndarray + P-value matrix for each edge. + summary : dict + Summary statistics. + """ + + Kreal = len(A_edges_real_list) + Kshuf = len(A_edges_shuf_list) + + # Stack into arrays: shape (K, N, N) + A_real = np.stack(A_edges_real_list, axis=0) + A_shuf = np.stack(A_edges_shuf_list, axis=0) + + N = A_real.shape[1] + + # Statistic: median magnitude across bootstraps for each edge + stat_obs = np.median(np.abs(A_real), axis=0) # shape (N, N) + + # Null distribution: pool shuffled bootstraps per row + E_pval = np.ones((N, N)) + E_edge_mask = np.zeros((N, N), dtype=bool) + #E_edge_mask is a binary mask of significant edges after row-wise FDR. + + # Statistics tracking + edges_per_row = [] + + for i in range(N): + # Pool null values for row i from shuffled data + null_vals_row = np.abs(A_shuf[:, i, :]).ravel() + # Compute p-values for all j in this row + for j in range(N): + if i == j: + continue + obs_val = stat_obs[i, j] + # Empirical p-value + pval = np.mean(null_vals_row >= obs_val) + E_pval[i, j] = pval + + # Apply FDR for this row + mask = np.ones(N, dtype=bool) + mask[i] = False + #...... per‑row FDR selection, not global FDR + reject, _ = fdrcorrection(E_pval[i, mask], alpha=alpha) + E_edge_mask[i, mask] = reject + + # Count selected edges for this row + edges_per_row.append(int(reject.sum())) + + # Additional statistics + edges_per_row = np.array(edges_per_row) + + # Count zeros in real data matrices + total_elements = A_real.size + zero_elements = np.sum(A_real == 0.0) + sparsity_A_real = zero_elements / total_elements + + # Compute sparsity of W_edge_mask (fraction of zeros in off-diagonal elements) + off_diag_elements = N * (N - 1) # Total off-diagonal elements + selected_edges = int(E_edge_mask.sum()) - N # Subtract diagonal (always True) + sparsity_E_mask = 1.0 - (selected_edges / off_diag_elements) + + # Compute pooled values from A_shuf dimensions: A_shuf shape is (Kshuf, N, N) + # For each row i, we pool A_shuf[:, i, :].ravel() which gives Kshuf*N values + pooled_values_per_row = Kshuf * N + pooled_values_total = N * pooled_values_per_row # N rows total + + summary = { + "alpha": alpha, + "num_bootstraps_real": Kreal, + "num_bootstraps_shuf": Kshuf, + "matrix_size": N, + "num_significant_edges": int(E_edge_mask.sum()) - N, # Exclude diagonal + "pooled_values_per_row": int(pooled_values_per_row), + "pooled_values_total": int(pooled_values_total), + "zero_elements_A_lasso": int(zero_elements), + "total_elements_A_lasso": int(total_elements), + "sparsity_A_lasso": float(sparsity_A_real), + "sparsity_E_mask": float(sparsity_E_mask), + "edges_per_row_avg": float(edges_per_row.mean()), + "edges_per_row_std": float(edges_per_row.std()), + "edges_per_row_min": int(edges_per_row.min()), + "edges_per_row_max": int(edges_per_row.max()), + } + + return E_edge_mask, E_pval, summary + + +def load_bootstrap_data(dataName, dataPath, K, verb=1): + """ + Load all bootstrap data from real and shuffled files. + + Parameters + ---------- + dataName : str + Base name for the dataset + dataPath : str + Path to the data directory + K : list + Number of bootstraps [Kreal, Kshuf] + verb : int + Verbosity level + + Returns + ------- + A_edges_real_list : list + List of real edge matrices (off-diagonal only) + A_edges_shuf_list : list + List of shuffled edge matrices (off-diagonal only) + A_real_list : list + List of full real A matrices + B_real_list : list + List of real B vectors + output_meta : dict + Metadata from first bootstrap file + """ + + A_edges_real_list = [] + A_edges_shuf_list = [] + A_real_list = [] + B_real_list = [] + + Kreal, Kshuf = K + + # Load real data + for k in range(Kreal): + real_file = os.path.join(dataPath, f"{dataName}-boots{k}.lassoFit.npz") + data_real, meta_real = read_data_npz(real_file, verb=(k==0 and verb>=2)) + A_real = data_real['A_lasso'] + B_real = data_real['B_lasso'] + + # Store full matrices for averaging + A_real_list.append(A_real) + B_real_list.append(B_real) + + # Store edges for FDR analysis + edges_real = A_real.copy() + np.fill_diagonal(edges_real, 0.0) + A_edges_real_list.append(edges_real) + + # Store metadata from first bootstrap for output + if k == 0: + output_meta = meta_real + output_big1=data_real + + # Load shuffled data + for k in range(Kshuf): + shuf_file = os.path.join(dataPath, f"{dataName}-desync{k}.lassoFit.npz") + data_shuf, _ = read_data_npz(shuf_file, verb=(k==0 and verb>=2)) + A_shuf = data_shuf['A_lasso'] + edges_shuf = A_shuf.copy() + np.fill_diagonal(edges_shuf, 0.0) + + A_edges_shuf_list.append(edges_shuf) + + if verb >= 1: + print(f"Loaded {Kreal} real and {Kshuf} shuffled bootstrap files") + + return A_edges_real_list, A_edges_shuf_list, A_real_list, B_real_list, output_meta,output_big1 + + +def load_auxiliary_plotting_data(fitMD, dataName, dataPath): + """ + Load auxiliary data needed for plotting (spike data, ground truth, metadata). + + Parameters + ---------- + fitMD : dict + Fit metadata dictionary + maskMD : dict + Mask metadata dictionary + dataName : str + Dataset name for short_name + dataPath : str + Path to data directory + alpha : float + FDR alpha value for metadata + + Returns + ------- + spikeD : dict + Spike data dictionary + trueD : dict or None + Ground truth data (None if not simDale) + MD : dict + Combined metadata for plotting + """ + + # Load spike data for frequency sorting + spikeF = fitMD['fit_lasso']['lassoFit_input_name'] + inpPath= fitMD['fit_lasso']['lassoFit_input_path'] + spikesFF = os.path.join(inpPath, f"{spikeF}.spikes.npz") + + spikeD, spikeMD = read_data_npz(spikesFF, verb=0) + + # Load ground truth data if simulated Dale data + trueD = None + if fitMD['data_type']=='simDale': + truthFF = os.path.join(inpPath, f"{spikeF}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=0) + # Combine metadata just for plotter + MD = {**fitMD, **trueMD, 'short_name': dataName} + else: + MD = {**fitMD, 'short_name': dataName} + + #XMD.update(maskMD) + + return spikeD, trueD, MD diff --git a/causal_net/stationaryLag1_ver6_lasso/UtilTorch.py b/causal_net/stationaryLag1_ver6_lasso/UtilTorch.py new file mode 100644 index 00000000..1ffa1dc7 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/UtilTorch.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +""" +PyTorch utilities for GPU setup, data preprocessing and model training. + +This module provides essential utilities for PyTorch-based neural network training, +specifically optimized for Poisson GLM fitting with distributed GPU support. +Main functionality includes: +- GPU availability checking and device configuration +- Data preprocessing with time decorrelation and shuffling options +- Custom dataset classes for paired neural data (Y_prev, Y_curr) +- Distributed training utilities with Poisson loss optimization +- Learning rate scheduling and training loop management + +Designed to work with multi-GPU setups using DistributedDataParallel +for efficient training of large-scale neural connectivity models. +""" + +import torch +import numpy as np +import time +import torch.optim as optim +import torch.distributed as dist +from torch.utils.data import Dataset, DataLoader +from torch.utils.data.distributed import DistributedSampler +from PoissonGLModel import poisson_nll_loss + +def check_gpu_availability(): + """Check GPU availability and set up device configuration.""" + if not torch.cuda.is_available(): + raise RuntimeError("This script requires a CUDA-enabled GPU environment.") + device = torch.device('cuda:0') + torch.cuda.set_device(0) + print(f"Using device: {torch.cuda.get_device_name(0)}") + return device + +def preprocess_data(Y, args): + import random + Nt, Nn = Y.shape + + # Apply time decorrelation if requested + if args.desyncTime: + if args.rank==0: print("\n=== Applying Time Decorrelation, it shifts time for each neuron ===") + seed = int(time.time() * 1000) % 1000000 + np.random.seed(seed) + shift_amounts = np.random.randint(1, Nt//4, size=Nn, dtype=np.int32) + if args.rank==0: print('Generated random shifts with seed=%d'%(seed),shift_amounts[:10],'...',flush=True) + Y_shifted = np.zeros_like(Y) + for neuron_idx in range(Nn): + shift_amount = int(shift_amounts[neuron_idx]) + Y_shifted[:, neuron_idx] = np.roll(Y[:, neuron_idx], shift_amount) + if args.rank==0: print(f"Applied time shifts, destroys temporal correlations between neurons") + Y = Y_shifted + + # Create consecutive pairs + max_pairs = Nt - 1 + num_samples = args.num_samples + if num_samples is None or num_samples > max_pairs: + num_samples = max_pairs + XY = np.stack([Y[:num_samples], Y[1:num_samples + 1]], axis=1) + + # Apply random data dropping if requested + if args.dropDataFrac > 0: + n_pairs = XY.shape[0] + drop_seed = (int(time.time() * 1000) + np.random.randint(0, 1000)) % (2**32) + np.random.seed(drop_seed) + random.seed(drop_seed) + n_keep = int(n_pairs * (1.0 - args.dropDataFrac)) + keep_indices = np.random.choice(n_pairs, size=n_keep, replace=False) + keep_indices = np.sort(keep_indices) + XY = XY[keep_indices] + if args.rank==0: print(f"Dropped {args.dropDataFrac:.1%} of data, keeping {XY.shape[0]} samples (seed={drop_seed})") + + return XY + + +def train_Poisson_model(model, device, train_loader, n_epochs, lr, L1_alpha=0.0, use_scheduler=False, firing_rates=None, train_sampler=None, print_every=20): + use_fused = (isinstance(device, torch.device) and device.type=='cuda' and torch.cuda.is_available()) + assert use_fused + optimizer = optim.Adam(model.parameters(), lr=lr, fused=True) + mdl = model.module if hasattr(model, 'module') else model + scheduler = optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.1, total_iters=n_epochs) if use_scheduler else None + + diag_mask = torch.eye(mdl.n_neurons, device=device).bool() + L1_weight_matrix = torch.ones(mdl.n_neurons, mdl.n_neurons, device=device) + L1_weight_matrix[diag_mask] = 0.0 + + firing_rates_tensor = torch.tensor(firing_rates, dtype=torch.float32, device=device) if firing_rates is not None else None + + train_losses_w_L1, train_losses_wo_L1, learning_rates = [], [], [] + train_epochs = [] + start_time = time.time() + + for epoch in range(n_epochs): + if train_sampler is not None: train_sampler.set_epoch(epoch) + model.train() + train_loss_w_L1 = 0 + train_loss_wo_L1 = 0 + for Y_prev, Y_curr in train_loader: + Y_prev, Y_curr = Y_prev.float().to(device, non_blocking=True), Y_curr.float().to(device, non_blocking=True) + optimizer.zero_grad(set_to_none=True) + spikes = model(Y_prev) + base_loss = poisson_nll_loss(spikes, Y_curr, firing_rates_tensor) + loss_with_L1 = base_loss.clone() + if L1_alpha > 0: + loss_with_L1 += L1_alpha * torch.mean(torch.abs(mdl.A) * L1_weight_matrix) + loss_with_L1.backward() + optimizer.step() + train_loss_w_L1 += loss_with_L1.item() + train_loss_wo_L1 += base_loss.item() + + # record train losses each epoch + train_losses_w_L1.append(train_loss_w_L1 / len(train_loader)) + train_losses_wo_L1.append(train_loss_wo_L1 / len(train_loader)) + train_epochs.append(epoch + 1) + learning_rates.append(optimizer.param_groups[0]['lr']) + + if scheduler: + scheduler.step() + + if (epoch + 1) % print_every == 0 and (not dist.is_initialized() or dist.get_rank()==0): + + print(f"Epoch {epoch+1}/{n_epochs}: Loss_Tot={train_losses_w_L1[-1]:.5g}, only_L1={(train_losses_w_L1[-1]-train_losses_wo_L1[-1]):.4g}, Elapsed={(time.time() - start_time):.1f}s") + + + + return train_losses_w_L1, train_losses_wo_L1, learning_rates, train_epochs + + +class NumpyPairDataset(Dataset): + def __init__(self, X_np, Y_np): + self.X = X_np + self.Y = Y_np + self.n = X_np.shape[0] + def __len__(self): + return self.n + def __getitem__(self, idx): + return torch.from_numpy(self.X[idx]).to(dtype=torch.float32), torch.from_numpy(self.Y[idx]).to(dtype=torch.float32) + +def make_loader(X, Yt, args, is_dist=False, shuffle=True): + dataset = NumpyPairDataset(X, Yt) + if is_dist: + sampler = DistributedSampler(dataset, shuffle=shuffle) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + return DataLoader(dataset, batch_size=max(1, args.batch_size//world_size), sampler=sampler, shuffle=False, drop_last=shuffle, + pin_memory=True, pin_memory_device='cuda', num_workers=8, persistent_workers=True, prefetch_factor=8) + else: + return DataLoader(dataset, batch_size=args.batch_size, shuffle=shuffle, drop_last=shuffle, pin_memory=True, pin_memory_device='cuda', num_workers=8, + persistent_workers=True, prefetch_factor=8) + + diff --git a/causal_net/stationaryLag1_ver6_lasso/Util_pseudospectra.py b/causal_net/stationaryLag1_ver6_lasso/Util_pseudospectra.py new file mode 100755 index 00000000..23fc3e1f --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/Util_pseudospectra.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +import numpy as np +import time + +def create_example_matrices(size=100): + A_base = np.diag(np.linspace(-5, -0.5, size)) + np.diag(np.ones(size - 2) * 4, k=2) + sparsity = 0.15 + random_connections = np.random.randn(size, size) + mask = np.random.rand(size, size) > sparsity + random_connections[mask] = 0 + A_true = A_base + random_connections * 0.1 + + noise = np.random.randn(size, size) * 0.10 + noise_mask = np.random.rand(size, size) > sparsity + noise[noise_mask] = 0 + A_fit = A_true + noise + + return A_true, A_fit + +def compute_pseudospectrum(A, npts, minY): + print(f'computing pseudospectrum A:{A.shape} .... ') + eigs = np.linalg.eigvals(A) + + # Calculate bounds + real_min, real_max = np.real(eigs).min(), np.real(eigs).max() + imag_min, imag_max = np.imag(eigs).min(), np.imag(eigs).max() + + real_pad = (real_max - real_min) * 0.2 + imag_pad = (imag_max - imag_min) * 0.2 + + bbox = [ + real_min - real_pad, + min(real_max + real_pad, 1.0), + imag_min, + imag_max + imag_pad + ] + + x_coords = np.linspace(bbox[0], bbox[1], npts) + y_coords = np.linspace(minY, bbox[3], npts) + #y_coords = np.linspace(bbox[2], bbox[3], npts) + X, Y = np.meshgrid(x_coords, y_coords) + Z = X + 1j * Y + + sigma_grid = np.zeros_like(Z, dtype=float) + I = np.eye(A.shape[0]) + + for i in range(npts): + for j in range(npts): + z = Z[i, j] + s = np.linalg.svd(z * I - A, compute_uv=False) + sigma_grid[i, j] = s[-1] + + return X, Y, sigma_grid, eigs + +def plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin): + levels = np.logspace(-2.5, -0.5, 10) # Use 10 contour levels + + # Plot the normal contour lines + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + + # Create a mask for the area greater than epsMin + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) # Masking areas greater than epsMin + + # Fill the area below epsMin + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + # Scatter plot for eigenvalues + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + # Vertical line at Re=0 + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + # Set titles and labels + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + # Additional formatting + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(bottom=minY) # Clip at minY + + +def main(size=100, minY=-0.9, epsMin=0.032): + A_true, A_fit = create_example_matrices(size) + + fig, axes = plt.subplots(1, 2, figsize=(16, 7)) + + matrices_to_plot = [ + ('True Matrix ($A_{true}$)', A_true), + ('Distorted Matrix ($A_{fit}$) ', A_fit) + ] + + print("Starting pseudospectrum calculations...", minY) + start_time = time.time() + + for i, (title, A) in enumerate(matrices_to_plot): + ax = axes[i] + print("Processing '{}'...".format(title)) + + npts = 80 # Grid resolution + X, Y, sigma_grid, eigs = compute_pseudospectrum(A, npts, minY) + + plot_pseudospectra(X, Y, sigma_grid, eigs, minY, title, ax, epsMin) # Pass epsMin + + total_time = time.time() - start_time + print(f"Calculation finished in {total_time:.2f} seconds.") + + fig.suptitle('Pseudospectrum Comparison (Clipped Imaginary Part)', fontsize=16) + plt.tight_layout(rect=[0, 0.03, 1, 0.95]) + + outFile = 'out/pseudospectra_comparison.png' + plt.savefig(outFile) + print(f"Plot saved as '{outFile}'") + plt.close(fig) + +# --- Main execution --- +if __name__ == '__main__': + + import matplotlib.pyplot as plt + + main(size=50, minY=-0.5, epsMin=0.024) diff --git a/causal_net/stationaryLag1_ver6_lasso/attic/UtilEdgeSelector.py b/causal_net/stationaryLag1_ver6_lasso/attic/UtilEdgeSelector.py new file mode 100644 index 00000000..5c17a8e1 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/attic/UtilEdgeSelector.py @@ -0,0 +1,449 @@ +""" +Edge selection utilities for bootstrap analysis. + +This module provides different statistical methods for selecting +significant edges from bootstrap connectivity matrix samples. +""" + +import numpy as np + +def edge_selector_s2n(A_edges_list, nSig, minW, verb=0): + """ + Select statistically significant edges from bootstrap samples using signal-to-noise ratio. + + Args: + A_edges_list: List of K A_edges arrays (off-diagonal elements) + nSig: Significance threshold (number of standard deviations) + minW: Minimum weight threshold + verb: Verbosity level + + Returns: + W_edges_mean: Mean values of significant edges (others set to 0) + W_edges_std: Std values of significant edges (others set to 0) + selection_stats: Dictionary with detailed selection statistics + """ + if len(A_edges_list) == 0: + raise ValueError("Empty A_edges_list provided") + + # Get shape from first array + shape = A_edges_list[0].shape + K = len(A_edges_list) + + # Initialize output arrays + W_edges_mean = np.zeros(shape) + W_edges_std = np.zeros(shape) + + # Initialize counters + stats = { + 'offdiag': 0, # Total off-diagonal elements processed + 'small_xmean': 0, # Rejected by first filter (abs(xmean) + nSig*xstd < minW) + 'small_nsig': 0, # Rejected by second filter (abs(xmean)/xstd < nSig) + 'acc_positive': 0, # Accepted edges with positive mean + 'acc_negative': 0 # Accepted edges with negative mean + } + + if verb >= 1: + print("=" * 60) + print("EDGE SELECTION: Signal-to-Noise Ratio Method") + print("=" * 60) + print(f"Parameters: nSig={nSig}, minW={minW}") + print(f"Criteria: |mean| + {nSig}*std >= {minW} AND |mean|/std >= {nSig}") + print(f"Processing {shape} edge positions across {K} bootstraps") + print() + + # Loop over all edge positions + for i in range(shape[0]): + for j in range(shape[1]): + # Skip diagonal elements (should already be removed, but safety check) + if i == j: + continue + + # Count off-diagonal elements + stats['offdiag'] += 1 + + # Collect values from all K bootstrap samples + xV = np.array([A_edges[i, j] for A_edges in A_edges_list]) + + # Compute statistics + xmean = np.mean(xV) + xstd = np.std(xV, ddof=1) # Use sample std (N-1 denominator) + + # Apply significance filters + abs_xmean = abs(xmean) + + # Filter 1: Mean + nSig*std must exceed minimum weight + if abs_xmean + nSig * xstd < minW: + stats['small_xmean'] += 1 + continue + + # Filter 2: Signal-to-noise ratio must exceed nSig + if xstd > 0 and abs_xmean / xstd < nSig: + stats['small_nsig'] += 1 + continue + + # Edge passed both filters - store it + W_edges_mean[i, j] = xmean + W_edges_std[i, j] = xstd + + # Count accepted edges by sign + if xmean >= 0: + stats['acc_positive'] += 1 + else: + stats['acc_negative'] += 1 + + if verb >= 3: + print(f"Significant edge ({i},{j}): mean={xmean:.4f}, std={xstd:.4f}, SNR={abs_xmean/xstd:.2f}") + + # Calculate total accepted + num_edges = stats['acc_positive'] + stats['acc_negative'] + + if verb >= 1: + print("RESULTS:") + print(f" Total off-diagonal elements: {stats['offdiag']}") + print(f" Rejected (small magnitude): {stats['small_xmean']} ({100*stats['small_xmean']/stats['offdiag']:.1f}%)") + print(f" Rejected (low significance): {stats['small_nsig']} ({100*stats['small_nsig']/stats['offdiag']:.1f}%)") + print(f" Accepted (positive): {stats['acc_positive']} ({100*stats['acc_positive']/stats['offdiag']:.1f}%)") + print(f" Accepted (negative): {stats['acc_negative']} ({100*stats['acc_negative']/stats['offdiag']:.1f}%)") + print(f" TOTAL SELECTED: {num_edges}/{stats['offdiag']} ({100*num_edges/stats['offdiag']:.1f}%)") + + # Verification + total_processed = stats['small_xmean'] + stats['small_nsig'] + stats['acc_positive'] + stats['acc_negative'] + if total_processed != stats['offdiag']: + print(f" WARNING: Count mismatch! Expected {stats['offdiag']}, got {total_processed}") + print() + + return W_edges_mean, W_edges_std, stats + +def edge_selector_fdr(A_edges_list, alpha=0.05, verb=0): + """ + Select statistically significant edges using row-wise FDR correction. + + Apply FDR correction independently for each neuron (each row of matrix A). + This accounts for different statistical accuracy due to varying firing rates. + + Args: + A_edges_list: List of K A_edges arrays (off-diagonal elements) + alpha: Significance level for FDR correction + verb: Verbosity level + + Returns: + W_edges_mean: Mean values of significant edges (others set to 0) + W_edges_std: Std values of significant edges (others set to 0) + selection_stats: Dictionary with detailed selection statistics + """ + if len(A_edges_list) == 0: + raise ValueError("Empty A_edges_list provided") + + try: + from scipy import stats + from statsmodels.stats.multitest import multipletests + except ImportError: + raise ImportError("FDR method requires scipy and statsmodels. Install with: pip install scipy statsmodels") + + # Get shape from first array + shape = A_edges_list[0].shape + K = len(A_edges_list) + M = shape[0] # Number of neurons + + # Initialize output arrays + W_edges_mean = np.zeros(shape) + W_edges_std = np.zeros(shape) + + # Initialize counters + stats_dict = { + 'offdiag': 0, # Total off-diagonal elements processed + 'tested': 0, # Elements tested (had enough variance for t-test) + 'low_variance': 0, # Skipped due to zero/low variance + 'significant': 0, # Passed FDR correction + 'acc_positive': 0, # Significant edges with positive mean + 'acc_negative': 0, # Significant edges with negative mean + 'neurons_tested': 0, # Number of neurons that had testable edges + 'neurons_with_edges': 0 # Number of neurons that got significant edges + } + + if verb >= 1: + print("=" * 60) + print("EDGE SELECTION: Row-wise FDR Multiple Testing Correction") + print("=" * 60) + print(f"Parameters: alpha={alpha}") + print(f"Method: Benjamini-Hochberg FDR applied independently for each neuron") + print(f"Processing {shape} edge positions across {K} bootstraps") + print(f"Applying FDR independently for each of {M} neurons") + print() + + # Process each neuron (row) independently + for neuron_i in range(M): + if verb >= 3: + print(f"Processing neuron {neuron_i+1}/{M}") + + # Collect data for this neuron's incoming connections + neuron_positions = [] + neuron_means = [] + neuron_stds = [] + neuron_p_values = [] + + # Collect all incoming edges for this neuron (row neuron_i) + for j in range(M): + if neuron_i == j: # Skip diagonal + continue + + stats_dict['offdiag'] += 1 + + # Collect values from all K bootstrap samples + xV = np.array([A_edges[neuron_i, j] for A_edges in A_edges_list]) + + # Compute statistics + xmean = np.mean(xV) + xstd = np.std(xV, ddof=1) + + neuron_positions.append((neuron_i, j)) + neuron_means.append(xmean) + neuron_stds.append(xstd) + + # Perform one-sample t-test against null hypothesis (mean = 0) + if xstd > 1e-10 and K > 1: + t_stat, p_val = stats.ttest_1samp(xV, 0) + neuron_p_values.append(p_val) + stats_dict['tested'] += 1 + else: + neuron_p_values.append(1.0) # Non-significant p-value for low variance + stats_dict['low_variance'] += 1 + + # Apply FDR correction for this neuron's edges only + if len(neuron_p_values) > 0: + stats_dict['neurons_tested'] += 1 + + # Count testable edges for this neuron + testable_edges = sum(1 for p in neuron_p_values if p < 1.0) + + if testable_edges > 0: + rejected, p_adj, alpha_sidak, alpha_bonf = multipletests( + neuron_p_values, alpha=alpha, method='fdr_bh' + ) + + neuron_significant_count = 0 + + # Store significant edges for this neuron + for k, (i, j) in enumerate(neuron_positions): + if rejected[k]: + W_edges_mean[i, j] = neuron_means[k] + W_edges_std[i, j] = neuron_stds[k] + stats_dict['significant'] += 1 + neuron_significant_count += 1 + + # Count by sign + if neuron_means[k] >= 0: + stats_dict['acc_positive'] += 1 + else: + stats_dict['acc_negative'] += 1 + + if verb >= 3: + print(f" Significant edge ({i},{j}): mean={neuron_means[k]:.4f}, std={neuron_stds[k]:.4f}, p_adj={p_adj[k]:.4e}") + + if neuron_significant_count > 0: + stats_dict['neurons_with_edges'] += 1 + + if verb >= 3: + print(f" Neuron {neuron_i}: {testable_edges} testable, {neuron_significant_count} significant") + + if verb >= 1: + print("RESULTS:") + print(f" Total off-diagonal elements: {stats_dict['offdiag']}") + print(f" Testable (sufficient var): {stats_dict['tested']} ({100*stats_dict['tested']/stats_dict['offdiag']:.1f}%)") + print(f" Low variance (skipped): {stats_dict['low_variance']} ({100*stats_dict['low_variance']/stats_dict['offdiag']:.1f}%)") + print(f" Accepted (positive): {stats_dict['acc_positive']} ({100*stats_dict['acc_positive']/stats_dict['offdiag']:.1f}%)") + print(f" Accepted (negative): {stats_dict['acc_negative']} ({100*stats_dict['acc_negative']/stats_dict['offdiag']:.1f}%)") + print(f" TOTAL SELECTED: {stats_dict['significant']}/{stats_dict['offdiag']} ({100*stats_dict['significant']/stats_dict['offdiag']:.1f}%)") + print(f" Neurons tested: {stats_dict['neurons_tested']} out of {M}") + print(f" Neurons with edges: {stats_dict['neurons_with_edges']} ({100*stats_dict['neurons_with_edges']/max(1,stats_dict['neurons_tested']):.1f}% of tested)") + + # Verification + print(f" Verification: {stats_dict['offdiag']} = {stats_dict['tested']} testable + {stats_dict['low_variance']} low-var") + print() + + return W_edges_mean, W_edges_std, stats_dict + +def edge_selector_fdr_effect_size(A_edges_list, alpha=0.05, effect_size_factor=1.5, verb=0): + """ + Select statistically significant edges using row-wise FDR with automatic minimum effect size. + + Apply FDR correction independently for each neuron (each row of matrix A) combined + with data-driven minimum effect size threshold computed per row. + + Args: + A_edges_list: List of K A_edges arrays (off-diagonal elements) + alpha: Significance level for FDR correction + effect_size_factor: Factor for effect size threshold (median + factor*MAD). Higher = more restrictive. + verb: Verbosity level + + Returns: + W_edges_mean: Mean values of significant edges (others set to 0) + W_edges_std: Std values of significant edges (others set to 0) + selection_stats: Dictionary with detailed selection statistics + """ + if len(A_edges_list) == 0: + raise ValueError("Empty A_edges_list provided") + + try: + from scipy import stats + from statsmodels.stats.multitest import multipletests + except ImportError: + raise ImportError("FDR method requires scipy and statsmodels. Install with: pip install scipy statsmodels") + + # Get shape from first array + shape = A_edges_list[0].shape + K = len(A_edges_list) + M = shape[0] # Number of neurons + + # Initialize output arrays + W_edges_mean = np.zeros(shape) + W_edges_std = np.zeros(shape) + + # Initialize counters + stats_dict = { + 'offdiag': 0, # Total off-diagonal elements processed + 'tested': 0, # Elements tested (had enough variance for t-test) + 'low_variance': 0, # Skipped due to zero/low variance + 'significant_fdr': 0, # Passed FDR correction + 'effect_filtered': 0, # Removed by effect size filter + 'significant': 0, # Final count after both filters + 'acc_positive': 0, # Significant edges with positive mean + 'acc_negative': 0, # Significant edges with negative mean + 'neurons_tested': 0, # Number of neurons that had testable edges + 'neurons_with_edges': 0 # Number of neurons that got significant edges + } + + if verb >= 1: + print("=" * 60) + print("EDGE SELECTION: Row-wise FDR + Minimum Effect Size") + print("=" * 60) + print(f"Parameters: alpha={alpha}, effect_size_factor={effect_size_factor}") + print(f"Method: Benjamini-Hochberg FDR + auto effect size per neuron") + print(f"Processing {shape} edge positions across {K} bootstraps") + print(f"Applying FDR + effect size independently for each of {M} neurons") + print() + + # Process each neuron (row) independently + for neuron_i in range(M): + if verb >= 3: + print(f"Processing neuron {neuron_i+1}/{M}") + + # Collect data for this neuron's incoming connections + neuron_positions = [] + neuron_means = [] + neuron_stds = [] + neuron_p_values = [] + + # Collect all incoming edges for this neuron (row neuron_i) + for j in range(M): + if neuron_i == j: # Skip diagonal + continue + + stats_dict['offdiag'] += 1 + + # Collect values from all K bootstrap samples + xV = np.array([A_edges[neuron_i, j] for A_edges in A_edges_list]) + + # Compute statistics + xmean = np.mean(xV) + xstd = np.std(xV, ddof=1) + + neuron_positions.append((neuron_i, j)) + neuron_means.append(xmean) + neuron_stds.append(xstd) + + # Perform one-sample t-test against null hypothesis (mean = 0) + if xstd > 1e-10 and K > 1: + t_stat, p_val = stats.ttest_1samp(xV, 0) + neuron_p_values.append(p_val) + stats_dict['tested'] += 1 + else: + neuron_p_values.append(1.0) # Non-significant p-value for low variance + stats_dict['low_variance'] += 1 + + # Apply FDR correction for this neuron's edges only + if len(neuron_p_values) > 0: + stats_dict['neurons_tested'] += 1 + + # Count testable edges for this neuron + testable_edges = sum(1 for p in neuron_p_values if p < 1.0) + + if testable_edges > 0: + # Stage 1: FDR correction + rejected, p_adj, alpha_sidak, alpha_bonf = multipletests( + neuron_p_values, alpha=alpha, method='fdr_bh' + ) + + # Stage 2: Compute minimum effect size for this neuron + neuron_means_array = np.array(neuron_means) + fdr_significant_mask = np.array(rejected) + + if np.any(fdr_significant_mask): + # Get effect sizes of FDR-significant edges for this neuron + significant_effects = np.abs(neuron_means_array[fdr_significant_mask]) + + if len(significant_effects) > 3: # Need enough edges for percentile + # Use median + factor*MAD as minimum effect size for this neuron + median_effect = np.median(significant_effects) + mad_effect = np.median(np.abs(significant_effects - median_effect)) + min_effect_size = median_effect + effect_size_factor * mad_effect + else: + # For few edges, use 75th percentile + min_effect_size = np.percentile(significant_effects, 75) if len(significant_effects) > 1 else significant_effects[0] * 0.8 + + if verb >= 3: + print(f" Neuron {neuron_i}: auto effect size threshold = {min_effect_size:.4f}") + else: + min_effect_size = 0.0 # No significant edges to filter + + neuron_fdr_count = 0 + neuron_final_count = 0 + + # Apply both FDR and effect size filters + for k, (i, j) in enumerate(neuron_positions): + if rejected[k]: # Passed FDR + stats_dict['significant_fdr'] += 1 + neuron_fdr_count += 1 + + # Check effect size + if abs(neuron_means[k]) >= min_effect_size: + W_edges_mean[i, j] = neuron_means[k] + W_edges_std[i, j] = neuron_stds[k] + stats_dict['significant'] += 1 + neuron_final_count += 1 + + # Count by sign + if neuron_means[k] >= 0: + stats_dict['acc_positive'] += 1 + else: + stats_dict['acc_negative'] += 1 + + if verb >= 3: + print(f" Final edge ({i},{j}): mean={neuron_means[k]:.4f}, std={neuron_stds[k]:.4f}, p_adj={p_adj[k]:.4e}") + else: + stats_dict['effect_filtered'] += 1 + + if neuron_final_count > 0: + stats_dict['neurons_with_edges'] += 1 + + if verb >= 3: + print(f" Neuron {neuron_i}: {testable_edges} testable, {neuron_fdr_count} FDR-sig, {neuron_final_count} final") + + if verb >= 1: + print("RESULTS:") + print(f" Total off-diagonal elements: {stats_dict['offdiag']}") + print(f" Testable (sufficient var): {stats_dict['tested']} ({100*stats_dict['tested']/stats_dict['offdiag']:.1f}%)") + print(f" Low variance (skipped): {stats_dict['low_variance']} ({100*stats_dict['low_variance']/stats_dict['offdiag']:.1f}%)") + print(f" Significant (FDR only): {stats_dict['significant_fdr']} ({100*stats_dict['significant_fdr']/stats_dict['offdiag']:.1f}%)") + print(f" Effect size filtered: {stats_dict['effect_filtered']} ({100*stats_dict['effect_filtered']/stats_dict['offdiag']:.1f}%)") + print(f" Accepted (positive): {stats_dict['acc_positive']} ({100*stats_dict['acc_positive']/stats_dict['offdiag']:.1f}%)") + print(f" Accepted (negative): {stats_dict['acc_negative']} ({100*stats_dict['acc_negative']/stats_dict['offdiag']:.1f}%)") + print(f" TOTAL SELECTED: {stats_dict['significant']}/{stats_dict['offdiag']} ({100*stats_dict['significant']/stats_dict['offdiag']:.1f}%)") + print(f" Neurons tested: {stats_dict['neurons_tested']} out of {M}") + print(f" Neurons with edges: {stats_dict['neurons_with_edges']} ({100*stats_dict['neurons_with_edges']/max(1,stats_dict['neurons_tested']):.1f}% of tested)") + + # Verification + print(f" Verification: {stats_dict['offdiag']} = {stats_dict['tested']} testable + {stats_dict['low_variance']} low-var") + print() + + return W_edges_mean, W_edges_std, stats_dict diff --git a/causal_net/stationaryLag1_ver6_lasso/attic/fit_regressPoisson.py b/causal_net/stationaryLag1_ver6_lasso/attic/fit_regressPoisson.py new file mode 100755 index 00000000..a3108b57 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/attic/fit_regressPoisson.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +import os +import time +import argparse +import numpy as np +import torch.optim as optim +import torch +from torch.utils.data import TensorDataset, DataLoader + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from PoissonGLModel import PoissonGLModel, poisson_nll_loss + +from UtilTorch import check_gpu_availability, preprocess_data, train_Poisson_model, NumpyPairDataset, make_loader +from pprint import pprint + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("-n", "--dataName", type=str, required=True, help="fitA NPZ file name") + parser.add_argument("--dataPath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="path to input and output files") + parser.add_argument("--num_samples", type=int, default=None, help="limit number of samples, None=all") + parser.add_argument("--n_epochs", type=int, default=None) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--lr", type=float, default=2e-4, help="learning rate") + parser.add_argument("--fitName", type=str, default=None, help="base name for output files") + parser.add_argument("--noise_scale", type=float, default=0.2, help="Random noise scale for A,B seed values (0.0=no noise, 0.1=moderate noise)") + + args = parser.parse_args() + args.desync_time=0 + + # DDP init + is_dist = (int(os.environ.get('WORLD_SIZE', '1')) > 1) or ('LOCAL_RANK' in os.environ) or ('RANK' in os.environ) + if is_dist: + import torch.distributed as dist + from torch.nn.parallel import DistributedDataParallel as DDP + from torch.utils.data.distributed import DistributedSampler + dist.init_process_group(backend='nccl') + local_rank = int(os.environ['LOCAL_RANK']) + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + local_rank = 0; rank = 0; world_size = 1 + device = check_gpu_availability() + if rank==0: + print("world_size=%d" % (world_size)) + print("Initial configuration:", vars(args)) + + # enable fast matmul paths + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision('high') + gpu_name = torch.cuda.get_device_name(device) if isinstance(device, torch.device) and device.type=='cuda' else str(device) + print("[rank %d] Using device %s : %s" % (rank, str(device), gpu_name)) + + fit1FF = os.path.join(args.dataPath, f"{args.dataName}.edgeMask.npz") + fitD, fitMD = read_data_npz(fit1FF,verb=rank==0) + A_init = fitD['A_init'] + B_init = fitD['B_init'] + + if rank==0: pprint(fitMD) + fmd=fitMD['fit_lasso'] + # Inherit hyperparams from stageA if not provided + if args.n_epochs is None: + args.n_epochs = fmd['n_epochs'] + if args.batch_size is None: + args.batch_size = fmd['batch_size'] + + if rank==0: print("Effective configuration:", vars(args)) + + spike_data_name = fmd['lassoFit_input_name'] + spikesFF = os.path.join(args.dataPath, f"{spike_data_name}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF,verb=rank==0) + dataYield, dataRates = spikeD['spikes'], np.clip(spikeD['single_rates'], 0.1, 40.0) + T, M = dataYield.shape + step_size = spikeMD['time_step_sec'] + + if args.fitName is None: + import hashlib + hash_str = hashlib.md5(os.urandom(32)).hexdigest()[:6] + args.fitName = f"{args.dataName}-{hash_str}" + + X_np, Yt_np = preprocess_data(dataYield, args) + n_pairs = X_np.shape[0] + + assert n_pairs >= args.batch_size, f"ERROR: Not enough samples ({n_pairs}) for batch size ({args.batch_size})." + + train_loader = make_loader(X_np, Yt_np, args, is_dist=is_dist) + + if rank==0: + print(f"Loaded T={T}, M={M}, using {n_pairs} pairs (all for training), world_size={world_size}, per_gpu_bs={args.batch_size//max(1,world_size)}") + + # --- Set up model for stage B --- + trainable_mask = fitD['mask.lasso.pass'] + if rank==0: + print(f"Trainable mask: {trainable_mask.sum()} out of {trainable_mask.size} A elements are trainable") + print(f"Trainable fraction: {trainable_mask.mean():.3f}") + + base_model = PoissonGLModel(M, A_init=A_init, B_init=B_init, trainable_mask=trainable_mask, noise_scale=args.noise_scale).to(device) + model = DDP(base_model, device_ids=[local_rank]) if is_dist else base_model + mdl = model.module if hasattr(model,'module') else model + + # Debug: check if noise was actually applied + if args.noise_scale > 0: + if rank==0: print(f"Checking noise application...") + A_reconstructed = mdl.A.detach().cpu().numpy() + A_diff = np.abs(A_reconstructed - A_init).max() + B_reconstructed = mdl.B.detach().cpu().numpy() + B_diff = np.abs(B_reconstructed - B_init).max() + print(f"Max A difference from init: {A_diff:.6f}") + print(f"Max B difference from init: {B_diff:.6f}") + + # Debug: count trainable parameters + total_params = sum(p.numel() for p in mdl.parameters()) + trainable_params = sum(p.numel() for p in mdl.parameters() if p.requires_grad) + expected_trainable = trainable_mask.sum() + M # masked A elements + B vector + if rank==0: + print(f"Total parameters: {total_params}, PyTorch trainable: {trainable_params}, Expected trainable: {expected_trainable}") + + if hasattr(mdl, 'is_stage_b') and mdl.is_stage_b: + if rank==0: print(f"Stage B efficient parameterization: C tensor has {mdl.C.numel()} parameters") + else: + if rank==0: print("Stage A standard parameterization") + + # Debug: check initial loss and parameter values + model.train() + + # Check initial loss before training + with torch.no_grad(): + for Y_prev, Y_curr in train_loader: + Y_prev, Y_curr = Y_prev.float().to(device), Y_curr.float().to(device) + spikes = model(Y_prev) + initial_loss = poisson_nll_loss(spikes, Y_curr, torch.tensor(dataRates, dtype=torch.float32, device=device)) + if rank==0: print(f"Initial loss with noise: {initial_loss:.6f}") + break + + # Check if C parameters actually changed from Stage A + if hasattr(mdl, 'C'): + C_values = mdl.C.detach().cpu() + if rank==0: print(f"C parameter stats: min={C_values.min():.6f}, max={C_values.max():.6f}, std={C_values.std():.6f}") + + # Check gradients after one step + for Y_prev, Y_curr in train_loader: + Y_prev, Y_curr = Y_prev.float().to(device), Y_curr.float().to(device) + spikes = model(Y_prev) + loss = poisson_nll_loss(spikes, Y_curr, torch.tensor(dataRates, dtype=torch.float32, device=device)) + loss.backward() + + # Check gradients for efficient parameterization + if hasattr(mdl, 'C') and mdl.C.grad is not None: + C_grad_norm = mdl.C.grad.norm().item() + C_grad_max = mdl.C.grad.abs().max().item() + if rank==0: print(f"C gradient: norm={C_grad_norm:.6f}, max={C_grad_max:.6f}") + break + + model.zero_grad() # Clear gradients before actual training + + start_time = time.time() + losses_total, _, learning_rates, losses_epochs = train_Poisson_model( + model, device, train_loader, args.n_epochs, lr=args.lr, L1_alpha=0.0, use_scheduler=True, firing_rates=dataRates, + train_sampler=train_loader.sampler if is_dist else None + ) + total_time = time.time() - start_time + if rank==0: + print(f"Training completed in {total_time:.1f} seconds") + + if rank==0: + mdl = model.module if hasattr(model,'module') else model + bigD = { 'A_regress': mdl.A.detach().cpu().numpy(), 'B_regress': mdl.B.detach().cpu().numpy(), 'losses_total': np.array(losses_total), 'losses_epochs': np.array(losses_epochs, dtype=np.int32), 'learning_rates': np.array(learning_rates), 'firing_rates': dataRates } + metaD = { 'regressFit_output_name': args.fitName, 'regressFit_input_name': args.dataName, + 'regressFit_input_name': args.dataName, + 'regressFit_output_name': args.fitName, + 'batch_size': args.batch_size, 'num_samples_used': n_pairs, 'n_epochs': args.n_epochs, 'num_train_samples': n_pairs, 'learning_rate': args.lr, 'step_size': step_size, 'training_time_sec': total_time, 'num_neurons': M, 'desync_time': args.desync_time } + fitMD['fit_regress']=metaD + fitMD['data_type']=spikeMD['data_type'] + fitMD['fit_type']='regress' + fitMD['short_name']=args.fitName + fit2FF = os.path.join(args.dataPath, f"{args.fitName}.regressFit.npz") + write_data_npz(bigD, fit2FF, metaD=fitMD) + print('\n ./eval_fitRegress.py --dataName %s -p a ' % (args.fitName)) + + # ensure distributed shutdown to avoid resource leak warning + if is_dist and dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + +if __name__ == "__main__": + main() diff --git a/causal_net/stationaryLag1_ver6_lasso/attic/readExp_npy.py b/causal_net/stationaryLag1_ver6_lasso/attic/readExp_npy.py new file mode 100755 index 00000000..d661a90b --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/attic/readExp_npy.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import sys,os,hashlib +import numpy as np +from pprint import pprint +import argparse +#...!...!.................. +def commandline_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v","--verb",type=int, help="increase debug verbosity", default=1) + parser.add_argument("--expPath",default='/global/cfs/cdirs/m2043/causal_inference/B6J/B6J/250902/M07036/Network',help="raw experimnetal data on CFS") + + parser.add_argument("--sessionName", default='000011/well000',help='raw data session name') + parser.add_argument("--dataPath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for any further data processing") + parser.add_argument("--outName", default=None,help='(optional) output file name - Is it needed?') + + # .... activity speciffic speciffic, + #parser.add_argument('--time_rebin', default=100, type=int, help='rebin of raw time axis') + parser.add_argument('--samp_freq', default=100, type=int, help='sets binning of time axis') + parser.add_argument('--freqRange', default=[0.1,30], type=float, nargs=2,help='rebin of raw time axis') + + args = parser.parse_args() + + for arg in vars(args): + print( 'myArgs:',arg, getattr(args, arg)) + + return args + +#...!...!.................... +def buildBioMeta(args): + pd={} # payload + pd['raw_bioexp_path']=args.expPath + pd['session_name']=args.sessionName + pd['culture_type']='my culture 77' + pd['recording_date']='19630417' + pd['chip_name']='M12345' + pd['run_number']='010203' + pd['well_no']='Well1234' + + sel={'freq_range':args.freqRange} + md={ 'bioexp':pd,'selector':sel} + myHN=hashlib.md5(os.urandom(32)).hexdigest()[:6] + md['hash']=myHN + if args.outName==None: + md['short_name']='%s-%s'%(args.sessionName,md['hash']) + else: + md['short_name']=args.outName + + if args.verb>1: print('\nBMD:');pprint(md) + return md + +def read_spike_npy(md,args): + pmd=md['bioexp'] + inpF=os.path.join(args.expPath,args.sessionName,'spike_times.npy') + print('inpF:',inpF) + assert os.path.exists(inpF) + # Load the dictionary from the .npy file + spike_dict = np.load(inpF, allow_pickle=True).item() + + #print('spike_dict',spike_dict);ok + print('spike_dict',sorted(spike_dict)) + #1raw_sampling_freq=10000 # Hz, number from Roy + #1assert raw_sampling_freq%args.time_rebin==0 + pmd['sampling_freq'] =args.samp_freq + + # neuron ID MEA chip + meaIdL=np.array(sorted(spike_dict)) # here order of feature_id is settled + maxFeat=len(meaIdL) + + if args.verb>1: print('RSN: meaID list:',meaIdL) + pmd['num_feature']=len(meaIdL) + + spikeT={} # spike times + spikeCntL=np.zeros(pmd['num_feature'],dtype=int) # num spikes per neuron + maxTbin=0 + + j=0 + for k in meaIdL: + rec=np.array(spike_dict[k])*args.samp_freq + #print(rec[:100],len(rec)) + spikeT[k]=rec.astype(int) # time-bin may repeat + spikeCntL[j]=len(rec) + j+=1 + if len(rec)==0: + continue + mxTb=np.max(rec) + if maxTbin< mxTb: maxTbin=mxTb + #exit(0) + pmd['num_time_bin']=int(maxTbin)+1 + pmd['max_time']=pmd['num_time_bin']/pmd['sampling_freq'] + chanFreq=spikeCntL/ pmd['max_time'] + rawD={'spikeT':spikeT,'chanFreq':chanFreq,'chanID':meaIdL} + return rawD + +#================================= +# M A I N +#================================= +#================================= +if __name__ == "__main__": + + args=commandline_parser() + np.set_printoptions(precision=3) + bioMD=buildBioMeta(args) + + # read raw data + rawD=read_spike_npy(bioMD,args) + pprint(bioMD) + print('rawD',sorted(rawD)) + print('chanFreq',rawD['chanFreq']) + + # REMAP MATRICES TO FREQUENCY-SORTED ORDER (PRIMARY INDEX) + neur_freqIdx = np.argsort(rawD['chanFreq']) # indices that sort chanFreq by value + + #neur_revFreqIdx = neur_freqIdx.copy() # freq_sorted_position → natural_index (original from estimate_rates) + neur_freqIdx = np.empty(len(neur_revFreqIdx), dtype=int) # natural_index → freq_sorted_position + neur_freqIdx[neur_revFreqIdx] = np.arange(len(neur_revFreqIdx)) + + diff --git a/causal_net/stationaryLag1_ver6_lasso/attic/readExp_pkl.py b/causal_net/stationaryLag1_ver6_lasso/attic/readExp_pkl.py new file mode 100755 index 00000000..fba2057d --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/attic/readExp_pkl.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import sys,os,hashlib +import numpy as np +import pickle +from pprint import pprint +import argparse +#...!...!.................. +def commandline_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v","--verb",type=int, help="increase debug verbosity", default=1) + parser.add_argument("--expPath",required=True,help="raw experimnetal data on CFS") + + parser.add_argument("--sessionName", default='HET_80k_1',help='raw data session name') + parser.add_argument("--dataPath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for any further data processing") + parser.add_argument("--outName", default=None,help='(optional) output file name - Is it needed?') + + # .... activity speciffic speciffic, + parser.add_argument('--time_rebin', default=100, type=int, help='rebin of raw time axis') + parser.add_argument('--freqRange', default=[0.1,30], type=float, nargs=2,help='rebin of raw time axis') + + args = parser.parse_args() + + for arg in vars(args): + print( 'myArgs:',arg, getattr(args, arg)) + + return args + +#...!...!.................... +def buildBioMeta(args): + pd={} # payload + pd['raw_bioexp_path']=args.expPath + pd['session_name']=args.sessionName + pd['culture_type']='my culture 77' + pd['recording_date']='19630417' + pd['chip_name']='M12345' + pd['run_number']='010203' + pd['well_no']='Well1234' + + sel={'freq_range':args.freqRange} + md={ 'bioexp':pd,'selector':sel} + myHN=hashlib.md5(os.urandom(32)).hexdigest()[:6] + md['hash']=myHN + if args.outName==None: + md['short_name']='%s-%s'%(args.sessionName,md['hash']) + else: + md['short_name']=args.outName + + if args.verb>1: print('\nBMD:');pprint(md) + return md + +def read_spike_dict(md,args): + pmd=md['bioexp'] + inpF=os.path.join(args.expPath,args.sessionName,'spike_dict.pkl') + print('inpF:',inpF) + assert os.path.exists(inpF) + # Load the dictionary from the .pkl file + with open(inpF, "rb") as f: + spike_dict = pickle.load(f) + + raw_sampling_freq=10000 # Hz, number from Roy + assert raw_sampling_freq%args.time_rebin==0 + pmd['sampling_freq'] =raw_sampling_freq/args.time_rebin + + # neuron ID MEA chip + meaIdL=np.array(sorted(spike_dict)) # here order of feature_id is settled + maxFeat=len(meaIdL) + + if args.verb>1: print('RSD: meaID list:',meaIdL) + pmd['num_feature']=len(meaIdL) + + spikeT={} # spike times + spikeCntL=np.zeros(pmd['num_feature'],dtype=int) # num spikes per neuron + maxTbin=0 + + j=0 + for k in meaIdL: + rec=np.array(spike_dict[k])/args.time_rebin + spikeT[k]=rec.astype(int) # time-bin may repeat + spikeCntL[j]=len(rec) + j+=1 + if len(rec)==0: + continue + mxTb=np.max(rec) + if maxTbin< mxTb: maxTbin=mxTb + + pmd['num_time_bin']=int(maxTbin)+1 + pmd['max_time']=pmd['num_time_bin']/pmd['sampling_freq'] + chanFreq=spikeCntL/ pmd['max_time'] + rawD={'spikeT':spikeT,'chanFreq':chanFreq,'chanID':meaIdL} + return rawD + +#================================= +# M A I N +#================================= +#================================= +if __name__ == "__main__": + + args=commandline_parser() + np.set_printoptions(precision=3) + bioMD=buildBioMeta(args) + + # read raw data + rawD=read_spike_dict(bioMD,args) + pprint(bioMD) + print('rawD',sorted(rawD)) + diff --git a/causal_net/stationaryLag1_ver6_lasso/attic/select_edges_naive.py b/causal_net/stationaryLag1_ver6_lasso/attic/select_edges_naive.py new file mode 100755 index 00000000..5ecd94c2 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/attic/select_edges_naive.py @@ -0,0 +1,318 @@ +#!/usr/bin/env python3 +""" +select_edges.py - Bootstrap edge selection for Lasso Poisson results + +This script processes multiple bootstrap results to identify statistically +significant edges in the learned connectivity matrix. + +FDR = False Discovery Rate - a statistical method for controlling errors when testing many hypotheses simultaneously. + +Original method: nSig was applied to signal-to-noise ratio of individual edges +FDR method: alpha controls family-wise error rate across all edges simultaneously +So FDR with α=0.05 is actually more stringent than the original method with nSig=2.0 because it accounts for testing thousands of edges at once! +Recommended starting points: +Exploratory: --alpha 0.05 (standard significance level) +Conservative: --alpha 0.01 (1% false positive rate) +Very conservative: --alpha 0.001 (0.1% false positive rate) + +""" + +import os +import argparse +import numpy as np +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from UtilEdgeSelector import edge_selector_s2n, edge_selector_fdr, edge_selector_fdr_effect_size +from PlotterFitEval import Plotter +from UtilDalePoisson import select_edges_from_fitLasso + +def main(): + parser = argparse.ArgumentParser(description="Bootstrap edge selection for Lasso Poisson results") + parser.add_argument("--dataName", type=str, required=True, help="Base name for the dataset") + parser.add_argument("--dataPath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="Path to the data directory") + parser.add_argument("-v", "--verb", type=int, default=1, help="Verbosity level (0=silent, 1=basic, 2=detailed, 3=debug)") + parser.add_argument("--num_bootstraps", type=int, required=True, help="Number of bootstrap files to process") + parser.add_argument("--method", type=str, default="s2n", choices=["s2n", "fdr1", "fdr2"], help="Edge selection method: s2n (signal-to-noise), fdr1 (basic FDR), fdr2 (FDR + effect size)") + parser.add_argument("--nSig", type=float, default=2.0, help="Significance threshold for s2n method (number of standard deviations)") + parser.add_argument("--minW", type=float, default=0.07, help="Minimum weight threshold for s2n method") + parser.add_argument("--alpha", type=float, default=0.05, help="Significance level for FDR methods (e.g., 0.05)") + parser.add_argument("--effect_size_factor", type=float, default=2.0, help="Effect size factor for fdr2 method (median + factor*MAD). Higher = more restrictive ") + parser.add_argument('-A',"--ampl_thres", type=float, default=[0.10],nargs='+', help=" inh< tht0, exct>th1 of accepted off-diagonal edge") + parser.add_argument('-p',"--showPlots", type=str,nargs='+', default="ab", help="Plot types to show: a=structure, b=distributions, c=reconstruction, d=category, d=A-matrix histograms") + parser.add_argument("--outPath", type=str, default=None, help="Output path for plots (defaults to dataPath)") + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + + args = parser.parse_args() + print(vars(args)) + # Validate method-specific parameters + if args.method == "s2n" and (args.nSig <= 0 or args.minW <= 0): + print("Error: s2n method requires positive nSig and minW values") + return 1 + elif args.method in ["fdr1", "fdr2"] and args.alpha <= 0: + print("Error: FDR methods require alpha > 0") + return 1 + + # Process plotting arguments + if len(args.ampl_thres)==1: + args.ampl_thres=[-args.ampl_thres[0],args.ampl_thres[0]] + + if args.outPath is None: + args.outPath = args.dataPath + args.showPlots=''.join(args.showPlots) + + if args.verb >= 1: + print("=" * 60) + print("BOOTSTRAP EDGE SELECTION") + print("=" * 60) + print(f"Data: {args.dataName}") + print(f"Bootstraps: {args.num_bootstraps}") + print(f"Method: {args.method}") + print() + + # Collect data from all bootstrap files + B_lasso_list = [] + A_lasso_list = [] + + if args.verb >= 2: + print("Loading bootstrap files:") + + for k in range(1, args.num_bootstraps + 1): + filename = f"{args.dataName}-boot{k}.lassoFit.npz" + filepath = os.path.join(args.dataPath, filename) + + if args.verb >= 2: + print(f" Loading {filename}") + + # Load data + data, metadata = read_data_npz(filepath, verb=0) + + # Extract arrays + B_lasso_list.append(data['B_lasso']) + A_lasso_list.append(data['A_lasso']) + + # Store metadata from first file + if k == 1: + first_metadata = metadata + + if args.verb >= 1: + print(f"Successfully loaded {len(B_lasso_list)} bootstrap files") + print() + + # Process B_lasso: compute element-wise mean and std + if args.verb >= 2: + print("Processing B_lasso arrays...") + + B_lasso_stack = np.stack(B_lasso_list, axis=0) # Shape: (K, M) + B_lasso_mean = np.mean(B_lasso_stack, axis=0) + B_lasso_std = np.std(B_lasso_stack, axis=0, ddof=1) + + if args.verb >= 2: + print(f"B_lasso shape: {B_lasso_mean.shape}") + print(f"B_lasso mean range: [{B_lasso_mean.min():.4f}, {B_lasso_mean.max():.4f}]") + print() + + # Process A_lasso: split into diagonal and off-diagonal + if args.verb >= 2: + print("Processing A_lasso arrays...") + + A_shape = A_lasso_list[0].shape + M = A_shape[0] + + # Extract diagonal elements from all bootstraps + A_diag_list = [] + A_edges_list = [] + + for A_lasso in A_lasso_list: + # Extract diagonal + A_diag = np.diag(A_lasso) + A_diag_list.append(A_diag) + + # Create off-diagonal matrix (set diagonal to 0) + A_edges = A_lasso.copy() + np.fill_diagonal(A_edges, 0) + A_edges_list.append(A_edges) + + # Process diagonal elements + A_diag_stack = np.stack(A_diag_list, axis=0) # Shape: (K, M) + A_diag_mean = np.mean(A_diag_stack, axis=0) + A_diag_std = np.std(A_diag_stack, axis=0, ddof=1) + + if args.verb >= 2: + print(f"A_diag shape: {A_diag_mean.shape}") + print(f"A_diag mean range: [{A_diag_mean.min():.4f}, {A_diag_mean.max():.4f}]") + + # Process off-diagonal elements with edge selection + if args.verb >= 2: + print(f"Applying edge selection to A_edges (shape: {A_shape})...") + print() + + # Choose edge selection method + if args.method == "s2n": + W_edges_mean, W_edges_std, edge_stats = edge_selector_s2n( + A_edges_list, args.nSig, args.minW, args.verb + ) + method_name = "s2n" + elif args.method == "fdr1": + W_edges_mean, W_edges_std, edge_stats = edge_selector_fdr( + A_edges_list, args.alpha, args.verb + ) + method_name = "fdr1" + elif args.method == "fdr2": + W_edges_mean, W_edges_std, edge_stats = edge_selector_fdr_effect_size( + A_edges_list, args.alpha, args.effect_size_factor, args.verb + ) + method_name = "fdr2" + else: + print(f"Error: Unknown method '{args.method}'") + return 1 + + num_edges = edge_stats['acc_positive'] + edge_stats['acc_negative'] + + # Combine diagonal and off-diagonal results + # Insert diagonal values back into the edge matrices + np.fill_diagonal(W_edges_mean, A_diag_mean) + np.fill_diagonal(W_edges_std, A_diag_std) + + # Prepare output data + output_data = { + 'B_lasso_mean': B_lasso_mean, + 'B_lasso_std': B_lasso_std, + 'W_edges_mean': W_edges_mean, + 'W_edges_std': W_edges_std + } + + # Update metadata with selection parameters + output_metadata = first_metadata + edge_selection_meta = { + 'num_bootstraps': args.num_bootstraps, + 'method': method_name, + 'num_selected_edges': int(num_edges), + 'total_off_diagonal': int(M * M - M), + 'selection_fraction': float(num_edges) / (M * M - M) + } + + # Add method-specific parameters + if args.method == "s2n": + edge_selection_meta['nSig'] = args.nSig + edge_selection_meta['minW'] = args.minW + elif args.method == "fdr1": + edge_selection_meta['alpha'] = args.alpha + elif args.method == "fdr2": + edge_selection_meta['alpha'] = args.alpha + edge_selection_meta['effect_size_factor'] = args.effect_size_factor + + output_metadata['fit_lasso']['edge_selection'] = edge_selection_meta + + # Add detailed edge selection statistics (method-dependent) + edg_select_stats = { + 'offdiag': int(edge_stats['offdiag']), + 'acc_positive': int(edge_stats['acc_positive']), + 'acc_negative': int(edge_stats['acc_negative']) + } + + if args.method == "s2n": + # S2N method counters + edg_select_stats.update({ + 'small_xmean': int(edge_stats['small_xmean']), + 'small_nsig': int(edge_stats['small_nsig']) + }) + elif args.method == "fdr1": + # FDR1 method counters + edg_select_stats.update({ + 'tested': int(edge_stats['tested']), + 'low_variance': int(edge_stats['low_variance']), + 'significant': int(edge_stats['significant']), + 'neurons_tested': int(edge_stats['neurons_tested']), + 'neurons_with_edges': int(edge_stats['neurons_with_edges']) + }) + elif args.method == "fdr2": + # FDR2 method counters (includes effect size filtering) + edg_select_stats.update({ + 'tested': int(edge_stats['tested']), + 'low_variance': int(edge_stats['low_variance']), + 'significant_fdr': int(edge_stats['significant_fdr']), + 'effect_filtered': int(edge_stats['effect_filtered']), + 'significant': int(edge_stats['significant']), + 'neurons_tested': int(edge_stats['neurons_tested']), + 'neurons_with_edges': int(edge_stats['neurons_with_edges']) + }) + + output_metadata['fit_lasso']['edg_select'] = edg_select_stats + + # Save results + output_filename = f"{args.dataName}-select{args.num_bootstraps}.lassoFit.npz" + output_filepath = os.path.join(args.dataPath, output_filename) + + if args.verb >= 1: + print(f"Saving results to: {output_filename}") + print(f"Selected {num_edges} significant edges out of {M*M-M} possible") + print() + + + if args.verb >= 1: + print("=" * 60) + print("BOOTSTRAP EDGE SELECTION COMPLETED") + print("=" * 60) + print(f"Output file: {output_filename}") + print(f"B_lasso: {B_lasso_mean.shape} (mean and std)") + print(f"A_matrix: {W_edges_mean.shape} (selected edges + diagonal)") + print(f"Total significant edges: {num_edges}/{M*M-M} ({100*num_edges/(M*M-M):.1f}%)") + print() + + write_data_npz(output_data, output_filepath, metaD=output_metadata) + + # Prepare plotting data (always executed) + # Create mask data using same structure as eval_fitLasso.py + fitD = output_data # Use our processed output data as fitD + fitMD = output_metadata + + # Rename records so select_edges_from_fitLasso() has the expected names + fitD['A_lasso'] = fitD['W_edges_mean'].copy() + fitD['B_lasso'] = fitD['B_lasso_mean'] + + # Set exactly 0 values to NaN so they are not displayed in plots + fitD['A_lasso'][fitD['A_lasso'] == 0.0] = np.nan + + maskD, maskMD = select_edges_from_fitLasso(fitD, args.ampl_thres) + maskMD['fit_lasso'] = fitMD['fit_lasso'] + + # Load spike data for frequency sorting (get spike file name from first bootstrap) + spikeF = fitMD['fit_lasso']['lassoFit_input_name'] + spikesFF = os.path.join(args.dataPath, f"{spikeF}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF) + + if fitMD['data_type']=='simDale': + # Load ground truth data + truthFF = os.path.join(args.dataPath, f"{spikeF}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF) + # Combine metadata just for plotter + MD = {**fitMD, **trueMD, 'short_name': args.dataName} + else: + MD = {**fitMD, 'short_name': args.dataName} + + MD.update(maskMD) + + # Add edge selection method to metadata for plotting + MD['edge_selection_method'] = method_name + + # Setup plotter + args.prjName = args.dataName + plot = Plotter(args) + + # Generate plots based on showPlots argument + if 'a' in args.showPlots: + # Plot correlation after threshold (requires simDale data type) + assert fitMD['data_type']=='simDale' + plot.correl_after_thresh(trueD,fitD,maskD,MD,figId=1) + + if 'e' in args.showPlots: + plot.experiment_eigen(fitD,MD, figId=5) + + if 'f' in args.showPlots: + plot.freqSortA_histos(fitD, MD, spikeD, figId=2) + + plot.display_all() + + return 0 + +if __name__ == "__main__": + exit(main()) diff --git a/causal_net/stationaryLag1_ver6_lasso/bootsFit.sh b/causal_net/stationaryLag1_ver6_lasso/bootsFit.sh new file mode 100755 index 00000000..2bdbeeed --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/bootsFit.sh @@ -0,0 +1,163 @@ +#!/bin/bash +# +# boostFit.sh - Simplified bootstrap wrapper for distributed Lasso Poisson training +# +# This script runs fitLasso.sh multiple times for bootstrap training + +# Parse arguments with better handling +DATANAME="" # mandatory parameter +NUM_BOOTSTRAPS=6 # default value +SHUFFLE_TIME=false +BOOTS_TAG='s1' +varArgs=() + +# Parse command line arguments +while [[ $# -gt 0 ]]; do + case $1 in + --data_name|--dataName) + DATANAME="$2" + varArgs+=("--dataName" "$2") + shift 2 + ;; + --num_bootstraps) + NUM_BOOTSTRAPS="$2" + shift 2 + ;; + --bootsTag) + BOOTS_TAG="$2" + shift 2 + ;; + --desyncTime) + DESYNC_TIME=true + varArgs+=("--desyncTime") + shift + ;; + *) + # Pass through all other arguments + varArgs+=("$1") + shift + ;; + esac +done + +# Validate mandatory parameter +if [[ -z "$DATANAME" ]]; then + echo "Error: --data_name (or --dataName) is mandatory" + echo "Usage: $0 --data_name [--num_bootstraps ] [--shuffleTime] [other_args...]" + exit 1 +fi +echo ${BOOTS_TAG} + +# Set naming pattern based on shuffleTime flag +FIT_SUFFIX="boots" +if [[ "$DESYNC_TIME" == true ]]; then + FIT_SUFFIX="desync" +fi + +# Convert varArgs array back to string for compatibility +varArgsStr="${varArgs[*]}" + +# Default/fixed arguments (removed --dropDataFrac since it can be passed via command line) +fixArgs="" + +# Validation +if [[ ! -f "./fitLasso4GPU.sh" ]]; then + echo "Error: fitLasso4GPU.sh not found in current directory" + exit 1 +fi + +# Print configuration +echo "=== Bootstrap Lasso Poisson Training ===" +echo "Extracted dataName: $DATANAME" +echo "FitName pattern: ${DATANAME}${BOOTS_TAG}-${FIT_SUFFIX}X (starting from 0)" +echo "Variable args: $varArgsStr" +echo "Default args: $fixArgs" +echo "Bootstraps: $NUM_BOOTSTRAPS" +echo "=========================================" +echo "" + +# Record start time +TOTAL_START_TIME=$(date +%s) + +# Bootstrap loop +for ((k=1; k<=NUM_BOOTSTRAPS; k++)); do + # Construct fitName using extracted dataName and suffix (count from 0) + FITNAME="${DATANAME}${BOOTS_TAG}-${FIT_SUFFIX}$((k-1))" + + echo "=== Bootstrap $k/$NUM_BOOTSTRAPS ===" + echo "fitName: $FITNAME" + echo "Starting at: $(date)" + echo "" + + # Record individual run start time + RUN_START_TIME=$(date +%s) + + # Run fitLasso.sh with simplified approach + ./fitLasso4GPU.sh $fixArgs --fitName "$FITNAME" $varArgsStr + + # Check exit status + EXIT_CODE=$? + + # Record individual run end time + RUN_END_TIME=$(date +%s) + RUN_DURATION=$((RUN_END_TIME - RUN_START_TIME)) + + if [[ $EXIT_CODE -eq 0 ]]; then + echo "" + echo "=== Bootstrap $k/$NUM_BOOTSTRAPS COMPLETED ===" + echo "Duration: ${RUN_DURATION}s ($(date -d@$RUN_DURATION -u +%H:%M:%S))" + echo "Finished at: $(date)" + echo "" + else + echo "" + echo "=== Bootstrap $k/$NUM_BOOTSTRAPS FAILED ===" + echo "Exit code: $EXIT_CODE" + echo "Duration: ${RUN_DURATION}s" + echo "" + + # Ask user if they want to continue + read -p "Continue with remaining bootstraps? (y/n): " -n 1 -r + echo + if [[ ! $REPLY =~ ^[Yy]$ ]]; then + echo "Stopping bootstrap execution." + exit $EXIT_CODE + fi + echo "" + fi + + # Add separator between runs (except after last run) + if [[ $k -lt $NUM_BOOTSTRAPS ]]; then + echo "################################################################" + echo "" + fi +done + +# Record total end time and calculate duration +TOTAL_END_TIME=$(date +%s) +TOTAL_DURATION=$((TOTAL_END_TIME - TOTAL_START_TIME)) +HOURS=$((TOTAL_DURATION / 3600)) +MINUTES=$(((TOTAL_DURATION % 3600) / 60)) +SECONDS=$((TOTAL_DURATION % 60)) + +echo "################################################################" +echo "=== ALL BOOTSTRAPS COMPLETED ===" +echo "Total bootstraps: $NUM_BOOTSTRAPS" +echo "Total time: ${TOTAL_DURATION}s (${HOURS}h ${MINUTES}m ${SECONDS}s)" +echo "Average per bootstrap: $((TOTAL_DURATION / NUM_BOOTSTRAPS))s" +echo "Finished at: $(date)" +echo " ./selectEdges_FDR.py --dataPath \$fitPath --dataName ${DATANAME}${BOOTS_TAG} --num_bootstraps $NUM_BOOTSTRAPS --alphaFDR 1e-3 -p ab " +echo "################################################################" + +echo "" +echo "Bootstrap results saved with fitNames:" +for ((k=1; k<=NUM_BOOTSTRAPS; k++)); do + echo " ${DATANAME}-${FIT_SUFFIX}$((k-1))" +done + + +# ./fit_lassoPoisson.py $fixArgs $varArgs # 1 GPU job, for testing + + +# ./fitLasso4GPU.sh --outPath $fitPath --inpPath $inpPath --dataName daleM140r1Hz --num_epochs 50 --num_samples 100_001 + +# ./bootsFit.sh --outPath $fitPath --inpPath $inpPath --dataName daleM140r1Hz --num_epochs 50 --num_samples 100_001 --dropDataFrac 0.33 --num_bootstraps 1 --bootsTag b1 diff --git a/causal_net/stationaryLag1_ver6_lasso/docs/Readme b/causal_net/stationaryLag1_ver6_lasso/docs/Readme new file mode 100644 index 00000000..f72e8cf0 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/docs/Readme @@ -0,0 +1,4 @@ +pdflatex ver6_dale_lasso.tex +pdflatex ver6_dale_lasso + +open ver6_dale_lasso diff --git a/causal_net/stationaryLag1_ver6_lasso/docs/stationary-ver6_dale_lasso.tex b/causal_net/stationaryLag1_ver6_lasso/docs/stationary-ver6_dale_lasso.tex new file mode 100644 index 00000000..1a284da7 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/docs/stationary-ver6_dale_lasso.tex @@ -0,0 +1,162 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} + +\geometry{margin=0.8in} + +\title{Stationary Lag-1 Poisson GLMs: Dale Simulation, LASSO Fit, and FDR Edge Selection (Ver6)} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +This note documents the \texttt{stationaryLag1\_ver6\_lasso} pipeline: (1)~\texttt{sim\_daleLag1Poisson.py} generates Dale-law spike trains from a stationary lag-1 Poisson GLM; (2)~\texttt{fit\_lassoPoisson.py} fits a fully connected Poisson GLM with $L_1$ regularization on off-diagonal connectivity; (3)~\texttt{selectEdges\_FDR.py} combines bootstrap LASSO estimates with time-desynchronized (null) fits and applies row-wise Benjamini--Hochberg FDR to produce a sparse inferred matrix~$\widehat{A}$. +\end{abstract} + +\tableofcontents +\newpage + +% --------------------------------------------------------------- +\section{Synthetic spikes: \texttt{sim\_daleLag1Poisson.py}} +\label{sec:sim} +% --------------------------------------------------------------- + +\subsection{Connectivity and Dale structure} + +For $N$ neurons with $N_{\mathrm{exc}}$ excitatory and $N_{\mathrm{inh}}=N-N_{\mathrm{exc}}$ inhibitory, target out-degrees $\rho_j$ are drawn uniformly from an interval set by \texttt{--edge\_prob} (fractions of $N$). An initial weight matrix is built in \texttt{gen\_init\_W}: for each presynaptic neuron $j$, $\rho_j$ distinct postsynaptic targets are chosen at random; excitatory rows receive positive weights drawn uniformly in a band, inhibitory rows receive negative weights scaled by a fixed ratio~$g$ (code default $g=2$). The diagonal is set to a negative value (default $-1$) to model self-inhibition. + +\subsection{Stabilization (continuous Lyapunov heuristic)} + +If the maximum real part of the eigenvalue spectrum of~$A$ is non-negative, \texttt{stabilize} iteratively adjusts \emph{inhibitory} entries using gradients derived from continuous Lyapunov equations (see SI of the referenced neuroscience paper in-code) until $\max_i \operatorname{Re}\lambda_i(A) < -\delta$ for a margin~$\delta$ (code uses \texttt{eigenGap}$=0.1$). Small inhibitory weights are zeroed to avoid numerical issues. + +\subsection{Lag-1 Poisson GLM simulation} + +Spike counts $Y_t \in \mathbb{Z}_{\ge 0}^N$ are generated for $t=0,\ldots,T-1$ with bin width $\Delta t$ (\texttt{--step\_size}). Let $A$ be the stabilized Dale matrix and $B$ the bias (log-rate scale). The code uses +\begin{equation} + \eta_t = A\,Y_{t-1} + B,\qquad + \tilde{\eta}_t = \mathrm{clip}(\eta_t,\,-5,\,5),\qquad + \mu_t = \exp(\tilde{\eta}_t)\,\Delta t,\qquad + Y_t \sim \mathrm{Poisson}(\mu_t), +\end{equation} +with $Y_0 \sim \mathrm{Poisson}(\exp(B)\,\Delta t)$. Here $(A\,Y_{t-1})_i = \sum_j A_{ij} Y_{t-1,j}$ matches the implementation \texttt{eta = A @ Y[t-1] + B}. + +Biases $B$ are drawn either uniformly in $\log$ of an idle rate range (\texttt{--idleRate}, default mode) or from a more structured frequency generator (\texttt{--expRate}). + +\subsection{Frequency sorting and outputs} + +Firing-rate statistics are computed with \texttt{UtilDalePoisson.estimate\_rates}. Neurons are then \emph{reordered by decreasing estimated firing rate}. All saved tensors use this \textbf{frequency-sorted} order: the permutation is recorded as \texttt{neur\_freqIdx} (natural $\to$ sorted index) and \texttt{neur\_revFreqIdx} (sorted $\to$ natural). The ground-truth matrix is $A' = P A P^\top$ and $B' = P B$ for the same permutation~$P$. + +\paragraph{\texttt{.simTruth.npz} (arrays).} +\texttt{A\_true}, \texttt{B\_true}, \texttt{neur\_freqIdx}, \texttt{neur\_revFreqIdx}. + +\paragraph{\texttt{.spikes.npz} (arrays).} +\texttt{spikes} (uint8, $T\times N$, frequency-sorted columns), \texttt{single\_rates}, \texttt{sigle\_rates\_var}, \texttt{sigle\_rates\_snr} (note spelling), \texttt{single\_fano\_fact}. + +\paragraph{Metadata.} +\texttt{dale\_conf} (network parameters), \texttt{evol\_conf} (steps, $\Delta t$, rate options), \texttt{short\_name}, \texttt{dale\_simu\_stats}; spike file includes \texttt{time\_step\_sec}, \texttt{data\_type}=\texttt{simDale}, \texttt{var\_time\_window\_sec}. + +% --------------------------------------------------------------- +\section{LASSO Poisson fitting: \texttt{fit\_lassoPoisson.py}} +\label{sec:lasso} +% --------------------------------------------------------------- + +\subsection{Input and design matrix} + +The script reads \texttt{/.spikes.npz}. Let $Y\in\mathbb{R}^{T\times N}$ be the spike matrix (already frequency-sorted). Training pairs are consecutive bins: +\begin{equation} + (x^{(t)}, y^{(t)}) = (Y_t,\, Y_{t+1}),\qquad t=0,\ldots,T-2, +\end{equation} +optionally truncated to \texttt{--num\_samples} pairs. Optional \texttt{--dropDataFrac} randomly subsamples pairs. + +\paragraph{Time desynchronization (null data for later FDR).} +If \texttt{--desyncTime} is used when fitting \emph{separate} runs, each neuron's time series is circularly shifted by an independent random amount, destroying cross-neuron temporal alignment while preserving marginal statistics. Those fits are saved with distinct filenames for use as nulls in Section~\ref{sec:fdr}. + +\subsection{Model and loss} + +\texttt{PoissonGLModel} parameterizes full $A\in\mathbb{R}^{N\times N}$ and $B\in\mathbb{R}^N$. For a batch of previous spikes $Y_{\mathrm{prev}}$ (rows = samples), predicted rates are +\begin{equation} + \mu = \exp\bigl( B^\top + Y_{\mathrm{prev}} A^\top \bigr)\,\Delta t, +\end{equation} +with the linear term clamped to $[-10,10]$ before the exponential (implementation uses \texttt{addmm} with $A^\top$). + +The training loss is a \emph{rate-weighted} Poisson negative log-likelihood (see \texttt{poisson\_nll\_loss}): weights are inversely proportional to per-neuron estimated firing rates (clipped to $[0.1,40]$~Hz from \texttt{single\_rates}), then normalized to mean~1. This down-weights very active neurons relative to a uniform NLL. + +With scalar \texttt{--L1\_alpha} $=\alpha_{\mathrm{L1}}$, the total objective adds mean absolute value of \emph{off-diagonal} entries of~$A$: +\begin{equation} + \mathcal{L} = \mathcal{L}_{\mathrm{NLL}} + \alpha_{\mathrm{L1}}\cdot \frac{1}{N(N-1)}\sum_{i\neq j}|A_{ij}|. +\end{equation} +Optimization uses Adam on all entries of $A$ and~$B$, with optional linear learning-rate decay (\texttt{train\_Poisson\_model}). Multi-GPU training uses PyTorch \texttt{DistributedDataParallel}. + +\subsection{Output file} + +Rank~0 writes \texttt{/.lassoFit.npz} (default \texttt{fitName} is \texttt{-}). Arrays include \texttt{A\_lasso}, \texttt{B\_lasso}, loss curves, and learning rates. Metadata embeds the spike file name/path, hyperparameters (\texttt{batch\_size}, \texttt{num\_epochs}, \texttt{L1\_alpha}, \texttt{step\_size}, etc.) and augments the spike metadata with \texttt{fit\_type}=\texttt{lasso}. + +% --------------------------------------------------------------- +\section{FDR edge selection: \texttt{selectEdges\_FDR.py}} +\label{sec:fdr} +% --------------------------------------------------------------- + +\subsection{Bootstrap inputs} + +The script expects a set of LASSO fits produced offline (e.g.\ bootstrap resamples of the spike train or repeated random seeds). Filenames follow: +\begin{itemize} + \item Real (aligned-time) bootstraps: \texttt{-boots\{k\}.lassoFit.npz} for $k=0,\ldots,K_{\mathrm{real}}-1$. + \item Null / desynchronized-time bootstraps: \texttt{-desync\{k\}.lassoFit.npz} for $k=0,\ldots,K_{\mathrm{desync}}-1$. +\end{itemize} +The counts $K_{\mathrm{real}}$ and $K_{\mathrm{desync}}$ are set by \texttt{--num\_bootstraps} (one integer duplicated for both, or two integers). + +\subsection{Test statistic and null} + +For each edge $(i,j)$, $i\neq j$, the observed statistic is the median of $|A^{(k)}_{ij}|$ across real bootstraps~$k$. For the null, all magnitudes $|\tilde{A}^{(k)}_{i' j}|$ from desync fits for fixed postsynaptic row~$i$ are pooled (all columns~$j$, all~$k$). An empirical $p$-value is +\begin{equation} + p_{ij} = \frac{\#\{\text{null values} \ge \text{observed median}\}}{\#\text{null values}}. +\end{equation} + +\subsection{Row-wise Benjamini--Hochberg and masked estimate} + +For each row~$i$, Benjamini--Hochberg FDR at level $\alpha$ (\texttt{--alphaFDR}, default $0.002$) is applied to $\{p_{ij}\}_{j\neq i}$. This yields a binary mask $M_{ij}$ on off-diagonals. Diagonal entries are always kept in the mask for downstream averaging. + +Across real bootstraps, the script computes elementwise means $\bar{A}$, $\bar{B}$ and standard deviations. The \textbf{reported} connectivity is $\bar{A}$ with non-selected off-diagonal entries zeroed: +\begin{equation} + \widehat{A}_{ij} = \bar{A}_{ij}\, M_{ij},\qquad i\neq j, +\end{equation} +with diagonal unmasked. + +\subsection{Output} + +\texttt{/.FDRselected.npz} contains \texttt{E\_mask}, \texttt{A\_avr}, \texttt{A\_std}, \texttt{B\_avr}, \texttt{B\_std}, \texttt{E\_pval}, \texttt{summary}, and loss arrays copied from the first bootstrap for convenience. Metadata records \texttt{edge\_selector}=\{FDR, $\alpha$\}. + +% --------------------------------------------------------------- +\section{Default hyperparameters (reference)} +\label{sec:defaults} +% --------------------------------------------------------------- + +\begin{table}[h] +\centering +\begin{tabular}{lll} +\toprule +Script & Parameter & Typical default \\ +\midrule +\texttt{sim\_daleLag1Poisson} & \texttt{--step\_size} & $0.01$ s \\ + & \texttt{--spectralR} & $1.0$ \\ + & \texttt{--edge\_prob} & $[0.05,\,0.2]$ \\ +\texttt{fit\_lassoPoisson} & \texttt{--lr} & $10^{-3}$ \\ + & \texttt{--L1\_alpha} & $10^{-3}$ \\ + & \texttt{--batch\_size} & $2048$ \\ + & \texttt{--num\_epochs} & $7$ \\ +\texttt{selectEdges\_FDR} & \texttt{--alphaFDR} & $0.002$ \\ +\bottomrule +\end{tabular} +\caption{Illustrative defaults; see each script's \texttt{argparse} for authoritative values.} +\end{table} + +\end{document} diff --git a/causal_net/stationaryLag1_ver6_lasso/docs/ver6_dale_lasso.pdf b/causal_net/stationaryLag1_ver6_lasso/docs/ver6_dale_lasso.pdf new file mode 100644 index 00000000..faf34867 Binary files /dev/null and b/causal_net/stationaryLag1_ver6_lasso/docs/ver6_dale_lasso.pdf differ diff --git a/causal_net/stationaryLag1_ver6_lasso/eval_fitLasso.py b/causal_net/stationaryLag1_ver6_lasso/eval_fitLasso.py new file mode 100755 index 00000000..ad7a715a --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/eval_fitLasso.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +""" +Evaluation and visualization tool for LASSO Poisson model fitting results. + +This script loads fitted LASSO connectivity matrices from Poisson GLM training +and provides comprehensive evaluation through statistical analysis and plotting. +Main functionality includes: +- Loading fitted model parameters (A, B matrices) from .lassoFit.npz files +- Computing network connectivity statistics and sparsity metrics +- Generating various visualizations (structure plots, distributions, reconstructions) +- Optionally comparing against ground truth for simulated data + +Usage: + ./eval_fitLasso.py --dataName mydata --dataPath /path/to/data/ -p ab +""" + +import numpy as np +import os +import argparse +import sys +from toolbox.Util_NumpyIO import read_data_npz +from PlotterFitEval import Plotter +#from UtilDalePoisson import select_edges_from_fitLasso +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + +from pprint import pprint + +######################### +# MAIN +######################### + +def main(): + parser = argparse.ArgumentParser(description="Evaluate and plot results from fit_poisson.py") + parser.add_argument("--dataName", type=str, default='dale_M120_3M', help="Base name for the dataset") + parser.add_argument("--dataPath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="Path to the data directory") + parser.add_argument('-p',"--showPlots", type=str,nargs='+', default="f", help="Plot types to show: a=structure, b=distributions, c=reconstruction, d=category, d=A-matrix histograms") + parser.add_argument("--outPath", type=str, default=None, help="Output path for plots (defaults to dataPath)") + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument('-v',"--verb", type=int, default=1, help="Verbosity level") + + args = parser.parse_args() + np.set_printoptions(precision=3) + + if args.outPath is None: args.outPath = args.dataPath + args.showPlots=''.join(args.showPlots) + print(vars(args)) + + # Load fit results + fitFF = os.path.join(args.dataPath, f"{args.dataName}.lassoFit.npz") + fitD, fitMD = read_data_npz(fitFF) + + if 0: # patch old data + #pprint(fitMD) + #fitMD['fit_type']='lasso' + fitMD['fit_lasso']['num_epochs']=fitMD['fit_lasso']['n_epochs'] + #pprint(fitMD) + + if args.verb>1: + pprint(fitMD); exit(1) + + # Load spike data for frequency sorting + spikeF = fitMD['fit_lasso']['lassoFit_input_name'] + inpPath= fitMD['fit_lasso']['lassoFit_input_path'] + spikesFF = os.path.join(inpPath, f"{spikeF}.spikes.npz") + + spikeD, spikeMD = read_data_npz(spikesFF) + #pprint(spikeMD) + + if fitMD['data_type']=='simDale': + truthFF = os.path.join(inpPath, f"{spikeF}.simTruth.npz") + trueD,trueMD = read_data_npz(truthFF) + # Combine metadata just for plotter + MD = {**fitMD, **trueMD, 'short_name': args.dataName} #, 'post': vars(args)} + else: + MD = {**fitMD, 'short_name': args.dataName} #, 'post': vars(args)} + + # Setup plotter + args.prjName = args.dataName + plot = Plotter(args) + + if 'a' in args.showPlots: + plot.summary_fitLasso(fitD,MD,figId=1) + + if 'c' in args.showPlots: + plot.freqSortA_histos(fitD, MD, spikeD, figId=2) + + plot.display_all() + +if __name__ == "__main__": + main() diff --git a/causal_net/stationaryLag1_ver6_lasso/eval_fitRegress.py b/causal_net/stationaryLag1_ver6_lasso/eval_fitRegress.py new file mode 100755 index 00000000..b09e820e --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/eval_fitRegress.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 + +import numpy as np +import os +import argparse +import sys +from toolbox.Util_NumpyIO import read_data_npz +from PlotterFitEval import Plotter +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from selectEdges_FDR import print_table_4_Yao +from UtilSelectFDR import eval_tagged_edges_4_simu, load_auxiliary_plotting_data + +from pprint import pprint + +######################### +# MAIN +######################### + +def main(): + parser = argparse.ArgumentParser(description="Evaluate and plot results from fit_poisson.py") + parser.add_argument("--dataName", type=str, default='dale_M120_3M', help="Base name for the dataset") + parser.add_argument("--dataPath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="path to input and output files") + + parser.add_argument('-p',"--showPlots", type=str,nargs='+', default="ab", help="Plot types to show") + parser.add_argument("--outPath", type=str, default=None, help="Output path for plots (defaults to dataPath)") + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument('-v',"--verb", type=int, default=1, help="Verbosity level") + + args = parser.parse_args() + np.set_printoptions(precision=3) + + if args.outPath is None: args.outPath = args.dataPath + args.showPlots=''.join(args.showPlots) + print(vars(args)) + + # Load fit results + fitFF = os.path.join(args.dataPath, f"{args.dataName}.regressFit.npz") + fitD, fitMD = read_data_npz(fitFF) + + if 0: # patch old data + #fitMD['data_type']='simDale' + #fitMD['fit_type']='regress' + fitD['single_rates']=fitD.pop('firing_rates') + + # Load auxiliary data needed for plotting + spikeD, trueD, MD = load_auxiliary_plotting_data(fitMD, args.dataName, args.dataPath) + + if args.verb>1: + pprint(fitMD); exit(1) + + if fitMD['data_type']=='simDale': + # tmp + fitD['A_avr']=fitD['A_regress'] + fitD['B_avr']=fitD['B_regress'] + evalD=eval_tagged_edges_4_simu(fitD,trueD) + print_table_4_Yao(evalD,fitMD) + + # Setup plotter + args.prjName = args.dataName + plot = Plotter(args) + fitType='regress' + + if 'a' in args.showPlots: + plot.summary_fitLasso(fitD,MD,figId=1) + + if 'xa' in args.showPlots: + xx1_fix_fig_a + plot.correl_after_thresh(trueD,fitD,maskD,MD,figId=1) + + if 'b' in args.showPlots: + assert fitMD['data_type']=='simDale' + plot.residuals(evalD,MD,figId=2) + + #plot.slicedA_histos(fitD, MD, spikeD, figId=2) + + if 'c' in args.showPlots: + plot.freqSortA_histos(fitD, MD, spikeD, figId=2) + + if 'xc' in args.showPlots: + plot.residuals(trueD,fitD,maskD,MD,figId=3) + + if 'xd' in args.showPlots: + plot.compare_eigen(trueD,fitD,MD, figId=4) + + if 'xe' in args.showPlots: + plot.experiment_eigen(fitD,MD, figId=5) + + plot.display_all() + + +if __name__ == "__main__": + main() diff --git a/causal_net/stationaryLag1_ver6_lasso/fitLasso4GPU.sh b/causal_net/stationaryLag1_ver6_lasso/fitLasso4GPU.sh new file mode 100755 index 00000000..ed8b44d4 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/fitLasso4GPU.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# +# fitLasso.sh - Simplified wrapper script for distributed Lasso Poisson training +# +# This script runs fit_lassoPoisson.py with optimized multi-GPU settings + +# Capture variable arguments as text +varArgs="$*" + +# Default/fixed arguments +fixArgs=" --batch_size 2048 --lr 1e-3 --L1_alpha 1e-3 " +# --dataPath /pscratch/sd/b/balewski/2025_causalNet_tmp/ + +# GPU configuration +NUM_GPUS=4 +CUDA_DEVICES="0,1,2,3" + +# Print configuration +echo "=== Distributed Lasso Poisson Training ===" +echo "Variable args: $varArgs" +echo "Fixed args: $fixArgs" +echo "GPUs: $NUM_GPUS ($CUDA_DEVICES)" +echo "===========================================" +echo "" + +# Set environment variables and run the training +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export CUDA_VISIBLE_DEVICES=$CUDA_DEVICES + +# Execute the training command +echo "Running: ./fit_lassoPoisson.py $fixArgs $varArgs" +echo "" + +time torchrun --standalone --nproc_per_node=4 ./fit_lassoPoisson.py $fixArgs $varArgs +# ./fit_lassoPoisson.py $fixArgs $varArgs # 1 GPU job, for testing diff --git a/causal_net/stationaryLag1_ver6_lasso/fit_lassoPoisson.py b/causal_net/stationaryLag1_ver6_lasso/fit_lassoPoisson.py new file mode 100755 index 00000000..06264688 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/fit_lassoPoisson.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +""" +Multi-GPU distributed training of Poisson GLM with LASSO regularization. + +This script implements distributed training of a Poisson Generalized Linear Model +for neural connectivity inference with L1 (LASSO) regularization. Key features: +- Multi-GPU support using PyTorch DistributedDataParallel +- LASSO regularization for sparse connectivity estimation +- Poisson negative log-likelihood loss optimized for spike count data +- Support for time decorrelation and data shuffling +- Efficient data loading with distributed sampling +- Automatic model checkpointing and metadata saving + +Usage: + Multi-GPU: OMP_NUM_THREADS=1 CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 ./fit_lassoPoisson.py --dataName mydata --num_epochs 100 + Single GPU: ./fit_lassoPoisson.py --dataName mydata --num_epochs 100 + +Example SLURM execution: + salloc -q interactive -C gpu -t 4:00:00 -A m2043 -N 1 + module load pytorch + OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 ./fit_lassoPoisson.py --dataName daleM600_443813 --num_epochs 5 --dataPath $dataPath +""" + +import os +import time +import random +import argparse +import numpy as np +from pprint import pprint +import torch.optim as optim +import torch +from torch.utils.data import TensorDataset, DataLoader + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from PoissonGLModel import PoissonGLModel, poisson_nll_loss + +from UtilTorch import check_gpu_availability, preprocess_data, train_Poisson_model, NumpyPairDataset, make_loader + +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data.distributed import DistributedSampler +from torch.utils.data import Dataset + +######################### +# MAIN +######################### + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dataName", type=str, default="dale_2aee70") + parser.add_argument("--inpPath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/") + parser.add_argument("--outPath", type=str, default=None) + parser.add_argument("--num_samples", type=int, default=None) + parser.add_argument("--num_epochs", type=int, default=7) + parser.add_argument("--batch_size", type=int, default=2048) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--L1_alpha", type=float, default=1e-3) + parser.add_argument("--fitName", type=str, default=None) + parser.add_argument("--desyncTime", action='store_true', help="If true completely shuffle time axis for input data, independently for all channels") + parser.add_argument("--dropDataFrac", type=float, default=0.0, help="Fraction of training samples to randomly drop per rank (0.0=use all data, 0.3=drop 30%%)") + + args = parser.parse_args() + + # DDP init + is_dist = (int(os.environ.get('WORLD_SIZE', '1')) > 1) or ('LOCAL_RANK' in os.environ) or ('RANK' in os.environ) + #print('is_dist;',is_dist,os.environ.get('WORLD_SIZE', '1'),'RANK' in os.environ) + + if is_dist: + dist.init_process_group(backend='nccl') + local_rank = int(os.environ['LOCAL_RANK']) + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + local_rank = 0; rank = 0; world_size = 1 + device = check_gpu_availability() + args.rank=rank + if rank==0: + if args.inpPath==None: args.inpPath= args.inpPath + print("FitLasso Config:", vars(args)) + print("world_size=%d" % (world_size)) + # enable fast matmul paths + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision('high') + gpu_name = torch.cuda.get_device_name(device) if isinstance(device, torch.device) and device.type=='cuda' else str(device) + print("[rank %d] Using device %s : %s" % (rank, str(device), gpu_name)) + + # --- data loading and preprocessing (rank 0 only) --- + if rank == 0: + spikesFF = os.path.join(args.inpPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=True) + dataYield = spikeD['spikes'] + dataRates = np.clip(spikeD['single_rates'], 0.1, 40.0) + Nt, Nn = dataYield.shape + pprint(spikeMD) + step_size = spikeMD['time_step_sec'] + XY_np = preprocess_data(dataYield, args) + n_pairs = XY_np.shape[0] + print(f"Rank 0 preprocessed data: XY shape={XY_np.shape}, n_pairs={n_pairs}") + else: + XY_np = None + Nn = None + dataRates = None + spikeMD = None + step_size = None + + # --- Broadcast data from rank 0 to all other ranks --- + if is_dist: + Nn_tensor = torch.tensor([Nn if rank==0 else 0], dtype=torch.int32, device='cuda') + dist.broadcast(Nn_tensor, src=0) + Nn = Nn_tensor.item() + if rank == 0: + shape_info = torch.tensor([XY_np.shape[0], XY_np.shape[1], XY_np.shape[2]], dtype=torch.int64, device='cuda') + else: + shape_info = torch.zeros(3, dtype=torch.int64, device='cuda') + dist.broadcast(shape_info, src=0) + if rank != 0: + XY_np = np.zeros((shape_info[0].item(), shape_info[1].item(), shape_info[2].item()), dtype=np.float32) + else: + XY_np = np.ascontiguousarray(XY_np, dtype=np.float32) + XY_tensor = torch.tensor(XY_np, dtype=torch.float32, device='cuda').contiguous() + dist.broadcast(XY_tensor, src=0) + XY_np = XY_tensor.cpu().numpy() + if rank != 0: + dataRates = torch.zeros(Nn, dtype=torch.float32, device='cuda') + step_size_tensor = torch.zeros(1, dtype=torch.float32, device='cuda') + else: + dataRates = torch.tensor(dataRates, dtype=torch.float32, device='cuda') + step_size_tensor = torch.tensor([step_size], dtype=torch.float32, device='cuda') + dist.broadcast(dataRates, src=0) + dist.broadcast(step_size_tensor, src=0) + dataRates = dataRates.cpu().numpy() + step_size = step_size_tensor.item() + if rank != 0: + print(f"[rank {rank}] Received broadcasted data: XY shape={XY_np.shape}") + + # Split XY into X and Y for all ranks + X_np = XY_np[:, 0, :] + Yt_np = XY_np[:, 1, :] + n_pairs = X_np.shape[0] + + assert n_pairs >= args.batch_size, f"ERROR: Not enough samples ({n_pairs}) for batch size ({args.batch_size}) after data dropping." + + train_loader = make_loader(X_np, Yt_np, args, is_dist=is_dist) + + if rank==0: + print(f"Loaded pairs={n_pairs/1000}k, Nn={Nn}, using {n_pairs/1000}k pairs (all for training), world_size={world_size}, per_gpu_bs={args.batch_size//max(1,world_size)}") + + # --- Original training logic from fit_poissonV4.py --- + base_model = PoissonGLModel(Nn).to(device) + model = DDP(base_model, device_ids=[local_rank]) if is_dist else base_model + start_time = time.time() + losses_total, losses_wo_L1, learning_rates, train_epochs = train_Poisson_model( + model, device, train_loader, args.num_epochs, lr=args.lr, L1_alpha=args.L1_alpha, firing_rates=dataRates, use_scheduler=True, + train_sampler=train_loader.sampler if isinstance(train_loader.sampler, DistributedSampler) else None + ) + total_time = time.time() - start_time + if rank==0: + print(f"Training completed in {total_time:.1f} seconds") + + if args.fitName is None: + import string + hash_str = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6)) + fit_core = f"{args.dataName}-{hash_str}" + else: + fit_core = args.fitName + + # saving from fit_Lasso --- + mdl = model.module if hasattr(model,'module') else model + lassoD = { 'A_lasso': mdl.A.detach().cpu().numpy(), 'B_lasso': mdl.B.detach().cpu().numpy(), 'losses_total': np.array(losses_total), 'losses_wo_L1': np.array(losses_wo_L1), 'losses_epochs': np.array(train_epochs, dtype=np.int32), 'learning_rates': np.array(learning_rates), 'single_rates': dataRates } + lassoMD = { 'lassoFit_output_name': fit_core, 'lassoFit_input_name': args.dataName, 'lassoFit_input_path': args.inpPath ,'batch_size': args.batch_size, 'num_samples_used': n_pairs, 'num_epochs': args.num_epochs, 'num_train_samples': n_pairs, 'learning_rate': args.lr, 'L1_alpha': args.L1_alpha, 'step_size': step_size, 'training_time_sec': total_time, 'num_neurons': Nn, 'dropDataFrac': args.dropDataFrac } + spikeMD['fit_type']='lasso' + spikeMD['fit_lasso']=lassoMD + spikeMD['edge_selector']={'selector_type':'None'} + + fitFF = os.path.join(args.outPath, f"{fit_core}.lassoFit.npz") + write_data_npz(lassoD, fitFF, metaD=spikeMD) + + if rank==0: + if spikeMD['data_type']=='simDale': flags=' -p a c ' + else: flags=' -p a c ' + print('\n ./eval_fitLasso.py --dataPath $fitPath --dataName %s %s ' % (fit_core,flags)) + print('\n ./fit_regressPoisson.py --dataName %s ' % (fit_core)) + print(' --dataPath '+args.outPath) + + # ensure distributed shutdown to avoid resource leak warning + if is_dist and dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + +if __name__ == "__main__": + main() diff --git a/causal_net/stationaryLag1_ver6_lasso/fit_regressPoisson.py b/causal_net/stationaryLag1_ver6_lasso/fit_regressPoisson.py new file mode 100755 index 00000000..d3e234fe --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/fit_regressPoisson.py @@ -0,0 +1,196 @@ +#!/usr/bin/env python3 +import os +import time +import argparse +import numpy as np +import torch.optim as optim +import torch +from torch.utils.data import TensorDataset, DataLoader + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from PoissonGLModel import PoissonGLModel, poisson_nll_loss + +from UtilTorch import check_gpu_availability, preprocess_data, train_Poisson_model, NumpyPairDataset, make_loader +from pprint import pprint + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("-n", "--dataName", type=str, required=True, help="fitA NPZ file name") + parser.add_argument("--dataPath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="path to input and output files") + parser.add_argument("--num_samples", type=int, default=None, help="limit number of samples, None=all") + parser.add_argument("--num_epochs", type=int, default=None) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--lr", type=float, default=2e-4, help="learning rate") + parser.add_argument("--fitName", type=str, default=None, help="base name for output files") + parser.add_argument("--noise_scale", type=float, default=0.2, help="Random noise scale for A,B seed values (0.0=no noise, 0.1=moderate noise)") + + args = parser.parse_args() + args.desyncTime=False + args.dropDataFrac=0 + + # DDP init + is_dist = (int(os.environ.get('WORLD_SIZE', '1')) > 1) or ('LOCAL_RANK' in os.environ) or ('RANK' in os.environ) + if is_dist: + import torch.distributed as dist + from torch.nn.parallel import DistributedDataParallel as DDP + from torch.utils.data.distributed import DistributedSampler + dist.init_process_group(backend='nccl') + local_rank = int(os.environ['LOCAL_RANK']) + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + local_rank = 0; rank = 0; world_size = 1 + device = check_gpu_availability() + if rank==0: + print("world_size=%d" % (world_size)) + print("Initial configuration:", vars(args)) + + # enable fast matmul paths + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision('high') + gpu_name = torch.cuda.get_device_name(device) if isinstance(device, torch.device) and device.type=='cuda' else str(device) + print("[rank %d] Using device %s : %s" % (rank, str(device), gpu_name)) + + fit1FF = os.path.join(args.dataPath, f"{args.dataName}.FDRselected.npz") + fitD, fitMD = read_data_npz(fit1FF,verb=rank==0) + A_init = fitD['A_avr'] + B_init = fitD['B_avr'] + + if rank==0: pprint(fitMD) + fmd=fitMD['fit_lasso'] + # Inherit hyperparams from stageA if not provided + if args.num_epochs is None: + args.num_epochs = fmd['num_epochs'] + if args.num_samples is None: + args.num_samples = fmd['num_samples_used'] + if args.batch_size is None: + args.batch_size = fmd['batch_size'] + + if rank==0: print("Effective configuration:", vars(args)) + + spikeF = fmd['lassoFit_input_name'] + inpPath= fitMD['fit_lasso']['lassoFit_input_path'] + spikesFF = os.path.join(inpPath, f"{spikeF}.spikes.npz") + + spikeD, spikeMD = read_data_npz(spikesFF,verb=rank==0) + dataYield, dataRates = spikeD['spikes'], np.clip(spikeD['single_rates'], 0.1, 40.0) + T, M = dataYield.shape + step_size = spikeMD['time_step_sec'] + + if args.fitName is None: + import hashlib + hash_str = hashlib.md5(os.urandom(32)).hexdigest()[:6] + args.fitName = f"{args.dataName}-{hash_str}" + + XY_np = preprocess_data(dataYield, args) + # Split XY into X and Y for all ranks + X_np = XY_np[:, 0, :] + Yt_np = XY_np[:, 1, :] + n_pairs = X_np.shape[0] + + assert n_pairs >= args.batch_size, f"ERROR: Not enough samples ({n_pairs}) for batch size ({args.batch_size})." + + train_loader = make_loader(X_np, Yt_np, args, is_dist=is_dist) + + if rank==0: + print(f"Loaded T={T}, M={M}, using {n_pairs} pairs (all for training), world_size={world_size}, per_gpu_bs={args.batch_size//max(1,world_size)}") + + # --- Set up model for L2-only fit --- + trainable_mask = fitD['E_mask'] + if rank==0: + print(f"Trainable mask: {trainable_mask.sum()} out of {trainable_mask.size} A elements are trainable") + print(f"Trainable fraction: {trainable_mask.mean():.3f}") + + base_model = PoissonGLModel(M, A_init=A_init, B_init=B_init, trainable_mask=trainable_mask, noise_scale=args.noise_scale).to(device) + model = DDP(base_model, device_ids=[local_rank]) if is_dist else base_model + mdl = model.module if hasattr(model,'module') else model + + # Debug: check if noise was actually applied + if args.noise_scale > 0: + if rank==0: print(f"Checking noise application...") + A_reconstructed = mdl.A.detach().cpu().numpy() + A_diff = np.abs(A_reconstructed - A_init).max() + B_reconstructed = mdl.B.detach().cpu().numpy() + B_diff = np.abs(B_reconstructed - B_init).max() + print(f"Max A difference from init: {A_diff:.6f}") + print(f"Max B difference from init: {B_diff:.6f}") + + # Debug: count trainable parameters + total_params = sum(p.numel() for p in mdl.parameters()) + trainable_params = sum(p.numel() for p in mdl.parameters() if p.requires_grad) + expected_trainable = trainable_mask.sum() + M # masked A elements + B vector + if rank==0: + print(f"Total parameters: {total_params}, PyTorch trainable: {trainable_params}, Expected trainable: {expected_trainable}") + + if hasattr(mdl, 'is_stage_b') and mdl.is_stage_b: + if rank==0: print(f"Stage B efficient parameterization: C tensor has {mdl.C.numel()} parameters") + else: + if rank==0: print("Stage A standard parameterization") + + # Debug: check initial loss and parameter values + model.train() + + # Check initial loss before training + with torch.no_grad(): + for Y_prev, Y_curr in train_loader: + Y_prev, Y_curr = Y_prev.float().to(device), Y_curr.float().to(device) + spikes = model(Y_prev) + initial_loss = poisson_nll_loss(spikes, Y_curr, torch.tensor(dataRates, dtype=torch.float32, device=device)) + if rank==0: print(f"Initial loss with noise: {initial_loss:.6f}") + break + + # Check if C parameters actually changed from Stage A + if hasattr(mdl, 'C'): + C_values = mdl.C.detach().cpu() + if rank==0: print(f"C parameter stats: min={C_values.min():.6f}, max={C_values.max():.6f}, std={C_values.std():.6f}") + + # Check gradients after one step + for Y_prev, Y_curr in train_loader: + Y_prev, Y_curr = Y_prev.float().to(device), Y_curr.float().to(device) + spikes = model(Y_prev) + loss = poisson_nll_loss(spikes, Y_curr, torch.tensor(dataRates, dtype=torch.float32, device=device)) + loss.backward() + + # Check gradients for efficient parameterization + if hasattr(mdl, 'C') and mdl.C.grad is not None: + C_grad_norm = mdl.C.grad.norm().item() + C_grad_max = mdl.C.grad.abs().max().item() + if rank==0: print(f"C gradient: norm={C_grad_norm:.6f}, max={C_grad_max:.6f}") + break + + model.zero_grad() # Clear gradients before actual training + + start_time = time.time() + losses_total, _, learning_rates, losses_epochs = train_Poisson_model( + model, device, train_loader, args.num_epochs, lr=args.lr, L1_alpha=0.0, use_scheduler=True, firing_rates=dataRates, + train_sampler=train_loader.sampler if is_dist else None + ) + total_time = time.time() - start_time + if rank==0: + print(f"Training completed in {total_time:.1f} seconds") + + if rank==0: + mdl = model.module if hasattr(model,'module') else model + bigD = { 'A_regress': mdl.A.detach().cpu().numpy(), 'B_regress': mdl.B.detach().cpu().numpy(), 'losses_total': np.array(losses_total), 'losses_epochs': np.array(losses_epochs, dtype=np.int32), 'learning_rates': np.array(learning_rates), 'single_rates': dataRates } + metaD = { 'regressFit_output_name': args.fitName, 'regressFit_input_name': args.dataName, + 'regressFit_input_name': args.dataName, + 'regressFit_output_name': args.fitName, + 'batch_size': args.batch_size, 'num_samples_used': n_pairs, 'num_epochs': args.num_epochs, 'num_train_samples': n_pairs, 'learning_rate': args.lr, 'step_size': step_size, 'training_time_sec': total_time, 'num_neurons': M, 'desyncTime': args.desyncTime } + fitMD['fit_regress']=metaD + fitMD['data_type']=spikeMD['data_type'] + fitMD['fit_type']='regress' + fitMD['short_name']=args.fitName + fit2FF = os.path.join(args.dataPath, f"{args.fitName}.regressFit.npz") + write_data_npz(bigD, fit2FF, metaD=fitMD) + print('\n ./eval_fitRegress.py --dataPath $fitPath --dataName %s -p b c ' % (args.fitName)) + + # ensure distributed shutdown to avoid resource leak warning + if is_dist and dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + +if __name__ == "__main__": + main() diff --git a/causal_net/stationaryLag1_ver6_lasso/mixSpikes.py b/causal_net/stationaryLag1_ver6_lasso/mixSpikes.py new file mode 100755 index 00000000..84270425 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/mixSpikes.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python3 +""" +Neuron mixing tool for simulated Dale Poisson network data. + +This script modifies simulated neural network data by selectively removing +neurons according to various strategies. It processes both connectivity +matrices and spike data to create modified datasets for testing robustness +and generalization of connectivity inference methods. + +Main functionality includes: +- Random neuron removal (dropAny): Randomly removes specified fraction of neurons +- Low-frequency neuron removal (dropLowFreq): Removes neurons with lowest firing rates +- Preserves data structure and updates index mappings +- Creates new datasets with modified metadata + +Usage: + ./mixSpikes.py --inpSimName daleM150_448b86 --action dropAny 0.2 --mixTag mix1 + ./mixSpikes.py --inpSimName daleM150_448b86 --action dropLowFreq 0.3 --mixTag lowFreqDrop +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os +import numpy as np +import argparse +from pprint import pprint +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser(description="Mix/modify simulated Dale Poisson spike data by removing neurons") + parser.add_argument("-v","--verbosity",type=int, help="increase output verbosity", default=1, dest='verb') + parser.add_argument("--dataPath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for input/output data") + parser.add_argument("--inpSimName", type=str, required=True, help="Input simulation base name (output from sim_dalePoisson.py)") + parser.add_argument("--action", nargs='+', required=True, help="Action to perform: dropAny X or dropLowFreq X where X is fraction (e.g., dropAny 0.2 for 20%%)") + parser.add_argument("--mixTag", type=str, required=True, help="Tag appended to output name (e.g., 'mix1')") + + args = parser.parse_args() + + print('myArg-program:',parser.prog) + for arg in vars(args): print('myArg:',arg, getattr(args, arg)) + + assert os.path.exists(args.dataPath) + return args + +#...!...!.................... +def parse_action(action_list): + if len(action_list) != 2: + raise ValueError(f"Action must be 'dropAny X' or 'dropLowFreq X', got: {action_list}") + action_type, fraction_str = action_list + try: + fraction = float(fraction_str) + except ValueError: + raise ValueError(f"Fraction must be a number, got: {fraction_str}") + if not 0 < fraction < 1: + raise ValueError(f"Fraction must be between 0 and 1, got: {fraction}") + return action_type, fraction + +#...!...!.................... +def select_neurons_to_keep(action_type, fraction, Nn, single_rates, num_excite): + n_drop = int(Nn * fraction) + n_keep = Nn - n_drop + print(f"\nSelecting neurons to keep: action={action_type}, fraction={fraction:.2f}") + print(f" Original neurons: {Nn} (excit={num_excite}, inhib={Nn-num_excite}), dropping: {n_drop}, keeping: {n_keep}") + + if action_type == 'dropAny': + keep_indices = np.sort(np.random.choice(Nn, size=n_keep, replace=False)) + print(f" Random selection: kept neuron indices range [{keep_indices[0]}, {keep_indices[-1]}]") + elif action_type == 'dropLowFreq': + freq_sorted_indices = np.argsort(single_rates) + keep_indices = np.sort(freq_sorted_indices[n_drop:]) + dropped_freq_range = [single_rates[freq_sorted_indices[0]], single_rates[freq_sorted_indices[n_drop-1]]] + kept_freq_range = [single_rates[keep_indices[0]], single_rates[keep_indices[-1]]] + print(f" Dropped frequency range: [{dropped_freq_range[0]:.2f}, {dropped_freq_range[1]:.2f}] Hz") + print(f" Kept frequency range: [{kept_freq_range[0]:.2f}, {kept_freq_range[1]:.2f}] Hz") + else: + raise ValueError(f"Unknown action type: {action_type}. Must be 'dropAny' or 'dropLowFreq'") + + num_excite_kept = np.sum(keep_indices < num_excite) + print(f" Kept neurons: excit={num_excite_kept}, inhib={n_keep-num_excite_kept}") + + return keep_indices, num_excite_kept + +#...!...!.................... +def remove_neurons(trueD, spikeD, keep_indices): + print("\n=== Removing neurons from data ===") + Nn_orig = trueD['A_true'].shape[0] + Nn_new = len(keep_indices) + + print(f"Original size: {Nn_orig} neurons") + print(f"New size: {Nn_new} neurons") + + # Remove from connectivity matrix (both rows and columns) + A_new = trueD['A_true'][np.ix_(keep_indices, keep_indices)] + print(f" A_true: {trueD['A_true'].shape} -> {A_new.shape}") + + # Remove from bias vector + B_new = trueD['B_true'][keep_indices] + print(f" B_true: {trueD['B_true'].shape} -> {B_new.shape}") + + # Remove from spike data (all time points, selected neurons) + spikes_new = spikeD['spikes'][:, keep_indices] + print(f" spikes: {spikeD['spikes'].shape} -> {spikes_new.shape}") + + # Remove from firing rates + rates_new = spikeD['single_rates'][keep_indices] + print(f" single_rates: {spikeD['single_rates'].shape} -> {rates_new.shape}") + + # Update index mappings - create new sequential indices + neur_revFreqIdx_new = np.arange(Nn_new) + neur_freqIdx_new = np.arange(Nn_new) + print(f" Index mappings recreated for {Nn_new} neurons") + + # Create output dictionaries + trueD_new = { + 'A_true': A_new, + 'B_true': B_new, + 'neur_freqIdx': neur_freqIdx_new, + 'neur_revFreqIdx': neur_revFreqIdx_new + } + + spikeD_new = { + 'spikes': spikes_new, + 'single_rates': rates_new + } + + + return trueD_new, spikeD_new + +#...!...!.................... +def update_metadata(trueMD, spikeMD, args, action_type, fraction, Nn_orig, Nn_new, num_excite_new): + print("\n=== Updating metadata ===") + + # Update simTruth metadata + trueMD_new = {} + trueMD_new['short_name'] = f"{args.inpSimName}-{args.mixTag}" + if 'evol_conf' in trueMD: + trueMD_new['evol_conf'] = trueMD['evol_conf'] + if 'dale_conf' in trueMD: + dale_conf_new = trueMD['dale_conf'].copy() + dale_conf_new['num_neurons'] = Nn_new + dale_conf_new['num_excite'] = int(num_excite_new) + trueMD_new['dale_conf'] = dale_conf_new + print(f" simTruth metadata: updated 'dale_conf' with num_neurons={Nn_new}, num_excite={num_excite_new}") + trueMD_new['input_mixer'] = { + 'inpSimName': args.inpSimName, + 'action': ' '.join(args.action), + 'mixTag': args.mixTag, + 'num_neurons_orig': Nn_orig, + 'num_neurons_new': Nn_new, + 'dropped_fraction': fraction, + 'dropped_count': Nn_orig - Nn_new + } + + print(f" simTruth metadata: removed 'dale_simu_stats'") + print(f" simTruth metadata: updated 'short_name' to {trueMD_new['short_name']}") + print(f" simTruth metadata: created 'input_mixer' with action info") + + # Update spike metadata + spikeMD_new = spikeMD.copy() + spikeMD_new['short_name'] = f"{args.inpSimName}-{args.mixTag}" + print(f" spike metadata: updated 'short_name' to {spikeMD_new['short_name']}") + #pprint(trueMD_new); aaa + return trueMD_new, spikeMD_new + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__=="__main__": + args = get_parser() + np.set_printoptions(precision=3) + + # Parse action list + action_type, fraction = parse_action(args.action) + + # Load input simulation data + print(f"\n=== Loading input data: {args.inpSimName} ===") + truthFF = os.path.join(args.dataPath, f"{args.inpSimName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb>0) + + spikesFF = os.path.join(args.dataPath, f"{args.inpSimName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb>0) + + if args.verb > 1: + print("\nOriginal simTruth metadata:") + pprint(trueMD) + print("\nOriginal spike metadata:") + pprint(spikeMD) + + # Get original neuron counts + Nn_orig = trueD['A_true'].shape[0] + num_excite_orig = trueMD.get('dale_conf', {}).get('num_excite', Nn_orig) + + # Select neurons to keep + keep_indices, num_excite_new = select_neurons_to_keep(action_type, fraction, Nn_orig, spikeD['single_rates'], num_excite_orig) + + # Remove neurons from data + trueD_new, spikeD_new = remove_neurons(trueD, spikeD, keep_indices) + + # Update metadata + Nn_new = len(keep_indices) + trueMD_new, spikeMD_new = update_metadata(trueMD, spikeMD, args, action_type, fraction, Nn_orig, Nn_new, num_excite_new) + + # Write output files + outName = f"{args.inpSimName}-{args.mixTag}" + print(f"\n=== Writing output files: {outName} ===") + + outTruthFF = os.path.join(args.dataPath, f"{outName}.simTruth.npz") + write_data_npz(trueD_new, outTruthFF, metaD=trueMD_new) + print(f" Wrote: {outTruthFF}") + + outSpikesFF = os.path.join(args.dataPath, f"{outName}.spikes.npz") + write_data_npz(spikeD_new, outSpikesFF, metaD=spikeMD_new) + print(f" Wrote: {outSpikesFF}") + + if args.verb > 1: + print("\nNew simTruth metadata:") + pprint(trueMD_new) + print("\nNew spike metadata:") + pprint(spikeMD_new) + + print("\nMixing completed successfully!") + print("\nNext step commands:") + print(f" ./view_dalePoisson.py --dataName {outName} -p a b c") + print(f" ./fit_lassoPoisson.py --dataName {outName} --num_epochs 50") + print(" --dataPath "+args.dataPath) + print("M:done") + diff --git a/causal_net/stationaryLag1_ver6_lasso/prep_bioexp.py b/causal_net/stationaryLag1_ver6_lasso/prep_bioexp.py new file mode 100755 index 00000000..d806457b --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/prep_bioexp.py @@ -0,0 +1,258 @@ +#!/usr/bin/env python3 +""" +Preprocessing pipeline for experimental neural data from Roy/Mandar laboratory. + +This script processes raw experimental neural recordings into standardized format +for connectivity analysis. The preprocessing pipeline includes: +- Raw data loading and format conversion +- Temporal binning and spike count extraction +- Data quality assessment and filtering +- Metadata extraction and session identification +- Output formatting for downstream analysis tools + +Session naming convention: +- B6J: cell line name +- 250619: recording date (YYMMDD) +- M08020: chip identifier +- 000093: run number +- Well000: well number + +The script generates .spikes.npz files with standardized spike count matrices +and associated metadata for further analysis. + +Usage: + ./prep_bioexp.py --sessionName B6J_250619_M08020_000093_Well000 --inputPath /path/to/raw/data/ +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" +import sys,os,hashlib +import numpy as np +import pickle +from pprint import pprint +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + +import argparse +#...!...!.................. +def commandline_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v","--verb",type=int, help="increase debug verbosity", default=1) + parser.add_argument("--expPath",required=True,help="raw experimnetal data on CFS") + + parser.add_argument("--sessionName", default='celllinename/dateofrecording/chipID/Assaytype/runnumber/wellnumber',help='raw data session name') + parser.add_argument("--dataPath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for any further data processing") + parser.add_argument("--shortName", default=None,help='(optional) output file name - Is it needed?') + + # .... activity speciffic speciffic, + parser.add_argument('--samp_freq', default=100, type=int, help='sets binning of time axis') + + parser.add_argument('--freqRange', default=[1.,50], type=float, nargs=2,help='rebin of raw time axis') + + args = parser.parse_args() + + for arg in vars(args): + print( 'myArgs:',arg, getattr(args, arg)) + + return args + +#...!...!.................... +def buildBioMeta(args): + pd={} # payload + pd['raw_bioexp_path']=args.expPath + pd['session_name']=args.sessionName + txtL=args.sessionName.split('/') + #print('tt',txtL); aa + pd['cell_line_name']=txtL[0] + pd['recording_date']=txtL[1] + pd['chip_ID']=txtL[2] + pd['run_num']=txtL[4] + pd['well_num']=txtL[5] + + sel={'freq_range':args.freqRange} + md={ 'bioexp':pd,'data_selector':sel} + myHN=hashlib.md5(os.urandom(32)).hexdigest()[:6] + md['hash']=myHN + if args.shortName==None: + md['short_name']='%s-%s'%(pd['recording_date'],md['hash']) + else: + md['short_name']=args.shortName + + if args.verb>1: print('\nBMD:');pprint(md) + return md + + +def read_spike_npy(md,args): + pmd=md['bioexp'] + inpF=os.path.join(args.expPath,args.sessionName,'spike_times.npy') + print('inpF:',inpF) + assert os.path.exists(inpF) + # Load the dictionary from the .npy file + spike_dict = np.load(inpF, allow_pickle=True).item() + + #print('spike_dict',spike_dict);ok + if args.verb>1: print('input spike_dict',sorted(spike_dict)) + pmd['sampling_freq'] =args.samp_freq + + # neuron ID MEA chip + meaIdL=np.array(sorted(spike_dict)) # here order of feature_id is settled + maxFeat=len(meaIdL) + + if args.verb>1: print('RSN: meaID list:',meaIdL) + pmd['num_feature']=len(meaIdL) + + spikeT={} # spike times + spikeCntL=np.zeros(pmd['num_feature'],dtype=int) # num spikes per neuron + maxTbin=0 + + j=0 + for k in meaIdL: + rec=np.array(spike_dict[k])*args.samp_freq + #print(rec[:100],len(rec)) + spikeT[k]=rec.astype(int) # time-bin may repeat + spikeCntL[j]=len(rec) + j+=1 + if len(rec)==0: + continue + mxTb=np.max(rec) + if maxTbin< mxTb: maxTbin=mxTb + #exit(0) + pmd['num_time_bin']=int(maxTbin)+1 + pmd['max_time']=pmd['num_time_bin']/pmd['sampling_freq'] + chanFreq=spikeCntL/ pmd['max_time'] + rawD={'spikeT':spikeT,'chanFreq':chanFreq,'MEA_idx':meaIdL} + return rawD + + +#...!...!.................... +def unroll_bioexp(rawD,bioMD): + pmd=bioMD['bioexp'] + sel=bioMD['data_selector'] + frLo, frHi = sel['freq_range'] + print('frLo, frHi',frLo, frHi) + assert frLo < frHi + chanFreq = np.asarray(rawD['chanFreq'], dtype=float) + # vectorized boolean mask for channels within (frLo, frHi) range + freqMask = (chanFreq >= frLo) & (chanFreq <= frHi) + sel['drop_neur_by_freq_range']=[ int(np.sum(chanFreq < frLo)), int(np.sum(chanFreq > frHi)) ] + print('freqMask all=%d , passed=%d'%(freqMask.shape[0],np.sum(freqMask))) + #print(sel);aaa + # --- drop channles out of freq range + chanFreq=rawD['chanFreq'][freqMask] + MEA_idx=rawD['MEA_idx'][freqMask] + + # .... REMAP MATRICES TO FREQUENCY-SORTED ORDER (PRIMARY INDEX) + neur_freqIdx = np.argsort(chanFreq) # indices that sort chanFreq by value + neur_revFreqIdx = np.empty(len(neur_freqIdx), dtype=int) # natural_index → freq_sorted_position + neur_revFreqIdx[neur_freqIdx] = np.arange(len(neur_freqIdx)) + + #--- reorder channles by frequency + chanFreq=chanFreq[neur_freqIdx] + MEA_idx=MEA_idx[neur_freqIdx] + print('chanFreq',chanFreq[:5],'...',chanFreq[-5:],'Hz') + + # create spike matrix: rows=time bins, cols=accepted channels + ntime=pmd['num_time_bin'] + nchan=MEA_idx.shape[0] + spikes2D=np.zeros((ntime,nchan),dtype=np.int32) + spikeT=rawD['spikeT'] + for ic, ch in enumerate(MEA_idx): + tV=spikeT[ch] + if len(tV)==0: continue + # tV holds time-bin indices where this channel fired one or more spikes + # bincount returns a length-ntime vector with spike counts per bin (zeros elsewhere) + # minlength=ntime guarantees the vector spans the full recording duration + cnt=np.bincount(tV, minlength=ntime) + spikes2D[:,ic]=cnt.astype(np.int32) + print('spikes2D shape',spikes2D.shape) + + # keep handy in meta for downstream + sel['num_chan']=nchan + sel['max_spike_per_bin']=int(np.max(spikes2D)) + + Y_uchar = np.clip(spikes2D, 0, 255).astype(np.uint8) + spikeD={'spikes':Y_uchar, + 'single_rates':chanFreq + } + + bioD={ } + bioD['neur_freqIdx']=neur_freqIdx + bioD['neur_revFreqIdx']=neur_revFreqIdx + bioD['MEA_idx']=MEA_idx + + #.... compute neural statistics for spikeMD + num_neurons = nchan + avg_rate = float(np.mean(chanFreq)) + std_rate = float(np.std(chanFreq)) + median_rate = float(np.median(chanFreq)) + min_rate = float(np.min(chanFreq)) + max_rate = float(np.max(chanFreq)) + + # Compute Fano factor (variance/mean) for each neuron + mean_counts_per_bin = np.mean(spikes2D, axis=0) + spike_variance = np.var(spikes2D, axis=0) + fano_factor = np.divide(spike_variance, mean_counts_per_bin, out=np.zeros_like(spike_variance), where=mean_counts_per_bin != 0) + avg_fano = float(np.mean(fano_factor)) + std_fano = float(np.std(fano_factor)) + + # Print summary statistics + print('Neural Statistics Summary:') + print('num neurons: %d, Avg Rate= %.2f±%.2f Hz, Avg Fano=%.2f±%.2f' % (num_neurons, avg_rate, std_rate, avg_fano, std_fano)) + print('Median rate %.2f Hz' % median_rate) + + #.... extract spikeMD for fitter + spikeMD={'time_step_sec': 1./pmd['sampling_freq'], 'short_name':bioMD['short_name'], 'data_type':'bioExp', 'num_neurons': num_neurons } + + bioMD['rate_summary']={ + 'avg_spike_rate': avg_rate, + 'std_spike_rate': std_rate, + 'avg_fano_factor': avg_fano, + 'std_fano_factor': std_fano, + 'median_spike_rate': median_rate, + 'min_spike_rate': min_rate, + 'max_spike_rate': max_rate + } + return bioD,spikeD,spikeMD + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__ == "__main__": + + args=commandline_parser() + np.set_printoptions(precision=3) + bioMD=buildBioMeta(args) + + # read raw data + #rawD=read_spike_dict(bioMD,args) + rawD=read_spike_npy(bioMD,args) + + #.... filter & unroll data + bioD,spikeD,spikeMD=unroll_bioexp(rawD,bioMD) + + #...... WRITE OUTPUT ......... + outFt = os.path.join(args.dataPath, bioMD['short_name'] + '.bioExp.npz') + write_data_npz(bioD, outFt, metaD=bioMD) + if args.verb>2: + print('\n bioD:',sorted(bioD)) + pprint(bioMD) + + outFs = outFt.replace('.bioExp.','.spikes.') + write_data_npz(spikeD, outFs, metaD=spikeMD) + if args.verb>2: + print('\nspikeD:',sorted(spikeD)) + pprint(spikeMD) + + print("\nNext step command:") + print(' ./view_bioexp.py --dataPath $dataPath --dataName %s -p a b -T 0 3550 '%(bioMD['short_name'] )) + print(" ./fit_lassoPoisson.py --dataPath $dataPath --dataName %s --num_epochs 10 " % bioMD['short_name'] ) + print(" ./fitLasso4GPU.sh --dataPath $dataPath --dataName %s --num_epochs 200 " % bioMD['short_name'] ) + print(" ./bootsFit.sh --dataName %s --num_epochs 250 --dropDataFrac 0.5 --num_bootstraps 7 " % bioMD['short_name'] ) + + print(" ./selectEdges_FDR.py --dataName %s --num_bootstraps 6 10 -p a c d " % bioMD['short_name'] ) + + print(' --dataPath '+args.dataPath) + + + diff --git a/causal_net/stationaryLag1_ver6_lasso/selectEdges_FDR.py b/causal_net/stationaryLag1_ver6_lasso/selectEdges_FDR.py new file mode 100755 index 00000000..4eb41537 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/selectEdges_FDR.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +""" +This script loads LASSO connectivity matrices from K bootstraps of real and desync data. +Diagonal elements are separated to compute mean and std across real-data bootstraps. +For off-diagonal edges, median magnitudes from real bootstraps are compared to a pooled null distribution +from desync bootstraps within each row to compute empirical p-values. +Row-wise Benjamini–Hochberg FDR is applied to select statistically significant edges at the given alpha level. +FDR = False Discovery Rate - a statistical method for controlling errors when testing many hypotheses simultaneously. +""" +import argparse +import os +import numpy as np +from pprint import pprint +from toolbox.Util_NumpyIO import write_data_npz +from PlotterFitEval import Plotter +from UtilSelectFDR import (summary_reco_neuronNet, compare_triplets, eval_tagged_edges_4_simu, + get_offdiag_triplets, edge_selector_fdr, load_bootstrap_data, + load_auxiliary_plotting_data) + +def print_table_4_Yao(evalD,fitMD): + + Nn=fitMD['fit_lasso'] ['num_neurons'] + Nedg=Nn*(Nn-1) + recL='#L,Nn,Nedg,' + recV='#prob,%d,%d,'%(Nn,Nedg) + recM='#M,' + recR='#rmse,' + for etype in ['neg','pos']: + tripV=evalD[etype] + TP,FP,FN=tripV + + p=TP.shape[0]/Nedg + if p>0 and p<1 : std=np.sqrt(p*(1-p)/Nedg) + else: std=1./Nedg + recL+=etype+'TP,std,' + recV+='%.4f,%.4f,'%(p,std) + + res=TP[:,3]-TP[:,2] + mean_val = np.mean(res) + std_val = np.std(res) + recM+=etype+'TP,std,' + recR+='%.2e,%.2e,'%(mean_val,std_val) + + p=FP.shape[0]/Nedg + if p>0 and p<1 : std=np.sqrt(p*(1-p)/Nedg) + else: std=1./Nedg + recL+=etype+'FP,std,' + recV+='%.4f,%.4f,'%(p,std) + + p=FN.shape[0]/Nedg + if p>0 and p<1 : std=np.sqrt(p*(1-p)/Nedg) + else: std=1./Nedg + recL+=etype+'FN,std,' + recV+='%.4f,%.4f,'%(p,std) + + for xx in ['diag','bterm']: + V=evalD[xx] + res=V[:,1]-V[:,0] + mean_val = np.mean(res) + std_val = np.std(res) + recM+=xx+',std,' + recR+='%.2e,%.2e,'%(mean_val,std_val) + + print(recL) + print(recV) + print(recM) + print(recR) + + +######################### +# MAIN +######################### + +def main(): + parser = argparse.ArgumentParser(description="Row-wise FDR edge selection from LASSO bootstraps") + parser.add_argument("--dataName", type=str, required=True, help="Base name for the dataset") + parser.add_argument("--basePath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input/output data") + parser.add_argument("--inpPath", type=str, default=None, help="alternative location of input, takes precedence") + parser.add_argument("--verb", "-v", type=int, default=1, help="Verbosity level") + parser.add_argument("--num_bootstraps", type=int, nargs='+', required=True, help="Number of bootstraps: 1 value (duplicated for real/desync) or 2 values [Kreal, Kdesync]") + parser.add_argument("--alphaFDR", type=float, default=0.002, help="FDR significance level") + + parser.add_argument('-p',"--showPlots", type=str,nargs='+', default="f", help="Plot types to show: a=structure, e=experiment_eigen, f=freqSortA_histos") + parser.add_argument("--outPath", type=str, default=None, help="Output path for plots (defaults to basePath/plots)") + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + args = parser.parse_args() + print(vars(args)) + + if args.inpPath is None: + args.inpPath = os.path.join(args.basePath, 'lassoFdrFit') + if args.outPath is None: + args.outPath = os.path.join(args.basePath, 'plots') + os.makedirs(args.outPath, exist_ok=True) + args.showPlots=''.join(args.showPlots) + + dataName = args.dataName + dataPath = args.inpPath + # Handle 1 or 2 values for num_bootstraps + if len(args.num_bootstraps) == 1: + K = [args.num_bootstraps[0], args.num_bootstraps[0]] # Duplicate single value + elif len(args.num_bootstraps) == 2: + K = args.num_bootstraps + else: + raise ValueError("--num_bootstraps must have 1 or 2 values") + + Kreal, Kdesync = K + alpha = args.alphaFDR + + # Load all bootstrap data + A_edges_real_list, A_edges_desync_list, A_real_list, B_real_list, output_meta,output_big1 = load_bootstrap_data( + dataName, dataPath, K, args.verb + ) + + # Apply row-wise FDR to edges + E_mask, E_pval, summary = edge_selector_fdr(A_edges_real_list, A_edges_desync_list, alpha=alpha) + + print("FDR selection summary:"); pprint(summary) + + # Compute averages and standard deviations from real data bootstraps + A_stack = np.stack(A_real_list, axis=0) # shape (K, N, N) + B_stack = np.stack(B_real_list, axis=0) # shape (K, N) + + A_avr = np.mean(A_stack, axis=0) + A_std = np.std(A_stack, axis=0) + B_avr = np.mean(B_stack, axis=0) + B_std = np.std(B_stack, axis=0) + + # Create full E_mask that applies only to off-diagonal elements + N = A_avr.shape[0] + E_mask_full = np.eye(N, dtype=bool) # Start with diagonal = True (keep diagonal) + # Apply FDR mask to off-diagonal elements only + off_diag_mask = ~np.eye(N, dtype=bool) + E_mask_full[off_diag_mask] = E_mask[off_diag_mask] + + A_avr[~E_mask_full]=0. # now none-existing edges are 0 + + print(f"Computed averages and std from {K} real bootstraps") + print(f"A_avr shape: {A_avr.shape}, B_avr shape: {B_avr.shape}") + print(f"Mask preserves diagonal and selects {E_mask.sum()} significant off-diagonal edges") + + # Prepare output data + output_data = { + 'E_mask': E_mask_full, + 'A_avr': A_avr, + 'A_std': A_std, + 'B_avr': B_avr, + 'B_std': B_std, + 'E_pval': E_pval, + 'summary': summary + } + for xx in [ 'losses_total', 'losses_epochs', 'losses_wo_L1']: + output_data[xx]=output_big1[xx] + + output_meta['edge_selector']={'selector_type':'FDR', 'alpha':args.alphaFDR} + + # Save results + output_file = os.path.join(dataPath, f"{dataName}.FDRselected.npz") + write_data_npz(output_data, output_file, metaD=output_meta) + print(f"FDR results saved to: {output_file}") + + print("\nNext step commands:") + print(f" ./fit_regressPoisson.py --dataPath $fitPath --dataName {dataName} ") + + + # Generate plots if requested + if args.showPlots: + print(f"\nGenerating plots: {args.showPlots}") + + # Prepare plotting data (compatible with eval_fitLasso.py structure) + fitD = output_data.copy() # Use our processed output data as fitD + fitMD = output_meta + + if 0: # patch old data + fitMD['fit_lasso']['num_epochs']=fitMD['fit_lasso']['n_epochs'] + + # Rename records so select_edges_from_fitLasso() has the expected names + fitD['A_lasso'] = A_avr.copy() # tmp + fitD['B_lasso'] = B_avr.copy() + + # Load auxiliary data needed for plotting + spikeD, trueD, MD = load_auxiliary_plotting_data(fitMD, dataName, dataPath) + # Add mask metadata and FDR method info + + MD['edge_selection_method'] = 'fdr' + MD['fdr_alpha'] = alpha + + if fitMD['data_type']=='simDale': + evalD=eval_tagged_edges_4_simu(fitD,trueD) + print_table_4_Yao(evalD,fitMD) + + edgeD=summary_reco_neuronNet(A_avr, A_std) + + # adjustment for plotting + + fitD['single_rates']=spikeD['single_rates'] + + # Setup plotter + args.prjName = dataName + plot = Plotter(args) + + # Generate plots based on showPlots argument + if 'a' in args.showPlots: + plot.summary_fitLasso(fitD,MD,figId=1) + + if 'b' in args.showPlots: + assert fitMD['data_type']=='simDale' + plot.residuals(evalD,MD,figId=2) + + if 'c' in args.showPlots: + plot.summary_network(fitD, edgeD,MD, figId=3) + + if 'd' in args.showPlots: + plot.freqSortA_histos(fitD, MD, spikeD, figId=4) + + + plot.display_all() + print("Plotting completed.") + +if __name__ == "__main__": + main() diff --git a/causal_net/stationaryLag1_ver6_lasso/sim_daleLag1Poisson.py b/causal_net/stationaryLag1_ver6_lasso/sim_daleLag1Poisson.py new file mode 100755 index 00000000..3ad4d1db --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/sim_daleLag1Poisson.py @@ -0,0 +1,448 @@ +#!/usr/bin/env python3 +""" +python gen_dalePoisson.py --num_neurons 10 --num_excite 6 --num_steps 1000 --dataName test_dale + + +This script simulates the activity of a recurrent neural network with biologically +inspired constraints. The key features of the simulation are: + +***Dale Poisson Simulator - Program Summary*** +This program simulates the activity of a recurrent neuronal network using a discrete-time +Poisson generalized linear model (GLM). The network obeys Dale's principle, meaning that +each neuron is either excitatory (producing only positive outgoing weights) or inhibitory +(producing only negative outgoing weights). + +Key Features: + +Connectivity Matrix Generation: +A spectral radius–based stabilization procedure ensuring max Re(lambda(A)) < 0, so the single-lag network decays in the absence of input noise. + +your script stabilizes using continuous Lyapunov equations (solve_continuous_lyapunov) which assumes continuous-time dynamics and only one lag. + + +The recurrent connectivity matrix A is generated with random weights that respect Dale's principle. +The parameter R scales the baseline synaptic strength and influences the initial spectral radius of A. +Regularization of A: + +A stability check is performed by computing the maximum real part of the eigenvalues. +Inhibitory weights (A_ij < 0) are iteratively adjusted using a gradient derived from Lyapunov equations. +Updates continue until the maximum real eigenvalue falls below a threshold (-delta), ensuring stable dynamics. +Poisson Process Simulation: + +At each time step t, the firing rate for each neuron is computed as: lambda(t, i) = exp( sum_j A_ij * Y(t-1, j) + B_i ) +Spike counts Y(t, i) are drawn from a Poisson distribution with rate lambda(t, i) * dt. +Additional customization is provided via command-line arguments, allowing you to set the +number of neurons, number of excitatory neurons, spectral radius (R), time step (dt), and +simulation duration. The program saves the generated spike data, firing rates, connectivity +matrix, and other simulation details for further analysis. + +""" + +import numpy as np +import time,hashlib +import scipy +import os +import sys + +import argparse +from pprint import pprint +from UtilDalePoisson import estimate_rates +from toolbox.Util_NumpyIO import write_data_npz + +###### Matrix generation ################## +# Generate an initial network connectivity matrix +def gen_init_W(num_neurons, num_excite, rho_target, gamma, R, varyW, minW,diag=0): + print(f" Generating initial connectivity matrix: N={num_neurons}, excit={num_excite}, gamma={gamma:.1f}, R={R:.1f}") + rand = np.random.default_rng() + + num_inhib = num_neurons - num_excite + Ainit = np.zeros((num_neurons, num_neurons)) + + # this scaling is a guess, may need adjustment + p_eff = np.mean(rho_target) / num_neurons + wC = R/np.sqrt(p_eff * (1 - p_eff) * (1 + gamma**2)/2) # central value + + # decide how much variation in weights + assert varyW<1 + wL=wC*varyW / np.sqrt(num_neurons) + wR=wC/varyW / np.sqrt(num_neurons) + if wL< minW: # shift both by the difference + delW=minW-wL + wL+=delW + wR+=delW + print(f" Weight parameters: wL={wL:.3f}, wR={wR:.3f}, p_eff={p_eff:.3f}") + + for j in range(num_neurons): + # Determine the number of connections for this neuron + num_connections = rho_target[j] + if num_connections == 0: + continue + + # Choose target neurons to connect to + targets = rand.choice(num_neurons, num_connections, replace=False) + + # Generate random weights for these connections + weights = np.random.uniform(wL, wR, size=num_connections) + + if j < num_excite: # Excitatory neuron + Ainit[j, targets] = weights + else: # Inhibitory neuron + Ainit[j, targets] = -gamma * weights + + # Setting diagonals to 0 initially + np.fill_diagonal(Ainit, diag) + excit_connections = np.sum(Ainit[:num_excite,:] > 1e-9) + inhib_connections = np.sum(Ainit[num_excite:,:] < -1e-9) + + print(f" Connections created: excit={excit_connections}, inhib={inhib_connections}") + print(f" Matrix stats: min={np.min(Ainit):.3f}, max={np.max(Ainit):.3f}, mean={np.mean(Ainit):.3f}") + return Ainit + +# Optimize the inhibitory weights of a matrix A to render it stable (i.e. max re lambda < 0) +# Implements the algorithm described here: https://epubs.siam.org/doi/abs/10.1137/070704034?journalCode=sjope8 +def stabilize(A, max_iter=3000, eta=50, C=3.0, B=1.0,delta=0.2,minW=0.1): + print(f" Stabilizing matrix: max_iter={max_iter}, eta={eta}, C={C}, B={B}, delta={delta}") + + # delta: limits real part of eigen values + # Regularization of the spectral absicca, described on pg. 8 of the supplement here: + # https://www.sciencedirect.com/science/article/pii/S0896627314003602?via%3Dihub#app2 + # delat is my margin from 0 toward negative + + alpha = np.max(np.real(np.linalg.eigvals(A))) + print(f" Initial spectral radius: {alpha:.3f}") + if alpha < -delta: + print(f" Matrix already stable (alpha={alpha:.3f} < -delta={-delta:.3f})") + return A + + iter_ = 0 + + while alpha > -delta and iter_ < max_iter: + if iter_ % 100 == 0: + print(f" Iteration {iter_}: alpha={alpha:.3f}") + + alpha_e = max(C * alpha, C * alpha + B) + Q = scipy.linalg.solve_continuous_lyapunov((A - alpha_e * np.eye(A.shape[0])).T, -2 * np.eye(A.shape[0])) + P = scipy.linalg.solve_continuous_lyapunov(A - alpha_e * np.eye(A.shape[0]), -2 * np.eye(A.shape[0])) + + grad = Q @ P/np.trace(Q @ P) + + # Adjust inhibitory weights + inh_idx = np.argwhere(A < 0) + weight_changes = 0 + for idx in inh_idx: + old_weight = A[idx[0], idx[1]] + A[idx[0], idx[1]] -= eta * grad[idx[0], idx[1]] + # Make sure no inhibitory weights got turned into excitatory weights + #if A[idx[0], idx[1]] > 0: + if A[idx[0], idx[1]] > -minW: # prevents too small weight for inhibitory cells + A[idx[0], idx[1]] = 0 + if abs(A[idx[0], idx[1]] - old_weight) > 1e-6: + weight_changes += 1 + + alpha = np.max(np.real(np.linalg.eigvals(A))) + iter_ += 1 + + if iter_ % 500 == 0: + print(f" Weight changes: {weight_changes}/{len(inh_idx)}") + + print(f" Stabilization complete after {iter_} iterations: final alpha={alpha:.3f}") + return A + +# Generate the full set of matrices for use in subsequent synthetic experiments +def gen_dale_matrics(conf, rho_target): + """Generates one stable Dale matrix based on configuration.""" + print(f"\n=== Generating Dale Matrix ===") + num_neurons, num_excite, g, r = conf['num_neurons'], conf['num_excite'], conf['g'], conf['R'] + print(f"Parameters: N={num_neurons}, excit={num_excite}, g={g:.1f}, r={r:.1f}") + + # additional configuration + varyW=0.6 # controls variation of excitatory weights + minW=0.1 # sets minimal value of any weights, also after stabilization + eigenGap=0.1 # sets threshold on Re(eigen value) after stabilization + + A = gen_init_W(num_neurons, num_excite, rho_target, g, r, varyW=varyW, minW=minW, diag=-1) + eig = np.linalg.eigvals(A) + print(f"Initial eigenvalues: real range [{np.min(np.real(eig)):.3f}, {np.max(np.real(eig)):.3f}]") + + if np.max(np.real(eig)) >= 0: + A = stabilize(A, eta=conf['eta'], C=conf['C'], B=conf['B'],delta=eigenGap,minW=minW) + eig = np.linalg.eigvals(A) + print(f"After stabilization: real range [{np.min(np.real(eig)):.3f}, {np.max(np.real(eig)):.3f}]") + else: + print("Matrix already stable") + + assert(np.max(np.real(eig)) < 0) + return A + +#################### Simulation ################## + +def set_flat_selfSpiking(args): + """Generate B_idle for flat (idleRate-based) self-spiking.""" + Nn = args.num_neurons + Ri_arg = np.array(args.idleRate) + Bi = np.log(Ri_arg) + B_idle = np.random.uniform(Bi[0], Bi[1], size=(Nn,)) + if args.exc_rate_dump != None: + B_idle[:args.num_excite] -= args.exc_rate_dump # reduce excite rate + return B_idle + +def set_real_selfSpiking(args): + """Generate B_idle from realistic frequency distribution; return rateGen_conf.""" + Nn = args.num_neurons + min_freq=args.idleRate[0] + assert min_freq>0.3 + from UtilGenExpFreqs import gen_realistic_freqs + exc_rateConf = { + 'min_freq': min_freq, + 'max_freq': 17, + 'trapezoid_height': 0.0, + 'trapezoid_rmin': 0.20, + 'sigma': 2 + } + inh_rateConf = { + 'min_freq': 4*min_freq, + 'max_freq': 90, + 'trapezoid_height': 0.15, + 'trapezoid_rmin': 0.3, + 'sigma': 6 + } + exc_freqs = gen_realistic_freqs(num_samples=args.num_excite, **exc_rateConf) + inh_freqs = gen_realistic_freqs(num_samples=Nn - args.num_excite, **inh_rateConf) + B_exc = np.log(exc_freqs) + B_inh = np.log(inh_freqs) + B_idle = np.concatenate([B_exc, B_inh], axis=0) + rateGen_conf = {'exc': exc_rateConf, 'inh': inh_rateConf} + if args.exc_rate_dump != None: + B_idle[:args.num_excite] -= args.exc_rate_dump # reduce excite rate + return B_idle, rateGen_conf + +def generate_lag1_poisson(num_steps, dt, A, B_intercept, num_excite, verb=0): + """ + Generates a multivariate Poisson VAR(1) process: + Y_t ~ Poisson(exp(A @ Y_{t-1} + B_intercept)) + + Args: + num_steps (int): Number of time steps for the simulation. + dt (float): Integration time step in seconds. + A (np.ndarray): N x N autoregressive coefficient matrix (the Dale matrix). + B_intercept (np.ndarray): N-dimensional intercept vector (bias). + num_excite (int): Number of excitatory neurons. + verb (int): Verbosity level for printing progress. + + Returns: + (Y, A, B_intercept): Tuple containing: + Y (np.ndarray): num_steps x N time series of spike counts. + A (np.ndarray): The input connectivity matrix. + B_intercept (np.ndarray): The input intercept vector. + """ + print(f"\n=== Generating Poisson Process ===") + print(f"Simulation parameters: steps={num_steps}, dt={dt:.3f}, neurons={A.shape[0]}, excit={num_excite}") + print(f"Matrix A stats: min={np.min(A):.3f}, max={np.max(A):.3f}, mean={np.mean(A):.3f}") + print(f"Bias B stats: min={np.min(B_intercept):.3f}, max={np.max(B_intercept):.3f}, mean={np.mean(B_intercept):.3f}") + + if A is None: + raise ValueError("Connectivity matrix A cannot be None.") + d=A.shape[0] + + Y = np.zeros((num_steps, d), dtype=int) + Y[0] = np.random.poisson(np.exp(B_intercept)*dt) # initial state + print(f"Initial state: total spikes={np.sum(Y[0])}, excit spikes={np.sum(Y[0][:num_excite])}, inhib spikes={np.sum(Y[0][num_excite:])}") + + if verb>0: + print('t=0 Y[t] sum=%d, Excit(first 3):%s, Inhib(first 3):%s'%(np.sum(Y[0]), Y[0][:3], Y[0][num_excite:num_excite+3])) + + kk=15 + # Main simulation loop + print("Starting main simulation loop...") + for t in range(1, num_steps): + eta = A @ Y[t-1] + B_intercept + #eta = B_intercept # use it to see idle rate only + lambda_t = np.exp(np.clip(eta, -5, 5)) # avoid overflow + Y[t] = np.random.poisson(lambda_t*dt) + + if verb>0 and t<5: + print('t=%d Y[t] sum=%d, Excit:%s, Inhib:%s'%(t, np.sum(Y[t]), Y[t][:kk], Y[t][num_excite:num_excite+kk])) + + if t % (num_steps // 10) == 0: + print(f" Progress: {t}/{num_steps} steps ({t/num_steps*100:.1f}%) - total spikes in this step: {np.sum(Y[t])}") + + print(f"Simulation complete. Final state: total spikes={np.sum(Y[-1])}") + print(f"Spike data stats: min={np.min(Y)}, max={np.max(Y)}, mean={np.mean(Y):.2f}, total spikes={np.sum(Y)}") + + return Y + +######################### +# MAIN +######################### + +def main(): + print("=" * 60) + print("DALE POISSON SIMULATION Lag=1") + print("=" * 60) + + parser = argparse.ArgumentParser(description="Simulate a recurrent neural network with Dale's principle.") + parser.add_argument("--num_neurons", type=int, default=50, help="Total number of neurons in the network.") + parser.add_argument("--num_excite", type=int, default=30, help="Number of excitatory neurons.") + parser.add_argument("--edge_prob", type=float, nargs=2, default=[0.05, 0.2], help="Range of edge probability [min, max] for rho_target generation.") + parser.add_argument("--num_steps", type=int, default=10_000, help="Number of time steps for simulation.") + parser.add_argument("--step_size", type=float, default=0.01, help="Integration time step size (dt) in seconds.") + parser.add_argument("--idleRate", type=float, nargs=2, default=[0.5, 10.], help="Range of idle firing rates [min, max] in Hz. Or min rate for 'expRate'") + parser.add_argument("--expRate", action="store_true", help="Switch to exponentially decausing rate") + parser.add_argument("--exc_rate_dump", type=float, default=None, help="boost B-value for inhibitory neurons") + parser.add_argument("--spectralR", type=float, default=1.0, help="Initial spectral radius (R).") + parser.add_argument('-v',"--verb", type=int, default=1, help="Verbosity level (0=quiet, 1=normal).") + parser.add_argument("--dataName", type=str, default=None, help="Base name for output files (default: dale_spikes_xx).") + parser.add_argument("--dataPath", type=str, default='/pscratch/sd/b/balewski/2025_causalNet_tmp/', help="Output directory for all files.") + + np.set_printoptions(precision=3, suppress=True) + + args = parser.parse_args() + # Determine output file prefix + if args.dataName is None: + args.dataName='daleM%d_'%args.num_neurons+hashlib.md5(os.urandom(32)).hexdigest()[:6] + + + print("\nStarting simulation with configuration:") + print(vars(args)) + + # Validation checks + print("Performing parameter validation...") + if args.num_excite >= args.num_neurons: + raise ValueError("Number of excitatory neurons must be less than the total number of neurons.") + + assert os.path.exists(args.dataPath) + print("Output directory exists: %s" % args.dataPath) + + Nn = args.num_neurons + # Generate rho_target Vector + print("\n=== Generating rho_target Vector ===") + probLo, probHi = args.edge_prob + min_rho = args.num_neurons * probLo + max_rho = args.num_neurons * probHi + rho_target = np.random.uniform(min_rho, max_rho, size=args.num_neurons) + rho_target = np.maximum(2.0, rho_target).astype(int) # + #1rho_target=np.linspace(1,args.num_neurons,args.num_neurons).astype(int) # testing only, linear growth + + print(f"Generated rho_target from range [%.2f, %.2f] with min value 5.0" % (min_rho, max_rho)) + print("edge_count_target stats: min=%d, max=%d, mean=%.2f" % (np.min(rho_target), np.max(rho_target), np.mean(rho_target))) + + dale_conf = { + 'num_neurons': args.num_neurons, + 'num_excite': args.num_excite, + 'g': 2, # Inhibitory-to-excitatory synaptic strength ratio + 'R': args.spectralR, # Initial spectral radius + 'eta': 10, # Learning rate for stabilization algorithm + 'C': 1.5, # Parameter for stabilization algorithm + 'B': 0.2, # Parameter for stabilization algorithm + 'edge_prob': args.edge_prob, + } + + if args.verb>1: + print("Dale configuration:"); pprint(dale_conf) + + # sanity checks + assert Nn>=10 + assert args.num_excite>=5 + assert args.num_excite=1000 + assert args.step_size>0.001 + assert args.spectralR>0.01 + if not args.expRate: + assert args.idleRate[0]>=0.5 + assert args.idleRate[1]>args.idleRate[0] + + # Generate the stable Dale connectivity matrix A + print("Generating stable Dale matrix for Nn=%d (%d Excit, %d Inhib)..." % (Nn, args.num_excite, Nn - args.num_excite)) + A_dale=gen_dale_matrics(dale_conf, rho_target) + + # Calculate and print sparsity of the generated matrix + total_connections = A_dale.size + zero_connections = np.sum(np.abs(A_dale) < 1e-10) # Count near-zero as zero + non_zero_connections = total_connections - zero_connections + sparsity = zero_connections / total_connections + + print(f'Matrix sparsity: {sparsity*100:.1f}% ({zero_connections}/{total_connections} connections are zero)') + print(f'Non-zero connections: {non_zero_connections} ({(1-sparsity)*100:.1f}%)') + + evol_conf={ + 'num_steps': args.num_steps, + 'step_size': args.step_size, + 'evol_time': args.num_steps*args.step_size, + 'expRate': args.expRate, + 'exc_rate_dump': args.exc_rate_dump + } + + if not args.expRate: + evol_conf['idleRate']= args.idleRate + B_idle = set_flat_selfSpiking(args) + else: + B_idle, rateGen_conf = set_real_selfSpiking(args) + evol_conf['rate_gen_conf']=rateGen_conf + + # print('B_idle:',B_idle) + # Generate spike data using the Poisson process + probLo, probHi = args.edge_prob + min_rho = args.num_neurons * probLo + max_rho = args.num_neurons * probHi + rho_target = np.random.uniform(min_rho, max_rho, size=args.num_neurons) + rho_target = np.maximum(4.0, rho_target).astype(int) + #print(f"Generated rho_target from range [{min_rho:.2f}, {max_rho:.2f}] with min value 5.0") + print("rho_target stats: min=%d, max=%d, mean=%.2f" % (np.min(rho_target), np.max(rho_target), np.mean(rho_target))) + + start_time = time.time() + Y = generate_lag1_poisson(num_steps=args.num_steps, dt=args.step_size, A=A_dale, B_intercept=B_idle, num_excite=args.num_excite, verb=args.verb) + sim_time = time.time() - start_time + print("Spike generation completed in %.1f seconds" % sim_time) + + # Evaluate spike stats and compute firing rates + varTwindow=5 #(sec) + stats_dict, rates_dict , neur_freqIdx= estimate_rates(Y, dt=args.step_size, num_excite=args.num_excite, max_samples=100000, varTwindow=varTwindow, mxNn=5) + + # REMAP MATRICES TO FREQUENCY-SORTED ORDER (PRIMARY INDEX) + neur_revFreqIdx = neur_freqIdx.copy() # freq_sorted_position → natural_index (original from estimate_rates) + neur_freqIdx = np.empty(len(neur_revFreqIdx), dtype=int) # natural_index → freq_sorted_position + neur_freqIdx[neur_revFreqIdx] = np.arange(len(neur_revFreqIdx)) + + # Reorder connectivity matrix and bias vector immediately + A_freq_sorted = A_dale[np.ix_(neur_revFreqIdx, neur_revFreqIdx)] # Reorder both rows and columns + B_freq_sorted = B_idle[neur_revFreqIdx] # Reorder bias vector + + # Store remapped matrices in trueD (primary data structure) + trueD = { 'A_true': A_freq_sorted, 'B_true': B_freq_sorted, 'neur_freqIdx': neur_freqIdx, 'neur_revFreqIdx': neur_revFreqIdx} + + # Create metadata for freq-sorted data (this is now primary) + trueMD = {'dale_conf': dale_conf, 'evol_conf': evol_conf,'short_name':args.dataName,'dale_simu_stats':stats_dict} + + # Prepare spike data for saving - reorder by frequency + Y_uchar = np.clip(Y, 0, 255).astype(np.uint8) + Y_freq_sorted = Y_uchar[:, neur_revFreqIdx] # Reorder neurons by frequency + rates_freq_sorted = rates_dict['single_rates'][neur_revFreqIdx] # Reorder rates by frequency + # Get variance, SNR, and Fano factor from estimate_rates and reorder + single_rates_var_sorted = rates_dict['sigle_rates_var'][neur_revFreqIdx] + single_rates_snr_sorted = rates_dict['sigle_rates_snr'][neur_revFreqIdx] + single_fano_fact_sorted = rates_dict['single_fano_fact'][neur_revFreqIdx] + + spikeD = { + 'spikes': Y_freq_sorted, + 'single_rates': rates_freq_sorted, + 'sigle_rates_var': single_rates_var_sorted, + 'sigle_rates_snr': single_rates_snr_sorted, + 'single_fano_fact': single_fano_fact_sorted + } + spikeMD={ 'short_name':args.dataName,'time_step_sec':args.step_size,'data_type':'simDale', 'var_time_window_sec':varTwindow } + + outFt = os.path.join(args.dataPath, args.dataName + '.simTruth.npz') + write_data_npz(trueD, outFt, metaD=trueMD) + if args.verb>1: pprint(trueMD) + outFs = os.path.join(args.dataPath, args.dataName + '.spikes.npz') + write_data_npz(spikeD, outFs, metaD=spikeMD) + if args.verb>1: pprint(spikeMD) + + print("\nSimulation completed successfully!") + print("\nNext step commands:") + print(" ./view_dalePoisson.py --dataPath $dataPath --dataName %s -p c d a b " % args.dataName) + print(" ./fit_lassoPoisson.py --dataPath $dataPath --dataName %s --num_epochs 50 " % args.dataName) + print(" --dataPath "+args.dataPath) + +if __name__ == '__main__': + main() diff --git a/causal_net/stationaryLag1_ver6_lasso/toolbox b/causal_net/stationaryLag1_ver6_lasso/toolbox new file mode 120000 index 00000000..c9203e9b --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/toolbox @@ -0,0 +1 @@ +../toolbox/ \ No newline at end of file diff --git a/causal_net/stationaryLag1_ver6_lasso/view_bioexp.py b/causal_net/stationaryLag1_ver6_lasso/view_bioexp.py new file mode 100755 index 00000000..560bebe5 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/view_bioexp.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +""" +Visualization tool for biological experiment input features and data quality. + +This script provides comprehensive visualization and analysis of experimental +neural data, focusing on data quality assessment and feature exploration. +Main functionality includes: +- Time series visualization of neural recordings with customizable time ranges +- Data quality metrics and statistical summaries +- Cluster detection and activity pattern analysis +- Interactive plotting with configurable display options + +Used primarily for exploratory data analysis of biological neural recordings +before further processing and connectivity analysis. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os + +from time import time +from pprint import pprint +import numpy as np +from PlotterBioExp import Plotter +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +from UtilBioExp import create_clusters_mask +import argparse +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("-v","--verbosity",type=int, help="increase output verbosity", default=1, dest='verb') + parser.add_argument("-p", "--showPlots", default='a', nargs='+',help="abcd-string listing shown plots") + + + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + + parser.add_argument("--dataPath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for any further data processing") + + parser.add_argument('-T','--time_range' , default=[0., 600], nargs=2, type=float, help='display data time range in seconds') + parser.add_argument('--burst_freq_thres' , default=40., type=float, help='tags high freq channels for burts detection') + parser.add_argument('--burst_chan_thres' , default=30, type=int, help='final thres eliminating time bins') + + parser.add_argument("--dataName", default='HET_80k_1-fc62ef',help='preprocessed session name') + + parser.add_argument('-R','--time_rebin2', default=50, type=int, help='rebin current time axis') + + args = parser.parse_args() + # make arguments more flexible + args.outPath='out/' + args.showPlots=''.join(args.showPlots) + + print( 'myArg-program:',parser.prog) + for arg in vars(args): print( 'myArg:',arg, getattr(args, arg)) + + if args.time_range!=None: assert args.time_range[0] < args.time_range[1] + assert os.path.exists(args.dataPath) + assert os.path.exists(args.outPath) + return args + + +#...!...!.................... +def detect_spike_bursts(spikeD, md): + rateThr2=args.burst_freq_thres + spikeYield = spikeD['spikes'] #, spikeD['single_rates'] + time_step=md['time_step_sec'] + tReb2=args.time_rebin2 + assert tReb2<101 # this would exceed 1 seconds + + #.... rebin time axis + ntime,nchan=spikeYield.shape + if ntime % tReb2 != 0: + ntime_c = ntime - (ntime % tReb2) + spikeYield = spikeYield[:ntime_c] # Clip the data + #print('ccl',ntime_c , ntime ,ntime % tReb2, tReb2) + #... sum yileds over tReb2 time bins + spikeYieldR= np.sum( spikeYield.reshape(-1, tReb2, nchan),axis=1) + time_step=md['time_step_sec'] + time_step2=time_step*tReb2 + #print('rr1',time_step2,spikeYieldR.shape,spikeYield.shape) + + rate2D=spikeYieldR/time_step2 + mask2D=rate2D>rateThr2 + highChan=np.sum(mask2D,axis=1) + + #.... compute running sume over K bins + K = 5 # Number of bins for the running sum + mCnt=args.burst_chan_thres # minimal number of highRate neurons to flag the cluster in time + # Create a kernel for the running sum + kernel = np.ones(K)/K + + # Compute the running sum using convolution + XS = np.convolve(highChan, kernel, mode='valid') + XS = np.pad(XS, (K-1, 0), mode='constant', constant_values=0) + XM=create_clusters_mask(XS,th=mCnt) + usableFrac=1-np.sum(XM)/XM.shape[0] + + rebD={'time_step2':time_step2,'rate_thres2':rateThr2,'smooth_kernel':K,'high_cnt_thres':mCnt,'usable_time_fract':usableFrac} + rebD['rate2D']=rate2D + rebD['mask2D']=mask2D + rebD['highChanCnt']=highChan + rebD['highChanSmooth']=XS + rebD['highChanMask']=XM + rebD['rate1D']=np.sum(rate2D,axis=1) + ntime//=tReb2 + rebD['timeV']= np.linspace(0, (ntime - 1) * time_step2, ntime) + + print('usable time frac:%.3f nchan=%d thr=%.1f Hz'%(rebD['usable_time_fract'],nchan, rateThr2)) + timeMask=np.repeat(XM, tReb2) + return rebD,timeMask + +#...!...!.................... +def filter_bursts(spikeD, md,timeMask): + pprint(md) + spikeD['spikes'][:len(timeMask)][timeMask]=0 + + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__=="__main__": + args=get_parser() + np.set_printoptions(precision=3) + + spikesFF = os.path.join(args.dataPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF) + if args.verb>1: pprint(spikeMD) + rebD,timeMask=detect_spike_bursts(spikeD, spikeMD) + + # ... for plotting + bioFF=spikesFF.replace('spikes','bioExp') + bioD, bioMD = read_data_npz(bioFF) + + #...... WRITE OUTPUT ......... + maskFF=spikesFF.replace('spikes','timeMask') + write_data_npz({'time_mask':timeMask}, maskFF, metaD=None) + filter_bursts(spikeD, spikeMD,timeMask) + + #-------------------------------- + # .... plotting ........ + args.prjName=spikeMD['short_name'] + spikeMD['plot']={} + spikeMD['plot']['time_rangeLR']=np.array(args.time_range) + spikeMD.update(**bioMD) + + plot=Plotter(args) + + if 'a' in args.showPlots: + plot.freq_histo(spikeD,spikeMD,figId=1) + if 'b' in args.showPlots: + plot.freq_vs_time(rebD,spikeMD,figId=2) + + plot.display_all() + print('M:done') + #pprint(expMD) #tmp diff --git a/causal_net/stationaryLag1_ver6_lasso/view_dalePoisson.py b/causal_net/stationaryLag1_ver6_lasso/view_dalePoisson.py new file mode 100755 index 00000000..579bbdf7 --- /dev/null +++ b/causal_net/stationaryLag1_ver6_lasso/view_dalePoisson.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +""" +Visualization tool for simulated Dale Poisson network data. + +This script provides comprehensive visualization of simulated Dale's principle +neural networks with Poisson spiking dynamics. It loads previously generated +simulation data and produces various analysis plots. + +Main functionality includes: +- Dale connectivity matrix visualization (natural and frequency-sorted order) +- Eigenvalue analysis showing network stability properties +- Weight and firing rate distribution histograms +- Neuron type (excitatory/inhibitory) analysis +- Comparison of different network orderings + +Used for post-simulation analysis and validation of Dale network properties +without re-running the full simulation. + +Usage: + ./view_dalePoisson.py --dataName daleM150_448b86 -p a b c +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os +from pprint import pprint +import numpy as np +from PlotterSimPoisson import Plotter +from toolbox.Util_NumpyIO import read_data_npz +import argparse + +#...!...!.................... +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize simulated Dale Poisson network data") + parser.add_argument("-v","--verbosity",type=int, help="increase output verbosity", default=1, dest='verb') + parser.add_argument("-p", "--showPlots", default='a b', nargs='+',help="abc-string listing shown plots: a=Dale_matrix_and_eigen, b=histo_weights_rates, c=histo_weights_rates_byFreq, d=rates_study") + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument("--dataPath",default='/pscratch/sd/b/balewski/2025_causalNet_tmp/',help="head dir for input data") + parser.add_argument("--dataName", default='daleM150_448b86',help='simulated Dale network base name') + parser.add_argument("--outPath", type=str, default=None, help="Output path for plots (defaults to dataPath)") + + args = parser.parse_args() + + # make arguments more flexible + if args.outPath is None: args.outPath = args.dataPath + args.showPlots=''.join(args.showPlots) + + print( 'myArg-program:',parser.prog) + for arg in vars(args): print( 'myArg:',arg, getattr(args, arg)) + + assert os.path.exists(args.dataPath) + return args + +#================================= +#================================= +# M A I N +#================================= +#================================= +if __name__=="__main__": + args=get_parser() + np.set_printoptions(precision=3) + + # Load simulation truth data (Dale matrices, biases, etc.) + truthFF = os.path.join(args.dataPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb>0) + if args.verb>1: + print("\nSimulation Truth Metadata:") + pprint(trueMD) + + # Load spike data (generated spike counts and rates) + spikesFF = os.path.join(args.dataPath, f"{args.dataName}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF, verb=args.verb>0) + if args.verb>1: + print("\nSpike Data Metadata:") + pprint(spikeMD) + + # Merge metadata for plotting + trueMD['short_name'] = args.dataName + + #-------------------------------- + # .... plotting ........ + args.prjName=args.dataName+'_view' + plot=Plotter(args) + + if 'a' in args.showPlots: + plot.Dale_matrix_and_eigen(trueD['A_true'],trueMD,trueD,figId=1) + + if 'b' in args.showPlots: + plot.histo_weights_rates(trueD,spikeD,trueMD,figId=2) + + if 'c' in args.showPlots: + plot.histo_weights_rates(trueD,spikeD,trueMD,byFreq=True,figId=3) + + if 'd' in args.showPlots: + plot.rates_study(trueD,spikeD,trueMD,figId=4) + + if 'e' in args.showPlots: + plot.Dale_matrix_pseudospectra(trueD['A_true'],trueMD,trueD,figId=5) + + plot.display_all() + print('M:done - view_dalePoisson completed successfully!') + diff --git a/causal_net/toolbox/PlotterBackbone.py b/causal_net/toolbox/PlotterBackbone.py new file mode 100644 index 00000000..7235eec5 --- /dev/null +++ b/causal_net/toolbox/PlotterBackbone.py @@ -0,0 +1,111 @@ +import os +import numpy as np +import time + +#...!...!.................. +def roys_fontset(plt): + print('load Roys fontest') + plt.rcParams['axes.spines.right'] = False + plt.rcParams['axes.spines.top'] = False + plt.rcParams['pdf.fonttype'] = 42 + plt.rcParams['ps.fonttype'] = 42 + #plt.rcParams['text.usetex'] = True #Needs new Docker image + + tick_major = 6 + tick_minor = 4 + plt.rcParams["xtick.major.size"] = tick_major + plt.rcParams["xtick.minor.size"] = tick_minor + plt.rcParams["ytick.major.size"] = tick_major + plt.rcParams["ytick.minor.size"] = tick_minor + + font_small = 12 + font_medium = 13 + font_large = 14 + plt.rc('font', size=font_small) # controls default text sizes + plt.rc('axes', titlesize=font_medium) # fontsize of the axes title + plt.rc('axes', labelsize=font_medium) # fontsize of the x and y labels + plt.rc('xtick', labelsize=font_small) # fontsize of the tick labels + plt.rc('ytick', labelsize=font_small) # fontsize of the tick labels + + plt.rc('figure', titlesize=font_large) # fontsize of the figure title + + # legend box + plt.rc('legend', frameon=False) # remove it the frame + plt.rc('legend', fontsize=font_small) # legend fontsize + +#............................ +#............................ +#............................ +class PlotterBackbone(object): + def __init__(self, args): + self.jobName=args.prjName + try: + self.venue=args.formatVenue + except: + self.venue='prod' + + import matplotlib as mpl + if args.noXterm: + if args.verb>0: print('disable Xterm') + mpl.use('Agg') # to plot w/o X-server + else: + mpl.use('TkAgg') + import matplotlib.pyplot as plt + + if args.verb>0: print(self.__class__.__name__,':','Graphics started') + plt.close('all') + self.plt=plt + self.args=args + self.figL=[] + self.outPath=args.outPath+'/' + assert os.path.exists(self.outPath) + if 'paper' in self.venue: + roys_fontset(plt) + + #............................ + def figId2name(self, fid): + figName='%s_f%d'%(self.jobName,fid) + return figName + + #............................ + def clear(self): + self.figL=[] + self.plt.close('all') + #............................ + def display_all(self, png=1): + if len(self.figL)<=0: + print('display_all - nothing top plot, quit') + return + + for fid in self.figL: + self.plt.figure(fid) + self.plt.tight_layout() + figName=self.outPath+self.figId2name(fid) + if png: figName+='.png' + else: figName+='.pdf' + print('Graphics display ',figName) + self.plt.savefig(figName) + self.plt.show() + +# figId=self.smart_append(figId) +#...!...!.................... + def smart_append(self,id): # increment id if re-used + while id in self.figL: id+=1 + self.figL.append(id) + return id + +#............................ + def blank_share2D(self,nrow=2,ncol=2, figsize=(6,6),figId=10): + figId=self.smart_append(figId) + kwargs={'num':figId,'facecolor':'white', 'figsize':figsize} + fig, axs = self.plt.subplots(nrow,ncol, sharex='col', sharey='row', + gridspec_kw={'hspace': 0, 'wspace': 0}, **kwargs) + return axs + +#............................ + def blank_separate2D(self,nrow=2,ncol=2, figsize=(6,6),figId=10): + figId=self.smart_append(figId) + kwargs={'num':figId,'facecolor':'white', 'figsize':figsize} + fig, axs = self.plt.subplots(nrow,ncol, **kwargs) + return axs + diff --git a/causal_net/toolbox/Util_H5io4.py b/causal_net/toolbox/Util_H5io4.py new file mode 100755 index 00000000..d32ae1e1 --- /dev/null +++ b/causal_net/toolbox/Util_H5io4.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +''' = = = = = HD5 advanced storage = = = +It can hold: +* python dictionaries which must pass: json.dumps(dict) +* single float or int variables w/o np-array packing. It is recovered as 1-value array +* arbitrary numpy array (a large size payloads) + - for an array of arbitrary strings must declare dtype='object' at write and use .decode("utf-8") to unpack +''' + +import numpy as np +import h5py, time, os +import json,time +from pprint import pprint + +#...!...!.................. +def write4_data_hdf5(dataD,outF,metaD=None,verb=1): + assert type(dataD)!=type(None) + assert len(outF)>0 + + if metaD!=None: + #pprint(metaD) + metaJ=json.dumps(metaD, default=str) + #print('meta.JSON:',metaJ) + dataD['meta.JSON']=metaJ + + dtvs = h5py.special_dtype(vlen=str) + h5f = h5py.File(outF, 'w') + if verb>0: + print('saving data as hdf5:',outF) + start = time.time() + for item in dataD: + rec=dataD[item] + if verb>1: print('x=',item,type(rec)) + if type(rec)==str: # special case + dset = h5f.create_dataset(item, (1,), dtype=dtvs) + dset[0]=rec + if verb>0:print('h5-write :',item, 'as string',dset.shape,dset.dtype) + continue + if type(rec)!=np.ndarray: # packs a single value into np-array + rec=np.array([rec]) + + h5f.create_dataset(item, data=rec) + if verb>0:print('h5-write :',item, rec.shape,rec.dtype) + h5f.close() + xx=os.path.getsize(outF)/1048576 + print('closed hdf5:',outF,' size=%.2f MB, elaT=%.1f sec'%(xx,(time.time() - start))) + + +#...!...!.................. +def read4_data_hdf5(inpF,verb=1): + if verb>0: + print('read data from hdf5:',inpF) + start = time.time() + h5f = h5py.File(inpF, 'r') + objD={} + for x in h5f.keys(): + if verb>1: print('\nitem=',x,type(h5f[x]),h5f[x].shape,h5f[x].dtype) + #if x in ['calTag','dataFile','date'] : continue + if h5f[x].dtype==object: + obj=h5f[x][:] + #print('bbb',type(obj),obj.dtype) + if verb>0: print('read str:',x,len(obj),type(obj)) + else: + obj=h5f[x][:] + if verb>0: print('read obj:',x,obj.shape,obj.dtype) + objD[x]=obj + try: + inpMD=json.loads(objD.pop('meta.JSON')[0]) + if verb>1: print(' recovered meta-data with %d keys'%len(inpMD)) + except: + inpMD=None + if verb>0: + print(' done h5, num rec:%d elaT=%.1f sec'%(len(objD),(time.time() - start))) + + h5f.close() + return objD,inpMD + + + +#================================= +#================================= +# U N I T T E S T +#================================= +#================================= + +if __name__=="__main__": + print('testing h5IO ver 3') + outF='abcTest.h5' + verb=1 + + var1=float(15) # single variable + one=np.zeros(shape=5,dtype=np.int16); one[3]=3 + two=np.zeros(shape=(2,3)); two[1,2]=4 + + three=np.empty((2), dtype='object') + three[0]='record aaaa' + three[1]='much longer record bbb' + # WARN: all decalred elements of three[] must be initialized before writeing H5 + + # this works too: + # three=np.array(['record aaaa','much longer record bbb'], dtype='object') + + text='This is text1' + + metaD={"age":17,"dom":"white","dates":[11,22,33]} + + outD={'one':one,'two':two,'var1':var1,'atext':text,'three':three} + + write4_data_hdf5(outD,outF,metaD=metaD,verb=verb) + + print('\nM: ***** verify by reading it back from',outF) + big,meta2=read4_data_hdf5(outF,verb=verb) + from pprint import pprint + print(' recovered meta-data'); pprint(meta2) + print('dump read-in data') + for x in big: + print('\nkey=',x); pprint(big[x]) + + #decode one string from string-array + rec2=big['three'][1].decode("utf-8") + print('rec2:',type(rec2),rec2) + print('\n check raw content: h5dump %s\n'%outF) diff --git a/causal_net/toolbox/Util_IOfunc.py b/causal_net/toolbox/Util_IOfunc.py new file mode 100644 index 00000000..da3a3f17 --- /dev/null +++ b/causal_net/toolbox/Util_IOfunc.py @@ -0,0 +1,107 @@ +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import numpy as np +import time, os +#import ruamel.yaml as yaml +import yaml +#import warnings +#warnings.simplefilter('ignore', yaml.error.MantissaNoDotYAML1_1Warning) + +from pprint import pprint +import csv + +#...!...!.................. +def read_yaml(ymlFn,verb=1): + if verb: print(' read yaml:',ymlFn,end='') + start = time.time() + ymlFd = open(ymlFn, 'r') + bulk=yaml.load( ymlFd, Loader=yaml.CLoader) + + ymlFd.close() + if verb>1: print(' done elaT=%.1f sec'%(time.time() - start)) + else: print() + return bulk + +#...!...!.................. +def write_yaml(rec,ymlFn,verb=1): + start = time.time() + ymlFd = open(ymlFn, 'w') + yaml.dump(rec, ymlFd, Dumper=yaml.CDumper) + ymlFd.close() + xx=os.path.getsize(ymlFn)/1024 + if verb: + print(' closed yaml:',ymlFn,' size=%.1f kB'%xx,' elaT=%.1f sec'%(time.time() - start)) + + +#...!...!.................. +def read_one_csv(fname,delim=','): + print('read_one_csv:',fname) + tabL=[] + with open(fname) as csvfile: + drd = csv.DictReader(csvfile, delimiter=delim) + print('see %d columns'%len(drd.fieldnames),drd.fieldnames) + for row in drd: + tabL.append(row) + + print('got %d rows \n'%(len(tabL))) + #print('LAST:',row) + return tabL,drd.fieldnames + +#...!...!.................. +def write_one_csv(fname,rowL,colNameL): + print('write_one_csv:',fname) + print('export %d columns'%len(colNameL), colNameL) + with open(fname,'w') as fou: + dw = csv.DictWriter(fou, fieldnames=colNameL)#, delimiter='\t' + dw.writeheader() + for row in rowL: + dw.writerow(row) + + +#...!...!.................. +def expand_dash_list(inpL): + # expand list if '-' are present + outL=[] + for x in inpL: + if '-' not in x: + outL.append(x) ; continue + # must contain string: xxxx[n1-n2] + ab,c=x.split(']') + assert len(ab)>3 + a,b=ab.split('[') + print('abc',a,b,c) + nL=b.split('-') + for i in range(int(nL[0]),int(nL[1])+1): + outL.append('%s%d%s'%(a,i,c)) + print('EDL:',inpL,' to ',outL) + return outL +''' - - - - - - - - - +Offset-aware time, usage: + +*) get current date: +t1=time.localtime() + +*) convert to string: +timeStr=dateT2Str(t1) + +*) revert to struct_time +t2=dateStr2T(timeStr) + +*) compute difference in sec: +t3=time.localtime() +delT=time.mktime(t3) - time.mktime(t1) +delSec=delT.total_seconds() +or delT is already in seconds +''' + +#...!...!.................. +def dateT2Str(xT): # --> string + nowStr=time.strftime("%Y%m%d_%H%M%S_%Z",xT) + return nowStr + +#...!...!.................. +def dateStr2T(xS): # --> datetime + yT = time.strptime(xS,"%Y%m%d_%H%M%S_%Z") + return yT + diff --git a/causal_net/toolbox/Util_NumpyIO.py b/causal_net/toolbox/Util_NumpyIO.py new file mode 100755 index 00000000..a1d4d689 --- /dev/null +++ b/causal_net/toolbox/Util_NumpyIO.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +''' = = = = = Numpy NPZ advanced storage = = = +It can hold: +* python dictionaries which must pass: json.dumps(dict) +* single float or int variables w/o np-array packing. It is recovered as 1-value array +* arbitrary numpy array (a large size payloads) + - for an array of arbitrary strings must declare dtype='object' at write and use .decode("utf-8") to unpack +''' + +import numpy as np +import time, os +import json,time +from pprint import pprint + +#...!...!.................. +def write_data_npz(dataD,outF,metaD=None,verb=1): + assert type(dataD)!=type(None) + assert len(outF)>0 + start = time.time() + + # Create a copy to avoid modifying original data + saveD = dataD.copy() + + if metaD!=None: + #pprint(metaD) + metaJ=json.dumps(metaD, default=str) + #print('meta.JSON:',metaJ) + saveD['meta.JSON']=np.array([metaJ], dtype='object') + + if verb>0: + print('saving data as npz:',outF) + + # Process data to ensure all items are numpy arrays + for item in list(saveD.keys()): + rec=saveD[item] + if verb>1: print('x=',item,type(rec)) + if type(rec)==str: # special case - convert string to object array + saveD[item] = np.array([rec], dtype='object') + if verb>0:print('npz-write :',item, 'as string array',saveD[item].shape,saveD[item].dtype) + continue + if type(rec)!=np.ndarray: # packs a single value into np-array + saveD[item]=np.array([rec]) + if verb>0:print('npz-write :',item, saveD[item].shape,saveD[item].dtype) + else: + if verb>0:print('npz-write :',item, rec.shape,rec.dtype) + + # Save to NPZ format + np.savez_compressed(outF, **saveD) + + xx=os.path.getsize(outF)/1048576 + if verb>0: + print('closed npz:',outF,' size=%.2f MB, elaT=%.1f sec'%(xx,(time.time() - start))) + + +#...!...!.................. +def read_data_npz(inpF,verb=1): + if verb>0: + print('read data from npz:',inpF) + start = time.time() + + npzData = np.load(inpF, allow_pickle=True) + objD={} + + for x in npzData.files: + obj = npzData[x] + if verb>1: print('\nitem=',x,type(obj),obj.shape,obj.dtype) + + if obj.dtype==object: + if verb>0: print('read obj:',x,len(obj),type(obj)) + else: + if verb>0: print('read obj:',x,obj.shape,obj.dtype) + objD[x]=obj + + npzData.close() + + try: + inpMD=json.loads(objD.pop('meta.JSON')[0]) + if verb>1: print(' recovered meta-data with %d keys'%len(inpMD)) + except: + inpMD=None + if verb>0: + print(' done npz, num rec:%d elaT=%.1f sec'%(len(objD),(time.time() - start))) + + return objD,inpMD + + + +#================================= +#================================= +# U N I T T E S T +#================================= +#================================= + +if __name__=="__main__": + print('testing npzIO ver 1') + outF='abcTest.npz' + verb=1 + + var1=float(15) # single variable + one=np.zeros(shape=5,dtype=np.int16); one[3]=3 + two=np.zeros(shape=(2,3)); two[1,2]=4 + + three=np.empty((2), dtype='object') + three[0]='record aaaa' + three[1]='much longer record bbb' + # WARN: all decalred elements of three[] must be initialized before writeing NPZ + + # this works too: + # three=np.array(['record aaaa','much longer record bbb'], dtype='object') + + text='This is text1' + + metaD={"age":17,"dom":"white","dates":[11,22,33]} + + outD={'one':one,'two':two,'var1':var1,'atext':text,'three':three} + + write_data_npz(outD,outF,metaD=metaD,verb=verb) + + print('\nM: ***** verify by reading it back from',outF) + big,meta2=read_data_npz(outF,verb=verb) + from pprint import pprint + print(' recovered meta-data'); pprint(meta2) + print('dump read-in data') + for x in big: + print('\nkey=',x); pprint(big[x]) + + #get one string from string-array + rec2=big['three'][1] + print('rec2:',type(rec2),rec2) + print('\n check raw content with: python -c "import numpy as np; data=np.load(\'%s\', allow_pickle=True); print(list(data.files)); data.close()"\n'%outF) diff --git a/causal_net/topoAna_ver4/PlotterTopoCorrel.py b/causal_net/topoAna_ver4/PlotterTopoCorrel.py new file mode 100644 index 00000000..3d5d6222 --- /dev/null +++ b/causal_net/topoAna_ver4/PlotterTopoCorrel.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +""" +Plotting utilities for topology-correlation inspection. +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from matplotlib.lines import Line2D +import numpy as np + +from toolbox.PlotterBackbone import PlotterBackbone + + +METRIC_PLOT_ORDER = ( + ("assortativity", "r_assort", "Directed assortativity"), + ("jaccard_index", "jaccard", "Mean Jaccard"), + ("cycle_decay", "cycle_decay", "Cycle decay"), + ("homology_k1", "hom_k1", "Homology k1"), +) + + +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self, args) + + def metrics_inspector(self, harvestD, nameTemplate, figId=1): + figId = self.smart_append(figId) + fig, axL = self.plt.subplots(1, 4, num=figId, facecolor="white", figsize=(18, 4.6)) + + colorL = ("tab:blue", "tab:orange", "tab:green", "tab:red", "tab:purple", "tab:brown", "tab:pink", "tab:gray") + markerL = ("o", "s", "^", "D", "v", "P", "X", "<") + + placement_LL = sorted({int(placement_L) for per_metric in harvestD.values() for placement_L in per_metric.keys()}) + styleD = {} + for idx, placement_L in enumerate(placement_LL): + styleD[placement_L] = { + "color": colorL[idx % len(colorL)], + "marker": markerL[idx % len(markerL)], + } + + for ax, (metric_name, y_label, title_txt) in zip(axL, METRIC_PLOT_ORDER): + metricD = harvestD.get(metric_name, {}) + for placement_L in sorted(metricD.keys()): + recL = metricD[placement_L] + if not recL: + continue + xV = np.asarray([rec[0] for rec in recL], dtype=float) + yV = np.asarray([rec[1] for rec in recL], dtype=float) + sty = styleD[placement_L] + ax.plot( + xV, + yV, + linestyle="none", + marker=sty["marker"], + color=sty["color"], + markersize=6.5, + markeredgewidth=0.7, + markeredgecolor="black", + alpha=0.9, + ) + + ax.set_title(title_txt) + ax.set_xlabel("placement_ker_delta") + ax.set_ylabel(y_label) + ax.grid(True, alpha=0.35) + + if metricD: + finite_y = [] + for recL in metricD.values(): + for _, y_val, _ in recL: + if np.isfinite(y_val): + finite_y.append(y_val) + if finite_y: + y_min = min(finite_y) + y_max = max(finite_y) + if y_min == y_max: + pad = 0.05 if y_min == 0 else 0.08 * abs(y_min) + ax.set_ylim(y_min - pad, y_max + pad) + + leg_handles = [ + Line2D( + [0], + [0], + linestyle="none", + marker=styleD[placement_L]["marker"], + color=styleD[placement_L]["color"], + markeredgecolor="black", + markeredgewidth=0.7, + markersize=7, + label="L=%d" % placement_L, + ) + for placement_L in placement_LL + ] + if leg_handles: + fig.legend(handles=leg_handles, loc="upper center", bbox_to_anchor=(0.5, 0.965), ncol=min(6, len(leg_handles))) + + fig.text(0.5, 0.995, "nameTemplate: %s" % nameTemplate, ha="center", va="top", fontsize=16) + fig.subplots_adjust(top=0.78, wspace=0.32) diff --git a/causal_net/topoAna_ver4/big_daleGenV1.sh b/causal_net/topoAna_ver4/big_daleGenV1.sh new file mode 100755 index 00000000..b49fd9d8 --- /dev/null +++ b/causal_net/topoAna_ver4/big_daleGenV1.sh @@ -0,0 +1,25 @@ +#!/bin/bash +set -u ; # exit if you try to use an uninitialized variable +set -e ; # bash exits if any statement returns a non-true return value + +#stop-me + +basePath="/pscratch/sd/b/balewski/2026_causalNet_tmp4/" + +N=100 +L=1 + +for dker in 0.1 0.5 1.0 1.5 2.0 2.5; do +#for dker in 0.1 1.0 2.0 2.5; do +#for dker in 2.5 ; do + + expName=aN${N}_L${L}_dker${dker} + #echo expName=$expName + + #./gen_daleMatrices4.py --spike_model B --time_kernel_q_tau 1.5 0.4 --basePath $basePath --placement_H_L_delta 1 $L $dker --num_neurons 200 --num_steps 8001 --dataName $expName + ./topoAna_daleMatrix4.py --basePath $basePath --dataName $expName |grep \#V +done + +echo all DONE + + diff --git a/causal_net/topoAna_ver4/big_daleGenV2.sh b/causal_net/topoAna_ver4/big_daleGenV2.sh new file mode 100755 index 00000000..e375cc1a --- /dev/null +++ b/causal_net/topoAna_ver4/big_daleGenV2.sh @@ -0,0 +1,23 @@ +#!/bin/bash +set -u ; # exit if you try to use an uninitialized variable +set -e ; # bash exits if any statement returns a non-true return value + +#stop-me + +basePath="/pscratch/sd/b/balewski/2026_causalNet_tmp4topo/" + +N=200 +L=10 # 1,2,5,10,20 + +for dker in $(seq -f "%.2f" 0.10 0.05 2.80); do + + expName=aN${N}_L${L}_dker${dker} + #echo expName=$expName + + #./gen_daleMatrices4.py --spike_model A --basePath $basePath --placement_H_L_delta 1 $L $dker --num_neurons $N --num_steps 1001 --dataName $expName + ./topoAna_daleMatrix4.py --basePath $basePath --dataName $expName |grep \#V + #break +done + +echo all DONE + diff --git a/causal_net/topoAna_ver4/docs/buildTopo.sh b/causal_net/topoAna_ver4/docs/buildTopo.sh new file mode 100755 index 00000000..cdac2ec5 --- /dev/null +++ b/causal_net/topoAna_ver4/docs/buildTopo.sh @@ -0,0 +1,25 @@ +#!/bin/bash + +set -e + +doc_name=network_ver4_topoDiscovery + +is_ubuntu=0 +if [ -f /etc/os-release ] && grep -qi 'ubuntu' /etc/os-release; then + is_ubuntu=1 +fi + +if [ "$is_ubuntu" -eq 1 ]; then + pdflatex "$doc_name" + pdflatex "$doc_name" + open "${doc_name}.pdf" +else + if ! command -v latex >/dev/null 2>&1; then + module load texlive + fi + latex "$doc_name" + latex "$doc_name" + pdflatex "$doc_name" # to produce .pdf for github + + xdvi -s 5 "${doc_name}.dvi" +fi diff --git a/causal_net/topoAna_ver4/docs/network_ver4_topoDiscovery.pdf b/causal_net/topoAna_ver4/docs/network_ver4_topoDiscovery.pdf new file mode 100644 index 00000000..8b75681d Binary files /dev/null and b/causal_net/topoAna_ver4/docs/network_ver4_topoDiscovery.pdf differ diff --git a/causal_net/topoAna_ver4/docs/network_ver4_topoDiscovery.tex b/causal_net/topoAna_ver4/docs/network_ver4_topoDiscovery.tex new file mode 100644 index 00000000..e4e00e2d --- /dev/null +++ b/causal_net/topoAna_ver4/docs/network_ver4_topoDiscovery.tex @@ -0,0 +1,477 @@ +\documentclass[11pt]{article} +\usepackage[utf8]{inputenc} +\usepackage{amsmath} +\usepackage{amsfonts} +\usepackage{amssymb} +\usepackage{bm} +\usepackage{geometry} +\usepackage{xcolor} +\usepackage{booktabs} +\usepackage{hyperref} + +\geometry{margin=1.0in} + +\title{Inferring Spatial Connectivity Decay Exponent\\ +from Graph Topology Alone} +\author{Jan Balewski, NERSC} +\date{\today} + +\begin{document} + +\maketitle + +\begin{abstract} +A spatially embedded neuronal network is constructed by sampling +connections with probability proportional to +$D_{ij}^{-\delta_{\mathrm{ker}}}$, where $D_{ij}$ is the Euclidean +distance between neurons $i$ and $j$ and +$\delta_{\mathrm{ker}}\in[0,3]$ is a real-valued decay exponent. +The limiting cases interpolate between uniform random connectivity +($\delta_{\mathrm{ker}}=0$) and strongly local, lattice-like wiring +($\delta_{\mathrm{ker}}\to 3$). We ask: given only the binary +adjacency matrix~$G$ --- without access to neuron positions or +distances --- can one estimate~$\delta_{\mathrm{ker}}$? +We present five purely graph-theoretic diagnostics, each producing a +single scalar that varies monotonically with $\delta_{\mathrm{ker}}$ +and therefore serves as a proxy for the hidden spatial decay. +The common principle is that large $\delta_{\mathrm{ker}}$ yields a +lattice-like graph (high local clustering, assortative degrees, low +spectral dimension), while small $\delta_{\mathrm{ker}}$ yields a +small-world or Erd\H{o}s--R\'enyi-like graph (diffuse neighborhoods, +weaker clustering, higher effective dimensionality). +\end{abstract} + +\tableofcontents +\newpage + +% =============================================================== +\section{Setup and Notation} +\label{sec:setup} +% =============================================================== + +\subsection{The Generative Model} +\label{sec:gen_model} + +Consider $N$ neurons placed on a regular two-dimensional grid within a +rectangular domain $[0,L]\times[0,H]$ with minimum spacing +$d_{\min}>0$. After random subsampling of grid sites and +lexicographic sorting, each neuron~$j$ occupies a position +$\mathbf{P}_j=(x_j,y_j)$. The pairwise Euclidean distance matrix is +$D_{ij}=\|\mathbf{P}_i-\mathbf{P}_j\|_2$. + +Connections are drawn according to a distance-dependent affinity +kernel. For every ordered pair $(i,j)$ with $i\neq j$, define the +affinity +\begin{equation} + V_{ij} = D_{ij}^{-\delta_{\mathrm{ker}}}, + \qquad i\neq j;\qquad V_{ii}=0, + \label{eq:affinity} +\end{equation} +where $\delta_{\mathrm{ker}}\in[0,3]$ is a real-valued \emph{decay +exponent} that controls how sharply connection probability falls off +with distance. Because grid nodes are separated by at least +$d_{\min}>0$, the affinity is everywhere finite. + +The exponent $\delta_{\mathrm{ker}}$ interpolates continuously between +qualitatively distinct connectivity regimes: + +\begin{center} +\renewcommand{\arraystretch}{1.3} +\begin{tabular}{cl} +\toprule +$\delta_{\mathrm{ker}}$ & Character of the resulting graph \\ +\midrule +$0$ & Erd\H{o}s--R\'enyi: connection probability independent of + distance \\ +$\lesssim 1$ & Weak spatial bias; long-range edges abundant; + small-world-like \\ +$\approx 1$ & Moderate locality; intermediate clustering \\ +$\approx 2$ & Strong locality; neighborhood structure clearly visible \\ +$\gtrsim 2$ & Very tight spatial coupling; lattice-like communities \\ +$3$ & Extreme locality; connections almost exclusively between + nearest neighbors \\ +\bottomrule +\end{tabular} +\end{center} + +\noindent +Each presynaptic neuron~$j$ independently samples an integer +out-degree $k_j\sim\mathrm{Uniform}(k_{\min},k_{\max})$ and draws +$k_j$ distinct postsynaptic targets $\mathcal{T}_j\subset +\{1,\dots,N\}\setminus\{j\}$ without replacement from the normalized +distribution +\begin{equation} + p_{ij} + = \frac{V_{ij}}{\displaystyle\sum_{i'\neq j}V_{i'j}} + = \frac{D_{ij}^{-\delta_{\mathrm{ker}}}} + {\displaystyle\sum_{i'\neq j}D_{i'j}^{-\delta_{\mathrm{ker}}}}, + \qquad i\neq j. + \label{eq:conn_prob} +\end{equation} + +\subsection{The Inference Problem} +\label{sec:inference_problem} + +We are given only the weighted connectivity matrix +$A_{\mathrm{true}}\in\mathbb{R}^{N\times N}$. Since $A$-matrix represents interaction between neurons which can be either excitatory or inhibitory, the values of $A_{ij}$ can be either positive or negative. + +Neuron positions +$\mathbf{P}$, pairwise distances~$D$, and the value of +$\delta_{\mathrm{ker}}$ are all \emph{unknown}. From +$A_{\mathrm{true}}$ we extract the binary adjacency matrix +\begin{equation} + G_{ij} = + \begin{cases} + 1 & A_{\mathrm{true},ij}^{\mathrm{off}}\neq 0,\;i\neq j,\\ + 0 & \text{otherwise}, + \end{cases} + \label{eq:adjacency} +\end{equation} +where $A^{\mathrm{off}}$ denotes the off-diagonal part of +$A_{\mathrm{true}}$ (self-loops are excluded). The goal is to +estimate the decay exponent $\delta_{\mathrm{ker}}$ from~$G$ alone, or +at minimum to construct scalar diagnostics that vary monotonically +with $\delta_{\mathrm{ker}}$ and therefore distinguish, e.g., +$\delta_{\mathrm{ker}}=1$ from $\delta_{\mathrm{ker}}=2$. + +The key qualitative insight is that as $\delta_{\mathrm{ker}}$ +increases, the graph becomes more spatially clustered: edges are +shorter, neighborhoods overlap more, the graph looks more like a +lattice. Conversely, as $\delta_{\mathrm{ker}}$ decreases toward +zero, edges span the full domain, neighborhoods are diffuse, and the +graph approaches an Erd\H{o}s--R\'enyi random graph. Each of the five +methods below detects a different combinatorial or spectral +consequence of this transition. + +\subsection{Additional Notation} + +Throughout, we use the following derived quantities from $G$: +\begin{equation} + k_i^{\mathrm{in}} = \sum_{j} G_{ij}, \qquad + k_j^{\mathrm{out}} = \sum_{i} G_{ij}, \qquad + |\mathcal{E}| = \sum_{i,j} G_{ij}, +\end{equation} +denoting the in-degree and out-degree of each neuron and the total +number of directed edges. The symmetrized adjacency is +$\tilde{G}=(G+G^{\!\top})/2$ with degree matrix +$\tilde{D}_{ii}=\sum_j\tilde{G}_{ij}$. Neighborhoods are defined as +$\mathcal{N}_{\mathrm{in}}(i)=\{j:G_{ij}=1\}$ and +$\mathcal{N}_{\mathrm{out}}(j)=\{i:G_{ij}=1\}$. + + +% =============================================================== +\section{Method 1: Degree Assortativity} +\label{sec:assortativity} +% =============================================================== + +\subsection{Definition} + +The Newman assortativity coefficient measures the Pearson correlation +between the degrees at either end of an edge. For the directed graph +$G$, define the in-degree of the target and the in-degree of the +source for each edge $(i,j)\in\mathcal{E}$ (meaning $G_{ij}=1$, +i.e., neuron~$j$ drives neuron~$i$): +\begin{equation} + r = \frac{ + |\mathcal{E}|^{-1}\!\displaystyle\sum_{(i,j)\in\mathcal{E}} + k_i^{\mathrm{in}}\,k_j^{\mathrm{in}} + \;-\; + \left[|\mathcal{E}|^{-1}\!\displaystyle\sum_{(i,j)\in\mathcal{E}} + k_i^{\mathrm{in}}\right] + \left[|\mathcal{E}|^{-1}\!\displaystyle\sum_{(i,j)\in\mathcal{E}} + k_j^{\mathrm{in}}\right] + }{ + \sigma_{\mathrm{in,target}}\;\sigma_{\mathrm{in,source}} + }, + \label{eq:assortativity} +\end{equation} +where $\sigma_{\mathrm{in,target}}$ and $\sigma_{\mathrm{in,source}}$ are the +standard deviations of $\{k_i^{\mathrm{in}}\}_{(i,j)\in\mathcal{E}}$ +and $\{k_j^{\mathrm{in}}\}_{(i,j)\in\mathcal{E}}$ over all edges. + +\subsection{Dependence on $\delta_{\mathrm{ker}}$} + +When $\delta_{\mathrm{ker}}$ is large, each neuron's targets are +concentrated in a tight spatial neighborhood. A neuron with high +in-degree sits in a locally dense region and receives input +predominantly from nearby neurons, which themselves receive input from +the same dense region --- hence they too have high in-degree. +High-degree nodes connect to high-degree nodes, producing +\emph{assortative} mixing ($r>0$). + +As $\delta_{\mathrm{ker}}$ decreases, edges reach farther, a neuron's +targets span diverse spatial neighborhoods with diverse degree +profiles, and the degree--degree correlation weakens. In the limit +$\delta_{\mathrm{ker}}\to 0$ (Erd\H{o}s--R\'enyi), degrees are +independent of position and $r\to 0$. + +The assortativity coefficient thus increases monotonically with +$\delta_{\mathrm{ker}}$: +\begin{equation} + \delta_{\mathrm{ker}}^{(1)} < \delta_{\mathrm{ker}}^{(2)} + \quad\Longrightarrow\quad + r^{(1)} \lesssim r^{(2)}. + \label{eq:r_monotone} +\end{equation} + + +% =============================================================== +\section{Method 2: Common-Neighbor Overlap (Jaccard Index)} +\label{sec:jaccard} +% =============================================================== + +\subsection{Definition} + +For each directed edge $j\to i$ (i.e.\ $G_{ij}=1$), the Jaccard +overlap between the in-neighborhood of the target and the +out-neighborhood of the source is +\begin{equation} + J_{ij} + = \frac{|\mathcal{N}_{\mathrm{in}}(i) + \cap \mathcal{N}_{\mathrm{out}}(j)|} + {|\mathcal{N}_{\mathrm{in}}(i) + \cup \mathcal{N}_{\mathrm{out}}(j)|}. + \label{eq:jaccard} +\end{equation} +This answers a purely combinatorial question: \emph{if $j$ drives $i$, +how many of $j$'s other targets also drive~$i$?} The mean Jaccard +index over all edges is +\begin{equation} + \bar{J} + = \frac{1}{|\mathcal{E}|} + \sum_{(i,j)\in\mathcal{E}} J_{ij}. + \label{eq:jaccard_mean} +\end{equation} + +\subsection{Dependence on $\delta_{\mathrm{ker}}$} + +The Jaccard index is the graph-theoretic analog of spatial locality +without ever defining locality explicitly. When +$\delta_{\mathrm{ker}}$ is large, each neuron's fan-out lands in a +compact spatial neighborhood; two neurons connected by an edge are +necessarily close, so their neighborhoods overlap substantially and +$J_{ij}$ is large. When $\delta_{\mathrm{ker}}$ is small, fan-out is +diffuse, common neighbors between any linked pair become rare, and +$\bar{J}$ shrinks. In the Erd\H{o}s--R\'enyi limit +($\delta_{\mathrm{ker}}=0$) with edge density $p$, the expected +Jaccard index scales as $\bar{J}\sim p/(2-p)$, independent of +spatial structure. + +The mean Jaccard index therefore increases monotonically with +$\delta_{\mathrm{ker}}$: +\begin{equation} + \delta_{\mathrm{ker}}^{(1)} < \delta_{\mathrm{ker}}^{(2)} + \quad\Longrightarrow\quad + \bar{J}^{(1)} \lesssim \bar{J}^{(2)}. + \label{eq:J_monotone} +\end{equation} + +% =============================================================== +\section{Method 3: Cycle Closure Decay Rate} +\label{sec:cycles} +\sectionmark{Method 3} +% =============================================================== + +\subsection{Definition} + +Powers of the adjacency matrix count directed walks: $(G^k)_{ij}$ is +the number of distinct walks of length~$k$ from~$j$ to~$i$. In +particular, $\operatorname{tr}(G^{k+1})$ counts closed walks +(directed cycles) of length $k{+}1$, while +$\mathbf{1}^{\!\top}G^k\mathbf{1}=\sum_{i,j}(G^k)_{ij}$ counts all +walks of length~$k$. Define the \emph{$k$-step transitivity} as the +fraction of length-$k$ walks whose endpoints are connected by a +direct edge (thus closing into a $(k{+}1)$-cycle): +\begin{equation} + \mathcal{T}_k + = \frac{\operatorname{tr}(G^{k+1})} + {\mathbf{1}^{\!\top}G^k\,\mathbf{1}}, + \qquad k=1,2,\dots + \label{eq:transitivity_k} +\end{equation} +At $k=1$ this reduces (up to normalization) to the directed +clustering coefficient. The \emph{cycle closure decay rate} is the +log-slope of $\mathcal{T}_k$ over a short range: +\begin{equation} + \xi + = -\frac{d\log\mathcal{T}_k}{dk} + \bigg|_{k=2,\dots,K}, + \label{eq:xi} +\end{equation} +estimated by linear regression of $\log\mathcal{T}_k$ on~$k$ for +$k=2,\dots,K$ (typically $K=5$ or $6$). + +\subsection{Dependence on $\delta_{\mathrm{ker}}$} + +For a graph embedded in two dimensions, the existence of short cycles +is a signature of spatial locality: a walk that stays in a compact +region is far more likely to return to its origin than one that +wanders across the domain. When $\delta_{\mathrm{ker}}$ is large, +neighborhoods are tight and densely interconnected, so walks of +moderate length close into cycles readily --- $\mathcal{T}_k$ remains +high across $k$ and $\xi$ is small. + +When $\delta_{\mathrm{ker}}$ is small, long-range edges carry walks +far from their starting point. A two-step walk $j\to i\to i'$ may +traverse the full domain, and the probability that $i'$ connects back +to~$j$ is no higher than the background edge density. Cycle closure +becomes increasingly unlikely with each additional step, so +$\mathcal{T}_k$ drops steeply and $\xi$ is large. + +The cycle closure decay rate therefore increases monotonically as +$\delta_{\mathrm{ker}}$ decreases: +\begin{equation} + \delta_{\mathrm{ker}}^{(1)} < \delta_{\mathrm{ker}}^{(2)} + \quad\Longrightarrow\quad + \xi^{(1)} \gtrsim \xi^{(2)}. + \label{eq:xi_monotone} +\end{equation} + +% =============================================================== +% =============================================================== +\section{Method 4: Persistent Homology} +\label{sec:persistent_homology} +\sectionmark{Method 4} +% =============================================================== + +This metric measures the summed lifetime of all loop features in the graph while sweeping from strong to weak edge weights. + +We probe the topological complexity of the graph by computing its +persistent homology. The symmetrized adjacency matrix +$\tilde{G}_{ij}=\tfrac{1}{2}(G_{ij}+G_{ji})$ removes any +directionality inherited from the original interaction matrix. From +$\tilde{G}$ we define edge weights $w_{ij}=|\tilde{G}_{ij}|$ +(interaction strength, irrespective of sign) and construct a filtered +Vietoris--Rips complex $\mathcal{R}(\epsilon)$: a simplex +$[v_0,\dots,v_k]$ is included whenever every pairwise weight exceeds +the threshold, +\begin{equation} + [v_0,\dots,v_k]\in\mathcal{R}(\epsilon) + \quad\Longleftrightarrow\quad + w_{v_i v_j}\geq\epsilon + \;\;\forall\;0\leq i postsynaptic i). + +Method 1 — degree assortativity (Newman, directed): Pearson r between +k_i^in at the target and k_j^out at the source over all edges (Eq. assortativity). + +Usage: + ./topoAna_daleMatrix4.py --basePath $basePath --dataName daleN200_fa3733 +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import os +from pprint import pprint +import numpy as np +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +import argparse + + +def adjacency_from_A_off(A_off, atol=0.0): + """ + Binary directed adjacency G: G_ij = 1 if |A_off_ij| > atol and i != j. + Rows index postsynaptic i, columns presynaptic j (matches gen_daleMatrices4 / TeX). + """ + A = np.asarray(A_off, dtype=np.float64) + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError("A_off must be square (N, N)") + G = (np.abs(A) > atol).astype(np.float64) + np.fill_diagonal(G, 0.0) + return G + + +def shuffle_neuron_order(A_off, perm=None): + """ + Reindex neurons with a shared permutation on rows and columns. + Rows are shuffled first; columns are reordered to the same new neuron order. + Returns (A_off_shuffled, perm). + """ + A = np.asarray(A_off, dtype=np.float64) + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError("A_off must be square (N, N)") + + n = A.shape[0] + if perm is None: + perm = np.random.permutation(n) + else: + perm = np.asarray(perm) + if perm.shape != (n,): + raise ValueError("perm must have shape (N,)") + return A[perm][:, perm], perm + + +def directed_degree_assortativity(G): + """ + Directed degree assortativity r per network_ver4_topoDiscovery.tex, Sec. 1. + For each edge j -> i (G_ij = 1): pair (k_i^in, k_j^out). + Uses |E|^{-1} moments (population cov / product of population std devs). + Returns (r, |E|, extra dict). r is nan if |E| < 2 or a std dev is 0. + """ + G = np.asarray(G, dtype=np.float64) + n = G.shape[0] + k_in = G.sum(axis=1) + k_out = G.sum(axis=0) + ei, ej = np.nonzero(G > 0) + m = int(ei.size) + if m < 2: + return float("nan"), m, {"n_nodes": n, "n_edges": m, "reason": "|E| < 2"} + + x = k_in[ei].astype(np.float64) # Target in-degrees + y = k_in[ej].astype(np.float64) # Source in-degrees (In-In correlation) + mx = x.mean() + my = y.mean() + cov = (x * y).mean() - mx * my + vx = ((x - mx) ** 2).mean() + vy = ((y - my) ** 2).mean() + sig_in_target = float(np.sqrt(vx)) + sig_in_source = float(np.sqrt(vy)) + if sig_in_target <= 0.0 or sig_in_source <= 0.0: + return float("nan"), m, { + "n_nodes": n, + "n_edges": m, + "sigma_in_target": sig_in_target, + "sigma_in_source": sig_in_source, + "reason": "zero std on edge multiset", + } + r = cov / (sig_in_target * sig_in_source) + return float(r), m, { + "n_nodes": n, + "n_edges": m, + "r_assortativity": float(r), + "mean_k_in_target": float(mx), + "mean_k_in_source": float(my), + "sigma_in_target": sig_in_target, + "sigma_in_source": sig_in_source, + "cov_edge": float(cov), + } + + +def run_method1(G, dmd, verb=1): + r, m_edges, info = directed_degree_assortativity(G) + d_ker = dmd["placement_ker_delta"] + + print("\n--- Method 1: directed degree assortativity ---") + print(" N=%d |E|=%d r=%s" % (info["n_nodes"], m_edges, repr(r) if np.isnan(r) else "%.6f" % r)) + if verb > 1: + print(" mean target k_in on edges: %.6f" % info["mean_k_in_target"]) + print(" mean source k_in on edges: %.6f" % info["mean_k_in_source"]) + print(" sigma target in: %.6f" % info["sigma_in_target"]) + print(" sigma source in: %.6f" % info["sigma_in_source"]) + print(" (truth metadata) placement_ker_delta = %s" % (d_ker,)) + + return info + + +def run_method2(G, dmd, verb=1): + G = np.asarray(G, dtype=np.float64) + G2 = G @ G + k_in = G.sum(axis=1) # Target nodes (rows) + k_out = G.sum(axis=0) # Source nodes (columns) + + ei, ej = np.nonzero(G > 0) + m = int(ei.size) + if m == 0: + return {"mean_jaccard": float("nan"), "n_edges": 0, "reason": "no edges"} + + num = G2[ei, ej] + den = k_in[ei] + k_out[ej] - num + # den is always >= 2 because i!=j, j in N_in(i), i in N_out(j) + j_idx = num / den + j_mean = float(np.mean(j_idx)) + + print("\n--- Method 2: common-neighbor overlap (Jaccard Index) ---") + print(" mean Jaccard index: %.6f" % j_mean) + + info = { + "mean_jaccard": j_mean, + "n_edges": m + } + return info + + +def run_method3(G, dmd, verb=1): + G = np.asarray(G, dtype=np.float64) + K = 6 + k_vals = np.arange(1, K + 1) # k = 1..K + + # Tk = trace(G^(k+1)) / sum(G^k) + transitivity = [] + Gk = G.copy() + for k in k_vals: + Gnext = Gk @ G + num = float(np.trace(Gnext)) + den = float(Gk.sum()) + if den > 0: + transitivity.append(num / den) + else: + transitivity.append(0.0) + Gk = Gnext + + transitivity = np.array(transitivity) + + # xi = -d log Tk / dk for k in 2..K + # Handle log compatibility + fit_mask = (k_vals >= 2) & (transitivity > 1e-16) + xi = float("nan") + if np.sum(fit_mask) >= 2: + slope, _ = np.polyfit(k_vals[fit_mask], np.log(transitivity[fit_mask]), 1) + xi = -float(slope) + + print("\n--- Method 4: cycle closure decay rate ---") + print(" decay rate xi: %.6f" % xi) + + return { + "cycle_decay": xi, + "transitivity_k": transitivity.tolist() + } + + +def run_method4(A_off, dmd, verb=1): + try: + import gudhi + except ImportError as exc: + raise ImportError( + "Method 4 requires GUDHI. Install it in this environment before running topology analysis." + ) from exc + A = np.abs(A_off) + # Symmetrized interaction strength (LaTeX Method 5, weighted variant) + W = (A + A.T) / 2.0 + N = W.shape[0] + + st = gudhi.SimplexTree() + # Add nodes at filtration 0 + for i in range(N): + st.insert([i], filtration=0.0) + + # Add edges with filtration -W_ij (strongest edges enter first) + ei, ej = np.nonzero(W > 1e-12) + for i, j in zip(ei, ej): + if i < j: + st.insert([i, j], filtration=-float(W[i, j])) + + # Expand to Rips complex of dimension 2 to capture H1 (loops) + st.expansion(2) + + # Compute persistence + st.persistence() + + # Total persistence Pi_1 (dim 1) + diag1 = st.persistence_intervals_in_dimension(1) + pi1 = 0.0 + if len(diag1) > 0: + # Filter out infinite death + valid1 = diag1[np.isfinite(diag1[:, 1])] + pi1 = float(np.sum(valid1[:, 1] - valid1[:, 0])) + + print("\n--- Method 4: persistent homology (H1) ---") + print(" total persistence Pi_1: %.6f" % pi1) + + return { + "homology_k1": pi1, + "n_features_k1": len(diag1) + } + + +def get_parser(): + parser = argparse.ArgumentParser(description="Topology analysis on Dale simTruth.npz (method 1: assortativity)") + parser.add_argument("-v", "--verbosity", type=int, help="increase output verbosity", default=1, dest="verb") + parser.add_argument("--basePath", default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input data") + parser.add_argument("--dataName", default="daleN150_448b86", help="simulated Dale network base name") + parser.add_argument("--shuffle", action="store_true", help="randomize neuron order of A_off before computing G") + + args = parser.parse_args() + + args.inpPath = os.path.join(args.basePath, "truthDale") + args.outPath = os.path.join(args.basePath, "topoAna") + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + assert os.path.exists(args.basePath) + os.makedirs(args.outPath, exist_ok=True) + return args + + +if __name__ == "__main__": + args = get_parser() + np.set_printoptions(precision=3) + + truthFF = os.path.join(args.inpPath, f"{args.dataName}.simTruth.npz") + trueD, trueMD = read_data_npz(truthFF, verb=args.verb > 0) + if args.verb > 1: + print("\nSimulation Truth Metadata:") + pprint(trueMD) + + trueMD["short_name"] = args.dataName + + A_off = np.asarray(trueD["A_off_true"], dtype=np.float64) + print("read obj: A_off_true %s %s" % (A_off.shape, A_off.dtype)) + neuron_perm = np.arange(A_off.shape[0], dtype=np.int64) + if args.shuffle: + A_off, neuron_perm = shuffle_neuron_order(A_off) + print("applied neuron shuffle to A_off, first 10 perm entries:", neuron_perm[:10]) + + G = adjacency_from_A_off(A_off) + dmd = trueMD["dale_conf"] + + info1 = run_method1(G, dmd, verb=args.verb) + info2 = run_method2(G, dmd, verb=args.verb) + info3 = run_method3(G, dmd, verb=args.verb) + info4 = run_method4(A_off, dmd, verb=args.verb) + + print("\nM:done topoAna_daleMatrix4") + + + # ══════════════════════════════════════════════════════════════ + # Save results + # ══════════════════════════════════════════════════════════════ + outMD ={"dale_conf":dmd} + outMD["shuffle"] = bool(args.shuffle) + outMD["provenance"] = {'A-input':args.dataName} + outMD['methods']={ + 'assortativity': info1, + 'jaccard_index': info2, + 'cycle_decay': info3, + 'homology_k1': info4 + } + + outD={'G': G,'A_off': A_off, 'neuron_perm': neuron_perm} + if args.verb>1: pprint(outMD) + outF = args.dataName + outFF = os.path.join(args.outPath, f"{outF}.topoAna.npz") + + write_data_npz(outD, outFF, metaD=outMD) + print(f"\nSaved: {outFF}") + print("\n#Summary: dataName, ker_delta, r_assort, jaccard, cycle_decay, hom_k1") + print("#Values: %s %.2f %.6f %.6f %.6f %.6f" % (args.dataName, dmd["placement_ker_delta"], info1["r_assortativity"], info2["mean_jaccard"], info3["cycle_decay"], info4["homology_k1"])) + diff --git a/causal_net/topoAna_ver4/view_daleMatrix4.py b/causal_net/topoAna_ver4/view_daleMatrix4.py new file mode 120000 index 00000000..8f0435d1 --- /dev/null +++ b/causal_net/topoAna_ver4/view_daleMatrix4.py @@ -0,0 +1 @@ +../nonStationary_ver4/view_daleMatrix4.py \ No newline at end of file diff --git a/causal_net/topoAna_ver4/view_topoCorrelation.py b/causal_net/topoAna_ver4/view_topoCorrelation.py new file mode 100755 index 00000000..bc727036 --- /dev/null +++ b/causal_net/topoAna_ver4/view_topoCorrelation.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +""" +Inspect how saved topology-analysis metrics vary with placement_ker_delta. + +Usage: + ./view_topoCorrelation.py --nameTemplate aN200 -p a +""" + +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +import argparse +import glob +import os + +import numpy as np + +from PlotterTopoCorrel import Plotter +from toolbox.Util_NumpyIO import read_data_npz + + +METRIC_SPECS = ( + ("assortativity", "r_assortativity", "r_assort"), + ("jaccard_index", "mean_jaccard", "jaccard"), + ("cycle_decay", "cycle_decay", "cycle_decay"), + ("homology_k1", "homology_k1", "hom_k1"), +) + + +def get_parser(): + parser = argparse.ArgumentParser(description="Visualize metric-vs-kernel correlations from *.topoAna.npz") + parser.add_argument("-v", "--verbosity", type=int, help="increase output verbosity", default=1, dest="verb") + parser.add_argument( + "-p", + "--showPlots", + default="a", + nargs="+", + help="plot letters: a=metric inspector", + ) + parser.add_argument("-X", "--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument("--basePath", default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="head dir for input data") + parser.add_argument("--inpPath", default=None, help="directory with *.topoAna.npz; defaults to basePath/topoAna") + parser.add_argument("--outPath", default=None, help="directory for plots; defaults to basePath/plots") + parser.add_argument("--nameTemplate", default="aN200", help="load files matching inpPath/[nameTemplate]*.topoAna.npz") + + args = parser.parse_args() + + if args.inpPath is None: + args.inpPath = os.path.join(args.basePath, "topoAna") + if args.outPath is None: + args.outPath = os.path.join(args.basePath, "plots") + + args.showPlots = "".join(args.showPlots) + + print("myArg-program:", parser.prog) + for arg in vars(args): + print("myArg:", arg, getattr(args, arg)) + + assert os.path.exists(args.inpPath), "missing input path: %s" % args.inpPath + os.makedirs(args.outPath, exist_ok=True) + return args + + +def _data_name_from_file(inpFF, inpMD): + provD = inpMD.get("provenance", {}) + if "A-input" in provD: + return provD["A-input"] + base = os.path.basename(inpFF) + return base.replace(".topoAna.npz", "") + + +def harvest_topology_metrics(args): + inpPattern = os.path.join(args.inpPath, "%s*.topoAna.npz" % args.nameTemplate) + inpFL = sorted(glob.glob(inpPattern)) + if not inpFL: + raise FileNotFoundError("no files found for pattern: %s" % inpPattern) + + if args.verb > 0: + print("\nScanning %d files from pattern:\n %s" % (len(inpFL), inpPattern)) + + harvestD = {metric_name: {} for metric_name, _, _ in METRIC_SPECS} + + print("\n#Summary: dataName, ker_delta, r_assort, jaccard, cycle_decay, hom_k1") + j=0 + for inpFF in inpFL: + _, inpMD = read_data_npz(inpFF, verb=j==0) + j += 1 + + dmd = inpMD["dale_conf"] + methodsD = inpMD["methods"] + placement_L = int(round(float(dmd["placement_L"]))) + ker_delta = float(dmd["placement_ker_delta"]) + data_name = _data_name_from_file(inpFF, inpMD) + + metric_values = {} + for method_name, observable_name, _ in METRIC_SPECS: + metric_val = float(methodsD[method_name][observable_name]) + metric_values[method_name] = metric_val + harvestD[method_name].setdefault(placement_L, []).append((ker_delta, metric_val, data_name)) + + print( + "#Values: %s %.2f %.6f %.6f %.6f %.6f" + % ( + data_name, + ker_delta, + metric_values["assortativity"], + metric_values["jaccard_index"], + metric_values["cycle_decay"], + metric_values["homology_k1"], + ) + ) + print("\nHarvested data for %d files." % j) + for metric_name in harvestD: + for placement_L in harvestD[metric_name]: + harvestD[metric_name][placement_L] = sorted(harvestD[metric_name][placement_L], key=lambda rec: (rec[0], rec[2])) + + return harvestD + + +if __name__ == "__main__": + args = get_parser() + np.set_printoptions(precision=3) + + harvestD = harvest_topology_metrics(args) + + args.prjName = args.nameTemplate + "_topoCorr" + plot = Plotter(args) + + if "a" in args.showPlots: + plot.metrics_inspector(harvestD, args.nameTemplate, figId=1) + + plot.display_all() + print("M:done - view_topoCorrelation completed successfully!") diff --git a/causal_net/toys/Readme b/causal_net/toys/Readme new file mode 100644 index 00000000..34da5ccd --- /dev/null +++ b/causal_net/toys/Readme @@ -0,0 +1,47 @@ +Simple programs + += = = = = = = = = + +./spikeMatrix.py + +fileN='spike_dict.pkl' + +# abs path +filepath = "/global/cfs/cdirs/m2043/causal_inference/DIV13/HET_80k_1/"+fileN + +Do some analysis of spike.pkl file and unrolls it to numpy array with spikes vs. time. Output is 15 GB + +Not used for anything. + + + += = = = = = = = = +dummy_mpi_lasso.py + + srun -n 8 ./dummy_mpi_lasso.py --num_feat 15 --num_samp 300 --lag 2 + +Initializes MPI and assigns ranks using MPI.COMM_WORLD. +Tracks MPI startup time from script start to when all ranks connect. +Uses argparse to accept command-line inputs: +--num_feat: Number of features (default: 20) +--num_samp: Number of samples (default: 300) +--lag: Lag value (default: 2) +Rank 0 generates dummy data (X, Y) and broadcasts it to all ranks. +All ranks receive the data using comm.bcast() to ensure synchronization. +Each rank initializes a dummy Lasso model (DummyLasso): +Simulates workload with sleep(np.random.uniform(2, 5)). +Generates random coefficients. +All ranks measure their fitting time and send results to rank 0 using comm.gather(). +Rank 0 collects and reports fit time statistics: +Average, min, max fit time across all ranks. +Total execution time from script start to end. +Rank 0 prints a structured summary including: +MPI startup time. +Data generation details. +Fit time statistics. +Total execution time. +Model coefficient dimensions. + + += = = = = = = == = = + diff --git a/causal_net/toys/dummy_mpi_lasso.py b/causal_net/toys/dummy_mpi_lasso.py new file mode 100755 index 00000000..4109ba7e --- /dev/null +++ b/causal_net/toys/dummy_mpi_lasso.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 + +''' use cases + ./dummy_mpi_lasso.py --num_feat 15 --num_samp 300 --lag 2 + + srun -n 8 ./dummy_mpi_lasso.py --num_feat 15 --num_samp 300 --lag 2 +''' +#...!...!.................... + + +import sys +import os +import argparse +import numpy as np +from mpi4py import MPI +from time import time, sleep + +# Record script start time +script_start_time = time() + +# Dummy Lasso class (Replacing UoI_Lasso) +class DummyLasso: + def __init__(self, n_real_features, fit_VAR, random_state, comm): + self.n_real_features = n_real_features + self.fit_VAR = fit_VAR + self.random_state = random_state + self.comm = comm + self.coef_ = None # Placeholder for model coefficients + + def fit(self, X, Y): + np.random.seed(self.random_state) + # Simulate computational workload + sleep(np.random.uniform(2, 5)) # Simulate varied fit time + self.coef_ = np.random.randn(X.shape[1]) # Dummy coefficients + +# Dummy data generation function +def generate_dummy_data(num_feat, num_samp, lag): + np.random.seed(42) + X = np.random.randn(num_samp, num_feat * lag) + Y = np.random.randn(num_samp, num_feat) + return X, Y + +def main(num_feat, num_samp, lag): + # Initialize MPI + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + num_ranks = comm.Get_size() + + # Synchronize all ranks and measure startup time + comm.barrier() # Ensure all ranks reach this point + startup_time = time() - script_start_time # Time from script start to this point + + if rank == 0: + print("------------------------------------------------------------") + print("MPI Startup Complete | numRanks=%d | Startup Time: %.3f sec" % (num_ranks, startup_time)) + print("------------------------------------------------------------", flush=True) + + print("Generating data... nFeat=%d, nSamp=%d, lag=%d" % (num_feat, num_samp, lag), flush=True) + + # Generate dummy data (Only rank 0) + X, Y = generate_dummy_data(num_feat, num_samp, lag) + else: + X = None + Y = None + + # Broadcast X, Y to all ranks + X = comm.bcast(X, root=0) + Y = comm.bcast(Y, root=0) + + # All ranks: Initialize and fit DummyLasso + lasso_model = DummyLasso(n_real_features=num_feat, fit_VAR=True, random_state=42, comm=comm) + + start_time = time() + lasso_model.fit(X, Y) + fit_time = time() - start_time + + # Rank 0 collects all fit times + fit_times = comm.gather(fit_time, root=0) + + if rank == 0: + avg_time = np.mean(fit_times) + min_time = np.min(fit_times) + max_time = np.max(fit_times) + total_time = time() - script_start_time # Total execution time from script start + + print("------------------------------------------------------------") + print("Fitting complete in %.1f sec | numRanks=%d" % (total_time, num_ranks)) + print("Avg Fit Time: %.3f sec | Min: %.3f sec | Max: %.3f sec" % (avg_time, min_time, max_time)) + print("Total Execution Time: %.3f sec" % total_time) + print("------------------------------------------------------------", flush=True) + + # Dummy output for model coefficients + B_model = lasso_model.coef_ + print("B_model: (%d,)" % (B_model.shape[0],)) + +if __name__ == "__main__": + # Use argparse for command-line inputs + parser = argparse.ArgumentParser(description="MPI-based Dummy Lasso fitting.") + parser.add_argument("--num_feat", type=int, default=20, help="Number of features (default: 20)") + parser.add_argument("--num_samp", type=int, default=300, help="Number of samples (default: 300)") + parser.add_argument("--lag", type=int, default=2, help="Lag value (default: 2)") + + args = parser.parse_args() + + main(args.num_feat, args.num_samp, args.lag) diff --git a/causal_net/toys/lapseDecision.py b/causal_net/toys/lapseDecision.py new file mode 100644 index 00000000..b781b531 --- /dev/null +++ b/causal_net/toys/lapseDecision.py @@ -0,0 +1,63 @@ +import numpy as np +import matplotlib.pyplot as plt +import argparse +import os + +def plot_function(gammaL, gammaR, b, w, ax): + """ + Draws the function P(yt = Right | xt) = γR + (1 - γR - γL) * (1 / (1 + e^-(wxt+b))). + + Args: + gammaL (float): The lower bound parameter (γL). + gammaR (float): The upper bound parameter (γR). + b (float): The bias parameter in the logistic term. + w (float): The weight/slope parameter in the logistic term. + ax (matplotlib.axes.Axes): The axes (canvas) to draw on. + """ + # Define the fixed x-range from -10 to 10 with a smooth curve + x = np.linspace(-10, 10, 400) + + # Calculate the logistic term (sigmoid function) + logistic_term = 1 / (1 + np.exp(-(w * x + b))) + + # Calculate the full function value + P_right = gammaR + (1 - gammaR - gammaL) * logistic_term + + # Plot the function on the provided axes + ax.plot(x, P_right, label=f'low: γR={gammaR}, high: γL={gammaL}, w={w}, b={b}') + ax.set_xlabel('$x_t$') + ax.set_ylabel('$P(y_t = Right | x_t)$') + ax.set_title('Classic lapse model for sensory decision-making') + ax.legend() + ax.grid(True) + ax.set_xlim([-10, 10]) # Ensure the fixed x-range is applied + ax.set_ylim() # Probabilities are typically between 0 and 1 + +def main(): + """ + Parses command line arguments and generates the plot. + """ + parser = argparse.ArgumentParser(description="Plot a modified logistic function with specified parameters.") + + # Add arguments for gammaL, gammaR, b, and w + parser.add_argument('--gammaL', type=float, default=0.2, help="Value for the lower bound parameter gammaL.") + parser.add_argument('--gammaR', type=float, default=0.7, help="Value for the upper bound parameter gammaR.") + parser.add_argument('--b', type=float, default=0.0, help="Value for the bias parameter b (default: 0.0).") + parser.add_argument('--w', type=float, default=1.0, help="Value for the weight/slope parameter w (default: 1.0).") + + args = parser.parse_args() + + # Create figure and axes + fig, ax = plt.subplots(figsize=(8, 5)) + + # Call the plotting function with user-provided arguments + plot_function(args.gammaL, args.gammaR, args.b, args.w, ax) + + # Define the filename and save the figure + filename = 'out/function_plot.png' + fig.savefig(filename) + print(f"Plot saved successfully as '{os.path.abspath(filename)}'") + +if __name__ == "__main__": + main() + diff --git a/causal_net/toys/mon.sh b/causal_net/toys/mon.sh new file mode 100755 index 00000000..39169413 --- /dev/null +++ b/causal_net/toys/mon.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash +while true; do + free -g + w + sleep 2 +done diff --git a/causal_net/toys/spikeMatrix.py b/causal_net/toys/spikeMatrix.py new file mode 100755 index 00000000..28a9d7e9 --- /dev/null +++ b/causal_net/toys/spikeMatrix.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 + +import pickle +import numpy as np +np.set_printoptions(precision=3) + +fileN='spike_dict.pkl' + +# abs path +filepath = "/global/cfs/cdirs/m2043/causal_inference/DIV13/HET_80k_1/"+fileN + +# Load the dictionary from the .pkl file +with open(filepath, "rb") as f: + spike_dict = pickle.load(f) + + + +max_value = max([max(v) for v in spike_dict.values()]) + +num_neurons = len(spike_dict) + +print('M: max_value=',max_value) +print('M: num_neurons',num_neurons) + +sampling_freq = 10000 # Hz +target_fs = 1000 +duration = max_value/sampling_freq +# Calculate the number of neurons and bins +num_neurons = len(spike_dict) +num_bins = int(duration * target_fs) + +# Create an empty binary matrix for the downsampled data +spike_matrix = np.zeros((num_neurons, num_bins), dtype=int) + +print('M: spike_matrix:',spike_matrix.shape) + +# Define the bin edges based on the target sampling frequency +bin_edges = np.linspace(0, duration, num_bins + 1) # num_bins + 1 to include the last edge + +print('M: bin_edges',bin_edges.shape) + +print('M: spike agregation') +print(' neuron_idx, id, num_samp spikes_in_seconds[:10]') +# Iterate over each neuron and populate the matrix +for neuron_idx, (neuron_id, spikes) in enumerate(spike_dict.items()): + # Convert spike times to seconds + spikes_in_seconds = np.array(spikes )/ sampling_freq + + if neuron_idx%100==0 or neuron_idx==num_neurons-1: + rec=spikes_in_seconds + print(' idx:',neuron_idx, ' id:',neuron_id, rec.shape, rec[:10]) + + + # Digitize the spike times into bins + bin_indices = np.digitize(spikes_in_seconds, bin_edges) - 1 # Convert to 0-based index + bin_indices = bin_indices[(bin_indices >= 0) & (bin_indices < num_bins)] # Remove out-of-range indices + + # Mark the bins where spikes occurred + spike_matrix[neuron_idx, np.unique(bin_indices)] = 1 # Use np.unique to avoid duplicate marking + + + +print("Binary Spike Matrix:") +print(spike_matrix) + +file_path = "out/spike.npy" + +# Save the matrix +np.save(file_path, spike_matrix) + +print(f"Spike matrix saved to {file_path}") diff --git a/causal_net/toys/uoi_mpi_admm_v5.py b/causal_net/toys/uoi_mpi_admm_v5.py new file mode 100755 index 00000000..927ac41d --- /dev/null +++ b/causal_net/toys/uoi_mpi_admm_v5.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 + +''' +IMG=nersc/causal-net:v4 # May 13 + export OMP_NUM_THREADS=2 + salloc -q interactive -C cpu --image=$IMG -t 4:00:00 -A m2043 -N 1 + +uoi_lasso = UoI_Lasso(fit_VAR = True, fit_intercept=False, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, n_lambdas = n_lambdas, random_state=None, comm = boot_comm, global_comm = comm, n_admm = n_admm, admm_rho = rho, max_iter = 50, solver='admm', estimation_solver = "ls") + +- max_iter = 50; enough convergence and for faster runtime +- estimation_solver = "ls": using standard Linear Regression library is faster for parameter estimation than ADMM +- REMOVED n_real_features variable, since it's obsolete now + + +NOTE: distribution=block:block forces the ADMM processes for each bootstrap to localize to a single compute node for efficient communication + +Minimal example testing code consistency: 38 sec + srun -n 1 --distribution=block:block shifter python uoi_mpi_admm_v5.py --num_feat 10 --num_samp 50 --num_admm 1 +>>> Fitting complete in 52.2 sec | numRanks=1 + +Doing meaningfull fit : 42 sec on N=1 +srun -n 128 --distribution=block:block shifter python uoi_mpi_admm_v5.py --num_feat 20 --num_samp 5000 --num_admm 16 + +''' +#...!...!.................... + + +import sys +import os +import argparse +import numpy as np +from mpi4py import MPI +from time import time, sleep + +# Record script start time +script_start_time = time() +omp_threads = os.environ.get("OMP_NUM_THREADS", "Not Set") +assert omp_threads=='2' + + +sys.path.append("/global/homes/b/balewski/prjs/2025_UoI-VAR/") +from examples.var_utils import * +from src.pyuoi.linear_model import * +sys.path.append("/global/homes/b/balewski/prjs/2025_UoI-VAR/src/pyuoi/linear_model") +from sparse_comm_util import build_bootstrap_comm + + +#...!...!.................... +def print_dale_matrix(A,nfeat=None): + if nfeat==None: nfeat=A.shape[0] + # Function to format values + def format_value(val): + if abs(val) < 0.01: + return " . " # Represent zero as '-' + return f"{val:+5.2f}" # Format as +0.12 or -0.23 + + col_indices = "feat " + " ".join(f"{i:2d}" for i in range(nfeat)) + print(col_indices) + # Print row index and formatted values + for i in range(nfeat): + row=A[i] + formatted_row = " ".join(format_value(row[j]) for j in range(nfeat) ) + print(f"{i:2d} {formatted_row}") # Row index + formatted values + + +#...!...!.................... +def generate_dummy_data(num_feat, num_samp, lag): + # Select spectral radius based on lag + rad_dict = {1: 0.98, 2: 0.70, 3: 0.30, 4: 0.20, 5: 0.20} + rad = rad_dict.get(lag, 0.20) + + # Generate data (Only rank 0) + data, transition_matrices, cov = generate_sparse_stationary_var_process( + num_feat, num_samp, lag=lag, sparsity=0.5, spectral_radius=rad, + process_type='gaussian', random_state=42 + ) + + return data + + +#...!...!.................... +def main(num_feat, num_samp, lag, inpName, n_admm): + # Initialize MPI + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + num_ranks = comm.Get_size() + + # Synchronize all ranks and measure startup time + comm.barrier() # Ensure all ranks reach this point + startup_time = time() - script_start_time # Time from script start to this point + + if rank == 0: + T1=time() + print("------------------------------------------------------------") + print("MPI Startup Complete | numRanks=%d | Startup Time: %.3f sec" % (num_ranks, startup_time)) + print("------------------------------------------------------------", flush=True) + + if inpName==None: + # Generate dummy data (Only rank 0) + print("Generating data... nFeat=%d, nSamp=%d, lag=%d" % (num_feat, num_samp, lag), flush=True) + mydata = generate_dummy_data(num_feat, num_samp, lag) + print(' runk0 generated data, elaT=%.1f min'%( (time()-T1)/60.), flush=True) + else: + print(' Load array back from file:',inpName, flush=True) + mydata = np.load(inpName) + num_samp,num_feat=mydata.shape + else: + mydata = None + + if rank == 0: + print('n_admm , num_ranks:', n_admm , num_ranks) + print('mydata:',mydata.shape,type(mydata),'start fit ...', flush=True) + + start_time = time() + + if use_admm: # very scalable + assert n_admm <= num_ranks + assert num_ranks % n_admm ==0 + boot_comm = build_bootstrap_comm(comm, n_admm) + + n_boots_sel = 12 + n_boots_est = 12 + selection_frac = 0.9 + n_lambdas = 48 + max_iter = 1000 + seed = 42 + rho_scaler = 1.5 # was 2.0 + imbalance_tolerance = 10. #was 0.1 + eps=1e-7 + uoi_lasso = UoI_Lasso(fit_VAR = True, fit_intercept=False, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, selection_frac = selection_frac, n_lambdas = n_lambdas, max_iter = max_iter, eps = eps, random_state=seed, comm = boot_comm, global_comm = comm, n_admm = n_admm, rho_scaler = rho_scaler, imbalance_tolerance = imbalance_tolerance, solver='admm', estimation_solver = "ls") + + + if boot_comm is not None: #if the global_rank is part of the boostrap distribution(not admm distribution) + if boot_comm.rank == 0: + uoi_lasso.fit(lag, data = mydata) + else: + uoi_lasso.fit(lag) + else: + if uoi_lasso.solver == "admm": + uoi_lasso.admm_queue() + + else: #original implemetation, not scalable + uoi_lasso = UoI_Lasso(n_real_features = num_feat, fit_VAR = True, fit_intercept=False, random_state=42, comm = comm) + + if comm.rank == 0: + uoi_lasso.fit(lag, data = mydata) + else: + uoi_lasso.fit(lag) + + fit_time = time() - start_time + + # Rank 0 collects all fit times + fit_times = comm.gather(fit_time, root=0) + + if rank != 0: return + + avg_time = np.mean(fit_times) + min_time = np.min(fit_times) + max_time = np.max(fit_times) + total_time = time() - script_start_time # Total execution time from script start + + print("------------------------------------------------------------") + print("Fitting complete in %.1f sec | numRanks=%d" % (total_time, num_ranks)) + print("Avg Fit Time: %.1f sec | Min: %.1f sec | Max: %.1f sec" % (avg_time, min_time, max_time)) + print("Total Execution Time: %.3f sec" % total_time) + print("------------------------------------------------------------", flush=True) + + # Extract model coefficients + # B_model = uoi_lasso.coef_ + # A_model = [B_model.reshape(num_feat, num_feat * lag).T[i * num_feat:(i + 1) * num_feat].T for i in range(lag)] + # A_model = np.array(A_model) + + A_model = uoi_lasso.VAR_coef_ + + print("A_model: (%d, %d, %d)" % (A_model.shape[0], A_model.shape[1], A_model.shape[2])) + # print("B_model: (%d,) | A_model: (%d, %d, %d)" % (B_model.shape[0], A_model.shape[0], A_model.shape[1], A_model.shape[2])) + A=A_model[0] + nfeat=min(20,A.shape[0]) + print_dale_matrix(A,nfeat) + +#================================= +# M A I N +#================================= + +if __name__ == "__main__": + # Use argparse for command-line inputs + parser = argparse.ArgumentParser(description="MPI-based Dummy Lasso fitting.") + parser.add_argument("--num_feat", type=int, default=20, help="Number of features ") + parser.add_argument("--num_samp", type=int, default=50_000, help="Number of samples") + parser.add_argument("--num_admm", type=int, default=32, help="num of processes per node to solve the bootstrap variable selection problem in a distributed fashion") + parser.add_argument("--lag", type=int, default=1, help="Lag value ") + parser.add_argument("--inpName", default=None,help='input name, will define num features and num samples') + + args = parser.parse_args() + use_admm = True # very scalable + #n_admm = 64 # n_process should be multiple of n_admm, and at MOST n_admm*n_boot*n_reg_param + rho = None + #rho = 1e10 + + main(args.num_feat, args.num_samp, args.lag, args.inpName, args.num_admm) diff --git a/causal_net/tutorial/actionPotential_model1.py b/causal_net/tutorial/actionPotential_model1.py new file mode 100755 index 00000000..5fb16317 --- /dev/null +++ b/causal_net/tutorial/actionPotential_model1.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +''' +Realistic Action Potential Simulator with argparse parameters and PNG saving + +Usage: +./action_potential_simulator.py --resting_pot -70 --peak_pot 40 --threshold_pot -55 --save_path neuron_trace.png +''' +import numpy as np +import matplotlib.pyplot as plt +import argparse + +def generate_action_potential( + time_ms, + resting_pot=-70, + peak_pot=40, + threshold_pot=-55, + spike_start=5, + depolar_dur=0.5, + peak_dur=0.5, + repolar_dur=1.0, + hyperpolar_dur=3.0, + hyperpolar_depth=10 +): + voltage_trace = np.full(time_ms.shape, resting_pot) + t_rise_end = spike_start + depolar_dur + t_peak_end = t_rise_end + peak_dur + t_repol_end = t_peak_end + repolar_dur + t_hyperpol_end = t_repol_end + hyperpolar_dur + + for i, t in enumerate(time_ms): + if spike_start <= t < t_rise_end: + voltage_trace[i] = resting_pot + (peak_pot - resting_pot) * (1 - np.exp(-(t - spike_start)/0.1)) + elif t_rise_end <= t < t_peak_end: + voltage_trace[i] = peak_pot + elif t_peak_end <= t < t_repol_end: + voltage_trace[i] = peak_pot - (peak_pot - resting_pot) * (1 - np.exp(-(t - t_peak_end)/0.3)) + elif t_repol_end <= t < t_hyperpol_end: + voltage_trace[i] = resting_pot - hyperpolar_depth * np.exp(-(t - t_repol_end)/1.0) + elif t >= t_hyperpol_end: + voltage_trace[i] = resting_pot + return voltage_trace + +def YYgenerate_action_potential(time_ms, resting_pot, peak_pot, threshold_pot, + spike_start, depolar_dur, peak_dur, repolar_dur, + hyperpolar_dur, hyperpolar_depth): + voltage_trace = np.full(time_ms.shape, resting_pot) + t_rise_end = spike_start + depolar_dur + t_peak_end = t_rise_end + peak_dur + t_repol_end = t_peak_end + repolar_dur + t_hyperpol_end = t_repol_end + hyperpolar_dur + + for i, t in enumerate(time_ms): + if spike_start <= t < t_rise_end: + voltage_trace[i] = resting_pot + (peak_pot - resting_pot) * (1 - np.exp(-(t - spike_start)/0.1)) + elif t_rise_end <= t < t_peak_end: + voltage_trace[i] = peak_pot + elif t_peak_end <= t < t_repol_end: + voltage_trace[i] = peak_pot - (peak_pot - resting_pot) * (1 - np.exp(-(t - t_peak_end)/0.3)) + elif t_repol_end <= t < t_hyperpol_end: + voltage_trace[i] = resting_pot - hyperpolar_depth * np.exp(-(t - t_repol_end)/1.0) + elif t >= t_hyperpol_end: + voltage_trace[i] = resting_pot + return voltage_trace + +def plot_action_potential(time_ms, voltage_trace, threshold_pot, save_path): + plt.figure(figsize=(12, 4)) + plt.plot(time_ms, voltage_trace, label='Action Potential') + plt.axhline(y=threshold_pot, color='gray', linestyle='--', label=f'Threshold Potential ({threshold_pot} mV)') + plt.title("Realistic Action Potential of Excitatory Neuron") + plt.xlabel("Time (ms)") + plt.ylabel("Membrane Potential (mV)") + plt.legend() + plt.grid(True) + plt.savefig(save_path, format='png') + plt.close() + +def main(): + parser = argparse.ArgumentParser(description="Generate realistic action potential trace.") + parser.add_argument('--resting_pot', type=float, default=-70, help='Resting potential (mV)') + parser.add_argument('--peak_pot', type=float, default=40, help='Peak potential (mV)') + parser.add_argument('--threshold_pot', type=float, default=-55, help='Threshold potential (mV)') + parser.add_argument('--spike_start', type=float, default=5, help='Spike start time (ms)') + parser.add_argument('--depolar_dur', type=float, default=0.5, help='Depolarization duration (ms)') + parser.add_argument('--peak_dur', type=float, default=0.5, help='Peak duration (ms)') + parser.add_argument('--repolar_dur', type=float, default=1.0, help='Repolarization duration (ms)') + parser.add_argument('--hyperpolar_dur', type=float, default=3.0, help='Hyperpolarization duration (ms)') + parser.add_argument('--hyperpolar_depth', type=float, default=10, help='Hyperpolarization depth below resting potential (mV)') + parser.add_argument('--total_time', type=float, default=20.0, help='Total simulation time (ms)') + parser.add_argument('--time_resolution', type=float, default=0.1, help='Time resolution (ms)') + parser.add_argument('--save_path', type=str, default='out/action_potential.png', help='Path to save PNG plot') + + args = parser.parse_args() + + time_ms = np.arange(0, args.total_time + args.time_resolution, args.time_resolution) + voltage_trace = generate_action_potential( + time_ms, + args.resting_pot, + args.peak_pot, + args.threshold_pot, + args.spike_start, + args.depolar_dur, + args.peak_dur, + args.repolar_dur, + args.hyperpolar_dur, + args.hyperpolar_depth + ) + #voltage_trace = generate_action_potential(time_ms) + + plot_action_potential(time_ms, voltage_trace, args.threshold_pot, args.save_path) + +if __name__ == "__main__": + main() diff --git a/causal_net/uoiPoissonLag1/PlotterFitEval2.py b/causal_net/uoiPoissonLag1/PlotterFitEval2.py new file mode 100644 index 00000000..c77beefd --- /dev/null +++ b/causal_net/uoiPoissonLag1/PlotterFitEval2.py @@ -0,0 +1,397 @@ +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +from toolbox.PlotterBackbone import PlotterBackbone +from matplotlib import cm as cmap +import matplotlib.ticker as ticker +from pprint import pprint +import numpy as np +import matplotlib.gridspec as gridspec +import matplotlib.colors as colors +from matplotlib.colors import TwoSlopeNorm +from scipy.stats import gennorm +from matplotlib.colors import LogNorm +from Util_poissonFdr import qa_Bfit, qa_Afit +import sys, os +sys.path.append(os.path.join(os.path.dirname(__file__), '../ver6')) +from Util_pseudospectra import compute_pseudospectrum + +#............................ +#............................ +#............................ +class Plotter(PlotterBackbone): + def __init__(self, args): + PlotterBackbone.__init__(self,args) + +#...!...!.................. + def summary_fitUoI(self,fitD, trueD,md,byFreq=False, figId=1): # p=a + figId=self.smart_append(figId) + nrow,ncol=1,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(14,3)) + + fmd=md['fit_uoi'] + cmd=md['conf_uoi'] + Nn=fmd['num_neurons'] + fitType=md['fit_type'] + isExp = md.get('data_type') == 'bioExp' # Automatically detect experimental data + + # Unpack arrays from bigD + A_fit = fitD['A_'+fitType] + B = fitD['B_'+fitType] + Freq=fitD['single_rates'] + A_true = trueD['A_true'] + + #...... Training curves + ax = self.plt.subplot(nrow,ncol,1) + iterSkip=3 # is x10 + l1_loss=fitD['l1_loss_sel'][:,iterSkip:] + ax.plot(l1_loss[0] ,l1_loss[1], marker='o') + ax.grid(True) + # Title and run summary + samples_k = int(fmd.get('num_samples_used', 0))//1000 + tit='UoI %s samples=%dk' % (md.get('short_name',''), samples_k) + ax.set(title=tit, + xlabel='UoI Iteration', + ylabel='UoI L1 Loss') + slurm_nodes = fmd.get('slurm_nodes', 0) + slurm_ranks = fmd.get('slurm_ranks', 0) + fit_min = fmd.get('training_time_sec', 0)/60.0 + n_boots_sel = cmd.get('n_boots_sel', fmd.get('n_boots', '?')) + n_boots_est = cmd.get('n_boots_est', fmd.get('n_boots', '?')) + info_txt = 'Slurm N=%s ranks=%s\n fit %.1f min\n nBoot %s, %s' % (str(slurm_nodes), str(slurm_ranks), fit_min, str(n_boots_sel), str(n_boots_est)) + ax.text(0.5, 0.95, info_txt, transform=ax.transAxes, va='top', ha='left') + + #...... off diagonal distribution + ax = self.plt.subplot(nrow,ncol,2) + # 1D histogram of non-zero off-diagonal A_fit values + N = A_fit.shape[0] + offdiag_mask = ~np.eye(N, dtype=bool) + A_off = A_fit[offdiag_mask] + nz = A_off[np.abs(A_off) > 0] + ax.hist(nz, bins=100, color='green', alpha=0.8, edgecolor=None) + tit='A_fit off-diagonal (N=%d)' % nz.size + ax.set(title=tit, + xlabel='A_fit value (off-diagonal, non-zero)', + ylabel='count') + ax.grid(True) + + #...... diagonal(A_fit) distribution + ax = self.plt.subplot(nrow,ncol,3) + diag_vals = np.diag(A_fit) + ax.hist(diag_vals, bins=100, color='brown', alpha=0.8, edgecolor=None) + tit='A_fit diagonal (N=%d)' % diag_vals.size + ax.set(title=tit, + xlabel='A_fit diag value', + ylabel='count') + ax.grid(True) + + #...... B (bias) distribution + ax = self.plt.subplot(nrow,ncol,4) + ax.hist(B, bins=100, color='salmon', alpha=0.8, edgecolor=None) + tit='B_fit values (N=%d)' % B.size + ax.set(title=tit, + xlabel='B_fit value', + ylabel='count') + ax.grid(True) + + # Figure title with short_name, samples, and fdr rate + fdr_rate = fmd.get('fdr_rate', md.get('conf_uoi', {}).get('fdr_rate', None)) + if fdr_rate is None: + fig_tit = '%s samples=%dk' % (md.get('short_name',''), int(fmd.get('num_samples_used', 0))//1000) + else: + fig_tit = '%s samples=%dk fdr=%.3f' % (md.get('short_name',''), int(fmd.get('num_samples_used', 0))//1000, float(fdr_rate)) + self.plt.suptitle(fig_tit) + self.plt.tight_layout(rect=[0,0,1,0.92]) + +#...!...!.................. + def correlations_fitUoI(self, fitD, trueD, md, figId=1): # p=b + figId=self.smart_append(figId) + nrow,ncol=3,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(14,9)) + + # Metadata + fmd = md['fit_uoi'] + data_short = md.get('short_name','') + samples_k = int(fmd.get('num_samples_used', 0))//1000 + fdr_rate = fmd.get('fdr_rate', md.get('conf_uoi', {}).get('fdr_rate', None)) + fig_tit = 'UoI-Poisson-FDR %s samples=%dk fdr=%.3f' % (data_short, samples_k, float(fdr_rate)) + # Extract arrays + fitType = md['fit_type'] + A_fit = fitD['A_'+fitType] + B_fit = fitD['B_'+fitType] + A_true = trueD['A_true'] + B_true = trueD['B_true'] + + # Build QA stats, borrowing style from sparePlot + qaD = qa_Afit(A_true, A_fit) + qaD['bterm'] = qa_Bfit(B_true, B_fit) + + colors = {'exc': 'red', 'inh': 'blue', 'diag': 'brown', 'bterm': 'salmon'} + names = ['inh', 'exc', 'diag', 'bterm'] + + for i, name in enumerate(names): + ax = self.plt.subplot(nrow, ncol, i+1) + stats = qaD[name] + color = colors[name] + tval = stats['tval'] + fval = stats['fval'] + + if tval.size > 0: + ax.scatter(tval, fval, alpha=0.6, s=10, c=color) + # y=x line + lims = [ + np.min([ax.get_xlim(), ax.get_ylim()]), + np.max([ax.get_xlim(), ax.get_ylim()]), + ] + ax.plot(lims, lims, 'k--', alpha=0.75, zorder=0,lw=0.8) + ax.set_aspect('equal', 'box') + ax.set_xlim(lims) + ax.set_ylim(lims) + # center of gravity cross + ax.plot(np.mean(tval), np.mean(fval), '+', c='black', markersize=18, markeredgewidth=3) + + ax.grid(True, alpha=0.3) + sub_tit=name.capitalize() + ax.set(title=sub_tit, xlabel='True Value', ylabel='Fitted Value') + + # mean/std box + ax.text(0.95, 0.07, 'resid N=%d\nmean=%.3f\nstd=%.3f'%(tval.shape[0], stats['res_mean'], stats['res_std']), transform=ax.transAxes, fontsize=9, va='bottom', ha='right') + if name in ['exc','inh']: + ax.text(0.05, 0.95, 'TP=%d\nFP=%d\nFN=%d'%(stats['TP'], stats['FP'], stats['FN']), + transform=ax.transAxes, fontsize=9, va='top', ha='left') + + # Second row: residuals vs single_rates with adaptive x-scale, lighter colors + rates = fitD.get('single_rates', None) + use_log_x = False + assert rates is not None #and np.any(rates > 0): + rmin = 0 # float(np.min(rates[rates > 0])) + rmax = float(np.max(rates)) + ratio = rmin / rmax if rmax > 0 else 1.0 + # If small dynamic range (min/max > 0.02), use linear; otherwise log + use_log_x = not (ratio > 0.02) + + for i, name in enumerate(names): + ax = self.plt.subplot(nrow, ncol, ncol + i + 1) + stats = qaD[name] + color = colors[name] + tval = stats['tval'] + fval = stats['fval'] + resid = fval - tval + if tval.size > 0 and rates is not None: + if name in ['exc','inh']: + diag_m = np.eye(A_true.shape[0], dtype=bool) + if name == 'exc': + true_m = (A_true > 0) & (~diag_m) + pred_m = (A_fit > 0) & (~diag_m) + else: + true_m = (A_true < 0) & (~diag_m) + pred_m = (A_fit < 0) & (~diag_m) + tp_m = true_m & pred_m + _, cols = np.where(tp_m) + xvals = rates[cols] + elif name == 'diag': + xvals = rates # one per neuron + else: # bterm + xvals = rates + epsx = 1e-6 + m = (xvals > epsx) if use_log_x else np.ones_like(xvals, dtype=bool) + ax.scatter(xvals[m], resid[:m.sum()], alpha=0.4, s=10, c=color) + ax.set_xscale('log' if use_log_x else 'linear') + ax.axhline(0.0, linestyle='--', color='gray', linewidth=0.8) + ax.grid(True, alpha=0.3) + sub_tit='%s residuals vs rate'%(name.capitalize()) + ax.set(title=sub_tit, + xlabel='single_rates (Hz)', + ylabel='fit - true') + + # Third row: histograms of residuals with stats and x=0 line + for i, name in enumerate(names): + #print('AAA',i,name) + ax = self.plt.subplot(nrow, ncol, 2*ncol + i + 1) + stats = qaD[name] + tval = stats['tval'] + fval = stats['fval'] + if tval.size > 0: + resid = fval - tval + ax.hist(resid, bins=60, color=colors[name], alpha=0.6, edgecolor=None) + ax.axvline(0.0, linestyle='--', color='black', linewidth=0.8) + # stats box (reuse top-row mean/std) + ax.text(0.95, 0.90, 'N=%d\nmean=%.3f\nstd=%.3f'%(tval.shape[0], stats['res_mean'], stats['res_std']), transform=ax.transAxes, fontsize=9, va='top', ha='right') + ax.grid(True, alpha=0.3) + sub_tit='%s residuals'%(name.capitalize()) + ax.set(title=sub_tit, xlabel='fit - true', ylabel='count') + + if 'resRange' in md: ax.set_xlim(md['resRange'][name]) + if name=='inh': ax.text(0.05,0.9,data_short, transform=ax.transAxes, fontsize=10,color='m') + + # Title for the figure + self.plt.suptitle(fig_tit) + self.plt.tight_layout(rect=[0,0,1,0.95]) + +#...!...!.................. + def correlation_for_kris(self, fitD, trueD, md, figId=1): #p=c + figId=self.smart_append(figId) + nrow,ncol=2,4 + fig=self.plt.figure(figId,facecolor='white', figsize=(14,6)) + + # Metadata and title + fmd = md['fit_uoi'] + data_short = md.get('short_name','') + samples_k = int(fmd.get('num_samples_used', 0))//1000 + fdr_rate = fmd.get('fdr_rate', md.get('conf_uoi', {}).get('fdr_rate', None)) + fig_tit = 'UoI-Poisson fit %s samples=%dk fdr=%.3f' % (data_short, samples_k, float(fdr_rate)) + + # Extract arrays + fitType = md['fit_type'] + A_fit = fitD['A_'+fitType] + B_fit = fitD['B_'+fitType] + A_true = trueD['A_true'] + B_true = trueD['B_true'] + + # QA stats + qaD = qa_Afit(A_true, A_fit) + qaD['bterm'] = qa_Bfit(B_true, B_fit) + + colors = {'exc': 'red', 'inh': 'blue', 'diag': 'brown', 'bterm': 'salmon'} + names = ['inh', 'exc', 'diag', 'bterm'] + + # Row 1: true vs fitted + for i, name in enumerate(names): + ax = self.plt.subplot(nrow, ncol, i+1) + stats = qaD[name] + color = colors[name] + tval = stats['tval']; fval = stats['fval'] + if tval.size > 0: + ax.scatter(tval, fval, alpha=0.6, s=10, c=color) + lims = [np.min([ax.get_xlim(), ax.get_ylim()]), np.max([ax.get_xlim(), ax.get_ylim()])] + ax.plot(lims, lims, 'k--', alpha=0.75, zorder=0,lw=0.8) + ax.set_aspect('equal', 'box'); ax.set_xlim(lims); ax.set_ylim(lims) + ax.plot(np.mean(tval), np.mean(fval), '+', c='black', markersize=18, markeredgewidth=3) + ax.grid(True, alpha=0.3) + sub_tit=name.capitalize() + ax.set(title=sub_tit, xlabel='True Value', ylabel='Fitted Value') + # lightweight stats text (no frame) + ax.text(0.95, 0.07, 'resid N=%d\nmean=%.3f\nstd=%.3f'%(tval.shape[0], stats['res_mean'], stats['res_std']), + transform=ax.transAxes, fontsize=9, va='bottom', ha='right') + if name in ['exc','inh']: + ax.text(0.05, 0.95, 'TP=%d\nFP=%d\nFN=%d'%(stats['TP'], stats['FP'], stats['FN']), + transform=ax.transAxes, fontsize=9, va='top', ha='left') + + # Row 2: residuals vs true B_idle + for i, name in enumerate(names): + ax = self.plt.subplot(nrow, ncol, ncol + i + 1) + stats = qaD[name] + color = colors[name] + tval = stats['tval']; fval = stats['fval'] + resid = fval - tval + if tval.size > 0: + if name in ['exc','inh']: + diag_m = np.eye(A_true.shape[0], dtype=bool) + if name == 'exc': + true_m = (A_true > 0) & (~diag_m) + pred_m = (A_fit > 0) & (~diag_m) + else: + true_m = (A_true < 0) & (~diag_m) + pred_m = (A_fit < 0) & (~diag_m) + tp_m = true_m & pred_m + _, cols = np.where(tp_m) + xvals = B_true[cols] + elif name == 'diag': + xvals = B_true + else: + xvals = B_true + ax.scatter(xvals, resid[:xvals.shape[0]], alpha=0.5, s=12, c=color) + ax.axhline(0.0, linestyle='--', color='gray', linewidth=0.8) + ax.grid(True, alpha=0.3) + sub_tit='%s residuals vs true B'%(name.capitalize()) + ax.set(title=sub_tit, xlabel='true B_idle', ylabel='fit - true') + + self.plt.suptitle(fig_tit) + self.plt.tight_layout(rect=[0,0,1,0.94]) + +#...!...!.................. + def eigenvalues_fitUoI(self, fitD, trueD, md, figId=1): # p=d + + figId=self.smart_append(figId) + nrow,ncol=1,2 + fig=self.plt.figure(figId,facecolor='white', figsize=(8,4)) + + fitType=md['fit_type'] + assert not md.get('data_type') == 'bioExp' # no truts for experimental data + + # Unpack arrays from bigD + A_fit = fitD['A_'+fitType] + A_true = trueD['A_true'] + + eigT = np.linalg.eigvals(A_true) + eigF = np.linalg.eigvals(A_fit) + + reT = np.real(eigT) + imT = np.imag(eigT) + + ax = self.plt.subplot(nrow,ncol,1) + ax.scatter(reT, imT, color='blue', marker='o', label='True', s=10) + + reF = np.real(eigF) + imF = np.imag(eigF) + ax.scatter(reF, imF, color='red', marker='o', facecolors='none', label='fit', s=20) + + ax.set_ylim(-0.1,) + #ax.set_xlim(right=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.grid(True, alpha=0.3) + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + ax.legend() + #ax.set_title(f'Eigenvalues for {outName} %d samples'%(args.samples)) + + def pseudospectra_fitUoI(self, fitD, trueD, md, figId=5): + figId = self.smart_append(figId) + fig = self.plt.figure(figId, facecolor='white', figsize=(8, 6)) + + fitType = md['fit_type'] + A_fit = fitD['A_' + fitType] + + npts = 80 + minY = -0.2 + epsMin = 0.03 + + X, Y, sigma_grid, eigs = compute_pseudospectrum(A_fit, npts, minY) + + ax = self.plt.subplot(1, 1, 1) + title = 'Pseudospectra, fit M%d, %s' % (A_fit.shape[0], md['short_name']) + + levels = np.logspace(-2.5, -0.5, 10) + contour = ax.contour(X, Y, sigma_grid, levels=levels, cmap='viridis', linewidths=0.8) + ax.clabel(contour, inline=True, fontsize=8, fmt='ε=%.3f') + ax.text(0.05, 0.85, 'green: ε<%.3f' % epsMin, transform=ax.transAxes) + mask = sigma_grid > epsMin + sigma_grid_masked = np.ma.array(sigma_grid, mask=mask) + + contour_fill = ax.contourf(X, Y, sigma_grid_masked, levels=np.linspace(0, epsMin, 10), + colors=['lightgreen'], alpha=0.5) + + ax.scatter(np.real(eigs), np.imag(eigs), color='red', s=15, zorder=3, label='Eigenvalues') + + ax.axvline(0, color='black', linestyle='--', lw=1.5) + ax.axhline(0, color='black', linestyle='--', lw=1.5) + + ax.set_title(title, fontsize=14) + ax.set_xlabel('Real Part') + ax.set_ylabel('Imaginary Part') + + ax.grid(True, linestyle=':', alpha=0.6) + ax.legend() + ax.set_ylim(bottom=minY) + ax.set_xlim(-3, ) + +#............................ +#............................ +#............................ + + + + + + diff --git a/causal_net/uoiPoissonLag1/README.Jan b/causal_net/uoiPoissonLag1/README.Jan new file mode 100644 index 00000000..027f7142 --- /dev/null +++ b/causal_net/uoiPoissonLag1/README.Jan @@ -0,0 +1,12 @@ +#IMG=nersc/causal-net:v4 # May 13 +IMG=nersc/causal-net:v5 # Nov 5 2025 +export OMP_NUM_THREADS=2 +salloc -q interactive -C cpu --image=$IMG -t 4:00:00 -A m2043 -N 4 + +One time install of UoI software insise new shifetr image + +Upgrade build tools inside the venv: +shifter python -m pip install -U pip setuptools wheel + +Do the editable install: +shifter python -m pip install -e /global/cfs/cdirs/mpccc/balewski/2025_UoI-VAR diff --git a/causal_net/uoiPoissonLag1/Util_poissonFdr.py b/causal_net/uoiPoissonLag1/Util_poissonFdr.py new file mode 100644 index 00000000..992d4e2b --- /dev/null +++ b/causal_net/uoiPoissonLag1/Util_poissonFdr.py @@ -0,0 +1,114 @@ +#...!...!.................. +import os +import numpy as np + + +#...!...!.................. +def qa_Bfit(B_truth, B_fit): + assert B_truth.shape == B_fit.shape + bterm_res = B_fit - B_truth + bterm_mean = np.mean(bterm_res) + bterm_std = np.std(bterm_res) + N = bterm_res.shape[0] + bterm_se_s = 0. + if N > 1: + bterm_se_s = bterm_std / np.sqrt(2 * (N - 1)) + print('B term: mean=%.2f std=%.3f +/- %.3f' % (bterm_mean, bterm_std, bterm_se_s)) + + out = { + 'tval': B_truth, + 'fval': B_fit, + 'res_mean': bterm_mean, + 'res_std': bterm_std + } + return out + + +#...!...!.................. +def qa_Afit(A_truth, A_fit): + assert A_truth.shape == A_fit.shape + assert A_truth.shape[0] == A_truth.shape[1] + + M = A_truth.shape[0] + diag_m = np.eye(M, dtype=bool) + out = {} + + # --- Diagonal elements + diag_tval = A_truth[diag_m] + diag_fval = A_fit[diag_m] + res_diag = diag_fval - diag_tval + diag_mean = np.mean(res_diag) if res_diag.size > 0 else 0 + diag_std = np.std(res_diag) if res_diag.size > 0 else 0 + out['diag'] = { + 'tval': diag_tval, + 'fval': diag_fval, + 'res_mean': diag_mean, + 'res_std': diag_std + } + + # --- Off-diagonal elements (Excitatory and Inhibitory) + for name in ['inh','exc']: + if name == 'exc': + true_m = (A_truth > 0) & (~diag_m) + pred_m = (A_fit > 0) & (~diag_m) + else: # inh + true_m = (A_truth < 0) & (~diag_m) + pred_m = (A_fit < 0) & (~diag_m) + + tp_m = true_m & pred_m + tval = A_truth[tp_m] + fval = A_fit[tp_m] + res = fval - tval + res_mean = np.mean(res) if res.size > 0 else 0 + res_std = np.std(res) if res.size > 0 else 0 + + TP = np.sum(tp_m) + FP = np.sum((~true_m) & pred_m) + FN = np.sum(true_m & (~pred_m)) + + out[name] = { + 'tval': tval, + 'fval': fval, + 'res_mean': res_mean, + 'res_std': res_std, + 'TP': TP, + 'FP': FP, + 'FN': FN + } + + # Keep printing for compatibility + for name, stats in [('exc', out['exc']), ('inh', out['inh'])]: + res_N = stats['tval'].shape[0] + se_s = 0. + if res_N > 1: + se_s = stats['res_std'] / np.sqrt(2 * (res_N - 1)) + + print('A type=%s TP:%d FP:%d FN:%d TP: mean=%.2f std=%.3f +/- %.3f'%(name, stats['TP'], stats['FP'], stats['FN'], stats['res_mean'], stats['res_std'], se_s)) + + diag_stats = out['diag'] + diag_N = diag_stats['tval'].shape[0] + diag_se_s = 0. + if diag_N > 1: + diag_se_s = diag_stats['res_std'] / np.sqrt(2 * (diag_N - 1)) + + print('A diag: mean=%.2f std=%.3f +/- %.3f'%(diag_stats['res_mean'], diag_stats['res_std'], diag_se_s)) + + return out + + +#...!...!.................. +def extract_loss_traces(model, loss_stride): + loss_iter = np.empty(0, dtype=np.int64) + loss_l1 = np.empty(0, dtype=np.float64) + + if loss_stride is None: + loss_stride = 10 + + if model is not None and hasattr(model, 'loss'): + loss = model.loss + if loss is not None and 'l1' in loss: + loss_l1 = np.asarray(loss['l1'], dtype=np.float64) + loss_iter = np.arange(loss_l1.size, dtype=np.int64) * loss_stride + + return loss_iter, loss_l1 + diff --git a/causal_net/uoiPoissonLag1/batchFitUoI.slr b/causal_net/uoiPoissonLag1/batchFitUoI.slr new file mode 100755 index 00000000..bfc8f4a4 --- /dev/null +++ b/causal_net/uoiPoissonLag1/batchFitUoI.slr @@ -0,0 +1,48 @@ +#!/bin/bash +#SBATCH -N 4 -C cpu --exclusive -A m2043 +# SBATCH --time=30:00 -q debug +#SBATCH --time 5:28:00 -q regular +#SBATCH --output=out/%j.out +#SBATCH --licenses=scratch +#SBATCH --image nersc/causal-net:v5 # Nov 5 2025 + +# - - - E N D O F SLURM C O M M A N D S +set -u ; # exit if you try to use an uninitialized variable + +export OMP_NUM_THREADS=2 + +# only Jan owns those 2 datasets: +inpPath=/pscratch/sd/b/balewski/2025Dec6_causalNetData/ +dataPath=/pscratch/sd/b/balewski/2025Dec6_causalNetFit/ +#dataPath=/pscratch/sd/b/balewski/2025_causalNet_tmp/ + +#dataName=daleM140r4Hz +#dataName=daleM140r2Hz +dataName=daleM140r1Hz +#dataName=daleM140r05Hz + +fdrRate=0.01 + +#samples=100_000 ; maxIter=50 +samples=400_000 ; maxIter=300 +#samples=1000_000 ; maxIter=200 ; fdrRate=0.002 +#samples=1600_000 ; maxIter=150 +#samples=2000_000 ; maxIter=100 + +N=${SLURM_NNODES} +jobId=${SLURM_JOBID} +ntask=$((32 * N)) +echo S: job dataName=$dataName samples=$samples ntask=$ntask maxIter $maxIter + +last_four=${jobId: -4} # Extract last 4 characters +fitName=${dataName}_uoi${last_four} + +time srun -n $ntask --distribution=block:block shifter python fit_uoiPoissonLag1.py --dataName $dataName --samples $samples --maxIter $maxIter --verb 2 --fdrRate $fdrRate --fitName $fitName --dataPath $inpPath --outPath $dataPath +echo +echo 'S:fit done' +echo + shifter ./eval_fitUoI.py --dataPath $dataPath --dataName $fitName -p a b c e -X + +echo +echo 'S:eval done' +echo diff --git a/causal_net/uoiPoissonLag1/eval_fitUoI.py b/causal_net/uoiPoissonLag1/eval_fitUoI.py new file mode 100755 index 00000000..bb5138d3 --- /dev/null +++ b/causal_net/uoiPoissonLag1/eval_fitUoI.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +""" + +""" + +import numpy as np +import os +import argparse +import sys +from toolbox.Util_NumpyIO import read_data_npz +from PlotterFitEval2 import Plotter +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz + +from pprint import pprint + +######################### +# MAIN +######################### + +def main(): + parser = argparse.ArgumentParser(description="Evaluate and plot results from fit_poisson.py") + parser.add_argument("--dataName", type=str, default='dale_M120_3M', help="Base name for the dataset") + parser.add_argument("--dataPath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="Path to the data directory") + parser.add_argument('-p',"--showPlots", type=str,nargs='+', default="a b", help="Plot types to show: a=summary, b=correlations, c=correlations_for_kris") + parser.add_argument("--outPath", type=str, default=None, help="Output path for plots (defaults to dataPath)") + parser.add_argument('-X',"--noXterm", action="store_true", help="Disable X terminal for plotting") + parser.add_argument('-v',"--verb", type=int, default=1, help="Verbosity level") + + args = parser.parse_args() + np.set_printoptions(precision=3) + + if args.outPath is None: args.outPath = args.dataPath + args.showPlots=''.join(args.showPlots) + print(vars(args)) + + # Load fit results + fitFF = os.path.join(args.dataPath, f"{args.dataName}.uoiFdr.npz") + fitD, fitMD = read_data_npz(fitFF) + MD={ **fitMD, 'short_name': args.dataName} + + if 0: # patch old data + #pprint(fitMD) + fitMD['fit_type']='uoiFdr' + #pprint(fitMD) + + if args.verb>1: + pprint(fitMD); exit(1) + + # Load spike data for frequency sorting + spikeF = fitMD['fit_uoi']['fit_input_name'] + inpPath= fitMD['fit_uoi']['fit_input_path'] + spikesFF = os.path.join(inpPath, f"{spikeF}.spikes.npz") + spikeD, spikeMD = read_data_npz(spikesFF) + #pprint(spikeMD) + + if fitMD['spike_data']['data_type']=='simDale': + truthFF = os.path.join(inpPath, f"{spikeF}.simTruth.npz") + trueD,trueMD = read_data_npz(truthFF) + #pprint(trueMD) + for xx in ['short_name']: + trueMD.pop(xx) + # Combine metadata just for plotter + MD.update( **trueMD) + + # Setup plotter + args.prjName = args.dataName + plot = Plotter(args) + + if 'a' in args.showPlots: + plot.summary_fitUoI(fitD,trueD,MD,figId=1) + if 'b' in args.showPlots: + MD['resRange']={'inh':(-0.15,0.1),'exc':(-0.07,0.11),'diag':(-0.2,0.2),'bterm':(-0.3,0.3)} + plot.correlations_fitUoI(fitD,trueD,MD,figId=2) + if 'c' in args.showPlots: + plot.correlation_for_kris(fitD,trueD,MD,figId=3) + if 'd' in args.showPlots: + plot.eigenvalues_fitUoI(fitD,trueD,MD,figId=4) + if 'e' in args.showPlots: + plot.pseudospectra_fitUoI(fitD,trueD,MD,figId=5) + + plot.display_all() + +if __name__ == "__main__": + main() diff --git a/causal_net/uoiPoissonLag1/fit_uoiPoissonLag1.py b/causal_net/uoiPoissonLag1/fit_uoiPoissonLag1.py new file mode 100755 index 00000000..47e3e023 --- /dev/null +++ b/causal_net/uoiPoissonLag1/fit_uoiPoissonLag1.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 + +''' + +IMG=nersc/causal-net:v5 # Nov 5 2025 +export OMP_NUM_THREADS=2 +salloc -q interactive -C cpu --image=$IMG -t 4:00:00 -A m2043 -N 4 + +time srun -n 128 --distribution=block:block shifter python fit_uoiPoissonLag1.py --samples 10_000 + +Existing samples @ /pscratch/sd/y/yxu2/data: + --dataName daleM20_746c4b 4 min 100k samp @ N=4 + --dataName daleM40_e33e89 9 min 100k samp @ N=4 + --dataName daleM80_285c84 75 min 300k samp @ N=4 + --dataName daleM150_448b86 62 min 300k samp @ N=4 + +Existing samples @ dataPath=/pscratch/sd/b/balewski/2025_causalNet_tmp/ + --dataName daleM300_aa61f5 +# hard case +--dataName daleM130_6eb245 + +# Yao: +n_process should be multiple of n_admm, and at MOST n_admm*n_boot*n_reg_param + Example run command on NERSC interactive compute node session: + srun -n 768 --ntasks-per-node=192 --distribution=block:block python uoi_var_poisson_test.py + + +UoI_Poisson parameters: + dt = 0.01 - sample time bin size for Poisson process + n_admm = 32 - number of ADMM processes + n_lambdas = 4 - number of L1 penalty values to test + manual_l1_range = [6e-7, 3e-6] - Hardcoded L1 penalty range + ??? imbalance_tolerance = 10 - tolerance for load imbalance + n_boots_sel = 6 - number of bootstrap samples for selection + n_boots_est = 6 - number of bootstrap samples for estimation + max_iter = 1000 - maximum number of iterations + selection_frac = 0.9 - Fraction of total data used for each bootstrap + lag = 1 - VAR lag order, must be 1 + rho_scaler = 1.0 - ADMM rho scaling factor + seed = 22 - random state seed + fdr_rate=0.05 - False Discovery Rate (FDR) , p-value for selection of existing edges +''' + +import os,sys +import numpy as np +import scipy.sparse as sparse +from numpy.linalg import norm +import argparse +import secrets + +from mpi4py import MPI +from time import time +from pprint import pprint +sys.path.append("/global/homes/b/balewski/prjs/2025_UoI-VAR/") +from examples.var_utils import * +from src.pyuoi.linear_model import * +sys.path.append("/global/homes/b/balewski/prjs/2025_UoI-VAR/src/pyuoi/linear_model") +from sparse_comm_util import build_bootstrap_comm +from var_utils import * +from Util_poissonFdr import extract_loss_traces , qa_Bfit, qa_Afit + +from toolbox.Util_NumpyIO import read_data_npz, write_data_npz +#...!...!.................. +def main(): + parser = argparse.ArgumentParser(description='UoI-VAR Poisson ADMM test') + parser.add_argument('--dataName', default='daleM20_746c4b', help='dataset name (e.g., daleM20_746c4b)') + parser.add_argument("--dataPath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/") + parser.add_argument("--outPath", type=str, default="/pscratch/sd/b/balewski/2025_causalNet_tmp/", help="output directory for fit artifacts") + + parser.add_argument('--samples', type=int, default=100000, help='number of data samples to use') + parser.add_argument('--fitName', default=None, help='output file name core') + parser.add_argument('--dropFreqWeight', action='store_true', help='disable frequency dependent weights; default uses frequency weights') + parser.add_argument('--maxIter', type=int, default=300, help='maximum number of iterations for ADMM') + parser.add_argument('--fdrRate', type=float, default=0.02, help='False discovery rate level for support selection') + parser.add_argument('--selectonFrac', type=float, default=0.8, help='fraction of data used per bootstrap selection') + parser.add_argument('--numBoots', type=int, default=5, help='number of bootstrap samples for selection/estimation') + parser.add_argument('--verb', '-v', type=int, default=1, help='Verbosity level') + args = parser.parse_args() + + confUoI = { + 'fit_VAR': True, + 'fit_intercept': False, + 'standardize': False, + #'manual_l1_range': [6e-7, 3e-6], 'n_lambdas': 4, + 'manual_l1_range': [6e-7, 7e-7], 'n_lambdas': 2, # needs to be changed with M-neurons + 'n_boots_sel': args.numBoots, 'n_boots_est': args.numBoots, + 'selection_frac': args.selectonFrac, + 'max_iter': args.maxIter, + 'rho_scaler': 1.0, 'imbalance_tolerance': 1, # do not change + 'n_admm': 32, + 'solver': 'admm', + 'dt': 0.01, # must be hardcoded - or use broadcasting to all ranks + 'fdr_rate': args.fdrRate, + 'estimation_solver': "lbfgs", + 'random_state': 22, + } + + + # end config print + lag = 1 #confMisc['model_lag'] + assert lag==1 + + rank = 0 + comm = MPI.COMM_WORLD + world_size = comm.Get_size() + if comm is not None: + rank = comm.rank + + if rank == 0: + for arg in vars(args): + print( 'myArgs:',arg, getattr(args, arg)) + + hash_str = secrets.token_hex(3) + outName = args.fitName + if outName is None: + outName = f'{args.dataName}_uoi{hash_str}' + + if args.verb > 1: + print('confUoI:'); pprint(confUoI) + #Xprint('confMisc:'); pprint(confMisc) + + print('Start dataName=%s samples=%d'%(args.dataName, args.samples)) + spikeF=os.path.join(args.dataPath, "{}.spikes.npz".format(args.dataName)) + spikeD, spikeMD = read_data_npz(spikeF, verb=True) + data=spikeD['spikes'][:args.samples].astype(np.double) + single_rates=spikeD['single_rates'] + data_pois = None + + #1confUoI['dt']=spikeMD['time_step_sec'] # can't do it now w/o broadcasting to all ranks + # Slurm context (assume run with srun) + slurm_nodes = int(os.environ.get('SLURM_NNODES') or os.environ.get('SLURM_JOB_NUM_NODES') or 0) + slumr_ranks = int(os.environ.get('SLURM_NTASKS') or world_size) + uoiMD = { 'fit_output_name': outName, 'fit_input_name': args.dataName, 'fit_input_path':args.dataPath,'num_samples_used': args.samples, 'max_iter': args.maxIter, 'fdr_rate': args.fdrRate, 'use_freq_weight': not args.dropFreqWeight, 'num_neurons': data.shape[1], 'sampl_selection_frac': args.selectonFrac, 'slurm_nodes': slurm_nodes, 'slurm_ranks': slumr_ranks } + outMD={'spike_data': spikeMD,'fit_uoi':uoiMD ,'fit_type':'uoiFdr','hash':hash_str,'conf_uoi':confUoI} + #pprint(outMD) + if args.dropFreqWeight: + w = np.ones(data.shape[1]) + else: + rates=np.mean(data,axis = 0)/confUoI['dt'] + rates = np.clip(rates, 0.1, 50) + w=1/rates + w/=np.sum(w) + w*=data.shape[1] + #1print('rates:',rates) + #1print('w',w) + + else: + w = None + + w = comm.bcast(w, root=0) + if rank == 0: print('%dk samples loaded to all %d ranks'%(args.samples//1000,world_size ),flush=True) + + #fitting with multiple processes + boot_comm = build_bootstrap_comm(comm, confUoI['n_admm']) + confUoI_fit = confUoI.copy() + confUoI_fit.pop('loss_stride', None) + + uoi_poisson = UoI_Poisson(**confUoI_fit, comm=boot_comm, global_comm=comm, weights=w) + assert uoi_poisson.solver == "admm" + + Tstart = time() + if boot_comm is not None: #if the global_rank is part of the boostrap distribution(not admm distribution) + if boot_comm.rank == 0: + uoi_poisson.fit(lag, data = data, data_pois = data_pois) + else: + uoi_poisson.fit(lag) + else: + uoi_poisson.admm_queue() + + if rank > 0: return + elaT= time()-Tstart + print("rank %d : Fitting complete in %d seconds."%(comm.rank, elaT), flush = True) + + #... finalize meta-data + confUoI['model_lag']=lag + uoiMD['training_time_sec']=elaT + outMD['short_name']=uoiMD['fit_output_name'] + A_fit = uoi_poisson.VAR_coef_[0] + B_fit = uoi_poisson.VAR_bias_ + + loss_stride =10 # hardcoded by Yao + sel_loss_iter, sel_loss_l1 = extract_loss_traces(getattr(uoi_poisson, '_selection_lm', None), loss_stride) + est_loss_iter, est_loss_l1 = extract_loss_traces(getattr(uoi_poisson, '_estimation_lm', None), loss_stride) + + l1_loss_sel = np.vstack((sel_loss_iter, sel_loss_l1)) if sel_loss_iter.size > 0 else np.empty((2, 0)) + l1_loss_est = np.vstack((est_loss_iter, est_loss_l1)) if est_loss_iter.size > 0 else np.empty((2, 0)) + + bigD={ 'A_uoiFdr':A_fit, 'B_uoiFdr':B_fit, 'l1_loss_sel':l1_loss_sel, 'l1_loss_est':l1_loss_est, 'single_rates':single_rates} + + outF = os.path.join(args.outPath, "{}.uoiFdr.npz".format(outName)) + + write_data_npz(bigD, outF, metaD=outMD) + print('saved output to:',outF) + #pprint(outMD) + + #.... evaluation of results + truthF=spikeF.replace('.spikes','.simTruth') + A_truth = np.load(truthF)["A_true"] + B_truth = np.load(truthF)["B_true"] + + TP, FP, TN, FN = matrix_comparison(A_truth, A_fit, threshold=0) + # these two adds up == real sparsity in B_truth + M=A_truth.shape[0]; M2=M*M + print('dataName=%s M=%d M^2=%d '%(args.dataName,M,M2)) + print("TP p=%.3e n=%d "%(TP,TP*M2)) + print("FN p=%.3e n=%d "%(FN,FN*M2)) + print("FP p=%.3e n=%d "%(FP,FP*M2)) + print("TN p=%.3e "% TN) + + + print('detailed QA %s M=%d M^2=%d samples=%d/k FDR rate=%.4f' %(args.dataName,M,M2,args.samples/1000,args.fdrRate)) + qaD=qa_Afit(A_truth, A_fit) + qaD['bterm']=qa_Bfit(B_truth, B_fit) + + #if spikeMD['data_type']=='simDale': flags=' -p a c ' + # else: + flags=' -p a c e ' + print('\n shifter ./eval_fitUoI.py --dataPath $fitPath --dataName %s %s -X ' % (outName,flags)) + + print(' --dataPath '+args.fitPath) + + +#...!...!.................. +if __name__ == "__main__": + main() + + + + + + + + + + diff --git a/causal_net/uoiPoissonLag1/toolbox b/causal_net/uoiPoissonLag1/toolbox new file mode 120000 index 00000000..c9203e9b --- /dev/null +++ b/causal_net/uoiPoissonLag1/toolbox @@ -0,0 +1 @@ +../toolbox/ \ No newline at end of file diff --git a/examples/README.Jan b/examples/README.Jan new file mode 100644 index 00000000..027f7142 --- /dev/null +++ b/examples/README.Jan @@ -0,0 +1,12 @@ +#IMG=nersc/causal-net:v4 # May 13 +IMG=nersc/causal-net:v5 # Nov 5 2025 +export OMP_NUM_THREADS=2 +salloc -q interactive -C cpu --image=$IMG -t 4:00:00 -A m2043 -N 4 + +One time install of UoI software insise new shifetr image + +Upgrade build tools inside the venv: +shifter python -m pip install -U pip setuptools wheel + +Do the editable install: +shifter python -m pip install -e /global/cfs/cdirs/mpccc/balewski/2025_UoI-VAR diff --git a/examples/UoI_VAR.ipynb b/examples/UoI_VAR.ipynb new file mode 100644 index 00000000..af6d9784 --- /dev/null +++ b/examples/UoI_VAR.ipynb @@ -0,0 +1,329 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "74afacf7", + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "8b97f65b-e58e-41c5-94c6-2a3242eb9864", + "metadata": {}, + "outputs": [], + "source": [ + "\"\"\"\n", + ".. _uoi_lasso:\n", + "\n", + "UoI-Lasso for sparse, minimal bias, regression\n", + "=============================r[i================\n", + "\n", + "This example with demonstrate the ability of UoI-Lasso to recover sparse\n", + "models with minimal bias.\n", + "\n", + "\"\"\"\n", + "\n", + "###############################################################################\n", + "# Load synthetic data\n", + "# -------------------\n", + "#\n", + "# The synthetic data will have 40 features, 10 of which are informative and\n", + "# 1 response variable.\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "b669ae76-1de6-4c1d-bda1-5a1f65b47f11", + "metadata": {}, + "outputs": [ + { + "ename": "ModuleNotFoundError", + "evalue": "No module named 'pandas'", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mModuleNotFoundError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[3], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpandas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mvar_utils\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mtqdm\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m tqdm\n", + "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'pandas'" + ] + } + ], + "source": [ + "import pandas as pd\n", + "from var_utils import *\n", + "from tqdm import tqdm" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fdbba236-1b3f-45fe-b053-cfe649798f48", + "metadata": {}, + "outputs": [], + "source": [ + "X = np.load(\"/global/cfs/cdirs/m2043/causal_inference/dale-gen-sim-may27/input_uoi/Xt_daleM20may27_simu.npy\")" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "9dc3cdfb-0530-4351-bfaf-5ba502a912a9", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(array([ 96., 861., 4796., 15252., 26232., 25179., 13233., 3732.,\n", + " 569., 50.]),\n", + " array([-1.03027344, -0.8203125 , -0.61083984, -0.40136719, -0.19140625,\n", + " 0.01855469, 0.22753906, 0.4375 , 0.64746094, 0.85742188,\n", + " 1.06640625]),\n", + " )" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + }, + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjkAAAGdCAYAAADwjmIIAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAJpFJREFUeJzt3X9QlPeBx/EPoCxq3KVGYWUk/oiNSiRiUHHTxCQn41pJGi7enRrHoCF6OpCJkqrYemhy1yNnm0ZzMXKZtNLOyFW9ieYqCYZi1EtEjSinksBES2I8s2hiYJUaVHjujw5P3Yo/QHHly/s1s1P3eb7P7vfZJ1ve87D7EGJZliUAAADDhAZ7AgAAAO2ByAEAAEYicgAAgJGIHAAAYCQiBwAAGInIAQAARiJyAACAkYgcAABgpC7BnkAwNTU16cSJE+rZs6dCQkKCPR0AAHAdLMvSmTNnFBMTo9DQK5+v6dSRc+LECcXGxgZ7GgAAoA2+/PJL9evX74rrO3Xk9OzZU9KfXySn0xnk2QAAgOvh9/sVGxtr/xy/kk4dOc2/onI6nUQOAAAdzLU+asIHjwEAgJGIHAAAYCQiBwAAGInIAQAARiJyAACAkYgcAABgJCIHAAAYicgBAABGInIAAICRiBwAAGAkIgcAABiJyAEAAEYicgAAgJGIHAAAYKQuwZ4AgI5hQHZhsKfQap+/nBLsKQAIIs7kAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjNSqyMnNzdXo0aPVs2dPRUVFKTU1VVVVVQFjHnnkEYWEhATc5s6dGzDm2LFjSklJUffu3RUVFaWFCxfq4sWLAWO2b9+u+++/Xw6HQ4MHD1Z+fv5l81m9erUGDBigiIgIJSUlae/eva3ZHQAAYLBW/RXyHTt2KCMjQ6NHj9bFixf1k5/8RBMmTNAnn3yiHj162ONmz56tl156yb7fvXt3+9+NjY1KSUmR2+3Wrl279NVXX+npp59W165d9a//+q+SpOrqaqWkpGju3Llat26dSkpK9Oyzz6pv377yer2SpPXr1ysrK0t5eXlKSkrSypUr5fV6VVVVpaioqBt6UQCYgb+cDnRuIZZlWW3d+NSpU4qKitKOHTs0btw4SX8+k5OQkKCVK1e2uM17772nxx57TCdOnFB0dLQkKS8vT4sXL9apU6cUHh6uxYsXq7CwUIcPH7a3mzp1qmpra1VUVCRJSkpK0ujRo/X6669LkpqamhQbG6vnnntO2dnZ1zV/v98vl8uluro6OZ3Otr4MQKfQEYOhIyJygGu73p/fN/SZnLq6OklSr169ApavW7dOvXv31vDhw7VkyRL96U9/steVlpYqPj7eDhxJ8nq98vv9qqiosMckJycHPKbX61Vpaakk6fz58yorKwsYExoaquTkZHtMSxoaGuT3+wNuAADATK36ddWlmpqaNH/+fP3gBz/Q8OHD7eVPPfWU+vfvr5iYGB08eFCLFy9WVVWV3n77bUmSz+cLCBxJ9n2fz3fVMX6/X+fOndO3336rxsbGFsdUVlZecc65ubl68cUX27rLAACgA2lz5GRkZOjw4cP68MMPA5bPmTPH/nd8fLz69u2r8ePH6+jRo7r77rvbPtObYMmSJcrKyrLv+/1+xcbGBnFGAACgvbQpcjIzM7Vlyxbt3LlT/fr1u+rYpKQkSdKRI0d09913y+12X/YtqJqaGkmS2+22/7d52aVjnE6nunXrprCwMIWFhbU4pvkxWuJwOORwOK5vJwEAQIfWqs/kWJalzMxMbdq0Sdu2bdPAgQOvuU15ebkkqW/fvpIkj8ejQ4cO6eTJk/aY4uJiOZ1OxcXF2WNKSkoCHqe4uFgej0eSFB4ersTExIAxTU1NKikpsccAAIDOrVVncjIyMlRQUKB33nlHPXv2tD9D43K51K1bNx09elQFBQWaNGmS7rzzTh08eFALFizQuHHjdN9990mSJkyYoLi4OM2YMUMrVqyQz+fT0qVLlZGRYZ9lmTt3rl5//XUtWrRIzzzzjLZt26YNGzaosPAv3+7IyspSWlqaRo0apTFjxmjlypWqr6/XrFmzbtZrAwAAOrBWRc6aNWsk/flr4pdau3atZs6cqfDwcP3hD3+wgyM2NlaTJ0/W0qVL7bFhYWHasmWL5s2bJ4/Hox49eigtLS3gujoDBw5UYWGhFixYoFWrVqlfv35666237GvkSNKUKVN06tQp5eTkyOfzKSEhQUVFRZd9GBkAAHRON3SdnI6O6+QA14/r5NwaXCcHuLZbcp0cAACA2xWRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIrYqc3NxcjR49Wj179lRUVJRSU1NVVVUVMOa7775TRkaG7rzzTt1xxx2aPHmyampqAsYcO3ZMKSkp6t69u6KiorRw4UJdvHgxYMz27dt1//33y+FwaPDgwcrPz79sPqtXr9aAAQMUERGhpKQk7d27tzW7AwAADNaqyNmxY4cyMjK0e/duFRcX68KFC5owYYLq6+vtMQsWLNDvf/97bdy4UTt27NCJEyf05JNP2usbGxuVkpKi8+fPa9euXfrNb36j/Px85eTk2GOqq6uVkpKiRx99VOXl5Zo/f76effZZbd261R6zfv16ZWVladmyZdq/f79GjBghr9erkydP3sjrAQAADBFiWZbV1o1PnTqlqKgo7dixQ+PGjVNdXZ369OmjgoIC/d3f/Z0kqbKyUsOGDVNpaanGjh2r9957T4899phOnDih6OhoSVJeXp4WL16sU6dOKTw8XIsXL1ZhYaEOHz5sP9fUqVNVW1uroqIiSVJSUpJGjx6t119/XZLU1NSk2NhYPffcc8rOzr6u+fv9frlcLtXV1cnpdLb1ZQA6hQHZhcGeQqfw+cspwZ4CcNu73p/fN/SZnLq6OklSr169JEllZWW6cOGCkpOT7TFDhw7VXXfdpdLSUklSaWmp4uPj7cCRJK/XK7/fr4qKCnvMpY/RPKb5Mc6fP6+ysrKAMaGhoUpOTrbHtKShoUF+vz/gBgAAzNTmyGlqatL8+fP1gx/8QMOHD5ck+Xw+hYeHKzIyMmBsdHS0fD6fPebSwGle37zuamP8fr/OnTunr7/+Wo2NjS2OaX6MluTm5srlctm32NjY1u84AADoENocORkZGTp8+LB+97vf3cz5tKslS5aorq7Ovn355ZfBnhIAAGgnXdqyUWZmprZs2aKdO3eqX79+9nK3263z58+rtrY24GxOTU2N3G63PeavvwXV/O2rS8f89Teyampq5HQ61a1bN4WFhSksLKzFMc2P0RKHwyGHw9H6HQYAAB1Oq87kWJalzMxMbdq0Sdu2bdPAgQMD1icmJqpr164qKSmxl1VVVenYsWPyeDySJI/Ho0OHDgV8C6q4uFhOp1NxcXH2mEsfo3lM82OEh4crMTExYExTU5NKSkrsMQAAoHNr1ZmcjIwMFRQU6J133lHPnj3tz7+4XC5169ZNLpdL6enpysrKUq9eveR0OvXcc8/J4/Fo7NixkqQJEyYoLi5OM2bM0IoVK+Tz+bR06VJlZGTYZ1nmzp2r119/XYsWLdIzzzyjbdu2acOGDSos/Mu3O7KyspSWlqZRo0ZpzJgxWrlyperr6zVr1qyb9doAAIAOrFWRs2bNGknSI488ErB87dq1mjlzpiTp1VdfVWhoqCZPnqyGhgZ5vV698cYb9tiwsDBt2bJF8+bNk8fjUY8ePZSWlqaXXnrJHjNw4EAVFhZqwYIFWrVqlfr166e33npLXq/XHjNlyhSdOnVKOTk58vl8SkhIUFFR0WUfRgYAAJ3TDV0np6PjOjnA9eM6ObcG18kBru2WXCcHAADgdkXkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADBSqyNn586devzxxxUTE6OQkBBt3rw5YP3MmTMVEhIScJs4cWLAmNOnT2v69OlyOp2KjIxUenq6zp49GzDm4MGDeuihhxQREaHY2FitWLHisrls3LhRQ4cOVUREhOLj4/Xuu++2dncAAIChurR2g/r6eo0YMULPPPOMnnzyyRbHTJw4UWvXrrXvOxyOgPXTp0/XV199peLiYl24cEGzZs3SnDlzVFBQIEny+/2aMGGCkpOTlZeXp0OHDumZZ55RZGSk5syZI0natWuXpk2bptzcXD322GMqKChQamqq9u/fr+HDh7d2t4BbZkB2YbCnAACdQohlWVabNw4J0aZNm5Sammovmzlzpmpray87w9Ps008/VVxcnD7++GONGjVKklRUVKRJkybp+PHjiomJ0Zo1a/TTn/5UPp9P4eHhkqTs7Gxt3rxZlZWVkqQpU6aovr5eW7ZssR977NixSkhIUF5e3nXN3+/3y+Vyqa6uTk6nsw2vANB6RA6u5vOXU4I9BeC2d70/v9vlMznbt29XVFSUhgwZonnz5umbb76x15WWlioyMtIOHElKTk5WaGio9uzZY48ZN26cHTiS5PV6VVVVpW+//dYek5ycHPC8Xq9XpaWlV5xXQ0OD/H5/wA0AAJjppkfOxIkT9dvf/lYlJSX6t3/7N+3YsUM//OEP1djYKEny+XyKiooK2KZLly7q1auXfD6fPSY6OjpgTPP9a41pXt+S3NxcuVwu+xYbG3tjOwsAAG5brf5MzrVMnTrV/nd8fLzuu+8+3X333dq+fbvGjx9/s5+uVZYsWaKsrCz7vt/vJ3QAADBUu3+FfNCgQerdu7eOHDkiSXK73Tp58mTAmIsXL+r06dNyu932mJqamoAxzfevNaZ5fUscDoecTmfADQAAmKndI+f48eP65ptv1LdvX0mSx+NRbW2tysrK7DHbtm1TU1OTkpKS7DE7d+7UhQsX7DHFxcUaMmSIvve979ljSkpKAp6ruLhYHo+nvXcJAAB0AK2OnLNnz6q8vFzl5eWSpOrqapWXl+vYsWM6e/asFi5cqN27d+vzzz9XSUmJnnjiCQ0ePFher1eSNGzYME2cOFGzZ8/W3r179dFHHykzM1NTp05VTEyMJOmpp55SeHi40tPTVVFRofXr12vVqlUBv2p6/vnnVVRUpFdeeUWVlZVavny59u3bp8zMzJvwsgAAgI6u1ZGzb98+jRw5UiNHjpQkZWVlaeTIkcrJyVFYWJgOHjyoH/3oR7rnnnuUnp6uxMRE/c///E/AtXLWrVunoUOHavz48Zo0aZIefPBBvfnmm/Z6l8ul999/X9XV1UpMTNQLL7ygnJwc+xo5kvTAAw+ooKBAb775pkaMGKH/+q//0ubNm7lGDgAAkHSD18np6LhODoKB6+TgarhODnBtQb1ODgAAQLAROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAI3UJ9gQAAH8xILsw2FNotc9fTgn2FIAWcSYHAAAYicgBAABGInIAAICRiBwAAGAkIgcAABiJyAEAAEYicgAAgJGIHAAAYCQiBwAAGInIAQAARiJyAACAkYgcAABgJCIHAAAYicgBAABGInIAAICRiBwAAGAkIgcAABiJyAEAAEYicgAAgJGIHAAAYCQiBwAAGInIAQAARiJyAACAkYgcAABgJCIHAAAYicgBAABGInIAAICRiBwAAGAkIgcAABiJyAEAAEYicgAAgJGIHAAAYCQiBwAAGInIAQAARiJyAACAkYgcAABgJCIHAAAYicgBAABGanXk7Ny5U48//rhiYmIUEhKizZs3B6y3LEs5OTnq27evunXrpuTkZH322WcBY06fPq3p06fL6XQqMjJS6enpOnv2bMCYgwcP6qGHHlJERIRiY2O1YsWKy+ayceNGDR06VBEREYqPj9e7777b2t0BAACGanXk1NfXa8SIEVq9enWL61esWKHXXntNeXl52rNnj3r06CGv16vvvvvOHjN9+nRVVFSouLhYW7Zs0c6dOzVnzhx7vd/v14QJE9S/f3+VlZXp5z//uZYvX64333zTHrNr1y5NmzZN6enpOnDggFJTU5WamqrDhw+3dpcAAICBQizLstq8cUiINm3apNTUVEl/PosTExOjF154QT/+8Y8lSXV1dYqOjlZ+fr6mTp2qTz/9VHFxcfr44481atQoSVJRUZEmTZqk48ePKyYmRmvWrNFPf/pT+Xw+hYeHS5Kys7O1efNmVVZWSpKmTJmi+vp6bdmyxZ7P2LFjlZCQoLy8vOuav9/vl8vlUl1dnZxOZ1tfBqBVBmQXBnsKwE31+cspwZ4COpnr/fl9Uz+TU11dLZ/Pp+TkZHuZy+VSUlKSSktLJUmlpaWKjIy0A0eSkpOTFRoaqj179thjxo0bZweOJHm9XlVVVenbb7+1x1z6PM1jmp+nJQ0NDfL7/QE3AABgppsaOT6fT5IUHR0dsDw6Otpe5/P5FBUVFbC+S5cu6tWrV8CYlh7j0ue40pjm9S3Jzc2Vy+Wyb7Gxsa3dRQAA0EF0qm9XLVmyRHV1dfbtyy+/DPaUAABAO7mpkeN2uyVJNTU1ActramrsdW63WydPngxYf/HiRZ0+fTpgTEuPcelzXGlM8/qWOBwOOZ3OgBsAADDTTY2cgQMHyu12q6SkxF7m9/u1Z88eeTweSZLH41Ftba3KysrsMdu2bVNTU5OSkpLsMTt37tSFCxfsMcXFxRoyZIi+973v2WMufZ7mMc3PAwAAOrdWR87Zs2dVXl6u8vJySX/+sHF5ebmOHTumkJAQzZ8/X//yL/+i//7v/9ahQ4f09NNPKyYmxv4G1rBhwzRx4kTNnj1be/fu1UcffaTMzExNnTpVMTExkqSnnnpK4eHhSk9PV0VFhdavX69Vq1YpKyvLnsfzzz+voqIivfLKK6qsrNTy5cu1b98+ZWZm3virAgAAOrwurd1g3759evTRR+37zeGRlpam/Px8LVq0SPX19ZozZ45qa2v14IMPqqioSBEREfY269atU2ZmpsaPH6/Q0FBNnjxZr732mr3e5XLp/fffV0ZGhhITE9W7d2/l5OQEXEvngQceUEFBgZYuXaqf/OQn+v73v6/Nmzdr+PDhbXohAACAWW7oOjkdHdfJQTBwnRyYhuvk4FYLynVyAAAAbhdEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAj3fTIWb58uUJCQgJuQ4cOtdd/9913ysjI0J133qk77rhDkydPVk1NTcBjHDt2TCkpKerevbuioqK0cOFCXbx4MWDM9u3bdf/998vhcGjw4MHKz8+/2bsCAAA6sHY5k3Pvvffqq6++sm8ffvihvW7BggX6/e9/r40bN2rHjh06ceKEnnzySXt9Y2OjUlJSdP78ee3atUu/+c1vlJ+fr5ycHHtMdXW1UlJS9Oijj6q8vFzz58/Xs88+q61bt7bH7gAAgA6oS7s8aJcucrvdly2vq6vTr371KxUUFOhv/uZvJElr167VsGHDtHv3bo0dO1bvv/++PvnkE/3hD39QdHS0EhIS9M///M9avHixli9frvDwcOXl5WngwIF65ZVXJEnDhg3Thx9+qFdffVVer7c9dgkAAHQw7XIm57PPPlNMTIwGDRqk6dOn69ixY5KksrIyXbhwQcnJyfbYoUOH6q677lJpaakkqbS0VPHx8YqOjrbHeL1e+f1+VVRU2GMufYzmMc2PcSUNDQ3y+/0BNwAAYKabHjlJSUnKz89XUVGR1qxZo+rqaj300EM6c+aMfD6fwsPDFRkZGbBNdHS0fD6fJMnn8wUETvP65nVXG+P3+3Xu3Lkrzi03N1cul8u+xcbG3ujuAgCA29RN/3XVD3/4Q/vf9913n5KSktS/f39t2LBB3bp1u9lP1ypLlixRVlaWfd/v9xM6AAAYqt2/Qh4ZGal77rlHR44ckdvt1vnz51VbWxswpqamxv4Mj9vtvuzbVs33rzXG6XReNaQcDoecTmfADQAAmKldPnh8qbNnz+ro0aOaMWOGEhMT1bVrV5WUlGjy5MmSpKqqKh07dkwej0eS5PF49LOf/UwnT55UVFSUJKm4uFhOp1NxcXH2mHfffTfgeYqLi+3HQOcxILsw2FMAANymbvqZnB//+MfasWOHPv/8c+3atUt/+7d/q7CwME2bNk0ul0vp6enKysrSBx98oLKyMs2aNUsej0djx46VJE2YMEFxcXGaMWOG/vd//1dbt27V0qVLlZGRIYfDIUmaO3eu/vjHP2rRokWqrKzUG2+8oQ0bNmjBggU3e3cAAEAHddPP5Bw/flzTpk3TN998oz59+ujBBx/U7t271adPH0nSq6++qtDQUE2ePFkNDQ3yer1644037O3DwsK0ZcsWzZs3Tx6PRz169FBaWppeeukle8zAgQNVWFioBQsWaNWqVerXr5/eeustvj4OAABsIZZlWcGeRLD4/X65XC7V1dXx+ZwOil9XAcH3+cspwZ4COpnr/fnN364CAABGInIAAICRiBwAAGAkIgcAABiJyAEAAEYicgAAgJGIHAAAYCQiBwAAGInIAQAARiJyAACAkdr9r5ADAMzWEf+8Cn+KonPgTA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADASkQMAAIzUJdgTwO1jQHZhsKcAAMBNw5kcAABgJCIHAAAYicgBAABG4jM5AIBOpyN+BvHzl1OCPYUOhzM5AADASEQOAAAwUoePnNWrV2vAgAGKiIhQUlKS9u7dG+wpAQCA20CHjpz169crKytLy5Yt0/79+zVixAh5vV6dPHky2FMDAABB1qEj55e//KVmz56tWbNmKS4uTnl5eerevbt+/etfB3tqAAAgyDrst6vOnz+vsrIyLVmyxF4WGhqq5ORklZaWtrhNQ0ODGhoa7Pt1dXWSJL/ff9PnN3zZ1pv+mACAzqs9flZ1VM2vhWVZVx3XYSPn66+/VmNjo6KjowOWR0dHq7KyssVtcnNz9eKLL162PDY2tl3mCADAzeJaGewZ3H7OnDkjl8t1xfUdNnLaYsmSJcrKyrLvNzU16fTp07rzzjsVEhJyS+bg9/sVGxurL7/8Uk6n85Y8J9qGY9WxcLw6Do5Vx3I7Hi/LsnTmzBnFxMRcdVyHjZzevXsrLCxMNTU1Actramrkdrtb3MbhcMjhcAQsi4yMbK8pXpXT6bxt/mPB1XGsOhaOV8fBsepYbrfjdbUzOM067AePw8PDlZiYqJKSEntZU1OTSkpK5PF4gjgzAABwO+iwZ3IkKSsrS2lpaRo1apTGjBmjlStXqr6+XrNmzQr21AAAQJB16MiZMmWKTp06pZycHPl8PiUkJKioqOiyDyPfThwOh5YtW3bZr81w++FYdSwcr46DY9WxdOTjFWJd6/tXAAAAHVCH/UwOAADA1RA5AADASEQOAAAwEpEDAACMROS0s5/97Gd64IEH1L179+u+8KBlWcrJyVHfvn3VrVs3JScn67PPPmvfiUKSdPr0aU2fPl1Op1ORkZFKT0/X2bNnr7rNI488opCQkIDb3Llzb9GMO5fVq1drwIABioiIUFJSkvbu3XvV8Rs3btTQoUMVERGh+Ph4vfvuu7dopmjNscrPz7/sPRQREXELZ9t57dy5U48//rhiYmIUEhKizZs3X3Ob7du36/7775fD4dDgwYOVn5/f7vNsKyKnnZ0/f15///d/r3nz5l33NitWrNBrr72mvLw87dmzRz169JDX69V3333XjjOFJE2fPl0VFRUqLi7Wli1btHPnTs2ZM+ea282ePVtfffWVfVuxYsUtmG3nsn79emVlZWnZsmXav3+/RowYIa/Xq5MnT7Y4fteuXZo2bZrS09N14MABpaamKjU1VYcPH77FM+98WnuspD9fTffS99AXX3xxC2fcedXX12vEiBFavXr1dY2vrq5WSkqKHn30UZWXl2v+/Pl69tlntXXrbfpHqS3cEmvXrrVcLtc1xzU1NVlut9v6+c9/bi+rra21HA6H9Z//+Z/tOEN88sknliTr448/tpe99957VkhIiPV///d/V9zu4Ycftp5//vlbMMPObcyYMVZGRoZ9v7Gx0YqJibFyc3NbHP8P//APVkpKSsCypKQk6x//8R/bdZ5o/bG63v9/RPuSZG3atOmqYxYtWmTde++9AcumTJlieb3edpxZ23Em5zZTXV0tn8+n5ORke5nL5VJSUpJKS0uDODPzlZaWKjIyUqNGjbKXJScnKzQ0VHv27LnqtuvWrVPv3r01fPhwLVmyRH/605/ae7qdyvnz51VWVhbwvggNDVVycvIV3xelpaUB4yXJ6/XyPmpnbTlWknT27Fn1799fsbGxeuKJJ1RRUXErpotW6mjvqw59xWMT+Xw+Sbrsqs3R0dH2OrQPn8+nqKiogGVdunRRr169rvraP/XUU+rfv79iYmJ08OBBLV68WFVVVXr77bfbe8qdxtdff63GxsYW3xeVlZUtbuPz+XgfBUFbjtWQIUP061//Wvfdd5/q6ur0i1/8Qg888IAqKirUr1+/WzFtXKcrva/8fr/OnTunbt26BWlmLeNMThtkZ2df9iG5v75d6c2MW6+9j9ecOXPk9XoVHx+v6dOn67e//a02bdqko0eP3sS9AMzl8Xj09NNPKyEhQQ8//LDefvtt9enTR//xH/8R7Kmhg+NMThu88MILmjlz5lXHDBo0qE2P7Xa7JUk1NTXq27evvbympkYJCQlteszO7nqPl9vtvuyDkRcvXtTp06ft43I9kpKSJElHjhzR3Xff3er54nK9e/dWWFiYampqApbX1NRc8di43e5WjcfN0ZZj9de6du2qkSNH6siRI+0xRdyAK72vnE7nbXcWRyJy2qRPnz7q06dPuzz2wIED5Xa7VVJSYkeN3+/Xnj17WvUNLfzF9R4vj8ej2tpalZWVKTExUZK0bds2NTU12eFyPcrLyyUpIFJxY8LDw5WYmKiSkhKlpqZKkpqamlRSUqLMzMwWt/F4PCopKdH8+fPtZcXFxfJ4PLdgxp1XW47VX2tsbNShQ4c0adKkdpwp2sLj8Vx2KYbb+n0V7E8+m+6LL76wDhw4YL344ovWHXfcYR04cMA6cOCAdebMGXvMkCFDrLffftu+//LLL1uRkZHWO++8Yx08eNB64oknrIEDB1rnzp0Lxi50KhMnTrRGjhxp7dmzx/rwww+t73//+9a0adPs9cePH7eGDBli7dmzx7Isyzpy5Ij10ksvWfv27bOqq6utd955xxo0aJA1bty4YO2CsX73u99ZDofDys/Ptz755BNrzpw5VmRkpOXz+SzLsqwZM2ZY2dnZ9viPPvrI6tKli/WLX/zC+vTTT61ly5ZZXbt2tQ4dOhSsXeg0WnusXnzxRWvr1q3W0aNHrbKyMmvq1KlWRESEVVFREaxd6DTOnDlj/1ySZP3yl7+0Dhw4YH3xxReWZVlWdna2NWPGDHv8H//4R6t79+7WwoULrU8//dRavXq1FRYWZhUVFQVrF66KyGlnaWlplqTLbh988IE9RpK1du1a+35TU5P1T//0T1Z0dLTlcDis8ePHW1VVVbd+8p3QN998Y02bNs264447LKfTac2aNSsgSKurqwOO37Fjx6xx48ZZvXr1shwOhzV48GBr4cKFVl1dXZD2wGz//u//bt11111WeHi4NWbMGGv37t32uocffthKS0sLGL9hwwbrnnvuscLDw617773XKiwsvMUz7rxac6zmz59vj42OjrYmTZpk7d+/Pwiz7nw++OCDFn9GNR+ftLQ06+GHH75sm4SEBCs8PNwaNGhQwM+v202IZVlWUE4hAQAAtCO+XQUAAIxE5AAAACMROQAAwEhEDgAAMBKRAwAAjETkAAAAIxE5AADASEQOAAAwEpEDAACMROQAAAAjETkAAMBIRA4AADDS/wMlQ8eraCGRDAAAAABJRU5ErkJggg==", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "plt.hist(X[:,1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3a70dcd0-1d9c-469d-9450-1667b483bd3f", + "metadata": {}, + "outputs": [], + "source": [ + "from pyuoi.linear_model import UoI_Poisson" + ] + }, + { + "cell_type": "markdown", + "id": "688690da", + "metadata": {}, + "source": [ + "### Gaussian VAR model" + ] + }, + { + "cell_type": "code", + "execution_count": 56, + "id": "71b226bd", + "metadata": {}, + "outputs": [], + "source": [ + "n_features = 10\n", + "n_samples = 15\n", + "lag = 2\n", + "\n", + "\n", + "data, transition_matrices, cov = generate_sparse_stationary_var_process(\n", + " n_features,\n", + " n_samples,\n", + " lag=lag,\n", + " sparsity=0.5,\n", + " spectral_radius=0.6, # Ensures stationarity\n", + " process_type='gaussian'\n", + "\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a914be9-25d3-4474-baad-1ceae70fc271", + "metadata": {}, + "outputs": [], + "source": [ + "uoi_var = UoI_Lasso(n_real_features = n_features, fit_VAR = True, fit_intercept = False, random_state=42)\n", + "uoi_var.fit(data)\n", + "B_model = uoi_var.VAR_coef_" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a3bdfda4", + "metadata": {}, + "outputs": [], + "source": [ + "# OBSOLETE: construct the lagged connectivity matrices from the model coefficient\n", + "A_model = [B_model.reshape(n_features,n_features*lag).T[i*n_features:(i+1)*n_features].T for i in range(lag)]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a01d3c2e", + "metadata": {}, + "outputs": [], + "source": [ + "# prediction\n", + "yhat = uoi_var.predict(X)\n", + "\n", + "# converted back to time series array\n", + "true_response = Y.reshape((n_samples-lag,n_features),order = \"F\") #shape: each row is a sample of length n_features\n", + "prediction = yhat.reshape((n_samples-lag,n_features),order = \"F\")" + ] + }, + { + "cell_type": "code", + "execution_count": 73, + "id": "fdf6eebd", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Text(0, 0.5, 'relative error between y_hat and y')" + ] + }, + "execution_count": 73, + "metadata": {}, + "output_type": "execute_result" + }, + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAAAkIAAAHFCAYAAAAe+pb9AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAABuyklEQVR4nO3dd1iTZ9sG8DMQkrD3RpYTxIl7Ic46Wqu2r6PuUftKa631q/rauqpia2vtwNHW1daqravVWrdYV917DxQVEEEh7JHc3x+U1AgowYQAOX/HwSG5n4cnV24fkot7SoQQAkREREQmyMzYARAREREZCxMhIiIiMllMhIiIiMhkMREiIiIik8VEiIiIiEwWEyEiIiIyWUyEiIiIyGQxESIiIiKTxUSIiIiITBYToVJauXIlJBJJsV8TJ040dngGt23bNsyYMaPU5w8bNkyrjmQyGapXr46JEydCqVSWKYYZM2ZAIpGU6WefFb+/vz+GDRtWpuuaqgMHDkAul+POnTsGuf6jR4/Qv39/uLm5QSKR4NVXXzXI8+h6X1dl9+7dw/jx4xEWFgYHBwdIJBKsXLmyxPN3796Nli1bwsrKCi4uLhg2bBgSExPLL+AqqLK/F3311Vdo0aIFXFxcIJfL4evri/79++PixYvFnv/111+jTp06kMvlCAgIwMyZM5GXl1fkvMTERAwbNgwuLi6wsrJCy5YtsWfPHq1z8vLyUL16dSxcuFD3wAWVyooVKwQAsWLFCnHkyBGtrzt37hg7PIOLiIgQutwuQ4cOFZaWlpo6+vPPP8XIkSMFANG5c+cyxTB9+nSdYnjSs+I/deqUuHHjRpmua4rUarVo3LixiIiIMNhzjB8/XshkMvHTTz+JI0eOiKtXrxrkeXS9r6uyffv2CRcXF9GpUycxYMAAzftdcaKjo4VUKhW9evUSO3fuFD/99JPw9vYWISEhIjs7u3wDr0L8/PzE0KFDjR1GmU2bNk3MmDFDbNq0SURHR4vly5eLWrVqCWtra3HlyhWtc2fPni0kEomYMmWK2Ldvn/j000+FTCYTo0eP1jovOztbhISECB8fH/HTTz+JnTt3il69egmpVCqio6O1zl25cqVwdHQUSUlJOsXNd4BSKkyEjh8/bpDrZ2RkGOS6+lKWRMja2rpIeXh4uAAgbt26pXMMhkqEKqOS7he1Wi0yMzNf6NqZmZlCrVaXeHzbtm0CQJE3Nn3q1KmTCAoKMtj1Cxnqvqjov8/FUalUmu+PHz/+zESoadOmIjg4WOTl5WnKDh06JACIRYsWGTrUKquyJ0LFuXTpkgAgPvroI01ZUlKSUCgU4s0339Q6d86cOUIikYiLFy9qyqKiogQAcfjwYU1ZXl6eCA4OFs2aNdP6+ZycHOHk5CTmzJmjU4zsGtOz33//XdNcbGtri86dO+PIkSNa5xR28Zw6dQqvvfYaHB0dUb16dQCAEAKLFi1Cw4YNYWlpCUdHR7z22mu4detWkefavn07OnbsCHt7e1hZWSEoKAiRkZGa4ydOnED//v3h7+8PS0tL+Pv7Y8CAAUW6MzIzMzFx4kQEBARAoVDAyckJTZo0wZo1awAUdHNFRUUBgFZ31+3bt3WunyZNmgAAHjx4oFW+bt06tGzZEtbW1rCxsUHXrl1x+vTp515v3bp16NKlCzw9PWFpaYmgoCBMnjwZGRkZmnOeF/+TzdEPHz6ETCbDRx99VOS5rly5AolEgq+++kpTlpCQgDFjxsDHxwcymUzTvJufn1+q+ijN6x42bBhsbGxw/vx5dOnSBba2tujYsaPm9bz99ttYsmQJgoKCIJfLsWrVKgDAwYMH0bFjR9ja2sLKygqtWrXCH3/8oXXtwi7fnTt3YsSIEXB1dYWVlRVycnJKjHnx4sVo2rQpateurVXu7++Pnj17Yvv27WjcuDEsLS1Rp04dLF++vFR1AQC3b9+GRCLB7t27cfnyZc3/VXR0NAAgNzcXs2fP1jSnu7q6Yvjw4Xj48GGRen2R+6IwjuK6hiQSiVZ3mj5+n0+fPo2ePXvCzc0NcrkcXl5e6NGjB+7du1fquntRZmal+zi4f/8+jh8/jsGDB0MqlWrKW7VqhVq1amHTpk1lev727dsjJCQEx48fR9u2bWFlZYXAwEDMmzcParVap2vt3bsX7du3h7OzMywtLeHr64u+ffsiMzNTc87MmTPRvHlzODk5wc7ODo0bN8ayZcsgntqHvPC+3rp1Kxo1aqS5n7Zu3Qqg4HcoKCgI1tbWaNasGU6cOKH184W/vxcvXkTHjh1hbW0NV1dXvP3221rxlESpVGren2UyGby9vTF+/HitexkAfv31VzRv3lzzeRAYGIgRI0boVG+G4OrqCgBa98r27duRnZ2N4cOHa507fPhwCCGwefNmTdmmTZtQu3ZttGzZUlMmlUoxaNAgHDt2DPfv39eUy2Qy9OvXD99++22R/8dnYSKkI5VKhfz8fK2vQj///DN69eoFOzs7rFmzBsuWLcPjx4/Rvn17HDx4sMi1+vTpgxo1auDXX3/FkiVLAABjxozB+PHj0alTJ2zevBmLFi3CxYsX0apVK63kYdmyZejevTvUajWWLFmCLVu2YNy4cVpvnLdv30bt2rWxcOFC7NixA5988gni4+PRtGlTJCUlac6bMGECFi9ejHHjxmH79u348ccf8frrryM5ORkA8NFHH+G1114DABw5ckTz5enpqXP9xcTEQCqVIjAwUFM2d+5cDBgwAMHBwfjll1/w448/Ii0tDW3btsWlS5eeeb3r16+je/fuWLZsGbZv347x48fjl19+wcsvv6w5R5f4XV1d0bNnT6xatarIm++KFSsgk8nwxhtvAChIgpo1a4YdO3Zg2rRp+PPPPzFy5EhERkZi9OjRz60LXV53bm4uXnnlFXTo0AG//fYbZs6cqTm2efNmLF68GNOmTcOOHTvQtm1b7N+/Hx06dEBqaiqWLVuGNWvWwNbWFi+//DLWrVtXJJYRI0bAwsICP/74I9avXw8LC4tiY87NzcXu3bsRHh5e7PGzZ8/i/fffx3vvvYfffvsN9evXx8iRI/HXX389tz4AwNPTE0eOHEGjRo0QGBio+b9q3Lgx1Go1evXqhXnz5mHgwIH4448/MG/ePOzatQvt27dHVlaW5jr6vi9Ko6y/zxkZGejcuTMePHiAqKgo7Nq1CwsXLoSvry/S0tKe+ZxqtbrI+1FxXyqVqkyvqTgXLlwAANSvX7/Isfr162uOl0VCQgLeeOMNDBo0CL///ju6deuGKVOm4Keffir1NW7fvo0ePXpAJpNh+fLl2L59O+bNmwdra2vk5uZqnTdmzBj88ssv2LhxI/r06YN33nkHH3/8cZFrnj17FlOmTMGkSZOwceNG2Nvbo0+fPpg+fTq+//57zJ07F6tXr0Zqaip69uypdS8CBeNXunfvjo4dO2Lz5s14++23sXTpUvTr1++ZryUzMxNhYWFYtWoVxo0bhz///BOTJk3CypUr8corr2g+7I8cOYJ+/fohMDAQa9euxR9//IFp06aV6g+y4j7TivvSJRlVqVTIycnBlStXMGrUKLi5uWklPYX3SL169bR+ztPTEy4uLlr30IULF0q81wAUGX/Uvn173LlzR7f7UKf2IxNW2DVW3FdeXp5QqVTCy8tL1KtXT6uJOS0tTbi5uYlWrVppygq7eKZNm6b1HEeOHBEAxOeff65VfvfuXWFpaSk++OADzTXt7OxEmzZtntmF8bT8/HyRnp4urK2txZdffqkpDwkJEa+++uozf7asXWN5eXkiLy9PJCUlicWLFwszMzPxv//9T3NebGyskEql4p133tH6+bS0NOHh4SH+85//aMqe1zWmVqtFXl6e2L9/vwAgzp49W6r4n26O/v333wUAsXPnTk1Zfn6+8PLyEn379tWUjRkzRtjY2BQZI/bZZ58JAFrNu0/T5XUPHTpUABDLly8vch0Awt7eXjx69EirvEWLFsLNzU2kpaVpvYbCvvbC+6bwvh4yZEiJsT7p6NGjAoBYu3ZtkWN+fn5CoVBo1UdWVpZwcnISY8aMKdX1C4WFhYm6detqla1Zs0YAEBs2bNAqL+zGKalLpiz3RUxMTIldQwDE9OnTNY9f9Pf5xIkTAoDYvHlzsfE/S+FzP+/Lz89Pp+s+q2ts9erVAoA4cuRIkWNvvvmmkMlkOr8OIQr+zwGIo0ePapUHBweLrl27lvo669evFwDEmTNnSv0zKpVK5OXliVmzZglnZ2et91U/Pz9haWkp7t27pyk7c+aMACA8PT21ukE3b94sAIjff/9dU1b4+/vke64QBd1AAMTBgwe1nuvJ96LIyEhhZmZWZEhG4Wvctm2bEOLf95yUlJRSv+ZChfX+vC9duuzkcrnm52rVqiUuXbqkdXz06NFCLpcX+7O1atUSXbp00Ty2sLAo9v3j8OHDAoD4+eeftcqvX78uAIjFixeXOt5/26qoVH744QcEBQVplUmlUly+fBlxcXEYP368VhOzjY0N+vbti6VLlyIzMxNWVlaaY3379tW6ztatWyGRSDBo0CCtTN7DwwMNGjTQdA8cPnwYSqUSY8eOfeYsqvT0dHz88cfYsGEDbt++rfVX4eXLlzXfN2vWDKtXr8bkyZPx0ksvoXnz5rC0tNStYoqRkZFRpGVhwIABmDNnjubxjh07kJ+fjyFDhmi9ZoVCgbCwMOzbt++Zz3Hr1i18+OGH2Lt3LxITE7WaQy9fvlzsXxLP061bN3h4eGDFihXo3LmzJs64uDitpuatW7ciPDwcXl5eWrF369YNEydOxP79+xEcHFzsc5TldT99vxTq0KEDHB0dNY8zMjJw9OhR/Pe//4WNjY2m3NzcHIMHD8akSZNw9epV1KlT57nXflpcXBwAwM3NrdjjDRs2hK+vr9brqVWrll5ml23duhUODg54+eWXteqsYcOG8PDwQHR0NP773/8CMMx98Txl/X2uUaMGHB0dMWnSJMTHx6Ndu3Yl3jdPe/PNN9GzZ8/nnieXy0v/QkqppPeess7sBArqplmzZlpl9evXx5kzZ0p9jYYNG0Imk+HNN9/E2LFj0bZtW60W6EJ79+7F3Llzcfz48SIzWRMTE+Hu7q51TW9vb83jws+A9u3ba72nF5YXd78XtiQXGjhwIKZOnYp9+/ahdevWxb6WrVu3IiQkBA0bNtS6h7p27arpMu7WrRuaNm0KAPjPf/6DkSNHonXr1lrxPsvSpUuf2/IIAC4uLqW6HlDwGZWbm4ubN2/iiy++QHh4OPbs2YO6detqznnWffL0MV3OLXxverLL7HmYCOkoKChIM87lSYXdSMU1q3t5eUGtVuPx48davzRPn/vgwQMIIbR+AZ9U+MtcOB7Cx8fnmbEOHDgQe/bswUcffYSmTZvCzs4OEokE3bt312q6/eqrr+Dj44N169bhk08+gUKhQNeuXTF//nzUrFnzmc/xLJaWlpoukYSEBHz++edYs2YN6tevj8mTJ2teMwDNL/LTnjVuIT09HW3btoVCocDs2bNRq1YtWFlZ4e7du+jTp0+R5unSkkqlGDx4ML7++mukpKTAwcEBK1euhKenJ7p27ao578GDB9iyZUuJ3UhPdj8+TdfXbWVlBTs7u2LPffo+evz4MYQQJd6LwL/3a0nXKElhnSoUimKPOzs7FymTy+Vl/r940oMHD5CSkgKZTFbs8cL6NtR98Txl/X22t7fH/v37MWfOHPzvf//D48eP4enpidGjR+PDDz8s8f4CChKHkpLSJ71IcvK0wv/jp+8hoGDZAycnpxe+9pN0vX+qV6+O3bt349NPP0VERAQyMjIQGBiIcePG4d133wUAHDt2DF26dEH79u3x3Xffacb4bd68GXPmzCnyfE+/psJ7sKTy7OxsrXKpVFrktXl4eAAovh4LPXjwADdu3Hjue0y7du2wefNmfPXVVxgyZAhycnJQt25dTJ06FQMGDCjx+kBBIi5KMZ6mtGPIAKBx48YAgBYtWuCVV15BjRo18L///Q+//fYbgIL/5+zs7CKNA0DBPRQaGqp57OzsXOK9BhT9Pyh8b9LlnmEipCeFN3l8fHyRY3FxcTAzM9P6qx0o+ubk4uICiUSiWaPlaYVlhYPPnjWQMjU1FVu3bsX06dM1SQcA5OTkaG6gQtbW1pg5cyZmzpyJBw8e4M8//8TkyZPx8ssv48qVK8962c9kZmamlTR27twZoaGhmDlzJt544w1Uq1ZN81fG+vXr4efnp9P19+7di7i4OERHRyMsLExTnpKSUuaYCw0fPhzz58/H2rVr0a9fP/z+++8YP348zM3NNee4uLigfv36Wi1cTypMOoqj6+vW5S8iR0dHmJmZlXgvPvn8pbn+kwp/7ul7qDy4uLjA2dkZ27dvL/a4ra0tAP3cF4Vvpk8PGn/Wh1ZZf5+BgrESa9euhRAC586dw8qVKzFr1ixYWlpq/f4+bdasWVrjxUri5+dXpskNxQkJCQEAnD9/Ht27d9c6dv78ec1xY2rbti3atm0LlUqFEydO4Ouvv8b48ePh7u6O/v37Y+3atbCwsMDWrVu1kvonB+nqU35+PpKTk7WSoYSEBADFJ3+FXFxcYGlpWeKEgyd/j3v16oVevXohJycHf//9NyIjIzFw4ED4+/trDTR+WseOHbF///7nvoahQ4c+c12pktja2qJOnTq4du2apqxwbND58+fRvHlzTXlCQgKSkpK07qF69erh/PnzRa5bWPb0/Vb43qRLCxYTIT2pXbs2vL298fPPP2PixImaN8WMjAxs2LBBM5PsWXr27Il58+bh/v37+M9//lPiea1atYK9vT2WLFmC/v37F/shJpFIIIQo8gb8/fffP3PgpLu7O4YNG4azZ89i4cKFmoy98DpZWVll7jaTy+WIiopC+/btMXv2bCxduhRdu3aFVCrFzZs3S909U6jwdT/9GpcuXVrsc+sSf1BQEJo3b44VK1ZoBv49PcOhZ8+e2LZtG6pXr14kyX2eF3ndz2NtbY3mzZtj48aN+OyzzzSvV61W46effoKPjw9q1apVpmsXNv3fvHlTb/GWVs+ePbF27VqoVCqtN8+n6eO+cHd3h0KhwLlz57TOL/yLtrTxlub3+enYGzRogC+++AIrV67EqVOnnnm+MbrGvL290axZM/z000+YOHGi5o+Dv//+G1evXsX48eP19lwvytzcHM2bN0edOnWwevVqnDp1SvOeKZVKtf6wycrKwo8//miwWFavXo1x48ZpHv/8888ACrrXStKzZ0/MnTsXzs7OCAgIKNXzyOVyzaKYO3bswOnTp5+ZCBmia+xJSUlJOH/+vFb330svvQSFQoGVK1dq/S4XzmJ9cgHV3r17Y+zYsTh69Kjm3Pz8fPz0009o3rx5kT84C2dklrZ7GWAipDdmZmb49NNP8cYbb6Bnz54YM2YMcnJyMH/+fKSkpGDevHnPvUbr1q3x5ptvYvjw4Thx4gTatWsHa2trxMfH4+DBg6hXr55m3Mfnn3+OUaNGoVOnThg9ejTc3d1x48YNnD17Ft988w3s7OzQrl07zJ8/Hy4uLvD398f+/fuxbNkyODg4aD1v8+bN0bNnT9SvXx+Ojo64fPkyfvzxR63krTCD/+STT9CtWzeYm5ujfv36JXZTlCQsLAzdu3fHihUrMHnyZAQEBGDWrFmYOnUqbt26hZdeegmOjo548OABjh07pmmtKk6rVq3g6OiIt956C9OnT4eFhQVWr16Ns2fPFjm3LPGPGDECY8aMQVxcHFq1alVkuvisWbOwa9cutGrVCuPGjUPt2rWRnZ2N27dvY9u2bViyZEmJ3Zf+/v5lft2lERkZic6dOyM8PBwTJ06ETCbDokWLcOHCBaxZs6bMXSU+Pj4IDAzE33//rfWmXh769++P1atXo3v37nj33XfRrFkzWFhY4N69e9i3bx969eqF3r176+2+GDRoEJYvX47q1aujQYMGOHbsmObDqzRK+/u8detWLFq0CK+++ioCAwMhhMDGjRuRkpKiGaNWEi8vr2e2POpq/fr1AP79MDlx4oRmnFnhDDugoL46d+6M119/HWPHjkViYiImT56MkJCQIn8w+Pv7A4DeWqSeZ8mSJdi7dy969OgBX19fZGdna1pUOnXqBADo0aMHFixYgIEDB+LNN99EcnIyPvvsM4OMpQIKusw+//xzpKeno2nTpjh8+DBmz56Nbt26oU2bNiX+3Pjx47Fhwwa0a9cO7733HurXrw+1Wo3Y2Fjs3LkT77//Ppo3b45p06bh3r176NixI3x8fJCSkoIvv/wSFhYWWq2ixXn6fa2sUlNT0blzZwwcOBA1a9aEpaUlrl27hi+//BI5OTmYPn265lwnJyd8+OGH+Oijj+Dk5IQuXbrg+PHjmDFjBkaNGqWVxIwYMQJRUVF4/fXXMW/ePLi5uWHRokW4evUqdu/eXSSOv//+G+bm5mjXrl3pgy/1sGoTV9oFFTdv3iyaN28uFAqFsLa2Fh07dhSHDh3SOqdwpsfDhw+Lvcby5ctF8+bNhbW1tbC0tBTVq1cXQ4YMESdOnNA6b9u2bSIsLExYW1sLKysrERwcLD755BPN8Xv37om+ffsKR0dHYWtrK1566SVx4cKFIjMTJk+eLJo0aSIcHR2FXC4XgYGB4r333tNanTMnJ0eMGjVKuLq6ColEIgCImJiYEuuhpAUVhRDi/PnzwszMTAwfPlyr3sLDw4WdnZ2Qy+XCz89PvPbaa2L37t1F6u1Jhw8fFi1bthRWVlbC1dVVjBo1Spw6darIjJdnxV/SImapqanC0tJSABDfffddsa/l4cOHYty4cSIgIEBYWFgIJycnERoaKqZOnSrS09NLrB9dXvez6hJAiSs8HzhwQHTo0EFzH7Vo0UJs2bJF65yyLBT60UcfCUdHxyIrCPv5+YkePXoUOT8sLEyEhYWV+vqFP/P0rDEhChZS++yzz0SDBg2EQqEQNjY2ok6dOmLMmDHi+vXrmvP0cV+kpqaKUaNGCXd3d2FtbS1efvllcfv27RJnjZX19/nKlStiwIABonr16sLS0lLY29uLZs2aiZUrV+pUZ/qAZ8waetrOnTtFixYthEKhEE5OTmLIkCHiwYMHRc5zcXERLVq0eO5zl/R/PnToUJ1mvR05ckT07t1b+Pn5CblcLpydnUVYWJjWTC4hCv5fateurXnPi4yMFMuWLSvy3lbSfV3c717hbMP58+drxW9tbS3OnTsn2rdvLywtLYWTk5P473//W+Q9orj3ovT0dPHhhx+K2rVrC5lMJuzt7UW9evXEe++9JxISEoQQQmzdulV069ZNeHt7C5lMJtzc3ET37t3FgQMHSl1vLyo7O1uMGjVKBAUFCRsbGyGVSoWPj48YNGhQiTNov/zyS1GrVi0hk8mEr6+vmD59usjNzS1yXkJCghgyZIhwcnISCoVCtGjRQuzatavYa7Zt21a8/PLLOsUuEUKHVYeIyOTFxcUhICAAP/zww3PXQSHTdunSJdStWxdbt25Fjx49jB2OUQwbNgzr169Henq6sUOp8m7evImaNWtix44dz21NfRIXVCQinXh5eWH8+PGYM2eOziv+kmnZt28fWrZsabJJEJWv2bNno2PHjjolQQDHCBFRGXz44YewsrLC/fv3Ua1atVL/3PNWujUzM9Npmi5VbBEREYiIiNDb9VQq1TOnekskEq0B0GQ68vPzUb16dUyZMkXnn2XXGBGVm+cN0i7rFF0yDf7+/s9cnDMsLEyzUCVRabFFiIjKzfHjx595vKxTdMk0bNmy5ZkbAheuJUWkC7YIERERkcliZzwRERGZLHaNPYdarUZcXBxsbW31ul8PERERGY4QAmlpafDy8nrmJAwmQs8RFxen06wYIiIiqjju3r37zE3KmQg9R+Hgu7t375a4+zcRERFVLEqlEtWqVXvuIHomQs9R2B1mZ2fHRIiIiKiSed6wFg6WJiIiIpPFRIiIiIhMFhMhIiIiMllMhIiIiMhkMREiIiIik8VEiIiIiEwWE6ESREVFITg4GE2bNjV2KERERGQg3HT1OZRKJezt7ZGamsp1hIiIiCqJ0n5+s0WIiIiITBYTISIiIjJZTISIiIjIZDERIiIiIpPFRIiIiIiM4szdFKRk5ho1BiZCREREVO7yVWpErD6FFpF7cOL2I6PFwUSIiIiIyt3OSw9wPyULVjIpQrztjRYHEyEiIiIqdysOxQAA3mjuC4WFudHiYCJERERE5ercvRQcv/0YFuYSDGrhZ9RYmAgRERFRuVpx6DYAoGd9L7jbKYwaCxMhIiIiKjcPlNnYei4OADC8tb9xgwETISIiIipHP/19B3kqgSZ+jqjv42DscJgIERERUfnIzlNh9dFYAMCINgFGjqYAEyEiIiIqF5tO38ejjFx4O1iiS7C7scMBwESIiIiIykFOvgrf7L0BoGBskNS8YqQgFSMKIiIiqtLWHI3F/ZQsuNvJjT5l/klMhIiIiMigMnPz8c2+mwCAdzrUNOoCik9jIkREREQGterwHSSl56CakyX+06SascPRwkSIiIiIDEaZnYcl+wtag8Z3rAWZtGKlHhUrGiIiIqoylNl5+GLXNaRm5aGGmw1ebeRt7JCKkBo7ACIiIqr8cvPVuJGYjsvxSlyMU+LY7WRcilNCLQqOv9+5FszNJMYNshgmkQj17t0b0dHR6NixI9avX2/scIiIiKqUD9afxabT95GnEkWO+TtboXcjH7wU4mGEyJ7PJBKhcePGYcSIEVi1apWxQyEiIqpSbj5Mxy8n7gEAbBVSBHnaIcjDFo39HNE8wBke9sbdVPV5TCIRCg8PR3R0tLHDICIiqnJ+P1OwgWq7Wq5YNbwpJJKK1/31LBV+sPRff/2Fl19+GV5eXpBIJNi8eXORcxYtWoSAgAAoFAqEhobiwIED5R8oERGRiRFCYMvZgkSodyOvSpcEAZUgEcrIyECDBg3wzTffFHt83bp1GD9+PKZOnYrTp0+jbdu26NatG2JjY8s5UiIiItNy4b4St5IyIJeaoXNwxRwD9DwVvmusW7du6NatW4nHFyxYgJEjR2LUqFEAgIULF2LHjh1YvHgxIiMjdX6+nJwc5OTkaB4rlUrdgyYiIjIBv5+9DwDoFOwOG3mFTymKVeFbhJ4lNzcXJ0+eRJcuXbTKu3TpgsOHD5fpmpGRkbC3t9d8VatWsVbAJCIiqgjUaoEtZ+MBAK808DJyNGVXqROhpKQkqFQquLu7a5W7u7sjISFB87hr1654/fXXsW3bNvj4+OD48eMlXnPKlClITU3VfN29e9dg8RMREVVWx24/QoIyG7YKKdrXdjV2OGVWOduxnvL04CwhhFbZjh07Sn0tuVwOuVyut9iIiIiqot/+mS3WLcQDcmnF2URVV5W6RcjFxQXm5uZarT8AkJiYWKSViIiIiPQjN1+NbecLusV6Nax422boolInQjKZDKGhodi1a5dW+a5du9CqVasXunZUVBSCg4PRtGnTF7oOERFRVXPg+kOkZuXB1VaOFoHOxg7nhVT4rrH09HTcuHFD8zgmJgZnzpyBk5MTfH19MWHCBAwePBhNmjRBy5Yt8e233yI2NhZvvfXWCz1vREQEIiIioFQqYW9v/6Ivg4iIqMrYcbGgJ6Z7iEeF3D9MFxU+ETpx4gTCw8M1jydMmAAAGDp0KFauXIl+/fohOTkZs2bNQnx8PEJCQrBt2zb4+fkZK2QiIqIqS60W2HvlIQBU2rWDniQRQhTdIY00CluEUlNTYWdnZ+xwiIiIjOrs3RT0ijoEG7kUpz7qDJm0Yo6yKe3nd8WMnoiIiCqkPZcfAADa1XKpsEmQLir/KzAQDpYmIiIqas+VRABAhzpVY3Y2E6ESRERE4NKlS89cfJGIiMiUxKdm4WKcEhIJEF6JF1F8EhMhIiIiKpW9/7QGNarmAGebqrH4MBMhIiIiKpU9lwsSoY5BVaNbDGAiRERERKWQlavCoRtJAICOQW5GjkZ/mAgRERHRcx2+mYScfDW8HSxR293W2OHoDROhEnDWGBERUQEhBHZeLJg23zHIrchm55UZF1R8Di6oSEREpkiZnYdF+27izN3HuByfhtSsPADAqhHNEFar4s8YK+3nd4XfYoOIiIjK349H7mDJ/puax+ZmErSq7oyWlXyT1acxESIiIqIiDlwv2E9scAs/9GtaDTXdbSCXmhs5Kv1jIkRERERasnJVOHUnBQAwvLU/Al1tjBuQAXGwNBEREWk5eecxclVqeNorEOBibexwDIqJUAk4a4yIiEzVoZsF6wW1qu5SpWaIFYeJUAm41xgREZmqwzcKE6GqNTC6OEyEiIiISCM1Kw/n76cCAFrXcDFyNIbHRIiIiIg0jt5KhloAga7W8LBXGDscg2MiRERERBqHbyYDAFpXr/qtQQATISIiInrCIRMaHwQwESIiIqJ/JKZl43piOiQSoCUToeJFR0cbIAwiIiIytiP/dIvV9bKDg5XMyNGUD50ToZdeegnVq1fH7NmzcffuXUPEVCFwHSEiIjI1hd1ipjI+CChDIhQXF4d3330XGzduREBAALp27YpffvkFubm5hojPaLiOEBERmRK1WuDg9YJEyFS6xYAyJEJOTk4YN24cTp06hRMnTqB27dqIiIiAp6cnxo0bh7NnzxoiTiIiIjKgvVcSEZeaDTuFFM0DmAiVSsOGDTF58mREREQgIyMDy5cvR2hoKNq2bYuLFy/qK0YiIiIysOWHYgAAA5r7wlJW9XaZL0mZEqG8vDysX78e3bt3h5+fH3bs2IFvvvkGDx48QExMDKpVq4bXX39d37ESERGRAVyOV+LwzWSYm0kwpKW/scMpV1Jdf+Cdd97BmjVrAACDBg3Cp59+ipCQEM1xa2trzJs3D/7+/noLkoiIiAxnxT+tQS+FeMDbwdLI0ZQvnROhS5cu4euvv0bfvn0hkxU/tc7Lywv79u174eCIiIjIsJLSc7D5TBwAYETrACNHU/50ToT27Nnz/ItKpQgLCytTQERERFR+fj4ai9x8NRpUc0BjXwdjh1PuuLI0ERGRicrJV+HHv+8AAEa09odEIjFyROWPiVAJuKAiERFVddsvJOBhWg7c7eToXs/T2OEYBROhEnBBRSIiquqirz4EAPRp7AMLc9NMCUzzVRMREZk4IQQO/rOlRtuaprOlxtOYCBEREZmg64npeJiWA4WFGRr7Oho7HKMp1awxR0fHUg+gevTo0QsFRERERIZXuK9YU38nKCxMZyXpp5UqEVq4cKHm++TkZMyePRtdu3ZFy5YtAQBHjhzBjh078NFHHxkkSCIiItKvwp3m29Qw3W4xAJAIIYQuP9C3b1+Eh4fj7bff1ir/5ptvsHv3bmzevFmf8RmdUqmEvb09UlNTYWdnZ+xwiIiIXlieSo2GM3ciI1eFre+0QYi3vbFD0rvSfn7rPEZox44deOmll4qUd+3aFbt379b1ckRERFTOzt1LQUauCo5WFgj2NO0/8nVOhJydnbFp06Yi5Zs3b4azs7NegiIiIiLDOXg9GQDQqoYLzMxMbxHFJ+m8xcbMmTMxcuRIREdHa8YI/f3339i+fTu+//57vQdIRERE+sXxQf/SOREaNmwYgoKC8NVXX2Hjxo0QQiA4OBiHDh1C8+bNDREjERER6UlGTj5OxT4GwEQIKEMiBADNmzfH6tWr9R0LERERGdixmEfIVwv4OlmhmpOVscMxujIlQmq1Gjdu3EBiYiLUarXWsXbt2uklMGOLiopCVFQUVCqVsUMhIiLSm8LVpFuzNQhAGRKhv//+GwMHDsSdO3fw9Mx7iURSZRKHiIgIREREaKbfERERVXb5KjX2XUkEwG6xQjonQm+99RaaNGmCP/74A56enqVecZqIiIiM66s913ErKQO2cikToX/onAhdv34d69evR40aNQwRDxERERnAkZvJ+HrfDQDAnD71YG9lYeSIKgad1xFq3rw5bty4YYhYiIiIyAAeZeRi/LrTEAL4TxMfvNLAy9ghVRg6twi98847eP/995GQkIB69erBwkI7o6xfv77egiMiIqIXI4TAB+vP4oEyB9VdrTHjlbrGDqlC0XmvMTOzoo1IEokEQogqNVi6EPcaIyKiykoIgc92XkXUvpuQSc2weWxrBHuZxmdZaT+/dW4RiomJeaHAiIiIqHx8uec6ovbdBADMfKWuySRButA5EfLz8zNEHERERKRHX++5joW7rwMAPuwRhAHNfI0cUcVUpgUVAeDSpUuIjY1Fbm6uVvkrr7zywkERERGR7rLzVDh15zH+OB+P1UdjAQBTutXBqLaBRo6s4tI5Ebp16xZ69+6N8+fPa8YGAdCsJ1TVxggRERFVVGq1wKV4JQ7eSMKhG0k4FvMIOfn/7vjwf11rY0xYdSNGWPHpnAi9++67CAgIwO7duxEYGIhjx44hOTkZ77//Pj777DNDxEhERGTyElKzMWvrRTzOyEO+Wo1clUBscgYeZ+ZpnedmK0ebGi54KcQDXep6GCnaykPnROjIkSPYu3cvXF1dYWZmBjMzM7Rp0waRkZEYN24cTp8+bYg4iYiITNrHWy9h2/mEIuXWMnO0rO6M1jVc0KaGC2q42XDXBx3onAipVCrY2NgAAFxcXBAXF4fatWvDz88PV69e1XuAREREpu7E7Uf443w8zCTA7FfrwdHKAlJzMzhZy1Dfxx4W5jqvj0z/0DkRCgkJwblz5xAYGIjmzZvj008/hUwmw7fffovAQA7GIiIi0ie1WuDjPy4DAPo1rYaBzTn7S590ToQ+/PBDZGRkAABmz56Nnj17om3btnB2dsa6dev0HiAREZEp+/1sHM7eTYG1zBzvda5l7HCqHJ0Toa5du2q+DwwMxKVLl/Do0SM4OjqyT5KIiEiPsnJV+GT7FQDA2PAacLNVGDmiqqfM6wg9ycnJSR+XqVCioqIQFRXF5QCIiKjc5avUuPogDauPxiI+NRte9gqMbBNg7LCqJJ33GjM13GuMiIjKy91HmfjfpvM4fvsRsvP+XQ/oy/4N0auhtxEjq3wMttcYERER6V92ngpjfjyJS/FKAICtQooGPg7oWtcdrzTwMnJ0VRcTISIiogpg5pZLuBSvhLO1DD+MbIYgDzuYmXHsraHpvPDAX3/9hfz8/CLl+fn5+Ouvv/QSFBERkSnZdPoe1hyLhUQCfNm/Eep62TMJKic6J0Lh4eF49OhRkfLU1FSEh4frJSgiIiJTcf1BGv638QIA4N2ONdGmpouRIzItOidCQohip8knJyfD2tpaL0ERERGZggPXH2LwsmPIylOhTQ0XvNOhprFDMjmlHiPUp08fAAW7zA8bNgxyuVxzTKVS4dy5c2jVqpX+IyQiIqpiMnPzMe/PK/jhyB0AQKCrNRb2bwhzdoeVu1InQvb29gAKWoRsbW1haWmpOSaTydCiRQuMHj1a/xESERFVEUII7L6ciLnbLiMmqWCXhiEt/TC5Wx1YyTh/yRhKXesrVqwAAPj7+2PixInsBiMiItLB4RtJ+HTHVZy5mwIA8LBT4NPX6qNdLVfjBmbiuKDic3BBRSIiehYhBJRZ+XiYno1EZQ4epufgYVoOEtNyEJ+ajfiULMSlZCEuNRsAoLAww7BWAfhv++qwt7QwcvRVl0EXVFy/fj1++eUXxMbGIjc3V+vYqVOnynJJIiKiCi87T4Ubiek4FfsYJ24/xtl7KYhPyUauSv3cn7Uwl2BgM19EhNeAmx33DKsodE6EvvrqK0ydOhVDhw7Fb7/9huHDh+PmzZs4fvw4IiIiDBEjERFRucnOU+HagzRcilPicrwSVxLS8ECZjaT0XKTnFF1Hr5CdQgpXWzlcbeVws1XAxUYOT3sFPB0U8HKwRKCLNRysZOX4Sqg0dE6EFi1ahG+//RYDBgzAqlWr8MEHHyAwMBDTpk0rdn0hIiKiikQIgbP3UnHuXgpy89XIVwvk5Klx82E6LsUrcethOtTPGDRiq5Cika8jQn0dEernCD9nK7jayqGwMC+/F0F6o3MiFBsbq5kmb2lpibS0NADA4MGD0aJFC3zzzTf6jZCIiEgPlNl5+O30faw+GosrCWnPPNfZWoYgTzsEedoiyNMOPo5WcLGRwcVWDlu5tNj19Khy0jkR8vDwQHJyMvz8/ODn54e///4bDRo0QExMDDjumoiIyoMQAvdTsnAlPg13H2ciLTsf6Tn5SMvOQ1p2/j9fef+U5WuOF5JJzdCmhgts5FJIzSWQmZvB19kKQZ52CPa0g5utnMmOidA5EerQoQO2bNmCxo0bY+TIkXjvvfewfv16nDhxQrPoIhERkb6lZObij/Px2HY+HufvpUKZXfJ4nZLUcLPBwGa+6NPYm+N1CEAZps+r1Wqo1WpIpQU51C+//IKDBw+iRo0aeOuttyCTVa0bi9PniYiM61TsYyyJvol9VxORp/r3I0tqJkENNxsEulrD3tICtgoL2MilsFVINd/bKaSw+eexrUIKZ2sZW3pMRGk/v7mO0HMwESIiMp7tFxIwbs1pzfT0YE87vNrIC21quKKGmw1kUp23zCQTYdB1hFJSUnDs2DEkJiZCrdZeO2HIkCFluSQREZGWX47fxeSN56AWQKcgd/xf19qo7WFr7LCoitE5EdqyZQveeOMNZGRkwNbWVquJUSKRMBEiIqIyy81XIy4lC1vPxeGzndcAAP2aVMPcPvW4ISkZhM6J0Pvvv48RI0Zg7ty5sLKyMkRMRERUBa0/eQ+/nbmPPJUaKrVAvlpApRbIUwmo1GqkZ+cjQZmttYbPmLBATH6pDsf1kMHonAjdv38f48aNq1RJ0NatW/H+++9DrVZj0qRJGDVqlLFDIiIyGWq1wCc7rmDp/lulOl9hYQYfRysMbuGHoa38DRscmTydE6GuXbvixIkTCAwMNEQ8epefn48JEyZg3759sLOzQ+PGjdGnTx84OTkZOzQioiovN1+ND9afxeYzcQCAMe0CEeJtD6mZBOZmEliYm8HcTAKpmQQKmTl8HC3hasM1fKj8lCoR+v333zXf9+jRA//3f/+HS5cuoV69erCw0N4595VXXtFvhC/o2LFjqFu3Lry9vQEA3bt3x44dOzBgwAAjR0ZEVPll5ORj/7WHSE7PQXaeGjn5Kq1/z99PxZm7KZCaSRDZpx5eb1LN2CETaSlVIvTqq68WKZs1a1aRMolEApVK9cJBPemvv/7C/PnzcfLkScTHx2PTpk1F4lm0aBHmz5+P+Ph41K1bFwsXLkTbtm0BAHFxcZokCAB8fHxw//59vcZIRGRKsvNUOHXnMdafuoftFxKQmfvs931rmTkWDQpFWC3XcoqQqPRKlQg9PUW+PGVkZKBBgwYYPnw4+vbtW+T4unXrMH78eCxatAitW7fG0qVL0a1bN1y6dAm+vr7FbvvBJlciotK7mpCG+Tuu4ObDDCSl5SDtqR3Y/Z2tEOxlB7nUHAoLM8il5pD/86+VzBxdgt0R6GpjpOiJnq1M6wiVRr169bBt2zZUq/ZizaDdunVDt27dSjy+YMECjBw5UjMAeuHChdixYwcWL16MyMhIeHt7a7UA3bt3D82bNy/xejk5OcjJydE8ViqVLxQ/EVFlpVILfHfgFhbsvKZZ0LCQnUKKHvW98FqoNxr7OvIPTKq0DJYI3b59G3l5eYa6PAAgNzcXJ0+exOTJk7XKu3TpgsOHDwMAmjVrhgsXLuD+/fuws7PDtm3bMG3atBKvGRkZiZkzZxo0biKiikgIgUcZuYhLyca9x5n47sAtnIpNAQB0qOOG0W0D4WYnh4uNHHYK7sBOVYPBEqHykJSUBJVKBXd3d61yd3d3JCQkAACkUik+//xzhIeHQ61W44MPPoCzs3OJ15wyZQomTJigeaxUKl+4VYuIqKIRQmDvlUTsuvQA91OycD8lC3EpWcjO0275sZFLMe3lYLwe6sPEh6qkSp0IFXr6l1MIoVX2yiuvlHo2m1wuh1wu12t8REQVyaEbSZi/4yrO3E0p9ribrRxeDpao5W6DcR1rwsex8qwbR6SrSp0Iubi4wNzcXNP6UygxMbFIKxERkSlJz8nHubspOH03BTFJGcjKVSErT4XEtGxcuF8w9lFhYYb+TX1R18sO3g6W8Ha0hIe9AnKpuZGjJyo/lToRkslkCA0Nxa5du9C7d29N+a5du9CrVy8jRkZEVDaZufnIylUhJ1+N7LyCf5/8XlOWp0L2P//m5KuRlp2PRGU2EpTZSEjNRkxyBoqZNAsAkJmbYWBzX4wNrw43W0X5vkCiCqbCJ0Lp6em4ceOG5nFMTAzOnDkDJycn+Pr6YsKECRg8eDCaNGmCli1b4ttvv0VsbCzeeuutF3reqKgoREVF6X1dJCKiJwkhcCleie0XErD9QgKuJ6br7dreDpZoWM0BQZ62sJFLYSWTQiEzR1N/R3jaW+rteYgqM4kobqGdZ4iJiUFAQMBzz/v555/Rq1cvWFtblzk4AIiOjkZ4eHiR8qFDh2LlypUAChZU/PTTTxEfH4+QkBB88cUXaNeu3Qs9byGlUgl7e3ukpqbCzs5OL9ckItORlp2HNcdicSc5Exk5+UjPUSE9Jw8ZOSqk5+QjNSsPjzJyi/yczNwMcqkZ5BbmkEvNtNbnUWjW6TGD4p/jVjIp3O0UcLeTw8NOgRpuNnCzY2sPma7Sfn7rnAiZm5ujXbt2GDlyJF577TUoFFX7F42JEBGVRZ5KjTXHYvHl7utILibReZLCwgzta7nhpRAPhNVyhb2lBczMOEOL6EUYLBG6cOECli9fjtWrVyMnJwf9+vXDyJEj0axZsxcOuiJiIkREpSWEwK2kDOy/+hA//X0Ht5IyAACBLtboWd8TtgoL2CiksJZLYSuXFnwvkyLAxRqWMg5QJtIngyVChfLz87FlyxasXLkSf/75J2rWrImRI0di8ODBcHWtOvvJMBEiIqBgFtbx24+QkJoNZVYelNkF3Vu5KjXy8tXIzlfjdOxj3HucpfkZZ2sZxneqif7NfGFhbmbE6IlMj8EToUI5OTlYtGgRpkyZgtzcXFhYWKBfv3745JNP4Onp+SKXNqonB0tfu3aNiRCRichXqZGYloP41ILZV9cepOHQjSScuZuCfPXz3y5l5mZoFuCE9rVd0a9pNdgqLMohaiJ6msEToRMnTmD58uVYu3YtrK2tMXToUIwcORJxcXGYNm0a0tLScOzYsTK/gIqCLUJEVZMQAlcS0nDg+kOcuZuCuJSCxCcxLRsl5Tt+zlao6WYDO0sL2CksYCOXQiY1g4W5GSzMJajuaoPmgU6wklX4CblEVV5pP791/m1dsGABVqxYgatXr6J79+744Ycf0L17d5iZFTT7BgQEYOnSpahTp07ZoyciegFCCDxQ5iApPQcZOfnIyM2HMisfD5TZeKDMQXxqFk7ceYyHaTnF/ryFuQTudgp42ivg42iF5gFOaF3DBdWcuMIyUVWjcyK0ePFijBgxAsOHD4eHh0ex5/j6+mLZsmUvHBwRUWndSc7A6qOxuBiXiktxSjzOfP6mz5YW5mgR6ISW1Z3h62QNLwcFPOwVcLGWc9YWkYl44TFCVR27xogqvgPXH2Ls6lNIy87XlJmbSeBiI4O17J9ZWgop3GzlcLdTwM1OgSAPW4T6O3I7CaIqymBdY6aCK0sTVQ4/HLmNmVsuQaUWaOTrgH5NqqGulz1quttAYcEkh4iejS1Cz8EWIaKKKS4lC1/tuY61x+8CAPo09kZkn3ps4SEiAGwRIqIqRq0WUGbn4czdFKw+Gos9lx9ALQCJBPigax28FRYIiYTjeohIN0yEiKhCSE7PwY3EdFxPTMed5AwkpeciKT0HSem5SE7PwaOM3CLr+LQIdEJEeA20rVl1FnElovLFRIiIDCo1Kw9Hbibh4I0k3H+cBWV2wUaj6dn5yFOpkatSIzdfjZx8damu52IjQ8/6XhjUwhc13GwNHD0RVXU6J0IPHjzAxIkTsWfPHiQmJuLpIUYcXExkWjJz83H9QTquPkjD9QdpeJSRh5x8FXLzC1ZoPncvpcQFCp8kkQA+jpao4WqDQFcbuNnK4Wwjh4uNDC42cjjbyOBkLeMYICLSK50ToWHDhiE2NhYfffQRPD09q2yfPGeNEf0rT6XG3UeZuJ2cgVsPM3A7OQMxSRm4nZSJuNQsPG/KRXVXa7St6YpgT7uCVZktpbCVW/yzKrMEFuZmcLGRc+NRIip3Os8as7W1xYEDB9CwYUMDhVSxcNYYVXXpOfk4cO0hjt9+jOSMHDzOzMPjjFyk5+QjMzcfmbkqZOTkP7NVx8VGjtoeNqjlbgt3OwXkUjPIpGawkUvRxN8J3g6W5feCiIhgwFlj1apVK9IdRkSVi1otsPH0ffx+Ng5/30xGrur543MsLczh72KNQBdr+LtYIcDFBgEu1ghwsYaTtawcoiYi0j+dE6GFCxdi8uTJWLp0Kfz9/Q0QEhEZ0o3EdEzecA4n7jzWlPk5WyG8thu8HSzhaC2Do5UF7CwtYGlhDiuZOWwUUrjayKtsVzgRmS6dE6F+/fohMzMT1atXh5WVFSwsLLSOP3r0SG/BEZH+5KnUWBJ9E1/vvYFclRrWMnOMCauO7vU8UN3VhkkOEZmkMrUIEVHlIoTAxF/P4rczcQCA9rVdMad3PY7dISKTp3MiNHToUEPEQUQGtOLQbfx2Jg5SMwnmv14frzb0ZgsQEREAs7L80M2bN/Hhhx9iwIABSExMBABs374dFy9e1GtwRPTijsU8wtxtlwEAU3sEoXcjHyZBRET/0DkR2r9/P+rVq4ejR49i48aNSE9PBwCcO3cO06dP13uAxhIVFYXg4GA0bdrU2KEQldkDZTbGrj6FfLVAr4ZeGNbK39ghERFVKDqvI9SyZUu8/vrrmDBhAmxtbXH27FkEBgbi+PHjePXVV3H//n1DxWoUXEeIKpP0nHxsOxePO48yEJeSjZN3HiP2USbqeNhi49hWsJJxVx0iMg0GW0fo/Pnz+Pnnn4uUu7q6Ijk5WdfLEZGeJKRmY+jyY7j6IE2r3E4hxZJBoUyCiIiKofM7o4ODA+Lj4xEQEKBVfvr0aXh7e+stMCIqvRuJaRiy7BjiUrPhaivHS3U94O1oCS8HS7QIdIKbrcLYIRIRVUg6J0IDBw7EpEmT8Ouvv0IikUCtVuPQoUOYOHEihgwZYogYiegZjsU8wugfTiA1Kw+BrtZYNbwZqjlZGTssIqJKQecxQnl5eRg2bBjWrl0LIQSkUilUKhUGDhyIlStXwty8am2ayDFCVBEIIaAWBf8KAEnpOfjjXDx+OxOH8/dTAQANqzlg+bCm3O6CiAil//zWOREqdPPmTZw+fRpqtRqNGjVCzZo1yxxsRcZEiIwpT6XGhpP3sCj6JmIfZRZ7jrmZBC/X98TcPvU4DoiI6B8GGyxdqFq1asjPz0f16tUhlfLNl0if8lVqrD95D9/su4F7j7OKPaeJnyN6NfRCt3qecLGRl3OERERVg84ZTGZmJt555x2sWrUKAHDt2jUEBgZi3Lhx8PLywuTJk/UeJJEpycpV4b+rTyL66kMAgIuNHP9tXx29GnpBaiaBBBJIzSWwlvMPECKiF6XzgopTpkzB2bNnER0dDYXi35konTp1wrp16/QanDFxQUUyhtTMPAxadhTRVx9CYWGGD3sE4cAH4RjZJgAuNnI4WMlgb2XBJIiISE90HiPk5+eHdevWoUWLFloLKt64cQONGzeGUqk0VKxGwTFCVF4eKLMxZFnBOkB2CimWD2uKJv5Oxg6LiKhSMtgYoYcPH8LNza1IeUZGBvcvIiqjKwlKjFp1AvceZ8HNVo4fRjZDHQ8m3kREhqZz11jTpk3xxx9/aB4XJj/fffcdWrZsqb/IiEzEjosJ6LPoMO49zoK/sxU2/LcVkyAionKic4tQZGQkXnrpJVy6dAn5+fn48ssvcfHiRRw5cgT79+83RIxEVZIQAl/vvYEFu64BAFrXcEbUwMZwsOI6QERE5UXnFqFWrVrh0KFDyMzMRPXq1bFz5064u7vjyJEjCA0NNUSMRFVObHImBi07qkmChrXyx8rhzZgEERGVszIvqGgqOFia9ClfpcaKQ7fx+a6ryM5TQy41w8xX6qJ/M19jh0ZEVKUYbLD0G2+8gfbt26N9+/ZVdjVporISQuDIrWScuZuCB6nZSFBm42FaDjJyVMjMy4cyKx+pWXkAgJaBzojsUw/+LtZGjpqIyHTpnAjZ2Njg888/x5gxY+Dh4YGwsDCEhYWhffv2qFOnjiFiJKoUTsU+xid/XsHRmEfPPM9OIcWHPYLxehMfzrQkIjKyMneNJSQkIDo6GtHR0di/fz+uXbsGNzc3xMfH6ztGo2LXGD3PA2U2Ptp8ATsvPQAAyKRm6FrXA9UcLeFhr4CbrRw2cgtYysxhJTOHr5MVF0QkIjIwg+81ZmtrC0dHRzg6OsLBwQFSqRQeHh5lvRxRpZSSmYs3vj+KG4npMJMAr4X64N1OteDtYGns0IiIqBR0ToQmTZqE/fv34+zZswgJCUG7du0wZcoUtGvXDg4ODgYIkahiyspVYeSqE7iRmA4POwVWjWiG2h62xg6LiIh0oHMiNH/+fLi6umL69Ono1asXgoKCDBGX0UVFRSEqKgoqlcrYoVAFlK9S4501p3DyzmPYKaT4YWQz1HJnEkREVNnoPEbo7Nmz2L9/P6Kjo3HgwAGYm5trBku3b9++yiVGHCNET0tOz8HMLZfw+9k4yKVm+GlUczTlnmBERBVKaT+/X3gdobNnz2LhwoX46aefoFarq1wLChMhKpScnoNvD9zCD4fvICtPBTMJsHhQKLrW5dg4IqKKxqCDpU+fPq2ZMXbgwAEolUo0bNgQ4eHhZQ6YqCLJzVdj2m8XcPBGElRqgXy1QGpmHnJVagBAfR97fNC1DtrUdDFypERE9CJ0ToQcHR2Rnp6OBg0aoH379hg9ejTatWvH1hKqUmZuuYi1x+8WKa/vY4/xnWoivLYb1wAiIqoCdE6EfvzxRyY+VKX9+PcdrD4aC4kE+KRvfQR72sHcTAJLC3P4OVsxASIiqkJ03nR148aNxX4QZGRkYMSIEXoJishYjtxMxszfLwIAPuhaB/9pUg0h3vYI8rSDv4s1kyAioipG50Ro1apVyMrKKlKelZWFH374QS9BEZW3PJUaey4/wNjVJ5GvFujV0AtvhQUaOywiIjKwUneNKZVKCCEghEBaWhoUCoXmmEqlwrZt2+Dm5maQIIkMQQiBs/dSsfHUPWw9F49HGbkAgHre9vikb322/hARmYBSJ0IODg6QSCSQSCSoVatWkeMSiQQzZ87Ua3BEhpCdp8KWs3H44cgdnL+fqil3sZHj5QaeeDu8BhQW5kaMkIiIykupE6F9+/ZBCIEOHTpgw4YNcHL6dwE5mUwGPz8/eHl5GSRIIn24+ygTP/19B+tO3EVKZh6Agg1Su4d4oHdjH7Su7gypuc69xUREVImVOhEKCwsDAMTExMDX15fdBlQpCCFw+GYylh+Mwd6riShcPtTbwRKDWvihX9NqcLKWGTdIIiIyGp2nz/v5+eHAgQNYunQpbt26hV9//RXe3t748ccfERAQgDZt2hgiTiKdnbj9CPN3XMXRmEeasrY1XTCkpT861HGDuRmTeSIiU6dzIrRhwwYMHjwYb7zxBk6dOoWcnBwAQFpaGubOnYtt27bpPUii50lKz8GV+DTEp2YhITUbJ+48xv5rDwEAMnMzDGhWDUNb+SPQ1cbIkRIRUUWicyI0e/ZsLFmyBEOGDMHatWs15a1atcKsWbP0GhzRs6jVAgduJGHN0VjsvvwA+WrtbfPMzST4TxMfvNOhJrwcLI0UJRERVWQ6J0JXr15Fu3btipTb2dkhJSVFHzERFSGEwP5rD3Es5hES03KQmJaDGw/SEJearTkn0NUaPo5W8LCTw8vBEr0aeiPAxdqIURMRUUWncyLk6emJGzduwN/fX6v84MGDCAzkAnSkfzcfpmPG7xdx4HpSkWO2Cin6NvZB/2bVUMeD274QEZFudE6ExowZg3fffRfLly+HRCJBXFwcjhw5gokTJ2LatGmGiNEooqKiEBUVBZVKZexQTFJGTj6uJKRh+4V4rDx8G3kqAZm5GV5t5AU/Z2u42srhYadAU38nWMq45g8REZWNRAghnn+atqlTp+KLL75AdnZBt4RcLsfEiRPx8ccf6z1AY1MqlbC3t0dqaio3mjWA+ylZ+HT7FSiz8pCvFshXCTxQZiMmOQNP3pkd67jho57B8GdXFxERlUJpP7/LlAgBQGZmJi5dugS1Wo3g4GDY2FTN2ThMhAxr/NrT2HwmrthjbrZy1PWyw+CWfuhQx72cIyMiosqstJ/fOneNFbKysoK7uzskEkmVTYLIsOJSsrD1XDwA4MMeQXC2kcHczAyOVhYI8rSDi43cyBESEVFVp3MilJ+fj5kzZ+Krr75Ceno6AMDGxgbvvPMOpk+fDgsLC70HSVXTysO3ka8WaBHohFFtOdCeiIjKn86J0Ntvv41Nmzbh008/RcuWLQEAR44cwYwZM5CUlIQlS5boPUiqetKy87DmaCwAYDSTICIiMhKdE6E1a9Zg7dq16Natm6asfv368PX1Rf/+/ZkIUamsO34XaTn5CHS1RnhtN2OHQ0REJkrnrbYVCkWRNYQAwN/fHzIZN6+k58tXqbHi0G0ABa1BZtzzi4iIjETnRCgiIgIff/yxZo8xAMjJycGcOXPw9ttv6zU4qpq2XUjA/ZQsOFvL0LuRt7HDISIiE1aqrrE+ffpoPd69ezd8fHzQoEEDAMDZs2eRm5uLjh076j9CqlJUaoFv/7oJABjc0g8KCy6GSERExlOqRMje3l7rcd++fbUeV6tWTX8RUZW24lAMLtxXwkpmjsEt/IwdDhERmbhSJUIrVqwwdBxkAm4kpuHTHVcBAB/2CIYz1wkiIiIj03mMEFFZ5KvUeP+Xs8jNV6NdLVcMaMZWRCIiMj4mQlQuluy/ibP3UmGnkOLTvvUhkXCmGBERGR8TITK4PZcf4Ms91wEAM3vVhYe9wsgRERERFSjzXmNEz5OYlo2ZWy7hj3/2E3uprgdebcjp8kREVHHo1CKUl5eH8PBwXLt2zVDxUCWXr1LjzN0UfL3nOjp+vh9/nIuHuZkEo9sG4It+DdklRkREFYpOLUIWFha4cOECP8yoiPspWZi15SIO30hGWk6+pry+jz3m9q6HEG/7Z/w0ERGRcejcNTZkyBAsW7YM8+bNM0Q8VAnl5Ksw5scTuHBfCQCwU0jRPNAZnYPc0TfUB+bcQoOIiCoonROh3NxcfP/999i1axeaNGkCa2trreMLFizQW3BUOcz94zIu3FfC0coCy4Y1RQMfByY/RERUKeicCF24cAGNGzcGgCJjhdhlZnq2nY/HqiN3AAAL+jVEY19HI0dERERUejonQvv27TNEHFQJ3UnOwKT15wAAb4VVR3htNyNHREREpJsXmj5/7949SCQSeHtzSrSpiE/NwoHrSTh8Iwl/XU9CWk4+mvg54v0utYwdGhERkc50XlBRrVZj1qxZsLe3h5+fH3x9feHg4ICPP/4YarXaEDG+sN69e8PR0RGvvfaasUOp1P48H4+2n+zDB+vPYfOZODzKyIW3gyW+GtAIFuZcm5OIiCofnVuEpk6dqpk11rp1awghcOjQIcyYMQPZ2dmYM2eOIeJ8IePGjcOIESOwatUqY4dSaT3OyMXUzReQrxao62WH9rVd0bq6Cxr7OUJhYW7s8IiIiMpE50Ro1apV+P777/HKK69oyho0aABvb2+MHTu2QiZC4eHhiI6ONnYYldrcbZfxKCMXtdxtsGlsa8ikbAEiIqLKT+dPs0ePHqFOnTpFyuvUqYNHjx7pHMBff/2Fl19+GV5eXpBIJNi8eXORcxYtWoSAgAAoFAqEhobiwIEDOj8Pld2Rm8n49eQ9AEBkn3pMgoiIqMrQ+ROtQYMG+Oabb4qUf/PNN2jQoIHOAWRkZJR4TQBYt24dxo8fj6lTp+L06dNo27YtunXrhtjYWM05oaGhCAkJKfIVFxenczykLTtPhambzgMA3mjui1A/JyNHREREpD86d419+umn6NGjB3bv3o2WLVtCIpHg8OHDuHv3LrZt26ZzAN26dUO3bt1KPL5gwQKMHDkSo0aNAgAsXLgQO3bswOLFixEZGQkAOHnypM7PW5KcnBzk5ORoHiuVSr1duzJaFH0Tt5Iy4GorxwcvFW0JJCIiqsx0bhEKCwvDtWvX0Lt3b6SkpODRo0fo06cPrl69irZt2+o1uNzcXJw8eRJdunTRKu/SpQsOHz6s1+cqFBkZCXt7e81XtWrVDPI8lcGNxDQsjr4BAJjxcl3YW1oYOSIiIiL90qlFKC8vD126dMHSpUvLZVB0UlISVCoV3N3dtcrd3d2RkJBQ6ut07doVp06dQkZGBnx8fLBp0yY0bdq02HOnTJmCCRMmaB4rlUqTTIbUaoH/bbyAPJVAhzpu6F7Pw9ghERER6V2l2H3+6ecTQugUw44dO0p9rlwuh1wuL/X5VdUvJ+7i2O1HsLQwx6xedbl9ChERVUk6d40V7j5fHlxcXGBubl6k9ScxMbFIKxHpz8O0HMzddhkA8H6XWvBxtDJyRERERIZRoXefl8lkCA0Nxa5du9C7d29N+a5du9CrVy+9PU9xoqKiEBUVBZVKZdDnqYg+3noJyux8hHjbYVgrf2OHQ0REZDBG330+PT0dN27c0DyOiYnBmTNn4OTkBF9fX0yYMAGDBw9GkyZN0LJlS3z77beIjY3FW2+9pfNz6SIiIgIRERFQKpWwt7c36HMZW1xKFo7cTMa5eyk4ey8VZ+6mwEwCRPauDym3ziAioipMp0RIpVJhxowZqFevHpyc9LOezIkTJxAeHq55XDhQeejQoVi5ciX69euH5ORkzJo1C/Hx8QgJCcG2bdvg5+enl+c3ValZefjzfDw2nb6PozFFF8J8p0NN1POp2gkgERGRRAghdPkBhUKBy5cvIyAgwFAxVSiFLUKpqamws7Mzdjh6cfRWMoavPI7M3H+7/Rr5OqBRNUc0qGaPhtUc4Ods/YwrEBERVWyl/fzWuWusXr16uHXrlskkQlVNvkqNDzdfQGauCtVdrfFaaDW80tAL3g6Wxg6NiIio3OmcCM2ZMwcTJ07Exx9/jNDQ0CKDpatKq0lVHSy99vhdXE9Mh6OVBTaObc1FEomIyKTp3DVmZvbv4NknB0cXru1T1RKHqtQ1pszOQ/j8aCRn5GLmK3UxlDPCiIioijJY19i+ffteKDAynkX7biI5IxeBrtYY2NzX2OEQEREZnc6JUFhYmCHiIAO7+ygTyw/GAACmdg+CBafFExER6b6yNAAcOHAAgwYNQqtWrXD//n0AwI8//oiDBw/qNTjSn3l/XkGuSo3WNZzRoY6bscMhIiKqEHROhDZs2ICuXbvC0tISp06dQk5ODgAgLS0Nc+fO1XuAxhIVFYXg4OASN2etTH47cx9/nI+HmQSY2j2Y+4YRERH9Q+dEaPbs2ViyZAm+++47WFj8O+OoVatWOHXqlF6DM6aIiAhcunQJx48fN3YoL+Tuo0x8uOkCgIJFEoO9KveAbyIiIn3SORG6evUq2rVrV6Tczs4OKSkp+oiJ9CRfpca4taeRlpOPJn6OeKdDDWOHREREVKHonAh5enpq7Q1W6ODBgwgMDNRLUKQfX+65jtOxKbBVSLGwf0PuG0ZERPQUnT8Zx4wZg3fffRdHjx6FRCJBXFwcVq9ejYkTJ2Ls2LGGiJHKYOfFBHyzryBhjexTDz6OVkaOiIiIqOLRefr8Bx98gNTUVISHhyM7Oxvt2rWDXC7HxIkT8fbbbxsiRtLRhpP38MGGcxAC+E8TH/Ss72XskIiIiCoknVeWLpSZmYlLly5BrVYjODgYNjY2+o6tQqhsK0t/f+AWZv9xGQDQp7E3Pulbn2sGERGRyTHYytKFrKys0KRJk7L+eIVX2fYay81XY/6OK/juQMGiiSPbBGBq9yCYmXGqPBERUUnK3CJkKipDi9DVhDS8t+4MLsUrAQD/17U2xravzvWCiIjIZBm8RYiMTwiB7w/EYP6Oq8hVqeFoZYE5veuhez1PY4dGRERUKTARqsT2XE7EnG0F44E61HHDvD714GanMHJURERElQcToUrs7L0UAEDP+p74ekAjdoURERHpqEzTiX788Ue0bt0aXl5euHPnDgBg4cKF+O233/QaHD3b/cdZAIBgLzsmQURERGWgcyK0ePFiTJgwAd27d0dKSopmVpWDgwMWLlyo7/joGe6lFCRC3g6WRo6EiIioctI5Efr666/x3XffYerUqTA3N9eUN2nSBOfPn9drcMZUGXafL2wR8nFkIkRERFQWOidCMTExaNSoUZFyuVyOjIwMvQRVEVT03efzVWokKLMBAN4O3D6DiIioLHROhAICAnDmzJki5X/++SeCg4P1EROVwoO0HKjUAhbmErjZyo0dDhERUaWk86yx//u//0NERASys7MhhMCxY8ewZs0aREZG4vvvvzdEjFSMwm4xT3tLrh5NRERURjonQsOHD0d+fj4++OADZGZmYuDAgfD29saXX36J/v37GyJGKsb9lEwAHChNRET0Isq0jtDo0aMxevRoJCUlQa1Ww83NTd9x0XMUtgh5c6A0ERFRmek8RmjmzJm4efMmAMDFxYVJkJHce8yp80RERC9K50Row4YNqFWrFlq0aIFvvvkGDx8+NERc9Bz3Uzh1noiI6EXpnAidO3cO586dQ4cOHbBgwQJ4e3uje/fu+Pnnn5GZmWmIGKkY7BojIiJ6cWXaYqNu3bqYO3cubt26hX379iEgIADjx4+Hh4eHvuMzmoq8oKIQ4t8WIa4hREREVGZlSoSeZG1tDUtLS8hkMuTl5ekjpgqhIi+omJSei5x8NSQSwMOeu80TERGVVZkSoZiYGMyZMwfBwcFo0qQJTp06hRkzZiAhIUHf8VExCluD3G0VkElfOJclIiIyWTpPn2/ZsiWOHTuGevXqYfjw4Zp1hKj8cHwQERGRfuicCIWHh+P7779H3bp1DREPlQIXUyQiItIPnROhuXPnGiIO0gFbhIiIiPSjVInQhAkT8PHHH8Pa2hoTJkx45rkLFizQS2BUssIxQmwRIiIiejGlSoROnz6tmRF2+vRpgwZEz3ePLUJERER6UapEaN++fcV+T8bx7xpCTISIiIhehM5zr0eMGIG0tLQi5RkZGRgxYoRegqKSKbPzkJadD4AtQkRERC9K50Ro1apVyMrKKlKelZWFH374QS9BUckKB0o7WctgJdN5rDsRERE9odSfpEqlEkIICCGQlpYGheLfFY1VKhW2bdvGnejLAXedJyIi0p9SJ0IODg6QSCSQSCSoVatWkeMSiQQzZ87Ua3DGFBUVhaioKKhUKmOHouX+Y64hREREpC+lToT27dsHIQQ6dOiADRs2wMnJSXNMJpPBz88PXl5eBgnSGCIiIhAREQGlUgl7e3tjh6OhmTrP8UFEREQvrNSJUFhYGICCfcaqVasGMzPucWUMXEOIiIhIf3Qebevn5wcAyMzMRGxsLHJzc7WO169fXz+RUbG4qjQREZH+6JwIPXz4EMOHD8eff/5Z7PGKNqamKll99A7O308FAAS4WBs5GiIiospP5/6t8ePH4/Hjx/j7779haWmJ7du3Y9WqVahZsyZ+//13Q8Ro8oQQ+GLXNUzddAFqAQxs7ota7rbGDouIiKjS07lFaO/evfjtt9/QtGlTmJmZwc/PD507d4adnR0iIyPRo0cPQ8RpslRqgQ83n8eaY3cBAOM61sR7nWoaOSoiIqKqQecWoYyMDM16QU5OTnj48CEAoF69ejh16pR+oyNsPHUPa47dhZkEmP1qCCZ0rgWJRGLssIiIiKoEnROh2rVr4+rVqwCAhg0bYunSpbh//z6WLFkCT09PvQdo6nZcTAAARITXwKAWfkaOhoiIqGrRuWts/PjxiI+PBwBMnz4dXbt2xerVqyGTybBy5Up9x2fSsvNUOHgjCQDQLYRJJhERkb7pnAi98cYbmu8bNWqE27dv48qVK/D19YWLi4tegzN1h24kITtPDS97BYI8OTiaiIhI3154104rKys0btxYH7HQU3ZfTgQAdAxy57ggIiIiAyhVIjRhwoRSX3DBggVlDob+JYTA3isPAAAdgriZLRERkSGUKhE6ffp0qS7GVgv9uRinxANlDqxk5mgZ6GzscIiIiKqkUiVC+/btM3Qc9JTdlwtag9rUcIHCwtzI0RAREVVNZd459caNG9ixYweysgr2vhJC6C0oAvb8Mz6oU5C7kSMhIiKqunROhJKTk9GxY0fUqlUL3bt310ylHzVqFN5//329B2iKHiizNXuKta/jauRoiIiIqi6dE6H33nsPFhYWiI2NhZWVlaa8X79+2L59u16DM6aoqCgEBwejadOm5f7ce68UtAY1qOYAN1tFuT8/ERGRqdB5+vzOnTuxY8cO+Pj4aJXXrFkTd+7c0VtgxhYREYGIiAgolUrY29uX63Pv+Wd8UKc6nC1GRERkSGXaa+zJlqBCSUlJkMvlegnK1J25W9At1rYWu8WIiIgMSedEqF27dvjhhx80jyUSCdRqNebPn4/w8HC9BmeKhBBIzcoFALjZMrEkIiIyJJ27xubPn4/27dvjxIkTyM3NxQcffICLFy/i0aNHOHTokCFiNCnZeWrkqQpm4NlZWhg5GiIioqpN5xah4OBgnDt3Ds2aNUPnzp2RkZGBPn364PTp06hevbohYjQpqVl5AABzMwmsZVw/iIiIyJB0ahHKy8tDly5dsHTpUsycOdNQMZm0wkTI3tKCK3UTEREZmE4tQhYWFrhw4QI/oA3oyUSIiIiIDEvnrrEhQ4Zg2bJlhoiF8G8ixPFBREREhqfzYOnc3Fx8//332LVrF5o0aQJra2ut49x9/sUoCxMhhc7/NURERKQjnT9tL1y4gMaNGwMArl27pnWMXWYvjl1jRERE5UfnRIg70RsWEyEiIqLyU+bd58kwmAgRERGVHyZCFYySiRAREVG5YSJUwSizOWuMiIiovDARqmDYNUZERFR+mAhVMEyEiIiIyg8ToQqGiRAREVH5YSJUwSiz8gEwESIiIioPTIQqkNx8NbLyVAAAOwUTISIiIkNjIlSBFHaLSSSALbfYICIiMjgmQhVIYSJkK5fCzIzblRARERkaE6EKRDNQ2ordYkREROWhyidCd+/eRfv27REcHIz69evj119/NXZIJdIspsjxQUREROWiyg9EkUqlWLhwIRo2bIjExEQ0btwY3bt3h7W1tbFDK4LbaxAREZWvKp8IeXp6wtPTEwDg5uYGJycnPHr0qEImQlxDiIiIqHwZvWvsr7/+wssvvwwvLy9IJBJs3ry5yDmLFi1CQEAAFAoFQkNDceDAgTI914kTJ6BWq1GtWrUXjNowUjOZCBEREZUnoydCGRkZaNCgAb755ptij69btw7jx4/H1KlTcfr0abRt2xbdunVDbGys5pzQ0FCEhIQU+YqLi9Ock5ycjCFDhuDbb781+GsqK7YIERERlS+jd41169YN3bp1K/H4ggULMHLkSIwaNQoAsHDhQuzYsQOLFy9GZGQkAODkyZPPfI6cnBz07t0bU6ZMQatWrZ57bk5OjuaxUqks7Ut5Ydx5noiIqHwZvUXoWXJzc3Hy5El06dJFq7xLly44fPhwqa4hhMCwYcPQoUMHDB48+LnnR0ZGwt7eXvNVnt1ohS1CTISIiIjKR4VOhJKSkqBSqeDu7q5V7u7ujoSEhFJd49ChQ1i3bh02b96Mhg0bomHDhjh//nyJ50+ZMgWpqamar7t3777Qa9AFu8aIiIjKl9G7xkpDItFeZVkIUaSsJG3atIFarS71c8nlcsjlcp3i05dUbrhKRERUrip0i5CLiwvMzc2LtP4kJiYWaSWqCriOEBERUfmq0ImQTCZDaGgodu3apVW+a9eu5w56flFRUVEIDg5G06ZNDfo8TypMhOy44SoREVG5MPonbnp6Om7cuKF5HBMTgzNnzsDJyQm+vr6YMGECBg8ejCZNmqBly5b49ttvERsbi7feesugcUVERCAiIgJKpRL29vYGfS4AUKkF0nLYNUZERFSejJ4InThxAuHh4ZrHEyZMAAAMHToUK1euRL9+/ZCcnIxZs2YhPj4eISEh2LZtG/z8/IwVskEUtgYBnDVGRERUXoyeCLVv3x5CiGeeM3bsWIwdO7acIjKOwhlj1jJzWJhX6B5LIiKiKoOfuBVE4WKK7BYjIiIqP0yESlDeg6W5mCIREVH5YyJUgoiICFy6dAnHjx8vl+djIkRERFT+mAhVEFxVmoiIqPwxEaogmAgRERGVPyZCFYSS22sQERGVOyZCJTDaYGkFEyEiIqLywkSoBOU9WPrffcaMvrQTERGRyWAiVEFoxghZsUWIiIiovDARqiA4WJqIiKj8MRGqILiyNBERUfljIlRBcLA0ERFR+WMiVAGo1eKJwdJMhIiIiMoLE6ESlOf0+fTcfKhFwffcYoOIiKj8MBEqQXlOn0/NLGgNkkvNoLAwN/jzERERUQEmQhUAB0oTEREZBxOhCoA7zxMRERkHE6EKgAOliYiIjIOJkJGdvPMY83dcBQC42MiMHA0REZFp4cZWRpKVq8JnO69i+aEYCAG42srxdnhNY4dFRERkUpgIlSAqKgpRUVFQqVR6v/bDtBy8tuQw7iRnAgD6NvbBRz2D4GDFFiEiIqLyJBFCCGMHUZEplUrY29sjNTUVdnZ2ermmEALDVx7H1YQ0zO1TD+G13fRyXSIiIipQ2s9vtggZgUQiwfzXGkBuYcYtNYiIiIyIiZCRuNrKjR0CERGRyeOsMSIiIjJZTISIiIjIZDERIiIiIpPFRIiIiIhMFhMhIiIiMllMhIiIiMhkMREqQVRUFIKDg9G0aVNjh0JEREQGwpWln8MQK0sTERGRYZX285stQkRERGSymAgRERGRyWIiRERERCaLiRARERGZLCZCREREZLK4+/xzFE6qUyqVRo6EiIiISqvwc/t5k+OZCD1HWloaAKBatWpGjoSIiIh0lZaWBnt7+xKPcx2h51Cr1YiLi4OtrS0kEonerqtUKlGtWjXcvXuX6xM9gfVSPNZL8VgvxWO9FI/1UryqWi9CCKSlpcHLywtmZiWPBGKL0HOYmZnBx8fHYNe3s7OrUjeevrBeisd6KR7rpXisl+KxXopXFevlWS1BhThYmoiIiEwWEyEiIiIyWUyEjEQul2P69OmQy+XGDqVCYb0Uj/VSPNZL8VgvxWO9FM/U64WDpYmIiMhksUWIiIiITBYTISIiIjJZTISIiIjIZDERIiIiIpPFRMhIFi1ahICAACgUCoSGhuLAgQPGDqncREZGomnTprC1tYWbmxteffVVXL16VescIQRmzJgBLy8vWFpaon379rh48aKRIjaOyMhISCQSjB8/XlNmqvVy//59DBo0CM7OzrCyskLDhg1x8uRJzXFTrJf8/Hx8+OGHCAgIgKWlJQIDAzFr1iyo1WrNOaZQL3/99RdefvlleHl5QSKRYPPmzVrHS1MHOTk5eOedd+Di4gJra2u88soruHfvXjm+Cv17Vr3k5eVh0qRJqFevHqytreHl5YUhQ4YgLi5O6xpVsV6KJajcrV27VlhYWIjvvvtOXLp0Sbz77rvC2tpa3Llzx9ihlYuuXbuKFStWiAsXLogzZ86IHj16CF9fX5Genq45Z968ecLW1lZs2LBBnD9/XvTr1094enoKpVJpxMjLz7Fjx4S/v7+oX7++ePfddzXlplgvjx49En5+fmLYsGHi6NGjIiYmRuzevVvcuHFDc44p1svs2bOFs7Oz2Lp1q4iJiRG//vqrsLGxEQsXLtScYwr1sm3bNjF16lSxYcMGAUBs2rRJ63hp6uCtt94S3t7eYteuXeLUqVMiPDxcNGjQQOTn55fzq9GfZ9VLSkqK6NSpk1i3bp24cuWKOHLkiGjevLkIDQ3VukZVrJfiMBEygmbNmom33npLq6xOnTpi8uTJRorIuBITEwUAsX//fiGEEGq1Wnh4eIh58+ZpzsnOzhb29vZiyZIlxgqz3KSlpYmaNWuKXbt2ibCwME0iZKr1MmnSJNGmTZsSj5tqvfTo0UOMGDFCq6xPnz5i0KBBQgjTrJenP/BLUwcpKSnCwsJCrF27VnPO/fv3hZmZmdi+fXu5xW5IxSWITzt27JgAoPmD3BTqpRC7xspZbm4uTp48iS5dumiVd+nSBYcPHzZSVMaVmpoKAHBycgIAxMTEICEhQauO5HI5wsLCTKKOIiIi0KNHD3Tq1Emr3FTr5ffff0eTJk3w+uuvw83NDY0aNcJ3332nOW6q9dKmTRvs2bMH165dAwCcPXsWBw8eRPfu3QGYbr08qTR1cPLkSeTl5Wmd4+XlhZCQEJOpJ6DgfVgikcDBwQGAadULN10tZ0lJSVCpVHB3d9cqd3d3R0JCgpGiMh4hBCZMmIA2bdogJCQEADT1UFwd3blzp9xjLE9r167FqVOncPz48SLHTLVebt26hcWLF2PChAn43//+h2PHjmHcuHGQy+UYMmSIydbLpEmTkJqaijp16sDc3BwqlQpz5szBgAEDAJju/fKk0tRBQkICZDIZHB0di5xjKu/J2dnZmDx5MgYOHKjZdNWU6oWJkJFIJBKtx0KIImWm4O2338a5c+dw8ODBIsdMrY7u3r2Ld999Fzt37oRCoSjxPFOrF7VajSZNmmDu3LkAgEaNGuHixYtYvHgxhgwZojnP1Opl3bp1+Omnn/Dzzz+jbt26OHPmDMaPHw8vLy8MHTpUc56p1UtxylIHplJPeXl56N+/P9RqNRYtWvTc86tivbBrrJy5uLjA3Ny8SEadmJhY5K+Wqu6dd97B77//jn379sHHx0dT7uHhAQAmV0cnT55EYmIiQkNDIZVKIZVKsX//fnz11VeQSqWa125q9eLp6Yng4GCtsqCgIMTGxgIw3fvl//7v/zB58mT0798f9erVw+DBg/Hee+8hMjISgOnWy5NKUwceHh7Izc3F48ePSzynqsrLy8N//vMfxMTEYNeuXZrWIMC06oWJUDmTyWQIDQ3Frl27tMp37dqFVq1aGSmq8iWEwNtvv42NGzdi7969CAgI0DoeEBAADw8PrTrKzc3F/v37q3QddezYEefPn8eZM2c0X02aNMEbb7yBM2fOIDAw0CTrpXXr1kWWV7h27Rr8/PwAmO79kpmZCTMz7bdwc3NzzfR5U62XJ5WmDkJDQ2FhYaF1Tnx8PC5cuFCl66kwCbp+/Tp2794NZ2dnreMmVS/GGqVtygqnzy9btkxcunRJjB8/XlhbW4vbt28bO7Ry8d///lfY29uL6OhoER8fr/nKzMzUnDNv3jxhb28vNm7cKM6fPy8GDBhQ5ab9lsaTs8aEMM16OXbsmJBKpWLOnDni+vXrYvXq1cLKykr89NNPmnNMsV6GDh0qvL29NdPnN27cKFxcXMQHH3ygOccU6iUtLU2cPn1anD59WgAQCxYsEKdPn9bMfipNHbz11lvCx8dH7N69W5w6dUp06NCh0k8Tf1a95OXliVdeeUX4+PiIM2fOaL0P5+TkaK5RFeulOEyEjCQqKkr4+fkJmUwmGjdurJk6bgoAFPu1YsUKzTlqtVpMnz5deHh4CLlcLtq1ayfOnz9vvKCN5OlEyFTrZcuWLSIkJETI5XJRp04d8e2332odN8V6USqV4t133xW+vr5CoVCIwMBAMXXqVK0PMlOol3379hX7fjJ06FAhROnqICsrS7z99tvCyclJWFpaip49e4rY2FgjvBr9eVa9xMTElPg+vG/fPs01qmK9FEcihBDl1/5EREREVHFwjBARERGZLCZCREREZLKYCBEREZHJYiJEREREJouJEBEREZksJkJERERkspgIERERkcliIkREVU50dDQkEglSUlKMFkNmZib69u0LOzs7o8dCRCVjIkRElcKMGTPQsGHDUp3bqlUrxMfHw97e3rBBPcOqVatw4MABHD582OixEFHJpMYOgIhIn/Ly8iCTyTQ7jxvLzZs3ERQUhJCQEKPGQUTPxhYhIjI4IQQ+/fRTBAYGwtLSEg0aNMD69es1xwu7svbs2YMmTZrAysoKrVq10uw6v3LlSsycORNnz56FRCKBRCLBypUrAQASiQRLlixBr169YG1tjdmzZxfbNXb48GG0a9cOlpaWqFatGsaNG4eMjAzN8UWLFqFmzZpQKBRwd3fHa6+99szXtGHDBtStWxdyuRz+/v74/PPPNcfat2+Pzz//HH/99RckEgnat29f5Odv374NMzMznDhxQqv866+/hp+fH7j7EVE5Me5WZ0RkCv73v/+JOnXqiO3bt4ubN2+KFStWCLlcLqKjo4UQ/24Q2bx5cxEdHS0uXrwo2rZtK1q1aiWEECIzM1O8//77om7duppdsjMzM4UQBZv4urm5iWXLlombN2+K27dva673+PFjIYQQ586dEzY2NuKLL74Q165dE4cOHRKNGjUSw4YNE0IIcfz4cWFubi5+/vlncfv2bXHq1Cnx5Zdflvh6Tpw4IczMzMSsWbPE1atXxYoVK4SlpaVm4+Dk5GQxevRo0bJlSxEfHy+Sk5OLvU7nzp3F2LFjtcoaNWokpk2bVua6JiLdMBEiIoNKT08XCoVCHD58WKt85MiRYsCAAUKIfxOh3bt3a47/8ccfAoDIysoSQggxffp00aBBgyLXByDGjx+vVfZ0IjR48GDx5ptvap1z4MABYWZmJrKyssSGDRuEnZ2dUCqVpXpNAwcOFJ07d9Yq+7//+z8RHBysefzuu++KsLCwZ15n3bp1wtHRUWRnZwshhDhz5oyQSCQiJiamVHEQ0Ytj1xgRGdSlS5eQnZ2Nzp07w8bGRvP1ww8/4ObNm1rn1q9fX/O9p6cnACAxMfG5z9GkSZNnHj958iRWrlyp9fxdu3aFWq1GTEwMOnfuDD8/PwQGBmLw4MFYvXo1MjMzS7ze5cuX0bp1a62y1q1b4/r161CpVM+Nt9Crr74KqVSKTZs2AQCWL1+O8PBw+Pv7l/oaRPRiOFiaiAxKrVYDAP744w94e3trHZPL5VqPLSwsNN9LJBKtn38Wa2vr58YwZswYjBs3rsgxX19fyGQynDp1CtHR0di5cyemTZuGGTNm4Pjx43BwcCjyM0IITXxPlulKJpNh8ODBWLFiBfr06YOff/4ZCxcu1Pk6RFR2TISIyKCCg4Mhl8sRGxuLsLCwMl9HJpPp1NrypMaNG+PixYuoUaNGiedIpVJ06tQJnTp1wvTp0+Hg4IC9e/eiT58+Rc4NDg7GwYMHtcoOHz6MWrVqwdzcXKfYRo0ahZCQECxatAh5eXnFPh8RGQ4TISIyKFtbW0ycOBHvvfce1Go12rRpA6VSicOHD8PGxgZDhw4t1XX8/f0RExODM2fOwMfHB7a2tkValEoyadIktGjRAhERERg9ejSsra1x+fJl7Nq1C19//TW2bt2KW7duoV27dnB0dMS2bdugVqtRu3btYq/3/vvvo2nTpvj444/Rr18/HDlyBN988w0WLVpU6nopFBQUhBYtWmDSpEkYMWIELC0tdb4GEZUdxwgRkcF9/PHHmDZtGiIjIxEUFISuXbtiy5YtCAgIKPU1+vbti5deegnh4eFwdXXFmjVrSv2z9evXx/79+3H9+nW0bdsWjRo1wkcffaQZh+Tg4ICNGzeiQ4cOCAoKwpIlS7BmzRrUrVu32Os1btwYv/zyC9auXYuQkBBMmzYNs2bNwrBhw0od05NGjhyJ3NxcjBgxokw/T0RlJxFl6dgmIiK9mTNnDtauXYvz588bOxQik8MWISIiI0lPT8fx48fx9ddfFzuQm4gMj4kQEZGRvP3222jTpg3CwsLYLUZkJOwaIyIiIpPFFiEiIiIyWUyEiIiIyGQxESIiIiKTxUSIiIiITBYTISIiIjJZTISIiIjIZDERIiIiIpPFRIiIiIhMFhMhIiIiMln/DwGVmThYRutSAAAAAElFTkSuQmCC", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "plt.plot(np.sort(np.abs((yhat - Y)/Y)))\n", + "plt.yscale(\"log\")\n", + "plt.title(\"Prediction Relative error\")\n", + "plt.xlabel(\"entries of y\")\n", + "plt.ylabel(\"relative error between y_hat and y\")" + ] + }, + { + "cell_type": "markdown", + "id": "d28d8596", + "metadata": {}, + "source": [ + "### Poisson VAR model" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "776fd2bf", + "metadata": {}, + "outputs": [], + "source": [ + "# Example usage\n", + "n_features = 10\n", + "n_samples = 10\n", + "lag = 1\n", + "\n", + "\n", + "base_intensity = np.random.randint(1, 10, n_features).astype(np.float64) # Base rates for each feature\n", + "\n", + "data, transition_matrices, _ = generate_sparse_stationary_var_process(\n", + " n_features,\n", + " n_samples,\n", + " lag=lag,\n", + " sparsity=0.5,\n", + " quench_factor = 1, \n", + " spectral_radius=0.9, # Ensures stationarity\n", + " process_type='poisson',\n", + " base_intensity=base_intensity\n", + "\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8af52915-53e2-43d8-a41e-aa89f07a17f3", + "metadata": {}, + "outputs": [], + "source": [ + "# have to set fit_intercept = False since the curretn VAR implementation doesn consider constant model terms\n", + "poisson = UoI_Poisson(n_real_features = n_features, fit_VAR =True, max_iter=2500, fit_intercept=False, random_state=42)\n", + "\n", + "poisson.fit(data)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "9e6d1b70", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "0.21311475409836067" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "B_model = poisson.VAR_coef_\n", + "#B_truth = np.vstack([m.toarray().T for m in transition_matrices]).T.flatten()\n", + "selection_accuracy(B_truth, B_model)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fac6020e", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "pyuoi", + "language": "python", + "name": "pyuoi" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/Util_NumpyIO.py b/examples/Util_NumpyIO.py new file mode 100755 index 00000000..0ab9b638 --- /dev/null +++ b/examples/Util_NumpyIO.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +__author__ = "Jan Balewski" +__email__ = "janstar1122@gmail.com" + +''' = = = = = Numpy NPZ advanced storage = = = +It can hold: +* python dictionaries which must pass: json.dumps(dict) +* single float or int variables w/o np-array packing. It is recovered as 1-value array +* arbitrary numpy array (a large size payloads) + - for an array of arbitrary strings must declare dtype='object' at write and use .decode("utf-8") to unpack +''' + +import numpy as np +import time, os +import json,time +from pprint import pprint + +#...!...!.................. +def write_data_npz(dataD,outF,metaD=None,verb=1): + assert type(dataD)!=type(None) + assert len(outF)>0 + + # Create a copy to avoid modifying original data + saveD = dataD.copy() + + if metaD!=None: + #pprint(metaD) + metaJ=json.dumps(metaD, default=str) + #print('meta.JSON:',metaJ) + saveD['meta.JSON']=np.array([metaJ], dtype='object') + + if verb>0: + print('saving data as npz:',outF) + start = time.time() + + # Process data to ensure all items are numpy arrays + for item in list(saveD.keys()): + rec=saveD[item] + if verb>1: print('x=',item,type(rec)) + if type(rec)==str: # special case - convert string to object array + saveD[item] = np.array([rec], dtype='object') + if verb>0:print('npz-write :',item, 'as string array',saveD[item].shape,saveD[item].dtype) + continue + if type(rec)!=np.ndarray: # packs a single value into np-array + saveD[item]=np.array([rec]) + if verb>0:print('npz-write :',item, saveD[item].shape,saveD[item].dtype) + else: + if verb>0:print('npz-write :',item, rec.shape,rec.dtype) + + # Save to NPZ format + np.savez_compressed(outF, **saveD) + + xx=os.path.getsize(outF)/1048576 + print('closed npz:',outF,' size=%.2f MB, elaT=%.1f sec'%(xx,(time.time() - start))) + + +#...!...!.................. +def read_data_npz(inpF,verb=1): + if verb>0: + print('read data from npz:',inpF) + start = time.time() + + npzData = np.load(inpF, allow_pickle=True) + objD={} + + for x in npzData.files: + obj = npzData[x] + if verb>1: print('\nitem=',x,type(obj),obj.shape,obj.dtype) + + if obj.dtype==object: + if verb>0: print('read obj:',x,len(obj),type(obj)) + else: + if verb>0: print('read obj:',x,obj.shape,obj.dtype) + objD[x]=obj + + npzData.close() + + try: + inpMD=json.loads(objD.pop('meta.JSON')[0]) + if verb>1: print(' recovered meta-data with %d keys'%len(inpMD)) + except: + inpMD=None + if verb>0: + print(' done npz, num rec:%d elaT=%.1f sec'%(len(objD),(time.time() - start))) + + return objD,inpMD + + + +#================================= +#================================= +# U N I T T E S T +#================================= +#================================= + +if __name__=="__main__": + print('testing npzIO ver 1') + outF='abcTest.npz' + verb=1 + + var1=float(15) # single variable + one=np.zeros(shape=5,dtype=np.int16); one[3]=3 + two=np.zeros(shape=(2,3)); two[1,2]=4 + + three=np.empty((2), dtype='object') + three[0]='record aaaa' + three[1]='much longer record bbb' + # WARN: all decalred elements of three[] must be initialized before writeing NPZ + + # this works too: + # three=np.array(['record aaaa','much longer record bbb'], dtype='object') + + text='This is text1' + + metaD={"age":17,"dom":"white","dates":[11,22,33]} + + outD={'one':one,'two':two,'var1':var1,'atext':text,'three':three} + + write_data_npz(outD,outF,metaD=metaD,verb=verb) + + print('\nM: ***** verify by reading it back from',outF) + big,meta2=read_data_npz(outF,verb=verb) + from pprint import pprint + print(' recovered meta-data'); pprint(meta2) + print('dump read-in data') + for x in big: + print('\nkey=',x); pprint(big[x]) + + #get one string from string-array + rec2=big['three'][1] + print('rec2:',type(rec2),rec2) + print('\n check raw content with: python -c "import numpy as np; data=np.load(\'%s\', allow_pickle=True); print(list(data.files)); data.close()"\n'%outF) diff --git a/examples/Util_poissonFdr.py b/examples/Util_poissonFdr.py new file mode 100644 index 00000000..80e61f32 --- /dev/null +++ b/examples/Util_poissonFdr.py @@ -0,0 +1,261 @@ +#...!...!.................. +import os +import numpy as np + +if 0: # pop-up canvas + import matplotlib as mpl + mpl.use('TkAgg') + +#...!...!.................. +def qa_Bfit(B_truth, B_fit): + assert B_truth.shape == B_fit.shape + bterm_res = B_fit - B_truth + bterm_mean = np.mean(bterm_res) + bterm_std = np.std(bterm_res) + N = bterm_res.shape[0] + bterm_se_s = 0. + if N > 1: + bterm_se_s = bterm_std / np.sqrt(2 * (N - 1)) + print('B term: mean=%.2f std=%.3f +/- %.3f' % (bterm_mean, bterm_std, bterm_se_s)) + + out = { + 'tval': B_truth, + 'fval': B_fit, + 'res_mean': bterm_mean, + 'res_std': bterm_std + } + return out + + +#...!...!.................. +def qa_Afit(A_truth, A_fit): + assert A_truth.shape == A_fit.shape + assert A_truth.shape[0] == A_truth.shape[1] + + M = A_truth.shape[0] + diag_m = np.eye(M, dtype=bool) + out = {} + + # --- Diagonal elements + diag_tval = A_truth[diag_m] + diag_fval = A_fit[diag_m] + res_diag = diag_fval - diag_tval + diag_mean = np.mean(res_diag) if res_diag.size > 0 else 0 + diag_std = np.std(res_diag) if res_diag.size > 0 else 0 + out['diag'] = { + 'tval': diag_tval, + 'fval': diag_fval, + 'res_mean': diag_mean, + 'res_std': diag_std + } + + # --- Off-diagonal elements (Excitatory and Inhibitory) + for name in ['inh','exc']: + if name == 'exc': + true_m = (A_truth > 0) & (~diag_m) + pred_m = (A_fit > 0) & (~diag_m) + else: # inh + true_m = (A_truth < 0) & (~diag_m) + pred_m = (A_fit < 0) & (~diag_m) + + tp_m = true_m & pred_m + tval = A_truth[tp_m] + fval = A_fit[tp_m] + res = fval - tval + res_mean = np.mean(res) if res.size > 0 else 0 + res_std = np.std(res) if res.size > 0 else 0 + + TP = np.sum(tp_m) + FP = np.sum((~true_m) & pred_m) + FN = np.sum(true_m & (~pred_m)) + + out[name] = { + 'tval': tval, + 'fval': fval, + 'res_mean': res_mean, + 'res_std': res_std, + 'TP': TP, + 'FP': FP, + 'FN': FN + } + + # Keep printing for compatibility + for name, stats in [('exc', out['exc']), ('inh', out['inh'])]: + res_N = stats['tval'].shape[0] + se_s = 0. + if res_N > 1: + se_s = stats['res_std'] / np.sqrt(2 * (res_N - 1)) + + print('A type=%s TP:%d FP:%d FN:%d TP: mean=%.2f std=%.3f +/- %.3f'%(name, stats['TP'], stats['FP'], stats['FN'], stats['res_mean'], stats['res_std'], se_s)) + + diag_stats = out['diag'] + diag_N = diag_stats['tval'].shape[0] + diag_se_s = 0. + if diag_N > 1: + diag_se_s = diag_stats['res_std'] / np.sqrt(2 * (diag_N - 1)) + + print('A diag: mean=%.2f std=%.3f +/- %.3f'%(diag_stats['res_mean'], diag_stats['res_std'], diag_se_s)) + + return out + + +#...!...!.................. +def plot_edges_correl(args, qaD, outName): + #import matplotlib as mpl + #mpl.use('TkAgg') + import matplotlib.pyplot as plt + fig, axes = plt.subplots(1, 4, figsize=(14, 4.)) + + colors = {'exc': 'red', 'inh': 'blue', 'diag': 'brown', 'bterm': 'green'} + + for i, name in enumerate(['inh', 'exc', 'diag', 'bterm']): + ax = axes[i] + stats = qaD[name] + color = colors[name] + + if stats['tval'].size > 0: + ax.scatter(stats['tval'], stats['fval'], alpha=0.5, s=8, c=color) + # Add y=x line + lims = [ + np.min([ax.get_xlim(), ax.get_ylim()]), + np.max([ax.get_xlim(), ax.get_ylim()]), + ] + ax.plot(lims, lims, 'r--', alpha=0.75, zorder=0) + ax.set_aspect('equal', 'box') + ax.set_xlim(lims) + ax.set_ylim(lims) + + # Add center of gravity cross + mean_t = np.mean(stats['tval']) + mean_f = np.mean(stats['fval']) + ax.plot(mean_t, mean_f, '+', c='black', markersize=20, markeredgewidth=3) + + ax.grid() + ax.set_title(name.capitalize()) + ax.set_xlabel('True Value') + if i == 0: + ax.set_ylabel('Fitted Value') + + # Add mean and std text + mean_val = stats['res_mean'] + std_val = stats['res_std'] + ax.text(0.95, 0.05, f'TP N={stats["tval"].shape[0]} \nmean={mean_val:.3f}\nstd={std_val:.3f}', + transform=ax.transAxes, fontsize=10, + verticalalignment='bottom', horizontalalignment='right', + bbox=dict(boxstyle='round,pad=0.3', fc='white', alpha=0.5)) + + if name in ['exc', 'inh']: + tp_val = stats['TP'] + fp_val = stats['FP'] + fn_val = stats['FN'] + ax.text(0.05, 0.95, f'TP={tp_val}\nFP={fp_val}\nFN={fn_val}', + transform=ax.transAxes, fontsize=10, + verticalalignment='top', horizontalalignment='left', + bbox=dict(boxstyle='round,pad=0.3', fc='white', alpha=0.5)) + + fig.suptitle(f'Fitted {outName} %d samples'%(args.samples)) + fig.tight_layout() + plotF = os.path.join(args.out, f'{outName}_corr.png') + fig.savefig(plotF) + print(f'Saved display {plotF}') + plt.show() + + +#...!...!.................. +def plot_eigenvalues(args, A_truth, A_fit, outName): + import matplotlib.pyplot as plt + eigT = np.linalg.eigvals(A_truth) + eigF = np.linalg.eigvals(A_fit) + + fig, ax = plt.subplots(1, 1, figsize=(8, 5)) + + reT = np.real(eigT) + imT = np.imag(eigT) + ax.scatter(reT, imT, color='blue', marker='o', label='True', s=10) + + reF = np.real(eigF) + imF = np.imag(eigF) + ax.scatter(reF, imF, color='red', marker='o', facecolors='none', label='fit', s=20) + + ax.set_ylim(-0.1,) + #ax.set_xlim(right=1) + ax.axhline(0, linestyle='--', color='k', linewidth=1) + ax.axvline(0, linestyle='--', color='k', linewidth=1) + ax.grid(True, alpha=0.3) + ax.set_xlabel("Real Part") + ax.set_ylabel("Imaginary Part") + ax.legend() + ax.set_title(f'Eigenvalues for {outName} %d samples'%(args.samples)) + + plotF = os.path.join(args.out, f'{outName}_eigen.png') + fig.savefig(plotF) + print(f'Saved display {plotF}') + plt.show() + + +#...!...!.................. +def plot_loss(args, A_truth, A_fit, outName, l1_loss_sel, loss_skip=0): + import matplotlib.pyplot as plt + + fig, axes = plt.subplots(1, 3, figsize=(15, 4.5)) + + # Panel A: L1 loss vs. iteration + if l1_loss_sel.shape[1] > 0: + mask = l1_loss_sel[0] >= loss_skip + axes[0].plot(l1_loss_sel[0, mask], l1_loss_sel[1, mask], marker='o') + #axes[0].set_yscale('log') + axes[0].set_title('Selection L1 loss') + axes[0].set_xlabel('Iteration') + axes[0].set_ylabel('L1 Loss') + axes[0].grid(True) + + # Prepare masks + M = A_truth.shape[0] + diag_mask = np.eye(M, dtype=bool) + exc_mask = (A_truth > 0) & (~diag_mask) + inh_mask = (A_truth < 0) & (~diag_mask) + + # Panel B: Excitatory distribution + for idx, name in enumerate(['exc', 'inh']): + ax = axes[idx + 1] + if name == 'exc': + mask = exc_mask + title = 'Excitatory edge weights' + else: + mask = inh_mask + title = 'Inhibitory edge weights' + + truth_vals = A_truth[mask] + fit_vals = A_fit[mask] + if truth_vals.size > 0 or fit_vals.size > 0: + ax.hist(truth_vals, bins=40, alpha=0.6, label='truth') + ax.hist(fit_vals, bins=40, alpha=0.6, label='fit') + ax.set_title(title) + ax.set_xlabel('Weight') + ax.set_ylabel('Count') + ax.legend() + + fig.suptitle(f'Fitted {outName} %d samples'%(args.samples)) + fig.tight_layout() + plotF = os.path.join(args.out, f'{outName}_loss_edges.png') + fig.savefig(plotF) + print('Saved display %s'%plotF) + plt.show() + + +#...!...!.................. +def extract_loss_traces(model, loss_stride): + loss_iter = np.empty(0, dtype=np.int64) + loss_l1 = np.empty(0, dtype=np.float64) + + if loss_stride is None: + loss_stride = 10 + + if model is not None and hasattr(model, 'loss'): + loss = model.loss + if loss is not None and 'l1' in loss: + loss_l1 = np.asarray(loss['l1'], dtype=np.float64) + loss_iter = np.arange(loss_l1.size, dtype=np.int64) * loss_stride + + return loss_iter, loss_l1 + diff --git a/examples/batchPoissonFDR.slr b/examples/batchPoissonFDR.slr new file mode 100755 index 00000000..65f7fa08 --- /dev/null +++ b/examples/batchPoissonFDR.slr @@ -0,0 +1,35 @@ +#!/bin/bash +#SBATCH -N 4 -C cpu --exclusive -A m2043 +# SBATCH --time=30:00 -q debug +#SBATCH --time 10:28:00 -q regular +#SBATCH --output=result/%j.out +#SBATCH --licenses=scratch +#SBATCH --image nersc/causal-net:v5 # Nov 5 2025 + +# - - - E N D O F SLURM C O M M A N D S +set -u ; # exit if you try to use an uninitialized variable + +export OMP_NUM_THREADS=2 + +#dataName=daleM20_746c4b +#dataName=daleM40_e33e89 +#dataName=daleM80_285c84 +#dataName=daleM150_448b86 + +# only Jan owns those 2 datasets: +#dataName=daleM300_aa61f5 +dataName=daleM130_6eb245 # hard case + +#samples=10_000 +samples=300_000 + +maxIter=500 + +N=${SLURM_NNODES} +jobId=${SLURM_JOBID} +ntask=$((32 * N)) +echo S: job dataName=$dataName samples=$samples ntask=$ntask maxIter $maxIter + +time srun -n $ntask --distribution=block:block shifter python uoi_var_poisson_fdr_test.py --dataName $dataName --samples $samples --maxIter $maxIter --dataPath /pscratch/sd/b/balewski/2025_causalNet_tmp/ --verb 2 --selectonFrac 0.66 --freqWeight + +# spare --freqWeight diff --git a/examples/uoi_var_mpi_check.py b/examples/uoi_var_mpi_check.py new file mode 100644 index 00000000..15ddff59 --- /dev/null +++ b/examples/uoi_var_mpi_check.py @@ -0,0 +1,178 @@ +import numpy as np + +from pyuoi.linear_model import * +from var_utils import * +import pdb, h5py, os +from mpi4py import MPI +from time import time +from pyuoi.linear_model.sparse_comm_util import build_bootstrap_comm + +# most important hyperparameters! +n_admm =16 +eps = 1e-3 # alpha grid min scaler: default is 1e-3 +n_lambdas = 48 + +# n_process should be multiple of n_admm, and at MOST n_admm*n_boot*n_reg_param +#srun -n 960 --ntasks-per-node=240 --distribution=block:block python uoi_var_mpi_check.py + +#srun -n 512 --ntasks-per-node=128 --distribution=block:block python uoi_var_mpi_check.py + +imbalance_tolerance = 10 + +n_boots_sel=12 +n_boots_est=12 + +max_iter = 400 + +selection_frac = 0.9 + + +use_admm = True + + +rank = 0 +comm = MPI.COMM_WORLD +if comm is not None: + rank = comm.rank + +# suffix = "20k" + +# if rank == 0: +# data = np.load("test_jan_admm_"+suffix+".npy") +# B_truth = np.load("jan_truth.npy") + +# 0.9 is the default bootstrap fraction +#assert(data.shape[0]*0.9/n_admm/data.shape[1] > 8, "Reduce n_admm so the n_samp/n_feat per ADMM process is optimal!") + +seed = 42 +lag = 1 +n_features = 200 +n_samples = 50000 + +if rank == 0: + # Select spectral radius based on lag + rad_dict = {1: 0.98, 2: 0.70, 3: 0.30, 4: 0.20, 5: 0.20} + rad = rad_dict.get(lag, 0.20) + + + generate_new_data = True + + if generate_new_data: + data, transition_matrices, bias_terms, covariance_matrix = generate_sparse_stationary_var_process( + n_features, n_samples, lag=lag, sparsity=0.8, spectral_radius=rad, + process_type='gaussian', random_state=seed + ) + + + with h5py.File('data/test_var_admm.h5', 'w') as f: + g = f.create_group('data') + g.create_dataset(name='data', data=data, compression="gzip") + g.create_dataset(name='transition_matrices', data=transition_matrices[0].toarray(), compression="gzip") + + B_truth = transition_matrices[0].toarray() + np.save("result/var_truth.npy", B_truth) + np.save("result/var_bias_truth.npy", bias_terms) + else: + with h5py.File('data/test_var_admm.h5', 'r') as f: + data = np.copy(f['data/data'][()]) + B_truth = np.copy(f['data/transition_matrices'][()]) + + + + # OBSOLETE: used for older version of VAR fit + # dense_matrices = [M.toarray() for M in transition_matrices] + # vecortized the ground truth transition matrices + # B_truth = np.vstack([m.T for m in dense_matrices]).T.flatten() + # X,Y = vectorization(data, lag) + +rho_list = 1.1+np.arange(4,5, dtype = int)/10 + +#fitting with multiple processes +for rho_scaler in rho_list: + if use_admm: + boot_comm = build_bootstrap_comm(comm, n_admm) + uoi_lasso = UoI_Lasso(fit_VAR = True, fit_intercept=True, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, selection_frac = selection_frac, n_lambdas = n_lambdas, max_iter = max_iter, eps = eps, random_state=seed, comm = boot_comm, global_comm = comm, n_admm = n_admm, rho_scaler = rho_scaler, imbalance_tolerance = imbalance_tolerance, solver='admm', estimation_solver = "ls") + + start = time() + if boot_comm is not None: #if the global_rank is part of the boostrap distribution(not admm distribution) + if boot_comm.rank == 0: + uoi_lasso.fit(lag, data = data) + else: + uoi_lasso.fit(lag) + else: + if uoi_lasso.solver == "admm": + uoi_lasso.admm_queue() + end = time() + + else: + uoi_lasso = UoI_Lasso(fit_VAR = True, fit_intercept=False, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, selection_frac = selection_frac, n_lambdas = n_lambdas, max_iter = max_iter, random_state=seed, comm = comm) + + start = time() + + if comm.rank == 0: + uoi_lasso.fit(lag, data = data) + else: + uoi_lasso.fit(lag) + + end = time() + + + + if rank == 0: + print("rank " +str(comm.rank)+": Fitting complete in "+str(end - start)+" seconds.", flush = True) + + #print(np.count_nonzero(uoi_lasso.VAR_coef_)) + if lag == 1: + + B_model = uoi_lasso.VAR_coef_[0] + model_bias = uoi_lasso.VAR_bias_ + + np.save("result/var_model_"+str(rho_scaler)+".npy", B_model) + np.save("result/var_model_bias_"+str(rho_scaler)+".npy", model_bias) + + + TP, FP, TN, FN = matrix_comparison(B_truth, B_model, threshold=0) + # these two adds up == real sparsity in B_truth + print("TP: ", TP) + print("FN: ", FN) + + print("TN: ", TN) + print("FP: ", FP) + + print("Selection accuracy: ", selection_accuracy(B_truth, B_model)) + print("Bias sparsity: ", np.count_nonzero(model_bias)/model_bias.shape[0]) + + # estimation error + est_mask = B_model != 0 + print("Estimation error: ", np.linalg.norm(B_truth*est_mask - B_model)**2) + + print(str(rho_scaler)+" solution sparsity: ", np.count_nonzero(B_model)/B_model.shape[0]**2) + print(uoi_lasso.intercept_) + + + +# indicator for whether comparing fitting results from single process fit vs multi-process fit +compare_to_single = False + +if compare_to_single: + #fitting with single process + if rank == 0: + uoi_lasso = UoI_Lasso(fit_VAR = True, fit_intercept = False, random_state=seed) + uoi_lasso.fit(lag, data = data) + B_model_single = uoi_lasso.VAR_coef_[0] + + + print("Selection accuracy: ", selection_accuracy(B_truth, B_model_single)) + # estimation error + est_mask = B_model != 0 + print("Estimation error: ", np.linalg.norm(B_truth*est_mask - B_model_single)**2) + + + print(np.allclose(B_model_single, B_model)) + + print( np.linalg.norm(B_model_single-B_model)) + + + + + diff --git a/examples/uoi_var_poisson_fdr_test.py b/examples/uoi_var_poisson_fdr_test.py new file mode 100755 index 00000000..6a00ed2e --- /dev/null +++ b/examples/uoi_var_poisson_fdr_test.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 + +''' + +IMG=nersc/causal-net:v4 # May 13 +export OMP_NUM_THREADS=2 +salloc -q interactive -C cpu --image=$IMG -t 4:00:00 -A m2043 -N 4 + +time srun -n 128 --distribution=block:block shifter python uoi_var_poisson_fdr_test.py --samples 10_000 + +Existing samples @ /pscratch/sd/y/yxu2/data: + --dataName daleM20_746c4b 4 min 100k samp @ N=4 + --dataName daleM40_e33e89 9 min 100k samp @ N=4 + --dataName daleM80_285c84 75 min 300k samp @ N=4 + --dataName daleM150_448b86 62 min 300k samp @ N=4 + +Existing samples @ dataPath=/pscratch/sd/b/balewski/2025_causalNet_tmp/ + --dataName daleM300_aa61f5 +# hard case +--dataName daleM130_6eb245 + +# Yao: +n_process should be multiple of n_admm, and at MOST n_admm*n_boot*n_reg_param + Example run command on NERSC interactive compute node session: + srun -n 768 --ntasks-per-node=192 --distribution=block:block python uoi_var_poisson_test.py + + +UoI_Poisson parameters: + dt = 0.01 - sample time bin size for Poisson process + n_admm = 32 - number of ADMM processes + n_lambdas = 4 - number of L1 penalty values to test + manual_l1_range = [6e-7, 3e-6] - Hardcoded L1 penalty range + ??? imbalance_tolerance = 10 - tolerance for load imbalance + n_boots_sel = 6 - number of bootstrap samples for selection + n_boots_est = 6 - number of bootstrap samples for estimation + max_iter = 1000 - maximum number of iterations + selection_frac = 0.9 - Fraction of total data used for each bootstrap + lag = 1 - VAR lag order, must be 1 + rho_scaler = 1.0 - ADMM rho scaling factor + seed = 22 - random state seed + fdr_rate=0.05 - False Discovery Rate (FDR) , p-value for selection of existing edges +''' + +import os,sys +import numpy as np +import scipy.sparse as sparse +from numpy.linalg import norm +#import importlib +import argparse +import secrets + +from mpi4py import MPI +from time import time +from pprint import pprint +sys.path.append("/global/homes/b/balewski/prjs/2025_UoI-VAR/") +from examples.var_utils import * +from src.pyuoi.linear_model import * +sys.path.append("/global/homes/b/balewski/prjs/2025_UoI-VAR/src/pyuoi/linear_model") +from sparse_comm_util import build_bootstrap_comm +from var_utils import * +from Util_poissonFdr import qa_Bfit, qa_Afit, plot_edges_correl, plot_eigenvalues, plot_loss, extract_loss_traces + +from Util_NumpyIO import read_data_npz, write_data_npz + +#...!...!.................. +def main(): + parser = argparse.ArgumentParser(description='UoI-VAR Poisson ADMM test') + parser.add_argument('--dataPath', default='/pscratch/sd/y/yxu2/data', help='path to data directory') + parser.add_argument('--dataName', default='daleM20_746c4b', help='dataset name (e.g., daleM20_746c4b)') + parser.add_argument('--out', default='result', help='output directory') + parser.add_argument('--samples', type=int, default=100000, help='number of data samples to use') + parser.add_argument('--outName', default=None, help='output file name core') + parser.add_argument('--freqWeight', action='store_true', help='use frequency dependent weights, default is False') + parser.add_argument('--maxIter', type=int, default=1000, help='maximum number of iterations for ADMM') + parser.add_argument('--fdrRate', type=float, default=0.01, help='False discovery rate level for support selection') + parser.add_argument('--selectonFrac', type=float, default=0.9, help='fraction of data used per bootstrap selection') + parser.add_argument('--verb', '-v', type=int, default=1, help='Verbosity level') + args = parser.parse_args() + + hash_str = None + outName = args.outName + if outName is None: + hash_str = secrets.token_hex(3) + outName = f'{args.dataName}_uoi{hash_str}' + + # most important hyperparameters! + confUoI = { + 'fit_VAR': True, + 'fit_intercept': False, + 'standardize': False, + 'manual_l1_range': [6e-7, 3e-6], + 'n_lambdas': 4, + 'n_boots_sel': 6, + 'n_boots_est': 6, + 'selection_frac': args.selectonFrac, + 'max_iter': args.maxIter, + 'random_state': 22, + 'rho_scaler': 1.0, + 'n_admm': 32, + 'imbalance_tolerance': 10, + 'solver': 'admm', + 'estimation_solver': "lbfgs", + 'dt': 0.01, # must be hardcoded - or use broadcasting to all ranks + 'loss_stride': 10, + 'fdr_rate': args.fdrRate, + } + + confMisc = { + 'model_lag': 1, + 'use_freq_weight': args.freqWeight, + 'num_samples': args.samples, + 'data_name': args.dataName, + 'uoi_hash': hash_str, + 'short_name': outName, + 'data_path': args.dataPath, + 'out_path': args.out, + } + + loss_stride = confUoI.get('loss_stride', 10) + + lag = confMisc['model_lag'] + assert lag==1 + + rank = 0 + comm = MPI.COMM_WORLD + world_size = comm.Get_size() + if comm is not None: + rank = comm.rank + + if rank == 0: + for arg in vars(args): + print( 'myArgs:',arg, getattr(args, arg)) + if args.verb > 1: + print('confUoI:'); pprint(confUoI) + print('confMisc:'); pprint(confMisc) + + print('Start dataName=%s samples=%d'%(confMisc['data_name'],confMisc['num_samples'])) + spikeF=os.path.join(confMisc['data_path'],f"{confMisc['data_name']}.spikes.npz") + #data = np.load(spikeF)['spikes'][:confMisc['num_samples']].astype(np.double) + spikeD, spikeMD = read_data_npz(spikeF, verb=True) + data=spikeD['spikes'][:confMisc['num_samples']].astype(np.double) + data_pois = None + pprint(spikeMD) + #confUoI['dt']=spikeMD['time_step_sec'] # can't do it now w/o broadcasting to all ranks + + if confMisc['use_freq_weight']: # enable freq-weighings + rates=np.mean(data,axis = 0)/confUoI['dt'] + rates = np.clip(rates, 0.1, 50) + w=1/rates + w/=np.sum(w) + w*=data.shape[1] + print('rates:',rates) + print('w',w) + else: + w = np.ones(data.shape[1]) + else: + w = None + + w = comm.bcast(w, root=0) + if rank == 0: print('%dk samples loaded to all %d ranks'%(confMisc['num_samples']//1000,world_size ),flush=True) + + #fitting with multiple processes + boot_comm = build_bootstrap_comm(comm, confUoI['n_admm']) + confUoI_fit = confUoI.copy() + confUoI_fit.pop('loss_stride', None) + + uoi_poisson = UoI_Poisson(**confUoI_fit, comm=boot_comm, global_comm=comm, weights=w) + assert uoi_poisson.solver == "admm" + + start = time() + if boot_comm is not None: #if the global_rank is part of the boostrap distribution(not admm distribution) + if boot_comm.rank == 0: + uoi_poisson.fit(lag, data = data, data_pois = data_pois) + else: + uoi_poisson.fit(lag) + else: + uoi_poisson.admm_queue() + end = time() + + if rank > 0: return + + print("rank " +str(comm.rank)+": Fitting complete in %d seconds."%(end - start), flush = True) + + A_fit = uoi_poisson.VAR_coef_[0] + B_fit = uoi_poisson.VAR_bias_ + + sel_loss_iter, sel_loss_l1 = extract_loss_traces(getattr(uoi_poisson, '_selection_lm', None), loss_stride) + est_loss_iter, est_loss_l1 = extract_loss_traces(getattr(uoi_poisson, '_estimation_lm', None), loss_stride) + + l1_loss_sel = np.vstack((sel_loss_iter, sel_loss_l1)) if sel_loss_iter.size > 0 else np.empty((2, 0)) + l1_loss_est = np.vstack((est_loss_iter, est_loss_l1)) if est_loss_iter.size > 0 else np.empty((2, 0)) + + + bigD={ 'A_fit':A_fit, 'B_fit':B_fit, 'l1_loss_sel':l1_loss_sel, 'l1_loss_est':l1_loss_est} + outF = os.path.join(confMisc['out_path'], f"{confMisc['short_name']}.uoiFdr.npz") + write_data_npz(bigD, outF, metaD=spikeMD) + print('saved output to:',outF) + + #.... evaluation of results + truthF=spikeF.replace('.spikes','.simTruth') + A_truth = np.load(truthF)["A_true"] + B_truth = np.load(truthF)["B_true"] + + TP, FP, TN, FN = matrix_comparison(A_truth, A_fit, threshold=0) + # these two adds up == real sparsity in B_truth + M=A_truth.shape[0]; M2=M*M + print('dataName=%s M=%d M^2=%d'%(confMisc['data_name'],M,M2)) + print("TP p=%.3e n=%d "%(TP,TP*M2)) + print("FN p=%.3e n=%d "%(FN,FN*M2)) + print("FP p=%.3e n=%d "%(FP,FP*M2)) + print("TN: ", TN) + + + print('detailed QA %s M=%d M^2=%d samples=%d/k' %(confMisc['data_name'],M,M2,confMisc['num_samples']/1000)) + qaD=qa_Afit(A_truth, A_fit) + qaD['bterm']=qa_Bfit(B_truth, B_fit) + + plot_loss(args, A_truth, A_fit, confMisc['short_name'], l1_loss_sel, loss_skip=50) + plot_edges_correl(args, qaD, confMisc['short_name']) + plot_eigenvalues(args, A_truth, A_fit,confMisc['short_name']) + + +#...!...!.................. +if __name__ == "__main__": + main() + + + + + + + + + + diff --git a/examples/uoi_var_poisson_test.py b/examples/uoi_var_poisson_test.py new file mode 100644 index 00000000..abf5ad2c --- /dev/null +++ b/examples/uoi_var_poisson_test.py @@ -0,0 +1,211 @@ +# Division is default behavior in Python 3, so this import is no longer needed +# from __future__ import division +import pdb, h5py, os +import numpy as np +import matplotlib.pyplot as plt +import scipy.sparse as sparse +from numpy.linalg import norm +import importlib +from sparse_randn import sprandn +from mpi4py import MPI +from time import time + +from sklearn.linear_model import LinearRegression + +#from pyuoi.linear_model.admm_mpi_poisson import * + + + +from pyuoi.linear_model import * +from var_utils import * +from pyuoi.linear_model.sparse_comm_util import build_bootstrap_comm + +# most important hyperparameters! +n_admm = 32 +n_lambdas = 4 +dt = 0.01 # sample time bin size for Poisson process +manual_l1_range = [6e-7, 3e-6] # Hardcoded L1 penalty range +fdr_rate = 0.05 + +eps = 1e-9 # Obsolete if using manual_l1_rage. alpha grid min scaler: default is 1e-3 +stability_selection=0.75 #Obsolete, this was for the occurence based parameter support seleciton method + + +# n_process should be multiple of n_admm, and at MOST n_admm*n_boot*n_reg_param +# Example run command on NERSC interactive compute node session: +# srun -n 768 --ntasks-per-node=192 --distribution=block:block python uoi_var_poisson_test.py + +imbalance_tolerance = 10 + +n_boots_sel=6 +n_boots_est=6 + +max_iter = 1000 +seed = 22 +selection_frac = 0.9 #Fraction of total data used for each bootstrap: default is 0.9 + +use_admm = True +generate_new_data = False + +# Example usage +# n_features = 20 +# n_samples = 20000 +lag = 1 + +rank = 0 +comm = MPI.COMM_WORLD +if comm is not None: + rank = comm.rank + + +# if rank == 0: + + +# if generate_new_data: + + +# base_intensity = np.random.randint(1, 10, n_features).astype(np.float64) # Base rates for each feature + +# data, transition_matrices, bias_terms, _ = generate_sparse_stationary_var_process( +# n_features, +# n_samples, +# lag=lag, +# sparsity=0.8, +# quench_factor = 1, +# spectral_radius=0.9, # Ensures stationarity +# process_type='poisson', +# base_intensity=base_intensity + +# ) + +# B_truth = transition_matrices[0].toarray() +# # np.save("result/poisson_truth.npy", B_truth) +# # np.save("result/poisson_bias_truth.npy", bias_terms) +# # np.save("data/test_poisson_admm.npy", data) + +# else: +# B_truth = np.load("data/poisson_truth.npy", allow_pickle=True) +# data = np.load("data/test_poisson_admm.npy") + +# data_pois = None + +# if rank == 0: +# with h5py.File('data/daleM40-a03d69a-c1edd9d.spike.h5', 'r') as f: +# #spike = np.copy(f['spikes_data'][()]) +# data = np.copy(f['stateVec_data'][()]).T[:15000].astype("float32") +# B_truth = np.copy(f['true_network_matrix'][()]).astype("float32") + +# data_pois = np.load("data/data_pois.npy").T[:15000] + + + + + +if rank == 0: + + # data = np.load('data/daleM120_12dff3.spikes.npz')['spikes'][:1000000].astype(np.double) + # B_truth = np.load('data/daleM120_12dff3.truth.npz')["A_true"] + + # M20_746c4b + # M40_e33e89 + # M80_285c84 + # M150_448b86 + + data = np.load('/pscratch/sd/y/yxu2/data/daleM80_285c84.spikes.npz')['spikes'][:100000].astype(np.double) + B_truth = np.load('/pscratch/sd/y/yxu2/data/daleM80_285c84.simTruth.npz')["A_true"] + data_pois = None + + # w = 1/np.maximum(np.mean(data,axis = 0), 0.1*np.ones(data.shape[1])) + # w = w/np.linalg.norm(w) * data.shape[1] + w = np.ones(data.shape[1]) +else: + w = None + +w = comm.bcast(w, root=0) + + + +rho_list = 1.5+np.arange(0,1, dtype = int)/10 + +#np.arange(3.759,3.761,0.0005) + +for l1_suppression in np.arange(1):#(14,15): + #fitting with multiple processes + for rho_scaler in rho_list: + if use_admm: + boot_comm = build_bootstrap_comm(comm, n_admm) + uoi_poisson = UoI_Poisson(fit_VAR = True, fit_intercept=False, standardize = False, manual_l1_range = manual_l1_range, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, selection_frac = selection_frac, stability_selection=stability_selection, n_lambdas = n_lambdas, max_iter = max_iter, eps = eps, random_state=seed, comm = boot_comm, global_comm = comm, n_admm = n_admm, rho_scaler = rho_scaler, imbalance_tolerance = imbalance_tolerance, l1_suppression= l1_suppression, solver='admm', estimation_solver = "lbfgs", weights = w, dt = dt, fdr_rate = fdr_rate) + + start = time() + if boot_comm is not None: #if the global_rank is part of the boostrap distribution(not admm distribution) + if boot_comm.rank == 0: + + uoi_poisson.fit(lag, data = data, data_pois = data_pois) + else: + uoi_poisson.fit(lag) + else: + if uoi_poisson.solver == "admm": + uoi_poisson.admm_queue() + end = time() + else: + uoi_poisson = UoI_Poisson(fit_VAR = True, fit_intercept=False, standardize = False, manual_l1_range = manual_l1_range, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, selection_frac = selection_frac, stability_selection=stability_selection, n_lambdas = n_lambdas, max_iter = max_iter, eps = eps, random_state=seed, comm = comm, rho_scaler = rho_scaler, imbalance_tolerance = imbalance_tolerance, l1_suppression= l1_suppression, weights = w, dt = dt, fdr_rate = fdr_rate) + + start = time() + + if comm.rank == 0: + uoi_poisson.fit(lag, data = data, data_pois = data_pois) + else: + uoi_poisson.fit(lag) + + end = time() + + + if rank == 0: + print("l1_suppression: ", l1_suppression, flush = True) + print("rank " +str(comm.rank)+": Fitting complete in "+str(end - start)+" seconds.", flush = True) + + #print(np.count_nonzero(uoi_poisson.VAR_coef_)) + if lag == 1: + B_model = uoi_poisson.VAR_coef_[0] + model_bias = uoi_poisson.VAR_bias_ + l1_support_count = uoi_poisson.self.l1_support_count + + # print("L1-loss: ", uoi_poisson.loss["l1"]) + + np.save("result/poisson_M80_"+str(rho_scaler)+"_intersect.npy", B_model) + np.save("result/poisson_M80_bias_"+str(rho_scaler)+"_intersect.npy", model_bias) + + TP, FP, TN, FN = matrix_comparison(B_truth, B_model, threshold=0) + # these two adds up == real sparsity in B_truth + print("TP: ", TP) + print("FN: ", FN) + + print("TN: ", TN) + print("FP: ", FP) + + # estimation error + est_mask = B_model != 0 + print("Estimation error: ", np.linalg.norm(B_truth*est_mask - B_model)**2) + + print(str(rho_scaler)+" solution sparsity: ", np.count_nonzero(B_model)/B_model.shape[0]**2) + + print("Bias sparsity: ", np.count_nonzero(model_bias)/model_bias.shape[0]) + + + + + + + + + + + + + + + + + + + diff --git a/examples/var_utils.py b/examples/var_utils.py new file mode 100644 index 00000000..3958b2ce --- /dev/null +++ b/examples/var_utils.py @@ -0,0 +1,532 @@ +""" +.. _uoi_lasso: + +UoI-Lasso for sparse, minimal bias, regression +=============================r[i================ + +This example with demonstrate the ability of UoI-Lasso to recover sparse +models with minimal bias. + +""" + +############################################################################### +# Load synthetic data +# ------------------- +# +# The synthetic data will have 40 features, 10 of which are informative and +# 1 response variable. + + +import numpy as np + +from sklearn.linear_model import LinearRegression, LassoCV + +#import pandas as pd +from scipy.linalg import solve_discrete_lyapunov +from scipy.special import expit # logistic function +from sklearn.metrics import precision_score, recall_score, f1_score + +from scipy import sparse + +def vectorization(data, lag): + # vectorize the VAR data for use with LASSO algorithm + # data: time series data in np array with shape n_samples X n_features + + # flipup so the last time sample in data is now first row + data = np.flipud(data) + n_samples = data.shape[0] + n_features = data.shape[1] + + # shape of Y: (T - D) X (n_features) *T - D: total number of samples - lag + Y = data[: n_samples-lag] + # ones = np.ones((n_samples-lag,1)) # for VAR porocess with intercept + Y = Y.T.flatten() + + # X.shape: (n_samples - lag) X (lag * n_features); + X_row = [np.hstack(data[i : lag+i]) for i in range(1,n_samples-lag+1)] + X = np.vstack(X_row) + # use kronecker product for vectorization of matrix multiplication + X = np.kron(np.eye(n_features), X) + + return X, Y + +def vectorization_mbb(data, n_boots): + # OBSOLETE for moving block boostraps: X_mbb.shape: n_boots X (lag * n_features); + + data = np.flipud(data) + Y = data[: n_samples-lag] + + data_mbb, boot_idx = mbbs_np(data, n_blocks=n_boots, block_length = lag) + + X_mbb_row = [np.hstack(data_mbb[lag*i : lag*(i+1)]) for i in range(n_boots)] + X_mbb = np.vstack(X_mbb_row) + # effectiver number of features is (lag * n_feat * n_feat), i.e. total number of entries in all transition matrices + X_mbb = np.kron(np.eye(n_features), X_mbb) + + # moving block bootstrap samples for Y: + Y_mbb = Y[boot_idx] + Y_mbb = Y_mbb.T.flatten() #vectorization + + return X_mbb, Y_mbb + + +def mbbs_df(data, n_blocks=10, block_length=None): + # moving block bootstrap for data in panda dataframe format + n = len(data) + + # assign a default value for block_length + if block_length is None: + block_length = n // n_blocks + + # Create empty DataFrame with correct size and columns + total_rows = n_blocks * block_length + mbbs_sample = pd.DataFrame(columns=data.columns, index=range(total_rows)) + + # generate starting indices for each block + block_indices = np.random.randint(low=0, high=n - block_length, size=n_blocks) + + # Create the bootstrap sample + for i in range(n_blocks): + start_index = block_indices[i] + end_index = start_index + block_length + mbbs_sample.iloc[i * block_length : (i + 1) * block_length] = data.iloc[start_index:end_index].values + + return mbbs_sample + +def mbbs_np(data, n_blocks=10, block_length=None): + # # moving block bootstrap for data in numpy array format + if not isinstance(data, np.ndarray): + raise TypeError("Input data must be a numpy array") + + # number of datapoints + n = len(data) + + # assign a default value for block_length + if block_length is None: + block_length = n // n_blocks + + # Create empty array with correct size + total_rows = n_blocks * block_length + mbbs_sample = np.empty((total_rows,) + data.shape[1:]) + + # generate starting indices for each block + block_indices = np.random.randint(low=0, high=n - block_length, size=n_blocks) + + # Create the bootstrap sample + for i in range(n_blocks): + start_index = block_indices[i] + end_index = start_index + block_length + mbbs_sample[i * block_length : (i + 1) * block_length] = data[start_index:end_index] + + return mbbs_sample, block_indices + + + + +def generate_sparse_stationary_var_process(n_features, n_samples, lag=1, sparsity=0.9, spectral_radius=0.8, quench_factor=1, + process_type='gaussian', random_state=None, + base_intensity=None, base_probability=None, + include_bias=True, bias_scale=1.0): + """ + Generate data from a stationary Vector Autoregressive (VAR) process with sparse transition matrices and bias terms. + + Parameters: + ----------- + n_features : int + Number of features (dimensions) in the VAR process + n_samples : int + Number of time points to generate + lag : int, default=1 + Order of the VAR process (number of past observations to use) + sparsity : float, default=0.9 + Desired sparsity level (proportion of zero elements) in transition matrices + spectral_radius : float, default=0.8 + Desired spectral radius of the VAR process (must be < 1 for stationarity) + quench_factor : float, default=1 + Additional scaling factor for spectral radius control + process_type : str, default='gaussian' + Type of VAR process. Either 'gaussian', 'poisson', or 'bernoulli' + random_state : int or None, default=None + Random seed for reproducibility + base_intensity : ndarray or None, default=None + Base intensity for Poisson process. Required if process_type='poisson' + base_probability : ndarray or None, default=None + Base probability for Bernoulli process. Required if process_type='bernoulli' + include_bias : bool, default=True + Whether to include bias terms for each node + bias_scale : float, default=1.0 + Scale factor for bias term generation + + Returns: + -------- + tuple: + - data: ndarray of shape (n_samples, n_features) + - transition_matrices: list of sparse matrices, each of shape (n_features, n_features) + - bias_terms: ndarray of shape (n_features,) - bias term for each node + - covariance_matrix: ndarray of shape (n_features, n_features), only for Gaussian process + """ + if random_state is not None: + np.random.seed(random_state) + + if process_type not in ['gaussian', 'poisson', 'bernoulli']: + raise ValueError("process_type must be either 'gaussian', 'poisson', or 'bernoulli'") + + if not 0 <= sparsity < 1: + raise ValueError("sparsity must be between 0 and 1") + + if not 0 <= spectral_radius < 1: + raise ValueError("spectral_radius must be between 0 and 1 for stationarity") + + if process_type == 'poisson' and base_intensity is None: + raise ValueError("base_intensity must be provided for Poisson process") + + if process_type == 'bernoulli': + if base_probability is None: + raise ValueError("base_probability must be provided for Bernoulli process") + if not np.all((base_probability >= 0) & (base_probability <= 1)): + raise ValueError("base_probability values must be between 0 and 1") + + def get_companion_spectral_radius(matrices): + """Calculate spectral radius of the companion matrix""" + n = matrices[0].shape[0] + p = len(matrices) + + # Convert sparse matrices to dense for companion matrix construction + dense_matrices = [M.toarray() for M in matrices] + companion = np.zeros((n * p, n * p)) + + companion[:n] = np.hstack(dense_matrices) + if p > 1: + companion[n:, :-n] = np.eye(n * (p - 1)) + + eigenvals = np.linalg.eigvals(companion) + return np.max(np.abs(eigenvals)) + + def generate_sparse_matrix(n, sparsity): + """Generate a sparse random matrix with controlled density""" + # Generate mask for non-zero elements + mask = np.random.random((n, n)) > sparsity + + # Generate random values for non-zero elements + values = np.random.randn(n, n) + values = values * mask + + return sparse.csr_matrix(values) + + # Generate bias terms for each node + if include_bias: + if process_type == 'gaussian': + bias_terms = np.random.randn(n_features) * bias_scale + elif process_type == 'poisson': + # For Poisson, base_intensity already serves as the bias term + bias_terms = base_intensity.copy() if base_intensity is not None else np.ones(n_features) + else: # Bernoulli + # For Bernoulli, logit(base_probability) already serves as the bias term + bias_terms = np.log(base_probability / (1 - base_probability)) if base_probability is not None else np.zeros(n_features) + else: + if process_type == 'gaussian': + bias_terms = np.zeros(n_features) + elif process_type == 'poisson': + # Even without explicit bias, base_intensity serves as bias + bias_terms = base_intensity.copy() if base_intensity is not None else np.ones(n_features) + else: # Bernoulli + # Even without explicit bias, logit(base_probability) serves as bias + bias_terms = np.log(base_probability / (1 - base_probability)) if base_probability is not None else np.zeros(n_features) + + # Generate transition matrices with controlled spectral radius + transition_matrices = [] + + # First generate matrices without scaling + for _ in range(lag): + matrix = generate_sparse_matrix(n_features, sparsity) + transition_matrices.append(matrix) + + # Calculate current spectral radius + current_radius = get_companion_spectral_radius(transition_matrices) + + # Scale matrices to achieve desired spectral radius + scaling_factor = spectral_radius / current_radius / quench_factor + transition_matrices = [matrix * scaling_factor for matrix in transition_matrices] + + # Verify stationarity + final_radius = get_companion_spectral_radius(transition_matrices) + assert final_radius < 1, "Failed to achieve stationarity" + + # Process-specific scaling + process_scale_factors = { + 'gaussian': 1.0, + 'poisson': 0.6, # More conservative for count data + 'bernoulli': 0.4 # Even more conservative for binary data + } + scale_factor = process_scale_factors[process_type] + transition_matrices = [matrix * scale_factor for matrix in transition_matrices] + + # Initialize process-specific parameters and initial data + if process_type == 'gaussian': + # Generate innovation covariance matrix (positive definite) + A = np.random.randn(n_features, n_features) + covariance_matrix = A @ A.T + np.eye(n_features) + + # For stationary initialization with bias, we need to solve: + # mu = (I - A1 - A2 - ... - Ap)^(-1} * bias + A_total = sum(M.toarray() for M in transition_matrices) + stationary_mean = np.linalg.solve(np.eye(n_features) - A_total, bias_terms) + + # Calculate stationary covariance + stationary_cov = solve_discrete_lyapunov(A_total, covariance_matrix) + + # Initialize data with stationary distribution + data = np.random.multivariate_normal( + mean=stationary_mean, + cov=stationary_cov, + size=lag + ) + elif process_type == 'poisson': + covariance_matrix = None + # base_intensity serves as the bias term for Poisson + data = np.random.poisson(lam=np.maximum(base_intensity, 0.1), size=(lag, n_features)) + else: # Bernoulli + covariance_matrix = None + # logit(base_probability) serves as the bias term + logit_prob = np.log(base_probability / (1 - base_probability)) + initial_prob = expit(logit_prob) + data = np.random.binomial(n=1, p=initial_prob, size=(lag, n_features)) + + # Generate the rest of the time series + for t in range(lag, n_samples): + if process_type == 'gaussian': + # Start with bias term + new_point = bias_terms.copy() + elif process_type == 'poisson': + # base_intensity serves as the bias term + new_point = base_intensity.copy() + else: # Bernoulli + # logit(base_probability) serves as the bias term + new_point = np.zeros(n_features) + + # Add contribution from each lag using sparse matrix multiplication + for i in range(lag): + new_point += transition_matrices[i].dot(data[t - i - 1]) + + if process_type == 'gaussian': + new_point += np.random.multivariate_normal( + mean=np.zeros(n_features), + cov=covariance_matrix + ) + elif process_type == 'poisson': + # Ensure non-negative intensities + new_point = np.maximum(new_point, 0.1) + new_point = np.random.poisson(lam=new_point) + else: # Bernoulli + # Add logit(base_probability) which serves as the bias term + new_point += np.log(base_probability / (1 - base_probability)) + probabilities = expit(new_point) + new_point = np.random.binomial(n=1, p=probabilities) + + data = np.vstack([data, new_point]) + + return data, transition_matrices, bias_terms, covariance_matrix + + +def print_process_info(data, transition_matrices, covariance_matrix): + """ + Print information about the generated VAR process. + + Parameters: + ----------- + data : ndarray + The generated time series data + transition_matrices : list of ndarrays + The transition matrices of the VAR process + covariance_matrix : ndarray + The innovation covariance matrix + """ + print(f"Data shape: {data.shape}") + print("\nTransition matrices:") + for i, matrix in enumerate(transition_matrices): + print(f"\nLag {i+1}:") + print(matrix) + print("\nInnovation covariance matrix:") + print(covariance_matrix) + + # Print some basic statistics + print("\nData statistics:") + print(f"Mean:\n{np.mean(data, axis=0)}") + print(f"\nEmpirical covariance:\n{np.cov(data.T)}") + + + + + + +def evaluate_sparse_var_estimation(true_matrices, estimated_matrices, threshold=1e-5): + """ + Compute various metrics for evaluating sparse VAR estimation accuracy. + + Parameters: + ----------- + true_matrices : list of sparse matrices or ndarrays + True transition matrices + estimated_matrices : list of sparse matrices or ndarrays + Estimated transition matrices + threshold : float, default=1e-5 + Threshold for considering an element as non-zero + + Returns: + -------- + dict: + Dictionary containing various evaluation metrics + """ + metrics = {} + + # Convert to dense if needed for calculations + true_dense = [M.toarray() if sparse.issparse(M) else M for M in true_matrices] + est_dense = [M.toarray() if sparse.issparse(M) else M for M in estimated_matrices] + + # Combine all matrices for overall metrics + true_combined = np.concatenate([M.flatten() for M in true_dense]) + est_combined = np.concatenate([M.flatten() for M in est_dense]) + + # Create binary masks for sparsity pattern + true_mask = np.abs(true_combined) > threshold + est_mask = np.abs(est_combined) > threshold + + # 1. Sparsity Pattern Metrics + metrics['precision'] = precision_score(true_mask, est_mask) + metrics['recall'] = recall_score(true_mask, est_mask) #The recall is intuitively the ability of the classifier to find all the positive samples. + metrics['f1'] = f1_score(true_mask, est_mask) + + # 2. Relative Frobenius Error +# frob_errors = [] +# for true_M, est_M in zip(true_dense, est_dense): +# frob_error = np.linalg.norm(true_M - est_M, 'fro') / np.linalg.norm(true_M, 'fro') +# frob_errors.append(frob_error) +# metrics['relative_frobenius'] = np.mean(frob_errors) + + # 3. Support Recovery Error + support_errors = [] + for true_M, est_M in zip(true_dense, est_dense): + true_supp = np.abs(true_M) > threshold + est_supp = np.abs(est_M) > threshold + support_error = np.sum(true_supp != est_supp) / true_M.size + support_errors.append(support_error) + metrics['support_error'] = np.mean(support_errors) + + # 4. Element-wise Metrics for Non-zero Elements + true_nonzero = true_combined[true_mask] + est_nonzero = est_combined[true_mask] + + if len(true_nonzero) > 0: + metrics['mae_nonzero'] = np.mean(np.abs(true_nonzero - est_nonzero)) + metrics['rmse_nonzero'] = np.sqrt(np.mean((true_nonzero - est_nonzero)**2)) + metrics['mape_nonzero'] = np.mean(np.abs((true_nonzero - est_nonzero) / true_nonzero)) * 100 + + # 5. Sparsity Level Comparison + metrics['true_sparsity'] = 1 - np.mean(true_mask) + metrics['estimated_sparsity'] = 1 - np.mean(est_mask) + + return metrics + +def visualize_matrix_comparison(true_matrix, estimated_matrix, threshold=1e-5): + """ + Create visualization comparing true and estimated matrices. + Returns plotting data that can be used with your preferred visualization library. + """ + # Convert to dense if sparse + true_dense = true_matrix.toarray() if sparse.issparse(true_matrix) else true_matrix + est_dense = estimated_matrix.toarray() if sparse.issparse(estimated_matrix) else estimated_matrix + + # Create masks for zero/nonzero elements + true_mask = np.abs(true_dense) > threshold + est_mask = np.abs(est_dense) > threshold + + # Categorize elements + comparison = { + 'TP': np.logical_and(true_mask, est_mask), + 'FP': np.logical_and(~true_mask, est_mask), + 'FN': np.logical_and(true_mask, ~est_mask), + 'TN': np.logical_and(~true_mask, ~est_mask), + 'error_magnitude': np.abs(true_dense - est_dense) + } + + return comparison + +# Example usage: +def example_usage(): + # Generate example matrices + n_features = 10 + true_matrices = [ + sparse.random(n_features, n_features, density=0.1).toarray() * 0.5 + for _ in range(2) # VAR(2) process + ] + + # Add some noise to create estimated matrices + estimated_matrices = [ + M + np.random.normal(0, 0.1, M.shape) + for M in true_matrices + ] + + # Compute metrics + metrics = evaluate_sparse_var_estimation(true_matrices, estimated_matrices) + + # Print results + print("\nEvaluation Metrics:") + print("------------------") + for metric, value in metrics.items(): + print(f"{metric}: {value:.4f}") + + return metrics + + +def selection_accuracy(B_truth, B_model): + comparison = visualize_matrix_comparison(B_truth, B_model, threshold=0) + FN = np.count_nonzero(comparison["FN"]) + FP = np.count_nonzero(comparison["FP"]) + M = np.count_nonzero(B_truth) + return 1-(FN+FP)/(M+FP) + + + +#stable transition matrices check: +def stability_check(dense_matrices): + n_features = dense_matrices[0].shape[0] + for _ in range(1000): + mat = np.eye(n_features) + for A in dense_matrices: + z = np.random.rand(n_features) + z = z/(np.linalg.norm(z)+10) + mat-=A@z + assert(np.linalg.det(mat)!=0) + + +def matrix_comparison(true_matrix, estimated_matrix, threshold=1e-5): + """ + Create visualization comparing true and estimated matrices. + Returns plotting data that can be used with your preferred visualization library. + """ + # Convert to dense if sparse + true_dense = true_matrix.toarray() if sparse.issparse(true_matrix) else true_matrix + est_dense = estimated_matrix.toarray() if sparse.issparse(estimated_matrix) else estimated_matrix + + # Create masks for zero/nonzero elements + true_mask = np.abs(true_dense) > threshold + est_mask = np.abs(est_dense) > threshold + + M = true_matrix.shape[0]**2 + + # Categorize elements + comparison = { + 'TP': np.logical_and(true_mask, est_mask), + 'FP': np.logical_and(~true_mask, est_mask), + 'FN': np.logical_and(true_mask, ~est_mask), + 'TN': np.logical_and(~true_mask, ~est_mask), + 'error_magnitude': np.abs(true_dense - est_dense) + } + + TP = np.count_nonzero(comparison["TP"])/M + FP = np.count_nonzero(comparison["FP"])/M + TN = np.count_nonzero(comparison["TN"])/M + FN = np.count_nonzero(comparison["FN"])/M + + return TP, FP, TN, FN + diff --git a/src/pyuoi/lbfgs/_lowlevel.pyx b/src/pyuoi/lbfgs/_lowlevel.pyx index 7afad383..918c364a 100644 --- a/src/pyuoi/lbfgs/_lowlevel.pyx +++ b/src/pyuoi/lbfgs/_lowlevel.pyx @@ -384,7 +384,7 @@ cdef class LBFGS(object): raise TypeError("progress must be callable, got %s" % type(f)) x0 = np.atleast_1d(x0) - n = np.product(x0.shape) + n = np.prod(x0.shape) cdef np.npy_intp tshape[1] tshape[0] = n diff --git a/src/pyuoi/linear_model/admm_mpi.py b/src/pyuoi/linear_model/admm_mpi.py new file mode 100644 index 00000000..a96942bf --- /dev/null +++ b/src/pyuoi/linear_model/admm_mpi.py @@ -0,0 +1,495 @@ +from __future__ import division +import pdb, time, h5py, os +import numpy as np +import scipy.sparse as sparse +from scipy.sparse.linalg import spsolve, splu +from numpy.linalg import norm, cholesky +from mpi4py import MPI +from optparse import OptionParser +import gc +from copy import deepcopy +from .sparse_comm_util import * +from time import time + + + +def adaptive_rho_update_boyd(r_res, s_res, rho, u, rho_scaler, imbalance_tolerance, n = None, p=None): + + if r_res > imbalance_tolerance * s_res: + rho *= rho_scaler + u /= rho_scaler + elif s_res > imbalance_tolerance * r_res: + rho /= rho_scaler + u *= rho_scaler + + return rho , u + +def adaptive_rho_update(r_res, s_res, primal_eps, dual_eps, rho, u, rho_scaler, imbalance_tolerance, n = None, p=None): + + if r_res/primal_eps > imbalance_tolerance * s_res/dual_eps: + rho *= rho_scaler + u /= rho_scaler + elif s_res/dual_eps > imbalance_tolerance * r_res/primal_eps: + rho /= rho_scaler + u *= rho_scaler + + return rho , u + + +def objective(X, y, alpha, x, z): + if alpha == 0: + return .5 * np.square(X.dot(x) - y).sum() + else: + return .5 * np.square(X.dot(x) - y).sum() + alpha * norm(z, 1) + +def factor(X, rho): + m, n = X.shape + if m >= n: + L = cholesky(X.T.dot(X) + rho * sparse.eye(n)) + else: + L = cholesky(sparse.eye(m) + 1. / rho * (X.dot(X.T))) + L = sparse.csc_matrix(L) + U = sparse.csc_matrix(L.T) + return L, U + +def sparse_factor(X, rho): + m, n = X.shape + + if m >= n: + A = X.T.dot(X) + rho * sparse.eye(n) + else: + A = sparse.eye(m) + 1. / rho * (X.dot(X.T)) + + # Use sparse LU(more general than Cholesky) factorization + # Note: scipy returns a factorized object, not the L matrix directly + factor = splu(A.tocsc()) # Convert to CSC for factorization + + return factor + +def soft_threshold(v, k, mask): + if mask is not None: + mask = mask.reshape(-1, 1) + k_masked = mask * k + else: + k_masked = k + + return np.sign(v) * np.maximum(np.abs(v) - k_masked, 0.0) + + +def soft_threshold_old(v, k): + return np.sign(v) * np.maximum(np.abs(v) - k, 0.0) + + +class ADMM_Lasso: + """ + MPI ADMM implementation of Linear Model trained with L1 prior as regularizer (aka the Lasso) + + The optimization objective for Lasso is: + + E(w) = 0.5 * ||y - Xw||^2_2 + alpha * ||w||_1 + + Parameters + ---------- + alpha : float, default=1.0 + Constant that multiplies the L1 term. alpha = 0 is equivalent to + ordinary least squares. For numerical reasons, using alpha = 0 is + not advised; instead, you should use Ridge or LinearRegression. + + fit_intercept : bool, default=True + Whether to calculate the intercept for this model. If set to False, + no intercept will be used in calculations. + + max_iter : int, default=1000 + The maximum number of iterations + + tol : float, default=1e-4 + The tolerance for the optimization. + + warm_start : bool, default=False + When set to True, reuse the solution of the previous call to fit as + initialization, otherwise, just erase the previous solution. + + processes: list + the integer indices of MPI processes that are alotted for the MPI-ADMM fit + + random_state : int, RandomState instance, default=None + The seed of the pseudo random number generator that selects a random + feature to update. + + Attributes + ---------- + coef_ : array, shape (n_features,) | (n_targets, n_features) + Parameter vector (w in the cost function formula). + + intercept_ : float | array, shape (n_targets,) + Independent term in decision function. + + n_iter_ : int | array-like, shape (n_targets,) + Number of iterations run by the coordinate descent solver to reach + the specified tolerance. + """ + + def __init__(self, comm, alpha=None, fit_intercept=False, max_iter=50, + abs_tol=1e-3, rel_tol = 1e-2,rho_scaler = 2, imbalance_tolerance = 10, warm_start=True, random_state=None): + self.alpha = alpha + self.fit_intercept = fit_intercept + self.max_iter = max_iter + self.abs_tol = abs_tol + self.rel_tol = rel_tol + self.warm_start = warm_start + self.random_state = random_state + self.comm = comm + self.coef_ = 0 + self.rho_scaler = rho_scaler + self.imbalance_tolerance = imbalance_tolerance + + def fit(self, X= None, y = None, z = None, rho = None, sparse_input = True, param_mask = None): + """ + Fit model with coordinate descent. + + Parameters + ---------- + X : {array-like, sparse matrix} of shape (n_samples, n_features) + Training data. Only fed into the root rank. + + y : array-like of shape (n_samples,) or (n_samples, n_targets) + Target values. Only fed into the root rank. + + + Returns + ------- + self : object + Returns self. + """ + self.comm_time = 0 + self.compute_time1 = 0 + self.compute_time2 = 0 + self.compute_time3 = 0 + self.total_time = 0 + + max_iter = self.max_iter + abs_tol = self.abs_tol + rel_tol = self.rel_tol + + comm = self.comm + size = comm.Get_size() + rank = comm.Get_rank() + + + N = size + + + ''' + Data + ''' + start_total_time = time() + start_time = time() + if rank == 0: + + m, n = X.shape + + + # scaled global alpha values, + # *m(number of samples in this bootstrap) accounts for the ADMM objective function being the total SSE + # /N(n_admm) accounts for the change in L1-penalty scale when the SSE term optimization is distributed + alpha = self.alpha * m / N + + + else: + n = np.zeros(1).astype('int') + m = np.zeros(1).astype('int') + alpha = np.zeros(1).astype('double') + + + self.terminate_selection = False + self.terminate_estimation = False + + n = comm.bcast(n, root=0) + + # using non-sensible features number as termination signals + if n == -1: + self.terminate_selection = True + return self + elif n == 0: + self.terminate_estimation = True + return self + + + m = comm.bcast(m, root=0) + + + alpha = comm.bcast(alpha, root=0) + + + + if rank != 0: + # X = np.empty((m, n), dtype=np.float64) + y = np.empty(m, dtype=np.float64) + + ''' + Broadcast data + ''' + comm.Barrier() + if rank == 0: + for dest_rank in range(1, size): + + X_segment = X[dest_rank::N, :] + # since it's sparse csr matrix, don't have to Send the segment shape first, send the actual segment + if sparse_input: + send_sparse_matrix(comm, X_segment, dest=dest_rank, tag = 11) + else: + # Send the segment shape first + segment_shape = np.array(X_segment.shape, dtype=np.int32) + comm.Send([segment_shape, MPI.INT], dest=dest_rank, tag=10) + + # Send the actual segment + comm.Send([np.ascontiguousarray(X_segment), MPI.DOUBLE], dest=dest_rank, tag=11) + if sparse_input: + X = X[0::N, :] + else: + X = np.ascontiguousarray(X[0::N, :]) + + gc.collect() + else: + if sparse_input: + X = recv_sparse_matrix(comm, source=0, tag=11) + else: + # First receive the shape of the incoming data + segment_shape = np.empty(2, dtype=np.int32) + comm.Recv([segment_shape, MPI.INT], source=0, tag=10) + + # Create a buffer to receive the segment + X = np.empty(segment_shape, dtype=np.float64) + + # Receive the segment + comm.Recv([X, MPI.DOUBLE], source=0, tag=11) + + + ''' + Select sample block + ''' + + # X = X[rank::N, :] + comm.Barrier() + + m, n = X.shape + comm.Bcast([y, MPI.DOUBLE]) + param_mask = comm.bcast(param_mask, root=0) + + # m is n_samp per ADMM process! + # this is for accomdating the definition of MSE term in ADMM-LASSO convention + #alpha *= m + # good heuristric is to start with rho = l1-penalty + if rho == None: + rho = alpha # admm parameter, modulating the constraint that aux variable equals the model variable + + # do the send-receisve again for y?? or integrate back into the last send-receive operation?? or just Bcast it like right now + y = np.ascontiguousarray(y.ravel()[rank::N].reshape((m, 1))) + + + + # save a matrix-vector multiply + Xty = X.T.dot(y) + + self.comm_time += time()-start_time + + # initialize ADMM solver + + if z is None: + + if self.warm_start and not np.all(self.coef_ == 0): + z = deepcopy(self.coef_) + else: + np.random.seed(self.random_state) + z = np.random.normal(scale=1, size=(n, 1)) + #z = np.zeros((n, 1)) + + x = deepcopy(z) + + + u = np.zeros((n, 1)) + + send = np.zeros(3) + recv = np.zeros(3) + + # cache the (Cholesky) factorization + start_compute_time = time() + if sparse_input: + factor_obj = sparse_factor(X, rho) + else: + L, U = factor(X, rho) + self.compute_time1 += time()-start_compute_time + + + + + objval = [] + r_norm = [] + s_norm = [] + eps_pri = [] + eps_dual = [] + + ''' + ADMM solver loop + ''' + + rho_history = [] + + # using scaled version where u = 1/y + for k in range(max_iter): # xrange -> range for Python 3 + + # u-update + # if k != 0: + # u += (x - z) + + # x-update + + q = Xty + rho * (z - u) # (temporary value) + + if sparse_input: + if m >= n: + x = factor_obj.solve(q) + else: + #print("CAUTION: n_samples < n_features", flush = True) + start_compute_time = time() + ULXq = factor_obj.solve(X.dot(q)) + self.compute_time2 += time()-start_compute_time + x = (q * 1. / rho) - ((X.T.dot(ULXq)) * 1. / (rho**2)) + else: + if m >= n: + x = spsolve(U, spsolve(L, q))[..., np.newaxis] + else: + ULXq = spsolve(U, spsolve(L, X.dot(q)))[..., np.newaxis] + x = (q * 1. / rho) - ((X.T.dot(ULXq)) * 1. / (rho**2)) + + + w = x + u + + zprev = np.copy(z) + + + start_time = time() + + comm.Barrier() + comm.Allreduce([w, MPI.DOUBLE], [z, MPI.DOUBLE]) # the resulting z is sum of N variants, so it has to be divided by N before all usage + self.comm_time += time()-start_time + + start_compute_time = time() + # z-update + if alpha == 0: #Linear regression case + z = z * 1. / N + else: + z = soft_threshold(z * 1. / N, alpha * 1. / (N * rho), mask = param_mask) + + r = x-z + u += r + self.compute_time3 += time()-start_compute_time + + + start_time = time() + send[0] = r.T.dot(r)[0][0] + send[1] = x.T.dot(x)[0][0] + #send[2] = u.T.dot(u)[0][0] / (rho**2) + send[2] = u.T.dot(u)[0][0] * (rho**2) + + + comm.Barrier() + comm.Allreduce([send, MPI.DOUBLE], [recv, MPI.DOUBLE]) + + self.comm_time += time()-start_time + + + r_res = np.sqrt(recv[0]) + s_res = np.sqrt(N) * rho * norm(z - zprev) + + primal_eps = np.sqrt(n * N) * abs_tol + rel_tol * np.maximum(np.sqrt(recv[1]), np.sqrt(N) * norm(z)) + dual_eps = np.sqrt(n * N) * abs_tol + rel_tol * np.sqrt(recv[2]) + + + # diagnostics, reporting, termination checks + objval.append(objective(X, y, alpha, x, z)) + r_norm.append(r_res) + s_norm.append(s_res) + eps_pri.append(primal_eps) + eps_dual.append(dual_eps) + + + if r_res < primal_eps and s_res < dual_eps and k > 0: + break + + + # adaptive rho selection based on residual + #rho, u = adaptive_rho_update_boyd(r_res, s_res, rho, u, self.rho_scaler, self.imbalance_tolerance) + rho, u = adaptive_rho_update(r_res, s_res, primal_eps, dual_eps, rho, u, self.rho_scaler, self.imbalance_tolerance) + + + # Compute residual + # r = x - z + if rank == 0: + rho_history.append(rho) + + # if rank == 0: + # np.save("rho_plot/rho_"+str(self.rho_scaler)+"_20k_160.npy", rho_history) + + + # Set attributes after fitting + self.coef_ = z + self.intercept_ = 0 + self.n_iter_ = k + + self.total_time += time()- start_total_time + + + return self + + def set_params(self, alpha): + self.alpha = alpha + + def predict(self, X): + """ + Predict using the linear model. + + Parameters + ---------- + X : array-like or sparse matrix, shape (n_samples, n_features) + Samples. + + Returns + ------- + C : array, shape (n_samples,) + Returns predicted values. + """ + y_pred = X @ self.coef_.ravel() + self.intercept_ + + return y_pred + + def score(self, X, y, sample_weight=None): + """ + Return the coefficient of determination R^2 of the prediction. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) + Test samples. + + y : array-like of shape (n_samples,) or (n_samples, n_outputs) + True values for X. + + sample_weight : array-like of shape (n_samples,), default=None + Sample weights. + + Returns + ------- + score : float + R^2 of self.predict(X) wrt. y. + """ + # Implementation would go here + + return None + + def _set_intercept(self, X_mean, y_mean, X_std): + """ + Set the intercept based on the fit. + """ + # Implementation would go here + pass diff --git a/src/pyuoi/linear_model/admm_mpi_dale.py b/src/pyuoi/linear_model/admm_mpi_dale.py new file mode 100644 index 00000000..f071995e --- /dev/null +++ b/src/pyuoi/linear_model/admm_mpi_dale.py @@ -0,0 +1,446 @@ +from __future__ import division +import pdb, time, h5py, os +import numpy as np +import scipy.sparse as sparse +from scipy.sparse.linalg import spsolve +from numpy.linalg import norm, cholesky +from mpi4py import MPI +from optparse import OptionParser +import gc +from copy import deepcopy + + + +def adaptive_rho_update(r_res, s_res, primal_eps, dual_eps, rho, u, rho_scaler, imbalance_tolerance, n = None, p=None): + + if r_res/primal_eps > imbalance_tolerance * s_res/dual_eps: + rho *= rho_scaler + u /= rho_scaler + elif s_res/dual_eps > imbalance_tolerance * r_res/primal_eps: + rho /= rho_scaler + u *= rho_scaler + + return rho , u + +def objective(X, y, alpha, x, z): + if alpha == 0: + return .5 * np.square(X.dot(x) - y).sum() + else: + return .5 * np.square(X.dot(x) - y).sum() + alpha * norm(z, 1) + +def factor(X, rho): + m, n = X.shape + if m >= n: + L = cholesky(X.T.dot(X) + rho * sparse.eye(n)) + else: + L = cholesky(sparse.eye(m) + 1. / rho * (X.dot(X.T))) + L = sparse.csc_matrix(L) + U = sparse.csc_matrix(L.T) + return L, U + + +def soft_threshold(v, k): + return np.sign(v) * np.maximum(np.abs(v) - k, 0.0) + + + +def EG_update(z, x_old, x, u, rho, alpha, rank, k): + """ + Parameters + ---------- + z : 1-dim np.array + + x_old : x(k) + + x : x(k+1) + """ + + min_exp = -70 + max_exp = 70 + + x_array = rho /alpha * (x-z+u) * np.sign(x) + + # Clip values to safe range + x_safe = np.clip(x_array, min_exp, max_exp) + + # if rank == 0: + # print(k, rho, np.max(x_safe), flush = True) + z_new = x_old * np.exp(x_safe) + return z_new + + + + +class ADMM_Dale: + """ + MPI ADMM implementation of Linear Model trained with L1 prior as regularizer (aka the Lasso) + + The optimization objective for Lasso is: + + E(w) = 0.5 * ||y - Xw||^2_2 + alpha * ||w||_1 + + Parameters + ---------- + alpha : float, default=1.0 + Constant that multiplies the L1 term. alpha = 0 is equivalent to + ordinary least squares. For numerical reasons, using alpha = 0 is + not advised; instead, you should use Ridge or LinearRegression. + + fit_intercept : bool, default=True + Whether to calculate the intercept for this model. If set to False, + no intercept will be used in calculations. + + max_iter : int, default=1000 + The maximum number of iterations + + tol : float, default=1e-4 + The tolerance for the optimization. + + warm_start : bool, default=False + When set to True, reuse the solution of the previous call to fit as + initialization, otherwise, just erase the previous solution. + + processes: list + the integer indices of MPI processes that are alotted for the MPI-ADMM fit + + random_state : int, RandomState instance, default=None + The seed of the pseudo random number generator that selects a random + feature to update. + + Attributes + ---------- + coef_ : array, shape (n_features,) | (n_targets, n_features) + Parameter vector (w in the cost function formula). + + intercept_ : float | array, shape (n_targets,) + Independent term in decision function. + + n_iter_ : int | array-like, shape (n_targets,) + Number of iterations run by the coordinate descent solver to reach + the specified tolerance. + """ + + def __init__(self, comm, rho = None, alpha=None, fit_intercept=False, max_iter=50, + abs_tol=1e-3, rel_tol = 1e-2, rho_scaler = 2, imbalance_tolerance = 10, warm_start=False, random_state=None): + self.alpha = alpha + self.fit_intercept = fit_intercept + self.max_iter = max_iter + self.abs_tol = abs_tol + self.rel_tol = rel_tol + self.warm_start = warm_start + self.random_state = random_state + self.comm = comm + self.rho_scaler = rho_scaler + self.imbalance_tolerance = imbalance_tolerance + + + def fit(self, X= None, y = None, z = None, rho_0_scaler = 1e-3, verbose = False): + """ + Fit model with coordinate descent. + + Parameters + ---------- + X : {array-like, sparse matrix} of shape (n_samples, n_features) + Training data. Only fed into the root rank. + + y : array-like of shape (n_samples,) or (n_samples, n_targets) + Target values. Only fed into the root rank. + + z : initial parameter + + + Returns + ------- + self : object + Returns self. + """ + max_iter = self.max_iter + abs_tol = self.abs_tol + rel_tol = self.rel_tol + + comm = self.comm + size = comm.Get_size() + rank = comm.Get_rank() + + + N = size + + + ''' + Data + ''' + if rank == 0: + + m, n = X.shape + + + # scaled global alpha values, + # *m(number of samples in this bootstrap) accounts for the ADMM objective function being the total SSE + # /N(n_admm) accounts for the change in L1-penalty scale when the SSE term optimization is distributed + alpha = self.alpha * m / N + + else: + n = np.zeros(1).astype('int') + m = np.zeros(1).astype('int') + alpha = np.zeros(1).astype('double') + + + + self.terminate_selection = False + self.terminate_estimation = False + + n = comm.bcast(n, root=0) + + # using non-sensibel features number as termination signals + if n == -1: + self.terminate_selection = True + return self + elif n == 0: + self.terminate_estimation = True + return self + + + m = comm.bcast(m, root=0) + + + alpha = comm.bcast(alpha, root=0) + + + + if rank != 0: + # X = np.empty((m, n), dtype=np.float64) + y = np.empty(m, dtype=np.float64) + z = np.empty((n,1), dtype=np.float64) + + ''' + Broadcast data + ''' + comm.Barrier() + if rank == 0: + for dest_rank in range(1, size): + + X_segment = X[dest_rank::N, :] + + # Send the segment shape first + segment_shape = np.array(X_segment.shape, dtype=np.int32) + comm.Send([segment_shape, MPI.INT], dest=dest_rank, tag=10) + + # Send the actual segment + comm.Send([np.ascontiguousarray(X_segment), MPI.DOUBLE], dest=dest_rank, tag=11) + + # comm.Bcast([X, MPI.DOUBLE]) + + + X = np.ascontiguousarray(X[0::N, :]) + + + gc.collect() + else: + # First receive the shape of the incoming data + segment_shape = np.empty(2, dtype=np.int32) + comm.Recv([segment_shape, MPI.INT], source=0, tag=10) + + # Create a buffer to receive the segment + X = np.empty(segment_shape, dtype=np.float64) + + # Receive the segment + comm.Recv([X, MPI.DOUBLE], source=0, tag=11) + + + ''' + Select sample block + ''' + + # X = X[rank::N, :] + comm.Barrier() + + m, n = X.shape + comm.Bcast([y, MPI.DOUBLE]) + comm.Bcast([z, MPI.DOUBLE]) + + rho = alpha*rho_0_scaler + + # do the send-receisve again for y?? or integrate back into the last send-receive operation?? or just Bcast it like right now + y = np.ascontiguousarray(y.ravel()[rank::N].reshape((m, 1))) + + + # save a matrix-vector multiply + Xty = X.T.dot(y) + + + # initialize ADMM solver + + if z is None: + z = np.random.normal(scale=1, size=(n, 1)) + #z = np.zeros((n, 1)) + z_original = deepcopy(z) + x = deepcopy(z) + + + u = np.zeros((n, 1)) + r = np.zeros((n, 1)) + + send = np.zeros(3) + recv = np.zeros(3) + + # cache the (Cholesky) factorization + L, U = factor(X, rho) + + # Saving state + if rank == 0 and verbose: + print('\n%3s\t%10s\t%10s\t%10s\t%10s\t%10s' % ('iter', + 'r norm', + 'eps pri', + 's norm', + 'eps dual', + 'objective')) + objval = [] + r_norm = [] + s_norm = [] + eps_pri = [] + eps_dual = [] + + ''' + ADMM solver loop + ''' + sign_preserved = [] + param_list = [] + for k in range(max_iter): # xrange -> range for Python 3 + + + x_old = deepcopy(x) + # u-update + if k != 0: + u += (x - z) + + # x-update + q = Xty + rho * (z - u) # (temporary value) + + if m >= n: + x = spsolve(U, spsolve(L, q))[..., np.newaxis] + else: + ULXq = spsolve(U, spsolve(L, X.dot(q)))[..., np.newaxis] + x = (q * 1. / rho) - ((X.T.dot(ULXq)) * 1. / (rho**2)) + + w = x + u + + + + send[0] = r.T.dot(r)[0][0] + send[1] = x.T.dot(x)[0][0] + #send[2] = u.T.dot(u)[0][0] / (rho**2) + send[2] = u.T.dot(u)[0][0] * (rho**2) + + zprev = np.copy(z) + + comm.Barrier() + comm.Allreduce([w, MPI.DOUBLE], [z, MPI.DOUBLE]) + comm.Allreduce([send, MPI.DOUBLE], [recv, MPI.DOUBLE]) + + # z-update + + if alpha == 0: #Linear regression case + z = z * 1. / N + else: + z = EG_update(z, x_old, x, u, rho, alpha, rank, k) + + + r_res = np.sqrt(recv[0]) + s_res = np.sqrt(N) * rho * norm(z - zprev) + + primal_eps = np.sqrt(n * N) * abs_tol + rel_tol * np.maximum(np.sqrt(recv[1]), np.sqrt(N) * norm(z)) + dual_eps = np.sqrt(n * N) * abs_tol + rel_tol * np.sqrt(recv[2]) + + + # diagnostics, reporting, termination checks + objval.append(objective(X, y, alpha, x, z)) + r_norm.append(r_res) + s_norm.append(s_res) + eps_pri.append(primal_eps) + eps_dual.append(dual_eps) + + if rank == 0 and verbose: + print('%4d\t%10.4f\t%10.4f\t%10.4f\t%10.4f\t%10.2f' % (k + 1, + r_norm[k], + eps_pri[k], + s_norm[k], + eps_dual[k], + objval[k])) + + if r_norm[k] < eps_pri[k] and s_norm[k] < eps_dual[k] and k > 0: + break + + rho, u = adaptive_rho_update(r_res, s_res, primal_eps, dual_eps, rho, u, self.rho_scaler, self.imbalance_tolerance) + + + + + # Compute residual + r = x - z + + # print(np.sum(np.sign(z)==np.sign(z_original)), flush = True) + sign_preserved.append(np.sum(np.sign(z)==np.sign(z_original))) + param_list.append(deepcopy(z)) + + # Set attributes after fitting + self.coef_ = z + self.intercept_ = 0 + self.n_iter_ = k + self.sign_preserved = sign_preserved + self.param_list = param_list + + + + return self + + def set_params(self, alpha): + self.alpha = alpha + + def predict(self, X): + """ + Predict using the linear model. + + Parameters + ---------- + X : array-like or sparse matrix, shape (n_samples, n_features) + Samples. + + Returns + ------- + C : array, shape (n_samples,) + Returns predicted values. + """ + y_pred = X @ self.coef_.ravel() + self.intercept_ + + return y_pred + + def score(self, X, y, sample_weight=None): + """ + Return the coefficient of determination R^2 of the prediction. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) + Test samples. + + y : array-like of shape (n_samples,) or (n_samples, n_outputs) + True values for X. + + sample_weight : array-like of shape (n_samples,), default=None + Sample weights. + + Returns + ------- + score : float + R^2 of self.predict(X) wrt. y. + """ + # Implementation would go here + + return None + + def _set_intercept(self, X_mean, y_mean, X_std): + """ + Set the intercept based on the fit. + """ + # Implementation would go here + pass \ No newline at end of file diff --git a/src/pyuoi/linear_model/admm_mpi_poisson.py b/src/pyuoi/linear_model/admm_mpi_poisson.py new file mode 100644 index 00000000..fd559de8 --- /dev/null +++ b/src/pyuoi/linear_model/admm_mpi_poisson.py @@ -0,0 +1,597 @@ +from __future__ import division +import pdb, time, h5py, os +import numpy as np +import scipy.sparse as sparse +from scipy.sparse.linalg import spsolve, splu +from numpy.linalg import norm, cholesky +from mpi4py import MPI +from optparse import OptionParser +import gc +from copy import deepcopy +from .sparse_comm_util import * +from time import time + + +def adaptive_rho_update_boyd(r_res, s_res, rho, u, v, rho_scaler, imbalance_tolerance, n = None, p=None): + + if r_res > imbalance_tolerance * s_res: + rho *= rho_scaler + u /= rho_scaler + v /= rho_scaler + elif s_res > imbalance_tolerance * r_res: + rho /= rho_scaler + u *= rho_scaler + v *= rho_scaler + + return rho, u, v + +def adaptive_rho_update(r_res, s_res, primal_eps, dual_eps, rho, u, v, rho_scaler, imbalance_tolerance, n = None, p=None): + + if r_res/primal_eps > imbalance_tolerance * s_res/dual_eps: + rho *= rho_scaler + u /= rho_scaler + v /= rho_scaler + elif s_res/dual_eps > imbalance_tolerance * r_res/primal_eps: + rho /= rho_scaler + u *= rho_scaler + v *= rho_scaler + + return rho, u, v + + +def objective(X, y, lamb, alpha, x, z, w): + if alpha == 0: + return 0 + # raise NotImplementedError("NO objectivefunction yet") + #return .5 * np.square(X.dot(x) - y).sum() + else: + return 0 + #return .5 * np.square(X.dot(x) - y).sum() + alpha * norm(z, 1) + + + +def sparse_factor(X, rho, lambda_2): + m, n = X.shape + + if m >= n: + A = lambda_2*sparse.eye(n) + rho * (X.T.dot(X) + sparse.eye(n)) + else: + raise NotImplementedError("Sparse factor not written for n_feature > n_sample") + #A = sparse.eye(m) + 1. / rho * (X.dot(X.T)) + + # Use sparse LU(more general than Cholesky) factorization + # Note: scipy returns a factorized object, not the L matrix directly + # try: + factor = splu(A.tocsc()) # Convert to CSC for factorization + # except Exception: + # print(A.shape, flush = True) + + + return factor + + +def soft_threshold(v, k, mask): + if mask is not None: + mask = mask.reshape(-1, 1) + k_masked = mask * k + else: + k_masked = k + + return np.sign(v) * np.maximum(np.abs(v) - k_masked, 0.0) + + +def old_update_z(y, a, rho, tol=1e-3, max_iter=50): + # Newton's method for convex scalar optimization(has overflow issue) + z = a.copy() + for i in range(len(y)): + for _ in range(max_iter): + exp_z = np.exp(z[i]) + grad = exp_z - y[i] + rho * (z[i] - a[i]) + hess = exp_z + rho + step = grad / hess + z[i] -= step + if abs(step) < tol: + break + return z + +def safe_exp(x): + # Avoid overflow by capping + return np.exp(np.clip(x, -100, 50)) # You can tune these bounds + +def update_z(y, a, rho, w = 1, dt = 1, tol=1e-3, max_iter=10, eps = 1e-10): + + """ + Vectorized Newton solver for the z-update in Poisson regression ADMM. + + Args: + y: (n,) count data + a: (n,) target linear predictor (Xβ - u) + w: (n,) weights for each variable + rho: scalar, penalty parameter + tol: stopping tolerance + max_iter: maximum Newton iterations + + Returns: + z: updated vector (n,) + """ + z = a.copy() # Initial guess + + + for i in range(max_iter): + ez = safe_exp(z) + stabilizer = (1+z)/(1+z+eps) + + grad = w * ez * dt - w * y * stabilizer + rho * (z - a) + hess = w * ez * dt + rho + step = grad / hess + + z_new = z - step + + + if np.max(np.abs(step/z)) < tol: + break + + z = z_new + + # print(i, flush = True) + + return z + +def poisson_loss(x, beta, y, lamb): + """ + Parameters + ---------- + x : array, n_sample X n_features + Design matrix/Predictor variable. + + beta : bool, default=True + Model parameter estimates + + y : 1D array + Predicted variable. + + """ + tse = np.sum(np.exp(x@beta)-y*(x@beta)) + l1_term = lamb * np.sum(np.abs(beta)) + + return tse, l1_term + +class ADMM_Poisson: + """ + MPI ADMM implementation of Generalized Linear Model trained with Elastic-net regularizer + + The optimization objective for Lasso is: + + E(w) = ***poisson stuff***0.5 * ||y - Xw||^2_2 + alpha (L1_regularizer * ||w||_1 + (1-L1_regularizer)* ||w||_2) + + Parameters + ---------- + alpha : float, default=1.0 + Constant that multiplies the L1 term. alpha = 0 is equivalent to + ordinary least squares. For numerical reasons, using alpha = 0 is + not advised; instead, you should use Ridge or LinearRegression. + + fit_intercept : bool, default=True + Whether to calculate the intercept for this model. If set to False, + no intercept will be used in calculations. + + max_iter : int, default=1000 + The maximum number of iterations + + tol : float, default=1e-4 + The tolerance for the optimization. + + warm_start : bool, default=False + When set to True, reuse the solution of the previous call to fit as + initialization, otherwise, just erase the previous solution. + + processes: list + the integer indices of MPI processes that are alotted for the MPI-ADMM fit + + random_state : int, RandomState instance, default=None + The seed of the pseudo random number generator that selects a random + feature to update. + + Attributes + ---------- + coef_ : array, shape (n_features,) | (n_targets, n_features) + Parameter vector (w in the cost function formula). + + intercept_ : float | array, shape (n_targets,) + Independent term in decision function. + + n_iter_ : int | array-like, shape (n_targets,) + Number of iterations run by the coordinate descent solver to reach + the specified tolerance. + """ + + def __init__(self, comm, rho = None, alpha=None, fit_intercept=False, max_iter=50, + abs_tol=1e-6, rel_tol = 1e-5,rho_scaler = 2, imbalance_tolerance = 10, warm_start=True, random_state=None, feature_weights = 1, dt = 1): + self.alpha = alpha + self.fit_intercept = fit_intercept + self.max_iter = max_iter + self.abs_tol = abs_tol + self.rel_tol = rel_tol + self.warm_start = warm_start + self.random_state = random_state + self.comm = comm + self.coef_ = 0 + self.rho_scaler = rho_scaler + self.imbalance_tolerance = imbalance_tolerance + self.feature_weights = feature_weights + self.dt = dt + + def fit(self, X= None, y = None, w = None, sparse_input = True, param_mask = None): + """ + Fit model with coordinate descent. + + Parameters + ---------- + X : {array-like, sparse matrix} of shape (n_samples, n_features) + Training data. Only fed into the root rank. + + y : array-like of shape (n_samples,) or (n_samples, n_targets) + Target values. Only fed into the root rank. + + + Returns + ------- + self : object + Returns self. + """ + max_iter = self.max_iter + abs_tol = self.abs_tol + rel_tol = self.rel_tol + + comm = self.comm + size = comm.Get_size() + rank = comm.Get_rank() + + + N = size + + # X_full = deepcopy(X) + # y_full = deepcopy(y) + + ''' + Data + ''' + if rank == 0: + + m, n = X.shape + + + # scaled global alpha values, + # *m(number of samples in this bootstrap) accounts for the ADMM objective function being the total SSE + # /N(n_admm) accounts for the change in L1-penalty scale when the SSE term optimization is distributed + lamb = self.lamb * m / N + alpha = self.alpha + + weights = np.tile(self.feature_weights, (int(len(y)/len(self.feature_weights)), 1)) + weights = weights.T.flatten() + + loss = {} + loss["tse"] = [] + loss["l1"] = [] + + + else: + n = np.zeros(1).astype('int') + m = np.zeros(1).astype('int') + lamb = np.zeros(1).astype('double') + alpha = np.zeros(1).astype('double') + + + self.terminate_selection = False + self.terminate_estimation = False + + n = comm.bcast(n, root=0) + + # using non-sensibel features number as termination signals + if n == -1: + self.terminate_selection = True + return self + elif n == 0: + self.terminate_estimation = True + return self + + + m = comm.bcast(m, root=0) + + + alpha = comm.bcast(alpha, root=0) + lamb = comm.bcast(lamb, root=0) + + + + if rank != 0: + # X = np.empty((m, n), dtype=np.float64) + y = np.empty(m, dtype=np.float64) + weights = np.empty(m, dtype=np.float64) + + + ''' + Broadcast data + ''' + comm.Barrier() + if rank == 0: + for dest_rank in range(1, size): + + X_segment = X[dest_rank::N, :] + # since it's sparse csr matrix, don't have to Send the segment shape first, send the actual segment + if sparse_input: + send_sparse_matrix(comm, X_segment, dest=dest_rank, tag = 11) + else: + # Send the segment shape first + segment_shape = np.array(X_segment.shape, dtype=np.int32) + comm.Send([segment_shape, MPI.INT], dest=dest_rank, tag=10) + + # Send the actual segment + comm.Send([np.ascontiguousarray(X_segment), MPI.DOUBLE], dest=dest_rank, tag=11) + if sparse_input: + X = X[0::N, :] + else: + X = np.ascontiguousarray(X[0::N, :]) + + gc.collect() + else: + if sparse_input: + X = recv_sparse_matrix(comm, source=0, tag=11) + else: + # First receive the shape of the incoming data + segment_shape = np.empty(2, dtype=np.int32) + comm.Recv([segment_shape, MPI.INT], source=0, tag=10) + + # Create a buffer to receive the segment + X = np.empty(segment_shape, dtype=np.float64) + + # Receive the segment + comm.Recv([X, MPI.DOUBLE], source=0, tag=11) + + + ''' + Select sample block + ''' + + # X = X[rank::N, :] + comm.Barrier() + + m, n = X.shape + comm.Bcast([y, MPI.DOUBLE]) + comm.Bcast([weights, MPI.DOUBLE]) + param_mask = comm.bcast(param_mask, root=0) + + # m is n_samp per ADMM process! + # this is for accomdating the definition of MSE term in ADMM-LASSO convention + #alpha *= m + # good heuristric is to start with rho = l1-penalty + + # rho = lamb * alpha / m + rho = lamb * alpha # admm parameter, modulating the constraint that aux variable equals the model variable + + # do the send-receive again for y?? or integrate back into the last send-receive operation?? or just Bcast it like right now + y = np.ascontiguousarray(y.ravel()[rank::N].reshape((m, 1))) + + weights = np.ascontiguousarray(weights.ravel()[rank::N].reshape((m, 1))) + + + + # initialize ADMM solver + if w is None: # w = None will be the default case!!! + if self.warm_start and not np.all(self.coef_ == 0): + w = deepcopy(self.coef_) + else: + np.random.seed(self.random_state) + w = np.random.normal(scale=1, size=(n, 1)) + x = deepcopy(w) + z = X.dot(x) + + u = np.zeros_like(z) + v = np.zeros((n, 1)) + + + send = np.zeros(7) + recv = np.zeros(7) + + + + + # cache the (Cholesky) factorization + if sparse_input: + factor_obj = sparse_factor(X, rho, lamb * (1-alpha)) + else: + L, U = factor(X, rho, lamb * (1-alpha)) + + + objval = [] + r_norm = [] + s_norm = [] + eps_pri = [] + eps_dual = [] + + ''' + ADMM solver loop + ''' + + rho_history = [] + for k in range(max_iter): # xrange -> range for Python 3 + + if rank == 0 and k%10 == 0: + tse, l1_term = poisson_loss(X, w, y, lamb) + loss["tse"].append(tse) + loss["l1"].append(l1_term) + + # x-update + q = rho * (X.T.dot(z + u) + (w + v)) # (temporary value) + + if sparse_input: + if m >= n: + x = factor_obj.solve(q) + else: + raise NotImplementedError() + # ULXq = factor_obj.solve(X.dot(q)) + # x = (q * 1. / rho) - ((X.T.dot(ULXq)) * 1. / (rho**2)) + else: + raise NotImplementedError() + # if m >= n: + # x = spsolve(U, spsolve(L, q))[..., np.newaxis] + # else: + # ULXq = spsolve(U, spsolve(L, X.dot(q)))[..., np.newaxis] + # x = (q * 1. / rho) - ((X.T.dot(ULXq)) * 1. / (rho**2)) + + + # z-update (newton's method) + a = X.dot(x) - u + zprev = np.copy(z) + z = update_z(y, a, rho, dt = self.dt, w = weights) + + w_input = x - v + + + + comm.Barrier() + comm.Allreduce([w_input, MPI.DOUBLE], [w, MPI.DOUBLE]) # the resulting w is sum of N variants, so it has to be divided by N before all usage + + wprev = np.copy(w) + + # w-update(all proesses does the same computation) + if lamb == 0: #Linear regression case + w = w * 1. / N + else: + w = soft_threshold(w * 1. / N, lamb * alpha * 1. / (N * rho), mask = param_mask) + + u += (z - X.dot(x)) # u is dual of z, so it's update is local + v += (w - x) # v is dual of w, so it's update should be global(but keep it local for now, since x is a vector that need to be Allreduced??!) + + r1 = X.dot(x) - z + r2 = x - w + + XTzz = X.T.dot(z-zprev) + XTu = X.T.dot(u) + Xx = X.dot(x) + + send[0] = (Xx).T.dot(Xx)[0][0] + send[1] = z.T.dot(z)[0][0] + send[2] = x.T.dot(x)[0][0] + send[3] = XTzz.T.dot(XTzz)[0][0] + send[4] = XTu.T.dot(XTu)[0][0] + send[5] = v.T.dot(v)[0][0] + send[6] = r1.T.dot(r1)[0][0] + r2.T.dot(r2)[0][0] + + + comm.Barrier() + comm.Allreduce([send, MPI.DOUBLE], [recv, MPI.DOUBLE]) + + r_res = np.sqrt(recv[6]) + s_res = rho * np.sqrt(recv[3] + N * norm(w - wprev)**2) + + primal_eps = np.sqrt(n * N) * abs_tol + rel_tol * np.max(np.array([np.sqrt(recv[0]), np.sqrt(recv[1]), np.sqrt(recv[2]), np.sqrt(N) * norm(w)])) + dual_eps = np.sqrt(n * N) * abs_tol + rel_tol * rho * np.maximum(np.sqrt(recv[4]), np.sqrt(recv[5])) + + + # diagnostics, reporting, termination checks + objval.append(objective(X, y, lamb, alpha, x, z, w)) + r_norm.append(r_res) + s_norm.append(s_res) + eps_pri.append(primal_eps) + eps_dual.append(dual_eps) + + + if r_res < primal_eps and s_res < dual_eps and k > 0: + break + + + # adaptive rho selection based on residual + #rho, u = adaptive_rho_update_boyd(r_res, s_res, rho, u, self.rho_scaler, self.imbalance_tolerance) + if k%20 == 0: + rho, u, v = adaptive_rho_update(r_res, s_res, primal_eps, dual_eps, rho, u, v, self.rho_scaler, self.imbalance_tolerance) + + + + if rank == 0: + rho_history.append(rho) + # if rank == 0: + # np.save("/global/homes/y/yxu2/packages/pyuoi/examples/rho_plot/rho_"+str(self.rho_scaler)+".npy", rho_history) + + # Set attributes after fitting + self.coef_ = w # we want w because it's a global variable, and converges to x in the limit + self.intercept_ = 0 + self.n_iter_ = None + if rank == 0: + self.loss = loss + + + return self + + def set_params(self, alpha, l1_ratio): + # + self.lamb = alpha + self.alpha = l1_ratio + + def predict(self, X): + """Predicts the response variable given a design matrix. The output is + the mode of the Poisson distribution. + + Parameters + ---------- + X : array_like, shape (n_samples, n_features) + Design matrix to predict on. + + Returns + ------- + mode : array_like, shape (n_samples) + The predicted response values, i.e. the modes. + """ + mu = self.predict_mean(X) + + mode = np.floor(mu) + return mode + + def predict_mean(self, X): + """Calculates the mean response variable given a design matrix. + + Parameters + ---------- + X : array_like, shape (n_samples, n_features) + Design matrix to predict on. + + Returns + ------- + mu : array_like, shape (n_samples) + The predicted response values, i.e. the conditional means. + """ + mu = np.exp(np.clip(self.intercept_ + np.dot(X, self.coef_), -5, 5)) * self.dt + + return mu + + + def score(self, X, y, sample_weight=None): + """ + Return the coefficient of determination R^2 of the prediction. + + Parameters + ---------- + X : array-like of shape (n_samples, n_features) + Test samples. + + y : array-like of shape (n_samples,) or (n_samples, n_outputs) + True values for X. + + sample_weight : array-like of shape (n_samples,), default=None + Sample weights. + + Returns + ------- + score : float + R^2 of self.predict(X) wrt. y. + """ + # Implementation would go here + + return None + + def _set_intercept(self, X_mean, y_mean, X_std): + """ + Set the intercept based on the fit. + """ + # Implementation would go here + pass \ No newline at end of file diff --git a/src/pyuoi/linear_model/base.py b/src/pyuoi/linear_model/base.py index ca93d160..2a7d896a 100644 --- a/src/pyuoi/linear_model/base.py +++ b/src/pyuoi/linear_model/base.py @@ -6,16 +6,324 @@ from sklearn.model_selection import train_test_split from sklearn.utils import check_X_y from sklearn.preprocessing import StandardScaler +import sys +from time import time -from scipy.sparse import issparse, csr_matrix +from scipy import sparse +from scipy.sparse import issparse, csr_matrix, csc_matrix, coo_matrix, kron, eye, block_diag from pyuoi import utils from pyuoi.mpi_utils import (Gatherv_rows, Bcast_from_root) -from .utils import stability_selection_to_threshold, intersection +from .utils import * from ..utils import check_logger +import gc +from copy import deepcopy +from mpi4py import MPI +def _build_design_matrix(X_lagged): + """ + Build block-diagonal design matrix for vectorized regression. + + For p dimensions and T-1 transitions: + - Each dimension i has its own coefficients: [b_i, A_{i,1}, ..., A_{i,p}] + - Design matrix has blocks for each dimension + + Parameters + ---------- + X_lagged : (T-1, p) array + Lagged observations X_{t-1} for t=1,...,T-1 + + Returns + ------- + design_matrix : ((T-1)*p, p*(p+1)) sparse array + Block diagonal design matrix + """ + T_minus_1, p = X_lagged.shape + + # Preserve the input dtype + dtype = X_lagged.dtype + + # Build blocks for each dimension + blocks = [] + for dim in range(p): + # For this dimension, create the feature matrix + # Features are [1, X_{t-1,1}, ..., X_{t-1,p}] for each t + block = np.column_stack([ + np.ones(T_minus_1, dtype=dtype), # intercept with matching dtype + X_lagged # all lagged dimensions + ]) + blocks.append(csr_matrix(block, dtype=dtype)) + + # Create block diagonal matrix + design_matrix = block_diag(blocks) + + return csr_matrix(design_matrix, dtype=dtype) + +def _vectorize_response(X_curr): + """ + Vectorize the response variable. + + Parameters + ---------- + X_curr : (T-1, p) array + Current observations X_t for t=1,...,T-1 + + Returns + ------- + y : ((T-1)*p,) array + Vectorized response + """ + # Stack all dimensions: first all t for dim 0, then all t for dim 1, etc. + return X_curr.T.ravel() + +def vectorization_bootstrap_test(X, sample_idx, lag, data_pois = None, feature_weights = None, has_bias = True, sparse_format = "csr"): + # Prepare data + X_lagged = X[:-1]#.astype(np.float64) # X_{t-1} for t=1,...,T-1 + X_curr = X[1:]#.astype(np.float64) # X_t for t=1,...,T-1 + + + design_matrix = _build_design_matrix(X_lagged) + y = _vectorize_response(X_curr) + + # sparse.save_npz("data/design_matrix.npz", design_matrix) + # np.save("data/y", y) + + + + return design_matrix, y, None + + +def _unpack_coefficients(coefficients, n_feat): + """ + Unpack coefficient vector into A matrix and b vector. + + Parameters + ---------- + coefficients : (p*(p+1),) array + Flattened coefficients [b_1, A_{1,:}, b_2, A_{2,:}, ...] + + Returns + ------- + A : (p, p) array + b : (p,) array + """ + p = n_feat + A = np.zeros((p, p)) + b = np.zeros(p) + coefficients= coefficients.ravel() + for i in range(p): + start_idx = i * (p + 1) + b[i] = coefficients[start_idx] + A[i, :] = coefficients[start_idx + 1 : start_idx + p + 1] + + return A, b + +def recover_VAR_parameter_from_vectorized_test(coef, lag, n_features, has_bias=True): + A, b = _unpack_coefficients(coef) + A = A[np.newaxis,...] + return A, b + +def parameter_mask(lag, n_features, has_bias = True, exclude_diagonal = True): + # to exclude bias terms and/or diagonal entries from L1-regularization shrinkage + # return a 1D vector that is the mask for the vectorized model coefficient with bias terms + + mask = np.ones((n_features,n_features)) + if exclude_diagonal: + np.fill_diagonal(mask, 0) + mask = np.tile(mask, (lag, 1)) + + if has_bias: + # padding a row of zeros that corresponds to bias terms + b = np.zeros(n_features) + mask = np.vstack((b,mask)) + + return mask.T.flatten() + + +def recover_VAR_parameter_from_vectorized(coef, lag, n_features, has_bias=True): + # reshape back to matrix form + # A: adjacency matrices + # b: intercept + coef = coef.reshape(n_features, lag * n_features + has_bias).T + + if has_bias: + b = coef[0, :] # shape (n_features,) + adjacency_flat = coef[1:, :] # shape (lag*n_features, n_features) + else: + b = None + adjacency_flat = coef # shape (lag*n_features, n_features) + + # reshape into adjacency matrices per lag + adjacency_matrices = adjacency_flat.reshape(lag, n_features, n_features) + + A = np.transpose(adjacency_matrices, (0, 2, 1)) + + return A, b + + +def vectorization_bootstrap(raw_data, sample_idx, lag, data_pois = None, feature_weights = None, has_bias = True, sparse_format = "csr"): + # vectorize the VAR bootstrap data for use with LASSO algorithm + # sample_idx: idx for sampling the response vectors of the VAR model + # data: raw time series data in np array with shape n_samples X n_features + # Construction strategy: first form X_boot, Y_boot, then vectorize both + + # flipup so the last time sample in data is now first row(no need anymore) + data = deepcopy(raw_data) + data = np.flipud(data) + n_samples = data.shape[0] + n_features = data.shape[1] + + # shape of Y: (n_boot_sample) X (n_features) + if data_pois is None: + Y = data[sample_idx] # sample_idx are in range: (0, n_samples-lag) + else: + Y = data_pois[sample_idx] + + if feature_weights is not None: + weights = np.tile(feature_weights, (Y.shape[0], 1)) + weights = weights.T.flatten() + else: + weights = None + + Y = Y.T.flatten() + + # X.shape: (n_boot_sample) X (lag * n_features); + X_row = [np.hstack(data[i+1 : lag+(i+1)]) for i in sample_idx] + X = np.vstack(X_row) + + if has_bias: + # adding internal bias term to fit + b = np.ones((X.shape[0],1)) + X = np.hstack((b,X)) + + # use kronecker product for vectorization of matrix multiplication + # use sparse kron operation of limit memory expansion + X = kron(eye(n_features), X, format=sparse_format) + + X, Y = check_X_y(X, Y, accept_sparse=['csr', 'csc', 'coo'], + y_numeric=True, multi_output=True) + return X, Y, weights + + +def vectorization_bootstrap_shuffle(raw_data, lag, boot_frac = 0.5, feature_weights = None, has_bias = True, sparse_format = "csr"): + # vectorize the VAR bootstrap data for use with LASSO algorithm + # sample_idx: idx for sampling the response vectors of the VAR model + # data: raw time series data in np array with shape n_samples X n_features + # Construction strategy: first form X_boot, Y_boot, then vectorize both + + # flipup so the last time sample in data is now first row(no need anymore) + data = deepcopy(raw_data) + for i in range(data.shape[1]): + np.random.shuffle(data[:, i]) + + data = np.flipud(data) + n_samples = data.shape[0] + n_features = data.shape[1] + + # sample_idx are in range: (0, n_samples-lag) + sample_idx = np.random.choice(data.shape[0]-lag, size = int(boot_frac*data.shape[0]), replace = False) + + # shape of Y: (n_boot_sample) X (n_features) + Y = data[sample_idx] + + + if feature_weights is not None: + weights = np.tile(feature_weights, (Y.shape[0], 1)) + weights = weights.T.flatten() + else: + weights = None + + Y = Y.T.flatten() + + + # X.shape: (n_boot_sample) X (lag * n_features); + X_row = [np.hstack(data[i+1 : lag+(i+1)]) for i in sample_idx] + X = np.vstack(X_row) + + if has_bias: + # adding internal bias term to fit + b = np.ones((X.shape[0],1)) + X = np.hstack((b,X)) + + # use kronecker product for vectorization of matrix multiplication + # use sparse kron operation of limit memory expansion + X = kron(eye(n_features), X, format=sparse_format) + + X, Y = check_X_y(X, Y, accept_sparse=['csr', 'csc', 'coo'], + y_numeric=True, multi_output=True) + return X, Y, weights + + +def vectorization(raw_data, lag, data_pois = None, feature_weights = None, has_bias = True): + # vectorize the full raw VAR data for use with LASSO algorithm + # data: time series data in np array with shape n_samples X n_features + + # flipup so the last time sample in data is now first row + data = deepcopy(raw_data) + data = np.flipud(data) + n_samples = data.shape[0] + n_features = data.shape[1] + + # shape of Y: (T - D) X (n_features) *T - D: total number of samples - lag + if data_pois is None: + Y = data[: n_samples-lag] + else: + Y = data_pois[: n_samples-lag] + + if feature_weights is not None: + weights = np.tile(feature_weights, (Y.shape[0], 1)) + weights = weights.T.flatten() + else: + weights = None + + # ones = np.ones((n_samples-lag,1)) # for VAR porocess with intercept + Y = Y.T.flatten() + + # X.shape: (n_samples - lag) X (lag * n_features); + X_row = [np.hstack(data[i : lag+i]) for i in range(1,n_samples-lag+1)] + X = np.vstack(X_row) + + if has_bias: + # adding internal bias term to fit + b = np.ones((X.shape[0],1)) + X = np.hstack((b,X)) + + # use kronecker product for vectorization of matrix multiplication + X = kron(eye(n_features), X, format="csr") + + X, Y = check_X_y(X, Y, accept_sparse=['csr', 'csc', 'coo'], + y_numeric=True, multi_output=True) + return X, Y, weights + +def intermediate_data(raw_data, lag, data_pois = None, has_bias = True): + # produce the linear system used to find regularization path + + # flipup so the last time sample in data is now first row + data = deepcopy(raw_data) + data = np.flipud(data) + n_samples = data.shape[0] + n_features = data.shape[1] + + # shape of Y: (T - D) X (n_features) *T - D: total number of samples - lag + if data_pois is None: + Y = data[: n_samples-lag] + else: + Y = data_pois[: n_samples-lag] + + # X.shape: (n_samples - lag) X (lag * n_features); + X_row = [np.hstack(data[i : lag+i]) for i in range(1,n_samples-lag+1)] + X = np.vstack(X_row) + + #should not include this for fingthe L1-penalty, since bias terms are not L1-constrained + # if has_bias: + # # adding internal bias term to fit + # b = np.ones((X.shape[0],1)) + # X = np.hstack((b,X)) + + return X, Y + class AbstractUoILinearModel(SparseCoefMixin, metaclass=_abc.ABCMeta): r"""An abstract base class for UoI ``linear_model`` classes. @@ -82,11 +390,14 @@ class AbstractUoILinearModel(SparseCoefMixin, metaclass=_abc.ABCMeta): for estimation for a given regularization parameter value (row). """ - def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, - estimation_frac=0.9, stability_selection=1., + def __init__(self, fit_VAR = False, n_boots_sel=12, n_boots_est=12, selection_frac=0.9, + estimation_frac=0.9, stability_selection=0.75, fdr_rate = 0.05, fit_intercept=True, standardize=True, shared_support=True, max_iter=None, tol=None, - random_state=None, comm=None, logger=None): + random_state=None, comm=None, admm_comm = None, logger=None): + + self.fit_VAR = fit_VAR + # data split fractions self.selection_frac = selection_frac self.estimation_frac = estimation_frac @@ -101,6 +412,10 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, self.max_iter = max_iter self.tol = tol self.comm = comm + self.admm_comm = admm_comm + self.output_dim = 1 # for VAR model: by vectorization construction + self.VAR_coef_ = None + self.fdr_rate = fdr_rate # preprocessing if isinstance(random_state, int): # make sure ranks use different seed @@ -120,6 +435,7 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, self.n_supports_ = None self._logger = check_logger(logger, 'uoi_linear_model', self.comm) + self.problem_size = None @_abc.abstractproperty def estimation_score(self): @@ -138,6 +454,55 @@ def intersect(self, coef, thresholds): """Intersect coefficients across all thresholds.""" pass + def _pre_fit_VAR(self, X): + """Perform z-score scaling for VAR raw data for fit().""" + if self.standardize: + if self.fit_intercept and issparse(X): + msg = ("Cannot center sparse matrices: " + "pass `fit_intercept=False`") + raise ValueError(msg) + + self._X_scaler = StandardScaler(with_mean=self.fit_intercept) + X = self._X_scaler.fit_transform(X) + + return X + + def _post_fit_VAR(self, lag, n_features): + """Perform VAR model coefficiant rescaling if the raw data is z-score scaled.""" + # construct the lagged connectivity matrices from the vectorized model coefficient + # A_model = np.array([self.coef_.reshape(n_features,n_features*lag).T[i*n_features:(i+1)*n_features].T for i in range(lag)]) + A_model, b_model = recover_VAR_parameter_from_vectorized(self.coef_, lag, n_features) + + if self.standardize: + sX = self._X_scaler + Sigma = np.outer(sX.scale_, 1/sX.scale_) + A_model *= Sigma + + # adjustment for mean shifts from the data z-score scaling + adjustment = (np.eye(n_features)-np.sum(A_model, axis = 0)) @ sX.mean_ + + + + b_model *= sX.scale_ + b_model += adjustment + + # adjustment for mean shifts from the data z-score scaling + if self.fit_intercept: + + adjustment = (np.eye(n_features)-np.sum(A_model, axis = 0)) @ sX.mean_ + + # check whether VAR model intercept is scalar or vector + if len(self.intercept_) == n_features: + self.intercept_ += adjustment + elif len(self.intercept_) == 1: + self.intercept_ += np.mean(adjustment) + + self.VAR_coef_ = A_model + self.VAR_bias_ = b_model + + + + def _pre_fit(self, X, y): """Perform class-specific setup for fit().""" if self.standardize: @@ -158,6 +523,7 @@ def _post_fit(self, X, y): if self.standardize: sX = self._X_scaler self.coef_ /= sX.scale_[np.newaxis] + @_abc.abstractmethod def _fit_intercept(self, X, y): @@ -176,17 +542,15 @@ def _fit_intercept_no_features(self, y): """ pass - def fit(self, X, y, stratify=None, verbose=False): + def fit(self, lag = 1, data = None, data_pois = None, stratify=None, verbose=False): """Fit data according to the UoI algorithm. Parameters ---------- - X : ndarray or scipy.sparse matrix, (n_samples, n_features) - The design matrix. - y : ndarray, shape (n_samples,) - Response vector. Will be cast to X's dtype if necessary. - Currently, this implementation does not handle multiple response - variables. + data : ndarray or scipy.sparse matrix, (n_boot_samples * (lag+1), n_features) + The boostrap design matrix. + lag : integer + Lag of the VAR model. stratify : array-like or None Ensures groups of samples are alloted to training/test sets proportionally. Labels for each group must be an int greater @@ -196,98 +560,224 @@ def fit(self, X, y, stratify=None, verbose=False): A switch indicating whether the fitting should print out messages displaying progress. """ + if verbose: self._logger.setLevel(logging.DEBUG) else: self._logger.setLevel(logging.WARNING) + + rank = 0 + size = 1 + if self.comm is not None: + rank = self.comm.rank + size = self.comm.size + + + # z-score scaling the raw data and picking the regularization parameters + if rank == 0: - X, y = self._pre_fit(X, y) + if self.fit_VAR: + + data = self._pre_fit_VAR(data) + + # only used for getting regularization parameters + # not fully vectorized X and y + X_all, y_all = intermediate_data(data, lag, data_pois = data_pois) + + # choose the regularization parameters for selection sweep + reg_params_ = self.get_reg_params(X_all, y_all) + + # clear memory + del X_all + del y_all + gc.collect() + param_mask = parameter_mask(lag, data.shape[1], has_bias = True) + + else: #fitting LASSO model(not VAR!) + X = data[0] + y = data[1] + X, y = self._pre_fit(X, y) + + X, y = check_X_y(X, y, accept_sparse=['csr', 'csc', 'coo'], + y_numeric=True, multi_output=True) + reg_params_ = self.get_reg_params(X, y) + + else: + reg_params_ = None + X = None + y = None + param_mask = None + - X, y = check_X_y(X, y, accept_sparse=['csr', 'csc', 'coo'], - y_numeric=True, multi_output=True) - # extract model dimensions - n_features = X.shape[1] - n_coef = self.get_n_coef(X, y) + + if size > 1: + if self.fit_VAR: + # Broadcast pre-fitted raw data to all ranks + data = self.comm.bcast(data, root=0) + param_mask = self.comm.bcast(param_mask, root=0) + data_pois = self.comm.bcast(data_pois, root=0) + + else: + X = self.comm.bcast(X, root=0) + y = self.comm.bcast(y, root=0) + reg_params_ = self.comm.bcast(reg_params_, root=0) + + + self.param_mask = param_mask + self.reg_params_ = reg_params_ + self.n_reg_params_ = len(self.reg_params_) + + if not self.fit_VAR: + #for LASSO fitting, convert X to sparse matrix format + X = csr_matrix(X) + + # y = csr_matrix(y) + + + # extract model dimensions + if self.fit_VAR: + # this is specifically for the vectorized VAR model + # including bias term as well + n_coef = lag * data.shape[1]**2 + data.shape[1] + n_features = n_coef + else: + n_features = X.shape[1] + n_coef = self.get_n_coef(X, y) + # check if the response variable is constant - if np.unique(y).size == 1: - self.coef_ = np.zeros((self.output_dim, n_features)) - self._fit_intercept(X, y) - self._post_fit(X, y) - return self + if self.fit_VAR: + + if np.unique(data).size == 1: + # self.coef_ = np.zeros((self.output_dim, n_features)) + # self._fit_intercept(X_all, y_all) + # self._post_fit(X_all, y_all) + return self + else: + if np.unique(y).size == 1: + self.coef_ = np.zeros((self.output_dim, n_features)) + self._fit_intercept(X, y) + self._post_fit(X, y) + return self + #################### # Selection Module # #################### - # choose the regularization parameters for selection sweep - self.reg_params_ = self.get_reg_params(X, y) - self.n_reg_params_ = len(self.reg_params_) - - rank = 0 - size = 1 - if self.comm is not None: - rank = self.comm.rank - size = self.comm.size # initialize selection - if size > self.n_boots_sel: - tasks = np.array_split(np.arange(self.n_boots_sel * + + # assigning tasks to each process(from total of n_boot_sel*n_reg_param tasks) + + if size > self.n_boots_sel: + # more process than n_boots_sel, so divvy the n_reg_param task up also + # each task is for a single boot and a single reg_parameter, so the num_tasks here will be Xn_reg_params more than the "else" clause + tasks = np.array_split(np.arange(self.n_boots_sel * self.n_reg_params_), size)[rank] selection_coefs = np.empty((tasks.size, n_coef)) - my_boots = dict((task_idx // self.n_reg_params_, None) + selection_coefs_shuffled = np.empty((tasks.size, n_coef)) + # but my_boots is still indexed by the boots, since the reg_param is picked separately when fitting is about to begin + my_boots = dict((task_idx // (self.n_reg_params_), None) for task_idx in tasks) - else: + else: + # less process than n_boots_sel, so some process will have multiple boots with all there reg_param runs as well + # we can see the selection_coeff will have space for all the reg_params + # split up bootstraps into processes tasks = np.array_split(np.arange(self.n_boots_sel), size)[rank] selection_coefs = np.empty((tasks.size, self.n_reg_params_, n_coef)) + + selection_coefs_shuffled = np.empty((tasks.size, self.n_reg_params_, + n_coef)) my_boots = dict((task_idx, None) for task_idx in tasks) + # this is where it assign the sample index to each bootstrap + + # with the rvals, rank 0 pick out the samples correponding rvals for each process, then sends them + # it will need to know what process has what "boot" tho, which can be obtained from "tasks" before taking the "[rank]", it's split into num_processes + # if size > self.n_boots_sel, then process "rank" will have boot=tasks[rank]//self.n_reg_params_ for boot in range(self.n_boots_sel): - if size > 1: + if size > 1: # num_MPI_process if rank == 0: - rvals = train_test_split(np.arange(X.shape[0]), + # selecting the samples for this particular bootstrap + # rvals[0] is the training_set_idx, rvals[1] is the test_set_idx + if self.fit_VAR: + #select from the (num of raw samples - lag) + idx_set = np.arange(data.shape[0]-lag) + else: + idx_set = np.arange(X.shape[0]) + + rvals = train_test_split(idx_set, test_size=1 - self.selection_frac, stratify=stratify, - random_state=self.random_state) + random_state=self.random_state) else: rvals = [None] * 2 + + # brocasting raw rvals to all processes and let them decide for themselves if rvals are theirs to keep rvals = [Bcast_from_root(rval, self.comm, root=0) for rval in rvals] + if boot in my_boots.keys(): + # the raw rvals are picking entries of the Y vector my_boots[boot] = rvals - else: + + else: #non-MPI + if self.fit_VAR: + idx_set = np.arange(data.shape[0]-lag) + else: + idx_set = np.arange(X.shape[0]) + my_boots[boot] = train_test_split( - np.arange(X.shape[0]), + idx_set, test_size=1 - self.selection_frac, stratify=stratify, - random_state=self.random_state) + random_state=self.random_state) + + # iterate over bootstraps(and reg_params) + # fitting happens here - # iterate over bootstraps + curr_boot_idx = None for ii, task_idx in enumerate(tasks): - if size > self.n_boots_sel: + if size > self.n_boots_sel: #in this case, my_reg_params has only 1 parameter boot_idx = task_idx // self.n_reg_params_ reg_idx = task_idx % self.n_reg_params_ my_reg_params = [self.reg_params_[reg_idx]] else: boot_idx = task_idx my_reg_params = self.reg_params_ + # Never warm start across bootstraps if (curr_boot_idx != boot_idx): if hasattr(self._selection_lm, 'coef_'): self._selection_lm.coef_ *= 0. if hasattr(self._selection_lm, 'intercept_'): self._selection_lm.intercept_ *= 0. + + # draw a resampled bootstrap, and vectorize it + # minimize repeated vectorization if using the same bootstrap + idxs_train, idxs_test = my_boots[boot_idx] + if self.fit_VAR: + start_time = time() + + X_rep, y_rep, feature_weights = vectorization_bootstrap(data, idxs_train, lag, data_pois = data_pois) + + X_shuffled, y_shuffled, feature_weights_shuffled = vectorization_bootstrap_shuffle(data, lag) + + # if self.kron_time is not None: + # self.kron_time += time()-start_time + if self.problem_size is None: + self.problem_size = sys.getsizeof(X_rep) + sys.getsizeof(y_rep) + else: + X_rep = X[idxs_train] + y_rep = y[idxs_train] + curr_boot_idx = boot_idx - # draw a resampled bootstrap - idxs_train, idxs_test = my_boots[boot_idx] - X_rep = X[idxs_train] - y_rep = y[idxs_train] - # fit the coefficients if size > self.n_boots_sel: msg = ("selection bootstrap %d, " @@ -297,22 +787,73 @@ def fit(self, X, y, stratify=None, verbose=False): else: self._logger.info("selection bootstrap %d" % (boot_idx)) + + # not warm starting + self._selection_lm.coef_ *= 0 selection_coefs[ii] = np.squeeze( self.uoi_selection_sweep(X_rep, y_rep, my_reg_params)) + # not warm starting + self._selection_lm.coef_ *= 0 + selection_coefs_shuffled[ii] = np.squeeze( + self.uoi_selection_sweep(X_shuffled, y_shuffled, my_reg_params)) + + + #print(np.count_nonzero(selection_coefs[ii])/selection_coefs[ii].size,flush = True) + # try: + # assert np.count_nonzero(selection_coefs[ii])/selection_coefs[ii].size > 0.01 + # except AssertionError: + # if self.solver == "admm": + # n = -1 + # n = self.admm_comm.bcast(n, root=0) + # if self.estimation_solver == "admm": + # n = 0 + # n = self.admm_comm.bcast(n, root=0) + # sys.exit("Producing trivial solution") + + + # if distributed, gather selection coefficients to 0, # perform intersection, and broadcast results if size > 1: selection_coefs = Gatherv_rows(selection_coefs, self.comm, root=0) + + selection_coefs_shuffled = Gatherv_rows(selection_coefs_shuffled, self.comm, root=0) if rank == 0: if size > self.n_boots_sel: selection_coefs = selection_coefs.reshape( self.n_boots_sel, self.n_reg_params_, n_coef) - supports = self.intersect( - selection_coefs, - self.selection_thresholds_).astype(int) + + selection_coefs_shuffled = selection_coefs_shuffled.reshape( + self.n_boots_sel, + self.n_reg_params_, + n_coef) + + # np.save("/pscratch/sd/y/yxu2/packages/pyuoi/examples/result/L0_support_intersect.npy", selection_coefs) + # np.save("/pscratch/sd/y/yxu2/packages/pyuoi/examples/result/L0_support_shuf_intersect.npy", selection_coefs_shuffled) + + # printing the total number of non-zero coef + # print(np.sum(selection_coefs > 0, axis=2)+np.sum(selection_coefs < 0, axis=2), flush = True) + + # Option 1: occurrence frequency based selection + # supports = self.intersect( + # selection_coefs, + # self.selection_thresholds_).astype(int) + + # Option 2: FDR selection + supports, count_list = intersection_FDR(selection_coefs, selection_coefs_shuffled, lag = lag, n_features = data.shape[1], fdr_rate = self.fdr_rate) + supports = supports.astype(int) + + + self.l1_support_count = count_list + + + + # np.save("/pscratch/sd/y/yxu2/packages/pyuoi/examples/result/L0_support_after_intersect.npy", supports) + + else: supports = None supports = Bcast_from_root(supports, self.comm, root=0) @@ -321,25 +862,47 @@ def fit(self, X, y, stratify=None, verbose=False): self.supports_ = self.intersect(selection_coefs, self.selection_thresholds_) + + self.n_supports_ = self.supports_.shape[0] + # if rank == 0: + # print(self.supports_.shape, flush = True) + # print(self.supports_, flush = True) + + if rank == 0: + #print(self.n_supports_ ,flush = True) self._logger.info("Found %d supports" % self.n_supports_) + # terminating condition for non-root admm ranks for selection module + + if self.solver == "admm": + n = -1 + n = self.admm_comm.bcast(n, root=0) + + + ##################### # Estimation Module # ##################### # set up data arrays + tasks = np.array_split(np.arange(self.n_boots_est * self.n_supports_), size)[rank] my_boots = dict((task_idx // self.n_supports_, None) for task_idx in tasks) estimates = np.zeros((tasks.size, n_coef)) - + + for boot in range(self.n_boots_est): if size > 1: if rank == 0: - rvals = train_test_split(np.arange(X.shape[0]), + if self.fit_VAR: + idx_set = np.arange(data.shape[0]-lag) + else: + idx_set = np.arange(X.shape[0]) + rvals = train_test_split(idx_set, test_size=1 - self.estimation_frac, stratify=stratify, random_state=self.random_state) @@ -350,26 +913,43 @@ def fit(self, X, y, stratify=None, verbose=False): if boot in my_boots.keys(): my_boots[boot] = rvals else: + if self.fit_VAR: + idx_set = np.arange(data.shape[0]-lag) + else: + idx_set = np.arange(X.shape[0]) my_boots[boot] = train_test_split( - np.arange(X.shape[0]), + idx_set, test_size=1 - self.estimation_frac, stratify=stratify, random_state=self.random_state) # score (r2/AIC/AICc/BIC) for each bootstrap for each support scores = np.zeros(tasks.size) - + # iterate over bootstrap samples and supports + + curr_boot_idx = None for ii, task_idx in enumerate(tasks): + boot_idx = task_idx // self.n_supports_ support_idx = task_idx % self.n_supports_ support = self.supports_[support_idx] # draw a resampled bootstrap - idxs_train, idxs_test = my_boots[boot_idx] - X_rep = X[idxs_train] - y_rep = y[idxs_train] + if curr_boot_idx != boot_idx: + idxs_train, idxs_test = my_boots[boot_idx] + if self.fit_VAR: + X_rep, y_rep, feature_weights = vectorization_bootstrap(data, idxs_train, lag, data_pois = data_pois) + else: + X_rep = X[idxs_train] + y_rep = y[idxs_train] + + curr_boot_idx = boot_idx + + self._logger.info("estimation bootstrap %d, support %d" % (boot_idx, support_idx)) + + if np.any(support): # compute the estimate and store the fitted coefficients @@ -377,29 +957,65 @@ def fit(self, X, y, stratify=None, verbose=False): self._estimation_lm.fit(X_rep[:, support], y_rep) estimates[ii, np.tile(support, self.output_dim)] = \ self._estimation_lm.coef_.ravel() + else: self._estimation_lm.fit(X_rep, y_rep, coef_mask=support) estimates[ii] = self._estimation_lm.coef_.ravel() + if self.estimation_score in ["r2", "acc", "log"]: # using test set + if self.fit_VAR: + X_score, y_score, feature_weights = vectorization_bootstrap(data, idxs_test, lag, data_pois = data_pois) + else: + X_score = X[idxs_test] + y_score = y[idxs_test] + + else: #using training set + X_score, y_score = X_rep, y_rep + scores[ii] = self._score_predictions( metric=self.estimation_score, fitter=self._estimation_lm, - X=X, y=y, - support=support, - boot_idxs=my_boots[boot_idx]) + X=X_score, y=y_score, + support=support) else: + fitter = self._fit_intercept_no_features(y_rep) - if issparse(X): - X_ = csr_matrix(X.shape, dtype=X.dtype) + + if self.estimation_score in ["r2", "acc", "log"]: # using test set + if self.fit_VAR: + + X_score, y_score, feature_weights = vectorization_bootstrap(data, idxs_test, lag, data_pois = data_pois) + else: + X_score = X[idxs_test] + y_score = y[idxs_test] + else: #using training set + X_score, y_score = X_rep, y_rep + + if issparse(X_score): + X_score = csr_matrix(X_score.shape, dtype=X_score.dtype) else: - X_ = np.zeros_like(X) + X_score = np.zeros_like(X_score) + scores[ii] = self._score_predictions( metric=self.estimation_score, fitter=fitter, - X=X_, y=y, - support=np.zeros(X.shape[1], dtype=bool), - boot_idxs=my_boots[boot_idx]) - + X=X_score, y=y_score, + support=np.zeros(X_score.shape[1], dtype=bool)) + + # terminating condition for non-root admm ranks for estimation module + if self.estimation_solver == "admm": + n = 0 + n = self.admm_comm.bcast(n, root=0) + + # clear memory + # if self.fit_VAR: + # del X_rep + # del y_rep + # del X_score + # del y_score + gc.collect() + + if size > 1: estimates = Gatherv_rows(send=estimates, comm=self.comm, root=0) @@ -420,13 +1036,23 @@ def fit(self, X, y, stratify=None, verbose=False): coef = np.median(best_estimates, axis=0).reshape(self.output_dim, n_features) self.coef_ = coef - self._fit_intercept(X, y) + if self.fit_VAR: + if self.fit_intercept: + # the full vectorzation won't do well with large dataset + X_all, y_all, feature_weights = vectorization(data, lag, data_pois = data_pois) + self._fit_intercept(X_all, y_all) + else: + self._fit_intercept(None, None) + else: + self._fit_intercept(X, y) + self.estimates_ = Bcast_from_root(estimates, self.comm, root=0) self.scores_ = Bcast_from_root(scores, self.comm, root=0) - self.coef_ = Bcast_from_root(coef, self.comm, root=0) + self.coef_ = Bcast_from_root(coef, self.comm, root=0) #it's done again down below for recaled VAR coef self.intercept_ = Bcast_from_root(self.intercept_, self.comm, root=0) self.rp_max_idx_ = self.comm.bcast(self.rp_max_idx_, root=0) + else: self.estimates_ = estimates.reshape(self.n_boots_est, self.n_supports_, n_coef) @@ -439,9 +1065,28 @@ def fit(self, X, y, stratify=None, verbose=False): # take the median across estimates for the final, bagged estimate self.coef_ = np.median(best_estimates, axis=0).reshape(self.output_dim, n_features) - self._fit_intercept(X, y) - self._post_fit(X, y) + if self.fit_VAR: + if self.fit_intercept: + X_all, y_all = vectorization(data, lag, data_pois = data_pois) + self._fit_intercept(X_all, y_all) + else: + self._fit_intercept(None, None) + else: + self._fit_intercept(X, y) + + if rank == 0: + if self.fit_VAR: + self._post_fit_VAR(lag, data.shape[1]) + else: + self._post_fit(X, y) + + if size > 1: + if self.fit_VAR: + self.VAR_coef_ = Bcast_from_root(self.VAR_coef_, self.comm, root=0) + else: + self.coef_ = Bcast_from_root(self.coef_, self.comm, root=0) + return self def uoi_selection_sweep(self, X, y, reg_param_values): @@ -472,10 +1117,13 @@ def uoi_selection_sweep(self, X, y, reg_param_values): # apply the selection regression to bootstrapped datasets for reg_param_idx, reg_params in enumerate(reg_param_values): # reset the regularization parameter + self._selection_lm.set_params(**reg_params) + #self._selection_lm.set_params(alpha = 0, l1_ratio= 0) #used to turn off L1-penalty for debugging # rerun fit self._selection_lm.fit(X, y) # store coefficients + coefs[reg_param_idx] = self._selection_lm.coef_.ravel() return coefs @@ -498,13 +1146,14 @@ class AbstractUoILinearRegressor(AbstractUoILinearModel, _default_est_targets = {'r2': 1, 'AIC': 0, 'AICc': 0, 'BIC': 0} - def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, - estimation_frac=0.9, stability_selection=1., + def __init__(self, fit_VAR = False, n_boots_sel=12, n_boots_est=12, selection_frac=0.9, + estimation_frac=0.9, stability_selection=0.75, estimation_score='r2', estimation_target=None, - copy_X=True, fit_intercept=True, + fdr_rate = 0.05, copy_X=True, fit_intercept=True, standardize=True, random_state=None, max_iter=None, tol=None, - comm=None, logger=None): + comm=None, admm_comm = None, logger=None): super(AbstractUoILinearRegressor, self).__init__( + fit_VAR = fit_VAR, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, selection_frac=selection_frac, @@ -514,8 +1163,10 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, standardize=standardize, max_iter=max_iter, tol=tol, + fdr_rate = fdr_rate, random_state=random_state, comm=comm, + admm_comm = admm_comm, logger=logger) if estimation_score not in self._valid_estimation_metrics: @@ -531,9 +1182,12 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, else: estimation_target = self._train_test_map[estimation_target] else: + estimation_target = self._default_est_targets[estimation_score] + self._estimation_target = estimation_target + def _pre_fit(self, X, y): X, y = super()._pre_fit(X, y) if y.ndim == 1: @@ -572,7 +1226,7 @@ def intersect(self, coef, thresholds): def estimation_score(self): return self.__estimation_score - def _score_predictions(self, metric, fitter, X, y, support, boot_idxs): + def _score_predictions(self, metric, fitter, X, y, support): """Score, according to some metric, predictions provided by a model. The resulting score will be negated if an information criterion is @@ -587,9 +1241,9 @@ def _score_predictions(self, metric, fitter, X, y, support, boot_idxs): fitter : object Must contain .predict and .predict_proba methods. X : array-like - The design matrix. + The selected design matrix for score prediction. y : array-like - Response vector. + The selected Response vector for score prediction. supports : array-like The value of the supports for the model boot_idxs : 2-tuple of array-like objects @@ -604,9 +1258,6 @@ def _score_predictions(self, metric, fitter, X, y, support, boot_idxs): The score. """ - # Select the data relevant for the estimation_score - X = X[boot_idxs[self._estimation_target]] - y = y[boot_idxs[self._estimation_target]] if y.ndim == 2: if y.shape[1] > 1: @@ -618,6 +1269,7 @@ def _score_predictions(self, metric, fitter, X, y, support, boot_idxs): '(n_samples, ) or (n_samples, 1).') y_pred = fitter.predict(X[:, support]) + if y.shape != y_pred.shape: raise ValueError('Targets and predictions are not the same shape.') @@ -677,13 +1329,14 @@ class AbstractUoIGeneralizedLinearRegressor(AbstractUoILinearModel, _default_est_targets = {'log': 1, 'AIC': 0, 'AICc': 0, 'BIC': 0, 'acc': 1} - def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, - estimation_frac=0.9, stability_selection=1., + def __init__(self, fit_VAR = False, n_boots_sel=12, n_boots_est=12, selection_frac=0.9, + estimation_frac=0.9, stability_selection=0.75, estimation_score='acc', estimation_target=None, - copy_X=True, fit_intercept=True, standardize=True, + fdr_rate = 0.05, copy_X=True, fit_intercept=True, standardize=True, random_state=None, max_iter=None, tol=None, shared_support=True, comm=None, logger=None): super(AbstractUoIGeneralizedLinearRegressor, self).__init__( + fit_VAR = fit_VAR, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, selection_frac=selection_frac, @@ -694,6 +1347,7 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, standardize=standardize, shared_support=shared_support, max_iter=max_iter, + fdr_rate = fdr_rate, tol=tol, comm=comm, logger=logger) @@ -710,6 +1364,7 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, else: estimation_target = self._train_test_map[estimation_target] else: + estimation_target = self._default_est_targets[estimation_score] self._estimation_target = estimation_target @@ -721,12 +1376,14 @@ def _post_fit(self, X, y): self.intercept_ += np.dot(sX.mean_ * sX.scale_, self.coef_.T) + def intersect(self, coef, thresholds): """Intersect coefficients accross all thresholds. This implementation will account for multi-class classification. """ supports = intersection(coef, thresholds) + if self.output_dim > 1 and self.shared_support: n_features = supports.shape[-1] // self.output_dim supports = supports.reshape((-1, self.output_dim, n_features)) @@ -738,7 +1395,7 @@ def intersect(self, coef, thresholds): def estimation_score(self): return self.__estimation_score - def _score_predictions(self, metric, fitter, X, y, support, boot_idxs): + def _score_predictions(self, metric, fitter, X, y, support): """Score, according to some metric, predictions provided by a model. The resulting score will be negated if an information criterion is @@ -753,9 +1410,9 @@ def _score_predictions(self, metric, fitter, X, y, support, boot_idxs): fitter : object Must contain .predict and .predict_proba methods. X : array-like - The design matrix. + The selected design matrix for score prediction. y : array-like - Response vector. + The selected Response vector for score prediction. supports : array-like The value of the supports for the model boot_idxs : 2-tuple of array-like objects @@ -770,10 +1427,6 @@ def _score_predictions(self, metric, fitter, X, y, support, boot_idxs): The score. """ - # Select the data relevant for the estimation_score - X = X[boot_idxs[self._estimation_target]] - y = y[boot_idxs[self._estimation_target]] - if metric == 'acc': if self.shared_support: y_pred = fitter.predict(X[:, support]) diff --git a/src/pyuoi/linear_model/lasso.py b/src/pyuoi/linear_model/lasso.py index 2bc148ee..030b465b 100644 --- a/src/pyuoi/linear_model/lasso.py +++ b/src/pyuoi/linear_model/lasso.py @@ -9,6 +9,8 @@ pycasso = None from .base import AbstractUoILinearRegressor +from .admm_mpi import ADMM_Lasso +from mpi4py import MPI class PycLasso(): @@ -150,7 +152,7 @@ class UoI_Lasso(AbstractUoILinearRegressor, LinearRegression): Objective used to choose the best estimates per bootstrap. estimation_target : string, "train" | "test" Decide whether to assess the estimation_score on the train - or test data across each bootstrap. By deafult, a sensible + or test data across each bootstrap. By default, a sensible choice is made based on the chosen estimation_score warm_start : bool When set to ``True``, reuse the solution of the previous call to fit as @@ -200,13 +202,14 @@ class UoI_Lasso(AbstractUoILinearRegressor, LinearRegression): boolean array indicating whether a given regressor (column) is selected for estimation for a given regularization parameter value (row). """ - def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, - estimation_frac=0.9, n_lambdas=48, stability_selection=1., + def __init__(self, fit_VAR = False, n_boots_sel=12, n_boots_est=12, selection_frac=0.9, + estimation_frac=0.9, n_lambdas=48, stability_selection=0.75, estimation_score='r2', estimation_target=None, eps=1e-3, warm_start=True, copy_X=True, fit_intercept=True, standardize=True, max_iter=1000, tol=1e-4, random_state=None, - comm=None, logger=None, solver='cd'): + comm=None, global_comm = None, n_admm = None, rho_scaler = 2, imbalance_tolerance = 10, logger=None, solver='cd', estimation_solver = "ls"): super(UoI_Lasso, self).__init__( + fit_VAR = fit_VAR, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, selection_frac=selection_frac, @@ -224,8 +227,19 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, logger=logger) self.n_lambdas = n_lambdas self.eps = eps - self.solver = solver + self.solver = solver # solver for selection module + self.estimation_solver = estimation_solver # solver for estimation module self.tol = tol + self.global_comm = global_comm + self.n_admm = n_admm + self.total_time = 0 + self.comm_time = 0 + self.compute_time1 = 0 + self.compute_time2 = 0 + self.compute_time3 = 0 + self.kron_time = 0 + + if solver == 'cd': self._selection_lm = Lasso( @@ -242,7 +256,58 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, selection_frac=0.9, max_iter=max_iter, tol=tol) - self._estimation_lm = LinearRegression(fit_intercept=fit_intercept) + elif solver == "admm": + + #every rank will have their corresponding admm_comm!!! + admm_comm_list = [] + admm_group_list = [] + + # global rank + rank = self.global_comm.rank + # idx of the correponding bootstrap + boot_rank = rank//n_admm + + + # all ranks of the corresponding admm group for that given bootstrap + admm_ranks_list = [np.arange(boot_rank*n_admm, (boot_rank+1)*n_admm) for boot_rank in range(int(self.global_comm.Get_size()//n_admm))] + + for admm_ranks in admm_ranks_list: + + admm_group_list.append(self.global_comm.group.Incl(admm_ranks)) + admm_comm_list.append(self.global_comm.Create(admm_group_list[-1])) + + + self.admm_comm = admm_comm_list[boot_rank] + + self._selection_lm = ADMM_Lasso( + comm = self.admm_comm, + max_iter=max_iter, + abs_tol=tol, + rel_tol = tol/10, + rho_scaler = rho_scaler, + imbalance_tolerance = imbalance_tolerance, + warm_start=warm_start, + random_state=random_state, + fit_intercept=fit_intercept) + + + + if estimation_solver == "admm": + self._estimation_lm = ADMM_Lasso( + comm = self.admm_comm, + max_iter=max_iter, + abs_tol=tol, + rel_tol = tol/10, + rho_scaler = rho_scaler, + imbalance_tolerance = imbalance_tolerance, + warm_start=False, + random_state=random_state, + fit_intercept=fit_intercept) + self._estimation_lm.set_params(alpha=0) + elif estimation_solver == "ls": + self._estimation_lm = LinearRegression(fit_intercept=fit_intercept) + + def get_reg_params(self, X, y): alphas = _alpha_grid( @@ -250,14 +315,38 @@ def get_reg_params(self, X, y): l1_ratio=1.0, fit_intercept=self.fit_intercept, eps=self.eps, - n_alphas=self.n_lambdas) + n_alphas=self.n_lambdas, copy_X = True) return [{'alpha': a} for a in alphas] + + + def admm_queue(self): + + if self.solver == "admm": + while True: + #the first bcast call in fit is blocking, so the non-root process will wait for the root to distribute data + self._selection_lm.fit() + #runtime scaling results + # if self.global_comm.rank == 100: + # self.total_time+=self._selection_lm.total_time + # self.comm_time+=self._selection_lm.comm_time + # self.compute_time+=self._selection_lm.compute_time + if self._selection_lm.terminate_selection: + break + + if self.estimation_solver == "admm": + while True: + self._estimation_lm.fit() + if self._estimation_lm.terminate_estimation: + break + + def uoi_selection_sweep(self, X, y, reg_param_values): """Overwrite base class selection sweep to accommodate pycasso path-wise solution""" + if self.solver == 'pyc': alphas = np.array([reg_param['alpha'] for reg_param in reg_param_values]) @@ -265,6 +354,34 @@ def uoi_selection_sweep(self, X, y, reg_param_values): self._selection_lm.fit(X, y) return self._selection_lm.coef_ + + elif self.solver == 'admm': + + n_param_values = len(reg_param_values) + n_coef = self.get_n_coef(X, y) + coefs = np.zeros((n_param_values, n_coef)) + + # apply the selection regression to bootstrapped datasets + for reg_param_idx, reg_params in enumerate(reg_param_values): + # reset the regularization parameter + + self._selection_lm.set_params(**reg_params) + # rerun fit + self._selection_lm.fit(X, y, param_mask = self.param_mask) + + #runtime scaling results + if self.comm.rank == 0: + self.total_time+=self._selection_lm.total_time + self.comm_time+=self._selection_lm.comm_time + self.compute_time1+=self._selection_lm.compute_time1 + self.compute_time2+=self._selection_lm.compute_time2 + self.compute_time3+=self._selection_lm.compute_time3 + + # store coefficients + coefs[reg_param_idx] = self._selection_lm.coef_.ravel() + + return coefs + else: return super(UoI_Lasso, self).uoi_selection_sweep(X, y, reg_param_values) diff --git a/src/pyuoi/linear_model/logistic.py b/src/pyuoi/linear_model/logistic.py index f3b1da7e..97b8550b 100644 --- a/src/pyuoi/linear_model/logistic.py +++ b/src/pyuoi/linear_model/logistic.py @@ -7,7 +7,7 @@ from sklearn.utils import (check_X_y, compute_class_weight, check_consistent_length, check_array) from sklearn.utils.multiclass import check_classification_targets -from sklearn.utils.extmath import safe_sparse_dot, log_logistic, squared_norm +from sklearn.utils.extmath import safe_sparse_dot, squared_norm from sklearn.preprocessing import StandardScaler from scipy.optimize import minimize @@ -20,6 +20,41 @@ from ..utils import sigmoid, softmax from ..lbfgs import fmin_lbfgs, AllZeroLBFGSError +def stable_log_logistic(X): + """ + Compute log(1 / (1 + exp(-x))) in a numerically stable way. + + This implementation handles both positive and negative inputs safely + by using the identity: log(1 / (1 + exp(-x))) = -log(1 + exp(-x)) + + For large negative values, we use the approximation log(1 + exp(-x)) ≈ -x + For large positive values, we use the approximation log(1 + exp(-x)) ≈ 0 + + Parameters + ---------- + X : array-like + Input array + + Returns + ------- + array-like + Log-logistic transformation of input + """ + X = np.asarray(X) + + # Initialize output array + out = np.zeros_like(X, dtype=np.float64) + + # Handle positive values + pos_mask = X > 0 + out[pos_mask] = -np.log1p(np.exp(-X[pos_mask])) + + # Handle negative values + neg_mask = ~pos_mask + out[neg_mask] = X[neg_mask] - np.log1p(np.exp(X[neg_mask])) + + return out + class UoI_L1Logistic(AbstractUoIGeneralizedLinearRegressor, LogisticRegression): r"""UoI\ :sub:`L1-Logistic` model. @@ -823,7 +858,7 @@ def _logistic_loss_and_grad(w, X, y, alpha, mask, sample_weight=None): sample_weight = np.ones(n_samples) # Logistic loss is the negative of the log of the logistic function. - out = -np.sum(sample_weight * log_logistic(yz)) / n_samples + out = -np.sum(sample_weight * stable_log_logistic(yz)) / n_samples out += .5 * alpha * np.dot(w, w) z = expit(yz) diff --git a/src/pyuoi/linear_model/poisson.py b/src/pyuoi/linear_model/poisson.py index 3fa31030..8fda96df 100644 --- a/src/pyuoi/linear_model/poisson.py +++ b/src/pyuoi/linear_model/poisson.py @@ -9,6 +9,58 @@ from sklearn.utils.validation import check_is_fitted from ..lbfgs import fmin_lbfgs +from scipy.optimize import fmin_l_bfgs_b +from .admm_mpi_poisson import ADMM_Poisson +from mpi4py import MPI +from scipy import sparse + +def _pack_coefficients(A, b): + """ + Pack A matrix and b vector into coefficient vector. + + Parameters + ---------- + A : (p, p) array + b : (p,) array + + Returns + ------- + coefficients : (p*(p+1),) array + """ + p = A.shape[0] + coefficients = np.zeros(p * (p + 1)) + + for i in range(p): + start_idx = i * (p + 1) + coefficients[start_idx] = b[i] + coefficients[start_idx + 1 : start_idx + p + 1] = A[i, :] + + return coefficients + +def _unpack_coefficients(coefficients): + """ + Unpack coefficient vector into A matrix and b vector. + + Parameters + ---------- + coefficients : (p*(p+1),) array + Flattened coefficients [b_1, A_{1,:}, b_2, A_{2,:}, ...] + + Returns + ------- + A : (p, p) array + b : (p,) array + """ + p = 20 + A = np.zeros((p, p)) + b = np.zeros(p) + + for i in range(p): + start_idx = i * (p + 1) + b[i] = coefficients[start_idx] + A[i, :] = coefficients[start_idx + 1 : start_idx + p + 1] + + return A, b class Poisson(BaseEstimator): @@ -52,7 +104,7 @@ class Poisson(BaseEstimator): The fitted intercept. """ def __init__(self, alpha=1.0, l1_ratio=1., fit_intercept=True, - standardize=False, max_iter=1000, tol=1e-5, warm_start=False, + standardize=False, max_iter=1000, tol=1e-5, warm_start=False, feature_weights = 1, dt = 1, solver='lbfgs'): self.alpha = alpha self.l1_ratio = l1_ratio @@ -62,6 +114,8 @@ def __init__(self, alpha=1.0, l1_ratio=1., fit_intercept=True, self.tol = tol self.warm_start = warm_start self.solver = solver + self.feature_weights = feature_weights + self.dt = dt def fit(self, X, y, sample_weight=None): """Fit the Poisson GLM. @@ -77,6 +131,11 @@ def fit(self, X, y, sample_weight=None): Array of weights assigned to the individual samples. If ``None``, then each sample is provided an equal weight. """ + + ## this is for fittign VAR only + weights = np.tile(self.feature_weights, (int(len(y)/len(self.feature_weights)), 1)) + sample_weight = weights.T.flatten() + self.n_samples, self.n_features = X.shape X, y = self._pre_fit(X, y) @@ -105,22 +164,80 @@ def fit(self, X, y, sample_weight=None): coef = np.append(coef, intercept) # create lbfgs function - def func(x, g, *args): + def func(x, *args): loss, grad = _poisson_loss_and_grad(x, *args) - g[:] = grad return loss + def grad(x, *args): + loss, grad = _poisson_loss_and_grad(x, *args) + return grad + l1_penalty = self.alpha * self.l1_ratio l2_penalty = self.alpha * (1 - self.l1_ratio) - # orthant-wise lbfgs optimization - coef = fmin_lbfgs(func, coef, - orthantwise_c=l1_penalty, - args=(X, y, l2_penalty, sample_weight), - max_iterations=self.max_iter, - epsilon=self.tol, - orthantwise_end=self.n_features) + # # create lbfgs function + # def func(x, g, *args): + # loss, grad = _poisson_loss_and_grad(x, *args) + # g[:] = grad + # return loss + # orthant-wise lbfgs optimization + # coef = fmin_lbfgs(func, coef, + # orthantwise_c=l1_penalty, + # args=(X, y, l2_penalty, sample_weight), + # max_iterations=self.max_iter, + # epsilon=self.tol, + # orthantwise_end=self.n_features) + + # Initialize coefficients + + + # p = 20 + # A_init = np.random.randn(p, p) * 0.01 + # b_init = np.zeros(p) - 1.0 + + # coefficients_init = _pack_coefficients(A_init, b_init) + + # # Set up bounds if not provided + + # bounds = [] + # for i in range(p): + # # Intercept bounds + # bounds.append((-10.0, 0)) + # # A matrix row bounds (encourage negative diagonal) + # for j in range(p): + # if i == j: + # bounds.append((-5.0, 0.5)) # Diagonal + # else: + # bounds.append((-5.0, 5.0)) # Off-diagonal + + # Optimize using L-BFGS-B + + + result = fmin_l_bfgs_b( + func=func, + x0=coef, + fprime=grad, # Providing gradient speeds up convergence + args=(X, y, l2_penalty, self.dt, sample_weight), + bounds=None, # No bounds for fully dense model + maxiter=self.max_iter, + maxfun=15000, + pgtol=1e-5, + factr=1e7, # Higher precision: use factr=10 for very high precision + ) + + + coef, min_nll, info = result + + + + # A,b = _unpack_coefficients(coef) + + #print(np.mean(X), np.mean(b), flush = True) + # np.save("result/real_b", b) + # np.save("result/real_A", A) + + if self.fit_intercept: self.coef_ = coef[:self.n_features] self.intercept_ = coef[-1] @@ -155,8 +272,14 @@ def predict(self, X): mode : array_like, shape (n_samples) The predicted response values, i.e. the modes. """ + if sparse.issparse(X): + dot = lambda X, coef: X.dot(coef) # sparse matrix-vector multiplication + else: + dot = lambda X, coef: np.dot(X, coef) + check_is_fitted(self, ['coef_', 'intercept_']) - mu = np.exp(self.intercept_ + np.dot(X, self.coef_)) + #mu = np.exp(np.clip(self.intercept_ + dot(X, self.coef_), -5, 5))*self.dt + mu = np.exp(self.intercept_ + dot(X, self.coef_))*self.dt mode = np.floor(mu) return mode @@ -173,8 +296,21 @@ def predict_mean(self, X): mu : array_like, shape (n_samples) The predicted response values, i.e. the conditional means. """ - check_is_fitted(self, ['coef_', 'intercept_']) - mu = np.exp(self.intercept_ + np.dot(X, self.coef_)) + + + if sparse.issparse(X): + dot = lambda X, coef: X.dot(coef) # sparse matrix-vector multiplication + else: + dot = lambda X, coef: np.dot(X, coef) + + if self.fit_intercept: + check_is_fitted(self, ['coef_', 'intercept_']) + #mu = np.exp(np.clip(self.intercept_ + dot(X, self.coef_), -5, 5))*self.dt + mu = np.exp(self.intercept_ + dot(X, self.coef_))*self.dt + else: + check_is_fitted(self, ['coef_']) + #mu = np.exp(np.clip(dot(X, self.coef_), -5, 5))*self.dt + mu = np.exp(dot(X, self.coef_))*self.dt return mu def _cd(self, X, y, sample_weight=None): @@ -258,7 +394,7 @@ def _cd_sweep(self, coef, X, w, z, active_idx, intercept=0): coef : ndarray, shape (n_features,) The current estimates of the parameters. X : ndarray, shape (n_samples, n_features) - The design matrix. + The design matrix.f w : ndarray, shape (n_samples,) The weights applied to each sample after linearization of the log-likelihood. @@ -310,7 +446,7 @@ def _cd_sweep(self, coef, X, w, z, active_idx, intercept=0): def _pre_fit(self, X, y): """Perform standardization, if needed, before fitting.""" if self.standardize: - self._X_scaler = StandardScaler() + self._X_scaler = StandardScaler(with_mean=self.fit_intercept) X = self._X_scaler.fit_transform(X) return X, y @@ -459,15 +595,17 @@ class UoI_Poisson(AbstractUoIGeneralizedLinearRegressor, Poisson): Boolean array indicating whether a given regressor (column) is selected for estimation for a given regularization parameter value (row). """ - def __init__(self, n_boots_sel=24, n_boots_est=24, n_lambdas=48, + def __init__(self, fit_VAR = False, n_boots_sel=12, n_boots_est=12, n_lambdas=48, alphas=np.array([1.]), selection_frac=0.8, - estimation_frac=0.8, stability_selection=1., + estimation_frac=0.8, stability_selection=0.75, + manual_l1_range = None, estimation_score='log', estimation_target=None, - solver='lbfgs', warm_start=True, - eps=1e-3, tol=1e-5, copy_X=True, fit_intercept=True, - standardize=True, max_iter=1000, - random_state=None, comm=None, logger=None): + solver='lbfgs', estimation_solver = 'lbfgs', warm_start=True, + eps=1e-3, tol=1e-8, fit_intercept=True, + standardize=False, max_iter=1000, + random_state=None, comm=None, global_comm = None, n_admm = None, rho_scaler = 2, imbalance_tolerance = 10, l1_suppression = 0, weights = 1, dt = 1, fdr_rate = 0.05, logger=None): super(UoI_Poisson, self).__init__( + fit_VAR = fit_VAR, n_boots_sel=n_boots_sel, n_boots_est=n_boots_est, selection_frac=selection_frac, @@ -475,8 +613,9 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, n_lambdas=48, stability_selection=stability_selection, estimation_score=estimation_score, estimation_target=estimation_target, - copy_X=copy_X, fit_intercept=fit_intercept, + fdr_rate = fdr_rate, + standardize=False, # hard code Z-score scaling to false b/c it's not neede for Poisson GLM random_state=random_state, comm=comm, logger=logger) @@ -485,23 +624,88 @@ def __init__(self, n_boots_sel=24, n_boots_est=24, n_lambdas=48, self.n_alphas = len(alphas) self.warm_start = warm_start self.eps = eps + self.solver = solver + self.estimation_solver = estimation_solver self.lambdas = None - self._selection_lm = Poisson( - fit_intercept=fit_intercept, - standardize=standardize, - max_iter=max_iter, - tol=tol, - warm_start=warm_start, - solver=solver) + self.global_comm = global_comm + self.n_admm = n_admm + self.l1_suppression = l1_suppression + self.dt =dt + self.manual_l1_range = manual_l1_range + + if solver == "admm": + + + #every rank will have their corresponding admm_comm!!! + admm_comm_list = [] + admm_group_list = [] + + # global rank + rank = self.global_comm.rank + # idx of the correponding bootstrap + boot_rank = rank//n_admm + + + # all ranks of the corresponding admm group for that given bootstrap + admm_ranks_list = [np.arange(boot_rank*n_admm, (boot_rank+1)*n_admm) for boot_rank in range(int(self.global_comm.Get_size()//n_admm))] + + for admm_ranks in admm_ranks_list: + + admm_group_list.append(self.global_comm.group.Incl(admm_ranks)) + admm_comm_list.append(self.global_comm.Create(admm_group_list[-1])) + + + self.admm_comm = admm_comm_list[boot_rank] + + self._selection_lm = ADMM_Poisson( + comm = self.admm_comm, + max_iter=max_iter, + abs_tol=tol, + rel_tol = tol/10, + rho_scaler = rho_scaler, + feature_weights = weights, + dt = dt, + imbalance_tolerance = imbalance_tolerance, + warm_start=warm_start, + random_state=random_state, + fit_intercept=fit_intercept) + else: + self._selection_lm = Poisson( + fit_intercept=fit_intercept, + max_iter=max_iter, + tol=tol, + warm_start=warm_start, + feature_weights = weights, + dt = dt, + solver=solver) + # estimation is a Poisson regression with no regularization - self._estimation_lm = Poisson( - alpha=0., - l1_ratio=1., - fit_intercept=fit_intercept, - max_iter=max_iter, - tol=tol, - warm_start=False, - solver=solver) + if estimation_solver == "admm": + self._estimation_lm = ADMM_Poisson( + comm = self.admm_comm, + max_iter=max_iter, + abs_tol=tol, + rel_tol = tol/10, + rho_scaler = rho_scaler, + feature_weights = weights, + dt = dt, + imbalance_tolerance = imbalance_tolerance, + warm_start=False, + random_state=random_state, + fit_intercept=fit_intercept) + self._estimation_lm.set_params(alpha = 0, l1_ratio= 0) + + else: + self._estimation_lm = Poisson( + alpha=0., + l1_ratio=1., + fit_intercept=fit_intercept, + max_iter=max_iter, + tol=tol, + warm_start=False, + feature_weights = weights, + dt = dt, + solver=estimation_solver) def get_reg_params(self, X, y): r"""Calculates the regularization parameters (alpha and lambda) to be @@ -525,18 +729,42 @@ def get_reg_params(self, X, y): """ n_samples = X.shape[0] if self.lambdas is None: - self.lambdas = np.zeros((self.n_alphas, self.n_lambdas)) + self.lambdas = np.zeros((self.n_alphas, self.n_lambdas-self.l1_suppression)) # a set of lambdas are generated for each alpha value (l1_ratio in # sci-kit learn parlance) for alpha_idx, alpha in enumerate(self.alphas): - # calculate upper bound for lambda sweep - ybar = y.mean() - lambda_max = np.max(np.abs(np.dot(X.T, y - ybar))) - lambda_max /= n_samples * alpha - self.lambdas[alpha_idx, :] = np.logspace( - start=np.log10(lambda_max), - stop=np.log10(self.eps * lambda_max), - num=self.n_lambdas) + if self.manual_l1_range is None: + # calculate upper bound for lambda sweep + ybar = y.mean() + lambda_max = np.max(np.abs(np.dot(X.T, y - ybar))) + lambda_max /= n_samples * alpha + + # Estimate zero-inflation level + observed_zeros = (X == 0).mean() # ~0.9 in your case + expected_poisson_zeros = np.exp(-X[X>0].mean()) + # Adjust λ_max based on excess zeros + excess_zeros = max(0, observed_zeros - expected_poisson_zeros) + lambda_adjustment = 1 + excess_zeros + # lambda_max = lambda_max * lambda_adjustment + + + + self.lambdas[alpha_idx, :] = np.logspace( + start=np.log10(lambda_max), + stop=np.log10(self.eps * lambda_max), + num=self.n_lambdas)[self.l1_suppression:] + + + else: + # Hard coded L1-path range + self.lambdas[alpha_idx, :] = np.logspace( + start=np.log10(self.manual_l1_range[1]), #max + stop=np.log10(self.manual_l1_range[0]), #min + num=self.n_lambdas)[self.l1_suppression:] + + + # self.lambdas[alpha_idx, :] /= X.shape[1]**2 + # print(self.lambdas, flush = True) # place the regularization parameters into a list of dictionaries reg_params = list() @@ -547,6 +775,21 @@ def get_reg_params(self, X, y): return reg_params + def admm_queue(self): + + if self.solver == "admm": + while True: + #the first bcast call in fit is blocking, so the non-root process will wait for the root to distribute data + self._selection_lm.fit() + if self._selection_lm.terminate_selection: + break + + if self.estimation_solver == "admm": + while True: + self._estimation_lm.fit() + if self._estimation_lm.terminate_estimation: + break + def _score_predictions(self, metric, fitter, X, y, support, boot_idxs=None): """Score, according to some metric, predictions provided by a model. @@ -580,12 +823,14 @@ def _score_predictions(self, metric, fitter, X, y, support, boot_idxs=None): The score. """ - # Select the train data - if boot_idxs is not None: - X = X[boot_idxs[self._estimation_target]] - y = y[boot_idxs[self._estimation_target]] + # # Select the train data + # OBSOLETE: from older version + # if boot_idxs is not None: + # X = X[boot_idxs[self._estimation_target]] + # y = y[boot_idxs[self._estimation_target]] # for Poisson, use predict_mean to calculate the "predicted" values + y_pred = fitter.predict_mean(X[:, support]) # calculate the log-likelihood ll = utils.log_likelihood_glm(model='poisson', y_true=y, y_pred=y_pred) @@ -655,6 +900,35 @@ def _fit_intercept_no_features(self, y): """ return PoissonInterceptFitterNoFeatures(y) + def uoi_selection_sweep(self, X, y, reg_param_values): + """Overwrite base class selection sweep to accommodate pycasso + path-wise solution""" + + + if self.solver == 'admm': + + n_param_values = len(reg_param_values) + n_coef = self.get_n_coef(X, y) + coefs = np.zeros((n_param_values, n_coef)) + + # apply the selection regression to bootstrapped datasets + for reg_param_idx, reg_params in enumerate(reg_param_values): + # reset the regularization parameter + + self._selection_lm.set_params(**reg_params) + # rerun fit + self._selection_lm.fit(X, y, param_mask = self.param_mask) + # store coefficients + coefs[reg_param_idx] = self._selection_lm.coef_.ravel() + + self.loss = self._selection_lm.loss + # np.save("result/loss_"+str(reg_params), loss) + + return coefs + + else: + return super(UoI_Poisson, self).uoi_selection_sweep(X, y, + reg_param_values) class PoissonInterceptFitterNoFeatures(object): def __init__(self, y): @@ -663,7 +937,7 @@ def __init__(self, y): else: self.intercept_ = -np.inf - def predict(self, X): + def predict(self, X, dt = 1): """Predicts the response variable given a design matrix. The output is the mode of the Poisson distribution. @@ -677,11 +951,11 @@ def predict(self, X): mode : array_like, shape (n_samples) The predicted response values, i.e. the modes. """ - mu = np.exp(self.intercept_) + mu = np.exp(np.clip(self.intercept_, -5, 5))*dt mode = np.floor(mu) return mode - def predict_mean(self, X): + def predict_mean(self, X, dt = 1): """Calculates the mean response variable given a design matrix. Parameters @@ -694,19 +968,20 @@ def predict_mean(self, X): mu : array_like, shape (n_samples) The predicted response values, i.e. the conditional means. """ - mu = np.exp(self.intercept_) + mu = np.exp(np.clip(self.intercept_, -5, 5))*dt return mu -def _poisson_loss_and_grad(coef, X, y, l2_penalty, sample_weight=None): - """Computes the Poisson loss and gradient. +def _poisson_loss_and_grad(coef, X, y, l2_penalty = 0, dt = 1, sample_weight=None): + """Computes the Poisson loss and gradient with sparse matrix support. + Parameters ---------- coef : ndarray, shape (n_features,) or (n_features + 1,) Coefficient vector. - X : array-like, shape (n_samples, n_features) - Design matrix. + X : array-like or scipy.sparse matrix, shape (n_samples, n_features) + Design matrix. Can be dense or sparse. y : ndarray, shape (n_samples,) Response vector. l2_penalty : float @@ -714,7 +989,7 @@ def _poisson_loss_and_grad(coef, X, y, l2_penalty, sample_weight=None): sample_weight : array-like, shape (n_samples,), default None Array of weights assigned to the individual samples. If None, then each sample is provided an equal weight. - + Returns ------- out : float @@ -723,29 +998,114 @@ def _poisson_loss_and_grad(coef, X, y, l2_penalty, sample_weight=None): Poisson gradient. """ n_samples, n_features = X.shape - grad = np.empty_like(coef) + n_samples = 1 # overriding to give total error instead of mean error + grad = np.empty_like(coef) + if sample_weight is None: - sample_weight = np.ones(n_samples) - + #sample_weight = np.ones(n_samples) + sample_weight = 1 + # extract intercept if grad.shape[0] > n_features: intercept = coef[-1] coef = coef[:n_features] else: intercept = 0 - + + # calculate linear predictor (eta) + if sparse.issparse(X): + eta = intercept + X.dot(coef) # sparse matrix-vector multiplication + else: + eta = intercept + np.dot(X, coef) + # calculate likelihood - eta = intercept + np.dot(X, coef) - out = -np.sum(sample_weight * (y * eta - np.exp(eta))) / n_samples + out = -np.sum(sample_weight * (y * eta + np.log(dt) - np.exp(eta)*dt)) / n_samples out += 0.5 * l2_penalty * np.dot(coef, coef) - + + # calculate residuals + y_res = sample_weight * (y - np.exp(eta)*dt) + # gradient of parameters - y_res = sample_weight * (y - np.exp(eta)) - grad[:n_features] = -np.dot(X.T, y_res) / n_samples + l2_penalty * coef - + if sparse.issparse(X): + # For sparse matrices, use X.T.dot() for efficient transpose multiplication + grad[:n_features] = -X.T.dot(y_res) / n_samples + l2_penalty * coef + else: + grad[:n_features] = -np.dot(X.T, y_res) / n_samples + l2_penalty * coef + # gradient of intercept if grad.shape[0] > n_features: grad[-1] = -np.mean(y_res) - + return out, grad + +def _negative_log_likelihood(coefficients, design_matrix, y): + """ + Compute negative log-likelihood in vectorized form. + + Parameters + ---------- + coefficients : (p*(p+1),) array + Model coefficients + design_matrix : sparse array + Design matrix + y : array + Vectorized response + + Returns + ------- + nll : float + Negative log-likelihood + """ + # Compute linear predictor + eta = design_matrix @ coefficients + + # Clip to prevent overflow + eta = np.clip(eta, -20, np.log(1e6)) + + # Compute rates + lam = np.exp(eta) + + # Negative log-likelihood: -y * eta + lambda + nll = np.sum(-y * eta + lam) + + + + return nll + +def _gradient(coefficients, design_matrix, y): + """ + Compute gradient of negative log-likelihood. + + Parameters + ---------- + coefficients : (p*(p+1),) array + Model coefficients + design_matrix : sparse array + Design matrix + y : array + Vectorized response + + Returns + ------- + grad : (p*(p+1),) array + Gradient vector + """ + # Compute linear predictor + eta = design_matrix @ coefficients + + # Clip to prevent overflow + eta = np.clip(eta, -20, np.log(1e6)) + + # Compute rates + lam = np.exp(eta) + + # Gradient w.r.t. eta: -y + lambda + grad_eta = -y + lam + + # Gradient w.r.t. coefficients + grad = design_matrix.T @ grad_eta + + + return grad + diff --git a/src/pyuoi/linear_model/sparse_comm_util.py b/src/pyuoi/linear_model/sparse_comm_util.py new file mode 100644 index 00000000..e4108fff --- /dev/null +++ b/src/pyuoi/linear_model/sparse_comm_util.py @@ -0,0 +1,72 @@ +import numpy as np +import scipy.sparse as sp +from mpi4py import MPI + +def send_sparse_matrix(comm, sparse_matrix, dest, tag=0): + """Send a sparse matrix to a destination process.""" + # Extract the components of the sparse matrix + data = sparse_matrix.data.astype(np.float64) # Ensure consistent dtype + indices = sparse_matrix.indices.astype(np.int32) + indptr = sparse_matrix.indptr.astype(np.int32) + shape = np.array(sparse_matrix.shape, dtype=np.int64) + + # Send shape first so destination knows what to expect + comm.Send([shape, MPI.INT64_T], dest=dest, tag=tag) + + # Send data array size, then data + data_size = np.array([data.size], dtype=np.int64) + comm.Send([data_size, MPI.INT64_T], dest=dest, tag=tag+1) + comm.Send([data, MPI.DOUBLE], dest=dest, tag=tag+2) + + # Send indices size, then indices + indices_size = np.array([indices.size], dtype=np.int64) + comm.Send([indices_size, MPI.INT64_T], dest=dest, tag=tag+3) + comm.Send([indices, MPI.INT32_T], dest=dest, tag=tag+4) + + # Send indptr size, then indptr + indptr_size = np.array([indptr.size], dtype=np.int64) + comm.Send([indptr_size, MPI.INT64_T], dest=dest, tag=tag+5) + comm.Send([indptr, MPI.INT32_T], dest=dest, tag=tag+6) + +def recv_sparse_matrix(comm, source, tag=0): + """Receive a sparse matrix from a source process.""" + # Receive shape first + shape = np.empty(2, dtype=np.int64) + comm.Recv([shape, MPI.INT64_T], source=source, tag=tag) + + # Receive data array size, then data + data_size = np.empty(1, dtype=np.int64) + comm.Recv([data_size, MPI.INT64_T], source=source, tag=tag+1) + data = np.empty(data_size[0], dtype=np.float64) + comm.Recv([data, MPI.DOUBLE], source=source, tag=tag+2) + + # Receive indices size, then indices + indices_size = np.empty(1, dtype=np.int64) + comm.Recv([indices_size, MPI.INT64_T], source=source, tag=tag+3) + indices = np.empty(indices_size[0], dtype=np.int32) + comm.Recv([indices, MPI.INT32_T], source=source, tag=tag+4) + + # Receive indptr size, then indptr + indptr_size = np.empty(1, dtype=np.int64) + comm.Recv([indptr_size, MPI.INT64_T], source=source, tag=tag+5) + indptr = np.empty(indptr_size[0], dtype=np.int32) + comm.Recv([indptr, MPI.INT32_T], source=source, tag=tag+6) + + # Reconstruct the sparse matrix + return sp.csr_matrix((data, indices, indptr), shape=tuple(shape)) + + + +def build_bootstrap_comm(comm, n_admm): + + size = comm.Get_size() # should be n_boot*n_admm (or better n_boot*n_reg*n_admm) + + boot_rank = np.arange(size)[::n_admm] + + boot_group = comm.group.Incl(list(boot_rank)) + # every rank in boot_rank will now have a different comm called boot_comm, the other ranks will have COMM_NULL for boot_comm + boot_comm = comm.Create(boot_group) + + if boot_comm == MPI.COMM_NULL: + boot_comm = None + return boot_comm \ No newline at end of file diff --git a/src/pyuoi/linear_model/utils.py b/src/pyuoi/linear_model/utils.py index 6b26a167..671c9439 100644 --- a/src/pyuoi/linear_model/utils.py +++ b/src/pyuoi/linear_model/utils.py @@ -1,5 +1,5 @@ import numpy as np - +from statsmodels.stats.multitest import fdrcorrection def stability_selection_to_threshold(stability_selection, n_boots): """Converts user inputted stability selection to an array of @@ -80,7 +80,7 @@ def stability_selection_to_threshold(stability_selection, n_boots): return selection_thresholds -def intersection(coefs, selection_thresholds=None): +def intersection(coefs, selection_thresholds=None, magnitude_threshold=0.01): """Performs the intersection operation on selection coefficients using stability selection criteria. @@ -127,11 +127,21 @@ def intersection(coefs, selection_thresholds=None): dtype=bool ) + # Count how many bootstraps have |coefficient| >= magnitude_threshold + # for each lambda and feature + significant_selections = np.sum( + np.abs(coefs) > magnitude_threshold, + axis=0 + ) + + + # significant_selections = np.count_nonzero(coefs, axis=0) + + # iterate over each stability selection threshold for thresh_idx, threshold in enumerate(selection_thresholds): # calculate the support given the specific selection threshold - supports[thresh_idx, ...] = \ - np.count_nonzero(coefs, axis=0) >= threshold + supports[thresh_idx, ...] = significant_selections >= threshold # unravel the dimension corresponding to selection thresholds @@ -143,3 +153,349 @@ def intersection(coefs, selection_thresholds=None): supports = np.unique(supports, axis=0) return supports + +def intersection_dirty(coefs, selection_thresholds=None, n_sig=2.5, min_w=0.01): + """Performs intersection operation using both minimum weight and SNR criteria. + + A feature is considered "selected" only if: + 1. Its upper confidence bound (mean + n_sig*std) exceeds min_w (not negligible) + 2. Its signal-to-noise ratio (mean/std) exceeds n_sig (consistent) + + Parameters + ---------- + coefs : np.ndarray, shape (# bootstraps, # lambdas, # features) + The coefficients obtained from the selection sweep. + + selection_thresholds: array-like, int + The selection thresholds to perform intersection across. + + n_sig : float, default=2.5 + Signal-to-noise ratio threshold and confidence interval multiplier. + Features must have |mean|/std > n_sig AND |mean| + n_sig*std > min_w. + + min_w : float, default=0.05 + Minimum weight threshold. Features with upper confidence bound + below this are considered negligible regardless of SNR. + + Returns + ------- + supports : np.ndarray, shape (# supports, # features), bool + Unique supports obtained by performing the intersection. + """ + + if selection_thresholds is None: + selection_thresholds = np.array([1]) + + n_selection_thresholds = len(selection_thresholds) + n_reg_params = coefs.shape[1] + n_features = coefs.shape[2] + supports = np.zeros( + (n_selection_thresholds, n_reg_params, n_features), + dtype=bool + ) + + # Calculate statistics across bootstraps + coef_mean = np.mean(coefs, axis=0) # Shape: (n_lambdas, n_features) + coef_std = np.std(coefs, axis=0) # Shape: (n_lambdas, n_features) + coef_mean_abs = np.abs(coef_mean) + + # Apply dual criteria for each lambda and feature + for lambda_idx in range(n_reg_params): + for feature_idx in range(n_features): + xmean = coef_mean_abs[lambda_idx, feature_idx] + xstd = coef_std[lambda_idx, feature_idx] + + # Condition 1: Check if upper confidence bound exceeds minimum weight + # This filters out consistently small coefficients + if xmean + n_sig * xstd < min_w: + continue + + # Condition 2: Check signal-to-noise ratio + # This filters out inconsistent coefficients + if xstd > 1e-10: # Avoid division by zero + if xmean / xstd < n_sig: + continue + else: + # If std is essentially zero, only check mean against min_w + if xmean < min_w: + continue + + # Both conditions passed - mark as selected + for thresh_idx in range(n_selection_thresholds): + supports[thresh_idx, lambda_idx, feature_idx] = True + + # Unravel and get unique supports + supports = np.squeeze(np.reshape( + supports, + (n_selection_thresholds * n_reg_params, n_features) + )) + + supports = np.unique(supports, axis=0) + + return supports + +def ie_type(coef, lag, num_feat, scheme = 2): + A_model = coef.reshape(coef.shape[:-1] + (num_feat,lag,num_feat)) # Reshape last axis from to the connectivity matrix + # shape: n_boot * n_reg_param * lag * n_feat * n_feat + A_model = np.transpose(A_model, (0,1,3,2,4)) #switch the axis on the lag and n_feat(because we column stacked) + + # shape: lag * n_boot * n_reg_param * n_feat * n_feat (transpose is so useful!!!) + A_model = np.transpose(A_model, (2,0,1,3,4)) + + p_count = np.sum(A_model > 0, axis=-1) + n_count = np.sum(A_model < 0, axis=-1) + + if scheme == 1: + # scheme #1: decision on aggregated counts + for i in range(lag): + p_count_sum = np.sum(p_count[i], axis=tuple(range(p_count[i].ndim-1))) + n_count_sum = np.sum(n_count[i], axis=tuple(range(n_count[i].ndim-1))) + + node_type = 2 * (p_count_sum > n_count_sum) -1 + + # for draws, take random pick + # this accounts for all zeros for a node; p_count==n_count for a node; + # all zeros for a node for all bootstrap(no support anyways); + print("draw: ",np.where(p_count_sum == n_count_sum)[0]) + for idx in np.where(p_count_sum == n_count_sum)[0]: + node_type[idx] = np.random.choice([-1,1]) + print("Lag "+str(i+1)+":", node_type) + elif scheme == 2: + # scheme #2: aggregate of individual bootstrap decisions + for i in range(lag): + tie = p_count[i] == n_count[i] + + p_comparison = p_count[i] > n_count[i] + n_comparison = p_count[i] < n_count[i] + + p_comparison_aggregate = np.sum(p_comparison, axis=tuple(range(p_comparison.ndim-1))) + n_comparison_aggregate = np.sum(n_comparison, axis=tuple(range(n_comparison.ndim-1))) + + node_type = 2 * (p_comparison_aggregate > n_comparison_aggregate)-1 + + # for draws, take random pick + # this accounts for all zeros for a node; p_count==n_count for a node; + # all zeros for a node for all bootstrap(no support anyways); I_candidate count == E_candidate count + print("draw: ",np.where(p_comparison_aggregate == n_comparison_aggregate)[0]) + for idx in np.where(p_comparison_aggregate == n_comparison_aggregate)[0]: + node_type[idx] = np.random.choice([-1,1]) + print("Lag "+str(i+1)+":", node_type) + + return node_type + + +def _unpack_coef(coef, lag, n_features, has_bias=True): + + # coef shape: (n_boot, n_reg, n_features * (lag * n_features + has_bias)) + # so we need to flip the first two axis first to align with the FDR indexing scheme + coef = np.transpose(coef, (1, 0, 2)) + # coef shape: (n_reg, n_boot, n_features * (lag * n_features + has_bias)) + n_reg, n_boot = coef.shape[:2] + + # reshape back to matrix form while preserving n_reg and n_boot dimensions + # reshape to (n_reg, n_boot, n_features, lag * n_features + has_bias) + coef = coef.reshape(n_reg, n_boot, n_features, lag * n_features + has_bias) + + # transpose last two dims: (n_reg, n_boot, lag * n_features + has_bias, n_features) + coef = np.transpose(coef, (0, 1, 3, 2)) + + if has_bias: + # b shape: (n_reg, n_boot, n_features) + b = coef[:, :, 0, :] + # adjacency_flat shape: (n_reg, n_boot, lag*n_features, n_features) + adjacency_flat = coef[:, :, 1:, :] + else: + b = None + # adjacency_flat shape: (n_reg, n_boot, lag*n_features, n_features) + adjacency_flat = coef + + # reshape into adjacency matrices per lag + # shape: (n_reg, n_boot, lag, n_features, n_features) + adjacency_matrices = adjacency_flat.reshape(n_reg, n_boot, lag, n_features, n_features) + + # transpose the last two dimensions for each lag + # final shape: (n_reg, n_boot, lag, n_features, n_features) + A = np.transpose(adjacency_matrices, (0, 1, 2, 4, 3)) + + return A, b + +def _pack_coef(A, b=None): + """ + Convert VAR parameters (adjacency matrices and intercept) to vectorized form. + + Parameters: + ----------- + A : np.ndarray + Adjacency matrices of shape (lag, n_features, n_features) + b : np.ndarray or None + Intercept vector of shape (n_features,) if present + + Returns: + -------- + coef : np.ndarray + Vectorized coefficients + """ + lag, n_features, _ = A.shape + has_bias = b is not None + + # Transpose A back: (lag, n_features, n_features) -> (lag, n_features, n_features) + adjacency_matrices = np.transpose(A, (0, 2, 1)) + + # Flatten adjacency matrices: (lag, n_features, n_features) -> (lag*n_features, n_features) + adjacency_flat = adjacency_matrices.reshape(lag * n_features, n_features) + + if has_bias: + # Stack intercept with adjacency matrices + coef = np.vstack([b.reshape(1, -1), adjacency_flat]) # shape: (lag*n_features + 1, n_features) + else: + coef = adjacency_flat # shape: (lag*n_features, n_features) + + # Transpose and flatten to match original input format + coef = coef.T.flatten() + + return coef + +def edge_selector_fdr(A_edges_real_list, A_edges_shuf_list, alpha=0.05): + """ + Apply row-wise FDR to select significant edges. + + Parameters + ---------- + A_edges_real_list : list of np.ndarray + List of Kreal real-data edge matrices (off-diagonals only). + A_edges_shuf_list : list of np.ndarray + List of Kshuf shuffled-data edge matrices (off-diagonals only). + alpha : float + FDR control level. + + Returns + ------- + W_edges : np.ndarray + Binary mask of significant edges (same shape as A_edges_real). + W_pval : np.ndarray + P-value matrix for each edge. + summary : dict + Summary statistics. + """ + + Kreal = len(A_edges_real_list) + Kshuf = len(A_edges_shuf_list) + + # Stack into arrays: shape (K, N, N) + A_real = np.stack(A_edges_real_list, axis=0) + A_shuf = np.stack(A_edges_shuf_list, axis=0) + + + N = A_real.shape[1] + + # Statistic: median magnitude across bootstraps for each edge + stat_obs = np.median(np.abs(A_real), axis=0) # shape (N, N) + + # Null distribution: pool shuffled bootstraps per row + W_pval = np.ones((N, N)) + W_edge_mask = np.zeros((N, N), dtype=bool) + #W_edge_mask is a binary mask of significant edges after row‑wise FDR. + + # Statistics tracking + edges_per_row = [] + + for i in range(N): + # Pool null values for row i from shuffled data + null_vals_row = np.abs(A_shuf[:, i, :]).ravel() + # Compute p-values for all j in this row + for j in range(N): + if i == j: + continue + obs_val = stat_obs[i, j] + # Empirical p-value + pval = np.mean(null_vals_row >= obs_val) + W_pval[i, j] = pval + + # Apply FDR for this row + mask = np.ones(N, dtype=bool) + mask[i] = False + #...... per‑row FDR selection, not global FDR + reject, _ = fdrcorrection(W_pval[i, mask], alpha=alpha) + W_edge_mask[i, mask] = reject + + # Count selected edges for this row + edges_per_row.append(int(reject.sum())) + + # Additional statistics + edges_per_row = np.array(edges_per_row) + + # Count zeros in real data matrices + total_elements = A_real.size + zero_elements = np.sum(A_real == 0.0) + sparsity_A_real = zero_elements / total_elements + + # Compute sparsity of W_edge_mask (fraction of zeros in off-diagonal elements) + off_diag_elements = N * (N - 1) # Total off-diagonal elements + selected_edges = int(W_edge_mask.sum()) - N # Subtract diagonal (always True) + sparsity_W_mask = 1.0 - (selected_edges / off_diag_elements) + + # Compute pooled values from A_shuf dimensions: A_shuf shape is (Kshuf, N, N) + # For each row i, we pool A_shuf[:, i, :].ravel() which gives Kshuf*N values + pooled_values_per_row = Kshuf * N + pooled_values_total = N * pooled_values_per_row # N rows total + + summary = { + "alpha": alpha, + "num_bootstraps_real": Kreal, + "num_bootstraps_shuf": Kshuf, + "matrix_size": N, + "num_significant_edges": int(W_edge_mask.sum()) - N, # Exclude diagonal + "pooled_values_per_row": int(pooled_values_per_row), + "pooled_values_total": int(pooled_values_total), + "zero_elements_A_lasso": int(zero_elements), + "total_elements_A_lasso": int(total_elements), + "sparsity_A_lasso": float(sparsity_A_real), + "sparsity_W_edge_mask": float(sparsity_W_mask), + "edges_per_row_avg": float(edges_per_row.mean()), + "edges_per_row_std": float(edges_per_row.std()), + "edges_per_row_min": int(edges_per_row.min()), + "edges_per_row_max": int(edges_per_row.max()), + } + + return W_edge_mask, W_pval, summary + +def intersection_FDR(coefs, coefs_shuf, lag, n_features, fdr_rate = 0.05): + + # _unpack_coef makes n_reg X n_boot X n_feat X n_feat... + A_real, _ = _unpack_coef(coefs, lag, n_features, has_bias=True) + A_shuf, _ = _unpack_coef(coefs_shuf, lag, n_features, has_bias=True) + + # in case if lag = 1, get rid of the redundant dimension + A_real = np.squeeze(A_real) + A_shuf = np.squeeze(A_shuf) + + # set diagonal to zero, since we are do doing variable selection on diagonals + A_real[:, :, np.arange(n_features), np.arange(n_features)] = 0 + A_shuf[:, :, np.arange(n_features), np.arange(n_features)] = 0 + + # also need to clear the digonals of the A_* + supports = [] + #iterating through the L1-penalty list + count_list = [] + for i_reg in range(A_real.shape[0]): + W_edge_mask, W_pval, summary = edge_selector_fdr(A_real[i_reg], A_shuf[i_reg], alpha = fdr_rate) + + + # W_edge_mask is n_feat X n_feat + np.fill_diagonal(W_edge_mask, 1) # add digonal back in + W_edge_mask = W_edge_mask[np.newaxis] # making the dimension for LAG + + # Count all nonzero - diagonal nonzero for each array + total_nonzero = np.count_nonzero(W_edge_mask, axis=(1, 2)) + diag_nonzero = np.count_nonzero(W_edge_mask.diagonal(axis1=1, axis2=2), axis=1) + offdiag_counts = total_nonzero - diag_nonzero + + count_list.append(offdiag_counts) # Array of shape (lag,) with counts for each matrix + + # the biase terms are always assumed to be fully dense, so there’s no selection on it + support_i = _pack_coef(W_edge_mask, b=np.ones(A_real.shape[-1])) + supports.append(support_i) + + + return np.array(supports), count_list # supports SHOULD have shape n_reg_params X n_coef!!! + \ No newline at end of file diff --git a/src/pyuoi/mpi_utils.py b/src/pyuoi/mpi_utils.py index 4c3f3e95..4a989b77 100644 --- a/src/pyuoi/mpi_utils.py +++ b/src/pyuoi/mpi_utils.py @@ -160,12 +160,15 @@ def Gatherv_rows(send, comm=None, root=0): if rank == root: rec_shape = (tot[0],) + shape[1:] rec = np.empty(rec_shape, dtype=dtype) - sizes = [size * np.prod(rec_shape[1:]) for size in rank_sizes] - disps = np.insert(np.cumsum(sizes), 0, 0)[:-1] + sizes = [(size * np.prod(rec_shape[1:])).astype(np.int64) for size in rank_sizes] + disps = (np.insert(np.cumsum(sizes), 0, 0)[:-1]).astype(np.int64) + else: rec = None sizes = None disps = None + comm.Gatherv(send, [rec, sizes, disps, _np2mpi[dtype]], root=0) return rec + diff --git a/src/pyuoi/utils.py b/src/pyuoi/utils.py index 1b3709d0..da063db9 100755 --- a/src/pyuoi/utils.py +++ b/src/pyuoi/utils.py @@ -61,7 +61,8 @@ def log_likelihood_glm(model, y_true, y_pred): else: ll = 0. else: - ll = np.mean(y_true * np.log(y_pred) - y_pred) + #ll = np.mean(y_true * np.log(y_pred) - y_pred) #old: using mean error, which was inconsistent with the poisson loss function for VAR case(that used 1/n_samples factor) + ll = -np.sum(y_true * np.log(y_pred) - y_pred) else: raise ValueError('Model is not available.') return ll diff --git a/tests/test_mpi/__init__.py b/tests/test_mpi/__init__.py deleted file mode 100644 index e69de29b..00000000