From 6760a9cd918910625563bbdfd6ec62e6c34ce3ba Mon Sep 17 00:00:00 2001 From: Michael Li Date: Mon, 5 Oct 2026 17:34:33 +0000 Subject: [PATCH 1/6] fix(docs): restore navigation and archived version switchers --- .github/repair_docs_switcher.py | 37 +++++++++++++++++ .github/tests/test_docs_switcher.py | 55 +++++++++++++++++++++++++ .github/workflows/release_and_docs.yaml | 2 + docs/conf.py | 1 + docs/index.rst | 1 - docs/installation.rst | 6 +++ 6 files changed, 101 insertions(+), 1 deletion(-) create mode 100644 .github/repair_docs_switcher.py create mode 100644 .github/tests/test_docs_switcher.py diff --git a/.github/repair_docs_switcher.py b/.github/repair_docs_switcher.py new file mode 100644 index 00000000..f8a7a7b0 --- /dev/null +++ b/.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") diff --git a/.github/tests/test_docs_switcher.py b/.github/tests/test_docs_switcher.py new file mode 100644 index 00000000..d0d07965 --- /dev/null +++ b/.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 = ( + "

Archived content

" + ) + 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('') + 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(), '') + + 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() diff --git a/.github/workflows/release_and_docs.yaml b/.github/workflows/release_and_docs.yaml index 812893a0..d4035b1f 100644 --- a/.github/workflows/release_and_docs.yaml +++ b/.github/workflows/release_and_docs.yaml @@ -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: diff --git a/docs/conf.py b/docs/conf.py index 73de1615..b525ba03 100644 --- a/docs/conf.py +++ b/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", diff --git a/docs/index.rst b/docs/index.rst index 81aa51c0..a22004bb 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -21,7 +21,6 @@ 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 diff --git a/docs/installation.rst b/docs/installation.rst index afdb811c..c8a4f82c 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -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 From 4706bbf5a3551ffc941e20caeb2e9de305b623d8 Mon Sep 17 00:00:00 2001 From: Rafael Brent <105883594+RafiBrent@users.noreply.github.com> Date: Mon, 5 Oct 2026 00:51:47 -0700 Subject: [PATCH 2/6] fix: remove retained hydrogens from Dimorphite products (#486) * fix: remove retained hydrogens from Dimorphite products * fix: configure existing Dimorphite hydrogen removal --- .../external/dimorphite_dl/dimorphite_dl.py | 9 +++++--- .../protonation/test_dimorphite.py | 21 +++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) create mode 100644 tests/experimental/protonation/test_dimorphite.py diff --git a/src/atomworks/experimental/protonation/external/dimorphite_dl/dimorphite_dl.py b/src/atomworks/experimental/protonation/external/dimorphite_dl/dimorphite_dl.py index 5dc9c7d3..7e5ffc73 100644 --- a/src/atomworks/experimental/protonation/external/dimorphite_dl/dimorphite_dl.py +++ b/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) diff --git a/tests/experimental/protonation/test_dimorphite.py b/tests/experimental/protonation/test_dimorphite.py new file mode 100644 index 00000000..d150aa96 --- /dev/null +++ b/tests/experimental/protonation/test_dimorphite.py @@ -0,0 +1,21 @@ +"""Regression tests for Dimorphite molecule outputs.""" + +import numpy as np +from rdkit import Chem +from rdkit.Chem import AllChem + +from atomworks.experimental.protonation.dimorphite import protonate_at_ph + + +def test_protonation_removes_hydrogens_on_non_tetrahedral_centers() -> None: + mol = Chem.MolFromSmiles("Cc1cn([C@H]2C[C@H](O)[C@@H](CO[P@SP2H](=O)O)O2)c(=O)[nH]c1=O") + for atom in mol.GetAtoms(): + atom.SetAtomMapNum(atom.GetIdx() + 1) + AllChem.Compute2DCoords(mol) + + product = protonate_at_ph(mol, 7.4) + + assert product is not None + assert [atom.GetAtomicNum() for atom in product.GetAtoms()] == [atom.GetAtomicNum() for atom in mol.GetAtoms()] + assert [atom.GetAtomMapNum() for atom in product.GetAtoms()] == list(range(1, mol.GetNumAtoms() + 1)) + np.testing.assert_array_equal(product.GetConformer().GetPositions(), mol.GetConformer().GetPositions()) From 9910dbf7d583d92080165c9842841acb47678d34 Mon Sep 17 00:00:00 2001 From: Jasper Butcher <66851659+Ubiquinone-dot@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:05:42 -0700 Subject: [PATCH 3/6] fix(protonation): preserve standard annotation aliases (#490) Co-authored-by: jiechen7 --- .../experimental/protonation/_placement.py | 8 +++++++- tests/experimental/protonation/test_regressions.py | 13 +++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/src/atomworks/experimental/protonation/_placement.py b/src/atomworks/experimental/protonation/_placement.py index 3c327c4b..274815fa 100644 --- a/src/atomworks/experimental/protonation/_placement.py +++ b/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 diff --git a/tests/experimental/protonation/test_regressions.py b/tests/experimental/protonation/test_regressions.py index 26173c24..398c9d31 100644 --- a/tests/experimental/protonation/test_regressions.py +++ b/tests/experimental/protonation/test_regressions.py @@ -347,3 +347,16 @@ def test_hydrogen_placement_preserves_pair_annotations_across_residue_reordering np.testing.assert_array_equal(annotation.values, [1.7]) assert annotation.n_atoms == len(placed) np.testing.assert_array_equal(atoms.get_annotation("restraint", n_body=2).pairs, [[0, 2]]) + + +def test_hydrogen_placement_preserves_standard_annotation_aliases(): + from atomworks.io.utils.atom_array_plus import as_atom_array_plus + from tests.experimental.protonation.test_assign import molecule_from_smiles + + atoms = as_atom_array_plus(molecule_from_smiles("C")) + atoms.set_annotation("S_ATM", np.ones(len(atoms), dtype=bool)) + + placed = place_hydrogens(assign_hydrogens(atoms)) + + assert "S_ATM" in placed.get_annotation_categories() + np.testing.assert_array_equal(placed.S_ATM[placed.element != "H"], atoms.S_ATM) From 269ff45752487a390f05f5c539898874db3e8249 Mon Sep 17 00:00:00 2001 From: Rafael Brent <105883594+RafiBrent@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:27:00 -0700 Subject: [PATCH 4/6] fix: separate SASA scoring and occlusion masks (#488) BREAKING CHANGE: replace atom_filter with occlusion_mask for environment filtering; use sasa_mask to select scored atoms without removing their occluders. --- src/atomworks/ml/transforms/sasa.py | 104 ++++++++++++++++------------ tests/ml/transforms/test_sasa.py | 18 +++++ 2 files changed, 76 insertions(+), 46 deletions(-) diff --git a/src/atomworks/ml/transforms/sasa.py b/src/atomworks/ml/transforms/sasa.py index 1adf8a89..f518d262 100644 --- a/src/atomworks/ml/transforms/sasa.py +++ b/src/atomworks/ml/transforms/sasa.py @@ -1,4 +1,3 @@ -import logging from typing import Any, Literal import biotite.structure as struc @@ -13,8 +12,6 @@ ) from atomworks.ml.transforms.base import Transform -logger = logging.getLogger("atomworks.ml") - # User-extension dict for VdW radii of elements not known to biotite. # Checked before biotite's vdw_radius_single in _get_element_radii and # _get_protor_element_fallback_radii. Populate at runtime to override or @@ -104,11 +101,12 @@ def calculate_atomwise_sasa( probe_radius: float = 1.4, atom_radii: str | np.ndarray = "auto", point_number: int = 100, - atom_filter: np.ndarray | None = None, + occlusion_mask: np.ndarray | None = None, + sasa_mask: np.ndarray | None = None, ) -> np.ndarray: """ Calculate the SASA for each atom in `atom_array`, excluding those - with NaN coordinates. The output will have the same length as the + with nonfinite coordinates. The output will have the same length as the input AtomArray, with NaN values for excluded (invalid) atoms. Args: @@ -125,20 +123,30 @@ def calculate_atomwise_sasa( receive NaN SASA (biotite default). Use "auto" for mixed structures. - np.ndarray: Custom radii array matching atom_array length. point_number: Number of points in the Shrake-Rupley algorithm. Defaults to 100. - atom_filter: Boolean mask for atoms to include. Defaults to None (all atoms). + occlusion_mask: Boolean mask of atoms that can occlude solvent. Defaults to all atoms. + Atoms without finite coordinates are excluded from calculation. + sasa_mask: Boolean mask of atoms to score. Must be a subset of occlusion_mask + (equality is allowed). Defaults to occlusion_mask. Unscored atoms return NaN. """ - has_resolved_coordinates = ~np.isnan(atom_array.coord).any(axis=-1) - if atom_filter is None: - atom_filter = has_resolved_coordinates - else: - atom_filter = atom_filter & has_resolved_coordinates - - valid_atom_array = atom_array[atom_filter] - - if len(valid_atom_array) == 0: - return np.full(atom_array.array_length(), np.nan, dtype=float) - - if atom_radii in ("element", "auto"): + n_atoms = len(atom_array) + if occlusion_mask is None: + occlusion_mask = np.ones(n_atoms, dtype=bool) + if sasa_mask is None: + sasa_mask = occlusion_mask + for name, mask in (("occlusion_mask", occlusion_mask), ("sasa_mask", sasa_mask)): + if not isinstance(mask, np.ndarray) or mask.shape != (n_atoms,) or mask.dtype != bool: + raise ValueError(f"{name} must be a boolean array of shape ({n_atoms},).") + if np.any(sasa_mask & ~occlusion_mask): + raise ValueError("sasa_mask must be a subset of occlusion_mask.") + + occlusion_mask = occlusion_mask & np.isfinite(atom_array.coord).all(axis=-1) + full_sasa = np.full(n_atoms, np.nan, dtype=float) + if not (sasa_mask & occlusion_mask).any(): + return full_sasa + valid_atom_array = atom_array[occlusion_mask] + valid_sasa_mask = sasa_mask[occlusion_mask] + + if isinstance(atom_radii, str) and atom_radii in ("element", "auto"): if atom_radii == "auto": radii = _get_protor_element_fallback_radii(valid_atom_array) else: @@ -148,17 +156,20 @@ def calculate_atomwise_sasa( probe_radius=probe_radius, vdw_radii=radii, point_number=point_number, + atom_filter=valid_sasa_mask, ignore_ions=False, ) else: - vdw_radii_arg = atom_radii[atom_filter] if isinstance(atom_radii, np.ndarray) else atom_radii + vdw_radii_arg = atom_radii[occlusion_mask] if isinstance(atom_radii, np.ndarray) else atom_radii valid_sasa = struc.sasa( - valid_atom_array, probe_radius=probe_radius, vdw_radii=vdw_radii_arg, point_number=point_number + valid_atom_array, + probe_radius=probe_radius, + vdw_radii=vdw_radii_arg, + point_number=point_number, + atom_filter=valid_sasa_mask, ) - full_sasa = np.full(atom_array.array_length(), np.nan, dtype=float) - full_sasa[atom_filter] = valid_sasa - + full_sasa[occlusion_mask] = valid_sasa return full_sasa @@ -167,7 +178,8 @@ def calculate_atomwise_rasa( probe_radius: float = 1.4, atom_radii: str | np.ndarray = "auto", point_number: int = 100, - atom_filter: np.ndarray | None = None, + occlusion_mask: np.ndarray | None = None, + sasa_mask: np.ndarray | None = None, ) -> np.ndarray: """ Calculate the Relative Solvent-Accessible Surface Area (RASA) for each atom. @@ -182,33 +194,33 @@ def calculate_atomwise_rasa( element fallback for the rest. "element" works for all atom types. "ProtOr" only works for standard protein residues (Tsai 1999). point_number: Number of Shrake-Rupley sample points. Defaults to 100. - atom_filter: Boolean mask for atoms to include. Defaults to None. + occlusion_mask: Boolean mask of atoms that can occlude solvent. Defaults to all atoms. + sasa_mask: Boolean mask of atoms to score; must be a subset of occlusion_mask + (equality is allowed). Defaults to occlusion_mask. Unscored or unresolved + atoms return NaN. Invalid masks and calculation errors raise. """ - try: - sasa = calculate_atomwise_sasa( - atom_array, - atom_filter=atom_filter, - probe_radius=probe_radius, - atom_radii=atom_radii, - point_number=point_number, - ) - except Exception as e: - logger.error(f"Error calculating SASA: {e}. Defaulting to NaN.") - return np.full(atom_array.array_length(), np.nan, dtype=float) - - try: + sasa = calculate_atomwise_sasa( + atom_array, + occlusion_mask=occlusion_mask, + sasa_mask=sasa_mask, + probe_radius=probe_radius, + atom_radii=atom_radii, + point_number=point_number, + ) + has_sasa = np.isfinite(sasa) + rasa = np.full(len(atom_array), np.nan, dtype=float) + if not has_sasa.any(): + return rasa + if isinstance(atom_radii, str): if atom_radii == "element": - vdw_radii_arr = _get_element_radii(atom_array) - elif atom_radii in ("ProtOr", "auto"): - vdw_radii_arr = _get_protor_element_fallback_radii(atom_array) + vdw_radii_arr = _get_element_radii(atom_array[has_sasa]) else: - vdw_radii_arr = atom_radii - except (ValueError, KeyError) as e: - logger.error(f"Error getting VdW radii for RASA normalization: {e}. Defaulting to NaN.") - return np.full(atom_array.array_length(), np.nan, dtype=float) + vdw_radii_arr = _get_protor_element_fallback_radii(atom_array[has_sasa]) + else: + vdw_radii_arr = atom_radii[has_sasa] max_value = 4 * np.pi * (vdw_radii_arr + probe_radius) ** 2 - rasa = sasa / max_value + rasa[has_sasa] = sasa[has_sasa] / max_value return rasa diff --git a/tests/ml/transforms/test_sasa.py b/tests/ml/transforms/test_sasa.py index 9bb16ff8..6498a05b 100644 --- a/tests/ml/transforms/test_sasa.py +++ b/tests/ml/transforms/test_sasa.py @@ -11,6 +11,7 @@ _get_element_radii, _get_protor_element_fallback_radii, calculate_atomwise_rasa, + calculate_atomwise_sasa, ) from atomworks.ml.utils.testing import cached_parse @@ -186,3 +187,20 @@ def test_element_radii_ligand_atoms(): assert not np.isnan(rasa).any(), "All atoms should have valid RASA" assert np.all(rasa >= 0) assert np.all(rasa <= 1) + + +@pytest.mark.parametrize("calculate", [calculate_atomwise_sasa, calculate_atomwise_rasa]) +def test_scoring_mask_preserves_occlusion(calculate): + atoms = AtomArray(3) + atoms.element[:] = "C" + atoms.coord[:] = [[0, 0, 0], [3, 0, 0], [np.nan, 0, 0]] + score = np.array([True, False, False]) + full = calculate(atoms, atom_radii="element") + masked = calculate(atoms, atom_radii="element", sasa_mask=score) + isolated = calculate(atoms, atom_radii="element", occlusion_mask=score) + assert 0 < masked[0] < isolated[0] + np.testing.assert_allclose(masked[0], full[0]) + assert np.isnan(masked[~score]).all() + assert np.isnan(full[2]) + with pytest.raises(ValueError, match="subset"): + calculate(atoms, occlusion_mask=score, sasa_mask=np.array([False, True, False])) From 550b2e56e1c269d08a46b0a4d866c8a307656bbb Mon Sep 17 00:00:00 2001 From: Kieran Didi <58345129+kierandidi@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:52:50 -0700 Subject: [PATCH 5/6] Stream indexed datasets directly from Hugging Face (#489) * feat: stream indexed blob records from pinned Hugging Face datasets * ci: test HF streaming and build its guide on GitHub runners * test: trim HF transport cases and explain dataset repository access * fix: avoid read-ahead for weighted HF blob reads * chore: trim HF reader CI and clarify shard caching * chore: keep HF reader PR focused on core code --- docs/huggingface.rst | 22 +++++++++++++ docs/index.rst | 1 + pyproject.toml | 1 + src/atomworks/ml/utils/blob_store.py | 11 +++++-- src/atomworks/ml/utils/huggingface.py | 45 +++++++++++++++++++++++++++ src/atomworks/ml/utils/io.py | 6 ++-- 6 files changed, 81 insertions(+), 5 deletions(-) create mode 100644 docs/huggingface.rst create mode 100644 src/atomworks/ml/utils/huggingface.py diff --git a/docs/huggingface.rst b/docs/huggingface.rst new file mode 100644 index 00000000..fdeae79c --- /dev/null +++ b/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. diff --git a/docs/index.rst b/docs/index.rst index a22004bb..222c098b 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -26,4 +26,5 @@ Welcome to **atomworks** — a toolkit for converting, parsing, and manipulating auto_examples/index contributor_guide mirrors + huggingface msa diff --git a/pyproject.toml b/pyproject.toml index 2c91e933..f104faaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dependencies = [ ] [project.optional-dependencies] +hf = ["huggingface_hub>=1.0,<2"] ml = [ # atomworks-ml dependencies "torch>=2.2.0", diff --git a/src/atomworks/ml/utils/blob_store.py b/src/atomworks/ml/utils/blob_store.py index 8ad97b7f..8f4205f3 100644 --- a/src/atomworks/ml/utils/blob_store.py +++ b/src/atomworks/ml/utils/blob_store.py @@ -5,6 +5,7 @@ import pyarrow as pa import pyarrow.feather as feather +from atomworks.ml.utils.huggingface import read_hf_bytes from atomworks.ml.utils.io import ( _decompress, build_feather_once, @@ -17,8 +18,8 @@ class BlobStore: """Read one compressed record by ``(shard, offset, length)`` from a local directory or a remote store. - A remote ``shard`` (an object-store URL) is fetched with a ranged ``GET`` via ``endpoint_url``; only - config is stored, so the store is picklable across DataLoader workers. + Remote records use ranged GETs; ``endpoint_url`` applies only to S3. HF paths + require an immutable revision. Only config is stored, so the store is picklable across workers. """ def __init__(self, data_dir: str, *, endpoint_url: str | None = None): @@ -28,8 +29,12 @@ def __init__(self, data_dir: str, *, endpoint_url: str | None = None): def get_bytes(self, shard: str, offset: int, length: int) -> bytes: """Fetch and decompress the record at ``[offset, offset+length)`` of ``shard``.""" url = f"{self.data_dir}/{shard}" - if "://" in url: + if url.startswith("hf://"): + raw = read_hf_bytes(url, offset=offset, length=length) + elif url.startswith("s3://"): raw = read_s3_bytes(url, offset=offset, length=length, endpoint_url=self.endpoint_url) + elif "://" in url: + raise ValueError("Blob stores support local paths, s3:// and revision-pinned hf://datasets/ URLs") else: with open(url, "rb") as f: f.seek(offset) diff --git a/src/atomworks/ml/utils/huggingface.py b/src/atomworks/ml/utils/huggingface.py new file mode 100644 index 00000000..89764f6f --- /dev/null +++ b/src/atomworks/ml/utils/huggingface.py @@ -0,0 +1,45 @@ +"""Authenticated, revision-pinned Hub metadata downloads and blob range reads.""" + +import os +import re +from functools import lru_cache +from typing import Any + + +def split_hf_url(url: str) -> tuple[str, str, str]: + """Parse ``hf://datasets/owner/repo@<40-character commit>/path``. + + Immutable revisions keep sampling tables, blob indices and shards consistent. + Credentials come from HF_TOKEN or the local Hugging Face login, never the URL. + """ + match = re.fullmatch(r"hf://datasets/([\w.-]+/[\w.-]+)@([a-f0-9]{40})/(.+)", url) + if not match or any(part in {"", ".", ".."} for part in match[3].split("/")): + raise ValueError("Use hf://datasets/owner/repo@/path for reproducible blob reads") + return match[1], match[2], match[3] + + +def download_hf_metadata(url: str) -> str: + """Download one pinned metadata file into the Hub's shared, locked local cache.""" + from huggingface_hub import hf_hub_download + + repo_id, revision, path = split_hf_url(url) + return hf_hub_download(repo_id, path, repo_type="dataset", revision=revision) + + +@lru_cache(maxsize=1) +def _filesystem(pid: int) -> Any: + """Keep pooled clients local to each DataLoader process, outside pickled state.""" + from huggingface_hub import HfFileSystem + + return HfFileSystem() + + +def read_hf_bytes(url: str, *, offset: int, length: int) -> bytes: + """Read exactly one compressed record without downloading its whole shard.""" + split_hf_url(url) + if offset < 0 or length <= 0: + raise ValueError("A blob range needs a nonnegative offset and positive length") + raw = _filesystem(os.getpid()).cat_file(url, start=offset, end=offset + length, cache_type="none") + if len(raw) != length: + raise OSError(f"Truncated HF blob range: expected {length} bytes, received {len(raw)}") + return raw diff --git a/src/atomworks/ml/utils/io.py b/src/atomworks/ml/utils/io.py index 09e900ed..08d934a7 100644 --- a/src/atomworks/ml/utils/io.py +++ b/src/atomworks/ml/utils/io.py @@ -24,6 +24,7 @@ from atomworks.common import string_to_md5_hash from atomworks.io.utils.io_utils import apply_sharding_pattern, build_sharding_pattern +from atomworks.ml.utils.huggingface import download_hf_metadata from atomworks.ml.utils.misc import logger try: @@ -300,8 +301,9 @@ def to_parquet_with_metadata(df: pd.DataFrame, filepath: PathLike, **kwargs: Any def _readable_source(filepath: PathLike) -> Any: - """A pandas/pyarrow-readable source for ``filepath``: the local path, or an in-memory buffer for an - ``s3://`` URL (downloaded once via boto3, since the bare readers don't accept ``s3://``).""" + """Resolve a local path, cached HF metadata file, or downloaded S3 buffer for pandas/pyarrow.""" + if str(filepath).startswith("hf://"): + return download_hf_metadata(str(filepath)) return io.BytesIO(read_s3_bytes(str(filepath))) if str(filepath).startswith("s3://") else filepath From 2ae2f6e2deaf0322dd3e83a4e3ba1fab779c7dff Mon Sep 17 00:00:00 2001 From: kierandidi Date: Thu, 8 Oct 2026 09:58:04 -0700 Subject: [PATCH 6/6] release: prepare AtomWorks 3.0.1 --- .github/workflows/release_and_docs.yaml | 2 +- README.md | 12 +++--------- docs/contributor_guide.rst | 14 +++++++------- docs/installation.rst | 2 +- pyproject.toml | 2 +- 5 files changed, 13 insertions(+), 19 deletions(-) diff --git a/.github/workflows/release_and_docs.yaml b/.github/workflows/release_and_docs.yaml index d4035b1f..ac079351 100644 --- a/.github/workflows/release_and_docs.yaml +++ b/.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] diff --git a/README.md b/README.md index 6cda31a2..c50a0377 100644 --- a/README.md +++ b/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). --- diff --git a/docs/contributor_guide.rst b/docs/contributor_guide.rst index b934e1ac..134fe9a5 100644 --- a/docs/contributor_guide.rst +++ b/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 `_. @@ -90,7 +90,7 @@ When contributing to this repository, please follow these steps: - Run autoformatting tools (``make format``) before committing. - Use commit messages like ``: ``. 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. diff --git a/docs/installation.rst b/docs/installation.rst index c8a4f82c..471bc628 100644 --- a/docs/installation.rst +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index f104faaa..86754e19 100644 --- a/pyproject.toml +++ b/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"