release: AtomWorks 3.0.1 with SASA fixes and HF reads

.github/repair_docs_switcher.py

@@ -0,0 +1,37 @@
+"""Point archived PyData theme pages at the deployed site's version manifest."""
+
+import argparse
+import json
+import re
+from pathlib import Path
+from urllib.parse import urlsplit
+
+SWITCHER_URL = re.compile(r"(DOCUMENTATION_OPTIONS\.theme_switcher_json_url\s*=\s*)([\"'])(.*?)(\2)")
+
+
+def repair_switcher_urls(pages: Path, site_url: str) -> int:
+    """Update only switcher URL assignments, preserving archived content and version labels."""
+    parsed = urlsplit(site_url)
+    if parsed.scheme not in {"http", "https"} or not parsed.netloc or parsed.query or parsed.fragment:
+        raise ValueError(f"Expected an HTTP(S) documentation site URL, got {site_url!r}")
+    manifest_url = site_url.rstrip("/") + "/latest/_static/switcher.json"
+    manifest = pages / "latest" / "_static" / "switcher.json"
+    entries = json.loads(manifest.read_text())
+    if not entries:
+        raise ValueError(f"Empty version manifest: {manifest}")
+    changed = 0
+    for page in pages.rglob("*.html"):
+        original = page.read_text()
+        updated = SWITCHER_URL.sub(lambda match: match[1] + json.dumps(manifest_url), original)
+        if updated != original:
+            page.write_text(updated)
+            changed += 1
+    return changed
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("pages", type=Path)
+    parser.add_argument("site_url")
+    args = parser.parse_args()
+    print(f"Updated version switcher URLs in {repair_switcher_urls(args.pages, args.site_url)} pages")

.github/tests/test_docs_switcher.py

@@ -0,0 +1,55 @@
+"""Regression checks for repairing archived documentation navigation."""
+
+import importlib.util
+import json
+import tempfile
+import unittest
+from pathlib import Path
+
+spec = importlib.util.spec_from_file_location(
+    "repair_docs_switcher", Path(__file__).parents[1] / "repair_docs_switcher.py"
+)
+repair = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(repair)
+
+
+class SwitcherRepairTests(unittest.TestCase):
+    def test_repairs_nested_archives_without_changing_content_or_version(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            pages = Path(directory)
+            manifest = pages / "latest/_static/switcher.json"
+            manifest.parent.mkdir(parents=True)
+            manifest.write_text(json.dumps([{"version": "2.2.1"}, {"version": "3.0.0"}]))
+            original = (
+                "<h1>Archived content</h1><script>"
+                "DOCUMENTATION_OPTIONS.theme_switcher_json_url = "
+                "'https://baker-laboratory.github.io/atomworks-dev/latest/_static/switcher.json';"
+                "DOCUMENTATION_OPTIONS.theme_switcher_version_match = '2.2.1';</script>"
+            )
+            expected = original.replace(
+                "'https://baker-laboratory.github.io/atomworks-dev/latest/_static/switcher.json'",
+                '"https://rosettacommons.github.io/atomworks/latest/_static/switcher.json"',
+            )
+            for name in ("v2.2.1/index.html", "v2.2.1/io/parser.html"):
+                page = pages / name
+                page.parent.mkdir(parents=True, exist_ok=True)
+                page.write_text(original)
+            untouched = pages / "index.html"
+            untouched.write_text('<meta http-equiv="refresh" content="0;url=latest/">')
+            self.assertEqual(repair.repair_switcher_urls(pages, "https://rosettacommons.github.io/atomworks/"), 2)
+            for page in (pages / "v2.2.1").rglob("*.html"):
+                self.assertEqual(page.read_text(), expected)
+            self.assertEqual(repair.repair_switcher_urls(pages, "https://rosettacommons.github.io/atomworks"), 0)
+            self.assertEqual(untouched.read_text(), '<meta http-equiv="refresh" content="0;url=latest/">')
+
+    def test_missing_manifest_prevents_archive_changes(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            page = Path(directory) / "index.html"
+            page.write_text("Archived content")
+            with self.assertRaises(FileNotFoundError):
+                repair.repair_switcher_urls(Path(directory), "https://example.com/docs")
+            self.assertEqual(page.read_text(), "Archived content")
+
+
+if __name__ == "__main__":
+    unittest.main()

.github/workflows/release_and_docs.yaml

@@ -2,7 +2,7 @@ name: CI and release
 
 on:
   push:
-    branches: [dev, main, production, release/atomworks-3-0]
+    branches: [production]
     tags: ['v*']
   pull_request:
     types: [opened, synchronize, reopened, ready_for_review]
@@ -189,6 +189,8 @@ jobs:
           ref: gh-pages
           path: gh-pages
       - run: python .github/release.py docs
+      - name: Repair archived version switchers
+        run: python .github/repair_docs_switcher.py gh-pages https://rosettacommons.github.io/atomworks
       - name: Publish versioned documentation
         uses: peaceiris/actions-gh-pages@v4
         with:

README.md

@@ -1,5 +1,5 @@
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
-[![CI](https://github.com/RosettaCommons/atomworks/actions/workflows/release_and_docs.yaml/badge.svg?branch=release%2Fatomworks-3-0)](https://github.com/RosettaCommons/atomworks/actions/workflows/release_and_docs.yaml)
+[![CI](https://github.com/RosettaCommons/atomworks/actions/workflows/release_and_docs.yaml/badge.svg?branch=production)](https://github.com/RosettaCommons/atomworks/actions/workflows/release_and_docs.yaml)
 [![PyPI version](https://img.shields.io/pypi/v/atomworks.svg)](https://pypi.org/project/atomworks/)
 [![Python versions](https://img.shields.io/pypi/pyversions/atomworks.svg)](https://pypi.org/project/atomworks/)
 [![Documentation Status](https://img.shields.io/badge/docs-latest-brightgreen.svg)](https://rosettacommons.github.io/atomworks/latest/)
@@ -14,11 +14,11 @@
 
 **atomworks** is an open-source platform that maximizes research velocity for biomolecular modeling tasks. Much like how [Torchvision](https://docs.pytorch.org/vision/stable/index.html) enables rapid prototyping within the vision domain, and [Torchaudio](https://docs.pytorch.org/audio/main/) within the audio domain, AtomWorks aims to accelerate development and experimentation within biomolecular modeling.
 
-Upgrading to 3.0? See the [migration guide](https://github.com/RosettaCommons/atomworks/blob/release/atomworks-3-0/docs/migration/migration-short.md).
+Upgrading to 3.0? See the [migration guide](https://github.com/RosettaCommons/atomworks/blob/production/docs/migration/migration-short.md).
 
 If you're looking for the models themselves (e.g., RF3, RFD3, MPNN) that integrate with AtomWorks rather than the underlying framework, check out [Foundry](https://github.com/RosettaCommons/foundry)
 
-> **💡 Note:** Start with the [AtomWorks examples](https://rosettacommons.github.io/atomworks/latest/auto_examples/index.html) or the [model-building tutorial](https://github.com/RosettaCommons/atomworks/blob/release/atomworks-3-0/docs/how_to_build_a_model/index.rst).
+> **💡 Note:** Start with the [AtomWorks examples](https://rosettacommons.github.io/atomworks/latest/auto_examples/index.html) or the [model-building tutorial](https://github.com/RosettaCommons/atomworks/blob/production/docs/how_to_build_a_model/index.rst).
 
 AtomWorks is composed of two symbiotic libraries:
 
@@ -93,12 +93,6 @@ If you are using [uv](https://docs.astral.sh/uv/reference/policies/versioning/)
 uv pip install "atomworks[ml,openbabel,dev]"
 ```
 
-To install the AtomWorks 3.0 release branch before it is published on PyPI:
-
-```shell
-pip install "atomworks[ml] @ git+https://github.com/RosettaCommons/atomworks.git@release/atomworks-3-0"
-```
-
 For more advanced setup options (including how to run workflows via apptainers) see the [full documentation](https://rosettacommons.github.io/atomworks/latest/index.html).
 
 ---

docs/conf.py

@@ -55,6 +55,7 @@
     "globaltoc_maxdepth": -1,  # Unlimited depth
     "header_links_before_dropdown": 8,
     "navbar_end": ["theme-switcher", "navbar-icon-links"],
+    "github_url": "https://github.com/" + os.environ.get("GITHUB_REPOSITORY", "RosettaCommons/atomworks"),
     "logo": {
         "image_light": "_static/atomworks_logo_light.svg",
         "image_dark": "_static/atomworks_logo_dark.svg",

docs/contributor_guide.rst

@@ -37,7 +37,7 @@ As you commit
 As you finalize a PR
 ---------------------
 
-1. Target **release/atomworks-3-0** for the AtomWorks 3.0 release.
+1. Target **production** for public AtomWorks changes.
 2. Keep overall PR under <400 LOC (lines of code) (Rule of thumb: 500 LOC takes about 1h to review).
 3. Read and fill out the `PR checklist <https://github.com/RosettaCommons/atomworks/blob/production/.github/pull_request_template.md>`_.
 
@@ -90,7 +90,7 @@ When contributing to this repository, please follow these steps:
    - Run autoformatting tools (``make format``) before committing.
    - Use commit messages like ``<type>: <description>``. Types: ``feat``, ``fix``, ``refactor``, ``docs``, ``chore``, ``wip``.
    - Example: ``git commit -m "docs: add contributing guidelines"``
-5. Open a pull request to ``release/atomworks-3-0`` and describe your changes.
+5. Open a pull request to ``production`` and describe your changes.
 6. Wait for review and merge your changes.
 
 
@@ -108,20 +108,20 @@ One-time repository setup:
 - Keep GitHub Pages Source as **Deploy from a branch**, using ``gh-pages`` and
   ``/ (root)``. The workflow preserves versioned pages on that branch.
 
-For 3.0.0, merge the reviewed PRs, wait for CI on the final release commit, and
-confirm that ``project.version`` is ``3.0.0`` and the tag/PyPI version are unused.
+For each release, merge the reviewed PRs, wait for CI on the final production commit,
+and confirm that ``project.version`` matches an unused tag and PyPI version.
 From a checkout whose ``origin`` is the public repository:
 
 .. code-block:: bash
 
    git fetch origin
-   git tag -a v3.0.0 origin/release/atomworks-3-0 -m "AtomWorks 3.0.0"
-   git push origin refs/tags/v3.0.0
+   git tag -a v3.0.1 origin/production -m "AtomWorks 3.0.1"
+   git push origin refs/tags/v3.0.1
 
 The tag runs package checks, installed-wheel checks, scientific CPU tests
 on GitHub-hosted runners and the docs build. After they pass, approve ``pypi`` to publish the checked
 archives; subsequent jobs create the GitHub Release and deploy documentation.
-Verify a fresh ``pip install atomworks==3.0.0``, ``aw --help`` and the versioned docs.
+Verify a fresh ``pip install atomworks==3.0.1``, ``aw --help`` and the versioned docs.
 
 Prereleases follow the same process without replacing stable ``latest`` docs.
 Never move a published tag or reuse a published package version.

docs/huggingface.rst

@@ -0,0 +1,22 @@
+Hugging Face dataset repositories
+=================================
+
+Install ``atomworks[ml,hf]`` to read an indexed blob dataset from a private or
+public HF dataset repository. Authenticate with ``hf auth login`` or ``HF_TOKEN``
+for private data. Pin every URL to the same full commit SHA:
+
+.. code-block:: python
+
+   from atomworks.ml.utils.blob_store import BlobIndex, BlobStore
+   from atomworks.ml.utils.io import read_parquet_with_metadata
+
+   root = "hf://datasets/org/repo@<40-character-sha>/protein"
+   examples = read_parquet_with_metadata(f"{root}/examples.parquet")
+   index = BlobIndex(f"{root}/blob/index.parquet", id_column="path")
+   store = BlobStore(f"{root}/blob/data")
+   cif = store.get_bytes(*index.lookup(examples.iloc[0]["path"]))
+
+Set ``id_column`` to the identifier used by your index. Sampling tables and
+indices use the HF file cache; blob records use exact ranged reads and are not
+cached. For repeated training, download the required shards once and use local
+paths. The existing S3 and local readers are unchanged.

docs/index.rst

@@ -21,10 +21,10 @@ Welcome to **atomworks** — a toolkit for converting, parsing, and manipulating
    how_to_build_a_model/index
    cookbook
    installation
-   migration/index
    glossary
    api_reference
    auto_examples/index
    contributor_guide
    mirrors
+   huggingface
    msa

docs/installation.rst

@@ -42,7 +42,7 @@ For development, use a virtual environment:
 
 .. code-block:: bash
 
-   git clone --branch release/atomworks-3-0 https://github.com/RosettaCommons/atomworks.git
+   git clone --branch production https://github.com/RosettaCommons/atomworks.git
    cd atomworks
    python -m venv .venv
    source .venv/bin/activate
@@ -74,3 +74,9 @@ For production use or training on the full PDB, you'll want complete mirrors rat
 * Setting up a full PDB mirror (~100 GB)
 * Setting up a CCD mirror (~2 GB)
 * Configuring environment variables for production use
+
+
+.. toctree::
+   :maxdepth: 1
+
+   migration/index

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "atomworks"
-version = "3.0.0"
+version = "3.0.1"
 description = "A research-oriented data toolkit for training biomolecular deep-learning foundation models"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -55,6 +55,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+hf = ["huggingface_hub>=1.0,<2"]
 ml = [
     # atomworks-ml dependencies
     "torch>=2.2.0",

src/atomworks/experimental/protonation/_placement.py

@@ -818,8 +818,14 @@ def place_hydrogens(atom_array: AtomArray) -> AtomArray:
     protonated.coord = np.concatenate([coord, h_coord[built]])[order]
     ensure_annotations(protonated, *recompute)
     for annot in recompute:
-        values = protonated.get_annotation(annot)
+        standard = next((a for a in STANDARD_ANNOTATIONS if annot in (a.name, a.full_name, *a.aliases)), None)
+        if standard is None:
+            values = protonated.get_annotation(annot)
+        else:
+            values = np.asarray(standard.annotation(protonated, default="raise")).copy()
         values[from_state] = atom_array.get_annotation(annot)[order[from_state]]
+        if standard is not None:
+            standard.set_annotation(protonated, values)
         protonated.set_annotation(annot, values)
     return protonated
 

src/atomworks/experimental/protonation/external/dimorphite_dl/dimorphite_dl.py

@@ -44,6 +44,9 @@
 
 _SOURCE_INDEX = "_atomworks_protonation_source_index"
 
+_REMOVE_HS_PARAMS = Chem.RemoveHsParameters()
+_REMOVE_HS_PARAMS.removeNontetrahedralNeighbors = True
+
 #: An aromatic nitrogen of formal charge -1.
 _AROMATIC_NITRANION = Chem.MolFromSmarts("[n-]")
 
@@ -256,7 +259,7 @@ def protonate_site(mols: list[Chem.Mol], site: tuple[int, str, str], *, remove_h
             nitrogen_charge = charge + int("*" not in prot_site_name)
             for mol in mols:
                 try:
-                    mol_copy = Chem.RemoveHs(mol) if remove_hs else Chem.Mol(mol)
+                    mol_copy = Chem.RemoveHs(mol, _REMOVE_HS_PARAMS) if remove_hs else Chem.Mol(mol)
                 except Exception:
                     continue
                 atom = mol_copy.GetAtomWithIdx(idx)
@@ -309,7 +312,7 @@ def protonate_mol_variants(
     if prepared is None:
         return []
     try:
-        prepared = Chem.RemoveHs(prepared)
+        prepared = Chem.RemoveHs(prepared, _REMOVE_HS_PARAMS)
     except Exception:
         return []
 
@@ -349,7 +352,7 @@ def protonate_mol_variants(
         if len(properly_formed_mols) == 1:
             new_mols = properly_formed_mols
     else:
-        mol_used_to_idx_sites = Chem.RemoveHs(mol_used_to_idx_sites)
+        mol_used_to_idx_sites = Chem.RemoveHs(mol_used_to_idx_sites, _REMOVE_HS_PARAMS)
         new_mols = [mol_used_to_idx_sites]
         properly_formed_mols.append(mol_used_to_idx_sites)