update to v0.4.2

Author: zdcao121

.gitignore

@@ -2,3 +2,7 @@
 job*
 *.out
 __pycache__/
+data/
+experimental/
+*.ipynb
+*.egg-info/

README.md

@@ -17,6 +17,7 @@ crystal space, which is crucial for data and compute efficient generative modeli
 
 - [Contents](#contents)
 - [Model card](#model-card)
+- [Status](#status)
 - [Get Started](#get-started)
 - [Installation](#installation)
   - [CPU installation](#cpu-installation)
@@ -27,6 +28,9 @@ crystal space, which is crucial for data and compute efficient generative modeli
   - [train](#train)
   - [sample](#sample)
   - [evaluate](#evaluate)
+- [Reinforcement Fine-tuning](#reinforcement-fine-tuning)
+  - [$E\_{hull}$ Reward](#e_hull-reward)
+  - [Dielectric FoM Reward](#dielectric-fom-reward)
 - [How to cite](#how-to-cite)
 
 ## Model card
@@ -44,6 +48,16 @@ The model is an autoregressive transformer for the space group conditioned cryst
 
 We 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.
 
+## Status
+
+Major milestones are summarized below.
+- v0.4.2 : Add implementation of direct preference optimization.
+- v0.4.1 : Replace the absolute positional embedding with the Rotary Positional Embedding (RoPE).
+- v0.4 : Add reinforcement learning (proximal policy optimization).
+- v0.3 : Add conditional generation in the plug-and-play manner.
+- v0.2 : Add Markov chain Monte Carlo (MCMC) sampling for template-based structure generation.
+- v0.1 : Initial implementations of crystalline material generation conditioned on the space group.
+
 ## Get Started
 
 **Notebooks**: The quickest way to get started with _CrystalFormer_ is our notebooks in the Google Colab and Bohrium (Chinese version) platforms:
@@ -88,7 +102,7 @@ pip install -r requirements.txt
 
 ## Available Weights
 
-We release the weights of the model trained on the MP-20 dataset. More details can be seen in the [model](./model/README.md) folder.
+We 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.
 
 ## How to run
 
@@ -163,10 +177,55 @@ Note that the training, test, and generated datasets should contain the structur
 
 More details about the post-processing can be seen in the [scripts](./scripts/README.md) folder.
 
+## Reinforcement Fine-tuning
+
+### $E_{hull}$ Reward
+
+```bash
+train_ppo --folder ./data/\
+          --restore_path YOUR_PATH\
+          --valid_path YOUR_PATH/alex_20/val.csv\
+          --test_path YOUR_PATH/alex_20/train.csv\
+          --reward ehull\
+          --convex_path YOUR_PATH/convex_hull_pbe_2023.12.29.json.bz2\
+          --mlff_model orb\
+          --mlff_path YOUR_PATH/orb-v2-20241011.ckpt
+```
+
+- `folder`: the folder to save the model and logs
+- `restore_path`: the path to the pre-trained model weights
+- `valid_path`: the path to the validation dataset
+- `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
+- `reward`: the reward function to use, `ehull` means the energy above the convex hull
+- `convex_path`: the path to the convex hull data, which is used to calculate the $E_{hull}$. Only used when the reward is `ehull`
+- `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
+- `mlff_path`: the path to load the checkpoint of the machine learning force field model
+
+### Dielectric FoM Reward
+
+```bash
+train_ppo --folder ./data/\
+          --restore_path YOUR_PATH\
+          --valid_path YOUR_PATH/alex_20/val.csv\
+          --test_path YOUR_PATH/alex_20/train.csv\
+          --reward dielectric\
+          --mlff_model matgl\
+          --mlff_path YOUR_PATH/model1,YOUR_PATH/model2
+```
+
+- `folder`: the folder to save the model and logs
+- `restore_path`: the path to the pre-trained model weights
+- `valid_path`: the path to the validation dataset
+- `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
+- `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
+- `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
+- `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 `,`
+
+
 ## How to cite
 
 ```bibtex
-@misc{cao2024space,
+@article{cao2024space,
       title={Space Group Informed Transformer for Crystalline Materials Generation}, 
       author={Zhendong Cao and Xiaoshan Luo and Jian Lv and Lei Wang},
       year={2024},
@@ -176,4 +235,16 @@ More details about the post-processing can be seen in the [scripts](./scripts/RE
 }
 ```
 
+```bibtex
+@article{cao2025crystalformerrl,
+      title={CrystalFormer-RL: Reinforcement Fine-Tuning for Materials Design}, 
+      author={Zhendong Cao and Lei Wang},
+      year={2025},
+      eprint={2504.02367},
+      archivePrefix={arXiv},
+      primaryClass={cond-mat.mtrl-sci},
+      url={https://arxiv.org/abs/2504.02367}, 
+}
+```
+
 **Note**: This project is unrelated to https://github.com/omron-sinicx/crystalformer with the same name.

crystalformer/cli/__init__.py

@@ -20,8 +20,36 @@ def get_labels(csv_file, label_col):
     labels = jnp.array(labels, dtype=float)
     return labels
 
+def GLXYZAW_from_sample(spg, test_path):
+    ### read from generated data
+    from ast import literal_eval
+    from crystalformer.src.wyckoff import mult_table
 
-if __name__  == "__main__":
+    test_data = pd.read_csv(test_path)
+    L, XYZ, A, W = test_data['L'], test_data['X'], test_data['A'], test_data['W']
+    L = L.apply(lambda x: literal_eval(x))
+    XYZ = XYZ.apply(lambda x: literal_eval(x))
+    A = A.apply(lambda x: literal_eval(x))
+    W = W.apply(lambda x: literal_eval(x))
+
+    # convert array of list to numpy ndarray
+    G = jnp.array([spg]*len(L))
+    L = jnp.array(L.tolist())
+    XYZ = jnp.array(XYZ.tolist())
+    A = jnp.array(A.tolist())
+    W = jnp.array(W.tolist())
+
+    M = jax.vmap(lambda g, w: mult_table[g-1, w], in_axes=(0, 0))(G, W) # (batchsize, n_max)
+    num_atoms = jnp.sum(M, axis=1)
+    length, angle = jnp.split(L, 2, axis=-1)
+    length = length/num_atoms[:, None]**(1/3)
+    angle = angle * (jnp.pi / 180) # to rad
+    L = jnp.concatenate([length, angle], axis=-1)
+
+    return G, L, XYZ, A, W
+
+
+def main():
 
     import argparse
     parser = argparse.ArgumentParser(description='')
@@ -30,6 +58,7 @@ def get_labels(csv_file, label_col):
     group.add_argument('--train_path', default='/data/zdcao/crystal_gpt/dataset/mp_20/train.csv', help='')
     group.add_argument('--valid_path', default='/data/zdcao/crystal_gpt/dataset/mp_20/val.csv', help='')
     group.add_argument('--test_path', default='/data/zdcao/crystal_gpt/dataset/mp_20/test.csv', help='')
+    group.add_argument('--spacegroup', type=int, default=None, help='The space group number')
     group.add_argument('--property', default='band_gap', help='The property to predict')
     group.add_argument('--num_io_process', type=int, default=40, help='number of io processes')
 
@@ -82,11 +111,13 @@ def get_labels(csv_file, label_col):
         valid_data = (*valid_data, valid_labels)
 
     else:
-        test_data = GLXYZAW_from_file(args.test_path, args.atom_types,
-                                      args.wyck_types, args.n_max, args.num_io_process)
-        test_labels = get_labels(args.test_path, args.property)
-
-        test_data = (*test_data, test_labels)
+        if args.spacegroup == None:
+            G, L, XYZ, A, W = GLXYZAW_from_file(args.test_path, args.atom_types,
+                                        args.wyck_types, args.n_max, args.num_io_process)
+            test_labels = get_labels(args.test_path, args.property)
+        
+        else:
+            G, L, XYZ, A, W = GLXYZAW_from_sample(args.spacegroup, args.test_path)
 
     ################### Model #############################
     transformer_params, state, transformer = make_transformer(key, args.Nf, args.Kx, args.Kl, args.n_max, 
@@ -146,12 +177,16 @@ def get_labels(csv_file, label_col):
         params, opt_state = train(subkey, optimizer, opt_state, loss_fn, params, state, epoch_finished, args.epochs, args.batchsize, train_data, valid_data, output_path)
 
     elif args.optimizer == 'none':
-        G, L, XYZ, A, W, labels = test_data
+        
         y = jax.vmap(forward_fn,
              in_axes=(None, None, None, 0, 0, 0, 0, 0, None)
              )(params, state, key, G, L, XYZ, A, W, False)
 
         jnp.save(args.output_path, y)
 
     else:
-        raise NotImplementedError(f"Optimizer {args.optimizer} not implemented")
+        raise NotImplementedError(f"Optimizer {args.optimizer} not implemented")
+
+
+if __name__ == "__main__":
+    main()

classifier.py crystalformer/cli/classifier.pyclassifier.py renamed to crystalformer/cli/classifier.py

@@ -17,7 +17,7 @@
 from crystalformer.src.transformer import make_transformer
 
 
-if __name__  == "__main__":
+def main():
 
     import argparse
     parser = argparse.ArgumentParser(description='')
@@ -248,3 +248,7 @@
     data.to_csv(filename, mode='a', index=False, header=header)
 
     print ("Wrote samples to %s"%filename)
+
+
+if __name__  == "__main__":
+    main()

cond_gen.py crystalformer/cli/cond_gen.pycond_gen.py renamed to crystalformer/cli/cond_gen.py

@@ -0,0 +1,60 @@
+import os
+import lmdb
+import pickle
+import numpy as np
+from crystalformer.src.utils import GLXYZAW_from_file
+import warnings
+warnings.filterwarnings("ignore")
+
+
+def csv_to_lmdb(csv_file, lmdb_file, args):
+    if os.path.exists(lmdb_file):
+        os.remove(lmdb_file)
+        print(f"Removed existing {lmdb_file}")
+
+    values = GLXYZAW_from_file(csv_file,
+                               atom_types=args.atom_types,
+                               wyck_types=args.wyck_types,
+                               n_max=args.n_max,
+                               num_workers=args.num_workers)
+    keys = np.arange(len(values[0]))
+
+    env = lmdb.open(
+        lmdb_file,
+        subdir=False,
+        readonly=False,
+        lock=False,
+        readahead=False,
+        meminit=False,
+        max_readers=1,
+        map_size=int(100e9),
+    )
+
+    with env.begin(write=True) as txn:
+        for key, value in zip(keys, values):
+            txn.put(str(key).encode("utf-8"), pickle.dumps(value))
+
+    print(f"Successfully converted {csv_file} to {lmdb_file}")
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--n_max', type=int, default=21, help='The maximum number of atoms in the cell')
+    parser.add_argument('--atom_types', type=int, default=119, help='Atom types including the padded atoms')
+    parser.add_argument('--wyck_types', type=int, default=28, help='Number of possible multiplicites including 0')
+
+    parser.add_argument("--path", type=str, required=True)
+    parser.add_argument("--num_workers", type=int, default=40)
+    args = parser.parse_args()
+
+    for i in ["test", "val", "train"]:
+        csv_to_lmdb(
+            os.path.join(args.path, f"{i}.csv"), 
+            os.path.join(args.path, f"{i}.lmdb"),
+            args
+        )
+
+
+if __name__ == "__main__":
+    main()