Blame for unionmind/external_sufficiency/data.py

"""Provider-defined physical features, fixed splits, and source-only discretization.

CPU work is restricted to archive I/O/parsing and provenance. All feature
aggregation, preprocessing, probability estimation and inference use CUDA.
"""
from __future__ import annotations

import hashlib
import io
import itertools
import json
from pathlib import Path
import time
import urllib.request
import zipfile

import numpy as np
import torch

ROOT = Path('artifacts/external_decision_sufficiency_01')
PROTOCOL = Path('configs/external_sufficiency_protocol_01.json')


def sha(path):
    h = hashlib.sha256()
    with Path(path).open('rb') as f:
        while chunk := f.read(8 << 20):
            h.update(chunk)
    return h.hexdigest()


def write(path, value):
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(json.dumps(value, indent=2)+'\n')


def split_group(name):
    return int(hashlib.sha256(('20101:'+name).encode()).hexdigest(), 16) % 2


def acquire(protocol):
    entries = []
    for name, spec in protocol['sources'].items():
        path = ROOT/'sources'/f'{name}.zip'
        assert not path.exists(), 'A source archive was already opened/downloaded.'
        path.parent.mkdir(parents=True, exist_ok=True)
        start = time.perf_counter()
        urllib.request.urlretrieve(spec['url'], path)
        elapsed = time.perf_counter()-start
        with zipfile.ZipFile(path) as z:
            inventory = [{'name':i.filename,'uncompressed_bytes':i.file_size,
                          'compressed_bytes':i.compress_size,'crc32':i.CRC}
                         for i in z.infolist()]
        entries.append(dict(name=name, url=spec['url'],doi=spec['doi'],license=spec['license'],
                            sha256=sha(path),bytes=path.stat().st_size,
                            download_seconds=elapsed,inventory=inventory))
        print(json.dumps({'downloaded':name,'bytes':path.stat().st_size,'seconds':elapsed}), flush=True)
    write(ROOT/'sources/manifest.json', {'archives':entries,'downloaded_after_lock':True})


def matrix(z, name):
    candidates = [n for n in z.namelist() if n.endswith('/'+name) or n == name]
    assert len(candidates) == 1, (name, candidates)
    # Parsing is source I/O, not model inference.
    return np.loadtxt(io.BytesIO(z.read(candidates[0])), dtype=np.float64)


def hydraulic(protocol, device):
    spec = protocol['sources']['hydraulic']
    with zipfile.ZipFile(ROOT/'sources/hydraulic.zip') as z:
        profile_cpu = matrix(z, 'profile.txt').astype(np.int64)
        profile = torch.as_tensor(profile_cpu, device=device)
        columns = []
        for name in spec['features']:
            raw = torch.as_tensor(matrix(z, name+'.txt'), device=device)
            assert raw.shape[0] == len(profile)
            columns.append(raw.mean(1))
            del raw
        x = torch.stack(columns, 1)
    groups = [','.join(map(str, p[:4])) for p in profile_cpu]
    side = torch.tensor([split_group(g) for g in groups], device=device)
    tasks = []
    for t, task in enumerate(spec['tasks']):
        context = spec['context_column'][t]
        source = profile[:, context] == spec['source_context'][t]
        h = (profile[:, t] != spec['normal'][t]).long()
        for value in spec['target_contexts'][t]:
            target = profile[:, context] == value
            tasks.append(dict(dataset='hydraulic',environment=f'context{context}={value}',
                              task=task,x=x,h=h,source=source,
                              adaptation=target & (side == 0),evaluation=target & (side == 1),
                              groups=groups,record_ids=list(range(len(x))),
                              context_description=f'profile column {context}; source optimal, target {value}'))
    return tasks


def gas(protocol, device):
    spec = protocol['sources']['gas']
    arrays, labels, batches, groups, ids = [], [], [], [], []
    with zipfile.ZipFile(ROOT/'sources/gas.zip') as z:
        for batch in spec['source_batches']+spec['target_batches']:
            names = [n for n in z.namelist() if n.endswith(f'batch{batch}.dat')]
            assert len(names) == 1
            rows, hs = [], []
            for i, line in enumerate(z.read(names[0]).decode().splitlines()):
                tokens = line.split()
                if not tokens:
                    continue
                hs.append(int(tokens[0]))
                values = dict((int(v.split(':')[0])-1,float(v.split(':')[1])) for v in tokens[1:])
                rows.append([values[j] for j in spec['features_zero_based']])
                groups.append(f'batch{batch}:block{i//16}')
                ids.append(f'batch{batch}:row{i}')
            arrays.extend(rows); labels.extend(hs); batches.extend([batch]*len(rows))
    x = torch.tensor(arrays, device=device, dtype=torch.float64)
    labels = torch.tensor(labels, device=device)
    batches = torch.tensor(batches, device=device)
    side = torch.tensor([split_group(g) for g in groups], device=device)
    source_batch = torch.zeros_like(side, dtype=torch.bool)
    for batch in spec['source_batches']:
        source_batch |= batches == batch
    tasks = []
    for a, b in itertools.combinations(range(1, 7), 2):
        relevant = (labels == a) | (labels == b)
        for batch in spec['target_batches']:
            target = relevant & (batches == batch)
            tasks.append(dict(dataset='gas',environment=f'batch{batch}',task=f'gas{a}_vs_{b}',
                              x=x,h=(labels == b).long(),source=source_batch & relevant,
                              adaptation=target & (side == 0),evaluation=target & (side == 1),
                              groups=groups,record_ids=ids,
                              context_description='provider temporal batch; source batches 1,2,3'))
    return tasks


def encode_records(x, h, threshold):
    assert x.device.type == 'cuda' and threshold.device.type == 'cuda'
    binary = (x > threshold).long()
    # Exact same event convention as the frozen controlled study.
    return h*256+(binary*x.new_tensor([128,64,32,16,8,4,2,1],dtype=torch.long)).sum(1)


def prepare(protocol):
    device = protocol['device']
    start = time.perf_counter()
    tasks = hydraulic(protocol, device)+gas(protocol, device)
    torch.cuda.synchronize(device)
    parse_seconds = time.perf_counter()-start
    destination = ROOT/'prepared'
    destination.mkdir(parents=True, exist_ok=True)
    summaries, excluded = [], []
    for index, task in enumerate(tasks):
        start = time.perf_counter()
        masks = [task[k] for k in ('source','adaptation','evaluation')]
        ns = [int(m.sum()) for m in masks]
        evaluation_h = task['h'][masks[2]]
        counts = [int((evaluation_h == h).sum()) for h in (0,1)]
        group_sets = [set(g for g, keep in zip(task['groups'],m.tolist()) if keep) for m in masks]
        assert not (group_sets[1] & group_sets[2]), 'Target configuration/block leakage'
        reasons = []
        eligibility = protocol['eligibility']
        if ns[0] < eligibility['source_minimum']: reasons.append('source_support')
        if ns[1] < 64: reasons.append('adaptation_support')
        if ns[2] < eligibility['evaluation_minimum']: reasons.append('evaluation_support')
        if min(counts, default=0) < eligibility['evaluation_minimum_per_hypothesis']: reasons.append('evaluation_hypothesis_support')
        if len(group_sets[2]) < eligibility['eligible_groups_minimum']: reasons.append('evaluation_group_support')
        summary = dict(index=index,dataset=task['dataset'],environment=task['environment'],task=task['task'],
                       counts=dict(zip(('source','adaptation','evaluation'),ns)),
                       evaluation_hypothesis_counts=counts,group_counts=[len(g) for g in group_sets],
                       exclusions=reasons,context_description=task['context_description'])
        if reasons:
            excluded.append(summary)
            continue
        threshold = task['x'][masks[0]].median(0).values
        events = encode_records(task['x'], task['h'], threshold)
        arrays = {k:events[m] for k,m in zip(('source','adaptation','evaluation'),masks)}
        # Stable original record IDs and condition-group IDs permit leakage audit.
        provenance = {k:[dict(record_id=i,group=g) for i,g,keep in
                         zip(task['record_ids'],task['groups'],m.tolist()) if keep]
                      for k,m in zip(('source','adaptation','evaluation'),masks)}
        assert torch.isfinite(task['x']).all()
        torch.cuda.synchronize(device)
        summary['preprocessing_seconds'] = time.perf_counter()-start
        summary['thresholds'] = threshold.tolist()
        summary['file'] = f'task_{index:03}.pt'
        torch.save({**arrays,'thresholds':threshold},destination/summary['file'])
        write(destination/f'task_{index:03}_provenance.json',provenance)
        summaries.append(summary)
    write(destination/'manifest.json',dict(tasks=summaries,excluded=excluded,
        total_declared_tasks=len(tasks),raw_parse_and_aggregation_seconds=parse_seconds,
        fixed_binary_adapter=True,provider_defined_sequential_task=False,
        target_configuration_or_contiguous_block_leakage=False,
        pipeline_note='All recorded future bits are evaluator/archival payloads. Policy construction sees only adaptation/source samples.'))
    print(json.dumps({'prepared':len(summaries),'excluded':len(excluded),'parse_seconds':parse_seconds}),flush=True)