113 lines
8.7 KiB
Python
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()
|