Metadata-Version: 2.4
Name: wyckoff-transformer
Version: 1.0.7
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.13,>=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
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
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: cdvae
Requires-Dist: cdvae-property-models; extra == 'cdvae'
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: 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'
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]
## Installation (PyPI)
WyFormer is published on PyPI. Be mindful of your PyTorch situation, and install:
```bash
pip install wyckoff-transformer
```

## Inference
The pre-trained models are published on [HuggingFace](https://huggingface.co/collections/SymmetryAdvantage/wyformer). To use a HuggingFace model run:
```bash
wyformer-generate <output-file.json.gz> --hf-model <model-name>
```
Using a local model directory (must contain `best_model_params.pt`, `config.yaml`, and `wyckoff_processor.json`):
```bash
wyformer-generate <output-file.json.gz> --model-path runs/<run-id>
```

See this [repository](https://github.com/SymmetryAdvantage/WyFormer-inference-demo) for a demo of standalone inference using the PyPI package as a library.

## WyFormer Generated Datasets
If you just need the generated datasets for benchmarking, they are available at Figshare, including both original and DFT: [WyFormer](https://figshare.com/articles/dataset/WyFormer_generated_structures/29094701); [WyFormer, DiffCSP(++), SymmCD, MiAD, WyCryst, CrystalFormer](https://figshare.com/articles/dataset/Generated_crystals_for_WyFormer_DiffCSP_DiffCSP_WyCryst_SymmCD_CrystalFormer_MiAD/29145101).

## Abstract
Crystal symmetry plays a fundamental role in determining its physical, chemical, and electronic properties such as electrical and thermal conductivity, optical and polarization behavior, and mechanical strength. Almost all known crystalline materials have internal symmetry. However, this is often inadequately addressed by existing generative models, making the consistent generation of stable and symmetrically valid crystal structures a significant challenge. We introduce WyFormer, a generative model that directly tackles this by formally conditioning on space group symmetry. It achieves this by using Wyckoff positions as the basis for an elegant, compressed, and discrete structure representation. To model the distribution, we develop a permutation-invariant autoregressive model based on the Transformer encoder and an absence of positional encoding. Extensive experimentation demonstrates WyFormer's compelling combination of attributes: it achieves best-in-class symmetry-conditioned generation, incorporates a physics-motivated inductive bias, produces structures with competitive stability, predicts material properties with competitive accuracy even without atomic coordinates, and exhibits unparalleled inference speed.

# Local development & training
## Installation
1. Clone the repository
2. Run `uv venv --python 3.12`
3. Install the dependencies, including torch. There are several options:
  - Manually install torch with your local flavour, e.g., `uv pip install torch --index-url https://download.pytorch.org/whl/cu130`, then run `uv pip install -e`
  - Configure `uv.toml` with your desired indices, see `uv.toml.local` and `uv.toml.cpu`
4. `wandb` library is used extensively and must be installed. Logging can be disabled via `WANDB_MODE=disabled`. Otherwise, log into Wandb. Internally, we use `WANDB_ENTITY=symmetry-advantage`.
## Running a pilot model
To verify that the installation is working, run a pilot model. Next token prediction:
```bash
python scripts/cache_a_dataset.py mp_20
python scripts/tokenise_a_dataset.py mp_20 yamls/tokenisers/mp_20_sg_multiplicity.yaml --new-tokenizer
python scripts/train.py yamls/models/NextToken/v6/base_sg.yaml mp_20 cuda --pilot
```
This will train a model, and save the results in the `runs` folder. The files are:
- `best_model_params.pt` - the model weights chosen by the validation loss
- `config.yaml` - the configuration used for training
- `wyckoff_processor.json` - tokenizers and preprocessing metadata (token engineers)
- `generated_wp_no_calibration.json.gz` - Wyckoff representation of the generated structures (if configured to evaluate generation)
- `generated_wp_temperature_calibration.json.gz` - Wyckoff representation of the generated structures, with the temperature calibration applied (if configured to evaluate generation)
## Training Data Preprocessing
The available datasets correspond to the folders in `data` and `cdvae/data`. Dataset idetifiers are the folder names, they are used throught the project. Note that some of the folders are symlinks.
Available datasets (in GitHub): `alex_mp_20`, `mp_20`, `mp_20_biternary` (binary and ternary structures from MP-20), `mpts_52`, `carbon_24`, `perov_5`. It is also possible to download and use `matbench_discovery_mp_2022` [notebook](scripts/data_preprocesssing/mp_2022.ipynb) and `matbench_discovery_mp_trj_full` [notebook](scripts/data_preprocesssing/mptrj_extract_all.ipynb).

For any data to be used for training, we need to do two preprocessing steps.
### Compute and cache symmetry information
```bash
python scripts/cache_a_dataset.py <dataset-name>
```
This will create a pickled representaiton of the dataset in `cache/<dataset-name>/data.pkl.gz`. The script supports setting symmetry tolerance and _this is not done automatically_, the datasets which include tolerance in their name were obtained by manually using the command-line option.
### Tokenization
The tokenization script serves two purposes: it produces the mapping from the real data to token ids, and saves the resulting tensors. To produce a new tokenizer:
```bash
python scripts/tokenise_a_dataset.py <dataset-name> <path-to-tokenizer-yaml> --new-tokenizer
```
Tokenizer configs are stored in `yamls/tokenisers`. The processor is saved to `cache/<dataset-names>/tokenisers/**.json`, preserving the folder structure of the config.

Alternatively, you can use a cached tokeniser. This is important when a model that was trained on one dataset is  applied to a different dataset.
```bash
python scripts/tokenise_a_dataset.py <dataset-name> <path-to-tokenizer-yaml> --tokenizer-path cache/<dataset-names>/tokenisers/<tokenizer-name>.json
```
## Training
```bash
python scripts/train.py <path-to-model-yaml> <dataset-name> <device>
```
The model weights are saved to `runs/<run-id>`, and to WanDB, along with the processor metadata. See [here](yamls/models/README.md) for the list of configs. Adding `--pilot` will run the model for a small number of epochs.
## Preparing Representative Checkpoints
To train and prepare representative checkpoints for datasets like `alex_mp_20` or `mp_20`, you can follow this end-to-end pipeline. Please replace `<dataset-name>` with your target dataset (e.g., `alex_mp_20` or `mp_20`).

First, cache and tokenize the dataset:

```bash
python scripts/cache_a_dataset.py <dataset-name>
python scripts/tokenise_a_dataset.py <dataset-name> yamls/tokenisers/<dataset-name>_sg_multiplicity.yaml --new-tokenizer
```
Before training, ensure that your model configuration file points to the correct tokenizer. For example, in `yamls/models/NextToken/v6/base_sg_schedule_free.yaml`, update the tokenizer name to match your dataset:

```yaml
tokeniser:
  name: <dataset-name>_sg_multiplicity
```
Then, initiate the training run:

```bash
python scripts/train.py yamls/models/NextToken/v6/base_sg_schedule_free.yaml <dataset-name> <device>
```
## Generating structures
### Wyckoff representations
Wyckoff representations are produced and stored in WanDB during model training. The `wyformer-generate` CLI (installed with the package) generates them from a trained model. Using a HuggingFace model:
```bash
wyformer-generate <output-file> --hf-model SymmetryAdvantage/<model-name>
```
Using a W&B run:
```bash
wyformer-generate <output-file> --wandb-run <wandb-id> --use-cached-tensors
```
Using a local model directory (must contain `best_model_params.pt`, `config.yaml`, and `wyckoff_processor.json`):
```bash
wyformer-generate <output-file> --model-path runs/<run-id>
```
Note that the code does not automatically download Wandb artifacts, you need to do it manually, and place them in the `runs` folder when restoring via run ID.

To constrain generation to specific elements:
```bash
wyformer-generate <output-file> --hf-model SymmetryAdvantage/<model-name> --required-elements Li-S --allowed-elements Li-S-P-O
```
To override the space group distribution from a cached dataset:
```bash
wyformer-generate <output-file> --hf-model SymmetryAdvantage/<model-name> --sg-dist mp_20
```

### 3D Structures
There are two ways to generate 3D structures from Wyckoff representations: DiffCSP++ and CHGNet. They later can be relaxed with CHGNet and/or DFT.
#### DiffCSP++
Wyckoffs can be relaxed with modified [DiffCSP++ code](https://github.com/kazeevn/DiffCSPNew/tree/master)
#### CrySPR + MACE
[CrySPR](https://chemrxiv.org/engage/chemrxiv/article-details/66b308a501103d79c5fd9b91) scheme using [`pyxtal`](https://pyxtal.readthedocs.io/en/latest/index.html) and a [MACE](https://github.com/ACEsuit/mace) ML force field is integrated directly into the package. Install the optional extra first:
```bash
pip install "wyckoff-transformer[relax]"
```
Then run from the command line:
```bash
wyformer-cryspr WyckoffTransformer_mp_20.json \
    --model https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model \
    --output-dir results/ --start 0 --end 1000
# model_name defaults to the model file stem
head results/2023-12-10-mace-128-L0_energy_epoch-249_results.csv
model,id,formula,energy,energy_per_atom
...
2023-12-10-mace-128-L0_energy_epoch-249,35,H6O8Si2,-97.98,...
```
To process multiple Wyckoff genes in parallel, increase the worker count:
```bash
wyformer-cryspr WyckoffTransformer_mp_20.json \
  --model https://github.com/ACEsuit/mace-foundations/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model \
  --output-dir results_parallel/ --start 0 --end 1000 --workers 8
```
URL-based models are downloaded once and cached in `~/.cache/wyckoff_transformer/mace_models/`. A local path is accepted too: `--model /path/to/model.model`.

With `--workers > 1`, the CLI uses spawned worker processes. Each worker constructs its own MACE calculator, and the process environment is forced to single-threaded BLAS/OpenMP settings (`OMP_NUM_THREADS=1`, `MKL_NUM_THREADS=1`, `OPENBLAS_NUM_THREADS=1`) to avoid CPU oversubscription. This is usually the right mode for CPU-bound batch relaxation. On GPU, start conservatively and scale `--workers` only if the device has enough memory for multiple concurrent calculators.

Output layout is identical to the CHGNet variant below. Key options:
- `--n-trials N` — number of random PyXtal trials per structure (default 6)
- `--fmax F` — force convergence criterion in eV/Å (default 0.01)
- `--model-name NAME` — label for the results CSV (default: model file stem)
- `--device auto|cpu|cuda` — PyTorch device selection (default: auto)
- `--workers W` — number of parallel worker processes; each worker builds its own calculator (default 1)

### CHGNet relaxation
The structures from all models can be optionally relaxed with CHGNet.
```bash
$ cp scripts/cryspr_chgnet.py mp_20/WyckoffLLM-naive/DiffCSP++/parsed_materials_10000_pyxtal.json_structures.json.gz /your/working/dir/
$ cd /your/working/dir/
$ gzip -d parsed_materials_10000_pyxtal.json_structures.json.gz
$ python ./cryspr_chgnet.py 0 -1 ./parsed_materials_10000_pyxtal.json_structures.json wylm-dcpp
$ head wylm-dcpp_id_formula_energy.csv
model,id,formula,energy
wylm-dcpp,0,Y4Ho4Ir8,-128.26178
wylm-dcpp,1,La4Cu4Si8,-89.13443
wylm-dcpp,2,Rb3Br3,-20.12497
wylm-dcpp,3,Gd1Ni2,-26.83080
wylm-dcpp,4,Pb3O18,-96.37995
wylm-dcpp,5,Tl2Au2S4Br4,-41.40260
wylm-dcpp,6,Ru1As2In3Ce3,-49.81073
wylm-dcpp,7,Pd4As4Ni4Se4,-73.34074
wylm-dcpp,8,Na4Lu4F16,-149.18192
```
- Output files are in the same format as above ([CrySPR + CHGNet](#cryspr--chgnet)).
- `${reduced_formula}_${full_formula}_cell+pos.cif` is the CHGNet relaxed structure.
### DFT relaxation
We followed the [Materials Project protocol](https://docs.materialsproject.org/methodology/materials-methodology/calculation-details), [`atomate2.vasp.flows.mp.MPGGADoubleRelaxStaticMaker`](https://materialsproject.github.io/atomate2/reference/atomate2.vasp.flows.mp.MPGGADoubleRelaxStaticMaker.html). There isn't much to add, as the rest of the details of running DFT, unfortunately, depend on the HPC setup, and VASP is not open source. [Here](https://github.com/kazeevn/NSCC-VASP-computer) is the code to run at ASPIRE2.
## Generated Data Analysis
### Storage
#### Public Figshare
Most analyzed datasets in a uniform format are available at [Figshare](https://figshare.com/articles/dataset/Generated_crystals_for_WyFormer_DiffCSP_DiffCSP_WyCryst_SymmCD_CrystalFormer_MiAD/29145101).
#### Private Dropbox
The raw files are stored in a private Dropbox. To pull:
```bash
rclone copy "NUS_Dropbox:/Nikita Kazeev/Wyckoff Transformer data/generated.tar.gz" . --progress
tar -xvf generated.tar.gz
```
To push:
```bash
tar --use-compress-program=pigz -cvf generated.tar.gz generated
rclone copy generated.tar.gz "NUS_Dropbox:/Nikita Kazeev/Wyckoff Transformer data/" --progress
```
Tar is used to handle the large number of small files, and `pigz` is used to speed up the compression. The Dropbox folder is private, if you are a collaborator, please contact for access.

### Preprocessing
In order to be analyzed the data must be preprocessed and cached. To preprocess all generated datasets in `generated/datasets.yaml`:
```bash
uv run python scripts/cache_generated_datasets.py
```
It supports filtering by dataset and transformations, e. g.:
```bash
uv run python scripts/cache_generated_datasets.py --dataset mp_20 --transformations DiffCSP++ DFT
```
Completing this step will enable loading the data with `evaluation.generated_dataset.GeneratedDataset.from_cache`

### Metric computation
The ICML 2025 results were computed by the notebooks in [ICML_eval](ICML_eval). They include, but not limited to the following metrics:
1. S.U.N. - the fraction of stable, unique, and novel structures.
2. S.S.U.N. - the fraction of symmetric, stable, unique, and novel structures.
3. Space Group $\chi^2$ - the $\chi^2$ statistic of the space group distribution between the generated and the test set.
4. P1 - the fraction of generated structures that lack internal symmetries, i. e. belong to space group P1.
5. Property similarity and naive validity from [Xie et al.](https://arxiv.org/abs/2110.06197)
