Skip to content

Commit 1b9f395

Browse files
authored
Update README.md
1 parent 0b6d5da commit 1b9f395

1 file changed

Lines changed: 152 additions & 3 deletions

File tree

README.md

Lines changed: 152 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
* [Introduction](#intro)
1818
* [Installation](#install)
1919
* [Examples](#example)
20+
* [Parameter information](#params)
2021
* [Pre-trained models](#pretrained)
2122
* [JARVIS-ALIGNN webapp](#webapp)
2223
* [ALIGNN-FF & ASE Calculator](#alignnff)
@@ -101,9 +102,9 @@ Examples
101102

102103
| Notebooks | Google Colab | Descriptions |
103104
| ---------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
104-
| [Regression task (grpah wise prediction)](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | [![Open in Google Colab]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | Examples for developing single output regression model for exfoliation energies of 2D materials. |
105-
| [Machine learning force-field training from scratch](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | [![Open in Google Colab]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | Examples of training a machine learning force field for Silicon. |
106-
| [ALIGNN-FF Relaxer+EV_curve+Phonons+Interface gamma_surface+Interface separation](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | [![Open in Google Colab]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | Examples of using pre-trained ALIGNN-FF force-field model. |
105+
| [Regression Example](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | [![Open in Google Colab]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | Examples for developing single output regression model for exfoliation energies of 2D materials. |
106+
| [Machine learning force-field training](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | [![Open in Google Colab]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | Examples of training a machine learning force field for Silicon. |
107+
| [Geometry optimization, EOS, Phonons, Interface](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | [![Open in Google Colab]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | Examples of using pre-trained ALIGNN-FF force-field model. |
107108
| [Scaling/timing comaprison](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Timing_uMLFF.ipynb) | [![Open in Google Colab]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Timing_uMLFF.ipynb) | Examples of analyzing scaling |
108109
| [Running MD for Melt-Quench](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) | [![Open in Google Colab]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) | Examples of making amorphous structure with moelcular dynamics. |
109110
| [Miscellaneous tasks](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) | [![Open in Google Colab]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) | Examples for developing single output (such as formation energy, bandgaps) or multi-output (such as phonon DOS, electron DOS) Regression or Classification (such as metal vs non-metal), Using several pretrained models. |
@@ -187,6 +188,154 @@ Additional example trainings for property prediction task: [2D-exfoliation energ
187188

188189
An example for training MLFF for Silicon is provided [here](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb). It is highly recommeded to get familiar with this example before developing a new model.
189190

191+
192+
193+
<a name="params"></a>
194+
ALIGNN Training Parameters Documentation
195+
----------------------------------------
196+
197+
198+
This document provides detailed descriptions of parameters used for training ALIGNN (Atomistic Line Graph Neural Network) models.
199+
200+
#### Basic Training Parameters
201+
202+
| Parameter | Default Value | Other Options | Description |
203+
|-----------|---------------|---------------|-------------|
204+
| `batch_size` | 64 | Any positive integer | Number of samples per batch during training |
205+
| `criterion` | 'mse' | 'l1', 'poisson', 'zig' | Loss function used for training |
206+
| `epochs` | 300 | Any positive integer | Total number of training epochs |
207+
| `learning_rate` | 0.01 | Any positive float | Initial learning rate for optimizer |
208+
| `optimizer` | 'adamw' | 'sgd' | Optimization algorithm |
209+
| `scheduler` | 'onecycle' | 'none' | Learning rate scheduler type |
210+
| `warmup_steps` | 2000 | Any non-negative integer | Number of warmup steps for learning rate |
211+
| `weight_decay` | 0 | Any non-negative float | L2 regularization factor |
212+
| `random_seed` | 123 | Any integer or None | Seed for random number generators to ensure reproducibility |
213+
| `progress` | True | False | Whether to display progress bars during training |
214+
215+
##### Data Configuration
216+
217+
| Parameter | Default Value | Other Options | Description |
218+
|-----------|---------------|---------------|-------------|
219+
| `dataset` | 'dft_3d' | 'jdft_3d-8-18-2021', 'dft_2d', 'megnet', 'megnet2', 'mp_3d_2020', 'qm9', 'qm9_dgl', 'qm9_std_jctc', 'user_data', 'oqmd_3d_no_cfid', 'edos_up', 'edos_pdos', 'qmof', 'qe_tb', 'hmof', 'hpov', 'pdbbind', 'pdbbind_core', 'tinnet_OH', 'tinnet_O', 'tinnet_N' | Dataset name or type |
220+
| `filename` | 'sample' | Any string | Base filename for saved models and outputs |
221+
| `target` | 'exfoliation_energy' | Multiple options including 'formation_energy_peratom', 'band_gap', etc. (see TARGET_ENUM) | Property to predict (column name in dataset) |
222+
| `id_tag` | 'jid' | 'id', '_oqmd_entry_id' | Identifier column in the dataset |
223+
| `atom_features` | 'cgcnn' | 'basic', 'atomic_number', 'cfid' | Type of atom features to use |
224+
| `n_train` | None | Any positive integer | Number of training samples |
225+
| `n_val` | None | Any positive integer | Number of validation samples |
226+
| `n_test` | None | Any positive integer | Number of test samples |
227+
| `train_ratio` | 0.8 | Any float between 0 and 1 | Fraction of data used for training |
228+
| `val_ratio` | 0.1 | Any float between 0 and 1 | Fraction of data used for validation |
229+
| `test_ratio` | 0.1 | Any float between 0 and 1 | Fraction of data used for testing |
230+
| `keep_data_order` | True | False | Whether to maintain the original order of data |
231+
| `use_canonize` | True | False | Whether to canonize molecular graphs |
232+
| `standard_scalar_and_pca` | False | True | Apply standardization and PCA to input features |
233+
| `target_multiplication_factor` | None | Any float | Factor to multiply target values |
234+
235+
#### Graph Construction Parameters
236+
237+
| Parameter | Default Value | Other Options | Description |
238+
|-----------|---------------|---------------|-------------|
239+
| `cutoff` | 8.0 | Any positive float | Cutoff radius for constructing graphs (Å) |
240+
| `cutoff_extra` | 3.0 | Any positive float | Secondary cutoff for special interactions (Å) |
241+
| `max_neighbors` | 12 | Any positive integer | Maximum number of neighbors per atom |
242+
| `neighbor_strategy` | 'k-nearest' | 'voronoi', 'radius_graph', 'radius_graph_jarvis' | Method to construct the graph |
243+
| `compute_line_graph` | True | False | Whether to precompute line graphs |
244+
245+
#### Hardware & Performance
246+
247+
| Parameter | Default Value | Other Options | Description |
248+
|-----------|---------------|---------------|-------------|
249+
| `device` | 'cpu' | 'cuda', etc. | Computing device |
250+
| `distributed` | False | True | Whether to use distributed training |
251+
| `data_parallel` | False | True | Whether to use DataParallel for multi-GPU training |
252+
| `dtype` | 'float32' | 'float64', etc. | Data type for tensors |
253+
| `num_workers` | 4 | Any non-negative integer | Number of data loading worker processes |
254+
| `pin_memory` | False | True | Whether to use pinned memory for faster GPU transfer |
255+
| `use_lmdb` | True | False | Use LMDB database for faster data loading |
256+
257+
#### Output & Checkpointing
258+
259+
| Parameter | Default Value | Other Options | Description |
260+
|-----------|---------------|---------------|-------------|
261+
| `output_dir` | Current directory | Any valid path | Directory for saving outputs |
262+
| `write_checkpoint` | True | False | Whether to save model checkpoints |
263+
| `write_predictions` | True | False | Whether to save model predictions |
264+
| `store_outputs` | True | False | Whether to store all outputs during training |
265+
| `save_dataloader` | False | True | Whether to save the dataloader for future use |
266+
| `log_tensorboard` | False | True | Whether to log metrics to TensorBoard |
267+
| `n_early_stopping` | None | Any positive integer | Number of epochs without improvement before early stopping |
268+
269+
#### Classification Settings
270+
271+
| Parameter | Default Value | Other Options | Description |
272+
|-----------|---------------|---------------|-------------|
273+
| `classification_threshold` | None | Any float | Threshold for binary classification |
274+
| `normalize_graph_level_loss` | False | True | Whether to normalize graph-level loss |
275+
276+
#### Model Architecture Parameters
277+
278+
The `model` dictionary contains parameters specific to the ALIGNN architecture:
279+
280+
#### Core Architecture
281+
282+
| Parameter | Default Value | Other Options | Description |
283+
|-----------|---------------|---------------|-------------|
284+
| `name` | 'alignn_atomwise' | 'alignn', 'ealignn_atomwise' | Model architecture name |
285+
| `gcn_layers` | 2 | Any positive integer | Number of graph convolution layers |
286+
| `alignn_layers` | 2 | Any positive integer | Number of ALIGNN layers |
287+
| `hidden_features` | 64 | Any positive integer | Size of hidden feature vectors |
288+
| `embedding_features` | 64 | Any positive integer | Size of embedding feature vectors |
289+
| `output_features` | 1 | Any positive integer | Size of output feature vectors |
290+
| `atom_input_features` | 1 | Any positive integer | Number of atom input features |
291+
| `edge_input_features` | 80 | Any positive integer | Number of edge input features |
292+
| `triplet_input_features` | 40 | Any positive integer | Number of triplet input features |
293+
294+
#### Atomistic Properties and Forces
295+
296+
| Parameter | Default Value | Other Options | Description |
297+
|-----------|---------------|---------------|-------------|
298+
| `calculate_gradient` | True | False | Whether to calculate force gradients |
299+
| `add_reverse_forces` | True | False | Whether to add reverse force contributions |
300+
| `gradwise_weight` | 1.0 | Any non-negative float | Weight for gradient (force) loss term |
301+
| `graphwise_weight` | 1.0 | Any non-negative float | Weight for graph-level (energy) loss term |
302+
| `stresswise_weight` | 0.0 | Any non-negative float | Weight for stress tensor loss term |
303+
| `atomwise_weight` | 0.0 | Any non-negative float | Weight for atom-level property loss term |
304+
| `grad_multiplier` | -1 | Any float | Multiplier for gradient values |
305+
| `stress_multiplier` | 1.0 | Any float | Multiplier for stress values |
306+
| `batch_stress` | True | False | Whether to calculate stress for batches |
307+
| `force_mult_natoms` | False | True | Whether to multiply forces by number of atoms |
308+
| `energy_mult_natoms` | True | False | Whether to multiply energy by number of atoms |
309+
310+
#### Advanced Features
311+
312+
| Parameter | Default Value | Other Options | Description |
313+
|-----------|---------------|---------------|-------------|
314+
| `backtracking` | True | False | Whether to use backtracking in edge updates |
315+
| `lighten_edges` | True | False | Whether to reduce edge feature dimensionality |
316+
| `lg_on_fly` | True | False | Calculate line graphs on the fly (saves memory) |
317+
| `link` | 'identity' | Various activation functions | Activation function type for linking layers |
318+
| `classification` | False | True | Whether the task is classification |
319+
| `zero_inflated` | False | True | Whether to use zero-inflated loss |
320+
| `use_cutoff_function` | False | True | Whether to use smooth cutoff functions |
321+
| `use_penalty` | True | False | Whether to apply penalties |
322+
| `penalty_factor` | 0.1 | Any positive float | Factor for penalty term |
323+
| `penalty_threshold` | 1.0 | Any positive float | Threshold for applying penalty |
324+
| `multiply_cutoff` | False | True | Whether to multiply features by cutoff function |
325+
| `inner_cutoff` | 4.0 | Any positive float | Inner cutoff radius (Å) |
326+
| `exponent` | 5 | Any positive integer | Exponent for cutoff function |
327+
| `include_pos_deriv` | False | True | Whether to include positional derivatives |
328+
| `additional_output_features` | 0 | Any non-negative integer | Number of additional output features |
329+
| `additional_output_weight` | 0.0 | Any non-negative float | Weight for additional output loss |
330+
| `atomwise_output_features` | 0 | Any non-negative integer | Number of atom-level output features |
331+
| `extra_features` | 0 | Any non-negative integer | Number of extra features |
332+
333+
#### Version Control
334+
335+
| Parameter | Default Value | Other Options | Description |
336+
|-----------|---------------|---------------|-------------|
337+
| `version` | Git commit hash or 'NA' | Any string | Git commit hash or version identifier |
338+
190339
<a name="pretrained"></a>
191340
Using pre-trained models
192341
-------------------------

0 commit comments

Comments
 (0)