Blame for PASTIS/main_method/code/foundation_transfer.py

"""Portable frozen TESSERA optical feature extraction, readout fitting and score replay.

Published historical research scripts are preserved separately. This entry point
never trains the encoder or chooses parameters using test labels.
"""
import argparse
import hashlib
import json
from datetime import timedelta
from pathlib import Path
import numpy as np
import torch

MAIN = Path(__file__).resolve().parents[1]
REPO = MAIN.parent
BUNDLE = MAIN / 'checkpoints/32X32_tessera_transfer_seed3407'
WORK = REPO / 'ablations/10_plateau_breakthrough/workspace'
ENCODER_SHA = '92619d4ffc2936895f145fdcf710140b58c23c860e02f01873828b10c6958d95'
READOUT_SHA = '5959d5f58fab2dfdff458281081f531b626965a8283458aaae61aa9857563753'
BAND_ORDER = [2, 0, 1, 6, 7, 3, 4, 5, 8, 9]


def digest(path):
    h = hashlib.sha256()
    with Path(path).open('rb') as stream:
        for block in iter(lambda: stream.read(8 * 1024 * 1024), b''):
            h.update(block)
    return h.hexdigest()


def verify():
    for filename, expected in [('frozen_satellite_encoder.pt', ENCODER_SHA),
                               ('optical_gaussian_classifier.npz', READOUT_SHA)]:
        actual = digest(BUNDLE / filename)
        if actual != expected:
            raise RuntimeError(f'Hash mismatch: {filename}')
    return {'encoder_sha256': ENCODER_SHA, 'readout_sha256': READOUT_SHA, 'passed': True}


def prepare(x, pos, days):
    """Undo legacy loader normalization, preserve real observations and map bands."""
    import tessera_student as student
    from dataset import get_norm, REF_DATE
    assert len(x) == 1 and x.shape[2] == 10
    valid = pos[0].gt(0)
    valid[0] = True  # Acquisition zero is real; subsequent zero positions pad.
    n = int(valid.sum())
    assert n > 0 and valid[:n].all() and not valid[n:].any()
    mean, std = get_norm()
    raw = x[0, :n].float() * torch.as_tensor(std + 1e-8, device=x.device)[None, :, None, None]
    raw = raw + torch.as_tensor(mean, device=x.device)[None, :, None, None]
    raw = raw[:, BAND_ORDER].permute(2, 3, 0, 1).reshape(-1, n, 10)
    doy = torch.tensor([(REF_DATE + timedelta(days=int(d))).timetuple().tm_yday
                        for d in days[0, :n].tolist()], device=x.device)
    remain = ((n + 7) // 8) * 8 - n
    if remain:
        groups = np.array_split(np.arange(n), remain)
        fill = np.array([g[len(g) // 2] for g in groups]) if remain <= n else np.arange(remain) % n
        indices = torch.as_tensor(np.concatenate([np.arange(n), fill]), device=x.device)
        raw, doy = raw[:, indices], doy[indices]
    bands = (raw - torch.as_tensor(student.S2_BAND_MEAN, device=x.device)) / (
        torch.as_tensor(student.S2_BAND_STD, device=x.device) + 1e-9)
    return torch.cat([bands, doy[None, :, None].expand(len(bands), -1, 1)], -1)


@torch.inference_mode()
def optical_features(model, inputs):
    chunks = []
    for start in range(0, len(inputs), 512):
        with torch.autocast('cuda', dtype=torch.float16, enabled=inputs.is_cuda):
            features = model.s2_backbone(inputs[start:start + 512])
        chunks.append(features.float().cpu().numpy())
    return np.concatenate(chunks)[:, None, :]


def predict(features, params):
    """Shared-covariance Gaussian discriminant scores, then calibrated softmax."""
    x = (np.asarray(features[:, 0], np.float32) - params['mean']) / params['std']
    scores = x @ params['weight'] + params['bias']
    p = np.zeros((len(x), 20), np.float32)
    p[:, 1:19] = torch.from_numpy(np.asarray(scores / float(params['temperature']), np.float32)).softmax(1).numpy()
    return p


def fit(features, labels, output):
    """Refit only the validation-frozen shrinkage=.1, temperature=4 configuration."""
    f = np.load(features, mmap_mode='r'); y = np.load(labels)
    keep = (y > 0) & (y < 19); y = y[keep].astype(np.int64) - 1
    x = np.asarray(f[:, 0], np.float32)[keep]
    mean, std = x.mean(0), x.std(0).clip(.01)
    x = (x - mean) / std
    centers = np.stack([x[y == c].mean(0) for c in range(18)])
    residual = x - centers[y]
    covariance = residual.T @ residual / len(x)
    values, vectors = np.linalg.eigh(covariance.astype(np.float64))
    inverse = (vectors / (.9 * values + .1 * values.mean()).clip(1e-6)) @ vectors.T
    weight = (inverse @ centers.T).astype(np.float32)
    bias = -.5 * (centers * weight.T).sum(1)
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, depth=0, mean=mean, std=std, type='linear', weight=weight, bias=bias, temperature=4.)
    return {'foreground_training_pixel_entries': len(x), 'output': str(output), 'encoder_updated': False}


def confusion(labels, prediction):
    valid = (labels > 0) & (labels < 19)
    return np.bincount(20 * labels[valid].astype(np.int64) + prediction[valid], minlength=400).reshape(20, 20)


def metrics(cm):
    tp = cm.diagonal()[1:19].astype(float); gt = cm.sum(1)[1:19]; pr = cm.sum(0)[1:19]
    iou = tp / np.maximum(gt + pr - tp, 1)
    f1 = 2 * tp / np.maximum(gt + pr, 1)
    total = cm.sum(); oa = cm.trace() / max(total, 1)
    chance = np.dot(cm.sum(0).astype(float), cm.sum(1)) / max(total * total, 1)
    return dict(miou=float(iou.mean()), oa=float(oa), macro_f1=float(f1.mean()),
                kappa=float((oa - chance) / max(1 - chance, 1e-12)), valid_pixels=int(total),
                per_class_iou=iou.tolist(), confusion_matrix=cm.tolist())


def evaluate_cached(check_raw_labels=False):
    """Replay nested probabilities exactly; optionally check every raw test target."""
    audit = json.loads((MAIN / 'configurations/32X32_tessera_transfer_seed3407/completion_audit.json').read_text())
    protocol = json.loads((WORK / 'research_goal_52/PROTOCOL.json').read_text())
    labels = np.load(WORK / 'convnext_research_20261010/test_labels.npy', mmap_mode='r')
    paths = ['convnext_research_20261010/recommended_test_probabilities.npy',
             'score_breakthrough_20261010/gaussian_test_probabilities.npy',
             'score_breakthrough_20261010/structured_test_probabilities.npy',
             'tessera_transfer_research_20261010/32X32_pretrained_optical_lda_seed3407/test_probabilities.npy']
    arrays = [np.load(WORK / p, mmap_mode='r') for p in paths]
    ids = protocol['roles']['test']['ordered_ids']
    assert len(ids) == len(labels) == 482 and len(set(ids)) == 482
    cm = np.zeros((20, 20), np.int64)
    if check_raw_labels:
        import config as C
    for i, pid in enumerate(ids):
        if check_raw_labels:
            assert np.array_equal(labels[i], np.load(C.ANNOT_DIR / f'TARGET_{pid}.npy')[0])
        p0, pg, ps, pt = [a[i] for a in arrays]
        assert all(p.shape == (128, 128, 20) and np.isfinite(p).all() for p in [p0, pg, ps, pt])
        incumbent = .65 * p0 + .25 * pg + .1 * ps
        cm += confusion(labels[i], (.5 * incumbent + .5 * pt).argmax(-1))
    result = metrics(cm)
    expected = audit['results']['pretrained_optical_lda']['test']
    assert result['confusion_matrix'] == expected['confusion_matrix']
    return {'passed': True, 'mode': 'saved_probability_replay', 'raw_labels_checked': check_raw_labels,
            'patches': 482, 'metrics': result}


@torch.inference_mode()
def extract(split, output, device):
    """Fresh raw-data optical inference. Train/val preserve original batch/crop order."""
    import tessera_student as student
    import config as C
    from torch.utils.data import DataLoader
    from dataset import PASTISPatchDataset, pad_collate, crop_ij, train_patch_ids, val_patch_ids, test_patch_ids
    verify(); torch.set_num_threads(4); torch.manual_seed(3407)
    torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True
    model = student.load_model(str(BUNDLE / 'frozen_satellite_encoder.pt'), torch.device(device)).eval().requires_grad_(False)
    ids = {'train': train_patch_ids, 'val': val_patch_ids, 'test': test_patch_ids}[split]()
    dl = DataLoader(PASTISPatchDataset(ids), batch_size=1 if split == 'test' else 2,
                    shuffle=False, num_workers=2, collate_fn=pad_collate)
    output.mkdir(parents=True, exist_ok=True)
    count = len(ids) * (16384 if split == 'test' else 2048)
    fw = np.lib.format.open_memmap(output / 'features.npy', mode='w+', dtype=np.float16, shape=(count, 1, 256))
    yw = np.lib.format.open_memmap(output / 'labels.npy', mode='w+', dtype=np.uint8, shape=(count,))
    pw = np.lib.format.open_memmap(output / 'probabilities.npy', mode='w+', dtype=np.float32, shape=(count, 20))
    params = np.load(BUNDLE / 'optical_gaussian_classifier.npz'); offset = 0; cm = np.zeros((20, 20), np.int64)
    for bi, ((x, pos, days), y) in enumerate(dl):
        for i, j in C.FT_CROP_IJ if split != 'test' else [(0, 0)]:
            crop, target = crop_ij(x, y, i, j, 4 if split != 'test' else 1)
            for sample in range(len(crop)):
                inputs = prepare(crop[sample:sample+1].to(device), pos[sample:sample+1].to(device), days[sample:sample+1])
                f = optical_features(model, inputs); n = len(f)
                # Fitting used half-precision feature caches; test used fresh float32 features.
                scored = f if split == 'test' else f.astype(np.float16).astype(np.float32)
                p = predict(scored, params); gt = target[sample].numpy().ravel()
                fw[offset:offset+n] = f; yw[offset:offset+n] = gt; pw[offset:offset+n] = p
                cm += confusion(gt, p.argmax(-1)); offset += n
        if bi % 25 == 0: print(split, bi, len(dl), flush=True)
    assert offset == count
    fw.flush(); yw.flush(); pw.flush()
    return {'split': split, 'patch_entries': len(ids), 'pixel_entries': count, 'standalone_metrics': metrics(cm)}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=['verify', 'evaluate-cached', 'extract', 'fit'])
    parser.add_argument('--check-raw-labels', action='store_true')
    parser.add_argument('--split', choices=['train', 'val', 'test'], default='val')
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--output', type=Path, default=MAIN / 'reproduction')
    parser.add_argument('--features', type=Path); parser.add_argument('--labels', type=Path)
    args = parser.parse_args()
    torch.set_num_threads(4)
    if args.action == 'verify': result = verify()
    elif args.action == 'evaluate-cached': result = evaluate_cached(args.check_raw_labels)
    elif args.action == 'extract': result = extract(args.split, args.output, args.device)
    else:
        if args.features is None or args.labels is None: parser.error('fit requires --features and --labels')
        result = fit(args.features, args.labels, args.output / 'refitted_readout.npz')
    print(json.dumps(result, indent=2))
    if args.action != 'verify':
        args.output.mkdir(parents=True, exist_ok=True)
        (args.output / f'{args.action}_{args.split}.json').write_text(json.dumps(result, indent=2))


if __name__ == '__main__': main()