Files
digital-drosophila/models/male-cns-lif/apl-pn-scope-audit-v1/localize.py
T

113 lines
8.7 KiB
Python

"""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()