Offline Visualization of Inference Traces

[archbilesherman]

Summary

This PR implements the backend reader layer for issue #41, which focuses on lightweight offline visualization of VizFold/OpenFold inference traces. The goal is to let visualization tools inspect saved inference traces without rerunning inference or depending on a heavy live backend.

The main contribution is a working ArchiveReader for standardized Zarr-style trace archives, alongside the existing LegacyTxtReader path for older text-based trace dumps.

What the code does

LegacyTxtReader

The LegacyTxtReader supports the older VizFold text-dump format. This is useful because some existing trace outputs are still stored as plain text files rather than standardized archives. Keeping this reader lets the offline visualization stack continue to support current/legacy outputs while the archive format from issue #39 stabilizes.

ArchiveReader

The new ArchiveReader supports Zarr-based inference trace archives. It can:

• open .zarr archives
• open zipped Zarr archives by extracting them to a temporary directory
• read metadata from archive attributes / metadata groups
• discover available attention types
• list available layers
• list available heads
• selectively load one attention head
• load all attention heads for a layer
• apply top_k filtering to attention connections
• load single representations
• load pair representations
• load structure data
• convert stored coordinates into basic PDB text as a lightweight visualization fallback

This is meant to provide a stable backend interface for the frontend visualization work. The frontend should be able to call the reader interface without needing to know the exact internal archive layout.

Why both readers exist

There are currently two reader paths because they support two different stages of the project:

• LegacyTxtReader supports existing text-based outputs.
• ArchiveReader supports the newer standardized Zarr archive direction from issue Design a Standardized Archive Format for Inference Traces #39.

This lets issue #41 support existing data while also moving toward the scalable archive format that will be needed for larger protein traces.

How to run locally

python -m pip install streamlit plotly py3Dmol matplotlib pandas zarr numcodecs fsspec
python -m pytest tests/test_archive_reader_contract.py tests/test_legacy_txt_reader.py
python -m streamlit run webui/app.py

[archbilesherman]

Updated the PR description with current backend functionality, test coverage, and remaining integration work. The reader now passes local tests and supports working Zarr archive loading, but still needs validation against a real standardized archive from the archive-writing pipeline.