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 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,3 @@
1
+ from .cgcnn_data import *
2
+ from .cgcnn_model import *
3
+ from .cgcnn_utils import *
@@ -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"