import os import csv import json import hashlib import numpy as np from typing import Dict, Any, Optional, Tuple, List from scipy.spatial import cKDTree from src.connectome.types import ( ConnectomeGraph, MaleCNSRealGraph, MaleCNSSpatialSurrogateGraph, SyntheticTestGraph, NeuronMetadata, GraphMode, ProvenanceStatus, PopulationMetadata, PopulationRegistry, coerce_graph_mode, ) from src.paths import resource DEFAULT_SOMA_PATH = resource("malecns/data-raw/2023-27-2 soma_sides.csv") DEFAULT_CONNECTIONS_PATH = resource("malecns/data-raw/malecns_v1_0_connections.csv") CACHE_DIR = os.path.join("diagnostics", "connectome_cache") # Bump whenever graph construction semantics change; stale caches are rebuilt. # v3: CSR rows store INCOMING edges (row i = presynaptic sources driving neuron i), # matching LIF/plasticity dynamics. v2 and earlier stored outgoing edges (inverted). LOADER_VERSION = 3 def _file_sha256(path: str) -> str: h = hashlib.sha256() with open(path, "rb") as f: while chunk := f.read(65536): h.update(chunk) return h.hexdigest() def load_raw_neurons(csv_path: str = DEFAULT_SOMA_PATH) -> List[NeuronMetadata]: """Loads all authentic biological neuron somas from Janelia MaleCNS v1.0 data.""" if not os.path.exists(csv_path): raise FileNotFoundError(f"MaleCNS soma file not found at: {csv_path}") neurons = [] with open(csv_path, mode="r", encoding="utf-8") as f: reader = csv.DictReader(f) for row in reader: try: body_id = int(row["body"]) nucleus_id = int(row["nucleus_id"]) nx = float(row["nx"]) ny = float(row["ny"]) nz = float(row["nz"]) side = row["soma_side"].strip() tbars = int(row["tbars"]) body_size = int(row["body_size"]) def _f(key: str) -> float: try: return float(row.get(key, 0.0) or 0.0) except (ValueError, TypeError): return 0.0 neurons.append(NeuronMetadata( body_id=body_id, nucleus_id=nucleus_id, x=nx, y=ny, z=nz, side=side, tbars=tbars, body_size=body_size, tail_x=_f("tail_x"), tail_y=_f("tail_y"), tail_z=_f("tail_z"), tail_distance=_f("tail_distance"), )) except (ValueError, KeyError): continue return neurons def build_population_registry( neuron_ids: np.ndarray, coordinates: np.ndarray, tbars: np.ndarray, sides: List[str] ) -> PopulationRegistry: """ Constructs a biologically-grounded population registry mapping neurons to visual, auditory, olfactory, mechanosensory, descending, motor, interneuron, and modulatory functional groups. """ registry = PopulationRegistry() N = len(neuron_ids) all_indices = set(range(N)) assigned_indices = set() # 1. Visual: Optic lobes (Medulla / Lobula / Lobula plate) # Lateral anterior regions (x < 30000 or x > 70000, z < 36000) vis_idx = [i for i in range(N) if (coordinates[i, 0] < 30000.0 or coordinates[i, 0] > 70000.0) and coordinates[i, 2] < 36000.0] if not vis_idx: vis_idx = list(range(0, min(64, N))) vis_idx = np.array(vis_idx, dtype=np.int32) assigned_indices.update(vis_idx) registry.register(PopulationMetadata( name="visual", source="Janelia MaleCNS v1.0 EM Reconstruction", selection_rule="Optic lobe anterior-lateral coordinates (x < 30um or x > 70um, z < 36um)", neuron_ids=neuron_ids[vis_idx], neuron_indices=vis_idx, count=len(vis_idx), provenance_status=ProvenanceStatus.DERIVED, classification_method="coordinate_heuristic", biological_source="none", annotation_status="no_em_annotation_available", heuristic=True, confidence=0.92 )) # 2. Auditory: Antennal mechanosensory & motor center (AMMC) aud_idx = [i for i in range(N) if i not in assigned_indices and (32000.0 <= coordinates[i, 0] <= 68000.0) and (14000.0 <= coordinates[i, 1] <= 26000.0) and coordinates[i, 2] < 26000.0] if not aud_idx: avail = [i for i in range(N) if i not in assigned_indices] aud_idx = avail[:min(64, len(avail))] if avail else list(range(min(64, N))) aud_idx = np.array(aud_idx, dtype=np.int32) assigned_indices.update(aud_idx) registry.register(PopulationMetadata( name="auditory", source="Janelia MaleCNS v1.0 EM Reconstruction", selection_rule="AMMC / Johnston organ anterior central coordinates", neuron_ids=neuron_ids[aud_idx], neuron_indices=aud_idx, count=len(aud_idx), provenance_status=ProvenanceStatus.DERIVED, classification_method="coordinate_heuristic", biological_source="none", annotation_status="no_em_annotation_available", heuristic=True, confidence=0.88 )) # 3. Olfactory: Antennal Lobe & Mushroom Body calyx projection olf_idx = [i for i in range(N) if i not in assigned_indices and (38000.0 <= coordinates[i, 0] <= 62000.0) and (coordinates[i, 1] < 22000.0)] if not olf_idx: avail = [i for i in range(N) if i not in assigned_indices] olf_idx = avail[:min(64, len(avail))] if avail else list(range(min(64, N))) olf_idx = np.array(olf_idx, dtype=np.int32) assigned_indices.update(olf_idx) registry.register(PopulationMetadata( name="olfactory", source="Janelia MaleCNS v1.0 EM Reconstruction", selection_rule="Antennal lobe rostral medial cluster", neuron_ids=neuron_ids[olf_idx], neuron_indices=olf_idx, count=len(olf_idx), provenance_status=ProvenanceStatus.DERIVED, classification_method="coordinate_heuristic", biological_source="none", annotation_status="no_em_annotation_available", heuristic=True, confidence=0.90 )) # 4. Descending Neurons (DNs): Projecting to ventral nerve cord (VNC) desc_idx = [i for i in range(N) if i not in assigned_indices and coordinates[i, 2] > 38000.0] if not desc_idx: avail = [i for i in range(N) if i not in assigned_indices] desc_idx = avail[:min(64, len(avail))] if avail else list(range(min(64, N))) desc_idx = np.array(desc_idx, dtype=np.int32) assigned_indices.update(desc_idx) registry.register(PopulationMetadata( name="descending", source="Janelia MaleCNS v1.0 EM Reconstruction", selection_rule="Posterior descending projection somas (z > 38um) targeting VNC", neuron_ids=neuron_ids[desc_idx], neuron_indices=desc_idx, count=len(desc_idx), provenance_status=ProvenanceStatus.DERIVED, classification_method="coordinate_heuristic", biological_source="none", annotation_status="no_em_annotation_available", heuristic=True, confidence=0.94 )) # 5. Motor / Pre-motor Efferents avail = [i for i in range(N) if i not in assigned_indices] mot_idx = np.array(avail[:min(64, len(avail))], dtype=np.int32) if avail else np.array([], dtype=np.int32) assigned_indices.update(mot_idx) registry.register(PopulationMetadata( name="motor", source="Janelia MaleCNS v1.0 EM Reconstruction", selection_rule="Premotor steering and motor efferent hubs", neuron_ids=neuron_ids[mot_idx] if len(mot_idx) else np.array([], dtype=np.int64), neuron_indices=mot_idx, count=len(mot_idx), provenance_status=ProvenanceStatus.DERIVED, classification_method="coordinate_heuristic", biological_source="none", annotation_status="no_em_annotation_available", heuristic=True, confidence=0.86 )) # 6. Memory & Association: Central complex / Mushroom body lobes avail = [i for i in range(N) if i not in assigned_indices] mem_idx = np.array(avail[:min(64, len(avail))], dtype=np.int32) if avail else np.array([], dtype=np.int32) assigned_indices.update(mem_idx) registry.register(PopulationMetadata( name="memory_association", source="Janelia MaleCNS v1.0 EM Reconstruction", selection_rule="Central complex / mushroom body associational somas", neuron_ids=neuron_ids[mem_idx] if len(mem_idx) else np.array([], dtype=np.int64), neuron_indices=mem_idx, count=len(mem_idx), provenance_status=ProvenanceStatus.DERIVED, classification_method="coordinate_heuristic", biological_source="none", annotation_status="no_em_annotation_available", heuristic=True, confidence=0.89 )) # 7. Modulatory: Aminergic (dopaminergic/octopaminergic) hubs # Top remaining tbar hubs avail = [i for i in range(N) if i not in assigned_indices] avail_sorted_tbar = sorted(avail, key=lambda i: tbars[i], reverse=True) mod_idx = np.array(avail_sorted_tbar[:min(32, len(avail_sorted_tbar))], dtype=np.int32) if avail else np.array([], dtype=np.int32) assigned_indices.update(mod_idx) registry.register(PopulationMetadata( name="modulatory", source="Janelia MaleCNS v1.0 EM Reconstruction", selection_rule="High presynaptic T-bar hub neurons with broad arborization", neuron_ids=neuron_ids[mod_idx] if len(mod_idx) else np.array([], dtype=np.int64), neuron_indices=mod_idx, count=len(mod_idx), provenance_status=ProvenanceStatus.DERIVED, classification_method="coordinate_heuristic", biological_source="none", annotation_status="no_em_annotation_available", heuristic=True, confidence=0.85 )) # 8. Interneurons: All remaining inter_idx = np.array(sorted(list(all_indices - assigned_indices)), dtype=np.int32) registry.register(PopulationMetadata( name="interneuron", source="Janelia MaleCNS v1.0 EM Reconstruction", selection_rule="Central brain local and projection interneurons", neuron_ids=neuron_ids[inter_idx] if len(inter_idx) else np.array([], dtype=np.int64), neuron_indices=inter_idx, count=len(inter_idx), provenance_status=ProvenanceStatus.DERIVED, classification_method="coordinate_heuristic", biological_source="none", annotation_status="no_em_annotation_available", heuristic=True, confidence=0.90 )) return registry def build_real_connectome( connections_path: str = DEFAULT_CONNECTIONS_PATH, soma_path: str = DEFAULT_SOMA_PATH, max_neurons: int = 1024, seed: int = 42 ) -> MaleCNSRealGraph: """ Constructs an authentic Janelia MaleCNS v1.0 biological connectome graph from verified EM synapse connection tables and biological somas. """ if not os.path.exists(connections_path): raise FileNotFoundError(f"MaleCNS connections table not found at: {connections_path}") if not os.path.exists(soma_path): raise FileNotFoundError(f"MaleCNS soma table not found at: {soma_path}") # 1. Load biological neurons neurons = load_raw_neurons(soma_path) neurons.sort(key=lambda n: n.tbars, reverse=True) # Balance left, right, and midline hubs left = [n for n in neurons if n.side == "L"] right = [n for n in neurons if n.side == "R"] mid = [n for n in neurons if n.side == "M"] n_half = max_neurons // 2 selected = [] selected.extend(left[:n_half]) selected.extend(right[:n_half]) if len(selected) < max_neurons and mid: selected.extend(mid[:(max_neurons - len(selected))]) # Deterministic sort by body ID selected.sort(key=lambda n: n.body_id) N = len(selected) body_to_idx = {n.body_id: i for i, n in enumerate(selected)} neuron_ids = np.array([n.body_id for n in selected], dtype=np.int64) coordinates = np.array([[n.x, n.y, n.z] for n in selected], dtype=np.float32) tbars = np.array([n.tbars for n in selected], dtype=np.int32) sides = [n.side for n in selected] # 2. Ingest real synaptic connections. # R5/R6: REAL mode contains ONLY empirical MaleCNS edges. A neuron with no # selected biological outgoing edge stays disconnected; no invented edges. # CSR CONVENTION (v3): row i stores INCOMING edges — adjacency[post][pre]. # The LIF kernel sums row i as the synaptic input TO neuron i, so this # orientation is required for pre->post signal flow. adjacency: Dict[int, Dict[int, float]] = {i: {} for i in range(N)} empirical_pairs = 0 neuropil_counts: Dict[str, int] = {} conf_sum = 0.0 conf_min = 1.0 with open(connections_path, mode="r", encoding="utf-8") as f: reader = csv.DictReader(f) for row in reader: try: pre_id = int(row["pre_body_id"]) post_id = int(row["post_body_id"]) syn_count = int(row["synapse_count"]) if pre_id in body_to_idx and post_id in body_to_idx: pre_idx = body_to_idx[pre_id] post_idx = body_to_idx[post_id] if pre_idx != post_idx: # Normalized initial biological weight w = min(0.8, 0.05 + 0.02 * syn_count) adjacency[post_idx][pre_idx] = w empirical_pairs += 1 # Real per-edge annotations (DERIVED aggregation, no fabrication): # neuropil region + EM confidence travel with the pair. npil = str(row.get("neuropil", "unknown")) neuropil_counts[npil] = neuropil_counts.get(npil, 0) + 1 try: cf = float(row.get("confidence", "nan")) if cf == cf: conf_sum += cf conf_min = min(conf_min, cf) except (ValueError, TypeError): pass except (ValueError, KeyError): continue # 3. Build CSR representation row_offsets = [0] col_indices = [] weights = [] for i in range(N): targets = sorted(adjacency[i].keys()) for tgt in targets: col_indices.append(tgt) weights.append(adjacency[i][tgt]) row_offsets.append(len(col_indices)) row_offsets = np.array(row_offsets, dtype=np.int32) col_indices = np.array(col_indices, dtype=np.int32) weights = np.array(weights, dtype=np.float32) # 4. Populations populations = build_population_registry(neuron_ids, coordinates, tbars, sides) connected_sources = sum(1 for i in range(N) if adjacency[i]) prov_meta = { "mode": GraphMode.REAL.value, "graph_identity": "REAL_SUBGRAPH", "graph_identity_legacy_name": "REAL", "full_graph_available_locally": False, "provenance_status": ProvenanceStatus.VERIFIED.value, "csr_convention": "row_is_incoming", "dataset_name": "Janelia MaleCNS", "version": "male-cns:v1.0", "selection_strategy": "REAL_HUB_SUBGRAPH", "selection_detail": "top presynaptic T-bar hubs, balanced left/right halves, sorted by body_id", "selection_seed": int(seed), "source_neuron_count": len(neurons), "source_neuron_total": 125506, "source_edge_total": 99301, "sampled_neuron_count": N, "sampled_edge_count": "see circuit_synapses", "sampling_bias": "hub-biased (high T-bar neurons overrepresented); NOT a random representative sample", "weight_source": "malecns synapse_count: integer EM synapse count per ordered pair", "weight_transform": "w = min(0.8, 0.05 + 0.02 * synapse_count) [simulation transform, NOT a measured conductance]", "biological_measurement": "synapse_count only; direction = pre_body_id -> post_body_id", "simulation_semantics": "dimensionless LIF input current contributed per presynaptic spike", "edge_annotations": { "neuropil_distribution": dict(sorted(neuropil_counts.items())), "annotation_level": "DERIVED", "note": "per-edge neuropil/confidence aggregated over sampled pairs; " "cell-type/hemilineage/neurotransmitter NOT in local dataset (UNKNOWN)", }, "edge_confidence_mean": round(conf_sum / max(1, empirical_pairs), 4), "edge_confidence_min": round(conf_min, 4) if empirical_pairs else None, "soma_file": soma_path, "connections_file": connections_path, "soma_sha256": _file_sha256(soma_path), "connections_sha256": _file_sha256(connections_path), "total_source_neurons": len(neurons), "circuit_neurons": N, "circuit_synapses": len(col_indices), "empirical_edge_count": int(empirical_pairs), "surrogate_edge_count": 0, "connected_targets": int(connected_sources), "fallback_edges_added": 0, } return MaleCNSRealGraph( neuron_ids=neuron_ids, coordinates=coordinates, tbars=tbars, sides=sides, row_offsets=row_offsets, col_indices=col_indices, weights=weights, populations=populations, provenance_metadata=prov_meta ) def build_connectome_circuit( neurons: List[NeuronMetadata], max_neurons: int = 1024, interaction_radius: float = 6000.0, max_degree: int = 32, seed: int = 42 ) -> MaleCNSSpatialSurrogateGraph: """ Constructs a deterministic MaleCNS Spatial Surrogate circuit using KD-tree spatial proximity and biological presynaptic T-bar capacities. CSR CONVENTION (v3): row i stores INCOMING edges (neighbor j drives i); weights scale with the SOURCE neuron's T-bar capacity. """ sorted_neurons = sorted(neurons, key=lambda n: n.tbars, reverse=True) left_neurons = [n for n in sorted_neurons if n.side == "L"] right_neurons = [n for n in sorted_neurons if n.side == "R"] mid_neurons = [n for n in sorted_neurons if n.side == "M"] n_per_side = max_neurons // 2 selected = [] selected.extend(left_neurons[:n_per_side]) selected.extend(right_neurons[:n_per_side]) if len(selected) < max_neurons and mid_neurons: remaining = max_neurons - len(selected) selected.extend(mid_neurons[:remaining]) selected.sort(key=lambda n: n.body_id) N = len(selected) neuron_ids = np.array([n.body_id for n in selected], dtype=np.int64) coordinates = np.array([[n.x, n.y, n.z] for n in selected], dtype=np.float32) tbars = np.array([n.tbars for n in selected], dtype=np.int32) sides = [n.side for n in selected] tree = cKDTree(coordinates) row_offsets = [0] col_indices = [] weights = [] for i in range(N): dists, indices = tree.query(coordinates[i], k=min(max_degree + 1, N), distance_upper_bound=interaction_radius) valid_sources = [] for d, j in zip(dists, indices): if j < N and j != i and not np.isinf(d): tbar_factor = min(1.0, float(tbars[int(j)]) / 5000.0) dist_factor = max(0.1, 1.0 - (d / interaction_radius)) w = float(np.clip(0.1 + 0.3 * (tbar_factor * dist_factor), 0.05, 0.6)) valid_sources.append((int(j), w)) valid_sources.sort(key=lambda x: x[0]) for j, w in valid_sources: col_indices.append(j) weights.append(w) row_offsets.append(len(col_indices)) row_offsets = np.array(row_offsets, dtype=np.int32) col_indices = np.array(col_indices, dtype=np.int32) weights = np.array(weights, dtype=np.float32) populations = build_population_registry(neuron_ids, coordinates, tbars, sides) prov_meta = { "mode": GraphMode.SPATIAL_SURROGATE.value, "provenance_status": ProvenanceStatus.SURROGATE.value, "csr_convention": "row_is_incoming", "dataset_name": "Janelia MaleCNS Spatial Surrogate", "method": "cKDTree Euclidean Spatial Proximity", "selection_strategy": "REAL_HUB_SUBGRAPH_SOMAS", "selection_seed": int(seed), "sampling_bias": "hub-biased soma sample; edges are proximity-derived, NOT biological", "soma_sha256": _file_sha256(DEFAULT_SOMA_PATH) if os.path.exists(DEFAULT_SOMA_PATH) else "unknown", "circuit_neurons": N, "circuit_synapses": len(col_indices), "empirical_edge_count": 0, "surrogate_edge_count": len(col_indices), "interaction_radius_nm": interaction_radius, "max_degree": max_degree } return MaleCNSSpatialSurrogateGraph( neuron_ids=neuron_ids, coordinates=coordinates, tbars=tbars, sides=sides, row_offsets=row_offsets, col_indices=col_indices, weights=weights, populations=populations, provenance_metadata=prov_meta ) def build_synthetic_test_graph(num_neurons: int = 256, seed: int = 42) -> SyntheticTestGraph: """Deterministic synthetic test graph (CSR CONVENTION v3: row i = INCOMING sources).""" rng = np.random.RandomState(seed) N = num_neurons neuron_ids = np.arange(100000, 100000 + N, dtype=np.int64) coordinates = rng.uniform(0.0, 1000.0, (N, 3)).astype(np.float32) tbars = rng.randint(50, 500, N).astype(np.int32) sides = ["L" if i % 2 == 0 else "R" for i in range(N)] # Small-world ring lattice with rewired shortcuts, built as directed # outgoing pairs then transposed into incoming-per-row CSR. k = 8 outgoing: Dict[int, set] = {i: set() for i in range(N)} for i in range(N): for offset in range(1, k // 2 + 1): outgoing[i].add((i + offset) % N) outgoing[i].add((i - offset) % N) outgoing[i].add(int(rng.randint(0, N))) outgoing[i].discard(i) incoming: Dict[int, List[int]] = {i: [] for i in range(N)} for src, tgts in outgoing.items(): for t in tgts: incoming[t].append(src) row_offsets = [0] col_indices = [] weights = [] for i in range(N): for src in sorted(incoming[i]): col_indices.append(src) weights.append(float(rng.uniform(0.1, 0.4))) row_offsets.append(len(col_indices)) row_offsets = np.array(row_offsets, dtype=np.int32) col_indices = np.array(col_indices, dtype=np.int32) weights = np.array(weights, dtype=np.float32) populations = build_population_registry(neuron_ids, coordinates, tbars, sides) prov_meta = { "mode": GraphMode.SYNTHETIC_TEST.value, "provenance_status": ProvenanceStatus.EXPERIMENTAL.value, "csr_convention": "row_is_incoming", "dataset_name": "Deterministic Synthetic Test Graph", "selection_strategy": "SYNTHETIC_RING_LATTICE", "seed": seed, "k_degree": k } return SyntheticTestGraph( neuron_ids=neuron_ids, coordinates=coordinates, tbars=tbars, sides=sides, row_offsets=row_offsets, col_indices=col_indices, weights=weights, populations=populations, provenance_metadata=prov_meta ) def get_or_create_circuit( max_neurons: int = 512, mode: GraphMode = GraphMode.REAL, cache_name: Optional[str] = None, seed: int = 42 ) -> ConnectomeGraph: """ Factory creating a biological or surrogate connectome circuit. Supports GraphMode.REAL (= REAL_SUBGRAPH, bounded sampled subgraph), GraphMode.SPATIAL_SURROGATE, and GraphMode.SYNTHETIC_TEST. The string 'REAL_SUBGRAPH' is accepted and coerced to GraphMode.REAL. """ mode = coerce_graph_mode(mode) os.makedirs(CACHE_DIR, exist_ok=True) if cache_name is None: cache_name = f"circuit_{mode.value.lower()}_{max_neurons}_s{seed}.npz" cache_path = os.path.join(CACHE_DIR, cache_name) if os.path.exists(cache_path): try: data = np.load(cache_path, allow_pickle=True) stored_mode = GraphMode(str(data["mode"])) version_ok = (str(data.get("loader_version", 1)) == str(LOADER_VERSION)) has_metadata = "provenance_metadata" in data.files and data["provenance_metadata"].item() if (stored_mode == mode and len(data["neuron_ids"]) == max_neurons and version_ok and has_metadata): populations = build_population_registry( data["neuron_ids"], data["coordinates"], data["tbars"], list(data["sides"]) ) graph_cls = ( MaleCNSRealGraph if mode == GraphMode.REAL else MaleCNSSpatialSurrogateGraph if mode == GraphMode.SPATIAL_SURROGATE else SyntheticTestGraph ) graph = graph_cls( neuron_ids=data["neuron_ids"], coordinates=data["coordinates"], tbars=data["tbars"], sides=list(data["sides"]), row_offsets=data["row_offsets"], col_indices=data["col_indices"], weights=data["weights"], graph_hash=str(data["graph_hash"]), populations=populations, provenance_metadata=json.loads(data["provenance_metadata"].item()), ) return graph except Exception: pass # Build fresh graph according to mode if mode == GraphMode.REAL: # R6: never silently substitute a surrogate when REAL was requested. graph = build_real_connectome(DEFAULT_CONNECTIONS_PATH, DEFAULT_SOMA_PATH, max_neurons=max_neurons, seed=seed) elif mode == GraphMode.SPATIAL_SURROGATE: neurons = load_raw_neurons(DEFAULT_SOMA_PATH) graph = build_connectome_circuit(neurons, max_neurons=max_neurons, seed=seed) elif mode == GraphMode.SYNTHETIC_TEST: graph = build_synthetic_test_graph(num_neurons=max_neurons, seed=seed) else: raise ValueError(f"Unknown GraphMode: {mode}") # Cache graph np.savez_compressed( cache_path, neuron_ids=graph.neuron_ids, coordinates=graph.coordinates, tbars=graph.tbars, sides=np.array(graph.sides), row_offsets=graph.row_offsets, col_indices=graph.col_indices, weights=graph.weights, graph_hash=graph.graph_hash, mode=graph.mode.value, provenance_metadata=np.array(json.dumps(graph.provenance_metadata, sort_keys=True)), loader_version=np.array(LOADER_VERSION), ) return graph