reproduction-bundle/src/run_real_diffusion_logic.py

#!/usr/bin/env python3
"""Exercise LOGDIFF's ME/NOT calculus on a public trained MNIST DDPM."""

from __future__ import annotations

import argparse
import csv
import hashlib
import importlib.util
import json
import math
from pathlib import Path

import torch
import torch.nn.functional as F
from diffusers import DDPMScheduler, UNet2DModel
from huggingface_hub import hf_hub_download
from torchvision import datasets, transforms
from torchvision.utils import save_image


ROOT = Path(__file__).resolve().parents[2]
MODEL_REPO = "Shiroi-max/ddpm-mnist-conditional-32"
MODEL_REVISION = "368673bd7cab30e3e04af9ce5ecb1da2bf18f30f"


def load_auxiliary_module():
    path = Path(__file__).with_name("train_public_mnist_classifiers.py")
    spec = importlib.util.spec_from_file_location("public_mnist_classifiers", path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def select_device(name: str) -> torch.device:
    if name != "auto":
        return torch.device(name)
    if torch.backends.mps.is_available():
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def subset_mixture(
    primitives: torch.Tensor, probabilities: torch.Tensor, indices: list[int]
) -> torch.Tensor:
    selected_p = probabilities[:, indices]
    selected = primitives[:, indices]
    denominator = selected_p.sum(1).clamp_min(1e-12)
    weights = selected_p / denominator.unsqueeze(1)
    return (weights[:, :, None, None, None] * selected).sum(1)


def recursive_or_me(
    primitives: torch.Tensor, probabilities: torch.Tensor, indices: list[int]
) -> torch.Tensor:
    if len(indices) != len(set(indices)):
        raise ValueError("OR-ME operands must be mutually exclusive; duplicate/overlapping atoms found")
    score = primitives[:, indices[0]]
    event_probability = probabilities[:, indices[0]]
    for index in indices[1:]:
        next_probability = probabilities[:, index]
        total = (event_probability + next_probability).clamp_min(1e-12)
        score = (
            event_probability[:, None, None, None] * score
            + next_probability[:, None, None, None] * primitives[:, index]
        ) / total[:, None, None, None]
        event_probability = total
    return score


def complement_mixture(
    primitives: torch.Tensor, probabilities: torch.Tensor, excluded: int
) -> tuple[torch.Tensor, torch.Tensor]:
    unconditional = (probabilities[:, :, None, None, None] * primitives).sum(1)
    p = probabilities[:, excluded]
    via_not = (
        unconditional - p[:, None, None, None] * primitives[:, excluded]
    ) / (1.0 - p).clamp_min(1e-12)[:, None, None, None]
    direct = subset_mixture(primitives, probabilities, [i for i in range(10) if i != excluded])
    return via_not, direct


@torch.inference_mode()
def primitive_predictions(
    unet: UNet2DModel, x: torch.Tensor, timestep: int | torch.Tensor
) -> torch.Tensor:
    batch = x.shape[0]
    if isinstance(timestep, int):
        t = torch.full((batch,), timestep, device=x.device, dtype=torch.long)
    else:
        t = timestep.to(x.device)
    repeated_x = x[:, None].expand(-1, 10, -1, -1, -1).reshape(
        batch * 10, *x.shape[1:]
    )
    repeated_t = t[:, None].expand(-1, 10).reshape(-1)
    labels = torch.arange(10, device=x.device).repeat(batch)
    prediction = unet(repeated_x, repeated_t, class_labels=labels).sample
    return prediction.reshape(batch, 10, *x.shape[1:])


@torch.inference_mode()
def fixed_timestep_audit(
    unet: UNet2DModel,
    classifier: torch.nn.Module,
    loader: torch.utils.data.DataLoader,
    scheduler: DDPMScheduler,
    device: torch.device,
    timesteps: list[int],
    max_examples: int,
    seed: int,
) -> tuple[list[dict[str, object]], dict[str, float | bool]]:
    metric_rows: list[dict[str, object]] = []
    identity = {
        "or_me_max_abs_error": 0.0,
        "not_max_abs_error": 0.0,
        "constant_weight_min_max_abs_difference": float("inf"),
        "overlap_control_rejected": False,
    }
    try:
        recursive_or_me(
            torch.zeros(1, 10, 1, 1, 1), torch.full((1, 10), 0.1), [1, 1]
        )
    except ValueError:
        identity["overlap_control_rejected"] = True

    alpha = scheduler.alphas_cumprod.float()
    for timestep in timesteps:
        generator = torch.Generator(device="cpu").manual_seed(seed + timestep)
        correct = total = 0
        nll_sum = 0.0
        processed_for_identity = 0
        for clean, target in loader:
            remaining = max_examples - total
            if remaining <= 0:
                break
            clean = clean[:remaining]
            target = target[:remaining]
            noise = torch.randn(clean.shape, generator=generator)
            a = alpha[timestep]
            noisy = a.sqrt() * clean + (1.0 - a).sqrt() * noise
            noisy_device = noisy.to(device)
            t = torch.full((clean.shape[0],), timestep, device=device, dtype=torch.long)
            logits = classifier(noisy_device, t)
            probabilities = logits.softmax(1)
            target_device = target.to(device)
            correct += int((logits.argmax(1) == target_device).sum())
            nll_sum += float(F.cross_entropy(logits, target_device, reduction="sum").cpu())
            total += target.numel()

            if processed_for_identity < 128:
                take = min(32, 128 - processed_for_identity, clean.shape[0])
                primitives = primitive_predictions(unet, noisy_device[:take], timestep)
                p = probabilities[:take]
                for subset in ([1, 4], [3, 7, 8], [0, 2, 5, 6]):
                    direct = subset_mixture(primitives, p, list(subset))
                    recursive = recursive_or_me(primitives, p, list(subset))
                    identity["or_me_max_abs_error"] = max(
                        float(identity["or_me_max_abs_error"]),
                        float((direct - recursive).abs().max().cpu()),
                    )
                    constant = primitives[:, list(subset)].mean(1)
                    identity["constant_weight_min_max_abs_difference"] = min(
                        float(identity["constant_weight_min_max_abs_difference"]),
                        float((direct - constant).abs().max().cpu()),
                    )
                for excluded in (0, 4, 9):
                    via_not, direct_not = complement_mixture(primitives, p, excluded)
                    identity["not_max_abs_error"] = max(
                        float(identity["not_max_abs_error"]),
                        float((via_not - direct_not).abs().max().cpu()),
                    )
                processed_for_identity += take

        metric_rows.append(
            {
                "timestep": timestep,
                "examples": total,
                "accuracy": correct / total,
                "nll": nll_sum / total,
            }
        )
    identity["passed"] = bool(
        float(identity["or_me_max_abs_error"]) < 2e-5
        and float(identity["not_max_abs_error"]) < 2e-5
        and float(identity["constant_weight_min_max_abs_difference"]) > 1e-5
        and identity["overlap_control_rejected"]
    )
    return metric_rows, identity


def entropy_from_counts(counts: list[int]) -> float:
    total = sum(counts)
    if not total:
        return 0.0
    return -sum((n / total) * math.log2(n / total) for n in counts if n)


@torch.inference_mode()
def generate(
    unet: UNet2DModel,
    posterior: torch.nn.Module,
    judge: torch.nn.Module,
    scheduler: DDPMScheduler,
    subset: list[int],
    method: str,
    samples: int,
    batch_size: int,
    steps: int,
    device: torch.device,
    seed: int,
    grid_path: Path,
) -> dict[str, object]:
    counts = [0] * 10
    saved: list[torch.Tensor] = []
    generated = 0
    generator = torch.Generator(device="cpu").manual_seed(seed)
    scheduler.set_timesteps(steps)
    while generated < samples:
        current = min(batch_size, samples - generated)
        x = torch.randn((current, 1, 32, 32), generator=generator).to(device)
        for timestep in scheduler.timesteps:
            t_value = int(timestep)
            primitives = primitive_predictions(unet, x, t_value)[:, subset]
            if method == "logdiff":
                t = torch.full((current,), t_value, device=device, dtype=torch.long)
                p = posterior(x, t).softmax(1)[:, subset]
                weights = p / p.sum(1, keepdim=True).clamp_min(1e-12)
            elif method == "constant":
                weights = torch.full(
                    (current, len(subset)), 1.0 / len(subset), device=device
                )
            else:
                raise ValueError(method)
            epsilon = (weights[:, :, None, None, None] * primitives).sum(1)
            x = scheduler.step(epsilon, timestep, x, generator=generator).prev_sample
        images = (x / 2.0 + 0.5).clamp(0, 1)
        predictions = judge(images).argmax(1).cpu()
        bincount = torch.bincount(predictions, minlength=10).tolist()
        counts = [left + right for left, right in zip(counts, bincount)]
        if len(saved) < 32:
            saved.extend(list(images.cpu()[: 32 - len(saved)]))
        generated += current
    grid_path.parent.mkdir(parents=True, exist_ok=True)
    save_image(torch.stack(saved), grid_path, nrow=8)
    conformity = sum(counts[index] for index in subset) / samples
    target_counts = [counts[index] for index in subset]
    return {
        "method": method,
        "subset": subset,
        "samples": samples,
        "steps": steps,
        "judge_counts_0_to_9": counts,
        "subset_conformity": conformity,
        "subset_label_entropy_bits": entropy_from_counts(target_counts),
        "grid": grid_path.name,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--classifier-dir",
        type=Path,
        default=ROOT / "outputs" / "public-mnist-classifiers",
    )
    parser.add_argument(
        "--data-root",
        type=Path,
        default=ROOT / "outputs" / "public-mnist-classifiers" / "data",
    )
    parser.add_argument(
        "--output-dir", type=Path, default=ROOT / "outputs" / "real-diffusion-logic"
    )
    parser.add_argument("--timesteps", type=int, nargs="+", default=[0, 50, 100, 250, 500, 750, 999])
    parser.add_argument("--audit-examples", type=int, default=10000)
    parser.add_argument("--generation-samples", type=int, default=2000)
    parser.add_argument("--generation-batch-size", type=int, default=128)
    parser.add_argument("--generation-steps", type=int, default=50)
    parser.add_argument("--subsets", nargs="+", default=["1,4", "3,7,8"])
    parser.add_argument("--device", default="auto")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    device = select_device(args.device)
    auxiliary = load_auxiliary_module()
    noisy_path = args.classifier_dir / "noisy_classifier.pt"
    judge_path = args.classifier_dir / "clean_judge.pt"
    if not noisy_path.is_file() or not judge_path.is_file():
        raise FileNotFoundError(
            "Run train_public_mnist_classifiers.py to completion before this audit"
        )
    posterior = auxiliary.NoisyDigitClassifier().to(device)
    posterior.load_state_dict(torch.load(noisy_path, map_location="cpu", weights_only=True))
    posterior.eval()
    judge = auxiliary.CleanDigitClassifier().to(device)
    judge.load_state_dict(torch.load(judge_path, map_location="cpu", weights_only=True))
    judge.eval()

    unet = UNet2DModel.from_pretrained(
        MODEL_REPO, subfolder="unet", revision=MODEL_REVISION
    ).eval().to(device)
    scheduler = DDPMScheduler.from_pretrained(
        MODEL_REPO, subfolder="scheduler", revision=MODEL_REVISION
    )
    transform = transforms.Compose(
        [transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]
    )
    test_set = datasets.MNIST(args.data_root, train=False, download=True, transform=transform)
    test_loader = torch.utils.data.DataLoader(
        test_set, batch_size=args.generation_batch_size, shuffle=False, num_workers=0
    )
    fixed_metrics, identity = fixed_timestep_audit(
        unet,
        posterior,
        test_loader,
        scheduler,
        device,
        args.timesteps,
        args.audit_examples,
        args.seed,
    )

    args.output_dir.mkdir(parents=True, exist_ok=True)
    with (args.output_dir / "fixed_timestep_classifier.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(fixed_metrics[0]))
        writer.writeheader()
        writer.writerows(fixed_metrics)
    generation_results: list[dict[str, object]] = []
    for subset_text in args.subsets:
        subset = [int(value) for value in subset_text.split(",")]
        if len(subset) != len(set(subset)) or any(value not in range(10) for value in subset):
            raise ValueError(f"Invalid mutually-exclusive digit subset: {subset_text}")
        for method in ("logdiff", "constant"):
            generation_results.append(
                generate(
                    unet,
                    posterior,
                    judge,
                    scheduler,
                    subset,
                    method,
                    args.generation_samples,
                    args.generation_batch_size,
                    args.generation_steps,
                    device,
                    args.seed + sum(subset),
                    args.output_dir / f"samples_{method}_{'-'.join(map(str, subset))}.png",
                )
            )

    weight_path = Path(
        hf_hub_download(
            MODEL_REPO,
            "unet/diffusion_pytorch_model.safetensors",
            revision=MODEL_REVISION,
        )
    )
    report = {
        "paper": "OAM1jJsMGp",
        "anchored_claim": 1,
        "scope": (
            "Real trained diffusion + full MNIST validation of mutually-exclusive disjunction "
            "and negation. CI conjunction/disjunction are covered by the separate exhaustive formal suite."
        ),
        "diffusion_model": {
            "repo": MODEL_REPO,
            "revision": MODEL_REVISION,
            "weight_sha256": sha256(weight_path),
            "parameters": sum(parameter.numel() for parameter in unet.parameters()),
        },
        "auxiliary_classifiers": {
            "noisy_checkpoint_sha256": sha256(noisy_path),
            "judge_checkpoint_sha256": sha256(judge_path),
            "noisy_parameters": sum(parameter.numel() for parameter in posterior.parameters()),
            "judge_parameters": sum(parameter.numel() for parameter in judge.parameters()),
        },
        "fixed_timestep_classifier": fixed_metrics,
        "exact_identity": identity,
        "generation": generation_results,
    }
    (args.output_dir / "real_diffusion_logic_report.json").write_text(
        json.dumps(report, indent=2) + "\n"
    )
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()