refresh_format2.py

"""Regenerate the DeepDeWedge FORMAT 2 package from its authoritative source.

This one-off maintenance converter is a package resource, not Scitomo runtime
code.  It intentionally refuses all historical Hugging Face payloads.
"""

from __future__ import annotations

import hashlib
import platform
import shutil
import subprocess
import sys
import types
from pathlib import Path

import pytorch_lightning
import safetensors
import torch

import scitomo as st
from scitomo.methods.restoration.deepdewedge.network_invocation import (
    invoke_deepdewedge_network,
)


SCRIPT_PATH = Path(__file__).resolve()
ROOT = (
    SCRIPT_PATH.parents[2]
    if SCRIPT_PATH.parent.name == "conversion"
    else SCRIPT_PATH.parents[1]
)
UPSTREAM = ROOT / "upstream"
CHECKPOINT = ROOT / "official" / "fitted_model.ckpt"
ARCHIVE = ROOT / "official" / "tutorial_data.zip"
TARGET = ROOT / "hf"
OUTPUT = ROOT / "package-fresh"
RUNTIME_VIEW = ROOT / "package-runtime-view"

UPSTREAM_REVISION = "072075692a44a8f17394214369e6e762abe52bc3"
CHECKPOINT_SIZE = 327952642
CHECKPOINT_SHA256 = "5262f6c11e85fd662b02e59efe936fa7b69913758e841235be2683f7bd03ec76"
ARCHIVE_SHA256 = "7c871342e51f5a66a773fe427d72944b5d2cc8ff41c5b7415ab38dbfc9ac6d58"
ARCHIVE_MD5 = "130264af7d96be6237351f8f51eda8c8"
PREVIOUS_HF_COMMIT = "87db06570dd874a99af1289e62b79ea99f87f006"


def _digest(path: Path, algorithm: str) -> str:
    hasher = hashlib.new(algorithm)
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            hasher.update(block)
    return hasher.hexdigest()


def _verify_authority() -> None:
    """Fail closed before the sole permitted Lightning deserialization."""

    if CHECKPOINT.stat().st_size != CHECKPOINT_SIZE:
        raise RuntimeError("Authoritative checkpoint byte size does not match.")
    if _digest(CHECKPOINT, "sha256") != CHECKPOINT_SHA256:
        raise RuntimeError("Authoritative checkpoint SHA-256 does not match.")
    if _digest(ARCHIVE, "sha256") != ARCHIVE_SHA256:
        raise RuntimeError("Authoritative archive SHA-256 does not match.")
    if _digest(ARCHIVE, "md5") != ARCHIVE_MD5:
        raise RuntimeError("Authoritative archive MD5 does not match.")
    revision = subprocess.check_output(
        ["git", "-C", str(UPSTREAM), "rev-parse", "HEAD"], text=True
    ).strip()
    if revision != UPSTREAM_REVISION:
        raise RuntimeError("Pinned upstream checkout revision does not match.")


def _load_upstream_model() -> torch.nn.Module:
    """Load the exact pinned model without executing its unrelated CLI package init."""

    ddw = types.ModuleType("ddw")
    ddw.__path__ = [str(UPSTREAM / "ddw")]
    sys.modules["ddw"] = ddw
    utils = types.ModuleType("ddw.utils")
    utils.__path__ = [str(UPSTREAM / "ddw" / "utils")]
    sys.modules["ddw.utils"] = utils

    # ``unet.py`` imports this training-only helper but conversion never calls it.
    # Providing the inert name avoids importing the upstream CLI-only dependency
    # chain (typer) while preserving the exact source model implementation.
    normalization = types.ModuleType("ddw.utils.normalization")
    normalization.get_avg_model_input_mean_and_std_from_dataloader = _unavailable
    sys.modules["ddw.utils.normalization"] = normalization

    from ddw.utils.unet import LitUnet3D

    return LitUnet3D.load_from_checkpoint(CHECKPOINT, map_location="cpu").eval()


def _unavailable(*args: object, **kwargs: object) -> None:
    del args, kwargs
    raise RuntimeError("Training-only upstream normalization is unavailable here.")


def _canonical_source_to_target(
    source: torch.nn.Module,
) -> tuple[
    st.network.ClosedDescribedNetwork,
    dict[str, torch.Tensor],
    tuple[st.artifacts.LearnedCheckpointTransformation, ...],
    st.methods.DeepDeWedgeFittedInference,
]:
    """Instantiate, strict-map, and lower the source model through Network authority."""

    params = dict(source.unet_params)
    expected = {
        "chans": 64,
        "num_downsample_layers": 3,
        "drop_prob": 0.0,
    }
    if {key: params.get(key) for key in expected} != expected:
        raise RuntimeError("Checkpoint U-Net parameters are not the audited tutorial architecture.")
    if set(params) != {
        "chans",
        "num_downsample_layers",
        "drop_prob",
        "normalization_loc",
        "normalization_scale",
    }:
        raise RuntimeError("Checkpoint contains unexpected U-Net parameter fields.")

    vendor = source.unet
    fitted = st.methods.DeepDeWedgeFittedInference(
        network_affine_loc=float(vendor.normalization_loc),
        network_affine_scale=float(vendor.normalization_scale),
    )
    described = st.network.build_described_network(
        st.network.UNet3D(
            initial_channels=params["chans"],
            num_downsampling_blocks=params["num_downsample_layers"],
        ),
        context=st.network.NetworkBuildContext(
            device=torch.device("cpu"), dtype=torch.float32, seed=0
        ),
    )
    source_state = vendor.state_dict()
    target_template = described.module.state_dict()
    affine_names = {"_normalization_loc", "_normalization_scale"}
    if set(source_state) - affine_names != {
        "bottleneck.2.bias" if name == "bottleneck.4.bias" else
        "bottleneck.2.weight" if name == "bottleneck.4.weight" else name
        for name in target_template
        if name not in affine_names
    }:
        raise RuntimeError("Source and target state namespaces are not the audited mapping.")

    mapped_target: dict[str, torch.Tensor] = {}
    target_to_source: dict[str, str] = {}
    for target_name in target_template:
        if target_name in affine_names:
            mapped_target[target_name] = source_state[target_name]
            continue
        source_name = (
            target_name.replace("bottleneck.4.", "bottleneck.2.")
            if target_name.startswith("bottleneck.4.")
            else target_name
        )
        tensor = source_state[source_name]
        if tuple(tensor.shape) != tuple(target_template[target_name].shape):
            raise RuntimeError(f"Mapped tensor shape differs for {target_name!r}.")
        mapped_target[target_name] = tensor
        target_to_source[target_name] = source_name
    incompatible = described.module.load_state_dict(mapped_target, strict=True)
    if incompatible.missing_keys or incompatible.unexpected_keys:
        raise RuntimeError("Strict mapped source state load failed.")

    closed = st.network.close_described_network(described)
    canonical = closed.state
    canonical_to_target: dict[str, str] = {}
    for canonical_name, tensor in canonical.items():
        matches = [
            target_name
            for target_name, target_tensor in described.module.state_dict().items()
            if target_name not in affine_names
            and target_tensor.data_ptr() == tensor.data_ptr()
            and tuple(target_tensor.shape) == tuple(tensor.shape)
            and target_tensor.dtype == tensor.dtype
        ]
        if len(matches) != 1:
            raise RuntimeError(f"Canonical state mapping is ambiguous for {canonical_name!r}.")
        canonical_to_target[canonical_name] = matches[0]
    if set(canonical_to_target.values()) != set(target_to_source):
        raise RuntimeError("Canonical state closure does not cover the source mapping.")

    transformations = tuple(
        st.artifacts.LearnedCheckpointTransformation(
            kind="identity" if target_to_source[target_name] == target_name else "rename",
            source=f"state_dict.unet.{target_to_source[target_name]}",
            target=canonical_name,
            details={"source_checkpoint": "tutorial_data/fitted_model.ckpt"},
        )
        for canonical_name, target_name in sorted(canonical_to_target.items())
    )
    return closed, canonical, transformations, fitted


def _profile(
    fitted: st.methods.DeepDeWedgeFittedInference,
) -> st.methods.DeepDeWedgeInferenceProfile:
    return st.methods.DeepDeWedgeInferenceProfile(
        contract=st.methods.DeepDeWedgeInferenceContract(
            missing_wedge_full_width_deg=50.0,
            full_tomogram_standardization=False,
            preconditioning_normalization_policy="recompute_patch_statistics",
            patch_shape=(96, 96, 96),
            overlap=(32, 32, 32),
        ),
        fitted=fitted,
    )


def _parity(
    source: torch.nn.Module,
    closed: st.network.ClosedDescribedNetwork,
    fitted: st.methods.DeepDeWedgeFittedInference,
) -> dict[str, float]:
    """Compare external-affine canonical Network inference on a non-symmetric input."""

    realized = st.network.realize_network(
        program=closed.program,
        state=closed.state,
        context=st.network.NetworkBuildContext(
            device=torch.device("cpu"), dtype=torch.float32, seed=19
        ),
    )
    value = torch.arange(1 * 1 * 16 * 16 * 16, dtype=torch.float32).reshape(
        1, 1, 16, 16, 16
    )
    value = value / 997.0 - 0.37
    with torch.no_grad():
        vendor_output = source.unet(value)
        format2_output = invoke_deepdewedge_network(
            realized.module, value, fitted=fitted
        )
    torch.testing.assert_close(vendor_output, format2_output, rtol=1.0e-5, atol=1.0e-6)
    difference = (vendor_output - format2_output).abs()
    relative_l2 = torch.linalg.vector_norm(difference) / torch.linalg.vector_norm(vendor_output)
    return {
        "input_elements": float(value.numel()),
        "max_abs_error": float(difference.max()),
        "relative_l2_error": float(relative_l2),
    }


def _resources(parity: dict[str, float]) -> dict[str, bytes]:
    card = f"""---
license: cc-by-4.0
library_name: scitomo
tags: [cryo-electron-tomography, deepdewedge, safetensors, scitomo, format-2]
---

# DeepDeWedge tutorial checkpoint — fresh Scitomo FORMAT 2 package

This is a fresh FORMAT 2 export from the authoritative original Lightning
checkpoint, not a migration of any earlier Hugging Face package. Normal runtime
uses Scitomo's generic FORMAT 2 loader and Safetensors only; it does not require
PyTorch Lightning or the upstream DeepDeWedge source checkout.

## Package identity

- package id: `deepdewedge_tutorial`; package revision: `3`
- learned-checkpoint format: `2`; manifest schema: `4`
- Scitomo conversion checkout: `2832957f69daff0d7baec5df17a7c54954623eed`
- minimum Scitomo version: `0.7.3`
- previous Hugging Face commit: `{PREVIOUS_HF_COMMIT}` — **HISTORICAL ONLY; NOT CONVERSION INPUT**

## Authoritative provenance

- upstream repository: <https://github.com/MLI-lab/DeepDeWedge>
- upstream revision: `{UPSTREAM_REVISION}`
- Figshare DOI: <https://doi.org/10.6084/m9.figshare.25043435.v1>; file id: `45582309`
- original archive SHA-256: `{ARCHIVE_SHA256}`
- original checkpoint member: `tutorial_data/fitted_model.ckpt`
- original checkpoint size: `{CHECKPOINT_SIZE}` bytes
- original checkpoint SHA-256: `{CHECKPOINT_SHA256}`

DeepDeWedge Tutorial Data is attributed to Simon Wiedemann and is distributed
under CC BY 4.0. The pinned DeepDeWedge implementation is BSD-2-Clause; its
license text is included below `LICENSES/`. See `ATTRIBUTION.md`.

## Scientific inference semantics

The pure persisted Network owns only the lowered U-Net architecture and its 54
canonical tensors. The fitted affine values remain outside Network state in the
typed `deepdewedge_inference` profile:

- `network_affine_loc`: `{_fmt(fitted_loc := -0.1489875167608261)}`
- `network_affine_scale`: `{_fmt(fitted_scale := 1.3237642049789429)}`
- input layout: `(..., Z, Y, X)`; Network layout: `(..., C, Z, Y, X)`
- paired halves are refined independently then averaged; full-width missing wedge: 50 degrees
- 96³ patches, 32³ overlap, trailing-reflection coverage, linear-ramp reassembly
- preconditioning recomputes patch statistics; output uses the checkpoint-fitted affine

## Fresh conversion and validation

`refresh_format2.py` is the exact one-off implementation and records
the verified source, explicit 54-tensor mapping, strict Network lowering, and
generic export. It was run with Python `{platform.python_version()}`, Torch
`{torch.__version__}`, Lightning `{pytorch_lightning.__version__}`, Safetensors
`{safetensors.__version__}`, and Scitomo `{st.__version__}` on `{platform.platform()}`.

The generic exporter freshly serializes `weights.safetensors`; no previous
Hugging Face Safetensors, manifest, construction, or inference record is read.
The conversion record lists every source checkpoint tensor to canonical target
mapping. The validation record binds package state closure, generic loader
reload, external-affine semantics, and deterministic forward parity.

For a deterministic directional, non-symmetric CPU float32 input of 4,096 elements,
authoritative upstream output versus FORMAT 2 pure-Network-plus-profile output
passed `rtol=1e-5`, `atol=1e-6`: maximum absolute error
`{parity['max_abs_error']:.9g}`, relative L2 error `{parity['relative_l2_error']:.9g}`.

## Files and closure

`manifest.json` is the authoritative, closed inventory of every package file,
with each fresh size and SHA-256. It declares only FORMAT 2 construction,
inference, Safetensors, conversion, validation, and documentation/license
resources; there is no format-1 or migration artifact. Validate and load with:

```python
import scitomo as st
loaded = st.api.load_learned_network("/path/to/package")
```

This operation uses the generic Scitomo FORMAT 2 loader and does not import
Lightning or DeepDeWedge. It is a checkpoint package, not a claim of scientific
approval for a new dataset or acquisition protocol.
"""
    attribution = f"""# Attribution and modification notice

## Original material

**DeepDeWedge Tutorial Data**
Creator: Simon Wiedemann
DOI: <https://doi.org/10.6084/m9.figshare.25043435.v1>
Figshare file id: `45582309`
Archive member: `tutorial_data/fitted_model.ckpt`
License: Creative Commons Attribution 4.0 International

The method is described by Simon Wiedemann and Reinhard Heckel, *A deep
learning method for simultaneous denoising and missing wedge reconstruction in
cryogenic electron tomography*, Nature Communications 15, 8255 (2024),
<https://doi.org/10.1038/s41467-024-51438-y>.

Pinned upstream code: <https://github.com/MLI-lab/DeepDeWedge/tree/{UPSTREAM_REVISION}>
(BSD-2-Clause).

## Changes in this package

On 2026-09-04 Scitomo freshly converted only the authoritative checkpoint
`official/fitted_model.ckpt`, after byte-size and SHA-256 verification, through
the exact pinned upstream source and current generic FORMAT 2 exporter. The
54 U-Net state tensors were explicitly mapped into canonical Network state.
The two fitted affine quantities were preserved as external DeepDeWedge
inference-profile state; they are not Network state. No old Hugging Face
Safetensors or format-1 package artifact was conversion input.

No endorsement by the cited authors, the Machine Learning and Information
Processing Laboratory, Figshare, or the rights holders is implied.
"""
    return {
        "README.md": card.encode("utf-8"),
        "ATTRIBUTION.md": attribution.encode("utf-8"),
        "LICENSES/DeepDeWedge-Code-BSD-2-Clause.txt": (UPSTREAM / "LICENSE").read_bytes(),
        "refresh_format2.py": Path(__file__).read_bytes(),
    }


def _fmt(value: float) -> str:
    return format(value, ".17g")


def _replace_hf_with_closed_package() -> None:
    if TARGET.resolve() != ROOT / "hf" or not (TARGET / ".git").is_dir():
        raise RuntimeError("Refusing to replace an unexpected Hugging Face working tree.")
    if not OUTPUT.is_dir() or OUTPUT.is_symlink():
        raise RuntimeError("Fresh output directory is unavailable for publication.")
    for child in TARGET.iterdir():
        if child.name == ".git":
            continue
        if child.is_dir() and not child.is_symlink():
            shutil.rmtree(child)
        else:
            child.unlink()
    shutil.move(str(OUTPUT), str(TARGET / ".package-fresh"))
    staged = TARGET / ".package-fresh"
    for child in staged.iterdir():
        shutil.move(str(child), str(TARGET / child.name))
    staged.rmdir()


def _runtime_view() -> Path:
    """Create a byte-identical package view excluding local Git administration."""

    if RUNTIME_VIEW.exists() or RUNTIME_VIEW.is_symlink():
        raise RuntimeError("Runtime validation view already exists.")
    shutil.copytree(TARGET, RUNTIME_VIEW, ignore=shutil.ignore_patterns(".git"))
    return RUNTIME_VIEW


def main() -> None:
    _verify_authority()
    if OUTPUT.exists() or OUTPUT.is_symlink():
        raise RuntimeError("Fresh output directory already exists before export.")
    source = _load_upstream_model()
    closed, canonical, mappings, fitted = _canonical_source_to_target(source)
    parity = _parity(source, closed, fitted)
    profile = _profile(fitted)
    owner = st.artifacts.LearnedCheckpointMethodOwner(
        family="restoration", method_kind="deepdewedge"
    )
    construction = st.artifacts.LearnedCheckpointConstructionRecordV2.from_program(
        closed.program
    )
    inference = st.artifacts.LearnedCheckpointInferenceRecordV3.from_profile(
        owner=owner, profile=profile
    )
    conversion = st.artifacts.LearnedCheckpointConversionEvidenceV2(
        source=st.artifacts.LearnedCheckpointConversionSourceV2(
            kind="figshare_checkpoint",
            project="DeepDeWedge Tutorial Data",
            identifier="45582309/tutorial_data/fitted_model.ckpt",
            url="https://doi.org/10.6084/m9.figshare.25043435.v1",
            revision=UPSTREAM_REVISION,
            sha256=CHECKPOINT_SHA256,
            metadata={
                "archive_sha256": ARCHIVE_SHA256,
                "archive_member": "tutorial_data/fitted_model.ckpt",
                "checkpoint_size_bytes": CHECKPOINT_SIZE,
                "upstream_repository": "https://github.com/MLI-lab/DeepDeWedge",
            },
        ),
        tool="deepdewedge_refresh_format2",
        tool_version="1",
        tensor_mappings=mappings,
        environment={
            "python": platform.python_version(),
            "torch": torch.__version__,
            "pytorch_lightning": pytorch_lightning.__version__,
            "safetensors": safetensors.__version__,
            "scitomo": st.__version__,
            "scitomo_commit": "2832957f69daff0d7baec5df17a7c54954623eed",
            "platform": platform.platform(),
        },
    )
    validation = st.artifacts.LearnedCheckpointValidationEvidenceV2(
        software=(
            st.artifacts.LearnedCheckpointSoftware(name="scitomo", version=st.__version__),
            st.artifacts.LearnedCheckpointSoftware(name="torch", version=torch.__version__),
            st.artifacts.LearnedCheckpointSoftware(name="pytorch_lightning", version=pytorch_lightning.__version__),
            st.artifacts.LearnedCheckpointSoftware(name="safetensors", version=safetensors.__version__),
        ),
        cases=(
            st.artifacts.LearnedCheckpointValidationCase(
                name="authoritative_source_mapping", kind="state_mapping", status="passed",
                metrics={"source_tensors": 56.0, "canonical_network_tensors": 54.0},
            ),
            st.artifacts.LearnedCheckpointValidationCase(
                name="external_fitted_affine_profile", kind="inference_profile", status="passed",
                metrics={"network_affine_loc": fitted.network_affine_loc, "network_affine_scale": fitted.network_affine_scale},
            ),
            st.artifacts.LearnedCheckpointValidationCase(
                name="deterministic_forward_parity", kind="forward_parity", status="passed",
                tolerances={"atol": 1.0e-6, "rtol": 1.0e-5}, metrics=parity,
            ),
        ),
    )
    exported = st.artifacts.export_learned_checkpoint_format2_package(
        canonical,
        construction=construction,
        inference=inference,
        package_id="deepdewedge_tutorial",
        package_revision=3,
        owner=owner,
        requirements=st.artifacts.LearnedCheckpointFormat2Requirements(
            minimum_scitomo_version="0.7.3"
        ),
        validation=validation,
        conversion=conversion,
        resources=_resources(parity),
        destination="package-fresh",
        write_scope=st.core.WriteScope(st.core.WriteScopeKind.MODELS, ROOT),
        provenance=st.artifacts.LearnedCheckpointProvenance(
            sources=(
                st.artifacts.LearnedCheckpointSource(
                    kind="upstream_repository", project="MLI-lab/DeepDeWedge",
                    identifier=UPSTREAM_REVISION,
                    url="https://github.com/MLI-lab/DeepDeWedge",
                    revision=UPSTREAM_REVISION,
                ),
            ),
            citations=(
                "https://doi.org/10.1038/s41467-024-51438-y",
                "https://doi.org/10.6084/m9.figshare.25043435.v1",
            ),
        ),
    )
    _replace_hf_with_closed_package()
    runtime_root = _runtime_view()
    try:
        package = st.artifacts.validate_learned_checkpoint_format2_package(runtime_root)
        loaded = st.api.load_learned_network(
            runtime_root,
            context=st.network.NetworkBuildContext(
                device=torch.device("cpu"), dtype=torch.float32, seed=29
            ),
            expected_owner=owner,
            expected_inference_profile=profile,
        )
        reloaded = st.network.canonical_network_state(loaded.network.module)
        if set(reloaded) != set(canonical) or any(
            not torch.equal(reloaded[name], canonical[name]) for name in canonical
        ):
            raise RuntimeError("Reloaded FORMAT 2 state is not closed over canonical state.")
        value = torch.arange(1 * 1 * 16 * 16 * 16, dtype=torch.float32).reshape(1, 1, 16, 16, 16)
        value = value / 997.0 - 0.37
        with torch.no_grad():
            expected = source.unet(value)
            actual = invoke_deepdewedge_network(
                loaded.network.module, value, fitted=profile.fitted
            )
        torch.testing.assert_close(expected, actual, rtol=1.0e-5, atol=1.0e-6)
        print(f"FORMAT 2 package validated at {TARGET}")
        print(f"weights_sha256={package.manifest.files.weights.sha256}")
        print(f"manifest_sha256={_digest(TARGET / 'manifest.json', 'sha256')}")
    finally:
        if RUNTIME_VIEW.exists():
            shutil.rmtree(RUNTIME_VIEW)


if __name__ == "__main__":
    main()