Digital Drosophila Archive
53 lines
2.4 KiB
Python
53 lines
2.4 KiB
Python
"""Static-only frozen encoder extraction and exact request-schedule generation."""
|
|
import ast
|
|
import hashlib
|
|
from pathlib import Path
|
|
from typing import Any
|
|
import numpy as np
|
|
|
|
|
|
def frozen_pn_drive(path, expected_sha):
|
|
"""Execute the exact one pure function, not unrelated historical experiments."""
|
|
import pandas as pd
|
|
source = Path(path).read_bytes()
|
|
assert hashlib.sha256(source).hexdigest() == expected_sha
|
|
tree = ast.parse(source, filename=str(path))
|
|
nodes = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == 'pn_drive']
|
|
assert len(nodes) == 1 and not nodes[0].decorator_list
|
|
module = ast.Module(body=nodes, type_ignores=[])
|
|
namespace = {'np': np, 'pd': pd, 'Any': Any}
|
|
exec(compile(module, str(path), 'exec'), namespace)
|
|
return namespace['pn_drive'], hashlib.sha256(ast.dump(nodes[0], include_attributes=False).encode()).hexdigest()
|
|
|
|
|
|
def requests_from_forced(forced, targets):
|
|
targets = np.asarray(targets)
|
|
assert targets.dtype == np.int32 and targets.ndim == 1 and len(np.unique(targets)) == len(targets)
|
|
assert forced.dtype == np.bool_ and forced.ndim == 2 and forced.shape[1] == len(targets)
|
|
ticks, columns = np.nonzero(forced)
|
|
requests = np.column_stack([ticks, targets[columns]]).astype(np.int32)
|
|
if len(requests):
|
|
requests = requests[np.lexsort((requests[:, 1], requests[:, 0]))]
|
|
assert np.all(np.diff(requests[:, 0].astype(np.int64) * 166700 + requests[:, 1]) > 0)
|
|
return requests
|
|
|
|
|
|
def make_requests(targets, rates, seed, steps=1000):
|
|
import torch
|
|
assert rates.dtype == np.float32 and rates.shape == targets.shape
|
|
assert np.isfinite(rates).all() and np.all(rates >= 0) and steps == 1000
|
|
original = rates * .1 / 1000.0
|
|
explicit = rates * np.float32(.1) / np.float32(1000.)
|
|
assert original.dtype == np.float32 and np.array_equal(original.view(np.uint32), explicit.view(np.uint32))
|
|
assert np.all(original <= 1)
|
|
generator = torch.Generator(device='cuda')
|
|
generator.manual_seed(int(seed))
|
|
probabilities = torch.as_tensor(original, dtype=torch.float32, device='cuda')
|
|
forced = (torch.rand((steps, len(targets)), generator=generator, device='cuda') < probabilities[None, :]).cpu().numpy()
|
|
return requests_from_forced(forced, targets)
|
|
|
|
|
|
def write_npz(path, **arrays):
|
|
with Path(path).open('xb') as stream:
|
|
np.savez_compressed(stream, **arrays)
|