atomic-datasets 0.1.0__py3-none-any.whl

atomic_datasets/__init__.py

@@ -0,0 +1,3 @@
+__version__ = "0.1.0"
+
+from .datatypes import Graph, MolecularDataset

atomic_datasets/datasets/__init__.py

@@ -0,0 +1,15 @@
+from .qm9 import QM9
+from .geom_drugs import GEOMDrugs
+from .tmqm import tmQM
+from .platonic_solids import PlatonicSolids
+from .tetris import Tetris
+from .proteins import (
+    CATHAlphaCarbons,
+    MiniproteinsAlphaCarbons,
+    MiniproteinsBackbone,
+    MiniproteinsBackboneNoAA,
+    Miniproteins,
+    get_amino_acids,
+)
+from .chembl3d import ChEMBL3D
+from .crossdocked import CrossDocked

atomic_datasets/datasets/chembl3d.py

@@ -0,0 +1,205 @@
+import os
+import json
+import zipfile
+from typing import Dict, Iterable, Optional, Sequence, Any
+
+import numpy as np
+
+from atomic_datasets import datatypes
+from atomic_datasets import utils
+
+
+# Zenodo URL for preprocessed data
+ChEMBL3D_ZENODO_URL = "https://zenodo.org/records/18488050/files/chembl3d_processed.zip"
+
+
+class ChEMBL3D(datatypes.MolecularDataset):
+    """
+    Dataset of ChEMBL3D structures from https://github.com/isayevlab/LoQI.
+    
+    Contains structures for a large subset of ChEMBL, pre-optimized or extracted
+    from 3D experimental data. This implementation uses a high-performance 
+    contiguous memory layout with memory-mapping support.
+
+    Args:
+        root_dir: Directory to store/load data
+        split: Which split to use ('train', 'val', 'test_small', 'test_rot_bonds', 'test_cremp')
+        start_index: Start index for slicing the dataset
+        end_index: End index for slicing the dataset
+        train_on_single_molecule: If True, use single molecule for all splits
+        train_on_single_molecule_index: Index of molecule to use if train_on_single_molecule=True
+        mmap_mode: Memory-map mode for numpy arrays ('r', 'r+', 'c', or None to load into memory)
+    """
+
+    ATOMIC_NUMBERS = np.asarray([
+        1, 5, 6, 7, 8, 9, 13, 14, 15, 16, 17, 33, 35, 53, 80, 83, 34
+    ], dtype=np.int32)
+
+    def __init__(
+        self,
+        root_dir: str,
+        split: str,
+        start_index: Optional[int] = None,
+        end_index: Optional[int] = None,
+        train_on_single_molecule: bool = False,
+        train_on_single_molecule_index: int = 0,
+        mmap_mode: Optional[str] = 'r',
+    ):
+        super().__init__(atomic_numbers=self.ATOMIC_NUMBERS)
+
+        self.root_dir = os.path.join(root_dir, "chembl3d")
+        self.split = split
+        self.start_index = start_index
+        self.end_index = end_index
+        self.train_on_single_molecule = train_on_single_molecule
+        self.train_on_single_molecule_index = train_on_single_molecule_index
+        self.mmap_mode = mmap_mode
+
+        self.preprocessed = False
+        
+        # Data storage - will be initialized as contiguous arrays or memory-mapped views
+        self._positions = None      # (total_atoms, 3) float32
+        self._species = None        # (total_atoms,) uint8
+        self._offsets = None        # (n_molecules + 1,) int64
+        self._charges = None        # (total_atoms,) int8
+        self._is_aromatic = None    # (total_atoms,) uint8
+        self._is_in_ring = None     # (total_atoms,) uint8
+        self._hybridization = None  # (total_atoms,) uint8
+        self._smiles = None         # list of str
+        self._chemblids = None      # list of str
+        self._indices = None        # indices after filtering/slicing
+
+        self.preprocess()
+
+    def _get_file_path(self, filename: str) -> str:
+        """Helper to get the absolute path to a preprocessed file."""
+        return os.path.join(self.root_dir, filename)
+
+    def _download_from_zenodo(self):
+        """Download and extract preprocessed data from Zenodo."""
+        os.makedirs(self.root_dir, exist_ok=True)
+        
+        zip_path = os.path.join(self.root_dir, "chembl3d_processed.zip")
+        
+        # Download
+        print(f"Downloading ChEMBL3D from Zenodo...")
+        print(f"  URL: {ChEMBL3D_ZENODO_URL}")
+        utils.download_url(ChEMBL3D_ZENODO_URL, self.root_dir)
+        
+        # Extract
+        print(f"Extracting to {self.root_dir}...")
+        with zipfile.ZipFile(zip_path, 'r') as zf:
+            # Extract and flatten directory structure
+            for member in zf.namelist():
+                # Skip directories
+                if member.endswith('/'):
+                    continue
+                
+                # Get the filename without the parent directory
+                filename = os.path.basename(member)
+                if not filename:
+                    continue
+                
+                # Extract to root_dir
+                source = zf.open(member)
+                target_path = os.path.join(self.root_dir, filename)
+                with open(target_path, 'wb') as target:
+                    target.write(source.read())
+                source.close()
+        
+        # Clean up zip file to save disk space
+        os.remove(zip_path)
+        print("Download complete.")
+
+    def preprocess(self):
+        """Load preprocessed numpy files and setup indices."""
+        if self.preprocessed:
+            return
+
+        prefix = self.split
+        
+        # Check that files exist, download if needed
+        positions_path = self._get_file_path(f"{prefix}_positions.npy")
+        if not os.path.exists(positions_path):
+            self._download_from_zenodo()
+        
+        print(f"Loading ChEMBL3D {self.split} split from: {self.root_dir}")
+        
+        # Load arrays (memory-mapped for large files to save RAM)
+        self._positions = np.load(self._get_file_path(f"{prefix}_positions.npy"), mmap_mode=self.mmap_mode)
+        self._species = np.load(self._get_file_path(f"{prefix}_species.npy"), mmap_mode=self.mmap_mode)
+        self._offsets = np.load(self._get_file_path(f"{prefix}_offsets.npy"))  # Small, load fully into RAM
+        self._charges = np.load(self._get_file_path(f"{prefix}_charges.npy"), mmap_mode=self.mmap_mode)
+        self._is_aromatic = np.load(self._get_file_path(f"{prefix}_is_aromatic.npy"), mmap_mode=self.mmap_mode)
+        self._is_in_ring = np.load(self._get_file_path(f"{prefix}_is_in_ring.npy"), mmap_mode=self.mmap_mode)
+        self._hybridization = np.load(self._get_file_path(f"{prefix}_hybridization.npy"), mmap_mode=self.mmap_mode)
+        
+        # Load metadata (SMILES and IDs)
+        metadata_path = self._get_file_path(f"{prefix}_metadata.json")
+        with open(metadata_path, 'r') as f:
+            metadata = json.load(f)
+        self._smiles = metadata["smiles"]
+        self._chemblids = metadata["chemblids"]
+        
+        n_molecules = len(self._offsets) - 1
+        n_atoms = len(self._positions)
+        print(f"Loaded {n_molecules:,} molecules, {n_atoms:,} total atoms")
+        
+        # Apply single molecule override for debugging or overfitting tests
+        if self.train_on_single_molecule:
+            self._indices = np.array([self.train_on_single_molecule_index])
+        else:
+            self._indices = np.arange(n_molecules)
+        
+        # Apply start/end index for subsetting
+        self._indices = self._indices[slice(self.start_index, self.end_index)]
+        
+        self.preprocessed = True
+
+    def __len__(self) -> int:
+        """Returns the number of molecules in the current split/slice."""
+        return len(self._indices)
+
+    def __getitem__(self, idx: int) -> datatypes.Graph:
+        """
+        Fast random access to a molecule using offset-based slicing.
+        
+        Returns:
+            A datatypes.Graph object containing positions, species, and atomic properties.
+        """
+        if idx < 0:
+            idx = len(self._indices) + idx
+        if idx < 0 or idx >= len(self._indices):
+            raise IndexError(f"Index {idx} out of range for dataset of size {len(self._indices)}")
+        
+        mol_idx = self._indices[idx]
+        start = self._offsets[mol_idx]
+        end = self._offsets[mol_idx + 1]
+        
+        # Slice arrays (copies from mmap into RAM for the return object)
+        species = np.array(self._species[start:end])
+        atomic_numbers = self.species_to_atomic_numbers(species)
+        atom_types = utils.atomic_numbers_to_symbols(atomic_numbers)
+        
+        return datatypes.Graph(
+            nodes=dict(
+                positions=np.array(self._positions[start:end]),
+                atomic_numbers=atomic_numbers,
+                species=species,
+                atom_types=atom_types,
+            ),
+            edges=None,
+            receivers=None,
+            senders=None,
+            globals=None,
+            n_node=np.asarray([end - start]),
+            n_edge=None,
+            properties={
+                "smiles": self._smiles[mol_idx],
+                "chemblid": self._chemblids[mol_idx],
+                "charges": np.array(self._charges[start:end]),
+                "is_aromatic": np.array(self._is_aromatic[start:end]),
+                "is_in_ring": np.array(self._is_in_ring[start:end]),
+                "hybridization": np.array(self._hybridization[start:end]),
+            },
+        )

atomic_datasets/datasets/crossdocked.py

@@ -0,0 +1,167 @@
+from typing import Iterable, Dict, Optional, Tuple
+
+import os
+import logging
+import json
+import zipfile
+
+import numpy as np
+
+from atomic_datasets import utils
+from atomic_datasets import datatypes
+
+# Zenodo URL for preprocessed data
+CROSSDOCKED_ZENODO_URL = "https://zenodo.org/records/18584578/files/crossdocked.zip"
+
+
+class CrossDocked(datatypes.MolecularDataset):
+    """
+    The CrossDocked dataset from https://pubs.acs.org/doi/full/10.1021/acs.jcim.0c00411
+    with splits from Luo et al. (https://proceedings.neurips.cc/paper/2021/hash/314450613369e0ee72d0da7f6fee773c-Abstract.html).
+
+    Loads preprocessed data from Zenodo (memory-mapped for efficiency).
+
+    Args:
+        root_dir: Directory to store/load data
+        split: Which split to use ('train', 'val', 'test')
+        start_index: Start index for slicing the dataset
+        end_index: End index for slicing the dataset
+        mmap_mode: Memory-map mode for numpy arrays ('r', 'r+', 'c', or None to load into memory)
+    """
+
+    ATOMIC_NUMBERS = np.asarray([
+        13, 33, 79,  5, 35,  6, 17, 27, 24, 29,  9, 26,  1, 80, 53,  3, 12,
+        42,  7,  8, 15, 44, 16, 21, 34, 14, 50, 23, 74, 39
+    ], dtype=np.int32)
+
+    def __init__(
+        self,
+        root_dir: str,
+        split: str = "train",
+        start_index: Optional[int] = None,
+        end_index: Optional[int] = None,
+        mmap_mode: Optional[str] = 'r',
+    ):
+        super().__init__(atomic_numbers=self.ATOMIC_NUMBERS)
+
+        self.root_dir = os.path.join(root_dir, "crossdocked")
+        self.split = split
+        self.start_index = start_index
+        self.end_index = end_index
+        self.mmap_mode = mmap_mode
+
+        self.preprocessed = False
+
+        # Data storage
+        self._positions = None       # (N_total, 3) memory-mapped
+        self._atom_types = None      # (N_total,) memory-mapped (indices into lookup)
+        self._offsets = None         # (n_complexes + 1,) start indices
+        self._n_atoms = None         # (n_complexes,) atoms per complex
+        self._atom_type_lookup = None  # (n_types,) symbol strings
+        self._properties = None      # List of dicts (pocket_file, starting_fragment_mask)
+        self._indices = None         # Indices after slicing
+
+        if split not in ("train", "val", "test"):
+            raise ValueError(f"split must be 'train', 'val', or 'test', got '{split}'")
+
+        self.preprocess()
+
+        readme_path = os.path.join(self.root_dir, "README.md")
+        if os.path.exists(readme_path):
+            print("Dataset description available at:", os.path.abspath(readme_path))
+
+    def preprocess(self):
+        """Initialize data access - downloads if needed, then loads."""
+        if self.preprocessed:
+            return
+
+        self._ensure_downloaded()
+        self._load_data()
+        self._setup_indices()
+
+        self.preprocessed = True
+
+    def _ensure_downloaded(self):
+        """Download and extract preprocessed files from Zenodo if not present."""
+        os.makedirs(self.root_dir, exist_ok=True)
+
+        # Check if data is already extracted
+        marker_file = os.path.join(self.root_dir, "crossdocked", "train_positions.npy")
+        if os.path.exists(marker_file):
+            return
+
+        zip_filename = "crossdocked.zip"
+        zip_path = os.path.join(self.root_dir, zip_filename)
+
+        if not os.path.exists(zip_path):
+            utils.download_url(CROSSDOCKED_ZENODO_URL, self.root_dir, filename=zip_filename)
+
+        print(f"Extracting {zip_filename}...")
+        with zipfile.ZipFile(zip_path, 'r') as zf:
+            zf.extractall(self.root_dir)
+
+        os.remove(zip_path)
+        print("Extraction complete.")
+
+    def _load_data(self):
+        """Load preprocessed data using memory mapping."""
+        prefix = self.split
+        print(f"Loading CrossDocked {self.split} split from {self.root_dir}")
+
+        self._positions = np.load(
+            os.path.join(self.root_dir, "crossdocked", f"{prefix}_positions.npy"),
+            mmap_mode=self.mmap_mode,
+        )
+        self._atom_types = np.load(
+            os.path.join(self.root_dir, "crossdocked", f"{prefix}_atom_types.npy"),
+            mmap_mode=self.mmap_mode,
+        )
+
+        self._offsets = np.load(os.path.join(self.root_dir, "crossdocked", f"{prefix}_offsets.npy"))
+        self._n_atoms = np.load(os.path.join(self.root_dir, "crossdocked", f"{prefix}_n_atoms.npy"))
+        self._atom_type_lookup = np.load(
+            os.path.join(self.root_dir, "crossdocked", f"{prefix}_atom_type_lookup.npy")
+        )
+
+        with open(os.path.join(self.root_dir, "crossdocked", f"{prefix}_properties.json")) as f:
+            self._properties = json.load(f)
+
+        n_complexes = len(self._n_atoms)
+        print(f"Loaded {n_complexes} complexes")
+
+    def _setup_indices(self):
+        """Setup indices with optional slicing."""
+        n_complexes = len(self._n_atoms)
+        self._indices = np.arange(n_complexes)
+        self._indices = self._indices[slice(self.start_index, self.end_index)]
+
+    def __len__(self) -> int:
+        return len(self._indices)
+
+    def __getitem__(self, idx: int) -> datatypes.Graph:
+        """Fast slice access via memory-mapped offsets."""
+        if idx < 0:
+            idx = len(self._indices) + idx
+
+        real_idx = self._indices[idx]
+        start, end = self._offsets[real_idx], self._offsets[real_idx + 1]
+
+        species = np.array(self._atom_types[start:end])
+        atom_types = self._atom_type_lookup[species]
+        atomic_numbers = utils.atomic_symbols_to_numbers(atom_types)
+
+        return datatypes.Graph(
+            nodes=dict(
+                positions=np.array(self._positions[start:end]),
+                atomic_numbers=atomic_numbers,
+                species=species,
+                atom_types=atom_types,
+            ),
+            edges=None,
+            senders=None,
+            receivers=None,
+            n_edge=None,
+            n_node=np.asarray([self._n_atoms[real_idx]]),
+            globals=None,
+            properties=self._properties[real_idx],
+        )

atomic_datasets/datasets/geom_drugs.py

@@ -0,0 +1,209 @@
+from typing import Dict, Iterable, Optional, List, Tuple, Union
+import os
+import logging
+import json
+import zipfile
+
+import numpy as np
+
+from atomic_datasets import datatypes
+from atomic_datasets import utils
+
+
+# Zenodo URL for preprocessed data
+GEOM_DRUGS_ZENODO_URL = "https://zenodo.org/record/18484634/files/geom_drugs_processed.zip"
+
+
+class GEOMDrugs(datatypes.MolecularDataset):
+    """
+    The GEOM (Drugs) dataset from https://www.nature.com/articles/s41597-022-01288-4.
+    
+    Loads preprocessed data from Zenodo (memory-mapped for efficiency).
+
+    Args:
+        root_dir: Directory to store/load data
+        split: Which split to use ('train', 'val', 'test')
+        start_index: Start index for slicing the dataset
+        end_index: End index for slicing the dataset
+        max_atoms: Filter out molecules with more atoms than this
+        conformer_selection: How to select conformers ('first', 'random', 'all')
+        random_seed: Random seed for conformer selection (if conformer_selection='random')
+        mmap_mode: Memory-map mode for numpy arrays ('r', 'r+', 'c', or None to load into memory)
+    """
+
+    # Atomic numbers present in GEOM-Drugs
+    ATOMIC_NUMBERS = np.asarray([1, 5, 6, 7, 8, 9, 13, 14, 15, 16, 17, 33, 35, 53, 80, 83])
+
+    def __init__(
+        self,
+        root_dir: str,
+        split: str = "train",
+        start_index: Optional[int] = None,
+        end_index: Optional[int] = None,
+        max_atoms: Optional[int] = None,
+        conformer_selection: str = "all",
+        random_seed: int = 0,
+        mmap_mode: Optional[str] = 'r',
+    ):
+        # Initialize the base class mapping logic
+        super().__init__(atomic_numbers=self.ATOMIC_NUMBERS)
+        
+        self.root_dir = os.path.join(root_dir, "geom_drugs")
+        self.split = split
+        self.start_index = start_index
+        self.end_index = end_index
+        self.max_atoms = max_atoms
+        self.conformer_selection = conformer_selection
+        self.random_seed = random_seed
+        self.mmap_mode = mmap_mode
+
+        self.preprocessed = False
+        
+        # Data storage
+        self._positions = None          # (N_total, 3) memory-mapped
+        self._atomic_numbers = None     # (N_total,) memory-mapped
+        self._offsets = None            # (n_conformers + 1,) start indices
+        self._n_atoms = None            # (n_conformers,) atoms per conformer
+        self._mol_indices = None        # (n_conformers,) molecule index
+        self._smiles = None             # List of SMILES strings
+        self._indices = None            # Indices into conformers (after filtering)
+        self._rng = None                
+
+        if split not in ("train", "val", "test"):
+            raise ValueError(f"split must be 'train', 'val', or 'test', got '{split}'")
+        
+        if conformer_selection not in ("first", "random", "all"):
+            raise ValueError(f"conformer_selection must be 'first', 'random', or 'all', got '{conformer_selection}'")
+
+        self.preprocess()
+
+    def preprocess(self):
+        """Initialize data access - downloads if needed, then loads."""
+        if self.preprocessed:
+            return
+        
+        # Download and extract if needed
+        self._ensure_downloaded()
+        
+        # Load data
+        self._load_data()
+        
+        # Setup indices based on conformer selection
+        self._rng = np.random.default_rng(self.random_seed)
+        self._setup_indices()
+        
+        self.preprocessed = True
+    
+    def _ensure_downloaded(self):
+        """Download and extract preprocessed files from Zenodo if not present."""
+        os.makedirs(self.root_dir, exist_ok=True)
+        
+        # Check if data is already extracted
+        marker_file = os.path.join(self.root_dir, "train_positions.npy")
+        if os.path.exists(marker_file):
+            return
+        
+        zip_filename = "geom_drugs_processed.zip"
+        zip_path = os.path.join(self.root_dir, zip_filename)
+        
+        if not os.path.exists(zip_path):
+            print(f"Downloading {zip_filename}...")
+            utils.download_url(GEOM_DRUGS_ZENODO_URL, self.root_dir, filename=zip_filename)
+        
+        print(f"Extracting {zip_filename}...")
+        with zipfile.ZipFile(zip_path, 'r') as zf:
+            zf.extractall(self.root_dir)
+        
+        os.remove(zip_path)
+        print("Extraction complete.")
+    
+    def _load_data(self):
+        """Load preprocessed data using memory mapping."""
+        prefix = self.split
+        print(f"Loading GEOM-Drugs {self.split} split from {self.root_dir}")
+        
+        self._positions = np.load(
+            os.path.join(self.root_dir, f"{prefix}_positions.npy"),
+            mmap_mode=self.mmap_mode
+        )
+        self._species = np.load(
+            os.path.join(self.root_dir, f"{prefix}_species.npy"),
+            mmap_mode=self.mmap_mode
+        )
+        
+        # Regular arrays for metadata/indexing
+        self._offsets = np.load(os.path.join(self.root_dir, f"{prefix}_offsets.npy"))
+        self._n_atoms = np.load(os.path.join(self.root_dir, f"{prefix}_n_atoms.npy"))
+        self._mol_indices = np.load(os.path.join(self.root_dir, f"{prefix}_mol_indices.npy"))
+        
+        with open(os.path.join(self.root_dir, f"{prefix}_smiles.json")) as f:
+            self._smiles = json.load(f)
+        
+        n_molecules = len(np.unique(self._mol_indices))
+        n_conformers = len(self._n_atoms)
+        print(f"Loaded {n_molecules} molecules with {n_conformers} total conformers")
+    
+    def _setup_indices(self):
+        """Setup indices based on conformer selection mode and dataset slicing."""
+        n_conformers = len(self._n_atoms)
+        
+        if self.conformer_selection == "all":
+            self._indices = np.arange(n_conformers)
+        else:
+            unique_mols, first_indices = np.unique(self._mol_indices, return_index=True)
+            
+            if self.conformer_selection == "first":
+                self._indices = first_indices
+            elif self.conformer_selection == "random":
+                counts = np.diff(np.append(first_indices, n_conformers))
+                offsets = np.array([self._rng.integers(0, c) for c in counts])
+                self._indices = first_indices + offsets
+        
+        # Filter by atom count
+        if self.max_atoms is not None:
+            mask = self._n_atoms[self._indices] <= self.max_atoms
+            self._indices = self._indices[mask]
+        
+        # Apply start/end slicing
+        self._indices = self._indices[slice(self.start_index, self.end_index)]
+
+    def __len__(self) -> int:
+        return len(self._indices)
+
+    def __getitem__(self, idx: int) -> datatypes.Graph:
+        """Fast slice access via memory-mapped offsets."""
+        if idx < 0:
+            idx = len(self._indices) + idx
+        
+        real_idx = self._indices[idx]
+        start, end = self._offsets[real_idx], self._offsets[real_idx + 1]
+        
+        # Extract data for this conformer
+        positions = np.array(self._positions[start:end])
+        species = np.array(self._species[start:end])
+        atomic_numbers = self.species_to_atomic_numbers(species)
+        atom_types = utils.atomic_numbers_to_symbols(atomic_numbers)
+
+        return datatypes.Graph(
+            nodes=dict(
+                positions=np.array(self._positions[start:end]),
+                atomic_numbers=atomic_numbers,
+                species=species,
+                atom_types=atom_types,
+            ),
+            edges=None,
+            senders=None,
+            receivers=None,
+            n_edge=None,
+            n_node=np.asarray([self._n_atoms[real_idx]]),
+            globals=None,
+            properties=dict(smiles=self._smiles[real_idx]),
+        )
+    
+    def get_num_conformers(self, mol_idx: int) -> int:
+        """Get number of conformers for a molecule by its molecule index."""
+        return np.sum(self._mol_indices == mol_idx)
+    
+    def get_molecule_indices(self) -> np.ndarray:
+        """Get array mapping each loaded conformer back to its molecule index."""
+        return self._mol_indices[self._indices]