update

Author: zdcao121

.github/workflows/tests.yml

@@ -24,7 +24,7 @@ jobs:
           python -m pip install --upgrade pip
           pip install pytest
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-          pip install --upgrade "jax[cpu]==0.4.25"
+          pip install -U jax
           pip install .
       - name: Test with pytest
         run: |

README.md

@@ -23,6 +23,7 @@ crystal space, which is crucial for data and compute efficient generative modeli
   - [CPU installation](#cpu-installation)
   - [CUDA (GPU) installation](#cuda-gpu-installation)
   - [install required packages](#install-required-packages)
+  - [command line tools](#command-line-tools)
 - [Available Weights](#available-weights)
 - [How to run](#how-to-run)
   - [train](#train)
@@ -85,17 +86,14 @@ pip install -U "jax[cpu]"
 
 ### CUDA (GPU) installation
 
-> [!CAUTION]
-> CrystalFormer requires JAX versions between v0.4.25 and v0.4.35. Please avoid using versions newer than v0.4.35 to prevent compatibility issues. If you have already installed a newer version, please uninstall it first and then install one within the required range.
-
 If 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:
 
 ```bash
 pip install --upgrade pip
 
 # NVIDIA CUDA 12 installation
 # Note: wheels only available on linux.
-pip install --upgrade "jax[cuda12]"
+pip install -U "jax[cuda12]"
 ```
 
 ### install required packages
@@ -104,6 +102,13 @@ pip install --upgrade "jax[cuda12]"
 pip install -r requirements.txt
 ```
 
+### command line tools
+To use the command line tools, you need to install the `crystalformer` package. You can use the following command to install the package:
+
+```bash
+pip install .
+```
+
 ## Available Weights
 
 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.
@@ -154,6 +159,9 @@ python ./scripts/awl2struct.py --output_path YOUR_PATH --label SPACE_GROUP  --nu
 - `label`: the label to save the `cif` files, which is the space group number `g`
 - `num_io_process`: the number of processes
 
+> [!IMPORTANT]
+> 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.
+
 Calculate the structure and composition validity of the generated structures:
 
 ```bash

requirements.txt

@@ -1,7 +1,3 @@
-dm-haiku==0.0.11
-optax==0.1.8
-pymatgen==2024.3.1
-pyxtal==0.6.3
-SMACT==2.5.5
-matbench-genmetrics==0.6.1
-matminer @ git+https://github.com/hackingmaterials/matminer.git
+dm-haiku==0.1.14
+optax==0.2.4
+pyxtal==1.0.7

scripts/README.md

@@ -9,6 +9,7 @@
   - [Relaxation](#relaxation)
   - [Energy Above the Hull](#energy-above-the-hull)
   - [Embedding Visualization](#embedding-visualization)
+  - [Stable, Unique and Novel Structures](#stable-unique-and-novel-structures)
   - [Structure Visualization](#structure-visualization)
 
 ### Transform
@@ -48,15 +49,18 @@ Note that the training, test, and generated datasets should contain the structur
 
 
 ### Relaxation
-`matgl_relax.py` is a script to relax the generated structures using the `matgl` package. You can install the `matgl` following the instructions in the [matgl repository](https://github.com/materialsvirtuallab/matgl?tab=readme-ov-file).
+`mlff_relax.py` is a script to relax the generated structures using pretrained machine learning force field. Now we support the [`orb`](https://github.com/orbital-materials/orb-models), [`MACE`](https://github.com/ACEsuit/mace), [`matgl`](https://github.com/materialsvirtuallab/matgl) and [`deepmd-kit`](https://github.com/deepmodeling/deepmd-kit) models. Please install corresponding packages before running the script.
+
 ```bash
-python matgl_relax.py --restore_path RESTORE_PATH --filename FILENAME --relaxation --model_path MODEL_PATH
+python mlff_relax.py --restore_path RESTORE_PATH --filename FILENAME --relaxation --model orb --model_path MODEL_PATH
 ```
 - `restore_path`: the path to the generated structures
 - `filename`: the filename of the generated structures
 - `relaxation`: whether to relax the structures, if not specified, the script will only predict the energy of the structures without relaxation
-- `model_path`: the path to the `matgl` model checkpoint
-
+- `model`: the model to use for relaxation, which can be `orb`, `mace`, `matgl` or `dp`
+- `model_path`: the path to the machine learning force field checkpoint
+- `primitive`: whether to convert the structures to primitive cells, if not specified, the script will only relax the structures without converting to primitive cells. This can be used to reduce the number of atoms in the structures and speed up the relaxation process
+- `fixsymmetry`: whether to fix the space group symmetry of the structures in the relaxation process
 
 ### Energy Above the Hull
 `e_above_hull.py` is a script to calculate the energy above the hull of the generated structures based on the Materials Project database. To calculate the energy above the hull, the API key of the Materials Project is required, which can be obtained from the [Materials Project website](https://next-gen.materialsproject.org/). Furthermore, the `mp_api` package should be installed.
@@ -70,6 +74,18 @@ python e_above_hull.py --restore_path RESTORE_PATH --filename FILENAME --api_key
 - `label`: the label to save the energy above the hull file
 - `relaxation`: whether to calculate the energy above the hull based on the relaxed structures
 
+`e_above_hull_alex.py` is a script to calculate the energy above the hull of the generated structures based on the Alexandria database. To calculate the energy above the hull, the Alexandria convex hull data is required, which can be obtained from the [Alexandria website](https://alexandria.icams.rub.de/).
+
+```bash
+python e_above_hull_alex.py --convex_path CONVEX_PATH --restore_path RESTORE_PATH --filename FILENAME --api_key API_KEY --label LABEL --relaxation
+```
+- `convex_path`: the path to the Alexandria convex hull data
+- `restore_path`: the path to the structures 
+- `filename`: the filename of the structures
+- `api_key`: the API key of the Materials Project
+- `label`: the label to save the energy above the hull file
+- `relaxation`: whether to calculate the energy above the hull based on the relaxed structures
+
 ### Embedding Visualization
 `plot_embeddings.py` is a script to visualize the correlation of the learned embedding vectors of different elements.
 
@@ -79,5 +95,16 @@ python plot_embeddings.py --restore_path RESTORE_PATH
 
 - `restore_path`: the path to the model checkpoint
 
+### Stable, Unique and Novel Structures
+`check_sun_materials.py` is a script to check the stable, unique and novel structures based on the given reference dataset.
+
+```bash
+python check_sun_materials.py --restore_path RESTORE_PATH --filename FILENAME --ref_path REF_PATH
+```
+
+- `restore_path`: the path to the generated structures
+- `filename`: the filename of the generated structures
+- `ref_path`: the path to the reference dataset
+
 ### Structure Visualization
 `structure_visualization.ipynb` is a notebook to visualize the generated structures.

scripts/mlff_relax.py

@@ -174,7 +174,7 @@ def main(args):
 
     import argparse
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model", type=str, choices=["orb", "matgl", "mace"], default="orb", help="choose the MLFF model")
+    parser.add_argument("--model", type=str, choices=["orb", "matgl", "mace", "dp"], default="orb", help="choose the MLFF model")
     parser.add_argument("--device", type=str, default="cuda", help="choose the device to run the model on")
     parser.add_argument("--model_path", type=str, default="./data/orb-v2-20241011.ckpt", help="path to the model checkpoint")
     parser.add_argument("--restore_path", type=str, default="./experimental/", help="")