Dataset Architecture#

AtomWorks provides a modern, composable dataset architecture that separates data loading, processing, and transformation concerns. This approach replaces the legacy parser-based system with functional loaders and transform pipelines.

Warning

The metadata parser system (atomworks.ml.datasets.parsers) is deprecated and will be removed in a future version. Use the new loader-based approach with FileDataset and PandasDataset instead.

Modern Dataset Architecture#

The current AtomWorks dataset system consists of three main components:

Datasets: Container classes that manage data access and indexing
Loaders: Functions that process raw data into transform-ready format
Transforms: Pipelines that convert loaded data into model inputs

Dataset Classes#

class atomworks.ml.datasets.ArrowMetadataIndex(*, data: DataFrame | PathLike, name: str, id_column: str = 'example_id', filters: list[str] | None = None, columns_to_load: list[str] | None = None, load_kwargs: dict | None = None, local_drive_mount: str = '/tmp', job_id_env_var: str = 'SLURM_JOB_ID')[source]#

Bases: object

Memory-mapped Arrow metadata with the same interface as MetadataIndex.

Converts the DataFrame to an uncompressed feather file on local storage and reads it back as a memory-mapped pyarrow.Table, keeping Python heap usage minimal. ID lookups are O(N) — acceptable because they are not on the DataLoader hot path.

get_column_values(column: str) → ndarray[source]#: Get values of a column as numpy array.

get_example_id(idx: int) → str[source]#: Get example ID by positional index.

get_row(idx: int) → Series[source]#: Get metadata row by positional index.

id_to_idx(example_id: str | list[str]) → int | list[int][source]#: Convert example ID(s) to positional index(es). O(N) per call.

idx_to_id(idx: int | list[int]) → str | ndarray[source]#: Convert positional index(es) to example ID(s).

class atomworks.ml.datasets.ConcatDatasetWithID(datasets: list[ExampleIDProtocol])[source]#

Bases: ConcatDataset

Equivalent to torch.utils.data.ConcatDataset but allows accessing examples by ID.

Provides ID-based access across multiple datasets that implement ExampleIDProtocol.

datasets: list[ExampleIDProtocol]#

get_dataset_by_id(example_id: str) → Dataset[source]#

Retrieves the dataset containing the example ID.

Parameters:: example_id – The ID to find.
Returns:: The sub-dataset containing the ID.

Warning

Assumes that the example ID is unique within the dataset. If not, the first occurrence of the example ID is returned.

get_dataset_by_idx(idx: int) → Dataset[source]#

Retrieves the dataset containing the index.

Parameters:: idx – The index to find.
Returns:: The sub-dataset containing the index.
Raises:: ValueError – If the index is out of bounds.

id_to_idx(example_id: str) → int[source]#

Retrieves the index corresponding to the example ID.

Parameters:: example_id – The ID to convert.
Returns:: The corresponding index.
Raises:: ValueError – If the example ID is not found.

Warning

Assumes that the example ID is unique within the dataset. If not, the first occurrence of the example ID is returned.

idx_to_id(idx: int) → str[source]#

Retrieves the example ID corresponding to the index.

Parameters:: idx – The index to convert.
Returns:: The corresponding example ID.
Raises:: ValueError – If the index is out of bounds.

class atomworks.ml.datasets.ExampleIDProtocol(*args, **kwargs)[source]#

Bases: Protocol

Structural interface for datasets that support ID-based access.

id_to_idx(example_id: str | list[str]) → int | list[int][source]#: Convert example ID(s) to index(es).

idx_to_id(idx: int | list[int]) → str | ndarray[source]#: Convert index(es) to example ID(s).

class atomworks.ml.datasets.FallbackDatasetWrapper(dataset: Dataset, fallback_dataset: Dataset)[source]#

Bases: Dataset

A wrapper around a dataset that allows for a fallback dataset to be used when an error occurs.

Meant to be used with a FallbackSamplerWrapper.

class atomworks.ml.datasets.FileDataset(*, file_paths: list[str | PathLike], name: str, filter_fn: Callable[[PathLike], bool] | None = None, **kwargs: Any)[source]#

Bases: MolecularDataset

Dataset that loads molecular data from individual files.

Each file represents one example in the dataset. If creating a dataset from a directory, use the from_directory() class method instead of the default constructor.

classmethod from_directory(*, directory: PathLike, name: str, max_depth: int = 3, **kwargs: Any) → FileDataset[source]#

Create a FileDataset by scanning a directory for files.

Parameters:

directory – Path to directory to scan for files.
name – Descriptive name for this dataset.
max_depth – Maximum depth to scan for files in subdirectories.
**kwargs – Additional arguments passed to FileDataset.

Returns:

FileDataset instance with files discovered from the directory.

Example

Create from directory:

>>> dataset = FileDataset.from_directory(directory="/path/to/files", name="my_dataset", max_depth=2)

classmethod from_file_list(*, file_paths: list[str | PathLike], name: str, **kwargs: Any) → FileDataset[source]#

Create a FileDataset from an explicit list of file paths.

This is an alias for the main constructor for clarity and consistency with from_directory().

Parameters:

file_paths – List of file paths for the dataset. Each file represents one example.
name – Descriptive name for this dataset.
**kwargs – Additional arguments passed to FileDataset.

Returns:

FileDataset instance with the provided file paths.

id_to_idx(example_id: str | list[str]) → int | list[int][source]#: Convert example ID(s) to index(es).

idx_to_id(idx: int | list[int]) → str | list[str][source]#: Convert index(es) to example ID(s).

class atomworks.ml.datasets.MetadataIndex(*, data: DataFrame | PathLike, name: str, id_column: str = 'example_id', filters: list[str] | None = None, columns_to_load: list[str] | None = None, load_kwargs: dict | None = None)[source]#

Bases: object

Parquet/DataFrame metadata with filtering and ID lookup.

Single source of truth for example IDs across all dataset backends. The id_column is set as the pandas index for O(1) .loc lookups.

get_column_values(column: str) → ndarray[source]#: Get values of a column as numpy array.

get_example_id(idx: int) → str[source]#: Get example ID by positional index.

get_row(idx: int) → Series[source]#: Get metadata row by positional index.

id_to_idx(example_id: str | list[str]) → int | list[int][source]#: Convert example ID(s) to positional index(es).

idx_to_id(idx: int | list[int]) → str | ndarray[source]#: Convert positional index(es) to example ID(s).

class atomworks.ml.datasets.MetadataIndexProtocol(*args, **kwargs)[source]#

Bases: Protocol

Structural interface shared by all metadata index backends.

get_column_values(column: str) → ndarray[source]#

get_example_id(idx: int) → str[source]#

get_row(idx: int) → Series[source]#

id_to_idx(example_id: str | list[str]) → int | list[int][source]#

idx_to_id(idx: int | list[int]) → str | ndarray[source]#

class atomworks.ml.datasets.MolecularDataset(*, name: str, transform: Callable | None = None, loader: Callable | None = None, save_failed_examples_to_dir: str | Path | None = None)[source]#

Bases: Dataset

Base class for AtomWorks molecular datasets.

Handles Transform pipelines and loader functionality for molecular data. Subclasses implement __getitem__() with their own data access patterns. Dictionary samples include dataset_name after transforms finish.

Bases: MolecularDataset

Dataset for tabular data stored as pandas DataFrames.

Delegates all metadata, filtering, and ID-mapping logic to MetadataIndex.

property data: DataFrame | Table#

The underlying metadata table.

Returns a pd.DataFrame for MetadataIndex or a pa.Table for ArrowMetadataIndex.

Raises:: AttributeError – If the metadata backend has no data attribute (e.g. SequentialMetadataIndex).

id_to_idx(example_id: str | list[str]) → int | list[int][source]#: Convert an example ID to the corresponding local index.

idx_to_id(idx: int | list[int]) → str | ndarray[source]#: Convert a local index to the corresponding example ID.

property metadata: MetadataIndexProtocol#: The metadata index.

class atomworks.ml.datasets.SequentialMetadataIndex(*, n_entries: int, idx_column: str)[source]#

Bases: object

Identity metadata mapping "0" … "n-1" to sequential indices.

Same interface as MetadataIndex / ArrowMetadataIndex for datasets that don’t need an external metadata file.

get_column_values(column: str) → ndarray[source]#: Return sequential indices for idx_column.

get_example_id(idx: int) → str[source]#

get_row(idx: int) → Series[source]#: Not available — SequentialMetadataIndex has no row data.

id_to_idx(example_id: str | list[str]) → int | list[int][source]#: Convert string index(es) to int.

idx_to_id(idx: int | list[int]) → str | ndarray[source]#: Convert int index(es) to string.

atomworks.ml.datasets.get_row_and_index_by_example_id(dataset: ExampleIDProtocol, example_id: str) → dict[source]#

Retrieve a row and its index from a nested dataset structure by its example ID.

Parameters:

dataset – The dataset or concatenated dataset to search. Must have the id_to_idx method.
example_id – The example ID to search for.

Returns:

Dictionary containing the row and global index corresponding to the example ID.

Functional Loaders#

Loaders are functions that process raw dataset output (e.g., pandas Series) into a Transform-ready format. They replace the legacy parser classes with a more flexible, functional approach.

Functional loader implementations for AtomWorks datasets.

Loaders are functions that process raw dataset output (e.g., pandas Series) into a Transform-ready format. E.g., converts what may be dataset-specific metadata into a standard format for use in AtomWorks Transform pipelines.

atomworks.ml.datasets.loaders.create_ase_materials_loader(*, wrap_fractional_coordinates: bool = False, include_atom_array: bool = False, keep_ase_atoms: bool = True) → Callable[[tuple], dict[str, Any]][source]#

Create a picklable periodic-materials loader for AseDBDataset.

Consumes (atoms_row, global_index, metadata_row). Dataset metadata indexing supplies example_id. Optional AtomArrays preserve the cell, PBC and ASE initial charges; they do not infer bonds or assign formal charges. They are coordinate containers, not structures prepared for molecular ML transforms.

atomworks.ml.datasets.loaders.create_base_loader(example_id_colname: str = 'example_id', path_colname: str = 'path', assembly_id_colname: str | None = 'assembly_id', altloc_seed_colname: str | None = None, attrs: dict | None = None, base_path: str = '', extension: str = '', sharding_pattern: str | None = None, parser_args: ParseConfig | dict | None = None) → Callable[[Series], dict[str, Any]][source]#: Deprecated; use create_structure_loader().

atomworks.ml.datasets.loaders.create_blob_cif_loader(blob_dir: str, *, endpoint_url: str | None = None, example_id_colname: str = 'example_id', record_id_colname: str | None = None, assembly_id_colname: str | None = 'assembly_id', altloc_seed_colname: str | None = None, parser_args: ParseConfig | dict | None = None) → Callable[[Series], dict[str, Any]][source]#: Deprecated; use create_structure_loader() with storage="blob".

atomworks.ml.datasets.loaders.create_cif_bytes_loader(parser_args: ParseConfig | dict | None = None, assembly_id_colname: str | None = 'assembly_id', altloc_seed_colname: str | None = None) → Callable[source]#: Deprecated; use create_structure_loader() with storage="bytes".

atomworks.ml.datasets.loaders.create_loader_with_interfaces_and_pn_units_to_score(example_id_colname: str = 'example_id', path_colname: str = 'path', assembly_id_colname: str | None = 'assembly_id', altloc_seed_colname: str | None = None, interfaces_to_score_colname: str | None = 'interfaces_to_score', pn_units_to_score_colname: str | None = 'pn_units_to_score', base_path: str = '', extension: str = '', sharding_pattern: str | None = None, attrs: dict | None = None, parser_args: ParseConfig | dict | None = None) → Callable[[Series], dict[str, Any]][source]#: Deprecated; use create_structure_loader() with column_mapping.

atomworks.ml.datasets.loaders.create_loader_with_query_pn_units(example_id_colname: str = 'example_id', path_colname: str = 'path', pn_unit_iid_colnames: str | list[str] | None = None, assembly_id_colname: str | None = 'assembly_id', altloc_seed_colname: str | None = None, base_path: str = '', extension: str = '', sharding_pattern: str | None = None, attrs: dict | None = None, parser_args: ParseConfig | dict | None = None) → Callable[[Series], dict[str, Any]][source]#: Deprecated; use create_structure_loader() with column_mapping.

atomworks.ml.datasets.loaders.create_structure_loader(storage: Literal['filesystem', 'bytes', 'blob'] = 'filesystem', *, column_mapping: Mapping[str, str | Sequence[str] | None] | None = None, **kwargs: Any) → Callable[source]#

Create a structure loader with consistent metadata projection across storage backends.

column_mapping maps output names to one metadata column or an ordered sequence of columns. Remaining arguments are forwarded to the selected backend factory. Blob loaders check the parsed cache before fetching records; shard paths must be immutable and globally identify their contents. Row metadata is applied after caching.

Basic Usage Examples#

File-based datasets (replacing simple file parsers):

from atomworks.ml.datasets import FileDataset
from atomworks.io import parse

def simple_loading_fn(raw_data) -> dict:
    """Simple loading function that parses structural data."""
    parse_output = parse(raw_data)
    return {"atom_array": parse_output["assemblies"]["1"][0]}

dataset = FileDataset.from_directory(
    directory="/path/to/structures",
    name="my_dataset",
    loader=simple_loading_fn
)

Tabular datasets (replacing metadata parsers):

from atomworks.ml.datasets import PandasDataset
from atomworks.ml.datasets.loaders import create_structure_loader

dataset = PandasDataset(
    data="metadata.parquet",
    name="interfaces_dataset",
    loader=create_structure_loader(
        column_mapping={"query_pn_unit_iids": ["pn_unit_1_iid", "pn_unit_2_iid"]}
    )
)

Custom loaders for specialized use cases:

def custom_loader(row: pd.Series) -> dict:
    """Custom loader with specific processing logic."""
    # Load structure
    structure_path = Path(row["path"])
    parse_output = parse(structure_path)

    # Extract specific metadata
    metadata = {
        "resolution": row.get("resolution", None),
        "method": row.get("method", "unknown"),
        "custom_field": row.get("custom_field", "default_value")
    }

    return {
        "atom_array": parse_output["assemblies"]["1"][0],
        "extra_info": metadata,
        "example_id": row["example_id"]
    }

dataset = PandasDataset(
    data=my_dataframe,
    name="custom_dataset",
    loader=custom_loader
)

Common Loader Patterns#

Base loader for standard structure loading:

from atomworks.ml.datasets.loaders import create_structure_loader

loader = create_structure_loader(
    example_id_colname="example_id",
    path_colname="path",
    assembly_id_colname="assembly_id",
    base_path="/data/structures",
    extension=".cif"
)

Interface loader for protein-protein interfaces:

from atomworks.ml.datasets.loaders import create_structure_loader

loader = create_structure_loader(
    column_mapping={"query_pn_unit_iids": ["pn_unit_1_iid", "pn_unit_2_iid"]},
    base_path="/data/pdb",
    extension=".cif.gz"
)

Validation loader with scoring targets:

from atomworks.ml.datasets.loaders import create_structure_loader

loader = create_structure_loader(
    column_mapping={
        "interfaces_to_score": "interfaces_to_score",
        "pn_units_to_score": "pn_units_to_score",
    }
)

Integration with Transform Pipelines#

Loaders work seamlessly with AtomWorks transform pipelines. The loader output becomes the input to the transform pipeline:

from atomworks.ml.transforms.base import Compose
from atomworks.ml.transforms.crop import CropSpatialLikeAF3
from atomworks.ml.transforms.atom_array import AddGlobalAtomIdAnnotation

# Create a transform pipeline
transform_pipeline = Compose([
    AddGlobalAtomIdAnnotation(),
    CropSpatialLikeAF3(crop_size=256),
])

# Create dataset with both loader and transforms
dataset = PandasDataset(
    data="metadata.parquet",
    name="my_dataset",
    loader=loader_with_query_pn_units(
        pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"]
    ),
    transform=transform_pipeline
)

# Access processed data
example = dataset[0]  # Returns transformed data ready for model input

Data Flow#

The complete data flow in the new architecture is:

Raw Data: File paths or DataFrame rows
Loader: Processes raw data into standardized format with AtomArray
Transform Pipeline: Converts loaded data into model-ready tensors
Model Input: Final processed data ready for training/inference

This separation allows for: - Reusable loaders across different datasets - Composable transforms that can be mixed and matched - Easy testing of individual components - Clear debugging when issues arise

Periodic materials#

Use the existing AseDBDataset with the materials loader for ASE LMDB shards:

from atomworks.ml.datasets.ase_dataset import AseDBDataset
from atomworks.ml.datasets.loaders import create_ase_materials_loader

dataset = AseDBDataset(
    lmdb_path="materials_shards/",
    name="materials",
    loader=create_ase_materials_loader(include_atom_array=True),
)
sample = dataset[0]

Install the ml and ase extras. Each sample includes Cartesian and fractional coordinates, cell vectors, lengths, angles, volume, and PBC. Available ASE row energy, forces and stress are retained. Molecular bond inference and preparation are not applied. Optional AtomArrays preserve the cell and initial charges without interpreting partial charges as formal charges.

space_group and parent_space_group are separate numeric labels supplied by the data source. Neither is computed from the coordinates, and parent/prototype labels are never promoted to current symmetry. Other metadata remains in extra_info. A full-rank cell is required; unwrapped fractional positions are returned by default. Wrapping affects only periodic axes.

Use the existing metadata example_id and lmdb_idx columns to select or reorder examples. Source ASE row IDs remain available as source_row_id.