#!/usr/bin/env python3 """ Olfactory System — Bilateral odor detection via antennal ORNs. Drosophila detects odors through Olfactory Receptor Neurons (ORNs) on the antennae. Different receptor types respond to different chemicals: - ORN_DM1 (~68 neurons, Or42b equivalent): food odors -> attraction - ORN_DA2 (~39 neurons, Or56a equivalent): geosmin -> aversion/escape Architecture: Virtual odor sources emit concentration gradients (inverse-square falloff). Left and right antenna positions are computed from fly position + heading. Bilateral concentration asymmetry enables chemotaxis (gradient navigation). """ import csv import numpy as np from pathlib import Path # ============================================================================ # Odor Source # ============================================================================ class OdorSource: """A point source emitting an odor with distance-dependent concentration. Parameters ---------- position : array-like [x, y, z] position in mm. odor_type : str 'attractive' or 'repulsive'. amplitude : float Peak concentration (0-1) at the source. spread : float Characteristic radius in mm (concentration halves at this distance). label : str Human-readable name. """ def __init__(self, position, odor_type, amplitude=1.0, spread=25.0, label=''): self.position = np.array(position[:3], dtype=np.float64) self.odor_type = odor_type # 'attractive' or 'repulsive' self.amplitude = float(amplitude) self.spread = float(spread) self.label = label or odor_type # ============================================================================ # Olfactory System # ============================================================================ class OlfactorySystem: """Bilateral olfactory detection using ORN populations from FlyWire. Parameters ---------- flyid2i : dict FlyWire neuron ID -> tensor index mapping. annotations_path : str or Path, optional Path to FlyWire annotations TSV. """ # Antenna offset from fly midline (mm). Exaggerated slightly # vs real anatomy (~0.15mm) to produce functional chemotaxis # gradients at simulation scale. ANTENNA_SPREAD = 2.0 # Firing rates ATTRACTIVE_MAX_RATE = 180.0 # Hz (Or42b food) REPULSIVE_MAX_RATE = 250.0 # Hz (Or56a danger) # Floor: below this concentration, no activation CONC_FLOOR = 0.02 # Repulsive escape threshold (normalized 0-1) REPULSION_ESCAPE_THRESH = 0.3 def __init__(self, flyid2i, annotations_path=None): self.flyid2i = flyid2i if annotations_path is None: annotations_path = Path(__file__).parent / 'data' / 'flywire_annotations.tsv' self._load_orn_populations(str(annotations_path)) # Runtime state self.conc_left_att = 0.0 self.conc_right_att = 0.0 self.conc_left_rep = 0.0 self.conc_right_rep = 0.0 self.attractive_rate_left = 0.0 self.attractive_rate_right = 0.0 self.repulsive_rate_left = 0.0 self.repulsive_rate_right = 0.0 self.active_source_label = '' # ── Population Loading ───────────────────────────────────────────────── def _load_orn_populations(self, annotations_path): """Load ORN_DM1 (attractive) and ORN_DA2 (repulsive) from annotations.""" att_left, att_right = [], [] rep_left, rep_right = [], [] with open(annotations_path, 'r') as f: reader = csv.DictReader(f, delimiter='\t') for row in reader: ct = row.get('cell_type', '') if ct not in ('ORN_DM1', 'ORN_DA2'): continue rid = int(row['root_id']) if rid not in self.flyid2i: continue idx = self.flyid2i[rid] side = row.get('side', '') if ct == 'ORN_DM1': if side == 'left': att_left.append(idx) elif side == 'right': att_right.append(idx) elif ct == 'ORN_DA2': if side == 'left': rep_left.append(idx) elif side == 'right': rep_right.append(idx) self.att_idx_left = np.array(att_left, dtype=np.int64) self.att_idx_right = np.array(att_right, dtype=np.int64) self.rep_idx_left = np.array(rep_left, dtype=np.int64) self.rep_idx_right = np.array(rep_right, dtype=np.int64) n_att = len(att_left) + len(att_right) n_rep = len(rep_left) + len(rep_right) print(f"[Olfactory] Attractive ORN (DM1/Or42b): {n_att} neurons " f"(L={len(att_left)}, R={len(att_right)})") print(f"[Olfactory] Repulsive ORN (DA2/Or56a): {n_rep} neurons " f"(L={len(rep_left)}, R={len(rep_right)})") # ── Concentration Computation ────────────────────────────────────────── @staticmethod def _concentration_at(pos, sources, odor_type): """Compute total concentration at a position. Uses inverse-square falloff: c = amplitude / (1 + (d/spread)^2) """ total = 0.0 for src in sources: if src.odor_type != odor_type: continue dist = np.linalg.norm(pos[:2] - src.position[:2]) conc = src.amplitude / (1.0 + (dist / src.spread) ** 2) total += conc return total # ── Main Processing ──────────────────────────────────────────────────── def process(self, fly_pos, fly_heading, odor_sources): """Compute bilateral ORN activation from odor sources. Parameters ---------- fly_pos : array-like, shape (3,) Fly position [x, y, z] in mm. fly_heading : float Fly yaw angle in radians. odor_sources : list of OdorSource """ if not odor_sources: self.conc_left_att = 0.0 self.conc_right_att = 0.0 self.conc_left_rep = 0.0 self.conc_right_rep = 0.0 self._compute_rates() return # Antenna positions: perpendicular to heading # Left antenna: +90 deg from forward perp_x = -np.sin(fly_heading) * self.ANTENNA_SPREAD perp_y = np.cos(fly_heading) * self.ANTENNA_SPREAD pos_left = np.array([ fly_pos[0] + perp_x, fly_pos[1] + perp_y, fly_pos[2]]) pos_right = np.array([ fly_pos[0] - perp_x, fly_pos[1] - perp_y, fly_pos[2]]) # Compute concentrations at each antenna self.conc_left_att = self._concentration_at( pos_left, odor_sources, 'attractive') self.conc_right_att = self._concentration_at( pos_right, odor_sources, 'attractive') self.conc_left_rep = self._concentration_at( pos_left, odor_sources, 'repulsive') self.conc_right_rep = self._concentration_at( pos_right, odor_sources, 'repulsive') self._compute_rates() # Find most active source for labeling self.active_source_label = '' max_conc = 0.0 for src in odor_sources: if src.odor_type == 'attractive': c = self.conc_left_att + self.conc_right_att else: c = self.conc_left_rep + self.conc_right_rep if c > max_conc: max_conc = c self.active_source_label = src.label def _compute_rates(self): """Map concentrations to ORN firing rates.""" # Attractive ORN_DM1 for attr_name, conc_attr in [ ('attractive_rate_left', 'conc_left_att'), ('attractive_rate_right', 'conc_right_att'), ]: c = getattr(self, conc_attr) if c > self.CONC_FLOOR: setattr(self, attr_name, min(c, 1.0) * self.ATTRACTIVE_MAX_RATE) else: setattr(self, attr_name, 0.0) # Repulsive ORN_DA2 for attr_name, conc_attr in [ ('repulsive_rate_left', 'conc_left_rep'), ('repulsive_rate_right', 'conc_right_rep'), ]: c = getattr(self, conc_attr) if c > self.CONC_FLOOR: setattr(self, attr_name, min(c, 1.0) * self.REPULSIVE_MAX_RATE) else: setattr(self, attr_name, 0.0) # ── Brain Injection ──────────────────────────────────────────────────── def get_rates(self): """Get combined (indices, rates) arrays for brain injection. Returns ------- indices : np.ndarray (int64) rates : np.ndarray (float64) """ all_idx = [] all_rates = [] if self.attractive_rate_left > 0.1 and len(self.att_idx_left) > 0: all_idx.append(self.att_idx_left) all_rates.append( np.full(len(self.att_idx_left), self.attractive_rate_left)) if self.attractive_rate_right > 0.1 and len(self.att_idx_right) > 0: all_idx.append(self.att_idx_right) all_rates.append( np.full(len(self.att_idx_right), self.attractive_rate_right)) if self.repulsive_rate_left > 0.1 and len(self.rep_idx_left) > 0: all_idx.append(self.rep_idx_left) all_rates.append( np.full(len(self.rep_idx_left), self.repulsive_rate_left)) if self.repulsive_rate_right > 0.1 and len(self.rep_idx_right) > 0: all_idx.append(self.rep_idx_right) all_rates.append( np.full(len(self.rep_idx_right), self.repulsive_rate_right)) if not all_idx: return np.array([], dtype=np.int64), np.array([], dtype=np.float64) return np.concatenate(all_idx), np.concatenate(all_rates) # ── Diagnostics ──────────────────────────────────────────────────────── @property def attractive_level(self): """Normalized attractive intensity [0-1].""" return (max(self.attractive_rate_left, self.attractive_rate_right) / self.ATTRACTIVE_MAX_RATE) @property def repulsive_level(self): """Normalized repulsive intensity [0-1].""" return (max(self.repulsive_rate_left, self.repulsive_rate_right) / self.REPULSIVE_MAX_RATE) @property def is_repulsive_escape(self): """True when repulsive concentration triggers escape.""" return self.repulsive_level > self.REPULSION_ESCAPE_THRESH @property def attraction_bias(self): """Orientation bias for attractive chemotaxis. +1 = more on right antenna -> turn right toward source. -1 = more on left -> turn left toward source. """ total = self.conc_left_att + self.conc_right_att if total < self.CONC_FLOOR * 2: return 0.0 return (self.conc_right_att - self.conc_left_att) / total @property def repulsion_bias(self): """Orientation bias for repulsive escape. +1 = more threat on right antenna. -1 = more threat on left. Bridge should turn AWAY from this direction. """ total = self.conc_left_rep + self.conc_right_rep if total < self.CONC_FLOOR * 2: return 0.0 return (self.conc_right_rep - self.conc_left_rep) / total def get_status_str(self): """One-line diagnostic string.""" parts = [] if self.attractive_level > 0.01: parts.append( f"FOOD={self.attractive_level:.2f} " f"[L={self.conc_left_att:.2f} R={self.conc_right_att:.2f}]") if self.repulsive_level > 0.01: parts.append( f"DANGER={self.repulsive_level:.2f} " f"[L={self.conc_left_rep:.2f} R={self.conc_right_rep:.2f}]") return " | ".join(parts) if parts else ""