Metadata-Version: 2.5
Name: wyckoff-transformer
Version: 2.4.0
Summary: Wyckoff Transformer is a machine learning model for generating crystal structures which are symmetric by design.
Project-URL: Homepage, https://github.com/SymmetryAdvantage/WyckoffTransformer
Project-URL: Issues, https://github.com/SymmetryAdvantage/WyckoffTransformer/issues
Author-email: Nikita Kazeev <kazeevn@gmail.com>
License-Expression: Apache-2.0
License-File: LICENSE
Classifier: Development Status :: 5 - Production/Stable
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: <3.15,>=3.12
Requires-Dist: datasets>=4.8.5
Requires-Dist: fsspec>=2026.2.0
Requires-Dist: huggingface-hub>=1.7.2
Requires-Dist: ipympl>=0.9.7
Requires-Dist: matminer>=0.10.1
Requires-Dist: monty>=2025.3.3
Requires-Dist: numpy>=2
Requires-Dist: omegaconf>=2.3.0
Requires-Dist: pandarallel>=1.6.5
Requires-Dist: pandas>=2.1.1
Requires-Dist: pyarrow>=24.0.0
Requires-Dist: pydantic>=2.12.5
Requires-Dist: pymatgen>=2025.10.7
Requires-Dist: pyxtal==1.1.5
Requires-Dist: safetensors>=0.7.0
Requires-Dist: schedulefree>=1.4.1
Requires-Dist: scikit-learn>=1.5.0
Requires-Dist: smact>=2.5.5
Requires-Dist: torch>=2.10.0
Requires-Dist: wandb>=0.25.0
Requires-Dist: wrapt-timeout-decorator>=1.5.1
Provides-Extra: cuequivariance-cu13
Requires-Dist: cuequivariance-ops-torch-cu13>=0.9; extra == 'cuequivariance-cu13'
Requires-Dist: cuequivariance-torch>=0.9; extra == 'cuequivariance-cu13'
Requires-Dist: triton>=3; extra == 'cuequivariance-cu13'
Provides-Extra: dev
Requires-Dist: pytest>=9.0.2; extra == 'dev'
Requires-Dist: ruff>=0.15.0; extra == 'dev'
Provides-Extra: mp
Requires-Dist: mp-api>=0.46.5; extra == 'mp'
Provides-Extra: nep
Requires-Dist: calorine>=3.5; extra == 'nep'
Requires-Dist: scipy>=1.15; extra == 'nep'
Provides-Extra: publish
Requires-Dist: build; extra == 'publish'
Requires-Dist: packaging>=25; extra == 'publish'
Requires-Dist: twine; extra == 'publish'
Provides-Extra: relax
Requires-Dist: ase>=3.23; extra == 'relax'
Requires-Dist: mace-torch>=0.3; extra == 'relax'
Requires-Dist: orb-models>=0.5.1; extra == 'relax'
Provides-Extra: research
Requires-Dist: ase; extra == 'research'
Requires-Dist: ipykernel; extra == 'research'
Requires-Dist: ipympl; extra == 'research'
Requires-Dist: ipywidgets; extra == 'research'
Requires-Dist: jupyter; extra == 'research'
Requires-Dist: matbench-discovery; extra == 'research'
Description-Content-Type: text/markdown

[![Pytest](https://github.com/SymmetryAdvantage/WyckoffTransformer/actions/workflows/pytest.yml/badge.svg)](https://github.com/SymmetryAdvantage/WyckoffTransformer/actions/workflows/pytest.yml)
[![PyPI - Version](https://img.shields.io/pypi/v/wyckoff-transformer)](https://pypi.org/project/wyckoff-transformer/)
[![ICML 2025](https://img.shields.io/badge/Paper-ICML%202025-blue)](https://icml.cc/virtual/2025/poster/44595)

# Wyckoff Transformer: Generation of Symmetric Crystals [ICML 2025]
WyFormer generates crystals as Wyckoff representations: a space group plus the occupied
Wyckoff positions and their elements. PyXtal then places the atoms, and an ML potential
relaxes the result into a 3D structure.

## Benchmarking against WyFormer
[Figshare](https://figshare.com/articles/dataset/Generated_crystals_for_WyFormer_DiffCSP_DiffCSP_WyCryst_SymmCD_CrystalFormer_MiAD/29145101) hosts the structures generated by WyFormer trained on MP-20 and MPTS-52, along with DiffCSP, DiffCSP++, WyCryst, SymmCD, CrystalFormer, MiAD, MatterGen, Crys-JEPA, and Composition in a unified format, along with Materials Project-compatible DFT relaxations: ~10k structures each for WyFormer, MatterGen and MiAD; 1k each for DiffCSP, SymmCD, Crys-JEPA and Composition.

## Installation
```bash
pip install "wyckoff-transformer[relax]"
```
The `[relax]` extra adds the ML potentials for the 3D step, skip it if you only need the Wyckoff representations.

## Inference
### Simple de novo generation
#### Step 1: Generate genes
```bash
wyformer-generate genes.json.gz --hf-model SymmetryAdvantage/WyFormer-LeMat
```
For a higher yield of new metastable structures, ask the e_hull-conditioned model for genes
near the hull instead. About 61% of its draws are formally valid at this guidance scale, so
draw twice as many as you keep:
```bash
wyformer-generate genes.json.gz --hf-model SymmetryAdvantage/WyFormer-LeMat-Ehull \
    --condition energy_above_hull=0.05 --guidance-scale 5 \
    --initial-n-samples 2000 --firm-n-samples 1000
```
Add `--device cuda` to the `wyformer-generate` and `wyformer-gene-screen` commands when a GPU is available. The
pre-trained models are in the [HuggingFace collection](https://huggingface.co/collections/SymmetryAdvantage/wyformer).
#### Step 2: Relax
```bash
wyformer-cryspr genes.json.gz --mlip orb_conserv_inf --output-dir relaxed --workers 8
```
Alternatively, the structures can be obtained with [GeoCSP](https://github.com/kazeevn/GeoCSP), our diffusion model based on DiffCSP++.
### Advanced generation, chemical system conditioning
Generate genes that contain only the given elements, rank them by predicted distance to the
PBE hull of [LeMat-Bulk-fixed](https://huggingface.co/datasets/kazeevn/LeMat-Bulk-fixed), and
relax the 100 best distinct ones of at most 40 atoms:
```bash
wyformer-generate batio_genes.json.gz --hf-model SymmetryAdvantage/WyFormer-LeMat-Chemsys \
    --chemical-system Ba-Ti-O --initial-n-samples 1100 --firm-n-samples 1000
wyformer-gene-screen batio_genes.json.gz --regressor-hf-model SymmetryAdvantage/WyFormer-LeMat-Energy \
    --hf-reference kazeevn/LeMat-Bulk-fixed \
    --max-atoms 40 --unique --top 100 --out batio_screen.csv --out-genes batio_top100.json.gz
wyformer-cryspr batio_top100.json.gz --mlip orb_conserv_inf --output-dir batio_relaxed --workers 8
```
`--chemical-system Ba-Ti-O` asks for compounds that contain all three elements; run TiO2 and
the other subsystems separately. `--max-atoms 40` drops genes with more than 40 atoms in the
conventional cell: the lowest predicted energies favour large cells, which are slow to relax.
`batio_screen.csv` holds the predicted energies, and
`batio_relaxed/orb_conserv_inf_results.csv` the relaxed energy and CIF of each gene.
ORB v3 is the potential LeMat-Bulk publishes its MLIP hull for, so those energies compare
directly with that hull. See [the pipeline notes](docs/generation.md#the-chemical-system-pipeline)
for timings and outputs.

## Documentation
- [Generating and relaxing structures](docs/generation.md): other model sources, generation
  options, PyXtal placement and MLIP relaxation, CHGNet and DFT.
- [Training](docs/training.md): installation from source, datasets, tokenization, training.
- [Screening and evaluation](docs/evaluation.md): DFT and novelty screens, the de novo
  ranking protocol, reproducing the original ICML paper metrics.
