StructureCloud 0.0.1__tar.gz → 0.0.2__tar.gz

structurecloud-0.0.2/PKG-INFO

@@ -0,0 +1,201 @@
+Metadata-Version: 2.4
+Name: StructureCloud
+Version: 0.0.2
+Summary: Standard formating and easy access to 3D structural datasets for machine learning. currently under development...
+Author-email: Ty Perez <ty.jperez@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/TyJPerez/StructureCloud
+Project-URL: Issues, https://github.com/TyJPerez/StructureCloud/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: torch_geometric
+Requires-Dist: huggingface_hub>=0.20
+Requires-Dist: numpy>=1.21
+Requires-Dist: tqdm>=4.60
+Provides-Extra: dev
+Requires-Dist: matplotlib>=3.5; extra == "dev"
+Requires-Dist: scipy>=1.7; extra == "dev"
+Requires-Dist: rdkit>=2022.3; extra == "dev"
+Requires-Dist: datasets>=2.0; extra == "dev"
+Provides-Extra: viz
+Requires-Dist: matplotlib>=3.5; extra == "viz"
+Requires-Dist: plotly>=5.0; extra == "viz"
+Dynamic: license-file
+
+# StructureCloud
+
+A collection of 3D atomistic (point-cloud) datasets for machine learning on molecules and
+materials — property prediction, representation learning, and generative modeling. Every
+dataset is exposed through a single, uniform loader that returns
+[PyTorch Geometric](https://pytorch-geometric.readthedocs.io/) `Data` objects, and is stored
+in a **memory-mapped, collated format** so that many large datasets can be **stacked and
+trained on together without exhausting RAM**.
+
+## Why StructureCloud
+
+Atomistic datasets ship in many incompatible formats (jsonl, LMDB, parquet, per-source
+schemas) and decoding them one row at a time is slow and RAM-hungry. StructureCloud solves
+this with three conventions shared across every dataset:
+
+- **One uniform sample schema.** Each item is a PyG `Data` with `pos` `[N,3]` (Å) and `z`
+  `[N]` (atomic numbers); periodic crystals add `cell` `[1,3,3]` and `pbc` `[1,3]`. Labels
+  carry short, descriptive keys — never a bare `y`/`dy`.
+- **Collated + memory-mapped storage.** Each split is preprocessed once into a single
+  `.pt` file where *everything is a tensor* (variable-length atom arrays are concatenated
+  with offset tables; strings are packed). `torch.load(mmap=True)` pages tensors in from
+  disk through the shared OS page cache, so stacking N datasets costs the combined working
+  set, not the sum of their sizes.
+- **Download-only loaders.** Preprocessed files live on the Hugging Face Hub under
+  `StructureCloud/{dataset}`; the loader downloads the split it needs and maps it — no
+  building from raw at load time.
+
+## Install
+
+```bash
+pip install StructureCloud
+```
+
+Or from source (editable), e.g. for development or to preprocess new datasets:
+
+```bash
+git clone https://github.com/TyJPerez/StructureCloud.git
+cd StructureCloud
+pip install -e .
+```
+
+**Dependencies.** The core install pulls in everything needed to load and stack datasets:
+`torch`, `torch_geometric`, `huggingface_hub`, `numpy`, and `tqdm`. Split files download
+automatically from the Hugging Face Hub on first use and are cached locally.
+
+> **Install `torch` and `torch_geometric` first, matched to your platform.** Their wheels are
+> tied to your CUDA/CPU setup, so they are intentionally left unpinned here — follow the
+> official [PyTorch](https://pytorch.org/get-started/locally/) and
+> [PyG](https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html)
+> install instructions for your system, then install StructureCloud.
+
+Two optional extras cover the tooling beyond loading:
+
+```bash
+pip install "StructureCloud[dev]"   # preprocessing + figures/stats: matplotlib, scipy, rdkit, datasets
+pip install "StructureCloud[viz]"   # 3D structure/atom plotting in StructureCloud.chem_tools: matplotlib, plotly
+```
+
+## Quickstart
+
+```python
+from StructureCloud.Datasets import QM9
+
+# Downloads and memory-maps StructureCloud/QM9 preprocessed 'train' split
+ds = QM9(split='train')          # pass mmap=False to load fully into RAM
+print(len(ds))                   # number of structures
+
+data = ds[0]                     # a torch_geometric.data.Data
+data.pos                         # [N, 3] float32 positions (Angstrom)
+data.z                           # [N]    int64 atomic numbers
+data.smiles                      # str metadata
+data.y                           # (19,) DFT property vector (QM9-specific)
+```
+
+Every loader takes the common arguments `split`, `transform`, `mmap`, and `root` (custom
+cache location); dataset-specific arguments (`label`, `subset`, `components`, `sample`, …)
+are documented in each loader's docstring and in that dataset's `README.md`. Batch with a
+standard PyG `DataLoader`:
+
+```python
+from torch_geometric.loader import DataLoader
+loader = DataLoader(ds, batch_size=64, shuffle=True)
+```
+
+## Available datasets
+
+Import any of these from `StructureCloud.Datasets`.
+
+### Small organic molecules (molecular, non-periodic)
+
+| Dataset | Size | Description | Splits |
+|---|---|---|---|
+| `QM9` | 130,831 molecules | CHONF, ≤9 heavy atoms; relaxed geometry + 19 DFT properties | `train` |
+| `PCQM4Mv2` (alias `PCQ`) | 3,378,550 molecules | DFT-relaxed geometry + HOMO–LUMO gap (OGB-LSC / PubChemQC) | `train` |
+| `GEOM` | ~37M conformers / ~450k molecules | Conformer **ensemble** (energies + Boltzmann weights); variants `GEOM`/`GEOM64`/`GEOM32`/`GEOM10`/`GEOM1`, collections `drugs`/`qm9`/`molnet` | per-collection |
+| `SPICE1` / `SPICE2` | 19,238 mol / 1.11M conf · 113,985 mol / 2.0M conf | QC single-points with per-atom **forces** + per-conformer energies (ωB97M-D3(BJ)/def2-TZVPPD); conformer ensemble | `all` |
+| `MD17` | 10 molecules × ~100k frames | rMD17 MD trajectories; energy + forces at PBE/def2-SVP; each molecule is its own split | per-molecule |
+
+### Materials (periodic crystals)
+
+| Dataset | Size | Description | Splits |
+|---|---|---|---|
+| `MP20` | 45,229 crystals | ≤20-atom Materials Project crystals; standard generation benchmark | `train`/`val`/`test` |
+| `AlexMP20` | 675,204 crystals | ≤20-atom crystals from Alexandria + MP; stability, band gap, elastic/magnetic properties | `train`/`val`/`test`/`all` |
+| `Carbon24` | 10,153 crystals | Carbon-only crystals from ab-initio random structure search; energy/atom label | `train`/`val`/`test` |
+| `Perov5` | 18,928 crystals | Perovskite ABX₃, fixed 5-atom cell; formation energy + band gap | `train`/`val`/`test` |
+| `WBM` | 256,963 crystals | Matbench Discovery test set; unrelaxed→relaxed formation energy / stability (85 elements) | `test` |
+| `Matbench` | 9 tasks | The 9 structure-based Matbench property tasks; fixed 5-fold CV (`mp_gap`, `mp_e_form`, `perovskites`, `phonons`, `dielectric`, `jdft2d`, `log_gvrh`, `log_kvrh`, `mp_is_metal`) | per-task |
+
+### Proteins and biomolecules
+
+| Dataset | Size | Description | Splits |
+|---|---|---|---|
+| `LBA` | 4,463 complexes | Protein–ligand binding affinity (Atom3D / PDBbind); full protein preserved, available as `protein`/`pocket`/`ligand` components | `train`/`val`/`test`/`all` |
+
+## Stacking datasets
+
+The core feature: concatenate arbitrary datasets into one memory-mapped training set with a
+uniform key schema. `StackedDataset` tags each sample's origin (`AddTag`) and standardizes
+keys (`StandardizeKeys` fills a `cell` bounding box / `pbc` for molecular samples and a
+unified `str_id`), so a mixed molecular + periodic batch collates through a single PyG
+`DataLoader`.
+
+```python
+from StructureCloud.Datasets import AlexMP20, GEOM, PCQM4Mv2, StackedDataset
+
+stk = StackedDataset(
+    [AlexMP20('all', label=None), GEOM('GEOM10', 'drugs'), PCQM4Mv2()],
+    tags=['alexmp20', 'geom10', 'pcq'],
+    keep_keys=['pos', 'z', 'cell', 'pbc', 'natoms', 'str_id', 'tag'],
+)
+```
+
+Because every dataset is memory-mapped, a stack that is several GB on disk stays at a small
+resident footprint, with cold pages evicted under memory pressure and shared across
+DataLoader workers.
+
+## Conformer-ensemble datasets
+
+`GEOM` and `SPICE` carry many conformers per molecule. Their loaders return one conformer
+per access with selectable sampling: `sample=True` (random conformer, default),
+`sample=False` (flatten — index every molecule/conformer pair), or `sample_fn=…` with a
+shared sampler from `StructureCloud.Datasets.samplers` (`sample_by_relative`,
+`sample_by_boltzmann`).
+
+## Roadmap
+
+- **OMol25** (Meta FAIR Open Molecules 2025) — large molecular DFT single-point dataset
+  (energy + forces + variable charge/spin at ωB97M-V/def2-TZVPD). *In progress.*
+- **QMOF** — periodic metal–organic-framework DFT properties (~20,372 MOFs; band gaps,
+  per-atom charges/magmoms, pore descriptors). *Planned.*
+- **3D objects** and **Gaussian splats** — planned.
+
+## Also in this repo
+
+Beyond the datasets, the package includes early/experimental modules that are not the focus
+of this README: `StructureCloud.models` (an equivariant-transformer architecture),
+`StructureCloud.utils` (graph construction, augmentation, structure viewing), and
+`StructureCloud.chem_tools` (3D atom/material plotting, atomic properties).
+
+## Contributing a dataset
+
+Dataset conventions, the collated/memory-mapped format, and the step-by-step process for
+adding a new dataset are documented under `.agents/docs/` — start with
+[`about.md`](.agents/docs/about.md) and
+[`adding_dataset/guidelines.md`](.agents/docs/adding_dataset/guidelines.md). New loaders
+subclass `CollatedStructureDataset` (`Datasets/collated.py`), ship a `README.md`, a
+distribution figure, and a `stats.json`, and register in `Datasets/__init__.py` and the
+verification harness (`scripts/verify_dataset.py`).
+
+## License
+
+MIT — see [LICENSE](LICENSE).

structurecloud-0.0.2/README.md

@@ -0,0 +1,173 @@
+# StructureCloud
+
+A collection of 3D atomistic (point-cloud) datasets for machine learning on molecules and
+materials — property prediction, representation learning, and generative modeling. Every
+dataset is exposed through a single, uniform loader that returns
+[PyTorch Geometric](https://pytorch-geometric.readthedocs.io/) `Data` objects, and is stored
+in a **memory-mapped, collated format** so that many large datasets can be **stacked and
+trained on together without exhausting RAM**.
+
+## Why StructureCloud
+
+Atomistic datasets ship in many incompatible formats (jsonl, LMDB, parquet, per-source
+schemas) and decoding them one row at a time is slow and RAM-hungry. StructureCloud solves
+this with three conventions shared across every dataset:
+
+- **One uniform sample schema.** Each item is a PyG `Data` with `pos` `[N,3]` (Å) and `z`
+  `[N]` (atomic numbers); periodic crystals add `cell` `[1,3,3]` and `pbc` `[1,3]`. Labels
+  carry short, descriptive keys — never a bare `y`/`dy`.
+- **Collated + memory-mapped storage.** Each split is preprocessed once into a single
+  `.pt` file where *everything is a tensor* (variable-length atom arrays are concatenated
+  with offset tables; strings are packed). `torch.load(mmap=True)` pages tensors in from
+  disk through the shared OS page cache, so stacking N datasets costs the combined working
+  set, not the sum of their sizes.
+- **Download-only loaders.** Preprocessed files live on the Hugging Face Hub under
+  `StructureCloud/{dataset}`; the loader downloads the split it needs and maps it — no
+  building from raw at load time.
+
+## Install
+
+```bash
+pip install StructureCloud
+```
+
+Or from source (editable), e.g. for development or to preprocess new datasets:
+
+```bash
+git clone https://github.com/TyJPerez/StructureCloud.git
+cd StructureCloud
+pip install -e .
+```
+
+**Dependencies.** The core install pulls in everything needed to load and stack datasets:
+`torch`, `torch_geometric`, `huggingface_hub`, `numpy`, and `tqdm`. Split files download
+automatically from the Hugging Face Hub on first use and are cached locally.
+
+> **Install `torch` and `torch_geometric` first, matched to your platform.** Their wheels are
+> tied to your CUDA/CPU setup, so they are intentionally left unpinned here — follow the
+> official [PyTorch](https://pytorch.org/get-started/locally/) and
+> [PyG](https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html)
+> install instructions for your system, then install StructureCloud.
+
+Two optional extras cover the tooling beyond loading:
+
+```bash
+pip install "StructureCloud[dev]"   # preprocessing + figures/stats: matplotlib, scipy, rdkit, datasets
+pip install "StructureCloud[viz]"   # 3D structure/atom plotting in StructureCloud.chem_tools: matplotlib, plotly
+```
+
+## Quickstart
+
+```python
+from StructureCloud.Datasets import QM9
+
+# Downloads and memory-maps StructureCloud/QM9 preprocessed 'train' split
+ds = QM9(split='train')          # pass mmap=False to load fully into RAM
+print(len(ds))                   # number of structures
+
+data = ds[0]                     # a torch_geometric.data.Data
+data.pos                         # [N, 3] float32 positions (Angstrom)
+data.z                           # [N]    int64 atomic numbers
+data.smiles                      # str metadata
+data.y                           # (19,) DFT property vector (QM9-specific)
+```
+
+Every loader takes the common arguments `split`, `transform`, `mmap`, and `root` (custom
+cache location); dataset-specific arguments (`label`, `subset`, `components`, `sample`, …)
+are documented in each loader's docstring and in that dataset's `README.md`. Batch with a
+standard PyG `DataLoader`:
+
+```python
+from torch_geometric.loader import DataLoader
+loader = DataLoader(ds, batch_size=64, shuffle=True)
+```
+
+## Available datasets
+
+Import any of these from `StructureCloud.Datasets`.
+
+### Small organic molecules (molecular, non-periodic)
+
+| Dataset | Size | Description | Splits |
+|---|---|---|---|
+| `QM9` | 130,831 molecules | CHONF, ≤9 heavy atoms; relaxed geometry + 19 DFT properties | `train` |
+| `PCQM4Mv2` (alias `PCQ`) | 3,378,550 molecules | DFT-relaxed geometry + HOMO–LUMO gap (OGB-LSC / PubChemQC) | `train` |
+| `GEOM` | ~37M conformers / ~450k molecules | Conformer **ensemble** (energies + Boltzmann weights); variants `GEOM`/`GEOM64`/`GEOM32`/`GEOM10`/`GEOM1`, collections `drugs`/`qm9`/`molnet` | per-collection |
+| `SPICE1` / `SPICE2` | 19,238 mol / 1.11M conf · 113,985 mol / 2.0M conf | QC single-points with per-atom **forces** + per-conformer energies (ωB97M-D3(BJ)/def2-TZVPPD); conformer ensemble | `all` |
+| `MD17` | 10 molecules × ~100k frames | rMD17 MD trajectories; energy + forces at PBE/def2-SVP; each molecule is its own split | per-molecule |
+
+### Materials (periodic crystals)
+
+| Dataset | Size | Description | Splits |
+|---|---|---|---|
+| `MP20` | 45,229 crystals | ≤20-atom Materials Project crystals; standard generation benchmark | `train`/`val`/`test` |
+| `AlexMP20` | 675,204 crystals | ≤20-atom crystals from Alexandria + MP; stability, band gap, elastic/magnetic properties | `train`/`val`/`test`/`all` |
+| `Carbon24` | 10,153 crystals | Carbon-only crystals from ab-initio random structure search; energy/atom label | `train`/`val`/`test` |
+| `Perov5` | 18,928 crystals | Perovskite ABX₃, fixed 5-atom cell; formation energy + band gap | `train`/`val`/`test` |
+| `WBM` | 256,963 crystals | Matbench Discovery test set; unrelaxed→relaxed formation energy / stability (85 elements) | `test` |
+| `Matbench` | 9 tasks | The 9 structure-based Matbench property tasks; fixed 5-fold CV (`mp_gap`, `mp_e_form`, `perovskites`, `phonons`, `dielectric`, `jdft2d`, `log_gvrh`, `log_kvrh`, `mp_is_metal`) | per-task |
+
+### Proteins and biomolecules
+
+| Dataset | Size | Description | Splits |
+|---|---|---|---|
+| `LBA` | 4,463 complexes | Protein–ligand binding affinity (Atom3D / PDBbind); full protein preserved, available as `protein`/`pocket`/`ligand` components | `train`/`val`/`test`/`all` |
+
+## Stacking datasets
+
+The core feature: concatenate arbitrary datasets into one memory-mapped training set with a
+uniform key schema. `StackedDataset` tags each sample's origin (`AddTag`) and standardizes
+keys (`StandardizeKeys` fills a `cell` bounding box / `pbc` for molecular samples and a
+unified `str_id`), so a mixed molecular + periodic batch collates through a single PyG
+`DataLoader`.
+
+```python
+from StructureCloud.Datasets import AlexMP20, GEOM, PCQM4Mv2, StackedDataset
+
+stk = StackedDataset(
+    [AlexMP20('all', label=None), GEOM('GEOM10', 'drugs'), PCQM4Mv2()],
+    tags=['alexmp20', 'geom10', 'pcq'],
+    keep_keys=['pos', 'z', 'cell', 'pbc', 'natoms', 'str_id', 'tag'],
+)
+```
+
+Because every dataset is memory-mapped, a stack that is several GB on disk stays at a small
+resident footprint, with cold pages evicted under memory pressure and shared across
+DataLoader workers.
+
+## Conformer-ensemble datasets
+
+`GEOM` and `SPICE` carry many conformers per molecule. Their loaders return one conformer
+per access with selectable sampling: `sample=True` (random conformer, default),
+`sample=False` (flatten — index every molecule/conformer pair), or `sample_fn=…` with a
+shared sampler from `StructureCloud.Datasets.samplers` (`sample_by_relative`,
+`sample_by_boltzmann`).
+
+## Roadmap
+
+- **OMol25** (Meta FAIR Open Molecules 2025) — large molecular DFT single-point dataset
+  (energy + forces + variable charge/spin at ωB97M-V/def2-TZVPD). *In progress.*
+- **QMOF** — periodic metal–organic-framework DFT properties (~20,372 MOFs; band gaps,
+  per-atom charges/magmoms, pore descriptors). *Planned.*
+- **3D objects** and **Gaussian splats** — planned.
+
+## Also in this repo
+
+Beyond the datasets, the package includes early/experimental modules that are not the focus
+of this README: `StructureCloud.models` (an equivariant-transformer architecture),
+`StructureCloud.utils` (graph construction, augmentation, structure viewing), and
+`StructureCloud.chem_tools` (3D atom/material plotting, atomic properties).
+
+## Contributing a dataset
+
+Dataset conventions, the collated/memory-mapped format, and the step-by-step process for
+adding a new dataset are documented under `.agents/docs/` — start with
+[`about.md`](.agents/docs/about.md) and
+[`adding_dataset/guidelines.md`](.agents/docs/adding_dataset/guidelines.md). New loaders
+subclass `CollatedStructureDataset` (`Datasets/collated.py`), ship a `README.md`, a
+distribution figure, and a `stats.json`, and register in `Datasets/__init__.py` and the
+verification harness (`scripts/verify_dataset.py`).
+
+## License
+
+MIT — see [LICENSE](LICENSE).

structurecloud-0.0.2/pyproject.toml

@@ -0,0 +1,51 @@
+[project]
+name = "StructureCloud"          # must be unique on PyPI!
+version = "0.0.2"
+authors = [
+  { name="Ty Perez", email="ty.jperez@gmail.com" },
+]
+
+description = "Standard formating and easy access to 3D structural datasets for machine learning. currently under development..."
+readme = "README.md"
+requires-python = ">=3.9"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+
+license = "MIT"
+license-files = ["LICEN[CS]E*"]
+
+# Core runtime dependencies: everything needed to load and stack datasets.
+# torch and torch_geometric are left unpinned on purpose: their wheels are
+# CUDA/platform-specific and users normally install them for their own system
+# first (see the README). Everything else gets a conservative lower-bound floor
+# (no upper caps) so the package coexists with a user's existing environment.
+dependencies = [
+    "torch",
+    "torch_geometric",
+    "huggingface_hub>=0.20",
+    "numpy>=1.21",
+    "tqdm>=4.60",
+]
+
+[project.optional-dependencies]
+# Preprocessing new datasets and generating distribution figures / stats.json.
+dev = [
+    "matplotlib>=3.5",
+    "scipy>=1.7",
+    "rdkit>=2022.3",
+    "datasets>=2.0",
+]
+# 3D structure/atom plotting helpers in StructureCloud.chem_tools.
+viz = [
+    "matplotlib>=3.5",
+    "plotly>=5.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/TyJPerez/StructureCloud"
+Issues = "https://github.com/TyJPerez/StructureCloud/issues"
+
+[tool.setuptools.packages.find]
+where = ["src"]  # <- Important to tell setuptools to find code inside src/

structurecloud-0.0.2/src/StructureCloud/Datasets/AlexMP20/alexmp20.py

@@ -0,0 +1,103 @@
+from typing import Callable, Optional
+
+import numpy as np
+import torch
+from torch.utils.data import ConcatDataset
+from torch_geometric.data import Data
+
+from ..collated import CollatedStructureDataset
+
+
+class AlexMP20(CollatedStructureDataset):
+    '''Fast, memory-mapped loader for the preprocessed Alex-MP-20 dataset.
+
+    Downloads a collated ``.pt`` per split from ``StructureCloud/AlexMP20`` and
+    memory-maps it (see ``CollatedStructureDataset``); pass ``mmap=False`` to load
+    fully into RAM.
+
+    args:
+        split : one of ['train', 'val', 'test', 'all']  ('all' concatenates the three)
+        label : target label, one of ``all_labels`` (or None for no target)
+        transform : optional callable applied to each Data object
+        mmap : memory-map the collated file (default True)
+    '''
+
+    # available labels and their keys in the collated dataset
+    all_labels = {
+        'space_group': 'space_group',
+        'chemical_system': 'chemical_system',
+        'energy_above_hull': 'energy_above_hull',
+        'band_gap': 'dft_band_gap',
+        'bulk_modulus': 'dft_bulk_modulus',
+        'mag_density': 'dft_mag_density',
+        'hhi_score': 'hhi_score',
+        'ml_bulk_modulus': 'ml_bulk_modulus',
+        'id': 'ids',
+    }
+
+    hf_repo = 'StructureCloud/AlexMP20'
+    available_splits = ['train', 'val', 'test', 'all']
+    string_fields = ('ids', 'space_group', 'chemical_system')
+
+    def __init__(self,
+                 split: str = 'train',
+                 label: Optional[str] = 'band_gap',
+                 transform: Optional[Callable] = None,
+                 mmap: bool = True,
+                 root: Optional[str] = None,
+                 ):
+        assert split in self.available_splits, (
+            f"Split {split} not recognized. Available splits: {self.available_splits}"
+        )
+        if label is not None:
+            assert label in self.all_labels, (
+                f"Label {label} not recognized. Available labels: {list(self.all_labels)}"
+            )
+        self.label = label
+
+        if split == 'all':
+            # concatenate the three splits, each individually memory-mapped
+            self.split = 'all'
+            self.transform = transform
+            self._parts = [
+                AlexMP20(s, label=label, transform=None, mmap=mmap, root=root)
+                for s in ['train', 'val', 'test']
+            ]
+            self._concat = ConcatDataset(self._parts)
+            self.num_structures = len(self._concat)
+            self.dataset = None
+            return
+
+        super().__init__(split, transform=transform, mmap=mmap, root=root)
+
+    def _label_value(self, idx):
+        key = self.all_labels[self.label]
+        if key in self.string_fields:
+            return self.decode(key, idx)
+        return self.dataset[key][idx].clone()
+
+    def __getitem__(self, idx):
+        if isinstance(idx, torch.Tensor):
+            idx = idx.item()
+        idx = int(idx)
+
+        if self.split == 'all':
+            data = self._concat[idx]           # parts carry no transform
+            if self.transform is not None:
+                data = self.transform(data)
+            return data
+
+        out_data = Data()
+        out_data.pos = self.slice_by_ptr('pos', idx)
+        out_data.cell = self.dataset['cell'][idx].reshape(1, 3, 3).clone()
+        out_data.z = self.slice_by_ptr('z', idx)
+        out_data.pbc = self.dataset['pbc'][idx].reshape(1, 3).clone()
+
+        out_data.id = self.decode('ids', idx)
+        if self.label is not None:
+            out_data.y = self._label_value(idx)
+
+        if self.transform is not None:
+            out_data = self.transform(out_data)
+
+        return out_data