"""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()