Proposal: Protein 3D Structure Visualization for Dataset Viewer

# Proposal: Protein 3D Structure Visualization for HuggingFace Dataset Viewer

## Executive Summary

This proposal outlines adding 3D protein structure visualization to the HuggingFace Dataset Viewer, enabling users to interactively view PDB and mmCIF molecular structures directly within the dataset preview interface.

---

## Data Type Support (Updated Architecture)

**Supported formats** (from recent PRs):
- **PDB** (PR #7926): `.pdb`, `.ent` extensions via `PdbFolder` builder
- **mmCIF** (PR #7925): `.cif`, `.mmcif` extensions via `MmcifFolder` builder

**New Implementation Pattern (One Row = One Structure)**:

Both PRs have been refactored to follow the **ImageFolder pattern**, where each row in the dataset contains one complete protein structure file. This is the recommended ML-friendly approach:

```python
>>> from datasets import load_dataset
>>> dataset = load_dataset("mmcif", data_dir="./structures")
>>> dataset[0]
{'structure': 'data_1ABC\n_entry.id 1ABC\n_atom_site...'}  # Complete mmCIF content

>>> from datasets import load_dataset  
>>> dataset = load_dataset("pdb", data_dir="./pdbs")
>>> dataset[0]
{'structure': 'HEADER  PROTEIN  01-JAN-20  1ABC\nATOM...'}  # Complete PDB content
```

**Key Components**:
- **ProteinStructure feature type**: New feature type supporting both PDB and mmCIF formats with lazy loading
- **PdbFolder builder** (PR #7926): Folder-based loader for PDB files with label and metadata support
- **MmcifFolder builder** (PR #7925): Folder-based loader for mmCIF files with label and metadata support

**What gets visualized**:
- 3D atomic coordinates (x, y, z)
- Chain structures
- Residue information
- Atom types and elements
- Secondary structure (helices, sheets)

**Not applicable** (1D sequence only):
- FASTA (PR #7923) - text sequences, no 3D coordinates
- FASTQ (PR #7924) - sequences with quality scores, no 3D coordinates

---

## Visualization Library Comparison

| Library | Bundle Size (minified) | Bundle Size (gzipped) | License | Pros | Cons |
|---------|------------------------|----------------------|---------|------|------|
| **3Dmol.js** | 512 KB | **~150 KB** | BSD-3 | Lightweight, easy integration, good docs | Fewer advanced features |
| **NGL Viewer** | 1.3 MB | ~350 KB | MIT | Excellent MMTF support, beautiful rendering | Moderate complexity |
| **Mol*** | 4.6 MB | ~1.3 MB | MIT | Industry standard, used by RCSB PDB, feature-rich | Heavy, complex |
| **PDBe Molstar** | 5.8 MB | ~1.6 MB | Apache 2.0 | EMBL-EBI maintained, simpler Mol* wrapper | Still very heavy |

*Bundle sizes verified by downloading actual distribution files from npm/CDN (January 2026)*

---

## Recommendation: 3Dmol.js

**Primary choice**: 3Dmol.js

**Rationale**:
1. **Bundle size**: ~150 KB gzipped - the lightest option by far, ideal for lazy loading
2. **Simple API**: Easy to integrate with React/Next.js
3. **BSD-3 License**: Compatible with HuggingFace licensing
4. **Active maintenance**: Regular updates, good community support
5. **Format support**: Native PDB and mmCIF parsing built-in
6. **Sufficient features**: Rotation, zoom, style switching (cartoon, stick, sphere)

**Why not Mol*?** As Georgia noted, Mol* is heavy (~1.3 MB gzipped). While it's the industry standard for RCSB PDB, it's overkill for a dataset preview where users just need to verify structure data looks correct.

**Alternative for power users**: If users need advanced features like density maps, ligand interactions, or sequence alignment overlay, consider PDBe Molstar as an optional "full viewer" mode.

---

## Summary

**Recommended approach**:
- Use **3Dmol.js** (~150 KB gzipped) with **lazy loading**
- Only loads when user views PDB/mmCIF datasets
- Simple integration, BSD-3 license, active community support

**Backend implementation** (Updated):
- PR #7925 (mmCIF): Uses **MmcifFolder** builder with **ProteinStructure** feature type
- PR #7926 (PDB): Uses **PdbFolder** builder with **ProteinStructure** feature type
- Both follow the **one-row-per-structure** pattern (like ImageFolder)
- Each row's `structure` column contains the complete file content ready for 3D rendering

---

## Next Steps

1. Get feedback on this proposal
2. Create proof-of-concept in a standalone demo if needed
3. Integrate into dataset-viewer once approach is approved


cc @georgia-hf - Following up on your question about protein visualization for the Dataset Viewer. This proposal recommends 3Dmol.js (~150KB gzipped) as a lightweight alternative to Mol* (~1.3MB gzipped).

Looking forward to your feedback!

Exciting ! cc @cfahlgren1 @severo for the Viewer part

For the `datasets` part I'll leave my feedbacks in the PRs :)

I don't know the JS libraries, but indeed, the lighter the better, as we don't require advanced features.

From a quick look at the PDB and mmCIF PRs I noticed that the dataset has one row = one atom. However I humbly believe that such datasets would be more practical to use if one row = one structure. This way each row is independent, which is practical in ML to perform train/test splits or dataset shuffling.

This would also make it easier to add labels and metadata for each structure, similar to what we already for images. E.g. you could group them per folder named after a label, or you can have a metadata.parquet file to add custom metadata per structure.

And this way in the Viewer it could show one 3D render per row.

What do you think ?

@lhoestq @severo @georgia-hf I will be waiting for all your comments; then, I will start implementing the final plan. 

adding some remarks from @0gust1 (feel free to add comments here!):

- NGLViewer works very well, and is used in production at [mabsilico](https://github.com/mabsilico)
- Molstar is big, but modular, so: it must be possible to use it without requiring the whole 1.3MB. Beware that the API and engine might lack documentation, so... some deep dive in the code might be needed.


Thanks all for the feedback — @lhoestq @severo @0gust1 @cfahlgren1. Synthesizing it into a concrete plan.

**Data model (settled + shipped).** Following @lhoestq's point that one row = one structure is more practical than one row = one atom, the PDB (#7926) and mmCIF (#7925) loaders now use exactly that: each row is one structure file (ImageFolder-style — folder names become labels, `metadata.csv`/`metadata.parquet` adds per-structure metadata), and the `structure` column is parsed into atom-level arrays under the standard PDBx/mmCIF `_atom_site` names. That makes train/test splits, shuffling, and per-structure labels work the way they do for images, and it gives the Viewer one renderable structure per row.

**Viewer library — trade-offs from the thread:**

| Library | Size | Notes |
|---|---|---|
| **3Dmol.js** | ~150 KB gzipped | Lightweight, simple API; my original recommendation |
| **NGL** | ~mid | "works very well, used in production at mabsilico" (@0gust1) |
| **Mol\*** | ~1.3 MB gzipped | Powerful but heavy; modular so a subset *might* be usable, though @0gust1 notes the API/engine can be under-documented and need a deep dive |

Given @severo's "the lighter the better, we don't need advanced features," my recommendation is **3Dmol.js as the default**, with **NGL as the proven fallback** if we hit a rendering gap — both are far below Mol\*'s footprint, and a slimmed Mol\* would need investigation effort that doesn't seem warranted for basic structure display.

**Proposed next step:** a minimal Viewer integration that renders the parsed `structure` column with 3Dmol.js, scoped to display-only (no editing/measurement tooling). Before I open that, one decision I'd like your steer on: should the Viewer render from the **parsed atom arrays** the loaders now produce, or fetch and parse the **raw file** client-side? The former keeps the JS thin; the latter keeps the row payload smaller. Happy to go whichever way fits the Viewer architecture best.


It seems 3Dmol.js also has an integration with IPython.display, enabling a viewer in jupyter notebooks as well. We need the other PRs to move forward though first :)