cgcnn2 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- cgcnn2-0.1.0/PKG-INFO +75 -0
- cgcnn2-0.1.0/README.md +45 -0
- cgcnn2-0.1.0/cgcnn2/__init__.py +3 -0
- cgcnn2-0.1.0/cgcnn2/cgcnn_data.py +392 -0
- cgcnn2-0.1.0/cgcnn2/cgcnn_model.py +372 -0
- cgcnn2-0.1.0/cgcnn2/cgcnn_utils.py +474 -0
- cgcnn2-0.1.0/pyproject.toml +39 -0
cgcnn2-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: cgcnn2
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Crystal Graph Convolutional Neural Networks
|
|
5
|
+
Home-page: https://github.com/jcwang587/cgcnn2/
|
|
6
|
+
License: MIT
|
|
7
|
+
Keywords: python,gnn,vasp,crystal
|
|
8
|
+
Author: Jiacheng Wang
|
|
9
|
+
Author-email: jiachengwang@umass.edu
|
|
10
|
+
Maintainer: Jiacheng Wang
|
|
11
|
+
Requires-Python: >=3.10,<3.12
|
|
12
|
+
Classifier: Development Status :: 1 - Planning
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
15
|
+
Classifier: Operating System :: OS Independent
|
|
16
|
+
Classifier: Programming Language :: Python :: 3
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
20
|
+
Requires-Dist: ase (==3.23.0)
|
|
21
|
+
Requires-Dist: numpy
|
|
22
|
+
Requires-Dist: pandas
|
|
23
|
+
Requires-Dist: pymatgen (==2024.10.3)
|
|
24
|
+
Requires-Dist: pymatviz (==0.13.2)
|
|
25
|
+
Requires-Dist: scikit-learn
|
|
26
|
+
Requires-Dist: torch (==2.5.1)
|
|
27
|
+
Project-URL: Repository, https://github.com/jcwang587/cgcnn2/
|
|
28
|
+
Description-Content-Type: text/markdown
|
|
29
|
+
|
|
30
|
+
# CGCNN2
|
|
31
|
+
|
|
32
|
+
As the original Crystal Graph Convolutional Neural Networks (CGCNN) repository is no longer actively maintained, this repository is a reproduction of [CGCNN](https://github.com/txie-93/cgcnn) by Xie et al. It includes necessary updates for deprecated components and a few additional functions to ensure smooth operation. Despite its age, CGCNN remains a straightforward and fast deep learning framework that is easy to learn and use.
|
|
33
|
+
|
|
34
|
+
The package provides following major functions:
|
|
35
|
+
|
|
36
|
+
- **Training** a CGCNN model with a customized dataset.
|
|
37
|
+
- **Predicting** material properties with a pre-trained CGCNN model.
|
|
38
|
+
- **Fine-tuning** a pre-trained CGCNN model on a new dataset.
|
|
39
|
+
|
|
40
|
+
## Installation
|
|
41
|
+
|
|
42
|
+
Make sure you have a Python interpreter, preferably version 3.10 or higher. Then, you can simply install xdatbus from
|
|
43
|
+
PyPI using `pip`:
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
pip install cgcnn2
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
If you'd like to use the latest unreleased version on the main branch, you can install it directly from GitHub:
|
|
50
|
+
|
|
51
|
+
```bash
|
|
52
|
+
pip install git+https://github.com/jcwang587/cgcnn2
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
## References
|
|
56
|
+
|
|
57
|
+
The original paper describes the details of the CGCNN framework:
|
|
58
|
+
|
|
59
|
+
```bibtex
|
|
60
|
+
@article{PhysRevLett.120.145301,
|
|
61
|
+
title = {Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties},
|
|
62
|
+
author = {Xie, Tian and Grossman, Jeffrey C.},
|
|
63
|
+
journal = {Phys. Rev. Lett.},
|
|
64
|
+
volume = {120},
|
|
65
|
+
issue = {14},
|
|
66
|
+
pages = {145301},
|
|
67
|
+
numpages = {6},
|
|
68
|
+
year = {2018},
|
|
69
|
+
month = {Apr},
|
|
70
|
+
publisher = {American Physical Society},
|
|
71
|
+
doi = {10.1103/PhysRevLett.120.145301},
|
|
72
|
+
url = {https://link.aps.org/doi/10.1103/PhysRevLett.120.145301}
|
|
73
|
+
}
|
|
74
|
+
```
|
|
75
|
+
|
cgcnn2-0.1.0/README.md
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
# CGCNN2
|
|
2
|
+
|
|
3
|
+
As the original Crystal Graph Convolutional Neural Networks (CGCNN) repository is no longer actively maintained, this repository is a reproduction of [CGCNN](https://github.com/txie-93/cgcnn) by Xie et al. It includes necessary updates for deprecated components and a few additional functions to ensure smooth operation. Despite its age, CGCNN remains a straightforward and fast deep learning framework that is easy to learn and use.
|
|
4
|
+
|
|
5
|
+
The package provides following major functions:
|
|
6
|
+
|
|
7
|
+
- **Training** a CGCNN model with a customized dataset.
|
|
8
|
+
- **Predicting** material properties with a pre-trained CGCNN model.
|
|
9
|
+
- **Fine-tuning** a pre-trained CGCNN model on a new dataset.
|
|
10
|
+
|
|
11
|
+
## Installation
|
|
12
|
+
|
|
13
|
+
Make sure you have a Python interpreter, preferably version 3.10 or higher. Then, you can simply install xdatbus from
|
|
14
|
+
PyPI using `pip`:
|
|
15
|
+
|
|
16
|
+
```bash
|
|
17
|
+
pip install cgcnn2
|
|
18
|
+
```
|
|
19
|
+
|
|
20
|
+
If you'd like to use the latest unreleased version on the main branch, you can install it directly from GitHub:
|
|
21
|
+
|
|
22
|
+
```bash
|
|
23
|
+
pip install git+https://github.com/jcwang587/cgcnn2
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
## References
|
|
27
|
+
|
|
28
|
+
The original paper describes the details of the CGCNN framework:
|
|
29
|
+
|
|
30
|
+
```bibtex
|
|
31
|
+
@article{PhysRevLett.120.145301,
|
|
32
|
+
title = {Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties},
|
|
33
|
+
author = {Xie, Tian and Grossman, Jeffrey C.},
|
|
34
|
+
journal = {Phys. Rev. Lett.},
|
|
35
|
+
volume = {120},
|
|
36
|
+
issue = {14},
|
|
37
|
+
pages = {145301},
|
|
38
|
+
numpages = {6},
|
|
39
|
+
year = {2018},
|
|
40
|
+
month = {Apr},
|
|
41
|
+
publisher = {American Physical Society},
|
|
42
|
+
doi = {10.1103/PhysRevLett.120.145301},
|
|
43
|
+
url = {https://link.aps.org/doi/10.1103/PhysRevLett.120.145301}
|
|
44
|
+
}
|
|
45
|
+
```
|
|
@@ -0,0 +1,392 @@
|
|
|
1
|
+
import csv
|
|
2
|
+
import functools
|
|
3
|
+
import json
|
|
4
|
+
import os
|
|
5
|
+
import random
|
|
6
|
+
import warnings
|
|
7
|
+
|
|
8
|
+
import numpy as np
|
|
9
|
+
import torch
|
|
10
|
+
from pymatgen.core.structure import Structure
|
|
11
|
+
from torch.utils.data import Dataset
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def collate_pool(dataset_list):
|
|
15
|
+
"""
|
|
16
|
+
Collate a list of data and return a batch for predicting crystal
|
|
17
|
+
properties.
|
|
18
|
+
|
|
19
|
+
Parameters
|
|
20
|
+
----------
|
|
21
|
+
|
|
22
|
+
dataset_list: list of tuples for each data point.
|
|
23
|
+
(atom_fea, nbr_fea, nbr_fea_idx, target)
|
|
24
|
+
|
|
25
|
+
atom_fea: torch.Tensor shape (n_i, atom_fea_len)
|
|
26
|
+
nbr_fea: torch.Tensor shape (n_i, M, nbr_fea_len)
|
|
27
|
+
nbr_fea_idx: torch.LongTensor shape (n_i, M)
|
|
28
|
+
target: torch.Tensor shape (1, )
|
|
29
|
+
cif_id: str or int
|
|
30
|
+
|
|
31
|
+
Returns
|
|
32
|
+
-------
|
|
33
|
+
N = sum(n_i); N0 = sum(i)
|
|
34
|
+
|
|
35
|
+
batch_atom_fea: torch.Tensor shape (N, orig_atom_fea_len)
|
|
36
|
+
Atom features from atom type
|
|
37
|
+
batch_nbr_fea: torch.Tensor shape (N, M, nbr_fea_len)
|
|
38
|
+
Bond features of each atom's M neighbors
|
|
39
|
+
batch_nbr_fea_idx: torch.LongTensor shape (N, M)
|
|
40
|
+
Indices of M neighbors of each atom
|
|
41
|
+
crystal_atom_idx: list of torch.LongTensor of length N0
|
|
42
|
+
Mapping from the crystal idx to atom idx
|
|
43
|
+
target: torch.Tensor shape (N, 1)
|
|
44
|
+
Target value for prediction
|
|
45
|
+
batch_cif_ids: list
|
|
46
|
+
"""
|
|
47
|
+
batch_atom_fea, batch_nbr_fea, batch_nbr_fea_idx = [], [], []
|
|
48
|
+
crystal_atom_idx, batch_target = [], []
|
|
49
|
+
batch_cif_ids = []
|
|
50
|
+
base_idx = 0
|
|
51
|
+
for i, ((atom_fea, nbr_fea, nbr_fea_idx), target, cif_id) in enumerate(
|
|
52
|
+
dataset_list
|
|
53
|
+
):
|
|
54
|
+
n_i = atom_fea.shape[0] # number of atoms for this crystal
|
|
55
|
+
batch_atom_fea.append(atom_fea)
|
|
56
|
+
batch_nbr_fea.append(nbr_fea)
|
|
57
|
+
batch_nbr_fea_idx.append(nbr_fea_idx + base_idx)
|
|
58
|
+
new_idx = torch.LongTensor(np.arange(n_i) + base_idx)
|
|
59
|
+
crystal_atom_idx.append(new_idx)
|
|
60
|
+
batch_target.append(target)
|
|
61
|
+
batch_cif_ids.append(cif_id)
|
|
62
|
+
base_idx += n_i
|
|
63
|
+
return (
|
|
64
|
+
(
|
|
65
|
+
torch.cat(batch_atom_fea, dim=0),
|
|
66
|
+
torch.cat(batch_nbr_fea, dim=0),
|
|
67
|
+
torch.cat(batch_nbr_fea_idx, dim=0),
|
|
68
|
+
crystal_atom_idx,
|
|
69
|
+
),
|
|
70
|
+
torch.stack(batch_target, dim=0),
|
|
71
|
+
batch_cif_ids,
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
class GaussianDistance(object):
|
|
76
|
+
"""
|
|
77
|
+
Expands the distance by Gaussian basis.
|
|
78
|
+
|
|
79
|
+
Unit: angstrom
|
|
80
|
+
"""
|
|
81
|
+
|
|
82
|
+
def __init__(self, dmin, dmax, step, var=None):
|
|
83
|
+
"""
|
|
84
|
+
Parameters
|
|
85
|
+
----------
|
|
86
|
+
|
|
87
|
+
dmin: float
|
|
88
|
+
Minimum interatomic distance
|
|
89
|
+
dmax: float
|
|
90
|
+
Maximum interatomic distance
|
|
91
|
+
step: float
|
|
92
|
+
Step size for the Gaussian filter
|
|
93
|
+
"""
|
|
94
|
+
assert dmin < dmax
|
|
95
|
+
assert dmax - dmin > step
|
|
96
|
+
self.filter = np.arange(dmin, dmax + step, step)
|
|
97
|
+
if var is None:
|
|
98
|
+
var = step
|
|
99
|
+
self.var = var
|
|
100
|
+
|
|
101
|
+
def expand(self, distances):
|
|
102
|
+
"""
|
|
103
|
+
Apply Gaussian distance filter to a numpy distance array
|
|
104
|
+
|
|
105
|
+
Parameters
|
|
106
|
+
----------
|
|
107
|
+
|
|
108
|
+
distances: np.ndarray
|
|
109
|
+
A distance matrix of any shape
|
|
110
|
+
|
|
111
|
+
Returns
|
|
112
|
+
-------
|
|
113
|
+
expanded_distance: shape (n+1)-d array
|
|
114
|
+
Expanded distance matrix with the last dimension of length
|
|
115
|
+
len(self.filter)
|
|
116
|
+
"""
|
|
117
|
+
return np.exp(-((distances[..., np.newaxis] - self.filter) ** 2) / self.var**2)
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
class AtomInitializer(object):
|
|
121
|
+
"""
|
|
122
|
+
Base class for initializing the vector representation for atoms.
|
|
123
|
+
|
|
124
|
+
!!! Use one AtomInitializer per dataset !!!
|
|
125
|
+
"""
|
|
126
|
+
|
|
127
|
+
def __init__(self, atom_types):
|
|
128
|
+
self.atom_types = set(atom_types)
|
|
129
|
+
self._embedding = {}
|
|
130
|
+
|
|
131
|
+
def get_atom_fea(self, atom_type):
|
|
132
|
+
assert atom_type in self.atom_types
|
|
133
|
+
return self._embedding[atom_type]
|
|
134
|
+
|
|
135
|
+
def load_state_dict(self, state_dict):
|
|
136
|
+
self._embedding = state_dict
|
|
137
|
+
self.atom_types = set(self._embedding.keys())
|
|
138
|
+
self._decodedict = {
|
|
139
|
+
idx: atom_type for atom_type, idx in self._embedding.items()
|
|
140
|
+
}
|
|
141
|
+
|
|
142
|
+
def state_dict(self):
|
|
143
|
+
return self._embedding
|
|
144
|
+
|
|
145
|
+
def decode(self, idx):
|
|
146
|
+
if not hasattr(self, "_decodedict"):
|
|
147
|
+
self._decodedict = {
|
|
148
|
+
idx: atom_type for atom_type, idx in self._embedding.items()
|
|
149
|
+
}
|
|
150
|
+
return self._decodedict[idx]
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
class AtomCustomJSONInitializer(AtomInitializer):
|
|
154
|
+
"""
|
|
155
|
+
Initialize atom feature vectors using a JSON file, which is a python
|
|
156
|
+
dictionary mapping from element number to a list representing the
|
|
157
|
+
feature vector of the element.
|
|
158
|
+
|
|
159
|
+
Parameters
|
|
160
|
+
----------
|
|
161
|
+
|
|
162
|
+
elem_embedding_file: str
|
|
163
|
+
The path to the .json file
|
|
164
|
+
"""
|
|
165
|
+
|
|
166
|
+
def __init__(self, elem_embedding_file):
|
|
167
|
+
with open(elem_embedding_file) as f:
|
|
168
|
+
elem_embedding = json.load(f)
|
|
169
|
+
elem_embedding = {int(key): value for key, value in elem_embedding.items()}
|
|
170
|
+
atom_types = set(elem_embedding.keys())
|
|
171
|
+
super(AtomCustomJSONInitializer, self).__init__(atom_types)
|
|
172
|
+
for key, value in elem_embedding.items():
|
|
173
|
+
self._embedding[key] = np.array(value, dtype=float)
|
|
174
|
+
|
|
175
|
+
|
|
176
|
+
class CIFData(Dataset):
|
|
177
|
+
"""
|
|
178
|
+
The CIFData dataset is a wrapper for a dataset where the crystal structures
|
|
179
|
+
are stored in the form of CIF files. The dataset should have the following
|
|
180
|
+
directory structure:
|
|
181
|
+
|
|
182
|
+
root_dir
|
|
183
|
+
├── id_prop.csv
|
|
184
|
+
├── atom_init.json
|
|
185
|
+
├── id0.cif
|
|
186
|
+
├── id1.cif
|
|
187
|
+
├── ...
|
|
188
|
+
|
|
189
|
+
id_prop.csv: a CSV file with two columns. The first column recodes a
|
|
190
|
+
unique ID for each crystal, and the second column recodes the value of
|
|
191
|
+
target property.
|
|
192
|
+
|
|
193
|
+
atom_init.json: a JSON file that stores the initialization vector for each
|
|
194
|
+
element.
|
|
195
|
+
|
|
196
|
+
ID.cif: a CIF file that recodes the crystal structure, where ID is the
|
|
197
|
+
unique ID for the crystal.
|
|
198
|
+
|
|
199
|
+
Parameters
|
|
200
|
+
----------
|
|
201
|
+
|
|
202
|
+
root_dir: str
|
|
203
|
+
The path to the root directory of the dataset
|
|
204
|
+
max_num_nbr: int
|
|
205
|
+
The maximum number of neighbors while constructing the crystal graph
|
|
206
|
+
radius: float
|
|
207
|
+
The cutoff radius for searching neighbors
|
|
208
|
+
dmin: float
|
|
209
|
+
The minimum distance for constructing GaussianDistance
|
|
210
|
+
step: float
|
|
211
|
+
The step size for constructing GaussianDistance
|
|
212
|
+
random_seed: int
|
|
213
|
+
Random seed for shuffling the dataset
|
|
214
|
+
|
|
215
|
+
Returns
|
|
216
|
+
-------
|
|
217
|
+
|
|
218
|
+
atom_fea: torch.Tensor shape (n_i, atom_fea_len)
|
|
219
|
+
nbr_fea: torch.Tensor shape (n_i, M, nbr_fea_len)
|
|
220
|
+
nbr_fea_idx: torch.LongTensor shape (n_i, M)
|
|
221
|
+
target: torch.Tensor shape (1, )
|
|
222
|
+
cif_id: str or int
|
|
223
|
+
"""
|
|
224
|
+
|
|
225
|
+
def __init__(
|
|
226
|
+
self, root_dir, max_num_nbr=12, radius=8, dmin=0, step=0.2, random_seed=123
|
|
227
|
+
):
|
|
228
|
+
self.root_dir = root_dir
|
|
229
|
+
self.max_num_nbr, self.radius = max_num_nbr, radius
|
|
230
|
+
assert os.path.exists(root_dir), "root_dir does not exist!"
|
|
231
|
+
id_prop_file = os.path.join(self.root_dir, "id_prop.csv")
|
|
232
|
+
assert os.path.exists(id_prop_file), "id_prop.csv does not exist!"
|
|
233
|
+
with open(id_prop_file) as f:
|
|
234
|
+
reader = csv.reader(f)
|
|
235
|
+
self.id_prop_data = [row for row in reader]
|
|
236
|
+
random.seed(random_seed)
|
|
237
|
+
random.shuffle(self.id_prop_data)
|
|
238
|
+
atom_init_file = os.path.join(self.root_dir, "atom_init.json")
|
|
239
|
+
assert os.path.exists(atom_init_file), "atom_init.json does not exist!"
|
|
240
|
+
self.ari = AtomCustomJSONInitializer(atom_init_file)
|
|
241
|
+
self.gdf = GaussianDistance(dmin=dmin, dmax=self.radius, step=step)
|
|
242
|
+
|
|
243
|
+
def __len__(self):
|
|
244
|
+
return len(self.id_prop_data)
|
|
245
|
+
|
|
246
|
+
@functools.lru_cache(maxsize=None) # Cache loaded structures
|
|
247
|
+
def __getitem__(self, idx):
|
|
248
|
+
cif_id, target = self.id_prop_data[idx]
|
|
249
|
+
crystal = Structure.from_file(os.path.join(self.root_dir, cif_id + ".cif"))
|
|
250
|
+
atom_fea = np.vstack(
|
|
251
|
+
[
|
|
252
|
+
self.ari.get_atom_fea(crystal[i].specie.number)
|
|
253
|
+
for i in range(len(crystal))
|
|
254
|
+
]
|
|
255
|
+
)
|
|
256
|
+
atom_fea = torch.Tensor(atom_fea)
|
|
257
|
+
all_nbrs = crystal.get_all_neighbors(self.radius, include_index=True)
|
|
258
|
+
all_nbrs = [sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs]
|
|
259
|
+
nbr_fea_idx, nbr_fea = [], []
|
|
260
|
+
for nbr in all_nbrs:
|
|
261
|
+
if len(nbr) < self.max_num_nbr:
|
|
262
|
+
warnings.warn(
|
|
263
|
+
"{} not find enough neighbors to build graph. "
|
|
264
|
+
"If it happens frequently, consider increase "
|
|
265
|
+
"radius.".format(cif_id)
|
|
266
|
+
)
|
|
267
|
+
nbr_fea_idx.append(
|
|
268
|
+
list(map(lambda x: x[2], nbr)) + [0] * (self.max_num_nbr - len(nbr))
|
|
269
|
+
)
|
|
270
|
+
nbr_fea.append(
|
|
271
|
+
list(map(lambda x: x[1], nbr))
|
|
272
|
+
+ [self.radius + 1.0] * (self.max_num_nbr - len(nbr))
|
|
273
|
+
)
|
|
274
|
+
else:
|
|
275
|
+
nbr_fea_idx.append(list(map(lambda x: x[2], nbr[: self.max_num_nbr])))
|
|
276
|
+
nbr_fea.append(list(map(lambda x: x[1], nbr[: self.max_num_nbr])))
|
|
277
|
+
nbr_fea_idx, nbr_fea = np.array(nbr_fea_idx), np.array(nbr_fea)
|
|
278
|
+
nbr_fea = self.gdf.expand(nbr_fea)
|
|
279
|
+
atom_fea = torch.Tensor(atom_fea)
|
|
280
|
+
nbr_fea = torch.Tensor(nbr_fea)
|
|
281
|
+
nbr_fea_idx = torch.LongTensor(nbr_fea_idx)
|
|
282
|
+
target = torch.Tensor([float(target)])
|
|
283
|
+
return (atom_fea, nbr_fea, nbr_fea_idx), target, cif_id
|
|
284
|
+
|
|
285
|
+
|
|
286
|
+
class CIFData_pred(Dataset):
|
|
287
|
+
"""
|
|
288
|
+
The CIFData dataset is a wrapper for a dataset where the crystal structures
|
|
289
|
+
are stored in the form of CIF files. The dataset should have the following
|
|
290
|
+
directory structure:
|
|
291
|
+
|
|
292
|
+
root_dir
|
|
293
|
+
├── id_prop.csv
|
|
294
|
+
├── atom_init.json
|
|
295
|
+
├── id0.cif
|
|
296
|
+
├── id1.cif
|
|
297
|
+
├── ...
|
|
298
|
+
|
|
299
|
+
id_prop.csv: a CSV file with two columns. The first column recodes a
|
|
300
|
+
unique ID for each crystal, and the second column recodes the value of
|
|
301
|
+
target property.
|
|
302
|
+
|
|
303
|
+
atom_init.json: a JSON file that stores the initialization vector for each
|
|
304
|
+
element.
|
|
305
|
+
|
|
306
|
+
ID.cif: a CIF file that recodes the crystal structure, where ID is the
|
|
307
|
+
unique ID for the crystal.
|
|
308
|
+
|
|
309
|
+
Parameters
|
|
310
|
+
----------
|
|
311
|
+
|
|
312
|
+
root_dir: str
|
|
313
|
+
The path to the root directory of the dataset
|
|
314
|
+
max_num_nbr: int
|
|
315
|
+
The maximum number of neighbors while constructing the crystal graph
|
|
316
|
+
radius: float
|
|
317
|
+
The cutoff radius for searching neighbors
|
|
318
|
+
dmin: float
|
|
319
|
+
The minimum distance for constructing GaussianDistance
|
|
320
|
+
step: float
|
|
321
|
+
The step size for constructing GaussianDistance
|
|
322
|
+
random_seed: int
|
|
323
|
+
Random seed for shuffling the dataset
|
|
324
|
+
|
|
325
|
+
Returns
|
|
326
|
+
-------
|
|
327
|
+
|
|
328
|
+
atom_fea: torch.Tensor shape (n_i, atom_fea_len)
|
|
329
|
+
nbr_fea: torch.Tensor shape (n_i, M, nbr_fea_len)
|
|
330
|
+
nbr_fea_idx: torch.LongTensor shape (n_i, M)
|
|
331
|
+
target: torch.Tensor shape (1, )
|
|
332
|
+
cif_id: str or int
|
|
333
|
+
"""
|
|
334
|
+
|
|
335
|
+
def __init__(
|
|
336
|
+
self, root_dir, max_num_nbr=12, radius=8, dmin=0, step=0.2, random_seed=123
|
|
337
|
+
):
|
|
338
|
+
self.root_dir = root_dir
|
|
339
|
+
self.max_num_nbr, self.radius = max_num_nbr, radius
|
|
340
|
+
assert os.path.exists(root_dir), "root_dir does not exist!"
|
|
341
|
+
id_prop_file = os.path.join(self.root_dir, "id_prop.csv")
|
|
342
|
+
assert os.path.exists(id_prop_file), "id_prop.csv does not exist!"
|
|
343
|
+
with open(id_prop_file) as f:
|
|
344
|
+
reader = csv.reader(f)
|
|
345
|
+
self.id_prop_data = [row for row in reader]
|
|
346
|
+
random.seed(random_seed)
|
|
347
|
+
atom_init_file = os.path.join(self.root_dir, "atom_init.json")
|
|
348
|
+
assert os.path.exists(atom_init_file), "atom_init.json does not exist!"
|
|
349
|
+
self.ari = AtomCustomJSONInitializer(atom_init_file)
|
|
350
|
+
self.gdf = GaussianDistance(dmin=dmin, dmax=self.radius, step=step)
|
|
351
|
+
|
|
352
|
+
def __len__(self):
|
|
353
|
+
return len(self.id_prop_data)
|
|
354
|
+
|
|
355
|
+
@functools.lru_cache(maxsize=None) # Cache loaded structures
|
|
356
|
+
def __getitem__(self, idx):
|
|
357
|
+
cif_id, target = self.id_prop_data[idx]
|
|
358
|
+
crystal = Structure.from_file(os.path.join(self.root_dir, cif_id + ".cif"))
|
|
359
|
+
atom_fea = np.vstack(
|
|
360
|
+
[
|
|
361
|
+
self.ari.get_atom_fea(crystal[i].specie.number)
|
|
362
|
+
for i in range(len(crystal))
|
|
363
|
+
]
|
|
364
|
+
)
|
|
365
|
+
atom_fea = torch.Tensor(atom_fea)
|
|
366
|
+
all_nbrs = crystal.get_all_neighbors(self.radius, include_index=True)
|
|
367
|
+
all_nbrs = [sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs]
|
|
368
|
+
nbr_fea_idx, nbr_fea = [], []
|
|
369
|
+
for nbr in all_nbrs:
|
|
370
|
+
if len(nbr) < self.max_num_nbr:
|
|
371
|
+
warnings.warn(
|
|
372
|
+
"{} not find enough neighbors to build graph. "
|
|
373
|
+
"If it happens frequently, consider increase "
|
|
374
|
+
"radius.".format(cif_id)
|
|
375
|
+
)
|
|
376
|
+
nbr_fea_idx.append(
|
|
377
|
+
list(map(lambda x: x[2], nbr)) + [0] * (self.max_num_nbr - len(nbr))
|
|
378
|
+
)
|
|
379
|
+
nbr_fea.append(
|
|
380
|
+
list(map(lambda x: x[1], nbr))
|
|
381
|
+
+ [self.radius + 1.0] * (self.max_num_nbr - len(nbr))
|
|
382
|
+
)
|
|
383
|
+
else:
|
|
384
|
+
nbr_fea_idx.append(list(map(lambda x: x[2], nbr[: self.max_num_nbr])))
|
|
385
|
+
nbr_fea.append(list(map(lambda x: x[1], nbr[: self.max_num_nbr])))
|
|
386
|
+
nbr_fea_idx, nbr_fea = np.array(nbr_fea_idx), np.array(nbr_fea)
|
|
387
|
+
nbr_fea = self.gdf.expand(nbr_fea)
|
|
388
|
+
atom_fea = torch.Tensor(atom_fea)
|
|
389
|
+
nbr_fea = torch.Tensor(nbr_fea)
|
|
390
|
+
nbr_fea_idx = torch.LongTensor(nbr_fea_idx)
|
|
391
|
+
target = torch.Tensor([float(target)])
|
|
392
|
+
return (atom_fea, nbr_fea, nbr_fea_idx), target, cif_id
|
|
@@ -0,0 +1,372 @@
|
|
|
1
|
+
"""
|
|
2
|
+
This module implements the Crystal Graph Convolutional Neural Network (CGCNN) for predicting
|
|
3
|
+
material properties based on their crystal structures.
|
|
4
|
+
|
|
5
|
+
Classes:
|
|
6
|
+
ConvLayer: Convolutional layer for graph data.
|
|
7
|
+
MaskedConvLayer: Convolutional layer with masking for padding indices.
|
|
8
|
+
CrystalGraphConvNet: CGCNN model for predicting material properties.
|
|
9
|
+
Normalizer: Utility class for normalizing tensors.
|
|
10
|
+
|
|
11
|
+
Usage:
|
|
12
|
+
Define your model by creating an instance of CrystalGraphConvNet with the desired parameters.
|
|
13
|
+
Use the Normalizer class to normalize your target properties during training.
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
import torch
|
|
17
|
+
import torch.nn as nn
|
|
18
|
+
from typing import List, Dict, Tuple
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class ConvLayer(nn.Module):
|
|
22
|
+
"""
|
|
23
|
+
Convolutional layer for graph data.
|
|
24
|
+
|
|
25
|
+
Performs a convolutional operation on graphs, updating atom features based on their neighbors.
|
|
26
|
+
"""
|
|
27
|
+
|
|
28
|
+
def __init__(self, atom_fea_len, nbr_fea_len):
|
|
29
|
+
"""
|
|
30
|
+
Initialize the ConvLayer.
|
|
31
|
+
|
|
32
|
+
Args:
|
|
33
|
+
atom_feature_len (int): Number of atom hidden features.
|
|
34
|
+
neighbor_feature_len (int): Number of bond (neighbor) features.
|
|
35
|
+
"""
|
|
36
|
+
super(ConvLayer, self).__init__()
|
|
37
|
+
self.atom_fea_len = atom_fea_len
|
|
38
|
+
self.nbr_fea_len = nbr_fea_len
|
|
39
|
+
self.fc_full = nn.Linear(
|
|
40
|
+
2 * self.atom_fea_len + self.nbr_fea_len, 2 * self.atom_fea_len
|
|
41
|
+
)
|
|
42
|
+
self.sigmoid = nn.Sigmoid()
|
|
43
|
+
self.softplus1 = nn.Softplus()
|
|
44
|
+
self.bn1 = nn.BatchNorm1d(2 * self.atom_fea_len)
|
|
45
|
+
self.bn2 = nn.BatchNorm1d(self.atom_fea_len)
|
|
46
|
+
self.softplus2 = nn.Softplus()
|
|
47
|
+
|
|
48
|
+
def forward(self, atom_in_fea, nbr_fea, nbr_fea_idx):
|
|
49
|
+
"""
|
|
50
|
+
Forward pass
|
|
51
|
+
|
|
52
|
+
N: Total number of atoms in the batch
|
|
53
|
+
M: Max number of neighbors
|
|
54
|
+
|
|
55
|
+
Parameters
|
|
56
|
+
----------
|
|
57
|
+
|
|
58
|
+
atom_in_fea: Variable(torch.Tensor) shape (N, atom_fea_len)
|
|
59
|
+
Atom hidden features before convolution
|
|
60
|
+
nbr_fea: Variable(torch.Tensor) shape (N, M, nbr_fea_len)
|
|
61
|
+
Bond features of each atom's M neighbors
|
|
62
|
+
nbr_fea_idx: torch.LongTensor shape (N, M)
|
|
63
|
+
Indices of M neighbors of each atom
|
|
64
|
+
|
|
65
|
+
Returns
|
|
66
|
+
-------
|
|
67
|
+
|
|
68
|
+
atom_out_fea: nn.Variable shape (N, atom_fea_len)
|
|
69
|
+
Atom hidden features after convolution
|
|
70
|
+
|
|
71
|
+
"""
|
|
72
|
+
N, M = nbr_fea_idx.shape
|
|
73
|
+
# convolution
|
|
74
|
+
atom_nbr_fea = atom_in_fea[nbr_fea_idx, :]
|
|
75
|
+
total_nbr_fea = torch.cat(
|
|
76
|
+
[
|
|
77
|
+
atom_in_fea.unsqueeze(1).expand(N, M, self.atom_fea_len),
|
|
78
|
+
atom_nbr_fea,
|
|
79
|
+
nbr_fea,
|
|
80
|
+
],
|
|
81
|
+
dim=2,
|
|
82
|
+
)
|
|
83
|
+
total_gated_fea = self.fc_full(total_nbr_fea)
|
|
84
|
+
total_gated_fea = self.bn1(
|
|
85
|
+
total_gated_fea.view(-1, self.atom_fea_len * 2)
|
|
86
|
+
).view(N, M, self.atom_fea_len * 2)
|
|
87
|
+
nbr_filter, nbr_core = total_gated_fea.chunk(2, dim=2)
|
|
88
|
+
nbr_filter = self.sigmoid(nbr_filter)
|
|
89
|
+
nbr_core = self.softplus1(nbr_core)
|
|
90
|
+
nbr_sumed = torch.sum(nbr_filter * nbr_core, dim=1)
|
|
91
|
+
nbr_sumed = self.bn2(nbr_sumed)
|
|
92
|
+
out = self.softplus2(atom_in_fea + nbr_sumed)
|
|
93
|
+
return out
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
class MaskedConvLayer(nn.Module):
|
|
97
|
+
"""
|
|
98
|
+
Convolutional operation on graphs
|
|
99
|
+
"""
|
|
100
|
+
|
|
101
|
+
def __init__(self, atom_fea_len, nbr_fea_len):
|
|
102
|
+
"""
|
|
103
|
+
Initialize ConvLayer.
|
|
104
|
+
|
|
105
|
+
Parameters
|
|
106
|
+
----------
|
|
107
|
+
|
|
108
|
+
atom_fea_len: int
|
|
109
|
+
Number of atom hidden features.
|
|
110
|
+
nbr_fea_len: int
|
|
111
|
+
Number of bond features.
|
|
112
|
+
"""
|
|
113
|
+
super(MaskedConvLayer, self).__init__()
|
|
114
|
+
self.atom_fea_len = atom_fea_len
|
|
115
|
+
self.nbr_fea_len = nbr_fea_len
|
|
116
|
+
self.fc_full = nn.Linear(
|
|
117
|
+
2 * self.atom_fea_len + self.nbr_fea_len, 2 * self.atom_fea_len
|
|
118
|
+
)
|
|
119
|
+
self.sigmoid = nn.Sigmoid()
|
|
120
|
+
self.softplus1 = nn.Softplus()
|
|
121
|
+
self.bn1 = nn.BatchNorm1d(2 * self.atom_fea_len)
|
|
122
|
+
self.bn2 = nn.BatchNorm1d(self.atom_fea_len)
|
|
123
|
+
self.softplus2 = nn.Softplus()
|
|
124
|
+
|
|
125
|
+
def forward(self, atom_in_fea, nbr_fea, nbr_fea_idx):
|
|
126
|
+
"""
|
|
127
|
+
Forward pass
|
|
128
|
+
|
|
129
|
+
N: Total number of atoms in the batch
|
|
130
|
+
M: Max number of neighbors
|
|
131
|
+
|
|
132
|
+
Parameters
|
|
133
|
+
----------
|
|
134
|
+
|
|
135
|
+
atom_in_fea: Variable(torch.Tensor) shape (N, atom_fea_len)
|
|
136
|
+
Atom hidden features before convolution
|
|
137
|
+
nbr_fea: Variable(torch.Tensor) shape (N, M, nbr_fea_len)
|
|
138
|
+
Bond features of each atom's M neighbors
|
|
139
|
+
nbr_fea_idx: torch.LongTensor shape (N, M)
|
|
140
|
+
Indices of M neighbors of each atom
|
|
141
|
+
|
|
142
|
+
Returns
|
|
143
|
+
-------
|
|
144
|
+
|
|
145
|
+
atom_out_fea: nn.Variable shape (N, atom_fea_len)
|
|
146
|
+
Atom hidden features after convolution
|
|
147
|
+
|
|
148
|
+
"""
|
|
149
|
+
N, M = nbr_fea_idx.shape
|
|
150
|
+
# Create a mask for valid neighbor indices
|
|
151
|
+
mask = nbr_fea_idx != 0 # Shape: (N, M)
|
|
152
|
+
# Get neighbor atom features
|
|
153
|
+
atom_nbr_fea = atom_in_fea[nbr_fea_idx, :] # Shape: (N, M, atom_fea_len)
|
|
154
|
+
# Zero out features corresponding to padding indices
|
|
155
|
+
mask = mask.unsqueeze(-1).float()
|
|
156
|
+
atom_nbr_fea = atom_nbr_fea * mask
|
|
157
|
+
nbr_fea = nbr_fea * mask
|
|
158
|
+
# Continue with your original code
|
|
159
|
+
total_nbr_fea = torch.cat(
|
|
160
|
+
[
|
|
161
|
+
atom_in_fea.unsqueeze(1).expand(N, M, self.atom_fea_len),
|
|
162
|
+
atom_nbr_fea,
|
|
163
|
+
nbr_fea,
|
|
164
|
+
],
|
|
165
|
+
dim=2,
|
|
166
|
+
)
|
|
167
|
+
total_gated_fea = self.fc_full(total_nbr_fea)
|
|
168
|
+
total_gated_fea = self.bn1(
|
|
169
|
+
total_gated_fea.view(-1, self.atom_fea_len * 2)
|
|
170
|
+
).view(N, M, self.atom_fea_len * 2)
|
|
171
|
+
nbr_filter, nbr_core = total_gated_fea.chunk(2, dim=2)
|
|
172
|
+
nbr_filter = self.sigmoid(nbr_filter)
|
|
173
|
+
nbr_core = self.softplus1(nbr_core)
|
|
174
|
+
# Zero out contributions from padding indices
|
|
175
|
+
nbr_filter = nbr_filter * mask
|
|
176
|
+
nbr_core = nbr_core * mask
|
|
177
|
+
nbr_sumed = torch.sum(nbr_filter * nbr_core, dim=1)
|
|
178
|
+
nbr_sumed = self.bn2(nbr_sumed)
|
|
179
|
+
out = self.softplus2(atom_in_fea + nbr_sumed)
|
|
180
|
+
return out
|
|
181
|
+
|
|
182
|
+
|
|
183
|
+
class CrystalGraphConvNet(nn.Module):
|
|
184
|
+
"""
|
|
185
|
+
Create a crystal graph convolutional neural network for predicting total
|
|
186
|
+
material properties.
|
|
187
|
+
"""
|
|
188
|
+
|
|
189
|
+
def __init__(
|
|
190
|
+
self,
|
|
191
|
+
orig_atom_fea_len,
|
|
192
|
+
nbr_fea_len,
|
|
193
|
+
atom_fea_len=64,
|
|
194
|
+
n_conv=3,
|
|
195
|
+
h_fea_len=128,
|
|
196
|
+
n_h=1,
|
|
197
|
+
classification=False,
|
|
198
|
+
):
|
|
199
|
+
"""
|
|
200
|
+
Initialize CrystalGraphConvNet.
|
|
201
|
+
|
|
202
|
+
Parameters
|
|
203
|
+
----------
|
|
204
|
+
|
|
205
|
+
orig_atom_fea_len: int
|
|
206
|
+
Number of atom features in the input.
|
|
207
|
+
nbr_fea_len: int
|
|
208
|
+
Number of bond features.
|
|
209
|
+
atom_fea_len: int
|
|
210
|
+
Number of hidden atom features in the convolutional layers
|
|
211
|
+
n_conv: int
|
|
212
|
+
Number of convolutional layers
|
|
213
|
+
h_fea_len: int
|
|
214
|
+
Number of hidden features after pooling
|
|
215
|
+
n_h: int
|
|
216
|
+
Number of hidden layers after pooling
|
|
217
|
+
"""
|
|
218
|
+
super(CrystalGraphConvNet, self).__init__()
|
|
219
|
+
self.classification = classification
|
|
220
|
+
self.embedding = nn.Linear(orig_atom_fea_len, atom_fea_len)
|
|
221
|
+
self.convs = nn.ModuleList(
|
|
222
|
+
[
|
|
223
|
+
ConvLayer(atom_fea_len=atom_fea_len, nbr_fea_len=nbr_fea_len)
|
|
224
|
+
for _ in range(n_conv)
|
|
225
|
+
]
|
|
226
|
+
)
|
|
227
|
+
self.conv_to_fc = nn.Linear(atom_fea_len, h_fea_len)
|
|
228
|
+
self.conv_to_fc_softplus = nn.Softplus()
|
|
229
|
+
if n_h > 1:
|
|
230
|
+
self.fcs = nn.ModuleList(
|
|
231
|
+
[nn.Linear(h_fea_len, h_fea_len) for _ in range(n_h - 1)]
|
|
232
|
+
)
|
|
233
|
+
self.softpluses = nn.ModuleList([nn.Softplus() for _ in range(n_h - 1)])
|
|
234
|
+
|
|
235
|
+
if self.classification:
|
|
236
|
+
self.fc_out = nn.Linear(h_fea_len, 2)
|
|
237
|
+
else:
|
|
238
|
+
self.fc_out = nn.Linear(h_fea_len, 1)
|
|
239
|
+
|
|
240
|
+
if self.classification:
|
|
241
|
+
self.logsoftmax = nn.LogSoftmax(dim=1)
|
|
242
|
+
self.dropout = nn.Dropout()
|
|
243
|
+
|
|
244
|
+
def forward(self, atom_fea, nbr_fea, nbr_fea_idx, crystal_atom_idx):
|
|
245
|
+
"""
|
|
246
|
+
Forward pass
|
|
247
|
+
|
|
248
|
+
N: Total number of atoms in the batch
|
|
249
|
+
M: Max number of neighbors
|
|
250
|
+
N0: Total number of crystals in the batch
|
|
251
|
+
|
|
252
|
+
Parameters
|
|
253
|
+
----------
|
|
254
|
+
|
|
255
|
+
atom_fea: Variable(torch.Tensor) shape (N, orig_atom_fea_len)
|
|
256
|
+
Atom features from atom type
|
|
257
|
+
nbr_fea: Variable(torch.Tensor) shape (N, M, nbr_fea_len)
|
|
258
|
+
Bond features of each atom's M neighbors
|
|
259
|
+
nbr_fea_idx: torch.LongTensor shape (N, M)
|
|
260
|
+
Indices of M neighbors of each atom
|
|
261
|
+
crystal_atom_idx: list of torch.LongTensor of length N0
|
|
262
|
+
Mapping from the crystal idx to atom idx
|
|
263
|
+
|
|
264
|
+
Returns
|
|
265
|
+
-------
|
|
266
|
+
|
|
267
|
+
prediction: nn.Variable shape (N, )
|
|
268
|
+
Atom hidden features after convolution
|
|
269
|
+
|
|
270
|
+
"""
|
|
271
|
+
atom_fea = self.embedding(atom_fea)
|
|
272
|
+
for conv_func in self.convs:
|
|
273
|
+
atom_fea = conv_func(atom_fea, nbr_fea, nbr_fea_idx)
|
|
274
|
+
crys_fea = self.pooling(atom_fea, crystal_atom_idx)
|
|
275
|
+
crys_fea = self.conv_to_fc(self.conv_to_fc_softplus(crys_fea))
|
|
276
|
+
crys_fea = self.conv_to_fc_softplus(crys_fea)
|
|
277
|
+
if self.classification:
|
|
278
|
+
crys_fea = self.dropout(crys_fea)
|
|
279
|
+
if hasattr(self, "fcs") and hasattr(self, "softpluses"):
|
|
280
|
+
for fc, softplus in zip(self.fcs, self.softpluses):
|
|
281
|
+
crys_fea = softplus(fc(crys_fea))
|
|
282
|
+
out = self.fc_out(crys_fea)
|
|
283
|
+
if self.classification:
|
|
284
|
+
out = self.logsoftmax(out)
|
|
285
|
+
return out, crys_fea
|
|
286
|
+
|
|
287
|
+
def pooling(self, atom_fea, crystal_atom_idx):
|
|
288
|
+
"""
|
|
289
|
+
Pooling the atom features to crystal features
|
|
290
|
+
|
|
291
|
+
N: Total number of atoms in the batch
|
|
292
|
+
N0: Total number of crystals in the batch
|
|
293
|
+
|
|
294
|
+
Parameters
|
|
295
|
+
----------
|
|
296
|
+
|
|
297
|
+
atom_fea: Variable(torch.Tensor) shape (N, atom_fea_len)
|
|
298
|
+
Atom feature vectors of the batch
|
|
299
|
+
crystal_atom_idx: list of torch.LongTensor of length N0
|
|
300
|
+
Mapping from the crystal idx to atom idx
|
|
301
|
+
"""
|
|
302
|
+
assert (
|
|
303
|
+
sum([len(idx_map) for idx_map in crystal_atom_idx])
|
|
304
|
+
== atom_fea.data.shape[0]
|
|
305
|
+
)
|
|
306
|
+
summed_fea = [
|
|
307
|
+
torch.mean(atom_fea[idx_map], dim=0, keepdim=True)
|
|
308
|
+
for idx_map in crystal_atom_idx
|
|
309
|
+
]
|
|
310
|
+
return torch.cat(summed_fea, dim=0)
|
|
311
|
+
|
|
312
|
+
|
|
313
|
+
class Normalizer:
|
|
314
|
+
"""
|
|
315
|
+
Normalizes a PyTorch tensor and allows restoring it later.
|
|
316
|
+
|
|
317
|
+
This class keeps track of the mean and standard deviation of a tensor and provides methods
|
|
318
|
+
to normalize and denormalize tensors using these statistics.
|
|
319
|
+
"""
|
|
320
|
+
|
|
321
|
+
def __init__(self, tensor: torch.Tensor):
|
|
322
|
+
"""
|
|
323
|
+
Initialize the Normalizer with a sample tensor to calculate mean and standard deviation.
|
|
324
|
+
|
|
325
|
+
Args:
|
|
326
|
+
tensor (torch.Tensor): Sample tensor to compute mean and standard deviation.
|
|
327
|
+
"""
|
|
328
|
+
self.mean: torch.Tensor = torch.mean(tensor)
|
|
329
|
+
self.std: torch.Tensor = torch.std(tensor)
|
|
330
|
+
|
|
331
|
+
def norm(self, tensor: torch.Tensor) -> torch.Tensor:
|
|
332
|
+
"""
|
|
333
|
+
Normalize a tensor using the stored mean and standard deviation.
|
|
334
|
+
|
|
335
|
+
Args:
|
|
336
|
+
tensor (torch.Tensor): Tensor to normalize.
|
|
337
|
+
|
|
338
|
+
Returns:
|
|
339
|
+
torch.Tensor: Normalized tensor.
|
|
340
|
+
"""
|
|
341
|
+
return (tensor - self.mean) / self.std
|
|
342
|
+
|
|
343
|
+
def denorm(self, normed_tensor: torch.Tensor) -> torch.Tensor:
|
|
344
|
+
"""
|
|
345
|
+
Denormalize a tensor using the stored mean and standard deviation.
|
|
346
|
+
|
|
347
|
+
Args:
|
|
348
|
+
normed_tensor (torch.Tensor): Normalized tensor to denormalize.
|
|
349
|
+
|
|
350
|
+
Returns:
|
|
351
|
+
torch.Tensor: Denormalized tensor.
|
|
352
|
+
"""
|
|
353
|
+
return normed_tensor * self.std + self.mean
|
|
354
|
+
|
|
355
|
+
def state_dict(self) -> Dict[str, torch.Tensor]:
|
|
356
|
+
"""
|
|
357
|
+
Returns the state dictionary containing the mean and standard deviation.
|
|
358
|
+
|
|
359
|
+
Returns:
|
|
360
|
+
Dict[str, torch.Tensor]: State dictionary.
|
|
361
|
+
"""
|
|
362
|
+
return {"mean": self.mean, "std": self.std}
|
|
363
|
+
|
|
364
|
+
def load_state_dict(self, state_dict: Dict[str, torch.Tensor]):
|
|
365
|
+
"""
|
|
366
|
+
Loads the mean and standard deviation from a state dictionary.
|
|
367
|
+
|
|
368
|
+
Args:
|
|
369
|
+
state_dict (Dict[str, torch.Tensor]): State dictionary containing 'mean' and 'std'.
|
|
370
|
+
"""
|
|
371
|
+
self.mean = state_dict["mean"]
|
|
372
|
+
self.std = state_dict["std"]
|
|
@@ -0,0 +1,474 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import csv
|
|
3
|
+
import sys
|
|
4
|
+
import glob
|
|
5
|
+
import torch
|
|
6
|
+
import argparse
|
|
7
|
+
|
|
8
|
+
import pandas as pd
|
|
9
|
+
import matplotlib.pyplot as plt
|
|
10
|
+
|
|
11
|
+
from datetime import datetime
|
|
12
|
+
from sklearn.metrics import mean_squared_error, r2_score
|
|
13
|
+
from pymatviz import density_hexbin
|
|
14
|
+
|
|
15
|
+
from torch.utils.data import DataLoader
|
|
16
|
+
from .cgcnn_data import CIFData_pred, collate_pool
|
|
17
|
+
from .cgcnn_model import CrystalGraphConvNet, Normalizer
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def output_id_gen():
|
|
21
|
+
"""
|
|
22
|
+
This function obtains the current date and time, formats it as 'mmdd_HHMM',
|
|
23
|
+
and prepends 'output_' to form a unique identifier. This can be useful
|
|
24
|
+
for creating distinct output folder names or filenames at runtime.
|
|
25
|
+
|
|
26
|
+
Returns:
|
|
27
|
+
- str: A string that represents the current date and time in the format of 'output_mmdd_HHMM'.
|
|
28
|
+
"""
|
|
29
|
+
|
|
30
|
+
now = datetime.now()
|
|
31
|
+
# Format time to match desired format (mmdd_HHMM)
|
|
32
|
+
timestamp = now.strftime("%m%d_%H%M")
|
|
33
|
+
# Prepend 'output_' to timestamp to form folder name
|
|
34
|
+
folder_name = f"output_{timestamp}"
|
|
35
|
+
|
|
36
|
+
return folder_name
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def get_lr(optimizer):
|
|
40
|
+
"""
|
|
41
|
+
This function iterates over the parameter groups of a given PyTorch optimizer,
|
|
42
|
+
extracting the learning rate from each group. The learning rates are then returned in a list.
|
|
43
|
+
|
|
44
|
+
Parameters:
|
|
45
|
+
- optimizer (torch.optim.Optimizer): The PyTorch optimizer to extract learning rates from.
|
|
46
|
+
|
|
47
|
+
Returns:
|
|
48
|
+
- list: A list of learning rates, one for each parameter group in the optimizer.
|
|
49
|
+
"""
|
|
50
|
+
|
|
51
|
+
return [param_group["lr"] for param_group in optimizer.param_groups]
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def extract_fea(model, loader, device):
|
|
55
|
+
"""
|
|
56
|
+
Applies a trained model to a dataset to extract learned feature
|
|
57
|
+
representations, targets, and CIF IDs, returning these as tensors.
|
|
58
|
+
|
|
59
|
+
Parameters:
|
|
60
|
+
- model (torch.nn.Module): The trained model.
|
|
61
|
+
- loader (torch.utils.data.DataLoader): DataLoader for the dataset.
|
|
62
|
+
- device (str): The device ('cuda' or 'cpu') to send tensors to.
|
|
63
|
+
|
|
64
|
+
Returns:
|
|
65
|
+
- tuple (torch.Tensor, torch.Tensor, list): A tuple where the first element is
|
|
66
|
+
the tensor of extracted features, the second element is the tensor of targets,
|
|
67
|
+
and the third is a list of CIF IDs.
|
|
68
|
+
"""
|
|
69
|
+
|
|
70
|
+
crys_fea_list, target_list, cif_id_list = [], [], []
|
|
71
|
+
|
|
72
|
+
with torch.no_grad():
|
|
73
|
+
for inputs, target, cif_id in loader:
|
|
74
|
+
inputs = [
|
|
75
|
+
item.to(device) if torch.is_tensor(item) else item for item in inputs
|
|
76
|
+
]
|
|
77
|
+
target = target.to(device)
|
|
78
|
+
|
|
79
|
+
_, crys_fea = model(*inputs)
|
|
80
|
+
|
|
81
|
+
crys_fea_list.append(crys_fea)
|
|
82
|
+
target_list.append(target)
|
|
83
|
+
cif_id_list.append(cif_id)
|
|
84
|
+
|
|
85
|
+
crys_fea = torch.cat(crys_fea_list, dim=0)
|
|
86
|
+
target = torch.cat(target_list, dim=0)
|
|
87
|
+
|
|
88
|
+
return crys_fea, target, cif_id_list
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def id_prop_gen(cif_dir):
|
|
92
|
+
cif_list = glob.glob(f"{cif_dir}/*.cif")
|
|
93
|
+
|
|
94
|
+
id_prop_cif = pd.DataFrame(
|
|
95
|
+
{
|
|
96
|
+
"id": [
|
|
97
|
+
os.path.basename(cif).split(".")[0] for cif in cif_list
|
|
98
|
+
],
|
|
99
|
+
"prop": [0 for _ in range(len(cif_list))],
|
|
100
|
+
}
|
|
101
|
+
)
|
|
102
|
+
|
|
103
|
+
id_prop_cif.to_csv(
|
|
104
|
+
f"{cif_dir}/id_prop.csv",
|
|
105
|
+
index=False,
|
|
106
|
+
header=False,
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def test_model(
|
|
111
|
+
model,
|
|
112
|
+
loader,
|
|
113
|
+
device,
|
|
114
|
+
plot_file="parity_plot.svg",
|
|
115
|
+
results_file="results.csv",
|
|
116
|
+
plot_mode=2,
|
|
117
|
+
):
|
|
118
|
+
"""
|
|
119
|
+
This function tests a trained machine learning model on a provided dataset, calculates the Mean Squared Error (
|
|
120
|
+
MSE) and R2 score, and prints these results. It also saves the prediction results as a CSV file and generates a
|
|
121
|
+
parity plot as an SVG file. The plot displays the model's predictions versus the actual values, color-coded by
|
|
122
|
+
the point density.
|
|
123
|
+
|
|
124
|
+
Parameters:
|
|
125
|
+
- model (torch.nn.Module): The trained model.
|
|
126
|
+
- loader (torch.utils.data.DataLoader): DataLoader for the dataset.
|
|
127
|
+
- device (str): The device ('cuda' or 'cpu') where the model will be run.
|
|
128
|
+
- plot_file (str, optional): The file path where the parity plot will be saved. Defaults to 'parity_plot.svg'.
|
|
129
|
+
- results_file (str, optional): The file path where the results will be saved as a CSV file. Defaults to 'results.csv'.
|
|
130
|
+
- plot_mode (int, optional): The mode for the parity plot. Set to 1 for scatter plot or 2 for density plot. Defaults to 2.
|
|
131
|
+
"""
|
|
132
|
+
|
|
133
|
+
model.eval()
|
|
134
|
+
targets_list = []
|
|
135
|
+
outputs_list = []
|
|
136
|
+
|
|
137
|
+
with torch.no_grad():
|
|
138
|
+
for input, target, cif_id in loader:
|
|
139
|
+
atom_fea, nbr_fea, nbr_fea_idx, crystal_atom_idx = input
|
|
140
|
+
atom_fea = atom_fea.to(device)
|
|
141
|
+
nbr_fea = nbr_fea.to(device)
|
|
142
|
+
nbr_fea_idx = nbr_fea_idx.to(device)
|
|
143
|
+
crystal_atom_idx = [idx_map.to(device) for idx_map in crystal_atom_idx]
|
|
144
|
+
target = target.to(device)
|
|
145
|
+
output, _ = model(atom_fea, nbr_fea, nbr_fea_idx, crystal_atom_idx)
|
|
146
|
+
targets_list.extend(target.cpu().numpy().ravel().tolist())
|
|
147
|
+
outputs_list.extend(output.cpu().numpy().ravel().tolist())
|
|
148
|
+
|
|
149
|
+
mse = mean_squared_error(targets_list, outputs_list)
|
|
150
|
+
r2 = r2_score(targets_list, outputs_list)
|
|
151
|
+
print(f"MSE: {mse:.4f}, R2 Score: {r2:.4f}")
|
|
152
|
+
|
|
153
|
+
# Save results to csv
|
|
154
|
+
with open(results_file, "w", newline="") as file:
|
|
155
|
+
writer = csv.writer(file)
|
|
156
|
+
writer.writerow(["cif_id", "Actual", "Predicted"])
|
|
157
|
+
writer.writerows(zip(cif_id, targets_list, outputs_list))
|
|
158
|
+
print(f"Prediction results have been saved to {results_file}")
|
|
159
|
+
|
|
160
|
+
# Generate parity plot
|
|
161
|
+
fig, ax = plt.subplots(figsize=(8, 6))
|
|
162
|
+
|
|
163
|
+
if plot_mode == 1:
|
|
164
|
+
ax.scatter(
|
|
165
|
+
targets_list, outputs_list, alpha=0.6, s=50, edgecolor="none", color="blue"
|
|
166
|
+
)
|
|
167
|
+
|
|
168
|
+
plt.plot(
|
|
169
|
+
[min(targets_list), max(targets_list)],
|
|
170
|
+
[min(targets_list), max(targets_list)],
|
|
171
|
+
"r--",
|
|
172
|
+
)
|
|
173
|
+
|
|
174
|
+
ax.xlabel("Actual", fontsize=14)
|
|
175
|
+
ax.ylabel("Predicted", fontsize=14)
|
|
176
|
+
ax.title(f"Parity Plot (R2={r2:.4f}, MSE={mse:.4f})", fontsize=16)
|
|
177
|
+
ax.grid(True)
|
|
178
|
+
|
|
179
|
+
elif plot_mode == 2:
|
|
180
|
+
# Density plot using pymatviz
|
|
181
|
+
df = pd.DataFrame({"Actual": targets_list, "Predicted": outputs_list})
|
|
182
|
+
density_hexbin("Actual", "Predicted", df=df, ax=ax, xlabel="Actual", ylabel="Predicted")
|
|
183
|
+
|
|
184
|
+
plt.tight_layout()
|
|
185
|
+
plt.savefig(plot_file, format="svg")
|
|
186
|
+
print(f"Parity plot has been saved to {plot_file}")
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
def predict_model(
|
|
190
|
+
model,
|
|
191
|
+
loader,
|
|
192
|
+
device,
|
|
193
|
+
verbose,
|
|
194
|
+
plot_file="parity_plot.svg",
|
|
195
|
+
results_file="results.csv",
|
|
196
|
+
):
|
|
197
|
+
"""
|
|
198
|
+
This function tests a trained machine learning model on a provided dataset, calculates the Mean Squared Error (
|
|
199
|
+
MSE) and R2 score, and prints these results. It also saves the prediction results as a CSV file and generates a
|
|
200
|
+
parity plot as an SVG file. The plot displays the model's predictions versus the actual values, color-coded by
|
|
201
|
+
the point density.
|
|
202
|
+
|
|
203
|
+
Parameters:
|
|
204
|
+
- model (torch.nn.Module): The trained model.
|
|
205
|
+
- loader (torch.utils.data.DataLoader): DataLoader for the dataset.
|
|
206
|
+
- device (str): The device ('cuda' or 'cpu') where the model will be run.
|
|
207
|
+
- plot_file (str, optional): The file path where the parity plot will be saved. Defaults to 'parity_plot.svg'.
|
|
208
|
+
- results_file (str, optional): The file path where the results will be saved as a CSV file. Defaults to 'results.csv'.
|
|
209
|
+
"""
|
|
210
|
+
|
|
211
|
+
model.eval()
|
|
212
|
+
targets_list = []
|
|
213
|
+
outputs_list = []
|
|
214
|
+
crys_feas_list = []
|
|
215
|
+
index = 0
|
|
216
|
+
|
|
217
|
+
with torch.no_grad():
|
|
218
|
+
for input, target, cif_id in loader:
|
|
219
|
+
atom_fea, nbr_fea, nbr_fea_idx, crystal_atom_idx = input
|
|
220
|
+
atom_fea = atom_fea.to(device)
|
|
221
|
+
nbr_fea = nbr_fea.to(device)
|
|
222
|
+
nbr_fea_idx = nbr_fea_idx.to(device)
|
|
223
|
+
crystal_atom_idx = [idx_map.to(device) for idx_map in crystal_atom_idx]
|
|
224
|
+
target = target.to(device)
|
|
225
|
+
|
|
226
|
+
output, crys_fea = model(atom_fea, nbr_fea, nbr_fea_idx, crystal_atom_idx)
|
|
227
|
+
|
|
228
|
+
targets_list.extend(target.cpu().numpy().ravel().tolist())
|
|
229
|
+
outputs_list.extend(output.cpu().numpy().ravel().tolist())
|
|
230
|
+
crys_feas_list.append(crys_fea.cpu().numpy())
|
|
231
|
+
|
|
232
|
+
index += 1
|
|
233
|
+
|
|
234
|
+
# Extract the actual values from cif_id and output tensor
|
|
235
|
+
cif_id_value = cif_id[0] if cif_id and isinstance(cif_id, list) else cif_id
|
|
236
|
+
prediction_value = output.item() if output.numel() == 1 else output.tolist()
|
|
237
|
+
|
|
238
|
+
if verbose >= 3:
|
|
239
|
+
print(
|
|
240
|
+
"index:",
|
|
241
|
+
index,
|
|
242
|
+
"| cif id:",
|
|
243
|
+
cif_id_value,
|
|
244
|
+
"| prediction:",
|
|
245
|
+
prediction_value,
|
|
246
|
+
)
|
|
247
|
+
|
|
248
|
+
return outputs_list, crys_feas_list
|
|
249
|
+
|
|
250
|
+
|
|
251
|
+
def cgcnn_pred(
|
|
252
|
+
model_path, all_set, verbose=3, cuda=False, num_workers=0
|
|
253
|
+
):
|
|
254
|
+
if not os.path.isfile(model_path):
|
|
255
|
+
raise FileNotFoundError(f"=> No model params found at '{model_path}'")
|
|
256
|
+
|
|
257
|
+
total_dataset = CIFData_pred(all_set)
|
|
258
|
+
|
|
259
|
+
checkpoint = torch.load(
|
|
260
|
+
model_path,
|
|
261
|
+
map_location=lambda storage, loc: storage if not cuda else None,
|
|
262
|
+
weights_only=False,
|
|
263
|
+
)
|
|
264
|
+
structures, _, _ = total_dataset[0]
|
|
265
|
+
orig_atom_fea_len = structures[0].shape[-1]
|
|
266
|
+
nbr_fea_len = structures[1].shape[-1]
|
|
267
|
+
model_args = argparse.Namespace(**checkpoint["args"])
|
|
268
|
+
model = CrystalGraphConvNet(
|
|
269
|
+
orig_atom_fea_len,
|
|
270
|
+
nbr_fea_len,
|
|
271
|
+
atom_fea_len=model_args.atom_fea_len,
|
|
272
|
+
n_conv=model_args.n_conv,
|
|
273
|
+
h_fea_len=model_args.h_fea_len,
|
|
274
|
+
n_h=model_args.n_h,
|
|
275
|
+
)
|
|
276
|
+
if cuda:
|
|
277
|
+
model.cuda()
|
|
278
|
+
|
|
279
|
+
normalizer = Normalizer(torch.zeros(3))
|
|
280
|
+
normalizer.load_state_dict(checkpoint["normalizer"])
|
|
281
|
+
model.load_state_dict(checkpoint["state_dict"])
|
|
282
|
+
|
|
283
|
+
if verbose >= 3:
|
|
284
|
+
print(
|
|
285
|
+
f"=> Loaded model from '{model_path}' (epoch {checkpoint['epoch']}, validation error {checkpoint['best_mae_error']})"
|
|
286
|
+
)
|
|
287
|
+
|
|
288
|
+
device = "cuda" if cuda else "cpu"
|
|
289
|
+
model.to(device).eval()
|
|
290
|
+
|
|
291
|
+
full_loader = DataLoader(
|
|
292
|
+
total_dataset,
|
|
293
|
+
batch_size=1,
|
|
294
|
+
shuffle=False,
|
|
295
|
+
num_workers=num_workers,
|
|
296
|
+
collate_fn=collate_pool,
|
|
297
|
+
pin_memory=cuda,
|
|
298
|
+
)
|
|
299
|
+
|
|
300
|
+
pred, last_layer = predict_model(model, full_loader, device, verbose)
|
|
301
|
+
|
|
302
|
+
return pred, last_layer
|
|
303
|
+
|
|
304
|
+
|
|
305
|
+
def parse_arguments():
|
|
306
|
+
"""
|
|
307
|
+
Parses command-line arguments for the script.
|
|
308
|
+
"""
|
|
309
|
+
parser = argparse.ArgumentParser(
|
|
310
|
+
description="Command-line interface for the Crystal Graph Convolutional Neural Network (CGCNN) model."
|
|
311
|
+
)
|
|
312
|
+
parser.add_argument(
|
|
313
|
+
"-mp",
|
|
314
|
+
"--model-path",
|
|
315
|
+
type=str,
|
|
316
|
+
help="Path to the file containing the trained model parameters.",
|
|
317
|
+
)
|
|
318
|
+
parser.add_argument(
|
|
319
|
+
"-mp2",
|
|
320
|
+
"--model-path2",
|
|
321
|
+
type=str,
|
|
322
|
+
help="Path to the file containing the trained model 2 parameters.",
|
|
323
|
+
)
|
|
324
|
+
parser.add_argument(
|
|
325
|
+
"-as",
|
|
326
|
+
"--total-set",
|
|
327
|
+
type=str,
|
|
328
|
+
help="Path to the directory containing all CIF files for the dataset.",
|
|
329
|
+
)
|
|
330
|
+
parser.add_argument(
|
|
331
|
+
"-trs",
|
|
332
|
+
"--train-set",
|
|
333
|
+
type=str,
|
|
334
|
+
help="Path to the directory containing CIF files for the train dataset.",
|
|
335
|
+
)
|
|
336
|
+
parser.add_argument(
|
|
337
|
+
"-vs",
|
|
338
|
+
"--valid-set",
|
|
339
|
+
type=str,
|
|
340
|
+
help="Path to the directory containing CIF files for the validation dataset.",
|
|
341
|
+
)
|
|
342
|
+
parser.add_argument(
|
|
343
|
+
"-ts",
|
|
344
|
+
"--test-set",
|
|
345
|
+
type=str,
|
|
346
|
+
help="Path to the directory containing CIF files for the test dataset.",
|
|
347
|
+
)
|
|
348
|
+
parser.add_argument(
|
|
349
|
+
"-trr",
|
|
350
|
+
"--train-ratio",
|
|
351
|
+
default=0.6,
|
|
352
|
+
type=float,
|
|
353
|
+
help="The ratio of the dataset to be used for training. Default: 0.6",
|
|
354
|
+
)
|
|
355
|
+
parser.add_argument(
|
|
356
|
+
"-vr",
|
|
357
|
+
"--valid-ratio",
|
|
358
|
+
default=0.2,
|
|
359
|
+
type=float,
|
|
360
|
+
help="The ratio of the dataset to be used for validation. Default: 0.2",
|
|
361
|
+
)
|
|
362
|
+
parser.add_argument(
|
|
363
|
+
"-tr",
|
|
364
|
+
"--test-ratio",
|
|
365
|
+
default=0.2,
|
|
366
|
+
type=float,
|
|
367
|
+
help="The ratio of the dataset to be used for testing. Default: 0.2",
|
|
368
|
+
)
|
|
369
|
+
parser.add_argument(
|
|
370
|
+
"-e",
|
|
371
|
+
"--epoch",
|
|
372
|
+
default=10000,
|
|
373
|
+
type=float,
|
|
374
|
+
help="Total epochs for training the model.",
|
|
375
|
+
)
|
|
376
|
+
parser.add_argument(
|
|
377
|
+
"-sp",
|
|
378
|
+
"--stop-patience",
|
|
379
|
+
default=100,
|
|
380
|
+
type=float,
|
|
381
|
+
help="Epochs for early stopping.",
|
|
382
|
+
)
|
|
383
|
+
parser.add_argument(
|
|
384
|
+
"-lrp",
|
|
385
|
+
"--lr-patience",
|
|
386
|
+
default=0,
|
|
387
|
+
type=float,
|
|
388
|
+
help="Epochs for reducing learning rate.",
|
|
389
|
+
)
|
|
390
|
+
parser.add_argument(
|
|
391
|
+
"-lrf",
|
|
392
|
+
"--lr-factor",
|
|
393
|
+
default=0.0,
|
|
394
|
+
type=float,
|
|
395
|
+
help="Factor for reducing learning rate.",
|
|
396
|
+
)
|
|
397
|
+
parser.add_argument(
|
|
398
|
+
"-tlfc",
|
|
399
|
+
"--train-last-fc",
|
|
400
|
+
default=0,
|
|
401
|
+
type=int,
|
|
402
|
+
help="Train on the last fully connected layer or all the fully connected layers",
|
|
403
|
+
)
|
|
404
|
+
parser.add_argument(
|
|
405
|
+
"-lrfc",
|
|
406
|
+
"--lr-fc",
|
|
407
|
+
default=0.01,
|
|
408
|
+
type=float,
|
|
409
|
+
help="Learning rate for fully connected layer.",
|
|
410
|
+
)
|
|
411
|
+
parser.add_argument(
|
|
412
|
+
"-lrnfc",
|
|
413
|
+
"--lr-non-fc",
|
|
414
|
+
default=0.001,
|
|
415
|
+
type=float,
|
|
416
|
+
help="Learning rate for non-fully connected layer.",
|
|
417
|
+
)
|
|
418
|
+
parser.add_argument(
|
|
419
|
+
"-rs", "--random-seed", default=123, type=int, help="Random seed."
|
|
420
|
+
)
|
|
421
|
+
parser.add_argument(
|
|
422
|
+
"-bs",
|
|
423
|
+
"--batch-size",
|
|
424
|
+
default=256,
|
|
425
|
+
type=int,
|
|
426
|
+
metavar="N",
|
|
427
|
+
help="The size of each batch during training or testing. Default: 256",
|
|
428
|
+
)
|
|
429
|
+
parser.add_argument(
|
|
430
|
+
"-j",
|
|
431
|
+
"--workers",
|
|
432
|
+
default=0,
|
|
433
|
+
type=int,
|
|
434
|
+
metavar="N",
|
|
435
|
+
help="The number of subprocesses to use for data loading. Default: 0",
|
|
436
|
+
)
|
|
437
|
+
parser.add_argument(
|
|
438
|
+
"--disable-cuda",
|
|
439
|
+
action="store_true",
|
|
440
|
+
help="Set this flag to disable CUDA, even if it is available.",
|
|
441
|
+
)
|
|
442
|
+
parser.add_argument(
|
|
443
|
+
"-m",
|
|
444
|
+
"--mode",
|
|
445
|
+
default=1,
|
|
446
|
+
type=int,
|
|
447
|
+
help="Set to 1 to train the model, or 0 to test the model. Default: 1",
|
|
448
|
+
)
|
|
449
|
+
parser.add_argument(
|
|
450
|
+
"-ji", "--job-id", default=None, type=str, help="The id of the current job."
|
|
451
|
+
)
|
|
452
|
+
parser.add_argument(
|
|
453
|
+
"-r",
|
|
454
|
+
"--replace",
|
|
455
|
+
default=1,
|
|
456
|
+
type=int,
|
|
457
|
+
help="Replace the training layer to restart.",
|
|
458
|
+
)
|
|
459
|
+
parser.add_argument(
|
|
460
|
+
"-bt",
|
|
461
|
+
"--bias-temperature",
|
|
462
|
+
default=0.0,
|
|
463
|
+
type=float,
|
|
464
|
+
help="Bias the loss function using a Boltzmann like factor.",
|
|
465
|
+
)
|
|
466
|
+
|
|
467
|
+
args = parser.parse_args(sys.argv[1:])
|
|
468
|
+
args.cuda = not args.disable_cuda and torch.cuda.is_available()
|
|
469
|
+
|
|
470
|
+
# Warning if train ratio and test ratio don't sum to 1
|
|
471
|
+
if abs(args.train_ratio + args.valid_ratio + args.test_ratio - 1) > 1e-6:
|
|
472
|
+
print("Warning: Train ratio, Valid ratio and Test ratio do not sum up to 1")
|
|
473
|
+
|
|
474
|
+
return args
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
[tool.poetry]
|
|
2
|
+
name = "cgcnn2"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Crystal Graph Convolutional Neural Networks"
|
|
5
|
+
authors = ["Jiacheng Wang <jiachengwang@umass.edu>"]
|
|
6
|
+
maintainers = ["Jiacheng Wang"]
|
|
7
|
+
license = "MIT"
|
|
8
|
+
readme = "README.md"
|
|
9
|
+
homepage = "https://github.com/jcwang587/cgcnn2/"
|
|
10
|
+
repository = "https://github.com/jcwang587/cgcnn2/"
|
|
11
|
+
|
|
12
|
+
classifiers = [
|
|
13
|
+
"Development Status :: 1 - Planning",
|
|
14
|
+
"Intended Audience :: Science/Research",
|
|
15
|
+
"Programming Language :: Python :: 3.10",
|
|
16
|
+
"Programming Language :: Python :: 3.11",
|
|
17
|
+
"Programming Language :: Python :: 3.12",
|
|
18
|
+
"Operating System :: OS Independent",
|
|
19
|
+
]
|
|
20
|
+
|
|
21
|
+
keywords = ["python", "gnn", "vasp", "crystal"]
|
|
22
|
+
|
|
23
|
+
[tool.poetry.dependencies]
|
|
24
|
+
python = ">=3.10,<3.12"
|
|
25
|
+
ase = "3.23.0"
|
|
26
|
+
numpy = "*"
|
|
27
|
+
pandas = "*"
|
|
28
|
+
scikit-learn = "*"
|
|
29
|
+
torch = "2.5.1"
|
|
30
|
+
pymatgen = "2024.10.3"
|
|
31
|
+
pymatviz = "0.13.2"
|
|
32
|
+
|
|
33
|
+
[tool.poetry.group.dev.dependencies]
|
|
34
|
+
pytest = "8.3.3"
|
|
35
|
+
pytest-cov = "5.0.0"
|
|
36
|
+
|
|
37
|
+
[build-system]
|
|
38
|
+
requires = ["poetry-core==1.9.0"]
|
|
39
|
+
build-backend = "poetry.core.masonry.api"
|