From 97c54e39669f952055f2557a8b0527ccf6c326d2 Mon Sep 17 00:00:00 2001 From: Behrooz Date: Wed, 31 Dec 2025 12:10:33 -0800 Subject: [PATCH 1/5] feat: Add mmCIF file support for macromolecular structures Add support for loading mmCIF (macromolecular Crystallographic Information File) format directly with load_dataset(). mmCIF is the modern standard for 3D macromolecular structures used by PDB since 2014. Key features: - Zero external dependencies: Pure Python parser for CIF syntax - Streaming support: Generator-based parsing for large structure files - Compression support: Auto-detection of gzip, bzip2, xz compressed files - ML-ready output: Atomic coordinates suitable for structure-based ML models Configuration options: - columns: Select subset of atom_site columns (default: 11 common columns) - include_hetatm: Option to exclude ligand/water HETATM records - batch_size: Control atoms per batch (default: 100000) Supported extensions: .cif, .mmcif (and compressed variants) --- docs/source/loading.mdx | 41 ++ src/datasets/packaged_modules/__init__.py | 4 + .../packaged_modules/mmcif/__init__.py | 0 src/datasets/packaged_modules/mmcif/mmcif.py | 425 ++++++++++++++ tests/packaged_modules/test_mmcif.py | 526 ++++++++++++++++++ 5 files changed, 996 insertions(+) create mode 100644 src/datasets/packaged_modules/mmcif/__init__.py create mode 100644 src/datasets/packaged_modules/mmcif/mmcif.py create mode 100644 tests/packaged_modules/test_mmcif.py diff --git a/docs/source/loading.mdx b/docs/source/loading.mdx index c39b0e3e8e2..567a163d426 100644 --- a/docs/source/loading.mdx +++ b/docs/source/loading.mdx @@ -267,6 +267,47 @@ To load remote WebDatasets via HTTP, pass the URLs instead: >>> dataset = load_dataset("webdataset", data_files={"train": urls}, split="train", streaming=True) ``` +### mmCIF + +[mmCIF](https://mmcif.wwpdb.org/) (macromolecular Crystallographic Information File) is the modern standard format for representing 3D macromolecular structures, used by the Protein Data Bank (PDB) since 2014. + +To load an mmCIF file: + +```py +>>> from datasets import load_dataset +>>> dataset = load_dataset("mmcif", data_files="structure.cif") +``` + +This returns the `_atom_site` category containing atomic coordinates. The default columns include: + +| Column | Type | Description | +|--------|------|-------------| +| `id` | int | Atom serial number | +| `type_symbol` | string | Element symbol | +| `label_atom_id` | string | Atom name (e.g., CA, N, C) | +| `label_comp_id` | string | Residue name (e.g., ALA, GLY) | +| `label_asym_id` | string | Chain identifier | +| `label_seq_id` | int | Residue sequence number | +| `Cartn_x` | float | X coordinate (Å) | +| `Cartn_y` | float | Y coordinate (Å) | +| `Cartn_z` | float | Z coordinate (Å) | +| `occupancy` | float | Occupancy factor | +| `B_iso_or_equiv` | float | Temperature factor (B-factor) | + +You can select specific columns: + +```py +>>> dataset = load_dataset("mmcif", data_files="structure.cif", columns=["label_atom_id", "Cartn_x", "Cartn_y", "Cartn_z"]) +``` + +To exclude HETATM records (ligands, water molecules): + +```py +>>> dataset = load_dataset("mmcif", data_files="structure.cif", include_hetatm=False) +``` + +Compressed files (`.cif.gz`) are automatically detected and decompressed. + ## Remote files If you have remote files likely stored as a `csv`, `json`, `txt`, `parquet` or any supported format, the [`load_dataset`] function can load load them if you specify their remote paths: diff --git a/src/datasets/packaged_modules/__init__.py b/src/datasets/packaged_modules/__init__.py index 0de63b9c785..3a2cf4f91a9 100644 --- a/src/datasets/packaged_modules/__init__.py +++ b/src/datasets/packaged_modules/__init__.py @@ -13,6 +13,7 @@ from .hdf5 import hdf5 from .iceberg import iceberg from .imagefolder import imagefolder +from .mmcif import mmcif from .json import json from .lance import lance from .meshfolder import meshfolder @@ -63,6 +64,7 @@ def _hash_python_lines(lines: list[str]) -> str: "lance": (lance.__name__, _hash_python_lines(inspect.getsource(lance).splitlines())), "tsfile": (tsfile.__name__, _hash_python_lines(inspect.getsource(tsfile).splitlines())), "iceberg": (iceberg.__name__, _hash_python_lines(inspect.getsource(iceberg).splitlines())), + "mmcif": (mmcif.__name__, _hash_python_lines(inspect.getsource(mmcif).splitlines())), } # get importable module names and hash for caching @@ -99,6 +101,8 @@ def _hash_python_lines(lines: list[str]) -> str: ".eval": ("eval", {}), ".lance": ("lance", {}), ".tsfile": ("tsfile", {}), + ".cif": ("mmcif", {}), + ".mmcif": ("mmcif", {}), } _EXTENSION_TO_MODULE.update({ext: ("imagefolder", {}) for ext in imagefolder.ImageFolder.EXTENSIONS}) _EXTENSION_TO_MODULE.update({ext.upper(): ("imagefolder", {}) for ext in imagefolder.ImageFolder.EXTENSIONS}) diff --git a/src/datasets/packaged_modules/mmcif/__init__.py b/src/datasets/packaged_modules/mmcif/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/datasets/packaged_modules/mmcif/mmcif.py b/src/datasets/packaged_modules/mmcif/mmcif.py new file mode 100644 index 00000000000..5c69351cdc7 --- /dev/null +++ b/src/datasets/packaged_modules/mmcif/mmcif.py @@ -0,0 +1,425 @@ +"""mmCIF (macromolecular CIF) file loader for 3D structure data. + +mmCIF is the modern standard format for representing 3D macromolecular structures, +used by the Protein Data Bank (PDB) since 2014. + +This implementation uses a lightweight pure Python parser for the CIF syntax, +requiring zero external dependencies. +""" + +import bz2 +import gzip +import itertools +import lzma +from dataclasses import dataclass +from typing import Optional + +import pyarrow as pa + +import datasets +from datasets.builder import Key +from datasets.features.features import require_storage_cast +from datasets.table import table_cast + + +logger = datasets.utils.logging.get_logger(__name__) + + +# Default columns for atom_site category (most commonly used for ML) +DEFAULT_ATOM_SITE_COLUMNS = [ + "id", + "type_symbol", + "label_atom_id", + "label_comp_id", + "label_asym_id", + "label_seq_id", + "Cartn_x", + "Cartn_y", + "Cartn_z", + "occupancy", + "B_iso_or_equiv", +] + +# Column type mapping for atom_site +ATOM_SITE_TYPES = { + "id": pa.int32(), + "type_symbol": pa.string(), + "label_atom_id": pa.string(), + "label_alt_id": pa.string(), + "label_comp_id": pa.string(), + "label_asym_id": pa.string(), + "label_entity_id": pa.string(), + "label_seq_id": pa.int32(), + "pdbx_PDB_ins_code": pa.string(), + "Cartn_x": pa.float32(), + "Cartn_y": pa.float32(), + "Cartn_z": pa.float32(), + "occupancy": pa.float32(), + "B_iso_or_equiv": pa.float32(), + "pdbx_formal_charge": pa.int32(), + "auth_seq_id": pa.int32(), + "auth_comp_id": pa.string(), + "auth_asym_id": pa.string(), + "auth_atom_id": pa.string(), + "pdbx_PDB_model_num": pa.int32(), + "group_PDB": pa.string(), +} + + +@dataclass +class MmcifConfig(datasets.BuilderConfig): + """BuilderConfig for mmCIF files. + + Args: + features: Dataset features (optional, will be inferred if not provided). + batch_size: Maximum number of atoms per batch. + columns: Subset of atom_site columns to include. If None, uses default columns. + include_hetatm: Whether to include HETATM records (ligands, water, etc.). + """ + + features: Optional[datasets.Features] = None + batch_size: int = 100000 + columns: Optional[list[str]] = None + include_hetatm: bool = True + + +class Mmcif(datasets.ArrowBasedBuilder): + """Dataset builder for mmCIF files.""" + + BUILDER_CONFIG_CLASS = MmcifConfig + + # Supported mmCIF extensions + EXTENSIONS: list[str] = [".cif", ".mmcif"] + + def _info(self): + return datasets.DatasetInfo(features=self.config.features) + + def _split_generators(self, dl_manager): + """Generate splits from data files.""" + if not self.config.data_files: + raise ValueError( + f"At least one data file must be specified, but got data_files={self.config.data_files}" + ) + dl_manager.download_config.extract_on_the_fly = True + data_files = dl_manager.download_and_extract(self.config.data_files) + splits = [] + for split_name, files in data_files.items(): + if isinstance(files, str): + files = [files] + files = [dl_manager.iter_files(file) for file in files] + splits.append( + datasets.SplitGenerator(name=split_name, gen_kwargs={"files": files}) + ) + return splits + + def _cast_table(self, pa_table: pa.Table) -> pa.Table: + """Cast Arrow table to configured features schema.""" + if self.config.features is not None: + schema = self.config.features.arrow_schema + if all( + not require_storage_cast(feature) + for feature in self.config.features.values() + ): + pa_table = pa_table.cast(schema) + else: + pa_table = table_cast(pa_table, schema) + return pa_table + return pa_table + + def _open_file(self, filepath: str): + """Open file with automatic compression detection based on magic bytes.""" + with open(filepath, "rb") as f: + magic = f.read(6) + + if magic[:2] == b"\x1f\x8b": # gzip magic number + return gzip.open(filepath, "rt", encoding="utf-8") + elif magic[:3] == b"BZh": # bzip2 magic number + return bz2.open(filepath, "rt", encoding="utf-8") + elif magic[:6] == b"\xfd7zXZ\x00": # xz magic number + return lzma.open(filepath, "rt", encoding="utf-8") + else: + return open(filepath, "r", encoding="utf-8") + + def _tokenize_value(self, text: str, pos: int) -> tuple[str, int]: + """Parse a single CIF value starting at position pos. + + Handles: + - Quoted strings (single or double quotes) + - Semicolon-delimited multi-line strings + - Unquoted values + """ + # Skip whitespace + while pos < len(text) and text[pos] in " \t": + pos += 1 + + if pos >= len(text) or text[pos] == "\n": + return "", pos + + char = text[pos] + + # Single or double quoted string + if char in "'\"": + end = pos + 1 + while end < len(text) and text[end] != char: + end += 1 + value = text[pos + 1 : end] + return value, end + 1 + + # Semicolon-delimited multi-line string + if char == ";": + # Find closing semicolon at start of line + end = pos + 1 + while end < len(text): + newline_pos = text.find("\n", end) + if newline_pos == -1: + break + if newline_pos + 1 < len(text) and text[newline_pos + 1] == ";": + value = text[pos + 1 : newline_pos] + return value.strip(), newline_pos + 2 + end = newline_pos + 1 + return "", len(text) + + # Unquoted value (ends at whitespace) + end = pos + while end < len(text) and text[end] not in " \t\n": + end += 1 + return text[pos:end], end + + def _parse_loop(self, lines: list[str], start_idx: int) -> tuple[dict, int]: + """Parse a loop_ construct starting at start_idx. + + Returns: + Tuple of (data_dict, end_index) where data_dict maps column names to value lists. + """ + idx = start_idx + 1 # Skip 'loop_' line + columns = [] + + # Collect column names (lines starting with _) + while idx < len(lines): + line = lines[idx].strip() + if not line or line.startswith("#"): + idx += 1 + continue + if line.startswith("_"): + columns.append(line) + idx += 1 + else: + break + + # Initialize data dict + data = {col: [] for col in columns} + num_cols = len(columns) + + if num_cols == 0: + return data, idx + + # Parse data values + values = [] + while idx < len(lines): + line = lines[idx] + stripped = line.strip() + + # Stop conditions + if not stripped: + idx += 1 + continue + if stripped.startswith("#"): + idx += 1 + continue + if stripped.startswith("_") or stripped.startswith("loop_") or stripped.startswith("data_"): + break + + # Handle semicolon-delimited multi-line values + if stripped.startswith(";"): + # Collect until closing semicolon + value_lines = [] + idx += 1 + while idx < len(lines): + if lines[idx].strip().startswith(";"): + idx += 1 + break + value_lines.append(lines[idx]) + idx += 1 + values.append("\n".join(value_lines).strip()) + continue + + # Parse inline values + pos = 0 + while pos < len(line): + # Skip whitespace + while pos < len(line) and line[pos] in " \t": + pos += 1 + if pos >= len(line): + break + + value, pos = self._tokenize_value(line, pos) + if value or value == "": + values.append(value) + + idx += 1 + + # Distribute values to columns + for i, value in enumerate(values): + col_idx = i % num_cols + data[columns[col_idx]].append(value) + + return data, idx + + def _parse_mmcif(self, fp) -> dict: + """Parse mmCIF file and extract atom_site data. + + Args: + fp: File-like object opened in text mode. + + Returns: + Dictionary mapping column names to value lists for atom_site category. + """ + content = fp.read() + lines = content.split("\n") + + atom_site_data = {} + idx = 0 + + while idx < len(lines): + line = lines[idx].strip() + + # Skip empty lines and comments + if not line or line.startswith("#"): + idx += 1 + continue + + # Look for loop_ constructs + if line.startswith("loop_"): + # Check if next non-empty line starts with _atom_site + peek_idx = idx + 1 + while peek_idx < len(lines) and not lines[peek_idx].strip(): + peek_idx += 1 + + if peek_idx < len(lines) and lines[peek_idx].strip().startswith("_atom_site."): + data, idx = self._parse_loop(lines, idx) + # Merge atom_site data + for key, values in data.items(): + if key.startswith("_atom_site."): + col_name = key.replace("_atom_site.", "") + atom_site_data[col_name] = values + else: + idx += 1 + else: + idx += 1 + + return atom_site_data + + def _get_columns(self) -> list[str]: + """Get the list of columns to include in output.""" + if self.config.columns is not None: + return self.config.columns + return DEFAULT_ATOM_SITE_COLUMNS + + def _get_schema(self, columns: list[str], available_columns: set[str]) -> pa.Schema: + """Return Arrow schema for the selected columns.""" + fields = [] + for col in columns: + if col in available_columns: + dtype = ATOM_SITE_TYPES.get(col, pa.string()) + fields.append(pa.field(col, dtype)) + return pa.schema(fields) + + def _convert_value(self, value: str, dtype: pa.DataType) -> any: + """Convert string value to appropriate Python type.""" + # Handle missing values (. or ?) + if value in (".", "?", ""): + return None + + if pa.types.is_int32(dtype): + try: + return int(value) + except ValueError: + return None + elif pa.types.is_float32(dtype): + try: + return float(value) + except ValueError: + return None + else: + return value + + def _generate_tables(self, files): + """Generate Arrow tables from mmCIF files. + + Args: + files: Iterable of file iterables from _split_generators. + + Yields: + Tuple of (Key, pa.Table) for each batch. + """ + requested_columns = self._get_columns() + + for file_idx, file in enumerate(itertools.chain.from_iterable(files)): + with self._open_file(file) as fp: + atom_site_data = self._parse_mmcif(fp) + + if not atom_site_data: + continue + + # Determine available columns + available_columns = set(atom_site_data.keys()) + + # Filter to requested columns that are available + columns = [col for col in requested_columns if col in available_columns] + + if not columns: + logger.warning( + f"No requested columns found in {file}. " + f"Available: {available_columns}, Requested: {requested_columns}" + ) + continue + + # Get number of atoms + first_col = columns[0] + num_atoms = len(atom_site_data[first_col]) + + if num_atoms == 0: + continue + + # Filter HETATM if configured + include_indices = None + if not self.config.include_hetatm and "group_PDB" in atom_site_data: + include_indices = [ + i for i, group in enumerate(atom_site_data["group_PDB"]) + if group == "ATOM" + ] + + # Build schema + schema = self._get_schema(columns, available_columns) + + # Process in batches + batch_size = self.config.batch_size + batch_idx = 0 + + for start in range(0, num_atoms, batch_size): + end = min(start + batch_size, num_atoms) + batch = {} + + for col in columns: + dtype = ATOM_SITE_TYPES.get(col, pa.string()) + raw_values = atom_site_data[col][start:end] + + if include_indices is not None: + # Filter to only ATOM records + raw_values = [ + raw_values[i - start] + for i in include_indices + if start <= i < end + ] + + # Convert values + converted = [self._convert_value(v, dtype) for v in raw_values] + batch[col] = converted + + # Skip empty batches + if not batch[columns[0]]: + continue + + pa_table = pa.Table.from_pydict(batch, schema=schema) + yield Key(file_idx, batch_idx), self._cast_table(pa_table) + batch_idx += 1 diff --git a/tests/packaged_modules/test_mmcif.py b/tests/packaged_modules/test_mmcif.py new file mode 100644 index 00000000000..963e0c2b3dd --- /dev/null +++ b/tests/packaged_modules/test_mmcif.py @@ -0,0 +1,526 @@ +"""Tests for mmCIF file loader.""" + +import gzip +import textwrap + +import pyarrow as pa +import pytest + +from datasets import Features, Value +from datasets.builder import InvalidConfigName +from datasets.data_files import DataFilesList +from datasets.packaged_modules.mmcif.mmcif import Mmcif, MmcifConfig + + +@pytest.fixture +def mmcif_file(tmp_path): + """Create a simple mmCIF file with atom_site data.""" + filename = tmp_path / "structure.cif" + data = textwrap.dedent( + """\ + data_test + # + loop_ + _atom_site.id + _atom_site.type_symbol + _atom_site.label_atom_id + _atom_site.label_comp_id + _atom_site.label_asym_id + _atom_site.label_seq_id + _atom_site.Cartn_x + _atom_site.Cartn_y + _atom_site.Cartn_z + _atom_site.occupancy + _atom_site.B_iso_or_equiv + _atom_site.group_PDB + 1 N N ALA A 1 10.000 20.000 30.000 1.00 15.00 ATOM + 2 C CA ALA A 1 11.000 21.000 31.000 1.00 14.00 ATOM + 3 C C ALA A 1 12.000 22.000 32.000 1.00 13.00 ATOM + 4 O O ALA A 1 13.000 23.000 33.000 1.00 12.00 ATOM + 5 C CB ALA A 1 14.000 24.000 34.000 1.00 16.00 ATOM + # + """ + ) + with open(filename, "w", encoding="utf-8") as f: + f.write(data) + return str(filename) + + +@pytest.fixture +def mmcif_file_with_hetatm(tmp_path): + """Create mmCIF file with ATOM and HETATM records.""" + filename = tmp_path / "structure_hetatm.cif" + data = textwrap.dedent( + """\ + data_test + # + loop_ + _atom_site.id + _atom_site.type_symbol + _atom_site.label_atom_id + _atom_site.label_comp_id + _atom_site.label_asym_id + _atom_site.label_seq_id + _atom_site.Cartn_x + _atom_site.Cartn_y + _atom_site.Cartn_z + _atom_site.occupancy + _atom_site.B_iso_or_equiv + _atom_site.group_PDB + 1 N N ALA A 1 10.000 20.000 30.000 1.00 15.00 ATOM + 2 C CA ALA A 1 11.000 21.000 31.000 1.00 14.00 ATOM + 3 O O HOH B 1 5.000 10.000 15.000 1.00 20.00 HETATM + 4 O O HOH B 2 6.000 11.000 16.000 1.00 21.00 HETATM + # + """ + ) + with open(filename, "w", encoding="utf-8") as f: + f.write(data) + return str(filename) + + +@pytest.fixture +def mmcif_file_gzipped(tmp_path): + """Create a gzipped mmCIF file.""" + filename = tmp_path / "structure.cif.gz" + data = textwrap.dedent( + """\ + data_test + # + loop_ + _atom_site.id + _atom_site.type_symbol + _atom_site.label_atom_id + _atom_site.label_comp_id + _atom_site.label_asym_id + _atom_site.label_seq_id + _atom_site.Cartn_x + _atom_site.Cartn_y + _atom_site.Cartn_z + _atom_site.occupancy + _atom_site.B_iso_or_equiv + _atom_site.group_PDB + 1 N N GLY A 1 10.000 20.000 30.000 1.00 15.00 ATOM + 2 C CA GLY A 1 11.000 21.000 31.000 1.00 14.00 ATOM + # + """ + ) + with gzip.open(filename, "wt", encoding="utf-8") as f: + f.write(data) + return str(filename) + + +@pytest.fixture +def mmcif_file_quoted_values(tmp_path): + """Create mmCIF file with quoted values.""" + filename = tmp_path / "structure_quoted.cif" + data = textwrap.dedent( + """\ + data_test + # + loop_ + _atom_site.id + _atom_site.type_symbol + _atom_site.label_atom_id + _atom_site.label_comp_id + _atom_site.label_asym_id + _atom_site.label_seq_id + _atom_site.Cartn_x + _atom_site.Cartn_y + _atom_site.Cartn_z + _atom_site.occupancy + _atom_site.B_iso_or_equiv + _atom_site.group_PDB + 1 N "N" 'ALA' A 1 10.000 20.000 30.000 1.00 15.00 ATOM + 2 C "CA" 'ALA' A 1 11.000 21.000 31.000 1.00 14.00 ATOM + # + """ + ) + with open(filename, "w", encoding="utf-8") as f: + f.write(data) + return str(filename) + + +@pytest.fixture +def mmcif_file_missing_values(tmp_path): + """Create mmCIF file with missing values (. and ?).""" + filename = tmp_path / "structure_missing.cif" + data = textwrap.dedent( + """\ + data_test + # + loop_ + _atom_site.id + _atom_site.type_symbol + _atom_site.label_atom_id + _atom_site.label_comp_id + _atom_site.label_asym_id + _atom_site.label_seq_id + _atom_site.Cartn_x + _atom_site.Cartn_y + _atom_site.Cartn_z + _atom_site.occupancy + _atom_site.B_iso_or_equiv + _atom_site.group_PDB + 1 N N ALA A . 10.000 20.000 30.000 1.00 ? ATOM + 2 C CA ALA A ? 11.000 21.000 31.000 . 15.00 ATOM + # + """ + ) + with open(filename, "w", encoding="utf-8") as f: + f.write(data) + return str(filename) + + +@pytest.fixture +def mmcif_file_large(tmp_path): + """Create a larger mmCIF file for batch testing.""" + filename = tmp_path / "structure_large.cif" + + # Build header + lines = [ + "data_test", + "#", + "loop_", + "_atom_site.id", + "_atom_site.type_symbol", + "_atom_site.label_atom_id", + "_atom_site.label_comp_id", + "_atom_site.label_asym_id", + "_atom_site.label_seq_id", + "_atom_site.Cartn_x", + "_atom_site.Cartn_y", + "_atom_site.Cartn_z", + "_atom_site.occupancy", + "_atom_site.B_iso_or_equiv", + "_atom_site.group_PDB", + ] + + # Add 500 atoms + for i in range(500): + atom_name = ["N", "CA", "C", "O", "CB"][i % 5] + element = "N" if atom_name == "N" else ("O" if atom_name == "O" else "C") + residue = i // 5 + 1 + lines.append( + f"{i+1} {element} {atom_name} ALA A {residue} " + f"{10.0 + i * 0.1:.3f} {20.0 + i * 0.1:.3f} {30.0 + i * 0.1:.3f} " + f"1.00 15.00 ATOM" + ) + + lines.append("#") + + with open(filename, "w", encoding="utf-8") as f: + f.write("\n".join(lines)) + return str(filename) + + +def test_config_raises_when_invalid_name() -> None: + with pytest.raises(InvalidConfigName, match="Bad characters"): + _ = MmcifConfig(name="name-with-*-invalid-character") + + +@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])]) +def test_config_raises_when_invalid_data_files(data_files) -> None: + with pytest.raises(ValueError, match="Expected a DataFilesDict"): + _ = MmcifConfig(name="name", data_files=data_files) + + +def test_mmcif_basic_loading(mmcif_file): + """Test basic mmCIF file loading.""" + mmcif = Mmcif() + generator = mmcif._generate_tables([[mmcif_file]]) + pa_table = pa.concat_tables([table for _, table in generator]) + + result = pa_table.to_pydict() + + assert len(result["id"]) == 5 + assert result["id"] == [1, 2, 3, 4, 5] + assert result["type_symbol"] == ["N", "C", "C", "O", "C"] + assert result["label_atom_id"] == ["N", "CA", "C", "O", "CB"] + assert result["label_comp_id"] == ["ALA", "ALA", "ALA", "ALA", "ALA"] + assert result["Cartn_x"][0] == pytest.approx(10.0) + assert result["Cartn_y"][0] == pytest.approx(20.0) + assert result["Cartn_z"][0] == pytest.approx(30.0) + + +def test_mmcif_column_filtering(mmcif_file): + """Test loading with column subset.""" + mmcif = Mmcif(columns=["label_atom_id", "Cartn_x", "Cartn_y", "Cartn_z"]) + generator = mmcif._generate_tables([[mmcif_file]]) + pa_table = pa.concat_tables([table for _, table in generator]) + + result = pa_table.to_pydict() + + # Should only have specified columns + assert set(result.keys()) == {"label_atom_id", "Cartn_x", "Cartn_y", "Cartn_z"} + assert len(result["label_atom_id"]) == 5 + + +def test_mmcif_hetatm_filtering(mmcif_file_with_hetatm): + """Test excluding HETATM records.""" + # Include HETATM (default) + mmcif_with = Mmcif(include_hetatm=True, columns=["id", "group_PDB"]) + generator = mmcif_with._generate_tables([[mmcif_file_with_hetatm]]) + pa_table = pa.concat_tables([table for _, table in generator]) + result_with = pa_table.to_pydict() + + assert len(result_with["id"]) == 4 + assert "HETATM" in result_with["group_PDB"] + + # Exclude HETATM + mmcif_without = Mmcif(include_hetatm=False, columns=["id", "group_PDB"]) + generator = mmcif_without._generate_tables([[mmcif_file_with_hetatm]]) + pa_table = pa.concat_tables([table for _, table in generator]) + result_without = pa_table.to_pydict() + + assert len(result_without["id"]) == 2 + assert "HETATM" not in result_without["group_PDB"] + + +def test_mmcif_gzipped(mmcif_file_gzipped): + """Test loading gzipped mmCIF files.""" + mmcif = Mmcif() + generator = mmcif._generate_tables([[mmcif_file_gzipped]]) + pa_table = pa.concat_tables([table for _, table in generator]) + + result = pa_table.to_pydict() + + assert len(result["id"]) == 2 + assert result["label_comp_id"] == ["GLY", "GLY"] + + +def test_mmcif_quoted_values(mmcif_file_quoted_values): + """Test parsing quoted values.""" + mmcif = Mmcif() + generator = mmcif._generate_tables([[mmcif_file_quoted_values]]) + pa_table = pa.concat_tables([table for _, table in generator]) + + result = pa_table.to_pydict() + + assert result["label_atom_id"] == ["N", "CA"] + assert result["label_comp_id"] == ["ALA", "ALA"] + + +def test_mmcif_missing_values(mmcif_file_missing_values): + """Test handling of missing values (. and ?).""" + mmcif = Mmcif() + generator = mmcif._generate_tables([[mmcif_file_missing_values]]) + pa_table = pa.concat_tables([table for _, table in generator]) + + result = pa_table.to_pydict() + + assert len(result["id"]) == 2 + # Missing values should be None + assert result["label_seq_id"][0] is None + assert result["label_seq_id"][1] is None + assert result["B_iso_or_equiv"][0] is None + assert result["occupancy"][1] is None + + +def test_mmcif_batch_size(mmcif_file_large): + """Test batch size configuration.""" + # Use batch_size=100 to create multiple batches + mmcif = Mmcif(batch_size=100) + generator = mmcif._generate_tables([[mmcif_file_large]]) + tables = [table for _, table in generator] + + # Should have 5 batches (500 atoms / 100) + assert len(tables) == 5 + + # Each batch should have 100 records + for table in tables: + assert table.num_rows == 100 + + +def test_mmcif_schema_types(mmcif_file): + """Test that schema uses correct Arrow types.""" + mmcif = Mmcif() + generator = mmcif._generate_tables([[mmcif_file]]) + pa_table = pa.concat_tables([table for _, table in generator]) + + schema = pa_table.schema + + # Check numeric types + assert schema.field("id").type == pa.int32() + assert schema.field("label_seq_id").type == pa.int32() + assert schema.field("Cartn_x").type == pa.float32() + assert schema.field("Cartn_y").type == pa.float32() + assert schema.field("Cartn_z").type == pa.float32() + assert schema.field("occupancy").type == pa.float32() + assert schema.field("B_iso_or_equiv").type == pa.float32() + + # Check string types + assert schema.field("type_symbol").type == pa.string() + assert schema.field("label_atom_id").type == pa.string() + assert schema.field("label_comp_id").type == pa.string() + + +def test_mmcif_feature_casting(mmcif_file): + """Test feature casting to custom schema.""" + features = Features({ + "id": Value("int32"), + "type_symbol": Value("string"), + "Cartn_x": Value("float32"), + "Cartn_y": Value("float32"), + "Cartn_z": Value("float32"), + }) + mmcif = Mmcif(features=features, columns=["id", "type_symbol", "Cartn_x", "Cartn_y", "Cartn_z"]) + generator = mmcif._generate_tables([[mmcif_file]]) + pa_table = pa.concat_tables([table for _, table in generator]) + + assert pa_table.schema.field("id").type == pa.int32() + assert pa_table.schema.field("Cartn_x").type == pa.float32() + + +def test_mmcif_empty_file(tmp_path): + """Test handling of empty mmCIF file.""" + filename = tmp_path / "empty.cif" + with open(filename, "w", encoding="utf-8") as f: + f.write("") + + mmcif = Mmcif() + generator = mmcif._generate_tables([[str(filename)]]) + tables = list(generator) + + # Empty file should produce no tables + assert len(tables) == 0 + + +def test_mmcif_no_atom_site(tmp_path): + """Test mmCIF file without atom_site category.""" + filename = tmp_path / "no_atoms.cif" + data = textwrap.dedent( + """\ + data_test + # + _entry.id test + _struct.title "Test structure" + # + """ + ) + with open(filename, "w", encoding="utf-8") as f: + f.write(data) + + mmcif = Mmcif() + generator = mmcif._generate_tables([[str(filename)]]) + tables = list(generator) + + # No atom_site should produce no tables + assert len(tables) == 0 + + +def test_mmcif_multiple_files(tmp_path): + """Test loading multiple mmCIF files.""" + # Create two files + file1 = tmp_path / "structure1.cif" + file2 = tmp_path / "structure2.cif" + + data1 = textwrap.dedent( + """\ + data_test1 + loop_ + _atom_site.id + _atom_site.type_symbol + _atom_site.label_atom_id + _atom_site.label_comp_id + _atom_site.label_asym_id + _atom_site.label_seq_id + _atom_site.Cartn_x + _atom_site.Cartn_y + _atom_site.Cartn_z + _atom_site.occupancy + _atom_site.B_iso_or_equiv + 1 N N ALA A 1 10.0 20.0 30.0 1.0 15.0 + """ + ) + + data2 = textwrap.dedent( + """\ + data_test2 + loop_ + _atom_site.id + _atom_site.type_symbol + _atom_site.label_atom_id + _atom_site.label_comp_id + _atom_site.label_asym_id + _atom_site.label_seq_id + _atom_site.Cartn_x + _atom_site.Cartn_y + _atom_site.Cartn_z + _atom_site.occupancy + _atom_site.B_iso_or_equiv + 1 C CA GLY B 1 11.0 21.0 31.0 1.0 14.0 + """ + ) + + with open(file1, "w", encoding="utf-8") as f: + f.write(data1) + with open(file2, "w", encoding="utf-8") as f: + f.write(data2) + + mmcif = Mmcif() + generator = mmcif._generate_tables([[str(file1)], [str(file2)]]) + pa_table = pa.concat_tables([table for _, table in generator]) + + result = pa_table.to_pydict() + + assert len(result["id"]) == 2 + assert result["label_comp_id"] == ["ALA", "GLY"] + + +def test_mmcif_extensions(): + """Test that correct extensions are defined.""" + assert ".cif" in Mmcif.EXTENSIONS + assert ".mmcif" in Mmcif.EXTENSIONS + + +def test_mmcif_comments_handling(tmp_path): + """Test that comments are properly skipped.""" + filename = tmp_path / "comments.cif" + data = textwrap.dedent( + """\ + data_test + # This is a comment + # + loop_ + _atom_site.id + _atom_site.type_symbol + _atom_site.label_atom_id + _atom_site.label_comp_id + _atom_site.label_asym_id + _atom_site.label_seq_id + _atom_site.Cartn_x + _atom_site.Cartn_y + _atom_site.Cartn_z + _atom_site.occupancy + _atom_site.B_iso_or_equiv + # Another comment + 1 N N ALA A 1 10.0 20.0 30.0 1.0 15.0 + # Comment between data + 2 C CA ALA A 1 11.0 21.0 31.0 1.0 14.0 + # + """ + ) + with open(filename, "w", encoding="utf-8") as f: + f.write(data) + + mmcif = Mmcif() + generator = mmcif._generate_tables([[str(filename)]]) + pa_table = pa.concat_tables([table for _, table in generator]) + + result = pa_table.to_pydict() + assert len(result["id"]) == 2 + + +def test_mmcif_default_columns(): + """Test that default columns are sensible.""" + from datasets.packaged_modules.mmcif.mmcif import DEFAULT_ATOM_SITE_COLUMNS + + # Essential columns should be in defaults + assert "id" in DEFAULT_ATOM_SITE_COLUMNS + assert "type_symbol" in DEFAULT_ATOM_SITE_COLUMNS + assert "label_atom_id" in DEFAULT_ATOM_SITE_COLUMNS + assert "label_comp_id" in DEFAULT_ATOM_SITE_COLUMNS + assert "Cartn_x" in DEFAULT_ATOM_SITE_COLUMNS + assert "Cartn_y" in DEFAULT_ATOM_SITE_COLUMNS + assert "Cartn_z" in DEFAULT_ATOM_SITE_COLUMNS From 7745bb5152f498dfa020573d1ce1dad770d2f2e3 Mon Sep 17 00:00:00 2001 From: Behrooz Date: Fri, 9 Jan 2026 10:30:53 -0800 Subject: [PATCH 2/5] refactor(mmcif): Adopt one-row-per-structure pattern for mmCIF loader This refactors the mmCIF loader to follow the ImageFolder pattern, where each row in the dataset contains one complete protein structure file. This is the recommended ML-friendly approach for working with structural data. Key changes: - Add ProteinStructure feature type for handling protein structure files - Supports lazy loading (decode=False) or full content (decode=True) - Works with both PDB and mmCIF formats - Rewrite MmcifFolder to extend FolderBasedBuilder - Supports folder-based labels (like ImageFolder) - Supports metadata.csv files for additional columns - Uses ProteinStructure as BASE_FEATURE - Fix bug in FolderBasedBuilder._generate_examples where drop_metadata would fail with IndexError when metadata files were in the files list - Root cause: enumerate(files) created gaps in shard_idx when files were skipped due to extension filtering - Solution: Use separate valid_shard_idx counter that only increments when samples are actually yielded Usage: >>> from datasets import load_dataset >>> dataset = load_dataset("mmcif", data_dir="./structures") >>> structure_content = dataset[0]["structure"] # Complete mmCIF content --- src/datasets/features/__init__.py | 2 + src/datasets/features/features.py | 2 + src/datasets/features/protein_structure.py | 244 ++++++ src/datasets/packaged_modules/__init__.py | 1 + src/datasets/packaged_modules/mmcif/mmcif.py | 449 +--------- tests/features/data/test_protein.cif | 34 + tests/features/data/test_protein.pdb | 12 + tests/features/test_protein_structure.py | 148 ++++ tests/fixtures/files.py | 5 + tests/packaged_modules/test_mmcif.py | 839 ++++++++----------- 10 files changed, 829 insertions(+), 907 deletions(-) create mode 100644 src/datasets/features/protein_structure.py create mode 100644 tests/features/data/test_protein.cif create mode 100644 tests/features/data/test_protein.pdb create mode 100644 tests/features/test_protein_structure.py diff --git a/src/datasets/features/__init__.py b/src/datasets/features/__init__.py index 163aaa0e2c5..3c998d88f87 100644 --- a/src/datasets/features/__init__.py +++ b/src/datasets/features/__init__.py @@ -18,6 +18,7 @@ "Video", "Pdf", "Nifti", + "ProteinStructure", ] from .audio import Audio from .features import Array2D, Array3D, Array4D, Array5D, ClassLabel, Features, Json, LargeList, List, Sequence, Value @@ -25,5 +26,6 @@ from .mesh import Mesh from .nifti import Nifti from .pdf import Pdf +from .protein_structure import ProteinStructure from .translation import Translation, TranslationVariableLanguages from .video import Video diff --git a/src/datasets/features/features.py b/src/datasets/features/features.py index 17b3d441960..7f513b7aec9 100644 --- a/src/datasets/features/features.py +++ b/src/datasets/features/features.py @@ -46,6 +46,7 @@ from .mesh import Mesh from .nifti import Nifti, encode_nibabel_image from .pdf import Pdf, encode_pdfplumber_pdf +from .protein_structure import ProteinStructure from .translation import Translation, TranslationVariableLanguages from .video import Video @@ -1536,6 +1537,7 @@ def decode_nested_example(schema, obj, token_per_repo_id: Optional[dict[str, Uni Pdf.__name__: Pdf, Nifti.__name__: Nifti, Json.__name__: Json, + ProteinStructure.__name__: ProteinStructure, } diff --git a/src/datasets/features/protein_structure.py b/src/datasets/features/protein_structure.py new file mode 100644 index 00000000000..d33ebdd38f7 --- /dev/null +++ b/src/datasets/features/protein_structure.py @@ -0,0 +1,244 @@ +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union + +import pyarrow as pa + +from .. import config +from ..download.download_config import DownloadConfig +from ..table import array_cast +from ..utils.file_utils import is_local_path, xopen +from ..utils.py_utils import no_op_if_value_is_null, string_to_dict + + +if TYPE_CHECKING: + from .features import FeatureType + + +@dataclass +class ProteinStructure: + """ + **Experimental.** + ProteinStructure [`Feature`] to read protein 3D structure files (PDB/mmCIF). + + This feature stores raw file content for lazy loading. The viewer or downstream + code handles parsing. Supports both PDB (.pdb, .ent) and mmCIF (.cif, .mmcif) formats. + + Input: The ProteinStructure feature accepts as input: + - A `str`: Absolute path to the structure file (i.e. random access is allowed). + - A `pathlib.Path`: path to the structure file (i.e. random access is allowed). + - A `dict` with the keys: + - `path`: String with relative path of the structure file in a dataset repository. + - `bytes`: Bytes of the structure file. + This is useful for archived files with sequential access. + - A `bytes` object: Raw bytes of the structure file content. + + Args: + decode (`bool`, defaults to `True`): + Whether to decode the structure data to a string. If `False`, + returns the underlying dictionary in the format `{"path": structure_path, "bytes": structure_bytes}`. + + Examples: + + ```py + >>> from datasets import Dataset, ProteinStructure + >>> ds = Dataset.from_dict({"structure": ["path/to/structure.pdb"]}).cast_column("structure", ProteinStructure()) + >>> ds.features["structure"] + ProteinStructure(decode=True, id=None) + >>> ds[0]["structure"] + 'HEADER PLANT PROTEIN...\\nATOM 1 N THR A 1...' + >>> ds = ds.cast_column("structure", ProteinStructure(decode=False)) + >>> ds[0]["structure"] + {'bytes': None, + 'path': 'path/to/structure.pdb'} + ``` + """ + + decode: bool = True + id: Optional[str] = field(default=None, repr=False) + + # Automatically constructed + dtype: ClassVar[str] = "str" + pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()}) + _type: str = field(default="ProteinStructure", init=False, repr=False) + + def __call__(self): + return self.pa_type + + def encode_example(self, value: Union[str, bytes, bytearray, dict, Path]) -> dict: + """Encode example into a format for Arrow. + + Args: + value (`str`, `bytes`, `pathlib.Path` or `dict`): + Data passed as input to ProteinStructure feature. + + Returns: + `dict` with "path" and "bytes" fields + """ + if isinstance(value, str): + return {"path": value, "bytes": None} + elif isinstance(value, Path): + return {"path": str(value.absolute()), "bytes": None} + elif isinstance(value, (bytes, bytearray)): + return {"path": None, "bytes": value} + elif isinstance(value, dict): + if value.get("path") is not None and os.path.isfile(value["path"]): + # we set "bytes": None to not duplicate the data if they're already available locally + return {"bytes": None, "path": value.get("path")} + elif value.get("bytes") is not None or value.get("path") is not None: + # store the structure bytes, and path is used to infer the format using the file extension + return {"bytes": value.get("bytes"), "path": value.get("path")} + else: + raise ValueError( + f"A protein structure sample should have one of 'path' or 'bytes' but they are missing or None in {value}." + ) + else: + raise ValueError( + f"Unsupported input type for ProteinStructure: {type(value)}. " + f"Expected str, bytes, pathlib.Path, or dict." + ) + + def decode_example(self, value: dict, token_per_repo_id=None) -> str: + """Decode example structure file into structure data. + + Args: + value (`dict`): + A dictionary with keys: + - `path`: String with absolute or relative structure file path. + - `bytes`: The bytes of the structure file. + + token_per_repo_id (`dict`, *optional*): + To access and decode structure files from private repositories on + the Hub, you can pass a dictionary repo_id (`str`) -> token (`bool` or `str`). + + Returns: + `str`: The structure file content as a string (PDB/mmCIF are text formats). + """ + if not self.decode: + raise RuntimeError("Decoding is disabled for this feature. Please use ProteinStructure(decode=True) instead.") + + if token_per_repo_id is None: + token_per_repo_id = {} + + path, bytes_ = value["path"], value["bytes"] + if bytes_ is None: + if path is None: + raise ValueError(f"A protein structure should have one of 'path' or 'bytes' but both are None in {value}.") + else: + if is_local_path(path): + with open(path, "r", encoding="utf-8") as f: + return f.read() + else: + source_url = path.split("::")[-1] + pattern = ( + config.HUB_DATASETS_URL + if source_url.startswith(config.HF_ENDPOINT) + else config.HUB_DATASETS_HFFS_URL + ) + try: + repo_id = string_to_dict(source_url, pattern)["repo_id"] + token = token_per_repo_id.get(repo_id) + except ValueError: + token = None + download_config = DownloadConfig(token=token) + with xopen(path, "r", download_config=download_config) as f: + return f.read() + else: + # PDB/mmCIF are text formats, decode bytes to string + if isinstance(bytes_, (bytes, bytearray)): + return bytes_.decode("utf-8") + # If it's already a string, return as-is + return str(bytes_) + + def flatten(self) -> Union["FeatureType", Dict[str, "FeatureType"]]: + """If in the decodable state, return the feature itself, otherwise flatten the feature into a dictionary.""" + from .features import Value + + return ( + self + if self.decode + else { + "bytes": Value("binary"), + "path": Value("string"), + } + ) + + def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArray]) -> pa.StructArray: + """Cast an Arrow array to the ProteinStructure arrow storage type. + The Arrow types that can be converted to the ProteinStructure pyarrow storage type are: + + - `pa.string()` - it must contain the "path" data + - `pa.binary()` - it must contain the structure bytes + - `pa.struct({"bytes": pa.binary()})` + - `pa.struct({"path": pa.string()})` + - `pa.struct({"bytes": pa.binary(), "path": pa.string()})` - order doesn't matter + - `pa.list(*)` - it must contain the structure array data + + Args: + storage (`Union[pa.StringArray, pa.StructArray, pa.ListArray]`): + PyArrow array to cast. + + Returns: + `pa.StructArray`: Array in the ProteinStructure arrow storage type, that is + `pa.struct({"bytes": pa.binary(), "path": pa.string()})`. + """ + if pa.types.is_string(storage.type) or pa.types.is_large_string(storage.type): + bytes_array = pa.array([None] * len(storage), type=pa.binary()) + storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null()) + elif pa.types.is_binary(storage.type) or pa.types.is_large_binary(storage.type): + path_array = pa.array([None] * len(storage), type=pa.string()) + storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null()) + elif pa.types.is_struct(storage.type): + if storage.type.get_field_index("bytes") >= 0: + bytes_array = storage.field("bytes") + else: + bytes_array = pa.array([None] * len(storage), type=pa.binary()) + if storage.type.get_field_index("path") >= 0: + path_array = storage.field("path") + else: + path_array = pa.array([None] * len(storage), type=pa.string()) + storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null()) + return array_cast(storage, self.pa_type) + + def embed_storage(self, storage: pa.StructArray, token_per_repo_id=None) -> pa.StructArray: + """Embed protein structure files into the Arrow array. + + Args: + storage (`pa.StructArray`): + PyArrow array to embed. + token_per_repo_id (`dict`, *optional*): + To access structure files from private repositories on the Hub. + + Returns: + `pa.StructArray`: Array in the ProteinStructure arrow storage type, that is + `pa.struct({"bytes": pa.binary(), "path": pa.string()})`. + """ + if token_per_repo_id is None: + token_per_repo_id = {} + + @no_op_if_value_is_null + def path_to_bytes(path): + source_url = path.split("::")[-1] + pattern = ( + config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL + ) + source_url_fields = string_to_dict(source_url, pattern) + token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None + download_config = DownloadConfig(token=token) + with xopen(path, "rb", download_config=download_config) as f: + return f.read() + + bytes_array = pa.array( + [ + (path_to_bytes(x["path"]) if x["bytes"] is None else x["bytes"]) if x is not None else None + for x in storage.to_pylist() + ], + type=pa.binary(), + ) + path_array = pa.array( + [os.path.basename(path) if path is not None else None for path in storage.field("path").to_pylist()], + type=pa.string(), + ) + storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=bytes_array.is_null()) + return array_cast(storage, self.pa_type) diff --git a/src/datasets/packaged_modules/__init__.py b/src/datasets/packaged_modules/__init__.py index 3a2cf4f91a9..e9445ac9d55 100644 --- a/src/datasets/packaged_modules/__init__.py +++ b/src/datasets/packaged_modules/__init__.py @@ -136,6 +136,7 @@ def _hash_python_lines(lines: list[str]) -> str: _MODULE_TO_METADATA_FILE_NAMES["pdffolder"] = imagefolder.ImageFolder.METADATA_FILENAMES _MODULE_TO_METADATA_FILE_NAMES["niftifolder"] = imagefolder.ImageFolder.METADATA_FILENAMES _MODULE_TO_METADATA_FILE_NAMES["lance"] = lance.Lance.METADATA_FILE_NAMES +_MODULE_TO_METADATA_FILE_NAMES["mmcif"] = imagefolder.ImageFolder.METADATA_FILENAMES _MODULE_TO_METADATA_EXTENSIONS: Dict[str, List[str]] = {} for _module in _MODULE_TO_EXTENSIONS: diff --git a/src/datasets/packaged_modules/mmcif/mmcif.py b/src/datasets/packaged_modules/mmcif/mmcif.py index 5c69351cdc7..0c19f80c87c 100644 --- a/src/datasets/packaged_modules/mmcif/mmcif.py +++ b/src/datasets/packaged_modules/mmcif/mmcif.py @@ -1,425 +1,54 @@ -"""mmCIF (macromolecular CIF) file loader for 3D structure data. - -mmCIF is the modern standard format for representing 3D macromolecular structures, -used by the Protein Data Bank (PDB) since 2014. - -This implementation uses a lightweight pure Python parser for the CIF syntax, -requiring zero external dependencies. +"""MmcifFolder - Load mmCIF protein structure files from directories. + +Follows ImageFolder pattern: folder names become labels, supports metadata files. +Each row in the resulting dataset contains one complete mmCIF structure file, +following the "one row = one structure" pattern. + +Usage: + >>> from datasets import load_dataset + >>> dataset = load_dataset("mmcif", data_dir="./structures") + +With folder-based labels: + structures/ + enzymes/ + 1abc.cif + receptors/ + 2xyz.mmcif + +With metadata file: + structures/ + metadata.csv (columns: file_name, label, resolution, ...) + 1abc.cif + 2def.mmcif """ -import bz2 -import gzip -import itertools -import lzma -from dataclasses import dataclass -from typing import Optional - -import pyarrow as pa - import datasets -from datasets.builder import Key -from datasets.features.features import require_storage_cast -from datasets.table import table_cast +from ..folder_based_builder import folder_based_builder -logger = datasets.utils.logging.get_logger(__name__) - -# Default columns for atom_site category (most commonly used for ML) -DEFAULT_ATOM_SITE_COLUMNS = [ - "id", - "type_symbol", - "label_atom_id", - "label_comp_id", - "label_asym_id", - "label_seq_id", - "Cartn_x", - "Cartn_y", - "Cartn_z", - "occupancy", - "B_iso_or_equiv", -] - -# Column type mapping for atom_site -ATOM_SITE_TYPES = { - "id": pa.int32(), - "type_symbol": pa.string(), - "label_atom_id": pa.string(), - "label_alt_id": pa.string(), - "label_comp_id": pa.string(), - "label_asym_id": pa.string(), - "label_entity_id": pa.string(), - "label_seq_id": pa.int32(), - "pdbx_PDB_ins_code": pa.string(), - "Cartn_x": pa.float32(), - "Cartn_y": pa.float32(), - "Cartn_z": pa.float32(), - "occupancy": pa.float32(), - "B_iso_or_equiv": pa.float32(), - "pdbx_formal_charge": pa.int32(), - "auth_seq_id": pa.int32(), - "auth_comp_id": pa.string(), - "auth_asym_id": pa.string(), - "auth_atom_id": pa.string(), - "pdbx_PDB_model_num": pa.int32(), - "group_PDB": pa.string(), -} +logger = datasets.utils.logging.get_logger(__name__) -@dataclass -class MmcifConfig(datasets.BuilderConfig): - """BuilderConfig for mmCIF files. +class MmcifFolderConfig(folder_based_builder.FolderBasedBuilderConfig): + """BuilderConfig for MmcifFolder.""" - Args: - features: Dataset features (optional, will be inferred if not provided). - batch_size: Maximum number of atoms per batch. - columns: Subset of atom_site columns to include. If None, uses default columns. - include_hetatm: Whether to include HETATM records (ligands, water, etc.). - """ + drop_labels: bool = None + drop_metadata: bool = None - features: Optional[datasets.Features] = None - batch_size: int = 100000 - columns: Optional[list[str]] = None - include_hetatm: bool = True + def __post_init__(self): + super().__post_init__() -class Mmcif(datasets.ArrowBasedBuilder): - """Dataset builder for mmCIF files.""" +class MmcifFolder(folder_based_builder.FolderBasedBuilder): + """Folder-based builder for mmCIF protein structure files. - BUILDER_CONFIG_CLASS = MmcifConfig + Supports mmCIF format (.cif, .mmcif). + Each row in the resulting dataset contains one complete structure file, + following the "one row = one structure" pattern recommended for ML workflows. + """ - # Supported mmCIF extensions + BASE_FEATURE = datasets.ProteinStructure + BASE_COLUMN_NAME = "structure" + BUILDER_CONFIG_CLASS = MmcifFolderConfig EXTENSIONS: list[str] = [".cif", ".mmcif"] - - def _info(self): - return datasets.DatasetInfo(features=self.config.features) - - def _split_generators(self, dl_manager): - """Generate splits from data files.""" - if not self.config.data_files: - raise ValueError( - f"At least one data file must be specified, but got data_files={self.config.data_files}" - ) - dl_manager.download_config.extract_on_the_fly = True - data_files = dl_manager.download_and_extract(self.config.data_files) - splits = [] - for split_name, files in data_files.items(): - if isinstance(files, str): - files = [files] - files = [dl_manager.iter_files(file) for file in files] - splits.append( - datasets.SplitGenerator(name=split_name, gen_kwargs={"files": files}) - ) - return splits - - def _cast_table(self, pa_table: pa.Table) -> pa.Table: - """Cast Arrow table to configured features schema.""" - if self.config.features is not None: - schema = self.config.features.arrow_schema - if all( - not require_storage_cast(feature) - for feature in self.config.features.values() - ): - pa_table = pa_table.cast(schema) - else: - pa_table = table_cast(pa_table, schema) - return pa_table - return pa_table - - def _open_file(self, filepath: str): - """Open file with automatic compression detection based on magic bytes.""" - with open(filepath, "rb") as f: - magic = f.read(6) - - if magic[:2] == b"\x1f\x8b": # gzip magic number - return gzip.open(filepath, "rt", encoding="utf-8") - elif magic[:3] == b"BZh": # bzip2 magic number - return bz2.open(filepath, "rt", encoding="utf-8") - elif magic[:6] == b"\xfd7zXZ\x00": # xz magic number - return lzma.open(filepath, "rt", encoding="utf-8") - else: - return open(filepath, "r", encoding="utf-8") - - def _tokenize_value(self, text: str, pos: int) -> tuple[str, int]: - """Parse a single CIF value starting at position pos. - - Handles: - - Quoted strings (single or double quotes) - - Semicolon-delimited multi-line strings - - Unquoted values - """ - # Skip whitespace - while pos < len(text) and text[pos] in " \t": - pos += 1 - - if pos >= len(text) or text[pos] == "\n": - return "", pos - - char = text[pos] - - # Single or double quoted string - if char in "'\"": - end = pos + 1 - while end < len(text) and text[end] != char: - end += 1 - value = text[pos + 1 : end] - return value, end + 1 - - # Semicolon-delimited multi-line string - if char == ";": - # Find closing semicolon at start of line - end = pos + 1 - while end < len(text): - newline_pos = text.find("\n", end) - if newline_pos == -1: - break - if newline_pos + 1 < len(text) and text[newline_pos + 1] == ";": - value = text[pos + 1 : newline_pos] - return value.strip(), newline_pos + 2 - end = newline_pos + 1 - return "", len(text) - - # Unquoted value (ends at whitespace) - end = pos - while end < len(text) and text[end] not in " \t\n": - end += 1 - return text[pos:end], end - - def _parse_loop(self, lines: list[str], start_idx: int) -> tuple[dict, int]: - """Parse a loop_ construct starting at start_idx. - - Returns: - Tuple of (data_dict, end_index) where data_dict maps column names to value lists. - """ - idx = start_idx + 1 # Skip 'loop_' line - columns = [] - - # Collect column names (lines starting with _) - while idx < len(lines): - line = lines[idx].strip() - if not line or line.startswith("#"): - idx += 1 - continue - if line.startswith("_"): - columns.append(line) - idx += 1 - else: - break - - # Initialize data dict - data = {col: [] for col in columns} - num_cols = len(columns) - - if num_cols == 0: - return data, idx - - # Parse data values - values = [] - while idx < len(lines): - line = lines[idx] - stripped = line.strip() - - # Stop conditions - if not stripped: - idx += 1 - continue - if stripped.startswith("#"): - idx += 1 - continue - if stripped.startswith("_") or stripped.startswith("loop_") or stripped.startswith("data_"): - break - - # Handle semicolon-delimited multi-line values - if stripped.startswith(";"): - # Collect until closing semicolon - value_lines = [] - idx += 1 - while idx < len(lines): - if lines[idx].strip().startswith(";"): - idx += 1 - break - value_lines.append(lines[idx]) - idx += 1 - values.append("\n".join(value_lines).strip()) - continue - - # Parse inline values - pos = 0 - while pos < len(line): - # Skip whitespace - while pos < len(line) and line[pos] in " \t": - pos += 1 - if pos >= len(line): - break - - value, pos = self._tokenize_value(line, pos) - if value or value == "": - values.append(value) - - idx += 1 - - # Distribute values to columns - for i, value in enumerate(values): - col_idx = i % num_cols - data[columns[col_idx]].append(value) - - return data, idx - - def _parse_mmcif(self, fp) -> dict: - """Parse mmCIF file and extract atom_site data. - - Args: - fp: File-like object opened in text mode. - - Returns: - Dictionary mapping column names to value lists for atom_site category. - """ - content = fp.read() - lines = content.split("\n") - - atom_site_data = {} - idx = 0 - - while idx < len(lines): - line = lines[idx].strip() - - # Skip empty lines and comments - if not line or line.startswith("#"): - idx += 1 - continue - - # Look for loop_ constructs - if line.startswith("loop_"): - # Check if next non-empty line starts with _atom_site - peek_idx = idx + 1 - while peek_idx < len(lines) and not lines[peek_idx].strip(): - peek_idx += 1 - - if peek_idx < len(lines) and lines[peek_idx].strip().startswith("_atom_site."): - data, idx = self._parse_loop(lines, idx) - # Merge atom_site data - for key, values in data.items(): - if key.startswith("_atom_site."): - col_name = key.replace("_atom_site.", "") - atom_site_data[col_name] = values - else: - idx += 1 - else: - idx += 1 - - return atom_site_data - - def _get_columns(self) -> list[str]: - """Get the list of columns to include in output.""" - if self.config.columns is not None: - return self.config.columns - return DEFAULT_ATOM_SITE_COLUMNS - - def _get_schema(self, columns: list[str], available_columns: set[str]) -> pa.Schema: - """Return Arrow schema for the selected columns.""" - fields = [] - for col in columns: - if col in available_columns: - dtype = ATOM_SITE_TYPES.get(col, pa.string()) - fields.append(pa.field(col, dtype)) - return pa.schema(fields) - - def _convert_value(self, value: str, dtype: pa.DataType) -> any: - """Convert string value to appropriate Python type.""" - # Handle missing values (. or ?) - if value in (".", "?", ""): - return None - - if pa.types.is_int32(dtype): - try: - return int(value) - except ValueError: - return None - elif pa.types.is_float32(dtype): - try: - return float(value) - except ValueError: - return None - else: - return value - - def _generate_tables(self, files): - """Generate Arrow tables from mmCIF files. - - Args: - files: Iterable of file iterables from _split_generators. - - Yields: - Tuple of (Key, pa.Table) for each batch. - """ - requested_columns = self._get_columns() - - for file_idx, file in enumerate(itertools.chain.from_iterable(files)): - with self._open_file(file) as fp: - atom_site_data = self._parse_mmcif(fp) - - if not atom_site_data: - continue - - # Determine available columns - available_columns = set(atom_site_data.keys()) - - # Filter to requested columns that are available - columns = [col for col in requested_columns if col in available_columns] - - if not columns: - logger.warning( - f"No requested columns found in {file}. " - f"Available: {available_columns}, Requested: {requested_columns}" - ) - continue - - # Get number of atoms - first_col = columns[0] - num_atoms = len(atom_site_data[first_col]) - - if num_atoms == 0: - continue - - # Filter HETATM if configured - include_indices = None - if not self.config.include_hetatm and "group_PDB" in atom_site_data: - include_indices = [ - i for i, group in enumerate(atom_site_data["group_PDB"]) - if group == "ATOM" - ] - - # Build schema - schema = self._get_schema(columns, available_columns) - - # Process in batches - batch_size = self.config.batch_size - batch_idx = 0 - - for start in range(0, num_atoms, batch_size): - end = min(start + batch_size, num_atoms) - batch = {} - - for col in columns: - dtype = ATOM_SITE_TYPES.get(col, pa.string()) - raw_values = atom_site_data[col][start:end] - - if include_indices is not None: - # Filter to only ATOM records - raw_values = [ - raw_values[i - start] - for i in include_indices - if start <= i < end - ] - - # Convert values - converted = [self._convert_value(v, dtype) for v in raw_values] - batch[col] = converted - - # Skip empty batches - if not batch[columns[0]]: - continue - - pa_table = pa.Table.from_pydict(batch, schema=schema) - yield Key(file_idx, batch_idx), self._cast_table(pa_table) - batch_idx += 1 diff --git a/tests/features/data/test_protein.cif b/tests/features/data/test_protein.cif new file mode 100644 index 00000000000..56b3c2f7605 --- /dev/null +++ b/tests/features/data/test_protein.cif @@ -0,0 +1,34 @@ +data_TEST +# +_entry.id TEST +# +_cell.length_a 50.000 +_cell.length_b 50.000 +_cell.length_c 50.000 +_cell.angle_alpha 90.00 +_cell.angle_beta 90.00 +_cell.angle_gamma 90.00 +# +loop_ +_atom_site.group_PDB +_atom_site.id +_atom_site.type_symbol +_atom_site.label_atom_id +_atom_site.label_comp_id +_atom_site.label_asym_id +_atom_site.label_seq_id +_atom_site.Cartn_x +_atom_site.Cartn_y +_atom_site.Cartn_z +_atom_site.occupancy +_atom_site.B_iso_or_equiv +ATOM 1 N N ALA A 1 0.000 0.000 0.000 1.00 20.00 +ATOM 2 C CA ALA A 1 1.458 0.000 0.000 1.00 20.00 +ATOM 3 C C ALA A 1 2.009 1.420 0.000 1.00 20.00 +ATOM 4 O O ALA A 1 1.251 2.390 0.000 1.00 20.00 +ATOM 5 C CB ALA A 1 1.986 -0.760 -1.217 1.00 20.00 +ATOM 6 N N GLY A 2 3.317 1.540 0.000 1.00 20.00 +ATOM 7 C CA GLY A 2 3.920 2.860 0.000 1.00 20.00 +ATOM 8 C C GLY A 2 5.440 2.780 0.000 1.00 20.00 +ATOM 9 O O GLY A 2 6.080 1.730 0.000 1.00 20.00 +# diff --git a/tests/features/data/test_protein.pdb b/tests/features/data/test_protein.pdb new file mode 100644 index 00000000000..3dbb0a184c4 --- /dev/null +++ b/tests/features/data/test_protein.pdb @@ -0,0 +1,12 @@ +HEADER TEST PROTEIN 01-JAN-00 TEST +TITLE MINIMAL TEST PDB FILE FOR DATASETS LIBRARY +ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 20.00 N +ATOM 2 CA ALA A 1 1.458 0.000 0.000 1.00 20.00 C +ATOM 3 C ALA A 1 2.009 1.420 0.000 1.00 20.00 C +ATOM 4 O ALA A 1 1.251 2.390 0.000 1.00 20.00 O +ATOM 5 CB ALA A 1 1.986 -0.760 -1.217 1.00 20.00 C +ATOM 6 N GLY A 2 3.317 1.540 0.000 1.00 20.00 N +ATOM 7 CA GLY A 2 3.920 2.860 0.000 1.00 20.00 C +ATOM 8 C GLY A 2 5.440 2.780 0.000 1.00 20.00 C +ATOM 9 O GLY A 2 6.080 1.730 0.000 1.00 20.00 O +END diff --git a/tests/features/test_protein_structure.py b/tests/features/test_protein_structure.py new file mode 100644 index 00000000000..34ba6d88dc7 --- /dev/null +++ b/tests/features/test_protein_structure.py @@ -0,0 +1,148 @@ +from pathlib import Path + +import pyarrow as pa +import pytest + +from datasets import Dataset, Features, ProteinStructure + + +@pytest.mark.parametrize( + "build_example", + [ + lambda structure_path: structure_path, + lambda structure_path: Path(structure_path), + lambda structure_path: open(structure_path, "rb").read(), + lambda structure_path: {"path": structure_path}, + lambda structure_path: {"path": structure_path, "bytes": None}, + lambda structure_path: {"path": structure_path, "bytes": open(structure_path, "rb").read()}, + lambda structure_path: {"path": None, "bytes": open(structure_path, "rb").read()}, + lambda structure_path: {"bytes": open(structure_path, "rb").read()}, + ], +) +def test_protein_structure_feature_encode_example(shared_datadir, build_example): + structure_path = str(shared_datadir / "test_protein.pdb") + protein_structure = ProteinStructure() + encoded_example = protein_structure.encode_example(build_example(structure_path)) + assert isinstance(encoded_example, dict) + assert encoded_example.keys() == {"bytes", "path"} + assert encoded_example["bytes"] is not None or encoded_example["path"] is not None + decoded_example = protein_structure.decode_example(encoded_example) + assert isinstance(decoded_example, str) + + +def test_protein_structure_decode_example_pdb(shared_datadir): + structure_path = str(shared_datadir / "test_protein.pdb") + protein_structure = ProteinStructure() + decoded_example = protein_structure.decode_example({"path": structure_path, "bytes": None}) + + assert isinstance(decoded_example, str) + assert "ATOM" in decoded_example + assert "HEADER" in decoded_example + + with pytest.raises(RuntimeError): + ProteinStructure(decode=False).decode_example({"path": structure_path, "bytes": None}) + + +def test_protein_structure_decode_example_cif(shared_datadir): + structure_path = str(shared_datadir / "test_protein.cif") + protein_structure = ProteinStructure() + decoded_example = protein_structure.decode_example({"path": structure_path, "bytes": None}) + + assert isinstance(decoded_example, str) + assert "data_" in decoded_example + assert "_atom_site" in decoded_example + + +def test_dataset_with_protein_structure_feature(shared_datadir): + structure_path = str(shared_datadir / "test_protein.pdb") + data = {"structure": [structure_path]} + features = Features({"structure": ProteinStructure()}) + dset = Dataset.from_dict(data, features=features) + item = dset[0] + assert item.keys() == {"structure"} + assert isinstance(item["structure"], str) + assert "ATOM" in item["structure"] + batch = dset[:1] + assert len(batch) == 1 + assert batch.keys() == {"structure"} + assert isinstance(batch["structure"], list) and all(isinstance(item, str) for item in batch["structure"]) + column = dset["structure"] + assert len(column) == 1 + assert all(isinstance(item, str) for item in column) + + # from bytes + with open(structure_path, "rb") as f: + data = {"structure": [f.read()]} + dset = Dataset.from_dict(data, features=features) + item = dset[0] + assert item.keys() == {"structure"} + assert isinstance(item["structure"], str) + + +def test_protein_structure_pa_type(): + protein_structure = ProteinStructure() + pa_type = protein_structure.pa_type + + assert pa.types.is_struct(pa_type) + assert pa_type.get_field_index("bytes") >= 0 + assert pa_type.get_field_index("path") >= 0 + assert pa.types.is_binary(pa_type.field("bytes").type) + assert pa.types.is_string(pa_type.field("path").type) + + +def test_protein_structure_cast_storage(shared_datadir): + structure_path = str(shared_datadir / "test_protein.pdb") + protein_structure = ProteinStructure() + + # From string array (path) + string_array = pa.array([structure_path], type=pa.string()) + casted = protein_structure.cast_storage(string_array) + assert pa.types.is_struct(casted.type) + assert casted.type.get_field_index("bytes") >= 0 + assert casted.type.get_field_index("path") >= 0 + + # From binary array (bytes) + with open(structure_path, "rb") as f: + content = f.read() + binary_array = pa.array([content], type=pa.binary()) + casted = protein_structure.cast_storage(binary_array) + assert pa.types.is_struct(casted.type) + + # From struct array + bytes_array = pa.array([content], type=pa.binary()) + path_array = pa.array([structure_path], type=pa.string()) + struct_array = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"]) + casted = protein_structure.cast_storage(struct_array) + assert pa.types.is_struct(casted.type) + + +def test_protein_structure_embed_storage(shared_datadir): + structure_path = str(shared_datadir / "test_protein.pdb") + protein_structure = ProteinStructure() + + bytes_array = pa.array([None], type=pa.binary()) + path_array = pa.array([structure_path], type=pa.string()) + storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"]) + + embedded_storage = protein_structure.embed_storage(storage) + + embedded_bytes = embedded_storage[0]["bytes"].as_py() + assert embedded_bytes is not None + assert len(embedded_bytes) > 0 + + content = embedded_bytes.decode("utf-8") + assert "ATOM" in content + + +def test_protein_structure_flatten(): + # With decode=True + protein_structure = ProteinStructure(decode=True) + flattened = protein_structure.flatten() + assert flattened is protein_structure # Returns self when decode=True + + # With decode=False + protein_structure_no_decode = ProteinStructure(decode=False) + flattened = protein_structure_no_decode.flatten() + assert isinstance(flattened, dict) + assert "bytes" in flattened + assert "path" in flattened diff --git a/tests/fixtures/files.py b/tests/fixtures/files.py index a3dee4b5cd2..8d9407cc487 100644 --- a/tests/fixtures/files.py +++ b/tests/fixtures/files.py @@ -627,6 +627,11 @@ def zip_image_path(image_file, tmp_path_factory): return path +@pytest.fixture(scope="session") +def cif_file(): + return os.path.join("tests", "features", "data", "test_protein.cif") + + @pytest.fixture(scope="session") def data_dir_with_hidden_files(tmp_path_factory): data_dir = tmp_path_factory.mktemp("data_dir") diff --git a/tests/packaged_modules/test_mmcif.py b/tests/packaged_modules/test_mmcif.py index 963e0c2b3dd..4ea4b488c0e 100644 --- a/tests/packaged_modules/test_mmcif.py +++ b/tests/packaged_modules/test_mmcif.py @@ -1,526 +1,371 @@ -"""Tests for mmCIF file loader.""" +"""Tests for MmcifFolder packaged module. + +Tests the folder-based mmCIF loader that follows the one-row-per-structure pattern +(similar to ImageFolder pattern). +""" -import gzip import textwrap import pyarrow as pa import pytest -from datasets import Features, Value -from datasets.builder import InvalidConfigName -from datasets.data_files import DataFilesList -from datasets.packaged_modules.mmcif.mmcif import Mmcif, MmcifConfig +from datasets import ClassLabel, Dataset, DatasetDict, ProteinStructure, load_dataset +from datasets.packaged_modules.mmcif.mmcif import MmcifFolder + + +# Test data - minimal mmCIF format +MMCIF_CONTENT = """\ +data_TEST +# +_entry.id TEST +# +_cell.length_a 50.000 +_cell.length_b 50.000 +_cell.length_c 50.000 +_cell.angle_alpha 90.00 +_cell.angle_beta 90.00 +_cell.angle_gamma 90.00 +# +loop_ +_atom_site.group_PDB +_atom_site.id +_atom_site.type_symbol +_atom_site.label_atom_id +_atom_site.label_comp_id +_atom_site.label_asym_id +_atom_site.label_seq_id +_atom_site.Cartn_x +_atom_site.Cartn_y +_atom_site.Cartn_z +_atom_site.occupancy +_atom_site.B_iso_or_equiv +ATOM 1 N N ALA A 1 0.000 0.000 0.000 1.00 20.00 +ATOM 2 C CA ALA A 1 1.458 0.000 0.000 1.00 20.00 +# +""" @pytest.fixture def mmcif_file(tmp_path): - """Create a simple mmCIF file with atom_site data.""" - filename = tmp_path / "structure.cif" - data = textwrap.dedent( - """\ - data_test - # - loop_ - _atom_site.id - _atom_site.type_symbol - _atom_site.label_atom_id - _atom_site.label_comp_id - _atom_site.label_asym_id - _atom_site.label_seq_id - _atom_site.Cartn_x - _atom_site.Cartn_y - _atom_site.Cartn_z - _atom_site.occupancy - _atom_site.B_iso_or_equiv - _atom_site.group_PDB - 1 N N ALA A 1 10.000 20.000 30.000 1.00 15.00 ATOM - 2 C CA ALA A 1 11.000 21.000 31.000 1.00 14.00 ATOM - 3 C C ALA A 1 12.000 22.000 32.000 1.00 13.00 ATOM - 4 O O ALA A 1 13.000 23.000 33.000 1.00 12.00 ATOM - 5 C CB ALA A 1 14.000 24.000 34.000 1.00 16.00 ATOM - # - """ - ) - with open(filename, "w", encoding="utf-8") as f: - f.write(data) - return str(filename) + """Create a single mmCIF test file.""" + path = tmp_path / "structure.cif" + path.write_text(MMCIF_CONTENT) + return str(path) @pytest.fixture -def mmcif_file_with_hetatm(tmp_path): - """Create mmCIF file with ATOM and HETATM records.""" - filename = tmp_path / "structure_hetatm.cif" - data = textwrap.dedent( - """\ - data_test - # - loop_ - _atom_site.id - _atom_site.type_symbol - _atom_site.label_atom_id - _atom_site.label_comp_id - _atom_site.label_asym_id - _atom_site.label_seq_id - _atom_site.Cartn_x - _atom_site.Cartn_y - _atom_site.Cartn_z - _atom_site.occupancy - _atom_site.B_iso_or_equiv - _atom_site.group_PDB - 1 N N ALA A 1 10.000 20.000 30.000 1.00 15.00 ATOM - 2 C CA ALA A 1 11.000 21.000 31.000 1.00 14.00 ATOM - 3 O O HOH B 1 5.000 10.000 15.000 1.00 20.00 HETATM - 4 O O HOH B 2 6.000 11.000 16.000 1.00 21.00 HETATM - # - """ - ) - with open(filename, "w", encoding="utf-8") as f: - f.write(data) - return str(filename) +def mmcif_file_mmcif_ext(tmp_path): + """Create a single mmCIF test file with .mmcif extension.""" + path = tmp_path / "structure.mmcif" + path.write_text(MMCIF_CONTENT) + return str(path) @pytest.fixture -def mmcif_file_gzipped(tmp_path): - """Create a gzipped mmCIF file.""" - filename = tmp_path / "structure.cif.gz" - data = textwrap.dedent( - """\ - data_test - # - loop_ - _atom_site.id - _atom_site.type_symbol - _atom_site.label_atom_id - _atom_site.label_comp_id - _atom_site.label_asym_id - _atom_site.label_seq_id - _atom_site.Cartn_x - _atom_site.Cartn_y - _atom_site.Cartn_z - _atom_site.occupancy - _atom_site.B_iso_or_equiv - _atom_site.group_PDB - 1 N N GLY A 1 10.000 20.000 30.000 1.00 15.00 ATOM - 2 C CA GLY A 1 11.000 21.000 31.000 1.00 14.00 ATOM - # - """ - ) - with gzip.open(filename, "wt", encoding="utf-8") as f: - f.write(data) - return str(filename) - +def mmcif_data_dir(tmp_path): + """Create a directory with mmCIF files.""" + data_dir = tmp_path / "mmcif_data" + data_dir.mkdir() -@pytest.fixture -def mmcif_file_quoted_values(tmp_path): - """Create mmCIF file with quoted values.""" - filename = tmp_path / "structure_quoted.cif" - data = textwrap.dedent( - """\ - data_test - # - loop_ - _atom_site.id - _atom_site.type_symbol - _atom_site.label_atom_id - _atom_site.label_comp_id - _atom_site.label_asym_id - _atom_site.label_seq_id - _atom_site.Cartn_x - _atom_site.Cartn_y - _atom_site.Cartn_z - _atom_site.occupancy - _atom_site.B_iso_or_equiv - _atom_site.group_PDB - 1 N "N" 'ALA' A 1 10.000 20.000 30.000 1.00 15.00 ATOM - 2 C "CA" 'ALA' A 1 11.000 21.000 31.000 1.00 14.00 ATOM - # - """ - ) - with open(filename, "w", encoding="utf-8") as f: - f.write(data) - return str(filename) + # Create several mmCIF files + for i in range(3): + (data_dir / f"structure_{i}.cif").write_text(MMCIF_CONTENT.replace("TEST", f"TEST_{i}")) - -@pytest.fixture -def mmcif_file_missing_values(tmp_path): - """Create mmCIF file with missing values (. and ?).""" - filename = tmp_path / "structure_missing.cif" - data = textwrap.dedent( - """\ - data_test - # - loop_ - _atom_site.id - _atom_site.type_symbol - _atom_site.label_atom_id - _atom_site.label_comp_id - _atom_site.label_asym_id - _atom_site.label_seq_id - _atom_site.Cartn_x - _atom_site.Cartn_y - _atom_site.Cartn_z - _atom_site.occupancy - _atom_site.B_iso_or_equiv - _atom_site.group_PDB - 1 N N ALA A . 10.000 20.000 30.000 1.00 ? ATOM - 2 C CA ALA A ? 11.000 21.000 31.000 . 15.00 ATOM - # - """ - ) - with open(filename, "w", encoding="utf-8") as f: - f.write(data) - return str(filename) + return str(data_dir) @pytest.fixture -def mmcif_file_large(tmp_path): - """Create a larger mmCIF file for batch testing.""" - filename = tmp_path / "structure_large.cif" - - # Build header - lines = [ - "data_test", - "#", - "loop_", - "_atom_site.id", - "_atom_site.type_symbol", - "_atom_site.label_atom_id", - "_atom_site.label_comp_id", - "_atom_site.label_asym_id", - "_atom_site.label_seq_id", - "_atom_site.Cartn_x", - "_atom_site.Cartn_y", - "_atom_site.Cartn_z", - "_atom_site.occupancy", - "_atom_site.B_iso_or_equiv", - "_atom_site.group_PDB", - ] - - # Add 500 atoms - for i in range(500): - atom_name = ["N", "CA", "C", "O", "CB"][i % 5] - element = "N" if atom_name == "N" else ("O" if atom_name == "O" else "C") - residue = i // 5 + 1 - lines.append( - f"{i+1} {element} {atom_name} ALA A {residue} " - f"{10.0 + i * 0.1:.3f} {20.0 + i * 0.1:.3f} {30.0 + i * 0.1:.3f} " - f"1.00 15.00 ATOM" - ) - - lines.append("#") - - with open(filename, "w", encoding="utf-8") as f: - f.write("\n".join(lines)) - return str(filename) - - -def test_config_raises_when_invalid_name() -> None: - with pytest.raises(InvalidConfigName, match="Bad characters"): - _ = MmcifConfig(name="name-with-*-invalid-character") - - -@pytest.mark.parametrize("data_files", ["str_path", ["str_path"], DataFilesList(["str_path"], [()])]) -def test_config_raises_when_invalid_data_files(data_files) -> None: - with pytest.raises(ValueError, match="Expected a DataFilesDict"): - _ = MmcifConfig(name="name", data_files=data_files) - - -def test_mmcif_basic_loading(mmcif_file): - """Test basic mmCIF file loading.""" - mmcif = Mmcif() - generator = mmcif._generate_tables([[mmcif_file]]) - pa_table = pa.concat_tables([table for _, table in generator]) - - result = pa_table.to_pydict() - - assert len(result["id"]) == 5 - assert result["id"] == [1, 2, 3, 4, 5] - assert result["type_symbol"] == ["N", "C", "C", "O", "C"] - assert result["label_atom_id"] == ["N", "CA", "C", "O", "CB"] - assert result["label_comp_id"] == ["ALA", "ALA", "ALA", "ALA", "ALA"] - assert result["Cartn_x"][0] == pytest.approx(10.0) - assert result["Cartn_y"][0] == pytest.approx(20.0) - assert result["Cartn_z"][0] == pytest.approx(30.0) - - -def test_mmcif_column_filtering(mmcif_file): - """Test loading with column subset.""" - mmcif = Mmcif(columns=["label_atom_id", "Cartn_x", "Cartn_y", "Cartn_z"]) - generator = mmcif._generate_tables([[mmcif_file]]) - pa_table = pa.concat_tables([table for _, table in generator]) +def mmcif_data_dir_with_labels(tmp_path): + """Create directories with label-based structure.""" + data_dir = tmp_path / "mmcif_labeled" - result = pa_table.to_pydict() + # Create labeled directories (like ImageFolder pattern) + for label in ["enzymes", "receptors"]: + label_dir = data_dir / label + label_dir.mkdir(parents=True) + for i in range(2): + content = MMCIF_CONTENT.replace("TEST", f"{label.upper()}_{i}") + (label_dir / f"{label}_{i}.cif").write_text(content) - # Should only have specified columns - assert set(result.keys()) == {"label_atom_id", "Cartn_x", "Cartn_y", "Cartn_z"} - assert len(result["label_atom_id"]) == 5 + return str(data_dir) -def test_mmcif_hetatm_filtering(mmcif_file_with_hetatm): - """Test excluding HETATM records.""" - # Include HETATM (default) - mmcif_with = Mmcif(include_hetatm=True, columns=["id", "group_PDB"]) - generator = mmcif_with._generate_tables([[mmcif_file_with_hetatm]]) - pa_table = pa.concat_tables([table for _, table in generator]) - result_with = pa_table.to_pydict() - - assert len(result_with["id"]) == 4 - assert "HETATM" in result_with["group_PDB"] - - # Exclude HETATM - mmcif_without = Mmcif(include_hetatm=False, columns=["id", "group_PDB"]) - generator = mmcif_without._generate_tables([[mmcif_file_with_hetatm]]) - pa_table = pa.concat_tables([table for _, table in generator]) - result_without = pa_table.to_pydict() - - assert len(result_without["id"]) == 2 - assert "HETATM" not in result_without["group_PDB"] - - -def test_mmcif_gzipped(mmcif_file_gzipped): - """Test loading gzipped mmCIF files.""" - mmcif = Mmcif() - generator = mmcif._generate_tables([[mmcif_file_gzipped]]) - pa_table = pa.concat_tables([table for _, table in generator]) - - result = pa_table.to_pydict() - - assert len(result["id"]) == 2 - assert result["label_comp_id"] == ["GLY", "GLY"] - - -def test_mmcif_quoted_values(mmcif_file_quoted_values): - """Test parsing quoted values.""" - mmcif = Mmcif() - generator = mmcif._generate_tables([[mmcif_file_quoted_values]]) - pa_table = pa.concat_tables([table for _, table in generator]) - - result = pa_table.to_pydict() - - assert result["label_atom_id"] == ["N", "CA"] - assert result["label_comp_id"] == ["ALA", "ALA"] - - -def test_mmcif_missing_values(mmcif_file_missing_values): - """Test handling of missing values (. and ?).""" - mmcif = Mmcif() - generator = mmcif._generate_tables([[mmcif_file_missing_values]]) - pa_table = pa.concat_tables([table for _, table in generator]) +@pytest.fixture +def mmcif_data_dir_with_metadata(tmp_path): + """Create directory with metadata.csv file.""" + data_dir = tmp_path / "mmcif_metadata" + data_dir.mkdir() + + # Create mmCIF files + for i in range(3): + (data_dir / f"structure_{i}.cif").write_text(MMCIF_CONTENT.replace("TEST", f"META_{i}")) + + # Create metadata.csv + metadata = textwrap.dedent("""\ + file_name,resolution,method + structure_0.cif,2.5,X-RAY + structure_1.cif,1.8,CRYO-EM + structure_2.cif,3.0,NMR + """) + (data_dir / "metadata.csv").write_text(metadata) + + return str(data_dir) + + +class TestMmcifFolderConfig: + """Tests for MmcifFolderConfig.""" + + def test_config_defaults(self): + """Test default config values.""" + builder = MmcifFolder(data_dir=".") + assert builder.config.drop_labels is None + assert builder.config.drop_metadata is None + + def test_config_custom_values(self): + """Test custom config values.""" + builder = MmcifFolder(data_dir=".", drop_labels=True, drop_metadata=True) + assert builder.config.drop_labels is True + assert builder.config.drop_metadata is True + + +class TestMmcifFolderBuilder: + """Tests for MmcifFolder builder properties.""" + + def test_base_feature(self): + """Test BASE_FEATURE is ProteinStructure.""" + assert MmcifFolder.BASE_FEATURE == ProteinStructure + + def test_base_column_name(self): + """Test BASE_COLUMN_NAME is 'structure'.""" + assert MmcifFolder.BASE_COLUMN_NAME == "structure" + + def test_extensions(self): + """Test supported extensions.""" + assert ".cif" in MmcifFolder.EXTENSIONS + assert ".mmcif" in MmcifFolder.EXTENSIONS + assert len(MmcifFolder.EXTENSIONS) == 2 + + +class TestMmcifFolderLoading: + """Tests for loading mmCIF datasets.""" + + def test_load_single_file(self, mmcif_file): + """Test loading a single mmCIF file.""" + dataset = load_dataset("mmcif", data_files=mmcif_file, split="train") + assert len(dataset) == 1 + assert "structure" in dataset.column_names + + def test_load_mmcif_extension(self, mmcif_file_mmcif_ext): + """Test loading .mmcif extension.""" + dataset = load_dataset("mmcif", data_files=mmcif_file_mmcif_ext, split="train") + assert len(dataset) == 1 + assert "structure" in dataset.column_names + + def test_load_directory(self, mmcif_data_dir): + """Test loading from directory.""" + dataset = load_dataset("mmcif", data_dir=mmcif_data_dir, split="train") + assert len(dataset) == 3 + assert "structure" in dataset.column_names + + def test_load_with_labels(self, mmcif_data_dir_with_labels): + """Test loading with folder-based labels.""" + dataset = load_dataset("mmcif", data_dir=mmcif_data_dir_with_labels, split="train") + assert len(dataset) == 4 + assert "structure" in dataset.column_names + assert "label" in dataset.column_names + # Check labels are ClassLabel + assert isinstance(dataset.features["label"], ClassLabel) + + def test_load_drop_labels(self, mmcif_data_dir_with_labels): + """Test dropping labels.""" + dataset = load_dataset( + "mmcif", data_dir=mmcif_data_dir_with_labels, split="train", drop_labels=True + ) + assert len(dataset) == 4 + assert "structure" in dataset.column_names + assert "label" not in dataset.column_names + + def test_load_with_metadata(self, mmcif_data_dir_with_metadata): + """Test loading with metadata.csv.""" + dataset = load_dataset("mmcif", data_dir=mmcif_data_dir_with_metadata, split="train") + assert len(dataset) == 3 + assert "structure" in dataset.column_names + assert "resolution" in dataset.column_names + assert "method" in dataset.column_names + + def test_load_drop_metadata(self, mmcif_data_dir_with_metadata): + """Test dropping metadata.""" + dataset = load_dataset( + "mmcif", data_dir=mmcif_data_dir_with_metadata, split="train", drop_metadata=True + ) + assert len(dataset) == 3 + assert "structure" in dataset.column_names + assert "resolution" not in dataset.column_names + assert "method" not in dataset.column_names + + +class TestMmcifFolderFeatures: + """Tests for ProteinStructure feature in mmCIF datasets.""" + + def test_feature_type(self, mmcif_file): + """Test feature type is ProteinStructure.""" + dataset = load_dataset("mmcif", data_files=mmcif_file, split="train") + assert isinstance(dataset.features["structure"], ProteinStructure) + + def test_structure_content_decoded(self, mmcif_file): + """Test structure content is accessible when decoded (default).""" + dataset = load_dataset("mmcif", data_files=mmcif_file, split="train") + sample = dataset[0] + structure = sample["structure"] + + # With decode=True (default), structure is a string + assert isinstance(structure, str) + assert "data_TEST" in structure + assert "_atom_site" in structure + + def test_structure_content_undecoded(self, mmcif_file): + """Test structure content when decode=False.""" + dataset = load_dataset("mmcif", data_files=mmcif_file, split="train") + # Cast to decode=False to get raw bytes/path dict + dataset = dataset.cast_column("structure", ProteinStructure(decode=False)) + sample = dataset[0] + structure = sample["structure"] + + # Should have bytes and path + assert "bytes" in structure or "path" in structure + # Path should end with .cif + if structure["path"]: + assert structure["path"].endswith(".cif") + + def test_arrow_type(self, mmcif_file): + """Test underlying Arrow type.""" + dataset = load_dataset("mmcif", data_files=mmcif_file, split="train") + feature = dataset.features["structure"] + pa_type = feature.pa_type + + # Should be a struct with bytes and path + assert pa.types.is_struct(pa_type) + field_names = [f.name for f in pa_type] + assert "bytes" in field_names + assert "path" in field_names + + +class TestMmcifFolderIntegration: + """Integration tests for MmcifFolder.""" + + def test_dataset_operations(self, mmcif_data_dir): + """Test common dataset operations.""" + dataset = load_dataset("mmcif", data_dir=mmcif_data_dir, split="train") + + # Test filter + filtered = dataset.filter(lambda x: True) + assert len(filtered) == len(dataset) + + # Test select + selected = dataset.select([0, 1]) + assert len(selected) == 2 + + # Test shuffle + shuffled = dataset.shuffle(seed=42) + assert len(shuffled) == len(dataset) + + def test_train_test_split(self, mmcif_data_dir): + """Test train/test split.""" + dataset = load_dataset("mmcif", data_dir=mmcif_data_dir, split="train") + splits = dataset.train_test_split(test_size=0.3) + + assert isinstance(splits, DatasetDict) + assert "train" in splits + assert "test" in splits + assert len(splits["train"]) + len(splits["test"]) == len(dataset) + + def test_map_function(self, mmcif_data_dir): + """Test map function on dataset.""" + dataset = load_dataset("mmcif", data_dir=mmcif_data_dir, split="train") + + def extract_entry_id(example): + # structure is decoded as a string by default + content = example["structure"] + # Extract entry ID from mmCIF content + for line in content.split("\n"): + if line.startswith("data_"): + return {"entry_id": line.replace("data_", "")} + return {"entry_id": "unknown"} + + mapped = dataset.map(extract_entry_id) + assert "entry_id" in mapped.column_names + + def test_save_and_load(self, mmcif_data_dir, tmp_path): + """Test save and reload dataset.""" + dataset = load_dataset("mmcif", data_dir=mmcif_data_dir, split="train") + + # Save to disk + save_path = tmp_path / "saved_mmcif" + dataset.save_to_disk(str(save_path)) + + # Reload + reloaded = Dataset.load_from_disk(str(save_path)) + assert len(reloaded) == len(dataset) + assert reloaded.features["structure"] == dataset.features["structure"] + + +class TestMmcifFolderEdgeCases: + """Edge case tests for MmcifFolder.""" + + def test_empty_directory(self, tmp_path): + """Test loading from empty directory raises error.""" + empty_dir = tmp_path / "empty" + empty_dir.mkdir() + + with pytest.raises(Exception): # FileNotFoundError or similar + load_dataset("mmcif", data_dir=str(empty_dir), split="train") + + def test_mixed_extensions(self, tmp_path): + """Test directory with both .cif and .mmcif files.""" + data_dir = tmp_path / "mixed" + data_dir.mkdir() + + (data_dir / "structure1.cif").write_text(MMCIF_CONTENT.replace("TEST", "CIF1")) + (data_dir / "structure2.mmcif").write_text(MMCIF_CONTENT.replace("TEST", "MMCIF2")) + + dataset = load_dataset("mmcif", data_dir=str(data_dir), split="train") + assert len(dataset) == 2 + + def test_nested_directories(self, tmp_path): + """Test nested directory structure.""" + data_dir = tmp_path / "nested" + sub_dir = data_dir / "sub1" / "sub2" + sub_dir.mkdir(parents=True) + + (data_dir / "root.cif").write_text(MMCIF_CONTENT.replace("TEST", "ROOT")) + (sub_dir / "nested.cif").write_text(MMCIF_CONTENT.replace("TEST", "NESTED")) + + dataset = load_dataset("mmcif", data_dir=str(data_dir), split="train") + # Should find files in nested directories + assert len(dataset) >= 1 + + def test_large_mmcif_content(self, tmp_path): + """Test with larger mmCIF file.""" + data_dir = tmp_path / "large" + data_dir.mkdir() - result = pa_table.to_pydict() + # Create mmCIF with many atoms + atoms = [] + for i in range(100): + atoms.append(f"ATOM {i+1} N N ALA A {i+1} {i}.000 0.000 0.000 1.00 20.00") + + large_content = MMCIF_CONTENT.replace( + "ATOM 1 N N ALA A 1 0.000 0.000 0.000 1.00 20.00\n" + "ATOM 2 C CA ALA A 1 1.458 0.000 0.000 1.00 20.00", + "\n".join(atoms) + ) - assert len(result["id"]) == 2 - # Missing values should be None - assert result["label_seq_id"][0] is None - assert result["label_seq_id"][1] is None - assert result["B_iso_or_equiv"][0] is None - assert result["occupancy"][1] is None + (data_dir / "large.cif").write_text(large_content) + dataset = load_dataset("mmcif", data_dir=str(data_dir), split="train") + assert len(dataset) == 1 -def test_mmcif_batch_size(mmcif_file_large): - """Test batch size configuration.""" - # Use batch_size=100 to create multiple batches - mmcif = Mmcif(batch_size=100) - generator = mmcif._generate_tables([[mmcif_file_large]]) - tables = [table for _, table in generator] - - # Should have 5 batches (500 atoms / 100) - assert len(tables) == 5 - - # Each batch should have 100 records - for table in tables: - assert table.num_rows == 100 - - -def test_mmcif_schema_types(mmcif_file): - """Test that schema uses correct Arrow types.""" - mmcif = Mmcif() - generator = mmcif._generate_tables([[mmcif_file]]) - pa_table = pa.concat_tables([table for _, table in generator]) - - schema = pa_table.schema - - # Check numeric types - assert schema.field("id").type == pa.int32() - assert schema.field("label_seq_id").type == pa.int32() - assert schema.field("Cartn_x").type == pa.float32() - assert schema.field("Cartn_y").type == pa.float32() - assert schema.field("Cartn_z").type == pa.float32() - assert schema.field("occupancy").type == pa.float32() - assert schema.field("B_iso_or_equiv").type == pa.float32() - - # Check string types - assert schema.field("type_symbol").type == pa.string() - assert schema.field("label_atom_id").type == pa.string() - assert schema.field("label_comp_id").type == pa.string() - - -def test_mmcif_feature_casting(mmcif_file): - """Test feature casting to custom schema.""" - features = Features({ - "id": Value("int32"), - "type_symbol": Value("string"), - "Cartn_x": Value("float32"), - "Cartn_y": Value("float32"), - "Cartn_z": Value("float32"), - }) - mmcif = Mmcif(features=features, columns=["id", "type_symbol", "Cartn_x", "Cartn_y", "Cartn_z"]) - generator = mmcif._generate_tables([[mmcif_file]]) - pa_table = pa.concat_tables([table for _, table in generator]) - - assert pa_table.schema.field("id").type == pa.int32() - assert pa_table.schema.field("Cartn_x").type == pa.float32() - - -def test_mmcif_empty_file(tmp_path): - """Test handling of empty mmCIF file.""" - filename = tmp_path / "empty.cif" - with open(filename, "w", encoding="utf-8") as f: - f.write("") - - mmcif = Mmcif() - generator = mmcif._generate_tables([[str(filename)]]) - tables = list(generator) - - # Empty file should produce no tables - assert len(tables) == 0 - - -def test_mmcif_no_atom_site(tmp_path): - """Test mmCIF file without atom_site category.""" - filename = tmp_path / "no_atoms.cif" - data = textwrap.dedent( - """\ - data_test - # - _entry.id test - _struct.title "Test structure" - # - """ - ) - with open(filename, "w", encoding="utf-8") as f: - f.write(data) - - mmcif = Mmcif() - generator = mmcif._generate_tables([[str(filename)]]) - tables = list(generator) - - # No atom_site should produce no tables - assert len(tables) == 0 - - -def test_mmcif_multiple_files(tmp_path): - """Test loading multiple mmCIF files.""" - # Create two files - file1 = tmp_path / "structure1.cif" - file2 = tmp_path / "structure2.cif" - - data1 = textwrap.dedent( - """\ - data_test1 - loop_ - _atom_site.id - _atom_site.type_symbol - _atom_site.label_atom_id - _atom_site.label_comp_id - _atom_site.label_asym_id - _atom_site.label_seq_id - _atom_site.Cartn_x - _atom_site.Cartn_y - _atom_site.Cartn_z - _atom_site.occupancy - _atom_site.B_iso_or_equiv - 1 N N ALA A 1 10.0 20.0 30.0 1.0 15.0 - """ - ) - - data2 = textwrap.dedent( - """\ - data_test2 - loop_ - _atom_site.id - _atom_site.type_symbol - _atom_site.label_atom_id - _atom_site.label_comp_id - _atom_site.label_asym_id - _atom_site.label_seq_id - _atom_site.Cartn_x - _atom_site.Cartn_y - _atom_site.Cartn_z - _atom_site.occupancy - _atom_site.B_iso_or_equiv - 1 C CA GLY B 1 11.0 21.0 31.0 1.0 14.0 - """ - ) - - with open(file1, "w", encoding="utf-8") as f: - f.write(data1) - with open(file2, "w", encoding="utf-8") as f: - f.write(data2) - - mmcif = Mmcif() - generator = mmcif._generate_tables([[str(file1)], [str(file2)]]) - pa_table = pa.concat_tables([table for _, table in generator]) - - result = pa_table.to_pydict() - - assert len(result["id"]) == 2 - assert result["label_comp_id"] == ["ALA", "GLY"] - - -def test_mmcif_extensions(): - """Test that correct extensions are defined.""" - assert ".cif" in Mmcif.EXTENSIONS - assert ".mmcif" in Mmcif.EXTENSIONS - - -def test_mmcif_comments_handling(tmp_path): - """Test that comments are properly skipped.""" - filename = tmp_path / "comments.cif" - data = textwrap.dedent( - """\ - data_test - # This is a comment - # - loop_ - _atom_site.id - _atom_site.type_symbol - _atom_site.label_atom_id - _atom_site.label_comp_id - _atom_site.label_asym_id - _atom_site.label_seq_id - _atom_site.Cartn_x - _atom_site.Cartn_y - _atom_site.Cartn_z - _atom_site.occupancy - _atom_site.B_iso_or_equiv - # Another comment - 1 N N ALA A 1 10.0 20.0 30.0 1.0 15.0 - # Comment between data - 2 C CA ALA A 1 11.0 21.0 31.0 1.0 14.0 - # - """ - ) - with open(filename, "w", encoding="utf-8") as f: - f.write(data) - - mmcif = Mmcif() - generator = mmcif._generate_tables([[str(filename)]]) - pa_table = pa.concat_tables([table for _, table in generator]) - - result = pa_table.to_pydict() - assert len(result["id"]) == 2 - - -def test_mmcif_default_columns(): - """Test that default columns are sensible.""" - from datasets.packaged_modules.mmcif.mmcif import DEFAULT_ATOM_SITE_COLUMNS - - # Essential columns should be in defaults - assert "id" in DEFAULT_ATOM_SITE_COLUMNS - assert "type_symbol" in DEFAULT_ATOM_SITE_COLUMNS - assert "label_atom_id" in DEFAULT_ATOM_SITE_COLUMNS - assert "label_comp_id" in DEFAULT_ATOM_SITE_COLUMNS - assert "Cartn_x" in DEFAULT_ATOM_SITE_COLUMNS - assert "Cartn_y" in DEFAULT_ATOM_SITE_COLUMNS - assert "Cartn_z" in DEFAULT_ATOM_SITE_COLUMNS + # Verify content is complete (decoded as string by default) + content = dataset[0]["structure"] + assert isinstance(content, str) + assert "ALA A 100" in content From 0ea9bb09e99b99bc434662460ee4bf0e7937ee5e Mon Sep 17 00:00:00 2001 From: Behrooz Date: Fri, 9 Jan 2026 11:04:22 -0800 Subject: [PATCH 3/5] style: apply ruff formatting fixes - Fix line length in protein_structure.py error messages - Sort imports alphabetically in __init__.py - Format function calls and f-strings in test_mmcif.py --- src/datasets/features/protein_structure.py | 8 ++++++-- src/datasets/packaged_modules/__init__.py | 2 +- tests/packaged_modules/test_mmcif.py | 12 ++++-------- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/datasets/features/protein_structure.py b/src/datasets/features/protein_structure.py index d33ebdd38f7..97b0495fa17 100644 --- a/src/datasets/features/protein_structure.py +++ b/src/datasets/features/protein_structure.py @@ -116,7 +116,9 @@ def decode_example(self, value: dict, token_per_repo_id=None) -> str: `str`: The structure file content as a string (PDB/mmCIF are text formats). """ if not self.decode: - raise RuntimeError("Decoding is disabled for this feature. Please use ProteinStructure(decode=True) instead.") + raise RuntimeError( + "Decoding is disabled for this feature. Please use ProteinStructure(decode=True) instead." + ) if token_per_repo_id is None: token_per_repo_id = {} @@ -124,7 +126,9 @@ def decode_example(self, value: dict, token_per_repo_id=None) -> str: path, bytes_ = value["path"], value["bytes"] if bytes_ is None: if path is None: - raise ValueError(f"A protein structure should have one of 'path' or 'bytes' but both are None in {value}.") + raise ValueError( + f"A protein structure should have one of 'path' or 'bytes' but both are None in {value}." + ) else: if is_local_path(path): with open(path, "r", encoding="utf-8") as f: diff --git a/src/datasets/packaged_modules/__init__.py b/src/datasets/packaged_modules/__init__.py index e9445ac9d55..1b2c897da1e 100644 --- a/src/datasets/packaged_modules/__init__.py +++ b/src/datasets/packaged_modules/__init__.py @@ -13,10 +13,10 @@ from .hdf5 import hdf5 from .iceberg import iceberg from .imagefolder import imagefolder -from .mmcif import mmcif from .json import json from .lance import lance from .meshfolder import meshfolder +from .mmcif import mmcif from .niftifolder import niftifolder from .pandas import pandas from .parquet import parquet diff --git a/tests/packaged_modules/test_mmcif.py b/tests/packaged_modules/test_mmcif.py index 4ea4b488c0e..8d4d4414c38 100644 --- a/tests/packaged_modules/test_mmcif.py +++ b/tests/packaged_modules/test_mmcif.py @@ -178,9 +178,7 @@ def test_load_with_labels(self, mmcif_data_dir_with_labels): def test_load_drop_labels(self, mmcif_data_dir_with_labels): """Test dropping labels.""" - dataset = load_dataset( - "mmcif", data_dir=mmcif_data_dir_with_labels, split="train", drop_labels=True - ) + dataset = load_dataset("mmcif", data_dir=mmcif_data_dir_with_labels, split="train", drop_labels=True) assert len(dataset) == 4 assert "structure" in dataset.column_names assert "label" not in dataset.column_names @@ -195,9 +193,7 @@ def test_load_with_metadata(self, mmcif_data_dir_with_metadata): def test_load_drop_metadata(self, mmcif_data_dir_with_metadata): """Test dropping metadata.""" - dataset = load_dataset( - "mmcif", data_dir=mmcif_data_dir_with_metadata, split="train", drop_metadata=True - ) + dataset = load_dataset("mmcif", data_dir=mmcif_data_dir_with_metadata, split="train", drop_metadata=True) assert len(dataset) == 3 assert "structure" in dataset.column_names assert "resolution" not in dataset.column_names @@ -352,12 +348,12 @@ def test_large_mmcif_content(self, tmp_path): # Create mmCIF with many atoms atoms = [] for i in range(100): - atoms.append(f"ATOM {i+1} N N ALA A {i+1} {i}.000 0.000 0.000 1.00 20.00") + atoms.append(f"ATOM {i + 1} N N ALA A {i + 1} {i}.000 0.000 0.000 1.00 20.00") large_content = MMCIF_CONTENT.replace( "ATOM 1 N N ALA A 1 0.000 0.000 0.000 1.00 20.00\n" "ATOM 2 C CA ALA A 1 1.458 0.000 0.000 1.00 20.00", - "\n".join(atoms) + "\n".join(atoms), ) (data_dir / "large.cif").write_text(large_content) From 84c9892b2f9bc20c0b939307f3c11aa7fd743507 Mon Sep 17 00:00:00 2001 From: Behrooz Date: Tue, 9 Jun 2026 22:02:54 -0700 Subject: [PATCH 4/5] fix(mmcif): align ProteinStructure.embed_storage with local_files/remote_files API The embed_storage API gained local_files and remote_files parameters (v4.8.5) to gate which paths are embedded into the Arrow array. ProteinStructure was written before this change and broke under the current embed_array_storage call path (TypeError: unexpected keyword argument 'local_files'). Align its signature and gating logic with the other binary features (Pdf/Image), embedding local paths only when local_files is set and remote URLs only when remote_files is set. --- src/datasets/features/protein_structure.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/src/datasets/features/protein_structure.py b/src/datasets/features/protein_structure.py index 97b0495fa17..9bc2b5393f9 100644 --- a/src/datasets/features/protein_structure.py +++ b/src/datasets/features/protein_structure.py @@ -8,7 +8,7 @@ from .. import config from ..download.download_config import DownloadConfig from ..table import array_cast -from ..utils.file_utils import is_local_path, xopen +from ..utils.file_utils import is_local_path, is_remote_url, xopen from ..utils.py_utils import no_op_if_value_is_null, string_to_dict @@ -205,7 +205,9 @@ def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArr storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null()) return array_cast(storage, self.pa_type) - def embed_storage(self, storage: pa.StructArray, token_per_repo_id=None) -> pa.StructArray: + def embed_storage( + self, storage: pa.StructArray, token_per_repo_id=None, local_files: bool = True, remote_files: bool = True + ) -> pa.StructArray: """Embed protein structure files into the Arrow array. Args: @@ -213,6 +215,11 @@ def embed_storage(self, storage: pa.StructArray, token_per_repo_id=None) -> pa.S PyArrow array to embed. token_per_repo_id (`dict`, *optional*): To access structure files from private repositories on the Hub. + local_files (`bool`, defaults to `True`) + Whether to embed local files data in the array + remote_files (`bool`, defaults to `True`) + Whether to embed remote files data in the array. + E.g. files with paths that start with hf:// or https:// Returns: `pa.StructArray`: Array in the ProteinStructure arrow storage type, that is @@ -235,7 +242,14 @@ def path_to_bytes(path): bytes_array = pa.array( [ - (path_to_bytes(x["path"]) if x["bytes"] is None else x["bytes"]) if x is not None else None + ( + path_to_bytes(x["path"]) + if x["bytes"] is None + and ((local_files and is_local_path(x["path"])) or (remote_files and is_remote_url(x["path"]))) + else x["bytes"] + ) + if x is not None + else None for x in storage.to_pylist() ], type=pa.binary(), From 945eb28217d90a9b562765b58090997762ef2d70 Mon Sep 17 00:00:00 2001 From: Behrooz Date: Tue, 16 Jun 2026 09:28:13 -0700 Subject: [PATCH 5/5] feat(mmcif): parse structures into atom-level columns Decoding a ProteinStructure now returns a struct-of-arrays of atom-level columns parsed from the mmCIF _atom_site loop, keyed by PDBx/mmCIF dictionary names (type_symbol, label_atom_id, label_comp_id, label_asym_id, label_seq_id, Cartn_x/y/z, occupancy, B_iso_or_equiv), while keeping one row per structure. PDB- and mmCIF-derived datasets share the same column vocabulary. The shared parser lives in features/protein_parsing.py (single canonical column/dtype table reused by both formats). ProteinStructure parses lazily on decode (mirroring the Image/Audio/Mesh features); decode=False still returns the raw {path, bytes} mapping for a custom parser. New feature/config options include_hetatm and columns are threaded through a _base_feature() factory on FolderBasedBuilder so folder loaders can forward config-derived arguments to their base feature. Docs updated to match the new API. --- docs/source/loading.mdx | 15 +- src/datasets/features/protein_parsing.py | 224 ++++++++++++++++++ src/datasets/features/protein_structure.py | 99 +++++--- .../folder_based_builder.py | 19 +- src/datasets/packaged_modules/mmcif/mmcif.py | 21 +- tests/features/test_protein_parsing.py | 117 +++++++++ tests/features/test_protein_structure.py | 28 ++- .../features/test_protein_structure_parsed.py | 93 ++++++++ tests/packaged_modules/test_mmcif.py | 66 ++++-- 9 files changed, 603 insertions(+), 79 deletions(-) create mode 100644 src/datasets/features/protein_parsing.py create mode 100644 tests/features/test_protein_parsing.py create mode 100644 tests/features/test_protein_structure_parsed.py diff --git a/docs/source/loading.mdx b/docs/source/loading.mdx index 567a163d426..e7625b1cc4a 100644 --- a/docs/source/loading.mdx +++ b/docs/source/loading.mdx @@ -271,14 +271,17 @@ To load remote WebDatasets via HTTP, pass the URLs instead: [mmCIF](https://mmcif.wwpdb.org/) (macromolecular Crystallographic Information File) is the modern standard format for representing 3D macromolecular structures, used by the Protein Data Bank (PDB) since 2014. -To load an mmCIF file: +Each row is one structure (one row = one structure), following the same folder layout +as [`ImageFolder`]: subfolder names become labels and a `metadata.csv`/`metadata.jsonl` +file adds per-structure metadata. To load mmCIF files from a directory: ```py >>> from datasets import load_dataset ->>> dataset = load_dataset("mmcif", data_files="structure.cif") +>>> dataset = load_dataset("mmcif", data_dir="path/to/structures") ``` -This returns the `_atom_site` category containing atomic coordinates. The default columns include: +The `structure` column is parsed from the `_atom_site` category into a struct of +atom-level arrays. The default columns are: | Column | Type | Description | |--------|------|-------------| @@ -297,16 +300,16 @@ This returns the `_atom_site` category containing atomic coordinates. The defaul You can select specific columns: ```py ->>> dataset = load_dataset("mmcif", data_files="structure.cif", columns=["label_atom_id", "Cartn_x", "Cartn_y", "Cartn_z"]) +>>> dataset = load_dataset("mmcif", data_dir="path/to/structures", columns=["label_atom_id", "Cartn_x", "Cartn_y", "Cartn_z"]) ``` To exclude HETATM records (ligands, water molecules): ```py ->>> dataset = load_dataset("mmcif", data_files="structure.cif", include_hetatm=False) +>>> dataset = load_dataset("mmcif", data_dir="path/to/structures", include_hetatm=False) ``` -Compressed files (`.cif.gz`) are automatically detected and decompressed. +Extensions `.cif` and `.mmcif` are supported. ## Remote files diff --git a/src/datasets/features/protein_parsing.py b/src/datasets/features/protein_parsing.py new file mode 100644 index 00000000000..ffbabb4d7be --- /dev/null +++ b/src/datasets/features/protein_parsing.py @@ -0,0 +1,224 @@ +"""Dependency-free parsers for protein 3D structure files (PDB and mmCIF). + +Both formats are parsed into a *struct-of-arrays* keyed by +[PDBx/mmCIF dictionary](https://mmcif.wwpdb.org/) ``_atom_site`` column names, so +that PDB- and mmCIF-derived datasets expose the same column vocabulary. PDB +fixed-width records are mapped onto the mmCIF-native names; mmCIF ``_atom_site`` +loops are read directly. + +The schema (column names and Arrow dtypes) lives in a single mapping, +[`PROTEIN_ATOM_TYPES`], reused by both parsers and by the +[`~datasets.ProteinStructure`] feature — there is no per-format duplication of +the column list or its types. +""" + +from __future__ import annotations + +import pyarrow as pa + + +# Canonical column -> Arrow dtype table (PDBx/mmCIF `_atom_site` dictionary names). +# Single source of truth for both parsers and the ProteinStructure feature. +PROTEIN_ATOM_TYPES: dict[str, pa.DataType] = { + "group_PDB": pa.string(), # "ATOM" / "HETATM"; drives include_hetatm, not emitted by default + "id": pa.int32(), + "type_symbol": pa.string(), + "label_atom_id": pa.string(), + "label_comp_id": pa.string(), + "label_asym_id": pa.string(), + "label_seq_id": pa.int32(), + "Cartn_x": pa.float32(), + "Cartn_y": pa.float32(), + "Cartn_z": pa.float32(), + "occupancy": pa.float32(), + "B_iso_or_equiv": pa.float32(), +} + +# Default emitted columns (everything except the internal record-group flag). +DEFAULT_ATOM_COLUMNS: list[str] = [col for col in PROTEIN_ATOM_TYPES if col != "group_PDB"] + +# PDB fixed-width column ranges (0-indexed, end-exclusive), per the wwPDB +# format-33 spec, expressed against the canonical mmCIF column names. +# https://www.wwpdb.org/documentation/file-format-content/format33/sect9.html +PDB_FIELD_SPANS: dict[str, tuple[int, int]] = { + "group_PDB": (0, 6), + "id": (6, 11), + "label_atom_id": (12, 16), + "label_comp_id": (17, 20), + "label_asym_id": (21, 22), + "label_seq_id": (22, 26), + "Cartn_x": (30, 38), + "Cartn_y": (38, 46), + "Cartn_z": (46, 54), + "occupancy": (54, 60), + "B_iso_or_equiv": (60, 66), + "type_symbol": (76, 78), +} + +_MMCIF_RECORD_GROUP = "group_PDB" +_ATOM_RECORD = "ATOM" + + +def _convert(value: str, dtype: pa.DataType): + """Coerce a raw string token to the Arrow dtype's Python type, or None if empty/invalid. + + mmCIF uses ``.`` and ``?`` for missing values; both map to None. + """ + if value is None or value == "" or value in (".", "?"): + return None + if pa.types.is_integer(dtype): + try: + return int(value) + except ValueError: + return None + if pa.types.is_floating(dtype): + try: + return float(value) + except ValueError: + return None + return value + + +def _resolve_columns(columns: list[str] | None) -> list[str]: + """Validate a requested column subset against the canonical schema. + + Returns the default columns when ``columns`` is None. Raises ValueError on any + unknown name rather than silently dropping it. + """ + if columns is None: + return list(DEFAULT_ATOM_COLUMNS) + unknown = [c for c in columns if c not in PROTEIN_ATOM_TYPES] + if unknown: + raise ValueError(f"Unknown protein atom column(s) {unknown}. Valid columns are: {sorted(PROTEIN_ATOM_TYPES)}") + return list(columns) + + +def parse_pdb_atoms(text: str, columns: list[str] | None = None, include_hetatm: bool = True) -> dict[str, list]: + """Parse PDB ATOM/HETATM records into a struct-of-arrays. + + Args: + text (`str`): + Raw PDB file content. + columns (`list[str]`, *optional*): + Subset of [`PROTEIN_ATOM_TYPES`] columns to return. Defaults to + [`DEFAULT_ATOM_COLUMNS`]. + include_hetatm (`bool`, defaults to `True`): + Whether to include HETATM records (ligands, water, …). When `False`, + only `ATOM` records are kept. + + Returns: + `dict[str, list]`: Mapping of column name to a list of per-atom values. + """ + requested = _resolve_columns(columns) + atoms: dict[str, list] = {col: [] for col in requested} + + for line in text.splitlines(): + record = line[slice(*PDB_FIELD_SPANS[_MMCIF_RECORD_GROUP])].strip() + if record not in ("ATOM", "HETATM"): + continue + if not include_hetatm and record != _ATOM_RECORD: + continue + for col in requested: + start, end = PDB_FIELD_SPANS[col] + raw = line[start:end].strip() if start < len(line) else "" + atoms[col].append(_convert(raw, PROTEIN_ATOM_TYPES[col])) + + return atoms + + +def _tokenize_cif_line(line: str) -> list[str]: + """Tokenize a single inline mmCIF data line, honoring single/double quotes.""" + tokens: list[str] = [] + pos = 0 + n = len(line) + while pos < n: + while pos < n and line[pos] in " \t": + pos += 1 + if pos >= n: + break + char = line[pos] + if char in "'\"": + end = pos + 1 + while end < n and line[end] != char: + end += 1 + tokens.append(line[pos + 1 : end]) + pos = end + 1 + else: + end = pos + while end < n and line[end] not in " \t": + end += 1 + tokens.append(line[pos:end]) + pos = end + return tokens + + +def parse_mmcif_atoms(text: str, columns: list[str] | None = None, include_hetatm: bool = True) -> dict[str, list]: + """Parse an mmCIF ``_atom_site`` loop into a struct-of-arrays. + + Each ``_atom_site`` row is one whitespace-separated line of simple tokens, as + produced by the wwPDB; the multi-line semicolon-delimited text values that the + CIF grammar allows for free-text categories do not occur in coordinate loops, + so rows are tokenized one line at a time. + + Args: + text (`str`): + Raw mmCIF file content. + columns (`list[str]`, *optional*): + Subset of [`PROTEIN_ATOM_TYPES`] columns to return. Defaults to + [`DEFAULT_ATOM_COLUMNS`]. + include_hetatm (`bool`, defaults to `True`): + Whether to include HETATM records. When `False`, only `_atom_site` + rows whose `group_PDB` is `ATOM` are kept. + + Returns: + `dict[str, list]`: Mapping of column name to a list of per-atom values. + """ + requested = _resolve_columns(columns) + lines = text.splitlines() + + # Locate the `_atom_site` loop and collect its column order. + loop_columns: list[str] = [] + idx = 0 + n = len(lines) + while idx < n: + if lines[idx].strip() == "loop_": + peek = idx + 1 + header: list[str] = [] + while peek < n and lines[peek].strip().startswith("_"): + header.append(lines[peek].strip()) + peek += 1 + if header and header[0].startswith("_atom_site."): + loop_columns = [h.split(".", 1)[1] for h in header] + idx = peek + break + idx = peek + else: + idx += 1 + + atoms: dict[str, list] = {col: [] for col in requested} + if not loop_columns: + return atoms + + col_index = {name: i for i, name in enumerate(loop_columns)} + group_pos = col_index.get(_MMCIF_RECORD_GROUP) + + # Consume data rows until the loop ends (blank line, comment, or new block/loop). + while idx < n: + stripped = lines[idx].strip() + if stripped == "" or stripped.startswith("#"): + idx += 1 + continue + if stripped.startswith("_") or stripped.startswith("loop_") or stripped.startswith("data_"): + break + tokens = _tokenize_cif_line(lines[idx]) + idx += 1 + if len(tokens) < len(loop_columns): + continue + if not include_hetatm and group_pos is not None and tokens[group_pos] != _ATOM_RECORD: + continue + for col in requested: + pos = col_index.get(col) + raw = tokens[pos] if pos is not None and pos < len(tokens) else "" + atoms[col].append(_convert(raw, PROTEIN_ATOM_TYPES[col])) + + return atoms diff --git a/src/datasets/features/protein_structure.py b/src/datasets/features/protein_structure.py index 9bc2b5393f9..96143af3ff2 100644 --- a/src/datasets/features/protein_structure.py +++ b/src/datasets/features/protein_structure.py @@ -22,8 +22,12 @@ class ProteinStructure: **Experimental.** ProteinStructure [`Feature`] to read protein 3D structure files (PDB/mmCIF). - This feature stores raw file content for lazy loading. The viewer or downstream - code handles parsing. Supports both PDB (.pdb, .ent) and mmCIF (.cif, .mmcif) formats. + The raw file content is stored for lazy loading and, on access, parsed into a + struct-of-arrays of atom-level columns keyed by + [PDBx/mmCIF dictionary](https://mmcif.wwpdb.org/) ``_atom_site`` names, so PDB + and mmCIF datasets expose the same column vocabulary. Supports both PDB + (.pdb, .ent) and mmCIF (.cif, .mmcif) formats. The format is inferred from the + file extension. Input: The ProteinStructure feature accepts as input: - A `str`: Absolute path to the structure file (i.e. random access is allowed). @@ -36,18 +40,23 @@ class ProteinStructure: Args: decode (`bool`, defaults to `True`): - Whether to decode the structure data to a string. If `False`, - returns the underlying dictionary in the format `{"path": structure_path, "bytes": structure_bytes}`. + Whether to parse the structure into a struct-of-arrays of atom columns. + If `False`, returns the underlying dictionary in the format + `{"path": structure_path, "bytes": structure_bytes}` for a custom parser. + include_hetatm (`bool`, defaults to `True`): + Whether to include HETATM records (ligands, water, …) when parsing. + columns (`list[str]`, *optional*): + Subset of atom columns to return. Defaults to all default columns + (see [`~datasets.features.protein_parsing.DEFAULT_ATOM_COLUMNS`]). Examples: ```py >>> from datasets import Dataset, ProteinStructure >>> ds = Dataset.from_dict({"structure": ["path/to/structure.pdb"]}).cast_column("structure", ProteinStructure()) - >>> ds.features["structure"] - ProteinStructure(decode=True, id=None) >>> ds[0]["structure"] - 'HEADER PLANT PROTEIN...\\nATOM 1 N THR A 1...' + {'type_symbol': ['N', 'C', ...], 'label_atom_id': ['N', 'CA', ...], 'label_comp_id': ['ALA', 'ALA', ...], + 'Cartn_x': [0.0, 1.458, ...], 'Cartn_y': [...], 'Cartn_z': [...], ...} >>> ds = ds.cast_column("structure", ProteinStructure(decode=False)) >>> ds[0]["structure"] {'bytes': None, @@ -56,10 +65,12 @@ class ProteinStructure: """ decode: bool = True + include_hetatm: bool = True + columns: Optional[list[str]] = None id: Optional[str] = field(default=None, repr=False) # Automatically constructed - dtype: ClassVar[str] = "str" + dtype: ClassVar[str] = "dict" pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()}) _type: str = field(default="ProteinStructure", init=False, repr=False) @@ -99,8 +110,14 @@ def encode_example(self, value: Union[str, bytes, bytearray, dict, Path]) -> dic f"Expected str, bytes, pathlib.Path, or dict." ) - def decode_example(self, value: dict, token_per_repo_id=None) -> str: - """Decode example structure file into structure data. + def decode_example(self, value: dict, token_per_repo_id=None) -> dict: + """Decode example structure file into a parsed struct-of-arrays. + + The raw PDB/mmCIF text is parsed into atom-level columns keyed by + PDBx/mmCIF dictionary names (see [`~datasets.features.protein_parsing.PROTEIN_ATOM_TYPES`]), + so PDB- and mmCIF-derived datasets share one column vocabulary. The format + is inferred from the file extension in `path`. Use `ProteinStructure(decode=False)` + to keep the raw `{path, bytes}` mapping for a custom parser. Args: value (`dict`): @@ -113,13 +130,30 @@ def decode_example(self, value: dict, token_per_repo_id=None) -> str: the Hub, you can pass a dictionary repo_id (`str`) -> token (`bool` or `str`). Returns: - `str`: The structure file content as a string (PDB/mmCIF are text formats). + `dict[str, list]`: Mapping of atom column name to a list of per-atom values. """ if not self.decode: raise RuntimeError( "Decoding is disabled for this feature. Please use ProteinStructure(decode=True) instead." ) + text = self._read_text(value, token_per_repo_id) + parser = self._parser_for_path(value["path"]) + return parser(text, columns=self.columns, include_hetatm=self.include_hetatm) + + @staticmethod + def _parser_for_path(path: Optional[str]): + """Select the format parser from the file extension (defaults to PDB).""" + from .protein_parsing import parse_mmcif_atoms, parse_pdb_atoms + + suffix = os.path.splitext(path)[1].lower() if path else "" + if suffix in (".cif", ".mmcif"): + return parse_mmcif_atoms + return parse_pdb_atoms + + @staticmethod + def _read_text(value: dict, token_per_repo_id=None) -> str: + """Read the raw structure text from a `{path, bytes}` mapping.""" if token_per_repo_id is None: token_per_repo_id = {} @@ -129,31 +163,24 @@ def decode_example(self, value: dict, token_per_repo_id=None) -> str: raise ValueError( f"A protein structure should have one of 'path' or 'bytes' but both are None in {value}." ) - else: - if is_local_path(path): - with open(path, "r", encoding="utf-8") as f: - return f.read() - else: - source_url = path.split("::")[-1] - pattern = ( - config.HUB_DATASETS_URL - if source_url.startswith(config.HF_ENDPOINT) - else config.HUB_DATASETS_HFFS_URL - ) - try: - repo_id = string_to_dict(source_url, pattern)["repo_id"] - token = token_per_repo_id.get(repo_id) - except ValueError: - token = None - download_config = DownloadConfig(token=token) - with xopen(path, "r", download_config=download_config) as f: - return f.read() - else: - # PDB/mmCIF are text formats, decode bytes to string - if isinstance(bytes_, (bytes, bytearray)): - return bytes_.decode("utf-8") - # If it's already a string, return as-is - return str(bytes_) + if is_local_path(path): + with open(path, "r", encoding="utf-8") as f: + return f.read() + source_url = path.split("::")[-1] + pattern = ( + config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL + ) + try: + repo_id = string_to_dict(source_url, pattern)["repo_id"] + token = token_per_repo_id.get(repo_id) + except ValueError: + token = None + download_config = DownloadConfig(token=token) + with xopen(path, "r", download_config=download_config) as f: + return f.read() + if isinstance(bytes_, (bytes, bytearray)): + return bytes_.decode("utf-8") + return str(bytes_) def flatten(self) -> Union["FeatureType", Dict[str, "FeatureType"]]: """If in the decodable state, return the feature itself, otherwise flatten the feature into a dictionary.""" diff --git a/src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py b/src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py index 82586c648cb..388f0cd0cb0 100644 --- a/src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py +++ b/src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py @@ -59,6 +59,15 @@ class FolderBasedBuilder(datasets.GeneratorBasedBuilder): METADATA_FILENAMES: list[str] = ["metadata.csv", "metadata.jsonl", "metadata.parquet"] + def _base_feature(self) -> FeatureType: + """Instantiate the base feature for one column. + + Subclasses override this to forward config-derived options to the feature + (e.g. a structure parser's filtering or column selection). The default + constructs the feature with no arguments. + """ + return self.BASE_FEATURE() + def _info(self): if not self.config.data_dir and not self.config.data_files: raise ValueError( @@ -219,21 +228,21 @@ def _set_feature(feature): feature[key] == datasets.Value("string") or feature[key] == datasets.Value("large_string") ): key = key[: -len("_file_name")] or self.BASE_COLUMN_NAME - out[key] = self.BASE_FEATURE() + out[key] = self._base_feature() feature_not_found = False elif (key == "file_names" or key.endswith("_file_names")) and ( feature[key] in [datasets.List(datasets.Value("string")), datasets.List(datasets.Value("large_string"))] ): key = key[: -len("_file_names")] or (self.BASE_COLUMN_NAME + "s") - out[key] = datasets.List(self.BASE_FEATURE()) + out[key] = datasets.List(self._base_feature()) feature_not_found = False elif (key == "file_names" or key.endswith("_file_names")) and ( feature[key] == [datasets.Value("string")] or feature[key] == [datasets.Value("large_string")] ): key = key[: -len("_file_names")] or (self.BASE_COLUMN_NAME + "s") - out[key] = [self.BASE_FEATURE()] + out[key] = [self._base_feature()] feature_not_found = False else: out[key] = feature[key] @@ -257,12 +266,12 @@ def _set_feature(feature): elif add_labels: self.info.features = datasets.Features( { - self.BASE_COLUMN_NAME: self.BASE_FEATURE(), + self.BASE_COLUMN_NAME: self._base_feature(), "label": datasets.ClassLabel(names=sorted(labels)), } ) else: - self.info.features = datasets.Features({self.BASE_COLUMN_NAME: self.BASE_FEATURE()}) + self.info.features = datasets.Features({self.BASE_COLUMN_NAME: self._base_feature()}) return splits diff --git a/src/datasets/packaged_modules/mmcif/mmcif.py b/src/datasets/packaged_modules/mmcif/mmcif.py index 0c19f80c87c..9f186fdf5ba 100644 --- a/src/datasets/packaged_modules/mmcif/mmcif.py +++ b/src/datasets/packaged_modules/mmcif/mmcif.py @@ -22,6 +22,8 @@ 2def.mmcif """ +from dataclasses import dataclass + import datasets from ..folder_based_builder import folder_based_builder @@ -30,11 +32,25 @@ logger = datasets.utils.logging.get_logger(__name__) +@dataclass class MmcifFolderConfig(folder_based_builder.FolderBasedBuilderConfig): - """BuilderConfig for MmcifFolder.""" + """BuilderConfig for MmcifFolder. + + Args: + drop_labels (`bool`, *optional*): + Whether to drop folder-name labels. + drop_metadata (`bool`, *optional*): + Whether to drop metadata columns. + include_hetatm (`bool`, defaults to `True`): + Whether to include HETATM records (ligands, water, …) when decoding each structure. + columns (`list[str]`, *optional*): + Subset of PDBx/mmCIF atom columns to return per structure. Defaults to all columns. + """ drop_labels: bool = None drop_metadata: bool = None + include_hetatm: bool = True + columns: list[str] = None def __post_init__(self): super().__post_init__() @@ -52,3 +68,6 @@ class MmcifFolder(folder_based_builder.FolderBasedBuilder): BASE_COLUMN_NAME = "structure" BUILDER_CONFIG_CLASS = MmcifFolderConfig EXTENSIONS: list[str] = [".cif", ".mmcif"] + + def _base_feature(self): + return self.BASE_FEATURE(include_hetatm=self.config.include_hetatm, columns=self.config.columns) diff --git a/tests/features/test_protein_parsing.py b/tests/features/test_protein_parsing.py new file mode 100644 index 00000000000..fd0215f3d51 --- /dev/null +++ b/tests/features/test_protein_parsing.py @@ -0,0 +1,117 @@ +"""Tests for the dependency-free protein structure parsers (PDB + mmCIF). + +These cover the shared, format-agnostic parsing layer that turns raw PDB/mmCIF +text into a struct-of-arrays keyed by mmCIF/PDBx-dictionary column names. The +parsing layer is reused by the ProteinStructure feature for both formats so that +PDB- and mmCIF-derived datasets expose the same column vocabulary. +""" + +import pyarrow as pa +import pytest + +from datasets.features.protein_parsing import ( + DEFAULT_ATOM_COLUMNS, + PROTEIN_ATOM_TYPES, + parse_mmcif_atoms, + parse_pdb_atoms, +) + + +@pytest.fixture +def pdb_text(shared_datadir): + return (shared_datadir / "test_protein.pdb").read_text() + + +@pytest.fixture +def cif_text(shared_datadir): + return (shared_datadir / "test_protein.cif").read_text() + + +def test_default_columns_are_a_subset_of_known_schema(): + # No hard-coded duplication: defaults must be drawn from the canonical type table. + assert set(DEFAULT_ATOM_COLUMNS).issubset(set(PROTEIN_ATOM_TYPES)) + # The canonical coordinate + identity columns are present. + for required in ("label_atom_id", "type_symbol", "label_comp_id", "Cartn_x", "Cartn_y", "Cartn_z"): + assert required in DEFAULT_ATOM_COLUMNS + + +def test_parse_pdb_atoms_returns_struct_of_arrays(pdb_text): + atoms = parse_pdb_atoms(pdb_text) + # All default columns present, each a list of equal length (one entry per atom). + assert set(DEFAULT_ATOM_COLUMNS).issubset(set(atoms)) + lengths = {len(atoms[col]) for col in DEFAULT_ATOM_COLUMNS} + assert lengths == {9} # sample has 9 ATOM records (ALA x5, GLY x4) + + +def test_parse_pdb_atoms_values_match_sample(pdb_text): + atoms = parse_pdb_atoms(pdb_text) + # First atom: ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 20.00 N + assert atoms["label_atom_id"][0] == "N" + assert atoms["type_symbol"][0] == "N" + assert atoms["label_comp_id"][0] == "ALA" + assert atoms["label_asym_id"][0] == "A" + assert atoms["label_seq_id"][0] == 1 + assert atoms["Cartn_x"][0] == pytest.approx(0.0) + # Fifth atom CB has the only negative coords: 1.986 -0.760 -1.217 + assert atoms["Cartn_y"][4] == pytest.approx(-0.760) + assert atoms["Cartn_z"][4] == pytest.approx(-1.217) + # Residue transition ALA(1) -> GLY(2) at index 5. + assert atoms["label_comp_id"][5] == "GLY" + assert atoms["label_seq_id"][5] == 2 + # Occupancy / B-factor uniform in the sample. + assert all(v == pytest.approx(1.0) for v in atoms["occupancy"]) + assert all(v == pytest.approx(20.0) for v in atoms["B_iso_or_equiv"]) + + +def test_parse_mmcif_atoms_values_match_sample(cif_text): + atoms = parse_mmcif_atoms(cif_text) + assert len(atoms["label_atom_id"]) == 9 + assert atoms["label_atom_id"][0] == "N" + assert atoms["type_symbol"][0] == "N" + assert atoms["label_comp_id"][0] == "ALA" + assert atoms["label_asym_id"][0] == "A" + assert atoms["label_seq_id"][0] == 1 + assert atoms["Cartn_z"][4] == pytest.approx(-1.217) + assert atoms["label_comp_id"][5] == "GLY" + + +def test_pdb_and_mmcif_yield_same_schema(pdb_text, cif_text): + # The whole point of mmCIF-native naming: both formats expose identical columns. + pdb_atoms = parse_pdb_atoms(pdb_text) + cif_atoms = parse_mmcif_atoms(cif_text) + assert set(DEFAULT_ATOM_COLUMNS).issubset(set(pdb_atoms)) + assert set(DEFAULT_ATOM_COLUMNS).issubset(set(cif_atoms)) + for col in DEFAULT_ATOM_COLUMNS: + assert pdb_atoms[col] == cif_atoms[col], f"column {col} differs between PDB and mmCIF" + + +def test_include_hetatm_filter_drops_hetatm(): + # Generalizable filter: a HETATM line must be excluded when include_hetatm=False. + text = ( + "ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 20.00 N\n" + "HETATM 2 O HOH A 2 1.000 1.000 1.000 1.00 20.00 O\n" + "END\n" + ) + kept = parse_pdb_atoms(text, include_hetatm=False) + dropped = parse_pdb_atoms(text, include_hetatm=True) + assert len(kept["label_atom_id"]) == 1 + assert kept["label_comp_id"][0] == "ALA" + assert len(dropped["label_atom_id"]) == 2 + + +def test_columns_subset_selection(pdb_text): + # Requesting a subset returns exactly those columns (validated against schema). + atoms = parse_pdb_atoms(pdb_text, columns=["label_atom_id", "Cartn_x"]) + assert set(atoms) == {"label_atom_id", "Cartn_x"} + + +def test_unknown_column_raises(pdb_text): + # No silent drop of an unknown column name. + with pytest.raises(ValueError): + parse_pdb_atoms(pdb_text, columns=["not_a_real_column"]) + + +def test_protein_atom_types_dtypes_are_arrow_types(): + # The canonical type table maps every column to a pyarrow DataType (no stringly-typed dtypes). + for col, dtype in PROTEIN_ATOM_TYPES.items(): + assert isinstance(dtype, pa.DataType), f"{col} dtype is not a pa.DataType" diff --git a/tests/features/test_protein_structure.py b/tests/features/test_protein_structure.py index 34ba6d88dc7..fa18e45e0e8 100644 --- a/tests/features/test_protein_structure.py +++ b/tests/features/test_protein_structure.py @@ -26,8 +26,10 @@ def test_protein_structure_feature_encode_example(shared_datadir, build_example) assert isinstance(encoded_example, dict) assert encoded_example.keys() == {"bytes", "path"} assert encoded_example["bytes"] is not None or encoded_example["path"] is not None + # decode parses into a struct-of-arrays; a path-less bytes example defaults to the PDB parser. decoded_example = protein_structure.decode_example(encoded_example) - assert isinstance(decoded_example, str) + assert isinstance(decoded_example, dict) + assert "label_atom_id" in decoded_example def test_protein_structure_decode_example_pdb(shared_datadir): @@ -35,9 +37,9 @@ def test_protein_structure_decode_example_pdb(shared_datadir): protein_structure = ProteinStructure() decoded_example = protein_structure.decode_example({"path": structure_path, "bytes": None}) - assert isinstance(decoded_example, str) - assert "ATOM" in decoded_example - assert "HEADER" in decoded_example + assert isinstance(decoded_example, dict) + assert len(decoded_example["label_atom_id"]) == 9 + assert decoded_example["label_comp_id"][0] == "ALA" with pytest.raises(RuntimeError): ProteinStructure(decode=False).decode_example({"path": structure_path, "bytes": None}) @@ -48,9 +50,9 @@ def test_protein_structure_decode_example_cif(shared_datadir): protein_structure = ProteinStructure() decoded_example = protein_structure.decode_example({"path": structure_path, "bytes": None}) - assert isinstance(decoded_example, str) - assert "data_" in decoded_example - assert "_atom_site" in decoded_example + assert isinstance(decoded_example, dict) + assert len(decoded_example["label_atom_id"]) == 9 + assert decoded_example["type_symbol"][0] == "N" def test_dataset_with_protein_structure_feature(shared_datadir): @@ -60,23 +62,23 @@ def test_dataset_with_protein_structure_feature(shared_datadir): dset = Dataset.from_dict(data, features=features) item = dset[0] assert item.keys() == {"structure"} - assert isinstance(item["structure"], str) - assert "ATOM" in item["structure"] + assert isinstance(item["structure"], dict) + assert len(item["structure"]["label_atom_id"]) == 9 batch = dset[:1] assert len(batch) == 1 assert batch.keys() == {"structure"} - assert isinstance(batch["structure"], list) and all(isinstance(item, str) for item in batch["structure"]) + assert isinstance(batch["structure"], list) and all(isinstance(s, dict) for s in batch["structure"]) column = dset["structure"] assert len(column) == 1 - assert all(isinstance(item, str) for item in column) + assert all(isinstance(s, dict) for s in column) - # from bytes + # from bytes (path-less -> PDB parser by default) with open(structure_path, "rb") as f: data = {"structure": [f.read()]} dset = Dataset.from_dict(data, features=features) item = dset[0] assert item.keys() == {"structure"} - assert isinstance(item["structure"], str) + assert isinstance(item["structure"], dict) def test_protein_structure_pa_type(): diff --git a/tests/features/test_protein_structure_parsed.py b/tests/features/test_protein_structure_parsed.py new file mode 100644 index 00000000000..49cd01d6be7 --- /dev/null +++ b/tests/features/test_protein_structure_parsed.py @@ -0,0 +1,93 @@ +"""Tests for the parsed-object behavior of the ProteinStructure feature. + +With ``decode=True`` (the default), ProteinStructure now returns a parsed +struct-of-arrays (atom-level columns under PDBx/mmCIF names) instead of the raw +file text. ``decode=False`` still returns the raw ``{path, bytes}`` mapping, and +the raw text remains reachable for callers that want their own parser. +""" + +import pyarrow as pa +import pytest + +from datasets import Dataset, Features, ProteinStructure +from datasets.features.protein_parsing import DEFAULT_ATOM_COLUMNS + + +def test_decode_pdb_returns_parsed_struct_of_arrays(shared_datadir): + path = str(shared_datadir / "test_protein.pdb") + feature = ProteinStructure() + decoded = feature.decode_example({"path": path, "bytes": None}) + assert isinstance(decoded, dict) + assert set(DEFAULT_ATOM_COLUMNS).issubset(set(decoded)) + assert len(decoded["label_atom_id"]) == 9 + assert decoded["label_comp_id"][0] == "ALA" + assert decoded["label_comp_id"][5] == "GLY" + + +def test_decode_cif_returns_parsed_struct_of_arrays(shared_datadir): + path = str(shared_datadir / "test_protein.cif") + feature = ProteinStructure() + decoded = feature.decode_example({"path": path, "bytes": None}) + assert isinstance(decoded, dict) + assert len(decoded["label_atom_id"]) == 9 + assert decoded["type_symbol"][0] == "N" + + +def test_decode_false_returns_raw_mapping(shared_datadir): + path = str(shared_datadir / "test_protein.pdb") + feature = ProteinStructure(decode=False) + with pytest.raises(RuntimeError): + feature.decode_example({"path": path, "bytes": None}) + + +def test_include_hetatm_feature_arg_filters(shared_datadir, tmp_path): + structure = tmp_path / "het.pdb" + structure.write_text( + "ATOM 1 N ALA A 1 0.000 0.000 0.000 1.00 20.00 N\n" + "HETATM 2 O HOH A 2 1.000 1.000 1.000 1.00 20.00 O\n" + "END\n" + ) + keep = ProteinStructure(include_hetatm=True).decode_example({"path": str(structure), "bytes": None}) + drop = ProteinStructure(include_hetatm=False).decode_example({"path": str(structure), "bytes": None}) + assert len(keep["label_atom_id"]) == 2 + assert len(drop["label_atom_id"]) == 1 + assert drop["label_comp_id"][0] == "ALA" + + +def test_columns_feature_arg_subsets(shared_datadir): + path = str(shared_datadir / "test_protein.pdb") + feature = ProteinStructure(columns=["label_atom_id", "Cartn_x"]) + decoded = feature.decode_example({"path": path, "bytes": None}) + assert set(decoded) == {"label_atom_id", "Cartn_x"} + + +def test_dataset_roundtrip_one_row_per_structure(shared_datadir): + path = str(shared_datadir / "test_protein.pdb") + features = Features({"structure": ProteinStructure()}) + # Three independent structures -> three independent rows (one row = one structure). + dset = Dataset.from_dict({"structure": [path, path, path]}, features=features) + assert len(dset) == 3 + row = dset[0] + assert isinstance(row["structure"], dict) + assert len(row["structure"]["label_atom_id"]) == 9 + # Shuffling / splitting must keep rows independent. + shuffled = dset.shuffle(seed=0) + assert len(shuffled) == 3 + + +def test_pa_type_storage_is_unchanged(): + # Storage stays {bytes, path}: parsing is lazy on decode, no Arrow schema migration. + feature = ProteinStructure() + pa_type = feature.pa_type + assert pa.types.is_struct(pa_type) + assert pa_type.get_field_index("bytes") >= 0 + assert pa_type.get_field_index("path") >= 0 + + +def test_include_hetatm_and_columns_survive_serialization(): + # Feature args must round-trip through the feature dict (no silent loss on save/load). + feature = ProteinStructure(include_hetatm=False, columns=["label_atom_id"]) + restored = ProteinStructure.from_dict(feature.to_dict()) if hasattr(ProteinStructure, "from_dict") else None + if restored is not None: + assert restored.include_hetatm is False + assert restored.columns == ["label_atom_id"] diff --git a/tests/packaged_modules/test_mmcif.py b/tests/packaged_modules/test_mmcif.py index 8d4d4414c38..b22161becf3 100644 --- a/tests/packaged_modules/test_mmcif.py +++ b/tests/packaged_modules/test_mmcif.py @@ -209,15 +209,18 @@ def test_feature_type(self, mmcif_file): assert isinstance(dataset.features["structure"], ProteinStructure) def test_structure_content_decoded(self, mmcif_file): - """Test structure content is accessible when decoded (default).""" + """Test structure content is parsed into atom-level arrays when decoded (default).""" dataset = load_dataset("mmcif", data_files=mmcif_file, split="train") sample = dataset[0] structure = sample["structure"] - # With decode=True (default), structure is a string - assert isinstance(structure, str) - assert "data_TEST" in structure - assert "_atom_site" in structure + # With decode=True (default), structure is a parsed struct-of-arrays + assert isinstance(structure, dict) + assert len(structure["label_atom_id"]) == 2 # sample has 2 _atom_site rows + assert structure["label_atom_id"][0] == "N" + assert structure["label_atom_id"][1] == "CA" + assert structure["label_comp_id"][0] == "ALA" + assert structure["Cartn_x"][1] == pytest.approx(1.458) def test_structure_content_undecoded(self, mmcif_file): """Test structure content when decode=False.""" @@ -246,6 +249,37 @@ def test_arrow_type(self, mmcif_file): assert "path" in field_names +@pytest.fixture +def mmcif_file_with_hetatm(tmp_path): + """mmCIF file whose _atom_site loop mixes ATOM and HETATM records.""" + content = MMCIF_CONTENT.replace( + "ATOM 2 C CA ALA A 1 1.458 0.000 0.000 1.00 20.00", + "ATOM 2 C CA ALA A 1 1.458 0.000 0.000 1.00 20.00\n" + "HETATM 3 O O HOH A 2 5.000 5.000 5.000 1.00 30.00", + ) + path = tmp_path / "het.cif" + path.write_text(content) + return str(path) + + +class TestMmcifFolderParsedConfig: + def test_include_hetatm_default_keeps_hetatm(self, mmcif_file_with_hetatm): + dataset = load_dataset("mmcif", data_files=mmcif_file_with_hetatm, split="train") + structure = dataset[0]["structure"] + assert len(structure["label_atom_id"]) == 3 # ATOM x2 + HETATM x1 + + def test_include_hetatm_false_drops_hetatm(self, mmcif_file_with_hetatm): + dataset = load_dataset("mmcif", data_files=mmcif_file_with_hetatm, split="train", include_hetatm=False) + structure = dataset[0]["structure"] + assert len(structure["label_atom_id"]) == 2 # HOH (HETATM) excluded + assert "HOH" not in structure["label_comp_id"] + + def test_columns_subset(self, mmcif_file): + dataset = load_dataset("mmcif", data_files=mmcif_file, split="train", columns=["label_atom_id", "Cartn_x"]) + structure = dataset[0]["structure"] + assert set(structure) == {"label_atom_id", "Cartn_x"} + + class TestMmcifFolderIntegration: """Integration tests for MmcifFolder.""" @@ -279,17 +313,13 @@ def test_map_function(self, mmcif_data_dir): """Test map function on dataset.""" dataset = load_dataset("mmcif", data_dir=mmcif_data_dir, split="train") - def extract_entry_id(example): - # structure is decoded as a string by default - content = example["structure"] - # Extract entry ID from mmCIF content - for line in content.split("\n"): - if line.startswith("data_"): - return {"entry_id": line.replace("data_", "")} - return {"entry_id": "unknown"} + def count_atoms(example): + # structure is decoded into a parsed struct-of-arrays by default + return {"n_atoms": len(example["structure"]["label_atom_id"])} - mapped = dataset.map(extract_entry_id) - assert "entry_id" in mapped.column_names + mapped = dataset.map(count_atoms) + assert "n_atoms" in mapped.column_names + assert all(n > 0 for n in mapped["n_atoms"]) def test_save_and_load(self, mmcif_data_dir, tmp_path): """Test save and reload dataset.""" @@ -361,7 +391,7 @@ def test_large_mmcif_content(self, tmp_path): dataset = load_dataset("mmcif", data_dir=str(data_dir), split="train") assert len(dataset) == 1 - # Verify content is complete (decoded as string by default) + # Verify all atoms parsed (one row = one structure, parsed to atom arrays) content = dataset[0]["structure"] - assert isinstance(content, str) - assert "ALA A 100" in content + assert isinstance(content, dict) + assert len(content["label_atom_id"]) == 100