diff --git a/docs/source/loading.mdx b/docs/source/loading.mdx index c39b0e3e8e2..e7625b1cc4a 100644 --- a/docs/source/loading.mdx +++ b/docs/source/loading.mdx @@ -267,6 +267,50 @@ 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. + +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_dir="path/to/structures") +``` + +The `structure` column is parsed from the `_atom_site` category into a struct of +atom-level arrays. The default columns are: + +| 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_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_dir="path/to/structures", include_hetatm=False) +``` + +Extensions `.cif` and `.mmcif` are supported. + ## 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/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_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 new file mode 100644 index 00000000000..96143af3ff2 --- /dev/null +++ b/src/datasets/features/protein_structure.py @@ -0,0 +1,289 @@ +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, is_remote_url, 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). + + 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). + - 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 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[0]["structure"] + {'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, + 'path': 'path/to/structure.pdb'} + ``` + """ + + 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] = "dict" + 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) -> 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`): + 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: + `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 = {} + + 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}." + ) + 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.""" + 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, local_files: bool = True, remote_files: bool = True + ) -> 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. + 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 + `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 + 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(), + ) + 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 0de63b9c785..1b2c897da1e 100644 --- a/src/datasets/packaged_modules/__init__.py +++ b/src/datasets/packaged_modules/__init__.py @@ -16,6 +16,7 @@ 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 @@ -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}) @@ -132,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/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/__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..9f186fdf5ba --- /dev/null +++ b/src/datasets/packaged_modules/mmcif/mmcif.py @@ -0,0 +1,73 @@ +"""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 +""" + +from dataclasses import dataclass + +import datasets + +from ..folder_based_builder import folder_based_builder + + +logger = datasets.utils.logging.get_logger(__name__) + + +@dataclass +class MmcifFolderConfig(folder_based_builder.FolderBasedBuilderConfig): + """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__() + + +class MmcifFolder(folder_based_builder.FolderBasedBuilder): + """Folder-based builder for mmCIF protein structure files. + + 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. + """ + + BASE_FEATURE = datasets.ProteinStructure + 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/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_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 new file mode 100644 index 00000000000..fa18e45e0e8 --- /dev/null +++ b/tests/features/test_protein_structure.py @@ -0,0 +1,150 @@ +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 + # 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, dict) + assert "label_atom_id" in decoded_example + + +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, 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}) + + +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, 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): + 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"], 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(s, dict) for s in batch["structure"]) + column = dset["structure"] + assert len(column) == 1 + assert all(isinstance(s, dict) for s in column) + + # 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"], dict) + + +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/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/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 new file mode 100644 index 00000000000..b22161becf3 --- /dev/null +++ b/tests/packaged_modules/test_mmcif.py @@ -0,0 +1,397 @@ +"""Tests for MmcifFolder packaged module. + +Tests the folder-based mmCIF loader that follows the one-row-per-structure pattern +(similar to ImageFolder pattern). +""" + +import textwrap + +import pyarrow as pa +import pytest + +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 single mmCIF test file.""" + path = tmp_path / "structure.cif" + path.write_text(MMCIF_CONTENT) + return str(path) + + +@pytest.fixture +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_data_dir(tmp_path): + """Create a directory with mmCIF files.""" + data_dir = tmp_path / "mmcif_data" + data_dir.mkdir() + + # Create several mmCIF files + for i in range(3): + (data_dir / f"structure_{i}.cif").write_text(MMCIF_CONTENT.replace("TEST", f"TEST_{i}")) + + return str(data_dir) + + +@pytest.fixture +def mmcif_data_dir_with_labels(tmp_path): + """Create directories with label-based structure.""" + data_dir = tmp_path / "mmcif_labeled" + + # 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) + + return str(data_dir) + + +@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 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 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.""" + 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 + + +@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.""" + + 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 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(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.""" + 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() + + # 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), + ) + + (data_dir / "large.cif").write_text(large_content) + + dataset = load_dataset("mmcif", data_dir=str(data_dir), split="train") + assert len(dataset) == 1 + + # Verify all atoms parsed (one row = one structure, parsed to atom arrays) + content = dataset[0]["structure"] + assert isinstance(content, dict) + assert len(content["label_atom_id"]) == 100