新建
晶体生成模型CrystalFormer
Russel
A.Tの
推荐镜像 :Basic Image:ubuntu:22.04-py3.10-cuda12.1
推荐机型 :c2_m4_cpu
赞 7
11
2
目录
CrystalFormer 快速上手指南
代码
文本
CrystalFormer 是一种基于Transformer架构的自回归模型，专为空间群控制的晶体材料生成而设计。利用空间群的对称性显著简化了晶体空间，可以高效地生成符合指定空间群约束的晶体构型。
©️ Copyright 2024 @ 曹振东，罗啸山，吕健， 王磊
日期：2024-4-XX
共享协议：本作品采用知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议进行许可。
快速开始：点击上方的 开始连接 按钮，推荐选择 Basic Image:ubuntu:22.04-py3.10-cuda12.1 镜像及 c2_m4_cpu 节点配置，连接之后即可运行。
代码
文本
理论部分
晶体表示 (Crystal Representation)
我们经常用下面的数学形式去表示一个晶体结构：
其中是表示晶胞内原子种类； 表示晶格常量以及它们之间的夹角； 表示晶胞内原子在晶胞内的分数坐标。
空间群 (Space Group)
空间群（Space groups）是晶体学中的一个概念，用于描述晶体中原子排列的三维对称性。空间群是对称操作的集合，这些操作包括旋转、反射、旋转加反射以及平移等。接下来我们可以看到，通过充分利用空间群，我们可以大大简化晶体结构的表示。
空间群对晶格的限制
空间群可以被分为7个晶系，分别是三斜（Triclinic）, 单斜（Monoclinic）, 正交（Orthorhombic）, 四方（Tetragonal）, 三方（Trigonal）, 六方（Hexagonal）, 立方（Cubic）。如下图所示，不同晶系下对于晶格参数 有着不同程度下的限制，比如对于立方晶系，三个晶格常数必须相等 （），且三个夹角必须为直角 （）。
空间群对原子种类和坐标的限制-Wyckoff位置 (Wyckoff Position)
空间群不仅仅对晶体结构的晶格参数做了限制，同时也对原子的种类和坐标做了一定限制，这就是Wyckoff位置。Wyckoff位置是晶体学中用来描述晶体中具有相同对称性的原子集合的一组特殊点。Wyckoff位置有助于简化晶体结构的分析，因为它们允许我们通过识别晶体的对称性来推断出原子的位置。Wyckoff位置一般包含以下信息：
Wyckoff字母：每个Wyckoff位置都有一个特定的字母标识，如 a、b、c 等。通常，从a开始，按字母顺序递增，直到最大的字母。最大的字母代表的是一般性位置（General position），它具有最高的多重性 （Multiplicity），意味着如果该位置上有原子，那么通过对称操作可以生成最多的等价原子。
位点对称性（Site Symmetry）：描述了Wyckoff位置上的点所具有的对称性质，例如，一个点可能具有旋转、反射或螺旋对称性
分数坐标：Wyckoff位置通常用分数坐标来表示，这些坐标反映了原子在晶胞中的位置
下图所示为167号空间群的Wyckoff位置的列表。也可以通过 Bilbao Server 来查找不同空间群下对应的Wyckoff位置。
利用Wykoff位置，我们可以极大的简化晶体的表示。如下图所示，对于一个空间群为167的CaCO结构，晶胞内有30个原子，如果采用常规的表示手段，我们需要30个原子的原子种类以及它们各自的分数坐标。而如果采用Wyckoff位置来表示晶体，我们只需要三个不等价的原子种类以及它们的Wyckoff位置做就可以表示晶胞的所有原子，这可以极大的简化晶体结构的表示，也更有利于深度学习。同时，Wyckoff位置有一个天然的对称性高低的顺序，用Wyckoff字母 表示，这十分切合自回归模型的假设。
代码
文本
模型 （CrystalFormer）
由上述可知，在给定空间群的情况下，我们可以利用Wyckoff位置来大大简化晶体的表示。对于无条件的晶体生成模型，概率分布为， 而将空间群作为条件之后，变成了条件生成概率 ，如下图所示，通过空间群的限制，我们用 来表示晶体结构，代替原本的 ，其中 是Wyckoff位置的集合。
因此，我们提出了CrystalFormer的晶体生成模型，利用Transformer的解码器（vanilla decoder-only transformer）来对晶体结构按照 的顺序进行生成。更准确的说，是在给定空间群之后，按照顺序对每个原子的 进行序列生成，如下图所示，预测的下一个原子的Wyckoff位置 为填充原子(padding atom)的时候，开始预测晶格 。下图中的字符串为 CaCO 结构的序列表示。
代码
文本
准备工作
代码
文本
下载CrystalFormer代码和安装依赖的python库
代码
文本
[1]
! git clone https://github.com/deepmodeling/CrystalFormer.git
import os
os.chdir('/CrystalFormer')
fatal: destination path 'CrystalFormer' already exists and is not an empty directory.
代码
文本
[2]
! git checkout 5c6103fb496cd2991243a492e6a41c3c524f1834
HEAD is now at 5c6103f add w_mask to control the sampling
代码
文本
[3]
! pip install -U "jax[cpu]"
! pip install dm-haiku==0.0.11
! pip install pymatgen==2024.3.1
! pip install pyxtal==0.6.3
! pip install hydra-core --upgrade
! pip install nglview
Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple Requirement already satisfied: jax[cpu] in /opt/mamba/lib/python3.10/site-packages (0.4.29) Requirement already satisfied: scipy>=1.9 in /opt/mamba/lib/python3.10/site-packages (from jax[cpu]) (1.11.4) Requirement already satisfied: opt-einsum in /opt/mamba/lib/python3.10/site-packages (from jax[cpu]) (3.3.0) Requirement already satisfied: ml-dtypes>=0.4.0 in /opt/mamba/lib/python3.10/site-packages (from jax[cpu]) (0.4.0) Requirement already satisfied: numpy>=1.22 in /opt/mamba/lib/python3.10/site-packages (from jax[cpu]) (1.26.2) Requirement already satisfied: jaxlib==0.4.29 in /opt/mamba/lib/python3.10/site-packages (from jax[cpu]) (0.4.29) WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple Requirement already satisfied: dm-haiku==0.0.11 in /opt/mamba/lib/python3.10/site-packages (0.0.11) Requirement already satisfied: absl-py>=0.7.1 in /opt/mamba/lib/python3.10/site-packages (from dm-haiku==0.0.11) (2.1.0) Requirement already satisfied: jmp>=0.0.2 in /opt/mamba/lib/python3.10/site-packages (from dm-haiku==0.0.11) (0.0.4) Requirement already satisfied: tabulate>=0.8.9 in /opt/mamba/lib/python3.10/site-packages (from dm-haiku==0.0.11) (0.9.0) Requirement already satisfied: flax>=0.7.1 in /opt/mamba/lib/python3.10/site-packages (from dm-haiku==0.0.11) (0.8.4) Requirement already satisfied: numpy>=1.18.0 in /opt/mamba/lib/python3.10/site-packages (from dm-haiku==0.0.11) (1.26.2) Requirement already satisfied: jax>=0.4.19 in /opt/mamba/lib/python3.10/site-packages (from flax>=0.7.1->dm-haiku==0.0.11) (0.4.29) Requirement already satisfied: optax in /opt/mamba/lib/python3.10/site-packages (from flax>=0.7.1->dm-haiku==0.0.11) (0.2.2) Requirement already satisfied: typing-extensions>=4.2 in /opt/mamba/lib/python3.10/site-packages (from flax>=0.7.1->dm-haiku==0.0.11) (4.9.0) Requirement already satisfied: tensorstore in /opt/mamba/lib/python3.10/site-packages (from flax>=0.7.1->dm-haiku==0.0.11) (0.1.61) Requirement already satisfied: rich>=11.1 in /opt/mamba/lib/python3.10/site-packages (from flax>=0.7.1->dm-haiku==0.0.11) (13.7.1) Requirement already satisfied: orbax-checkpoint in /opt/mamba/lib/python3.10/site-packages (from flax>=0.7.1->dm-haiku==0.0.11) (0.5.16) Requirement already satisfied: msgpack in /opt/mamba/lib/python3.10/site-packages (from flax>=0.7.1->dm-haiku==0.0.11) (1.0.8) Requirement already satisfied: PyYAML>=5.4.1 in /opt/mamba/lib/python3.10/site-packages (from flax>=0.7.1->dm-haiku==0.0.11) (6.0.1) Requirement already satisfied: opt-einsum in /opt/mamba/lib/python3.10/site-packages (from jax>=0.4.19->flax>=0.7.1->dm-haiku==0.0.11) (3.3.0) Requirement already satisfied: ml-dtypes>=0.4.0 in /opt/mamba/lib/python3.10/site-packages (from jax>=0.4.19->flax>=0.7.1->dm-haiku==0.0.11) (0.4.0) Requirement already satisfied: scipy>=1.9 in /opt/mamba/lib/python3.10/site-packages (from jax>=0.4.19->flax>=0.7.1->dm-haiku==0.0.11) (1.11.4) Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /opt/mamba/lib/python3.10/site-packages (from rich>=11.1->flax>=0.7.1->dm-haiku==0.0.11) (2.17.2) Requirement already satisfied: markdown-it-py>=2.2.0 in /opt/mamba/lib/python3.10/site-packages (from rich>=11.1->flax>=0.7.1->dm-haiku==0.0.11) (3.0.0) Requirement already satisfied: chex>=0.1.86 in /opt/mamba/lib/python3.10/site-packages (from optax->flax>=0.7.1->dm-haiku==0.0.11) (0.1.86) Requirement already satisfied: jaxlib>=0.1.37 in /opt/mamba/lib/python3.10/site-packages (from optax->flax>=0.7.1->dm-haiku==0.0.11) (0.4.29) Requirement already satisfied: etils[epath,epy] in /opt/mamba/lib/python3.10/site-packages (from orbax-checkpoint->flax>=0.7.1->dm-haiku==0.0.11) (1.7.0) Requirement already satisfied: nest_asyncio in /opt/mamba/lib/python3.10/site-packages (from orbax-checkpoint->flax>=0.7.1->dm-haiku==0.0.11) (1.5.8) Requirement already satisfied: protobuf in /opt/mamba/lib/python3.10/site-packages (from orbax-checkpoint->flax>=0.7.1->dm-haiku==0.0.11) (5.27.1) Requirement already satisfied: toolz>=0.9.0 in /opt/mamba/lib/python3.10/site-packages (from chex>=0.1.86->optax->flax>=0.7.1->dm-haiku==0.0.11) (0.12.0) Requirement already satisfied: mdurl~=0.1 in /opt/mamba/lib/python3.10/site-packages (from markdown-it-py>=2.2.0->rich>=11.1->flax>=0.7.1->dm-haiku==0.0.11) (0.1.2) Requirement already satisfied: zipp in /opt/mamba/lib/python3.10/site-packages (from etils[epath,epy]->orbax-checkpoint->flax>=0.7.1->dm-haiku==0.0.11) (3.19.2) Requirement already satisfied: fsspec in /opt/mamba/lib/python3.10/site-packages (from etils[epath,epy]->orbax-checkpoint->flax>=0.7.1->dm-haiku==0.0.11) (2024.6.0) Requirement already satisfied: importlib_resources in /opt/mamba/lib/python3.10/site-packages (from etils[epath,epy]->orbax-checkpoint->flax>=0.7.1->dm-haiku==0.0.11) (6.4.0) WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple Requirement already satisfied: pymatgen==2024.3.1 in /opt/mamba/lib/python3.10/site-packages (2024.3.1) Requirement already satisfied: sympy in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (1.12.1) Requirement already satisfied: monty>=2024.2.2 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (2024.5.24) Requirement already satisfied: matplotlib>=1.5 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (3.9.0) Requirement already satisfied: palettable>=3.1.1 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (3.3.3) Requirement already satisfied: tabulate in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (0.9.0) Requirement already satisfied: joblib in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (1.4.2) Requirement already satisfied: pandas in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (2.1.4) Requirement already satisfied: ruamel.yaml>=0.17.0 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (0.18.6) Requirement already satisfied: scipy>=1.5.0 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (1.11.4) Requirement already satisfied: plotly>=4.5.0 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (5.22.0) Requirement already satisfied: requests in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (2.31.0) Requirement already satisfied: pybtex in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (0.24.0) Requirement already satisfied: spglib>=2.0.2 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (2.4.0) Requirement already satisfied: uncertainties>=3.1.4 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (3.2.1) Requirement already satisfied: networkx>=2.2 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (3.3) Requirement already satisfied: tqdm in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (4.64.1) Requirement already satisfied: numpy>=1.25.0 in /opt/mamba/lib/python3.10/site-packages (from pymatgen==2024.3.1) (1.26.2) Requirement already satisfied: python-dateutil>=2.7 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=1.5->pymatgen==2024.3.1) (2.8.2) Requirement already satisfied: contourpy>=1.0.1 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=1.5->pymatgen==2024.3.1) (1.2.1) Requirement already satisfied: pillow>=8 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=1.5->pymatgen==2024.3.1) (10.3.0) Requirement already satisfied: kiwisolver>=1.3.1 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=1.5->pymatgen==2024.3.1) (1.4.5) Requirement already satisfied: cycler>=0.10 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=1.5->pymatgen==2024.3.1) (0.12.1) Requirement already satisfied: fonttools>=4.22.0 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=1.5->pymatgen==2024.3.1) (4.53.0) Requirement already satisfied: packaging>=20.0 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=1.5->pymatgen==2024.3.1) (23.2) Requirement already satisfied: pyparsing>=2.3.1 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=1.5->pymatgen==2024.3.1) (3.1.2) Requirement already satisfied: tenacity>=6.2.0 in /opt/mamba/lib/python3.10/site-packages (from plotly>=4.5.0->pymatgen==2024.3.1) (8.3.0) Requirement already satisfied: ruamel.yaml.clib>=0.2.7 in /opt/mamba/lib/python3.10/site-packages (from ruamel.yaml>=0.17.0->pymatgen==2024.3.1) (0.2.8) Requirement already satisfied: tzdata>=2022.1 in /opt/mamba/lib/python3.10/site-packages (from pandas->pymatgen==2024.3.1) (2023.3) Requirement already satisfied: pytz>=2020.1 in /opt/mamba/lib/python3.10/site-packages (from pandas->pymatgen==2024.3.1) (2023.3.post1) Requirement already satisfied: latexcodec>=1.0.4 in /opt/mamba/lib/python3.10/site-packages (from pybtex->pymatgen==2024.3.1) (3.0.0) Requirement already satisfied: PyYAML>=3.01 in /opt/mamba/lib/python3.10/site-packages (from pybtex->pymatgen==2024.3.1) (6.0.1) Requirement already satisfied: six in /opt/mamba/lib/python3.10/site-packages (from pybtex->pymatgen==2024.3.1) (1.16.0) Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/mamba/lib/python3.10/site-packages (from requests->pymatgen==2024.3.1) (1.26.11) Requirement already satisfied: certifi>=2017.4.17 in /opt/mamba/lib/python3.10/site-packages (from requests->pymatgen==2024.3.1) (2022.9.24) Requirement already satisfied: charset-normalizer<4,>=2 in /opt/mamba/lib/python3.10/site-packages (from requests->pymatgen==2024.3.1) (2.1.1) Requirement already satisfied: idna<4,>=2.5 in /opt/mamba/lib/python3.10/site-packages (from requests->pymatgen==2024.3.1) (3.4) Requirement already satisfied: mpmath<1.4.0,>=1.1.0 in /opt/mamba/lib/python3.10/site-packages (from sympy->pymatgen==2024.3.1) (1.3.0) WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple Requirement already satisfied: pyxtal==0.6.3 in /opt/mamba/lib/python3.10/site-packages (0.6.3) Requirement already satisfied: pymatgen>=2023.10.3 in /opt/mamba/lib/python3.10/site-packages (from pyxtal==0.6.3) (2024.3.1) Requirement already satisfied: ase>=3.18.0 in /opt/mamba/lib/python3.10/site-packages (from pyxtal==0.6.3) (3.23.0) Requirement already satisfied: scipy>=1.7.3 in /opt/mamba/lib/python3.10/site-packages (from pyxtal==0.6.3) (1.11.4) Requirement already satisfied: pandas>=2.0.2 in /opt/mamba/lib/python3.10/site-packages (from pyxtal==0.6.3) (2.1.4) Requirement already satisfied: spglib>=1.10.4 in /opt/mamba/lib/python3.10/site-packages (from pyxtal==0.6.3) (2.4.0) Requirement already satisfied: networkx>=2.3 in /opt/mamba/lib/python3.10/site-packages (from pyxtal==0.6.3) (3.3) Requirement already satisfied: importlib-metadata>=1.4 in /opt/mamba/lib/python3.10/site-packages (from pyxtal==0.6.3) (7.1.0) Requirement already satisfied: matplotlib>=3.3.4 in /opt/mamba/lib/python3.10/site-packages (from ase>=3.18.0->pyxtal==0.6.3) (3.9.0) Requirement already satisfied: numpy>=1.18.5 in /opt/mamba/lib/python3.10/site-packages (from ase>=3.18.0->pyxtal==0.6.3) (1.26.2) Requirement already satisfied: zipp>=0.5 in /opt/mamba/lib/python3.10/site-packages (from importlib-metadata>=1.4->pyxtal==0.6.3) (3.19.2) Requirement already satisfied: pytz>=2020.1 in /opt/mamba/lib/python3.10/site-packages (from pandas>=2.0.2->pyxtal==0.6.3) (2023.3.post1) Requirement already satisfied: tzdata>=2022.1 in /opt/mamba/lib/python3.10/site-packages (from pandas>=2.0.2->pyxtal==0.6.3) (2023.3) Requirement already satisfied: python-dateutil>=2.8.2 in /opt/mamba/lib/python3.10/site-packages (from pandas>=2.0.2->pyxtal==0.6.3) (2.8.2) Requirement already satisfied: plotly>=4.5.0 in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (5.22.0) Requirement already satisfied: monty>=2024.2.2 in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (2024.5.24) Requirement already satisfied: tabulate in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (0.9.0) Requirement already satisfied: ruamel.yaml>=0.17.0 in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (0.18.6) Requirement already satisfied: palettable>=3.1.1 in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (3.3.3) Requirement already satisfied: requests in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (2.31.0) Requirement already satisfied: joblib in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (1.4.2) Requirement already satisfied: tqdm in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (4.64.1) Requirement already satisfied: pybtex in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (0.24.0) Requirement already satisfied: sympy in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (1.12.1) Requirement already satisfied: uncertainties>=3.1.4 in /opt/mamba/lib/python3.10/site-packages (from pymatgen>=2023.10.3->pyxtal==0.6.3) (3.2.1) Requirement already satisfied: kiwisolver>=1.3.1 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=3.3.4->ase>=3.18.0->pyxtal==0.6.3) (1.4.5) Requirement already satisfied: pyparsing>=2.3.1 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=3.3.4->ase>=3.18.0->pyxtal==0.6.3) (3.1.2) Requirement already satisfied: fonttools>=4.22.0 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=3.3.4->ase>=3.18.0->pyxtal==0.6.3) (4.53.0) Requirement already satisfied: pillow>=8 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=3.3.4->ase>=3.18.0->pyxtal==0.6.3) (10.3.0) Requirement already satisfied: packaging>=20.0 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=3.3.4->ase>=3.18.0->pyxtal==0.6.3) (23.2) Requirement already satisfied: cycler>=0.10 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=3.3.4->ase>=3.18.0->pyxtal==0.6.3) (0.12.1) Requirement already satisfied: contourpy>=1.0.1 in /opt/mamba/lib/python3.10/site-packages (from matplotlib>=3.3.4->ase>=3.18.0->pyxtal==0.6.3) (1.2.1) Requirement already satisfied: tenacity>=6.2.0 in /opt/mamba/lib/python3.10/site-packages (from plotly>=4.5.0->pymatgen>=2023.10.3->pyxtal==0.6.3) (8.3.0) Requirement already satisfied: six>=1.5 in /opt/mamba/lib/python3.10/site-packages (from python-dateutil>=2.8.2->pandas>=2.0.2->pyxtal==0.6.3) (1.16.0) Requirement already satisfied: ruamel.yaml.clib>=0.2.7 in /opt/mamba/lib/python3.10/site-packages (from ruamel.yaml>=0.17.0->pymatgen>=2023.10.3->pyxtal==0.6.3) (0.2.8) Requirement already satisfied: latexcodec>=1.0.4 in /opt/mamba/lib/python3.10/site-packages (from pybtex->pymatgen>=2023.10.3->pyxtal==0.6.3) (3.0.0) Requirement already satisfied: PyYAML>=3.01 in /opt/mamba/lib/python3.10/site-packages (from pybtex->pymatgen>=2023.10.3->pyxtal==0.6.3) (6.0.1) Requirement already satisfied: idna<4,>=2.5 in /opt/mamba/lib/python3.10/site-packages (from requests->pymatgen>=2023.10.3->pyxtal==0.6.3) (3.4) Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/mamba/lib/python3.10/site-packages (from requests->pymatgen>=2023.10.3->pyxtal==0.6.3) (1.26.11) Requirement already satisfied: certifi>=2017.4.17 in /opt/mamba/lib/python3.10/site-packages (from requests->pymatgen>=2023.10.3->pyxtal==0.6.3) (2022.9.24) Requirement already satisfied: charset-normalizer<4,>=2 in /opt/mamba/lib/python3.10/site-packages (from requests->pymatgen>=2023.10.3->pyxtal==0.6.3) (2.1.1) Requirement already satisfied: mpmath<1.4.0,>=1.1.0 in /opt/mamba/lib/python3.10/site-packages (from sympy->pymatgen>=2023.10.3->pyxtal==0.6.3) (1.3.0) WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple Requirement already satisfied: hydra-core in /opt/mamba/lib/python3.10/site-packages (1.3.2) Requirement already satisfied: omegaconf<2.4,>=2.2 in /opt/mamba/lib/python3.10/site-packages (from hydra-core) (2.3.0) Requirement already satisfied: packaging in /opt/mamba/lib/python3.10/site-packages (from hydra-core) (23.2) Requirement already satisfied: antlr4-python3-runtime==4.9.* in /opt/mamba/lib/python3.10/site-packages (from hydra-core) (4.9.3) Requirement already satisfied: PyYAML>=5.1.0 in /opt/mamba/lib/python3.10/site-packages (from omegaconf<2.4,>=2.2->hydra-core) (6.0.1) WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple Requirement already satisfied: nglview in /opt/mamba/lib/python3.10/site-packages (3.1.2) Requirement already satisfied: notebook>=7 in /opt/mamba/lib/python3.10/site-packages (from nglview) (7.2.1) Requirement already satisfied: ipywidgets>=8 in /opt/mamba/lib/python3.10/site-packages (from nglview) (8.1.3) Requirement already satisfied: numpy in /opt/mamba/lib/python3.10/site-packages (from nglview) (1.26.2) Requirement already satisfied: jupyterlab-widgets in /opt/mamba/lib/python3.10/site-packages (from nglview) (3.0.11) Requirement already satisfied: widgetsnbextension~=4.0.11 in /opt/mamba/lib/python3.10/site-packages (from ipywidgets>=8->nglview) (4.0.11) Requirement already satisfied: traitlets>=4.3.1 in /opt/mamba/lib/python3.10/site-packages (from ipywidgets>=8->nglview) (5.14.0) Requirement already satisfied: ipython>=6.1.0 in /opt/mamba/lib/python3.10/site-packages (from ipywidgets>=8->nglview) (8.18.1) Requirement already satisfied: comm>=0.1.3 in /opt/mamba/lib/python3.10/site-packages (from ipywidgets>=8->nglview) (0.2.0) Requirement already satisfied: notebook-shim<0.3,>=0.2 in /opt/mamba/lib/python3.10/site-packages (from notebook>=7->nglview) (0.2.3) Requirement already satisfied: jupyterlab-server<3,>=2.27.1 in /opt/mamba/lib/python3.10/site-packages (from notebook>=7->nglview) (2.27.2) Requirement already satisfied: jupyter-server<3,>=2.4.0 in /opt/mamba/lib/python3.10/site-packages (from notebook>=7->nglview) (2.12.1) Requirement already satisfied: tornado>=6.2.0 in /opt/mamba/lib/python3.10/site-packages (from notebook>=7->nglview) (6.4) Requirement already satisfied: jupyterlab<4.3,>=4.2.0 in /opt/mamba/lib/python3.10/site-packages (from notebook>=7->nglview) (4.2.2) Requirement already satisfied: decorator in /opt/mamba/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets>=8->nglview) (5.1.1) Requirement already satisfied: exceptiongroup in /opt/mamba/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets>=8->nglview) (1.2.0) Requirement already satisfied: pygments>=2.4.0 in /opt/mamba/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets>=8->nglview) (2.17.2) Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.41 in /opt/mamba/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets>=8->nglview) (3.0.42) Requirement already satisfied: matplotlib-inline in /opt/mamba/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets>=8->nglview) (0.1.6) Requirement already satisfied: jedi>=0.16 in /opt/mamba/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets>=8->nglview) (0.19.1) Requirement already satisfied: stack-data in /opt/mamba/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets>=8->nglview) (0.6.3) Requirement already satisfied: pexpect>4.3 in /opt/mamba/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets>=8->nglview) (4.9.0) Requirement already satisfied: jupyter-core!=5.0.*,>=4.12 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (5.5.0) Requirement already satisfied: nbconvert>=6.4.4 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (7.12.0) Requirement already satisfied: pyzmq>=24 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (25.1.2) Requirement already satisfied: jinja2 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (3.1.2) Requirement already satisfied: websocket-client in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (1.7.0) Requirement already satisfied: prometheus-client in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.19.0) Requirement already satisfied: terminado>=0.8.3 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.18.0) Requirement already satisfied: send2trash>=1.8.2 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (1.8.2) Requirement already satisfied: nbformat>=5.3.0 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (5.9.2) Requirement already satisfied: jupyter-events>=0.9.0 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.9.0) Requirement already satisfied: overrides in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (7.4.0) Requirement already satisfied: anyio>=3.1.0 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (4.1.0) Requirement already satisfied: packaging in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (23.2) Requirement already satisfied: argon2-cffi in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (23.1.0) Requirement already satisfied: jupyter-client>=7.4.4 in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (8.6.0) Requirement already satisfied: jupyter-server-terminals in /opt/mamba/lib/python3.10/site-packages (from jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.5.0) Requirement already satisfied: ipykernel>=6.5.0 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (6.27.1) Requirement already satisfied: setuptools>=40.1.0 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (65.5.0) Requirement already satisfied: tomli>=1.2.2 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (2.0.1) Requirement already satisfied: httpx>=0.25.0 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (0.27.0) Requirement already satisfied: jupyter-lsp>=2.0.0 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (2.2.1) Requirement already satisfied: async-lru>=1.0.0 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (2.0.4) Requirement already satisfied: babel>=2.10 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (2.14.0) Requirement already satisfied: json5>=0.9.0 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (0.9.14) Requirement already satisfied: jsonschema>=4.18.0 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (4.20.0) Requirement already satisfied: requests>=2.31 in /opt/mamba/lib/python3.10/site-packages (from jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (2.31.0) Requirement already satisfied: idna>=2.8 in /opt/mamba/lib/python3.10/site-packages (from anyio>=3.1.0->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (3.4) Requirement already satisfied: sniffio>=1.1 in /opt/mamba/lib/python3.10/site-packages (from anyio>=3.1.0->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (1.3.0) Requirement already satisfied: typing-extensions>=4.0.0 in /opt/mamba/lib/python3.10/site-packages (from async-lru>=1.0.0->jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (4.9.0) Requirement already satisfied: certifi in /opt/mamba/lib/python3.10/site-packages (from httpx>=0.25.0->jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (2022.9.24) Requirement already satisfied: httpcore==1.* in /opt/mamba/lib/python3.10/site-packages (from httpx>=0.25.0->jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (1.0.5) Requirement already satisfied: h11<0.15,>=0.13 in /opt/mamba/lib/python3.10/site-packages (from httpcore==1.*->httpx>=0.25.0->jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (0.14.0) Requirement already satisfied: debugpy>=1.6.5 in /opt/mamba/lib/python3.10/site-packages (from ipykernel>=6.5.0->jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (1.8.0) Requirement already satisfied: psutil in /opt/mamba/lib/python3.10/site-packages (from ipykernel>=6.5.0->jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (5.9.6) Requirement already satisfied: nest-asyncio in /opt/mamba/lib/python3.10/site-packages (from ipykernel>=6.5.0->jupyterlab<4.3,>=4.2.0->notebook>=7->nglview) (1.5.8) Requirement already satisfied: parso<0.9.0,>=0.8.3 in /opt/mamba/lib/python3.10/site-packages (from jedi>=0.16->ipython>=6.1.0->ipywidgets>=8->nglview) (0.8.3) Requirement already satisfied: MarkupSafe>=2.0 in /opt/mamba/lib/python3.10/site-packages (from jinja2->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (2.1.3) Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /opt/mamba/lib/python3.10/site-packages (from jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (2023.11.2) Requirement already satisfied: rpds-py>=0.7.1 in /opt/mamba/lib/python3.10/site-packages (from jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (0.13.2) Requirement already satisfied: referencing>=0.28.4 in /opt/mamba/lib/python3.10/site-packages (from jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (0.32.0) Requirement already satisfied: attrs>=22.2.0 in /opt/mamba/lib/python3.10/site-packages (from jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (23.1.0) Requirement already satisfied: python-dateutil>=2.8.2 in /opt/mamba/lib/python3.10/site-packages (from jupyter-client>=7.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (2.8.2) Requirement already satisfied: platformdirs>=2.5 in /opt/mamba/lib/python3.10/site-packages (from jupyter-core!=5.0.*,>=4.12->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (4.1.0) Requirement already satisfied: rfc3986-validator>=0.1.1 in /opt/mamba/lib/python3.10/site-packages (from jupyter-events>=0.9.0->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.1.1) Requirement already satisfied: python-json-logger>=2.0.4 in /opt/mamba/lib/python3.10/site-packages (from jupyter-events>=0.9.0->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (2.0.7) Requirement already satisfied: rfc3339-validator in /opt/mamba/lib/python3.10/site-packages (from jupyter-events>=0.9.0->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.1.4) Requirement already satisfied: pyyaml>=5.3 in /opt/mamba/lib/python3.10/site-packages (from jupyter-events>=0.9.0->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (6.0.1) Requirement already satisfied: beautifulsoup4 in /opt/mamba/lib/python3.10/site-packages (from nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (4.12.2) Requirement already satisfied: defusedxml in /opt/mamba/lib/python3.10/site-packages (from nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.7.1) Requirement already satisfied: jupyterlab-pygments in /opt/mamba/lib/python3.10/site-packages (from nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.3.0) Requirement already satisfied: mistune<4,>=2.0.3 in /opt/mamba/lib/python3.10/site-packages (from nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (3.0.2) Requirement already satisfied: bleach!=5.0.0 in /opt/mamba/lib/python3.10/site-packages (from nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (6.1.0) Requirement already satisfied: nbclient>=0.5.0 in /opt/mamba/lib/python3.10/site-packages (from nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.9.0) Requirement already satisfied: pandocfilters>=1.4.1 in /opt/mamba/lib/python3.10/site-packages (from nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (1.5.0) Requirement already satisfied: tinycss2 in /opt/mamba/lib/python3.10/site-packages (from nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (1.2.1) Requirement already satisfied: fastjsonschema in /opt/mamba/lib/python3.10/site-packages (from nbformat>=5.3.0->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (2.19.0) Requirement already satisfied: ptyprocess>=0.5 in /opt/mamba/lib/python3.10/site-packages (from pexpect>4.3->ipython>=6.1.0->ipywidgets>=8->nglview) (0.7.0) Requirement already satisfied: wcwidth in /opt/mamba/lib/python3.10/site-packages (from prompt-toolkit<3.1.0,>=3.0.41->ipython>=6.1.0->ipywidgets>=8->nglview) (0.2.12) Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/mamba/lib/python3.10/site-packages (from requests>=2.31->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (1.26.11) Requirement already satisfied: charset-normalizer<4,>=2 in /opt/mamba/lib/python3.10/site-packages (from requests>=2.31->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (2.1.1) Requirement already satisfied: argon2-cffi-bindings in /opt/mamba/lib/python3.10/site-packages (from argon2-cffi->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (21.2.0) Requirement already satisfied: executing>=1.2.0 in /opt/mamba/lib/python3.10/site-packages (from stack-data->ipython>=6.1.0->ipywidgets>=8->nglview) (2.0.1) Requirement already satisfied: asttokens>=2.1.0 in /opt/mamba/lib/python3.10/site-packages (from stack-data->ipython>=6.1.0->ipywidgets>=8->nglview) (2.4.1) Requirement already satisfied: pure-eval in /opt/mamba/lib/python3.10/site-packages (from stack-data->ipython>=6.1.0->ipywidgets>=8->nglview) (0.2.2) Requirement already satisfied: six>=1.12.0 in /opt/mamba/lib/python3.10/site-packages (from asttokens>=2.1.0->stack-data->ipython>=6.1.0->ipywidgets>=8->nglview) (1.16.0) Requirement already satisfied: webencodings in /opt/mamba/lib/python3.10/site-packages (from bleach!=5.0.0->nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (0.5.1) Requirement already satisfied: jsonpointer>1.13 in /opt/mamba/lib/python3.10/site-packages (from jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (2.4) Requirement already satisfied: webcolors>=1.11 in /opt/mamba/lib/python3.10/site-packages (from jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (1.13) Requirement already satisfied: fqdn in /opt/mamba/lib/python3.10/site-packages (from jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (1.5.1) Requirement already satisfied: isoduration in /opt/mamba/lib/python3.10/site-packages (from jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (20.11.0) Requirement already satisfied: uri-template in /opt/mamba/lib/python3.10/site-packages (from jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (1.3.0) Requirement already satisfied: cffi>=1.0.1 in /opt/mamba/lib/python3.10/site-packages (from argon2-cffi-bindings->argon2-cffi->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (1.15.1) Requirement already satisfied: soupsieve>1.2 in /opt/mamba/lib/python3.10/site-packages (from beautifulsoup4->nbconvert>=6.4.4->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (2.5) Requirement already satisfied: pycparser in /opt/mamba/lib/python3.10/site-packages (from cffi>=1.0.1->argon2-cffi-bindings->argon2-cffi->jupyter-server<3,>=2.4.0->notebook>=7->nglview) (2.21) Requirement already satisfied: arrow>=0.15.0 in /opt/mamba/lib/python3.10/site-packages (from isoduration->jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (1.3.0) Requirement already satisfied: types-python-dateutil>=2.8.10 in /opt/mamba/lib/python3.10/site-packages (from arrow>=0.15.0->isoduration->jsonschema>=4.18.0->jupyterlab-server<3,>=2.27.1->notebook>=7->nglview) (2.8.19.14) WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
代码
文本
加载训练好的模型
注意这里加载的模型为arxiv文章中模型的mini版本 （参数量约为原本的十分之一）
代码
文本
[4]
import sys
sys.path.append('./src/')
import jax
import jax.numpy as jnp
from jax.flatten_util import ravel_pytree
from hydra import initialize, compose
import checkpoint
from transformer import make_transformer
with initialize(version_base=None, config_path="./CrystalFormer/model"):
args = compose(config_name="config")
print(args)
key = jax.random.PRNGKey(42)
params, transformer = make_transformer(key, args.Nf, args.Kx, args.Kl, args.n_max,
args.h0_size,
4, 8,
32, args.model_size, args.embed_size,
args.atom_types, args.wyck_types,
0.3)
print("\n========== Load checkpoint==========")
restore_path = "/share/"
ckpt_filename, epoch_finished = checkpoint.find_ckpt_filename(restore_path)
if ckpt_filename is not None:
print("Load checkpoint file: %s, epoch finished: %g" %(ckpt_filename, epoch_finished))
ckpt = checkpoint.load_data(ckpt_filename)
params = ckpt["params"]
else:
print("No checkpoint file found. Start from scratch.")
print ("# of transformer params", ravel_pytree(params)[0].size)
{'epochs': 10000, 'batchsize': 100, 'lr': 0.0001, 'lr_decay': 0.0, 'weight_decay': 0.0, 'clip_grad': 1.0, 'optimizer': 'adam', 'folder': './', 'restore_path': None, 'training_name': '${folder}${optimizer}_bs_${batchsize}_lr_${lr}_decay_${lr_decay}_clip_${clip_grad}', 'train_path': '/data/zdcao/crystal_gpt/dataset/mp_20/train.csv', 'valid_path': '/data/zdcao/crystal_gpt/dataset/mp_20/val.csv', 'test_path': '/data/zdcao/crystal_gpt/dataset/mp_20/test.csv', 'Nf': 5, 'Kx': 16, 'Kl': 4, 'h0_size': 256, 'transformer_layers': 16, 'num_heads': 16, 'key_size': 64, 'model_size': 64, 'embed_size': 32, 'dropout_rate': 0.5, 'transformer_name': 'Nf_${Nf}_Kx_${Kx}_Kl_${Kl}_h0_${h0_size}_l_${transformer_layers}_H_${num_heads}_k_${key_size}_m_${model_size}_e_${embed_size}_drop_${dropout_rate}', 'lamb_a': 1.0, 'lamb_w': 1.0, 'lamb_l': 1.0, 'loss_name': 'a_${lamb_a}_w_${lamb_w}_l_${lamb_l}', 'n_max': 21, 'atom_types': 119, 'wyck_types': 28, 'physics_name': 'A_${atom_types}_W_${wyck_types}_N_${n_max}', 'spacegroup': None, 'elements': None, 'top_p': 1.0, 'temperature': 1.0, 'num_io_process': 40, 'num_samples': 1000, 'use_foriloop': True, 'output_filename': 'output.csv'}
========== Load checkpoint==========
Load checkpoint file: /share/epoch_009800.pkl, epoch finished: 9800
# of transformer params 461159
代码
文本
加载生成晶体的函数
代码
文本
[5]
import numpy as np
from pymatgen.core import Structure, Lattice
from time import time
from pymatgen.io.ase import AseAtomsAdaptor
from ase.visualize import view
from sample import sample_crystal
from elements import element_dict, element_list
from scripts.awl2struct import get_struct_from_lawx
jax.config.update("jax_enable_x64", True) # to get off compilation warning, and to prevent sample nan lattice
代码
文本
生成晶体
代码
文本
运行下面代码即可得到生成的晶体的3D结构，运行成功结果如下图所示：
代码
文本
[6]
def generate_and_visualize(spacegroup, elements, temperature, seed):
print(f"Generating with spacegroup={spacegroup}, elements={elements}, temperature={temperature}")
top_p = 1
n_sample = 1
elements = elements.split()
if elements is not None:
idx = [element_dict[e] for e in elements]
atom_mask = [1] + [1 if a in idx else 0 for a in range(1, args.atom_types)]
atom_mask = jnp.array(atom_mask)
# print ('sampling structure formed by these elements:', elements)
# print (atom_mask)
else:
atom_mask = jnp.zeros((args.atom_types), dtype=int) # we will do nothing to a_logit in sampling
key = jax.random.PRNGKey(seed)
key, subkey = jax.random.split(key)
start_time = time()
XYZ, A, W, M, L = sample_crystal(subkey, transformer, params, args.n_max, n_sample, args.atom_types, args.wyck_types, args.Kx, args.Kl, spacegroup, None, atom_mask, top_p, temperature, temperature, args.use_foriloop)
end_time = time()
print("executation time:", end_time - start_time)
XYZ = np.array(XYZ)
A = np.array(A)
W = np.array(W)
L = np.array(L)
G = np.array([spacegroup for i in range(len(L))])
structures = [get_struct_from_lawx(g, l, a, w, xyz) for g, l, a, w, xyz in zip(G, L, A, W, XYZ)]
structures = [Structure.from_dict(_) for _ in structures]
atoms_list = [AseAtomsAdaptor().get_atoms(struct) for struct in structures]
return view(atoms_list[0], viewer='ngl')
代码
文本
[7]
# ============= params to control the generation =============
spacegroup = 225 # 设置生成的晶体的空间群，范围为1-230
elements = "" # 限制生成晶体所包含的元素种类，每个元素需要用空格隔开，比如 "Ba Ti O"
temperature = 1.0 # 控制transformer生成的温度，温度越高生成的novelty越高，推荐值为 0.5到1.5
seed = 42 # 随机种子
# =============== generate and visualization =================
generate_and_visualize(spacegroup, elements, temperature, seed)
Generating with spacegroup=225, elements=, temperature=1.0 executation time: 30.423662424087524
代码
文本
已赞7
本文被以下合集收录
材料+深度学习
Hayden
更新于 2024-08-24
6 篇0 人关注
晶体合成
yuxiangc22
更新于 2024-06-26
1 篇0 人关注
推荐阅读
公开
基于生成模型的晶体结构预测案例
2043899742@qq.com
发布于 2023-07-30
3 赞3 转存文件1 评论
公开
基于生成模型的晶体结构预测案例副本CyrusZHOU
发布于 2024-03-16
评论
SchrodingersCat
Russel