Digital Drosophila Archive
132 lines
7.6 KiB
Python
132 lines
7.6 KiB
Python
"""Exact frozen L rate transform and full-shape CUDA sampler, no brain/update.
|
|
|
|
Static table derivation is CPU/read-only. Sampling is a separate explicit job.
|
|
Common-random aliases share only saved immutable CS bytes, never brain outputs.
|
|
"""
|
|
from __future__ import annotations
|
|
import argparse
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import time
|
|
import numpy as np
|
|
from normal_io import *
|
|
|
|
PINS = {
|
|
'conditioning_inputs.py': '19382e066e450140bb653a3661f980ffe7cc62f412020ef57df0b6e6b3cbb8ba',
|
|
'conditioning_preflight_inputs.py': '846d2141b472f18b1bb728eed4aa33876896d4b33dece364787c81204f62995f',
|
|
'trial_plan.py': '6fed3f494d075339a7e5b001d43b2598640b189e038952fbd4d41200bacac587',
|
|
'seed_registry.py': '0da033b132095ec3087b5d9f64e53c2e391fa50ab0b7653b31296eeb84799b7e',
|
|
}
|
|
|
|
def dependencies():
|
|
verify({str(LAB/n): s for n,s in PINS.items()})
|
|
require(not any(n in sys.modules for n in ('runtime_light','conditioning_epoch')), 'No neural modules in sampler')
|
|
sys.path.insert(0, str(LAB))
|
|
import conditioning_preflight_inputs as source
|
|
import conditioning_inputs as inputs
|
|
require(Path(source.__file__) == LAB/'conditioning_preflight_inputs.py' and
|
|
Path(inputs.__file__) == LAB/'conditioning_inputs.py', 'Wrong input module identity')
|
|
return source, inputs
|
|
|
|
def derive_rates(root=ROOT):
|
|
"""Calls the original pinned pn_drive for both registered modes."""
|
|
import pandas as pd
|
|
import yaml
|
|
source, inputs = dependencies()
|
|
loaded = source.load_rate_tables(root)
|
|
static = Path(root)/source.STATIC
|
|
selection = read(static/'odor-encoder/selection.json')
|
|
panel = pd.read_parquet(static/'odor-encoder/frozen-odor-panel.parquet')
|
|
mapping = pd.read_parquet(static/'receptor-mapping/olfactory-mapping.parquet')
|
|
config = yaml.safe_load((Path(root)/source.CONFIG_REL).read_text())
|
|
drive, ast_sha = source._extract_drive((Path(root)/source.SOURCE_REL).read_bytes(), str(Path(root)/source.SOURCE_REL))
|
|
vectors = npz(static/'odor-encoder/mapped-glomerular-vectors.npz')
|
|
result = {}
|
|
for mode in ('physiological','matched_energy'):
|
|
for cue in ('A','B'):
|
|
original = loaded['tables'][cue]
|
|
targets, rates, frame = drive(original['odor_id'], vectors[original['odor_id']], mapping,
|
|
selection['glomeruli_order'], mode, loaded['audit']['median_mapped_response_sum'], config)
|
|
inputs.probabilities_f32(targets, rates)
|
|
require(np.array_equal(frame.neuron_idx.to_numpy(np.int32), targets) and
|
|
frame.rate_hz.to_numpy(np.float32).tobytes() == rates.tobytes(), 'Original PN order differs')
|
|
if mode == 'physiological':
|
|
require(all(original[n].tobytes() == a.tobytes() for n,a in (('targets',targets),('rates_hz',rates))),
|
|
'Physiological reconstruction differs')
|
|
prior = npz(Path(root)/f'experiments/phase4ao-conditioning-interface/inputs/static/{cue}-rates.npz')
|
|
require(all(prior[n].dtype == a.dtype and prior[n].shape == a.shape and prior[n].tobytes() == a.tobytes()
|
|
for n,a in (('targets',targets),('rates_hz',rates))), 'Prior AO rate bytes differ')
|
|
result[mode+'--'+cue] = dict(targets=targets, rates_hz=rates)
|
|
verify({str(Path(root)/p):s for p,s in source.SOURCE_PINS.items()})
|
|
return result, dict(original_source_sha256=source.SOURCE_PINS, original_AST_sha256=ast_sha,
|
|
physiological_matches_original_AO=True, original_PN_column_order_preserved=True,
|
|
modes=['physiological','matched_energy'], clipping=False, gain_tuning=False,
|
|
historical_dynamic_results_read=False, random_draws=0)
|
|
|
|
def descriptor(task):
|
|
r = task['recipe']
|
|
require(task['steps'] == 1200 and task['mode'] in ('physiological','matched_energy'), 'Wrong normal clock/mode')
|
|
if task['kind'] == 'train':
|
|
cs = r['cs']
|
|
if cs is None: return None
|
|
cue, seed, start, stop = cs['cue'],cs['seed'],cs['start_tick'],cs['stop_tick_exclusive']
|
|
stream = cs['source_stream_id']
|
|
else:
|
|
require(task['kind'] == 'probe' and r['US_requested'] is False and r['weight_updates_enabled'] is False,
|
|
'Probe contains US/update')
|
|
cue, seed = r['cue'], r['input']['seed']
|
|
start, stop, stream = r['cs_start_tick'], r['cs_stop_tick_exclusive'], r['input']['stream_id']
|
|
require(cue in ('A','B') and type(seed) is int and 0 <= seed < 2**63 and (start,stop) == (0,1000),
|
|
'CS timing/full seed differs')
|
|
return dict(mode=task['mode'],cue=cue,seed=seed,steps=1200,cs_start_tick=start,cs_stop_tick_exclusive=stop,
|
|
source_stream_id=stream,table_id=task['mode']+'--'+cue)
|
|
|
|
def inventory_from_plan(plan):
|
|
descriptors, task_inputs = {}, {}
|
|
for task in plan['tasks']:
|
|
d = descriptor(task)
|
|
key = None if d is None else content_sha(d)
|
|
if d is not None:
|
|
require(key not in descriptors or descriptors[key] == d, 'Conflicting shared CS descriptor')
|
|
descriptors[key] = d
|
|
require(task['task_id'] not in task_inputs, 'Duplicate planned task')
|
|
task_inputs[task['task_id']] = key
|
|
require(len(descriptors) == 4080, 'Expected 4,000 training + 80 evaluation streams')
|
|
return dict(schema='normal-CS-input-inventory-v1',descriptors=descriptors,task_inputs=task_inputs,
|
|
brain_output_deduplication=False,CS_only_saved=True,US_composed_from_each_original_epoch=True)
|
|
|
|
def sample(reg_sha):
|
|
require(sys.dont_write_bytecode, 'Use -B')
|
|
file = EXP/'protocol/preregistered.json'; require(digest(file) == reg_sha, 'Registration differs')
|
|
reg = read(file); verify(reg['source_sha256']); verify(reg['static_input_sha256'])
|
|
require(reg['scientific_scope'] == 'normal-acquisition-only', 'Wrong scientific scope')
|
|
out = unlinked(EXP/'inputs/generated'); require(not out.exists(), 'Sampler cannot retry existing output')
|
|
h = load_source('_normal_sampler_admission', AP/'isolated_workers.py', AP_SHA)
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = GPU
|
|
admission = h._admit_brain()
|
|
source, inputs = dependencies()
|
|
inv = read(EXP/'protocol/input-inventory.json')
|
|
require(digest(EXP/'protocol/input-inventory.json') == reg['prepared_sha256'][str(EXP/'protocol/input-inventory.json')],
|
|
'Input inventory differs')
|
|
out.mkdir(); start = time.perf_counter(); completed = {}
|
|
for index,(key,d) in enumerate(inv['descriptors'].items()):
|
|
if index % 100 == 0: h._admit_brain()
|
|
table = npz(EXP/f"inputs/static/{d['table_id']}.npz")
|
|
require(digest(EXP/f"inputs/static/{d['table_id']}.npz") ==
|
|
reg['prepared_sha256'][str(EXP/f"inputs/static/{d['table_id']}.npz")], 'Rate file changed')
|
|
result = inputs.generate_cs_requests(table['targets'],table['rates_hz'],seed=d['seed'],steps=1200,
|
|
cs_start_tick=0,cs_stop_tick=1000,device='cuda:0',max_uniform_draws=reg['budget']['max_uniform_draws'])
|
|
save_npz(out/(key+'.npz'),result['requests'])
|
|
save(out/(key+'.json'),dict(descriptor=d,audit=result['audit'],requests_sha256=digest(out/(key+'.npz'))))
|
|
completed[key] = dict(requests=pin(out/(key+'.npz')),audit=pin(out/(key+'.json')))
|
|
verify(reg['source_sha256']); verify(reg['static_input_sha256'])
|
|
save(EXP/'receipts/input-generation.json',dict(complete=True,registration_sha256=reg_sha,
|
|
input_inventory_sha256=digest(EXP/'protocol/input-inventory.json'),inputs=completed,
|
|
descriptors=len(completed),wall_seconds=time.perf_counter()-start,admission=admission,
|
|
brain_executed=False,weight_updates_executed=False))
|
|
|
|
if __name__ == '__main__':
|
|
p=argparse.ArgumentParser();p.add_argument('--registration-sha256',required=True)
|
|
sample(p.parse_args().registration_sha256)
|