{"id":15713519,"url":"https://github.com/deepmodeling/crystalformer","last_synced_at":"2025-05-01T09:17:08.358Z","repository":{"id":228735770,"uuid":"773712197","full_name":"deepmodeling/CrystalFormer","owner":"deepmodeling","description":"Space Group Informed Transformer for Crystalline Materials Generation","archived":false,"fork":false,"pushed_at":"2025-04-26T10:04:04.000Z","size":1700,"stargazers_count":92,"open_issues_count":4,"forks_count":12,"subscribers_count":5,"default_branch":"main","last_synced_at":"2025-05-01T09:16:54.258Z","etag":null,"topics":["autoregressive","crystal","generative-model","reinforcement-finetuning","transformer"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/deepmodeling.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2024-03-18T09:08:02.000Z","updated_at":"2025-04-29T09:34:04.000Z","dependencies_parsed_at":"2024-05-28T08:28:45.718Z","dependency_job_id":"6831f273-0591-4bf3-a553-32c20f54ea23","html_url":"https://github.com/deepmodeling/CrystalFormer","commit_stats":null,"previous_names":["zdcao121/crystalformer","deepmodeling/crystalformer"],"tags_count":2,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/deepmodeling%2FCrystalFormer","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/deepmodeling%2FCrystalFormer/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/deepmodeling%2FCrystalFormer/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/deepmodeling%2FCrystalFormer/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/deepmodeling","download_url":"https://codeload.github.com/deepmodeling/CrystalFormer/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":251850182,"owners_count":21653978,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["autoregressive","crystal","generative-model","reinforcement-finetuning","transformer"],"created_at":"2024-10-03T21:31:53.449Z","updated_at":"2025-05-01T09:17:08.351Z","avatar_url":"https://github.com/deepmodeling.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cdiv align=\"center\"\u003e\n\t\u003cimg align=\"middle\" src=\"imgs/crystalformer.png\" width=\"400\" alt=\"logo\"/\u003e\n  \u003ch2\u003e Crystal Generation with Space Group Informed Transformer\u003c/h2\u003e \n\u003c/div\u003e\n\n[![arXiv](https://img.shields.io/badge/arXiv-2403.15734-b31b1b.svg)](https://arxiv.org/abs/2403.15734) [![arXiv](https://img.shields.io/badge/arXiv-2504.02367-b31b1b.svg)](https://arxiv.org/abs/2504.02367)\n\n_CrystalFormer_ is a transformer-based autoregressive model specifically designed for space group-controlled generation of crystalline materials. The space group symmetry significantly simplifies the\ncrystal space, which is crucial for data and compute efficient generative modeling of crystalline materials.\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg align=\"middle\" src=\"imgs/output.gif\" width=\"400\"\u003e\n  \u003ch3\u003e Generating Cs\u003csub\u003e2\u003c/sub\u003eZnFe(CN)\u003csub\u003e6\u003c/sub\u003e Crystal (\u003ca href=https://next-gen.materialsproject.org/materials/mp-570545\u003emp-570545\u003c/a\u003e) \u003c/h3\u003e\n\u003c/div\u003e\n\n## Contents\n\n- [Contents](#contents)\n- [Model card](#model-card)\n- [Status](#status)\n- [Get Started](#get-started)\n- [Installation](#installation)\n  - [CPU installation](#cpu-installation)\n  - [CUDA (GPU) installation](#cuda-gpu-installation)\n  - [install required packages](#install-required-packages)\n  - [command line tools](#command-line-tools)\n- [Available Weights](#available-weights)\n- [How to run](#how-to-run)\n  - [train](#train)\n  - [sample](#sample)\n  - [evaluate](#evaluate)\n- [Reinforcement Fine-tuning](#reinforcement-fine-tuning)\n  - [$E\\_{hull}$ Reward](#e_hull-reward)\n  - [Dielectric FoM Reward](#dielectric-fom-reward)\n- [How to cite](#how-to-cite)\n\n## Model card\n\nThe model is an autoregressive transformer for the space group conditioned crystal probability distribution `P(C|g) = P (W_1 | ... ) P ( A_1 | ... ) P(X_1| ...) P(W_2|...) ... P(L| ...)`, where\n\n- `g`: space group number 1-230\n- `W`: Wyckoff letter ('a', 'b',...,'A')\n- `A`: atom type ('H', 'He', ..., 'Og')\n- `X`: factional coordinates\n- `L`: lattice vector [a,b,c, alpha, beta, gamma]\n- `P(W_i| ...)` and `P(A_i| ...)`  are categorical distributuions.\n- `P(X_i| ...)` is the mixture of von Mises distribution.\n- `P(L| ...)`  is the mixture of Gaussian distribution.\n\nWe only consider symmetry inequivalent atoms. The remaining atoms are restored based on the space group and Wyckoff letter information. Note that there is a natural alphabetical ordering for the Wyckoff letters, starting with 'a' for a position with the site-symmetry group of maximal order and ending with the highest letter for the general position. The sampling procedure starts from higher symmetry sites (with smaller multiplicities) and then goes on to lower symmetry ones (with larger multiplicities). Only for the cases where discrete Wyckoff letters can not fully determine the structure, one needs to further consider factional coordinates in the loss or sampling.\n\n## Status\n\nMajor milestones are summarized below.\n- v0.4.2 : Add implementation of direct preference optimization.\n- v0.4.1 : Replace the absolute positional embedding with the Rotary Positional Embedding (RoPE).\n- v0.4 : Add reinforcement learning (proximal policy optimization).\n- v0.3 : Add conditional generation in the plug-and-play manner.\n- v0.2 : Add Markov chain Monte Carlo (MCMC) sampling for template-based structure generation.\n- v0.1 : Initial implementations of crystalline material generation conditioned on the space group.\n\n## Get Started\n\n**Notebooks**: The quickest way to get started with _CrystalFormer_ is our notebooks in the Google Colab and Bohrium (Chinese version) platforms:\n\n- CrystalFormer Quickstart [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1IMQV6OQgIGORE8FmSTmZuC5KgQwGCnDx?usp=sharing) [![Open In Bohrium](https://cdn.dp.tech/bohrium/web/static/images/open-in-bohrium.svg)](https://nb.bohrium.dp.tech/detail/68177247598): GUI notebook demonstrating the conditional generation of crystalline materials with _CrystalFormer_\n- CrystalFormer Application [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1QdkELaQXAHR1zEu2fcdfgabuoP61_wbU?usp=sharing): Generating stable crystals with a given structure prototype. This workflow can be applied to tasks that are dominated by element substitution\n- CrystalFormer-RL [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1ojSqMQzdnlWZRPOQP20nTvvIh67HXdwp#scrollTo=lKOZgUczOAxE) [![Open In Bohrium](https://cdn.dp.tech/bohrium/web/static/images/open-in-bohrium.svg)](https://bohrium.dp.tech/notebooks/52828216135): Reinforcement fine-tuning for materials design\n\n## Installation\n\nCreate a new environment and install the required packages, we recommend using python `3.10.*` and conda to create the environment:\n\n```bash\n  conda create -n crystalgpt python=3.10\n  conda activate crystalgpt\n```\n\nBefore installing the required packages, you need to install `jax` and `jaxlib` first.\n\n### CPU installation\n\n```bash\npip install -U \"jax[cpu]\"\n```\n\n### CUDA (GPU) installation\n\nIf you intend to use CUDA (GPU) to speed up the training, it is important to install the appropriate version of `jax` and `jaxlib`. It is recommended to check the [jax docs](https://github.com/google/jax?tab=readme-ov-file#installation) for the installation guide. The basic installation command is given below:\n\n```bash\npip install --upgrade pip\n\n# NVIDIA CUDA 12 installation\n# Note: wheels only available on linux.\npip install -U \"jax[cuda12]\"\n```\n\n### install required packages\n\n```bash\npip install -r requirements.txt\n```\n\n### command line tools\nTo use the command line tools, you need to install the `crystalformer` package. You can use the following command to install the package:\n\n```bash\npip install .\n```\n\n## Available Weights\n\nWe release the weights of the model trained on the MP-20 dataset and Alex-20 dataset. More details can be seen in the [model](./model/README.md) folder.\n\n## How to run\n\n### train\n\n```bash\npython ./main.py --folder ./data/ --train_path YOUR_PATH/mp_20/train.csv --valid_path YOUR_PATH/mp_20/val.csv\n```\n\n- `folder`: the folder to save the model and logs\n- `train_path`: the path to the training dataset\n- `valid_path`: the path to the validation dataset\n- `test_path`: the path to the test dataset\n\n### sample\n\n```bash\npython ./main.py --optimizer none --test_path YOUR_PATH/mp_20/test.csv --restore_path YOUR_MODEL_PATH --spacegroup 160 --num_samples 1000  --batchsize 1000 --temperature 1.0\n```\n\n- `optimizer`: the optimizer to use, `none` means no training, only sampling\n- `restore_path`: the path to the model weights\n- `spacegroup`: the space group number to sample\n- `num_samples`: the number of samples to generate\n- `batchsize`: the batch size for sampling\n- `temperature`: the temperature for sampling\n\nYou can also use the `elements` to sample the specific element. For example, `--elements La Ni O` will sample the structure with La, Ni, and O atoms. The sampling results will be saved in the `output_LABEL.csv` file, where the `LABEL` is the space group number `g` specified in the command `--spacegroup`.\n\nThe input for the `elements` can be also the `json` file which specifies the atom mask in each Wyckoff site and the constraints. An example `atoms.json` file can be seen in the [data](./data/atoms.json) folder. There are two keys in the `atoms.json` file:\n\n- `atom_mask`: set the atom list for each Wyckoff position, the element can only be selected from the list in the corresponding Wyckoff position\n- `constraints`: set the constraints for the Wyckoff sites in the sampling, you can specify the pair of Wyckoff sites that should have the same elements\n\n\n### evaluate\n\nBefore evaluating the generated structures, you need to transform the generated `g, W, A, X, L` to the `cif` format. You can use the following command to transform the generated structures to the `cif` format and save as the `csv` file:\n\n```bash\npython ./scripts/awl2struct.py --output_path YOUR_PATH --label SPACE_GROUP  --num_io_process 40\n```\n\n- `output_path`: the path to read the generated `L, W, A, X` and save the `cif` files\n- `label`: the label to save the `cif` files, which is the space group number `g`\n- `num_io_process`: the number of processes\n\n\u003e [!IMPORTANT]\n\u003e The following evaluation script requires the [`SMACT`](https://github.com/WMD-group/SMACT), [`matminer`](https://github.com/hackingmaterials/matminer), and [`matbench-genmetrics`](https://github.com/sparks-baird/matbench-genmetrics) packages. We recommend installing them in a separate environment to avoid conflicts with other packages.\n\nCalculate the structure and composition validity of the generated structures:\n\n```bash\npython ./scripts/compute_metrics.py --root_path YOUR_PATH --filename YOUR_FILE --num_io_process 40\n```\n\n- `root_path`: the path to the dataset\n- `filename`: the filename of the generated structures\n- `num_io_process`: the number of processes\n\nCalculate the novelty and uniqueness of the generated structures:\n\n```bash\npython ./scripts/compute_metrics_matbench.py --train_path TRAIN_PATH --test_path TEST_PATH --gen_path GEN_PATH --output_path OUTPUT_PATH --label SPACE_GROUP --num_io_process 40\n```\n\n- `train_path`: the path to the training dataset\n- `test_path`: the path to the test dataset\n- `gen_path`: the path to the generated dataset\n- `output_path`: the path to save the metrics results\n- `label`: the label to save the metrics results, which is the space group number `g`\n- `num_io_process`: the number of processes\n\nNote that the training, test, and generated datasets should contain the structures within the **same** space group `g` which is specified in the command `--label`.\n\nMore details about the post-processing can be seen in the [scripts](./scripts/README.md) folder.\n\n## Reinforcement Fine-tuning\n\n\u003e [!IMPORTANT]\n\u003e Before running the reinforcement fine-tuning, please make sure you have installed the corresponding machine learning force field model or property prediction model. The `mlff_model` and `mlff_path` arguments in the command line should be set according to the model you are using. Now we support the[`orb`](https://github.com/orbital-materials/orb-models) and [`MACE`](https://github.com/ACEsuit/mace) models for the $E_{hull}$ reward, and the [`matgl`](https://github.com/materialsvirtuallab/matgl) model for the dielectric FoM reward.\n\n### $E_{hull}$ Reward\n\n```bash\ntrain_ppo --folder ./data/\\\n          --restore_path YOUR_PATH\\\n          --valid_path YOUR_PATH/alex_20/val.csv\\\n          --test_path YOUR_PATH/alex_20/train.csv\\\n          --reward ehull\\\n          --convex_path YOUR_PATH/convex_hull_pbe_2023.12.29.json.bz2\\\n          --mlff_model orb\\\n          --mlff_path YOUR_PATH/orb-v2-20241011.ckpt\n```\n\n- `folder`: the folder to save the model and logs\n- `restore_path`: the path to the pre-trained model weights\n- `valid_path`: the path to the validation dataset\n- `test_path`: the path to the test dataset. The space group distribution will be loaded from this dataset and used for the sampling in the reinforcement learning fine-tuning\n- `reward`: the reward function to use, `ehull` means the energy above the convex hull\n- `convex_path`: the path to the convex hull data, which is used to calculate the $E_{hull}$. Only used when the reward is `ehull`\n- `mlff_model`: the machine learning force field model to predict the total energy. We support [`orb`](https://github.com/orbital-materials/orb-models) and [`MACE`](https://github.com/ACEsuit/mace) models for the $E_{hull}$ reward\n- `mlff_path`: the path to load the checkpoint of the machine learning force field model\n\n### Dielectric FoM Reward\n\n```bash\ntrain_ppo --folder ./data/\\\n          --restore_path YOUR_PATH\\\n          --valid_path YOUR_PATH/alex_20/val.csv\\\n          --test_path YOUR_PATH/alex_20/train.csv\\\n          --reward dielectric\\\n          --mlff_model matgl\\\n          --mlff_path YOUR_PATH/model1,YOUR_PATH/model2\n```\n\n- `folder`: the folder to save the model and logs\n- `restore_path`: the path to the pre-trained model weights\n- `valid_path`: the path to the validation dataset\n- `test_path`: the path to the test dataset. The space group distribution will be loaded from this dataset and used for the sampling in the reinforcement learning fine-tuning\n- `reward`: the reward function to use, `dielectric` means the dielectric figure of merit (FoM), which is the product of the total dielectric constant and the band gap\n- `mlff_model`: the machine learning force field model to predict the total energy. We only support models in [`matgl`](https://github.com/materialsvirtuallab/matgl) for the dielectric reward\n- `mlff_path`: the path to load the checkpoint of the machine learning force field model. Note that you need to provide the model paths for the total dielectric constant and band gap, separated by the `,`\n\n\n## How to cite\n\n```bibtex\n@article{cao2024space,\n      title={Space Group Informed Transformer for Crystalline Materials Generation}, \n      author={Zhendong Cao and Xiaoshan Luo and Jian Lv and Lei Wang},\n      year={2024},\n      eprint={2403.15734},\n      archivePrefix={arXiv},\n      primaryClass={cond-mat.mtrl-sci}\n}\n```\n\n```bibtex\n@article{cao2025crystalformerrl,\n      title={CrystalFormer-RL: Reinforcement Fine-Tuning for Materials Design}, \n      author={Zhendong Cao and Lei Wang},\n      year={2025},\n      eprint={2504.02367},\n      archivePrefix={arXiv},\n      primaryClass={cond-mat.mtrl-sci},\n      url={https://arxiv.org/abs/2504.02367}, \n}\n```\n\n**Note**: This project is unrelated to https://github.com/omron-sinicx/crystalformer with the same name.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdeepmodeling%2Fcrystalformer","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdeepmodeling%2Fcrystalformer","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdeepmodeling%2Fcrystalformer/lists"}