mirror of
https://github.com/rndlabsoy/fly-brain-full.git
synced 2026-09-10 11:20:43 +02:00
erojasoficial-byte
c1d089dc3b
138,639 LIF neurons (FlyWire v783, GPU/PyTorch) connected to NeuroMechFly v2 (MuJoCo) with compound-eye vision, olfaction, gustation, somatosensation, courtship song, virtual flight, and proboscis extension. Includes benchmark suite (Brian2, Brian2CUDA, PyTorch, NEST GPU), project website, and technical papers (EN/ES).
336 lines
12 KiB
Python
336 lines
12 KiB
Python
#!/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 ""
|