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).
