"""Official paired-synapse localization; numerical counts must reconcile to CSR.""" from pathlib import Path import hashlib import json import time import numpy as np import pandas as pd import pyarrow as pa import pyarrow.compute as pc from scipy.spatial import cKDTree from audit import save,digest ROOT=Path('/data/digital-drosophila');EXP=ROOT/'experiments/phase4aa-apl-pn-scope' def main(): pa.set_cpu_count(4);pa.set_io_thread_count(4);start=time.perf_counter() out=EXP/'localization';out.mkdir(exist_ok=False) official=EXP/'external-resources/synaptic-partners' verified=json.loads((official/'download-verification.json').read_text());assert verified['pass'] path=Path(verified['path']);assert path.stat().st_size==verified['bytes'] source_hash=digest(path);assert source_hash==verified['sha256'] save(out/'protocol.json',{'source_sha256':source_hash,'source_generation':verified['generation'], 'source_script_sha256':digest(Path(__file__)),'protocol_sha256':digest(Path(__file__).with_name('LOCALIZATION-PROTOCOL.md')), 'coordinate_unit_nm':8,'distance_to_microns_multiplier':0.008, 'nearest_neighbor_is_not_bouton_or_functional_identity':True,'before_analysis':True}) canonical=ROOT/'datasets/male-cns-v1.0/processed/neuron-only' m=pd.read_parquet(canonical/'neuron-index/unified-metadata.parquet').sort_values('neuron_idx').reset_index(drop=True) ptr=np.load(canonical/'gpu/csr-indptr.int64.npy',mmap_mode='r');dst=np.load(canonical/'gpu/csr-indices.int32.npy',mmap_mode='r');w=np.load(canonical/'gpu/csr-weights.int64.npy',mmap_mode='r') body=m.body_id.to_numpy();apl=body[m.annotation_type.eq('APL')];kc=body[m.annotation_class.eq('Kenyon_Cell')];alln=body[m.annotation_class.eq('ALLN')] pnpairs=pd.read_parquet(EXP/'connectivity/APL-PN-reciprocity.parquet');pn=pnpairs.target_body_id.unique() watchidx=json.loads((ROOT/'experiments/phase4z-pn-feedback-audit/protocol/watched.json').read_text())['all_incoming_ALPNs'] watch=body[watchidx] groups={k:[] for k in ('APL-incident','APL-target-PN-to-KC','watched-PN-to-ALLN','ALLN-to-watched-PN')} sets={name:pa.array(arr.astype(np.int64)) for name,arr in [('APL',apl),('KC',kc),('ALLN',alln),('PN',pn),('watched',watch)]} total=0 with pa.memory_map(str(path),'r') as f: reader=pa.ipc.open_file(f);names=reader.schema.names required=['x_pre','y_pre','z_pre','body_pre','conf_pre','x_post','y_post','z_post','body_post','conf_post','primary_post'] save(out/'actual-schema.json',{'fields':[{'name':x.name,'dtype':str(x.type)} for x in reader.schema],'record_batches':reader.num_record_batches}) assert set(required).issubset(names),names for i in range(reader.num_record_batches): batch=reader.get_batch(i);total+=len(batch) pre=batch.column(names.index('body_pre'));post=batch.column(names.index('body_post')) masks={ 'APL-incident':pc.or_(pc.is_in(pre,value_set=sets['APL']),pc.is_in(post,value_set=sets['APL'])), 'APL-target-PN-to-KC':pc.and_(pc.is_in(pre,value_set=sets['PN']),pc.is_in(post,value_set=sets['KC'])), 'watched-PN-to-ALLN':pc.and_(pc.is_in(pre,value_set=sets['watched']),pc.is_in(post,value_set=sets['ALLN'])), 'ALLN-to-watched-PN':pc.and_(pc.is_in(pre,value_set=sets['ALLN']),pc.is_in(post,value_set=sets['watched']))} for name,mask in masks.items(): sub=batch.filter(mask) if len(sub):groups[name].append(pa.Table.from_batches([sub])) if (i+1)%500==0:print('streamed_batches',i+1,'rows',total,flush=True) lookup=pd.Series(m.neuron_idx.to_numpy(),index=body) dfs={};checks=[] keys=['body_pre','x_pre','y_pre','z_pre','body_post','x_post','y_post','z_post'] for name,tables in groups.items(): a=pa.concat_tables(tables).to_pandas();assert not a.duplicated(keys).any(),name+' duplicate partner row' a['source_neuron_idx']=a.body_pre.map(lookup);a['target_neuron_idx']=a.body_post.map(lookup) canonical_rows=a.source_neuron_idx.notna()&a.target_neuron_idx.notna() pairs=a.loc[canonical_rows].groupby(['source_neuron_idx','target_neuron_idx']).size() problems=[] for (s,t),count in pairs.items(): s,t=int(s),int(t);eid=np.arange(ptr[s],ptr[s+1]);found=eid[dst[eid]==t] if len(found)!=1 or int(w[found[0]])!=count:problems.append({'source':s,'target':t,'actual_partner_count':int(count),'CSR_weight':int(w[found[0]]) if len(found)==1 else None}) result={'group':name,'rows':len(a),'canonical_rows':int(canonical_rows.sum()),'canonical_pairs':len(pairs),'duplicate_partner_rows':0,'mismatches':problems,'pass':not problems} checks.append(result);save(out/(name+'-count-check.json'),result);assert result['pass'],name+' count mismatch' a.to_parquet(out/(name+'.parquet'),index=False);dfs[name]=a print(json.dumps(result),flush=True) a=dfs['APL-incident'];forward=a[a.body_pre.isin(apl)&a.body_post.isin(pn)].copy();reverse=a[a.body_pre.isin(pn)&a.body_post.isin(apl)].copy() assert len(forward)==int(pnpairs.raw_synapse_count.sum())==4600 frames={'APL-to-PN':forward,'PN-to-APL':reverse,'PN-to-KC':dfs['APL-target-PN-to-KC'], 'watched-PN-to-ALLN':dfs['watched-PN-to-ALLN'],'ALLN-to-watched-PN':dfs['ALLN-to-watched-PN']} roi_rows=[];pair_rows=[] for name,frame in frames.items(): for roi,g in frame.groupby('primary_post',dropna=False,observed=True): roi_rows.append({'group':name,'primary_post':str(roi),'contacts':len(g),'fraction':len(g)/len(frame),'source_bodies':g.body_pre.nunique(),'target_bodies':g.body_post.nunique()}) for (pre,post,roi),g in frame.groupby(['body_pre','body_post','primary_post'],dropna=False,observed=True): pair_rows.append({'group':name,'body_pre':int(pre),'body_post':int(post),'primary_post':str(roi),'contacts':len(g)}) rois=pd.DataFrame(roi_rows);rois.to_csv(out/'population-ROI-summary.csv',index=False);rois.to_parquet(out/'population-ROI-summary.parquet',index=False) pd.DataFrame(pair_rows).to_parquet(out/'pair-ROI-counts.parquet',index=False) # Neuropil annotation is a location label, not evidence for receptor or inhibitory coefficient. distances=[];pnkc=dfs['APL-target-PN-to-KC'] for b,g in forward.groupby('body_post'): q=pnkc[pnkc.body_pre==b] if len(q): sites=q[['x_pre','y_pre','z_pre']].drop_duplicates().to_numpy(np.float64) query=g[['x_post','y_post','z_post']].to_numpy(np.float64) dist=cKDTree(sites).query(query,workers=1)[0]*0.008 distances.append({'PN_body_id':int(b),'PN_neuron_idx':int(lookup.loc[b]),'APL_postsynaptic_contacts':len(g),'PN_to_KC_unique_presynaptic_sites':len(sites), 'nearest_distance_microns_min':float(dist.min()),'nearest_distance_microns_median':float(np.median(dist)), 'nearest_distance_microns_p95':float(np.quantile(dist,.95)),'nearest_distance_microns_max':float(dist.max()), 'bouton_identity_confirmed':False,'watched':int(lookup.loc[b]) in watchidx}) else:distances.append({'PN_body_id':int(b),'PN_neuron_idx':int(lookup.loc[b]),'APL_postsynaptic_contacts':len(g),'PN_to_KC_unique_presynaptic_sites':0, 'bouton_identity_confirmed':False,'watched':int(lookup.loc[b]) in watchidx}) pd.DataFrame(distances).to_csv(out/'spatial-neighborhoods.csv',index=False) watched_rows=[] for node in watchidx: b=body[node];row={'neuron_idx':node,'body_id':int(b),'type':str(m.iloc[node].annotation_type)} for name,frame in frames.items(): sub=frame[frame.body_post==b] if name in ('APL-to-PN','ALLN-to-watched-PN') else frame[frame.body_pre==b] row[name]={str(k):int(v) for k,v in sub.groupby('primary_post',observed=True).size().items()} watched_rows.append(row) save(out/'watched-PN-ROI.json',watched_rows) summary={'pass':all(x['pass'] for x in checks),'source_sha256':source_hash,'full_source_rows_streamed':total, 'groups':checks,'APL_to_PN_contacts_localized':len(forward),'APL_to_PN_pairs_localized':len(pnpairs), 'APL_target_PNs_localized':int(forward.body_post.nunique()),'ROI_counts':rois.to_dict('records'), 'spatial_location_from_official_partners':True,'bouton_identity_confirmed':False,'target_receptor_confirmed':False, 'terminal_inhibition_coefficient_known':False,'new_dynamics':False,'whole_brain_functional_intervention_run':False, 'wall_seconds':time.perf_counter()-start} save(out/'summary.json',summary);print(json.dumps(summary),flush=True) if __name__=='__main__':main()